"""Command-line interface for deep mutational scanning."""

import logging
import warnings
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Union

import click
import numpy as np
import pandas as pd

from kdpred.kd_prediction import GaussianProcessPredictor, ProdigyPredictor
from kdpred.mutations import MutationGenerator
from kdpred.structure import ColabFoldPredictor
from kdpred.utils import (
    STANDARD_AMINO_ACIDS,
    read_lines,
)


__all__ = [
    "setup_logging",
    "deep_mutational_scanning_pipeline",
]


def setup_logging(verbose: bool = False):
    """Setup logging configuration.

    Parameters
    ----------
    verbose : bool, optional
        Enable verbose (DEBUG) logging, by default False.
    """
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def deep_mutational_scanning_pipeline(
    protein_seq_fpath: Union[Path, str],
    mutation_config_fpath: Union[Path, str],
    output_dir: Union[Path, str],
    residue_list: Optional[List[str]] = None,
    protein_name: str = "",
    colabfold_cmd: str = "colabfold_batch",
    job_type: str = "gp_multiple",  # "prodigy", "gp_single", or "gp_multiple"
    num_recycles: int = 3,
    num_models: int = 5,
    temperature: float = 25.0,
    distance_cutoff: float = 5.5,
    acc_threshold: float = 0.05,
    summary_results_fpath: Optional[Union[Path, str]] = None,
):
    """Run the deep mutational scanning pipeline.

    Parameters
    ----------
    protein_sequences : List[str]
        List of protein sequences.
    mutation_strings : List[str]
        List of mutation strings.
    output_dir : Path
        Output directory for results.
    residue_list : Optional[List[str]], optional
        List of residues for saturation mutations, by default None.
    colabfold_cmd : str, optional
        ColabFold command, by default "colabfold_batch".
    prodigy_cmd : str, optional
        Prodigy command, by default "prodigy".
    num_recycles : int, optional
        Number of recycles for ColabFold, by default 3.
    num_models : int, optional
        Number of models for ColabFold, by default 5.
    job_type : str, optional
        Type of Kd prediction model to use, by default "gp". The other option is "prodigy".

    """
    output_dir = Path(output_dir).expanduser().resolve()

    if residue_list is None:
        residue_list = STANDARD_AMINO_ACIDS

    # step 1: generate mutation sequences
    protein_seq = read_lines(Path(protein_seq_fpath))
    mutation_strings = read_lines(Path(mutation_config_fpath))
    # generate the mutated sequences and save as FASTA files
    seq_gen = MutationGenerator(protein_sequence=protein_seq, protein_name=protein_name)
    mutation_list = seq_gen.expand_saturation_mutations(
        mutation_strings=mutation_strings,
        residue_list=residue_list,
    )
    fasta_paths = seq_gen.generate_all_fastas(
        mutation_list=mutation_list,
        output_dir=output_dir,
        residue_list=residue_list,
    )

    # step 2: structure prediction with colabfold for each mutant
    # mutation info
    # pdb structure location directory
    # kd_pred_mean, kd_pred_std

    df_results = pd.DataFrame(
        columns=[
            "mutation",
            "pdb_file_paths",
            "kd_mean_prodigy_μM",
            "kd_std_prodigy_μM",
            "kd_mean_gp_μM",
            "kd_std_gp_μM",
            "AF2_status",
        ]
    )
    for idx, fasta_path in enumerate(fasta_paths):
        df_results.loc[idx, "mutation"] = fasta_path.stem

        colabfold_predictor = ColabFoldPredictor(
            colabfold_command=colabfold_cmd,
            num_recycles=num_recycles,
            num_models=num_models,
            mutation_name=fasta_path.stem,
        )
        colabfold_status = colabfold_predictor.predict_structure(
            fasta_path=fasta_path,
            output_dir=Path(output_dir).resolve(),
        )

        if colabfold_status:
            df_results.loc[idx, "AF2_status"] = "Success"
            logging.info("Structure prediction succeeded for %s", fasta_path.stem)
            pdb_fpaths = colabfold_predictor.get_pdb_files(output_dir=output_dir)
            # record the location of the pdb files
            df_results.loc[idx, "pdb_file_paths"] = str(pdb_fpaths[0].parent)

            # step 3: Kd prediction with prodigy and ML models
            if job_type == "prodigy":
                prodigy_predictor = ProdigyPredictor(
                    temperature=temperature,
                    distance_cutoff=distance_cutoff,
                    acc_threshold=acc_threshold,
                )
                kd_values = prodigy_predictor.predict_kd_multiple(pdb_paths=pdb_fpaths)
                kd_pred_mean = np.mean(kd_values)
                kd_pred_std = np.std(kd_values)
                df_results.loc[idx, "kd_mean_prodigy_μM"] = kd_pred_mean
                df_results.loc[idx, "kd_std_prodigy_μM"] = kd_pred_std

            elif job_type == "gp_multiple":
                gp_predictor = GaussianProcessPredictor(
                    wdir=output_dir / fasta_path.stem,
                    model_type=job_type,
                    temperature=temperature,
                    distance_cutoff=distance_cutoff,
                    acc_threshold=acc_threshold,
                )
                kd_pred_mean, kd_pred_std = gp_predictor.predict_kd(pdb_fpath=None)
                df_results.loc[idx, "kd_mean_gp_μM"] = kd_pred_mean
                df_results.loc[idx, "kd_std_gp_μM"] = kd_pred_std
            elif job_type == "gp_single":
                gp_predictor = GaussianProcessPredictor(
                    wdir=output_dir / fasta_path.stem,
                    model_type=job_type,
                    temperature=temperature,
                    distance_cutoff=distance_cutoff,
                    acc_threshold=acc_threshold,
                )
                kd_pred_mean, kd_pred_std = gp_predictor.predict_kd(pdb_fpath=None)
                df_results.loc[idx, "kd_mean_gp_μM"] = kd_pred_mean
                df_results.loc[idx, "kd_std_gp_μM"] = kd_pred_std
            else:
                raise ValueError(
                    f"Unknown job type: {job_type}. Supported types are 'prodigy' and 'gp'."
                )

        # when the structure prediction fails
        else:
            df_results.loc[idx, "AF2_status"] = "Failed"
            df_results.loc[idx, "kd_mean_prodigy_μM"] = np.nan
            df_results.loc[idx, "kd_std_prodigy_μM"] = np.nan
            df_results.loc[idx, "kd_mean_gp_μM"] = np.nan
            df_results.loc[idx, "kd_std_gp_μM"] = np.nan
            logging.warning("Structure prediction failed for %s", fasta_path.stem)
            continue

    # step 5: save the summary results
    # when no path is provided, save to output_dir/dms_results_<date>.csv
    if summary_results_fpath is None:
        summary_results_fpath = output_dir / f"dms_results_{datetime.now().strftime('%Y%b%d')}.csv"
    # when a path is provided
    else:
        summary_results_fpath = Path(summary_results_fpath).expanduser().resolve()
        if summary_results_fpath.suffix.lower() in [".csv", ".tsv", ".txt"]:
            sep = "\t" if summary_results_fpath.suffix.lower() in [".tsv", ".txt"] else ","
            df_results.to_csv(
                summary_results_fpath,
                index=False,
                sep=sep,
            )
        elif summary_results_fpath.suffix.lower() in [".xlsx", ".xls"]:
            df_results.to_excel(
                summary_results_fpath,
                index=False,
                engine="openpyxl",
            )
        else:
            warnings.warn(
                f"Unrecognized file extension {summary_results_fpath.suffix}. Saving as CSV format."
            )
            df_results.to_csv(
                summary_results_fpath.with_suffix(".csv"),
                index=False,
                sep="\t",
            )

    logging.info("Summary results saved to %s", summary_results_fpath)

    return df_results


@click.command()
@click.option(
    "--protein-seq-fpath",
    type=click.Path(exists=True, dir_okay=False, path_type=Path),
    required=True,
    help="Path to the file containing the protein sequence.",
)
@click.option(
    "--mutation-config-fpath",
    type=click.Path(exists=True, dir_okay=False, path_type=Path),
    required=True,
    help="Path to the file containing mutation configurations.",
)
@click.option(
    "--output-dir",
    type=click.Path(file_okay=False, path_type=Path),
    required=True,
    help="Output directory for results.",
)
@click.option(
    "--residue-list",
    type=str,
    default=None,
    help="Comma-separated list of residues for saturation mutations. "
    "If not provided, standard amino acids will be used.",
)
@click.option(
    "--protein-name",
    type=str,
    default="",
    help="Name of the protein.",
)
@click.option(
    "--colabfold-cmd",
    type=str,
    default="colabfold_batch",
    help="ColabFold command to use.",
)
@click.option(
    "--job-type",
    type=click.Choice(["prodigy", "gp_single", "gp_multiple"]),
    default="gp_multiple",
    help="Type of Kd prediction model to use.",
)
@click.option(
    "--num-recycles",
    type=int,
    default=3,
    help="Number of recycles for ColabFold.",
)
@click.option(
    "--num-models",
    type=int,
    default=5,
    help="Number of models for ColabFold.",
)
@click.option(
    "--temperature",
    type=float,
    default=25.0,
    help="Temperature for Kd prediction.",
)
@click.option(
    "--distance-cutoff",
    type=float,
    default=5.5,
    help="Distance cutoff for Kd prediction.",
)
@click.option(
    "--acc-threshold",
    type=float,
    default=0.05,
    help="Accessibility threshold for Kd prediction.",
)
@click.option(
    "--summary-results-fpath",
    type=click.Path(dir_okay=False, path_type=Path),
    default=None,
    help="Path to save the summary results file.",
)
def main(
    protein_seq_fpath,
    mutation_config_fpath,
    output_dir,
    residue_list,
    protein_name,
    colabfold_cmd,
    job_type,
    num_recycles,
    num_models,
    temperature,
    distance_cutoff,
    acc_threshold,
    summary_results_fpath,
):
    """Main CLI entry point."""
    _ = deep_mutational_scanning_pipeline(
        protein_seq_fpath=protein_seq_fpath,
        mutation_config_fpath=mutation_config_fpath,
        output_dir=output_dir,
        residue_list=residue_list,
        protein_name=protein_name,
        colabfold_cmd=colabfold_cmd,
        # "prodigy", "gp_single", or "gp_multiple"
        job_type=job_type,
        num_recycles=num_recycles,
        num_models=num_models,
        temperature=temperature,
        distance_cutoff=distance_cutoff,
        acc_threshold=acc_threshold,
        summary_results_fpath=summary_results_fpath,
    )


if __name__ == "__main__":
    main()
