import argparse
import contextlib
import datetime
import errno
import json
import os
import pathlib
import tempfile
import typing
import sys

import anndata
import rich.console
import rich.progress
import numpy
import pandas
import scipy.sparse
from scipy.cluster.hierarchy import fcluster

try:
    import argcomplete
except ImportError as err:
    argcomplete = err

try:
    from rich_argparse import ArgumentDefaultsRichHelpFormatter as HelpFormatter
except ImportError:
    from argparse import ArgumentDefaultsHelpFormatter as HelpFormatter

from .. import __version__
from ..dataset.base import BaseDataset
from ..dataset.list import DatasetList
from ..mmseqs import MMSeqs, Database, Clustering
from ..pipeline import PipelineParameters, Pipeline
from ..clustering import default_strategy, ClusteringStrategy, LinearClustering, HierarchicalClustering
from .._utils import Stopwatch
from .inputs import (
    GenBankInput, 
    GenBankListInput, 
    AntiSMASHGenBankInput, 
    AntiSMASHZipInput,
    DefenseFinderTSV,
)


def build_parser(argv: typing.List[str]) -> argparse.ArgumentParser:
    show_all = '--help-all' in argv
    params = PipelineParameters.default()
    strategy = default_strategy()

    parser = argparse.ArgumentParser(
        prog="igua",
        formatter_class=HelpFormatter,
        add_help=False,
        description=(
            "A method for content-agnostic high-throughput identification of "
            "Gene Cluster Families (GCFs) from gene clusters of genomic and "
            "metagenomic origin. "
        ),
        epilog="Use --help-all to see all available options including advanced parameters." if not show_all else None,
    )

    def extended_help_text(text: str) -> str:
        return text if show_all else argparse.SUPPRESS

    def extended_parser_group(name, help) -> str:
        if show_all:
            return parser.add_argument_group(name, help)
        else:
            return parser

    parser.add_argument(
        "-h",
        "--help",
        help="Display this help message and exit.",
        action="help",
    )
    parser.add_argument(
        "--help-all",
        help="Display all options including advanced parameters.",
        action="help",
    )

    parser.add_argument(
        "-V",
        "--version",
        help="Display the program version and exit.",
        action="version",
        version=f"igua v{__version__}",
    )
    parser.add_argument(
        "-j",
        "--jobs",
        help="The number of threads to use in parallel sections.",
        type=int,
        default=os.cpu_count(),
        metavar="N",
    )

    group_input = parser.add_argument_group(
        "Input",
        "Mandatory input files required by the pipeline."
    )
    group_input.add_argument(
        "-g",
        "--gbk",
        help="Input path to GenBank files containing generic gene clusters.",
        metavar="FILE",
        dest="inputs",
        action="append",
        type=GenBankInput,
        default=[],
    )
    group_input.add_argument(
        "-l",
        "--gbk-list",
        help="Path to a file containing a list of GenBank files to process.",
        metavar="FILE",
        dest="inputs",
        action="append",
        type=GenBankListInput,
        default=[],
    )
    group_input.add_argument(
        "--antismash-gbk",
        help="Input path to GenBank files containing antiSMASH predictions.",
        metavar="FILE",
        dest="inputs",
        action="append",
        type=AntiSMASHGenBankInput,
        default=[],
    )
    group_input.add_argument(
        "--antismash-zip",
        help="Input path to Zip files containing antiSMASH predictions.",
        metavar="FILE",
        dest="inputs",
        action="append",
        type=AntiSMASHZipInput,
        default=[],
    )
    group_input.add_argument(
        "--defense-finder-tsv",
        help="Input path to TSV metadata file for DefenseFinder predictions.",
        metavar="FILE",
        dest="inputs",
        action="append",
        type=DefenseFinderTSV,
        default=[],
    )

    group_output = parser.add_argument_group(
        "Output", "Output files generated by the pipeline."
    )
    group_output.add_argument(
        "-o",
        "--output",
        help="The name of the output table to generate.",
        default=pathlib.Path("gcfs.tsv"),
        type=pathlib.Path,
    )
    group_output.add_argument(
        "-C",
        "--compositions",
        help=extended_help_text("A file where to write compositional data for GCF representatives."),
        type=pathlib.Path,
    )
    group_output.add_argument(
        "-F",
        "--features",
        help=extended_help_text("A file where to write protein cluster representatives in FASTA format."),
        type=pathlib.Path,
    )

    group_parameters = parser.add_argument_group(
        "Parameters", "General purpose parameters."
    )
    group_parameters.add_argument(
        "-w",
        "--workdir",
        help=extended_help_text("A folder to use for temporary data produced by MMSeqs2."),
    )
    group_parameters.add_argument(
        "--prefix",
        help="The prefix for GCF identifiers generated by the pipeline.",
        default="GCF",
    )

    group_mmseqs_dedup = extended_parser_group(
        "MMSeqs2 Deduplication",
        "Parameters for the first nucleotide clustering step (exact/near-exact deduplication).",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-identity",
        help=extended_help_text("Sequence identity threshold for deduplication step."),
        type=float,
        default=params.nuc1["sequence_identity"],
        metavar="FLOAT",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-coverage",
        help=extended_help_text("Coverage threshold for deduplication step."),
        type=float,
        default=params.nuc1["coverage"],
        metavar="FLOAT",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-evalue",
        help=extended_help_text("E-value threshold for deduplication step."),
        type=float,
        default=params.nuc1["e_value"],
        metavar="FLOAT",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-cluster-mode",
        help=extended_help_text("Clustering mode for deduplication: 0=SetCover, 1=Connected component, 2=Greedy incremental."),
        type=int,
        default=params.nuc1["cluster_mode"],
        choices=[0, 1, 2],
        metavar="INT",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-coverage-mode",
        help=extended_help_text("Coverage mode for deduplication: 0=target, 1=query, 2=both, 3=length of target, 4=length of query, 5=length of both."),
        type=int,
        default=params.nuc1["coverage_mode"],
        choices=[0, 1, 2, 3, 4, 5],
        metavar="INT",
    )
    group_mmseqs_dedup.add_argument(
        "--dedup-spaced-kmer-mode",
        help=extended_help_text("Spaced k-mer mode for deduplication: 0=use ungapped k-mers, 1=use spaced k-mers."),
        type=int,
        default=params.nuc1["spaced_kmer_mode"],
        choices=[0, 1],
        metavar="INT",
    )

    group_mmseqs_nuc = extended_parser_group(
        "MMSeqs2 Nucleotide Clustering",
        "Parameters for the second nucleotide clustering step (relaxed clustering of representatives).",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-identity",
        help=extended_help_text("Sequence identity threshold for nucleotide clustering step."),
        type=float,
        default=params.nuc2["sequence_identity"],
        metavar="FLOAT",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-coverage",
        help=extended_help_text("Coverage threshold for nucleotide clustering step."),
        type=float,
        default=params.nuc2["coverage"],
        metavar="FLOAT",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-evalue",
        help=extended_help_text("E-value threshold for nucleotide clustering step."),
        type=float,
        default=params.nuc2["e_value"],
        metavar="FLOAT",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-cluster-mode",
        help=extended_help_text("Clustering mode for nucleotide step: 0=SetCover, 1=Connected component, 2=Greedy incremental."),
        type=int,
        default=params.nuc2["cluster_mode"],
        choices=[0, 1, 2],
        metavar="INT",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-coverage-mode",
        help=extended_help_text("Coverage mode for nucleotide step: 0=target, 1=query, 2=both, 3=length of target, 4=length of query, 5=length of both."),
        type=int,
        default=params.nuc2["coverage_mode"],
        choices=[0, 1, 2, 3, 4, 5],
        metavar="INT",
    )
    group_mmseqs_nuc.add_argument(
        "--nuc-spaced-kmer-mode",
        help=extended_help_text("Spaced k-mer mode for nucleotide step: 0=use ungapped k-mers, 1=use spaced k-mers."),
        type=int,
        default=params.nuc2["spaced_kmer_mode"],
        choices=[0, 1],
        metavar="INT",
    )

    group_mmseqs_prot = extended_parser_group(
        "MMSeqs2 Protein Clustering",
        "Parameters for protein clustering step (used for compositional analysis).",
    )
    group_mmseqs_prot.add_argument(
        "--prot-identity",
        help=extended_help_text("Sequence identity threshold for protein clustering step."),
        type=float,
        default=params.prot["sequence_identity"],
        metavar="FLOAT",
    )
    group_mmseqs_prot.add_argument(
        "--prot-coverage",
        help=extended_help_text("Coverage threshold for protein clustering step."),
        type=float,
        default=params.prot["coverage"],
        metavar="FLOAT",
    )
    group_mmseqs_prot.add_argument(
        "--prot-evalue",
        help=extended_help_text("E-value threshold for protein clustering step."),
        type=float,
        default=params.prot["e_value"],
        metavar="FLOAT",
    )
    group_mmseqs_prot.add_argument(
        "--prot-coverage-mode",
        help=extended_help_text("Coverage mode for protein step: 0=target, 1=query, 2=both, 3=length of target, 4=length of query, 5=length of both."),
        type=int,
        default=params.prot["coverage_mode"],
        choices=[0, 1, 2, 3, 4, 5],
        metavar="INT",
    )

    group_clustering = extended_parser_group(
        "Clustering", "Parameters to control the hierarchical clustering."
    )
    group_clustering.add_argument(
        "--clustering-method",
        help=extended_help_text("The hierarchical clustering method to use for protein-level clustering."),
        default=strategy.method,
        choices=[
            "average",
            "single",
            "complete",
            "weighted",
            "centroid",
            "median",
            "ward",
            "linclust",
            "none",
        ],
    )
    group_clustering.add_argument(
        "--clustering-distance",
        help=extended_help_text("The distance threshold after which to stop merging clusters."),
        type=float,
        default=strategy.distance,
    )
    group_clustering.add_argument(
        "--clustering-precision",
        help=extended_help_text("The numerical precision to use for computing distances for hierarchical clustering."),
        default=strategy.precision,
        choices=["half", "single", "double"],
    )
    group_clustering.add_argument(
        "--clustering-weight",
        help=extended_help_text("The method to use to weigh dimensions in the distance matrix for hierarchical clustering."),
        default="protein",
        choices=["protein", "none"],
    )

    return parser

def _build_parser_sphinx():
    return build_parser(["igua", "--help-all"])


def create_dataset(
    progress: rich.progress.Progress,
    args: argparse.Namespace,
) -> BaseDataset:
    """Create dataset handler with appropriate adapter.
    """
    datasets = []
    for input_ in args.inputs:
        dataset = input_.to_dataset(args)
        datasets.append(dataset)
    return DatasetList(datasets)


def get_mmseqs_params(args: argparse.Namespace) -> PipelineParameters:
    """Build MMSeqs2 parameter dictionaries from command-line arguments."""

    params_nuc1 = dict(
        e_value=args.dedup_evalue,
        sequence_identity=args.dedup_identity,
        coverage=args.dedup_coverage,
        cluster_mode=args.dedup_cluster_mode,
        coverage_mode=args.dedup_coverage_mode,
        spaced_kmer_mode=args.dedup_spaced_kmer_mode,
    )

    params_nuc2 = dict(
        e_value=args.nuc_evalue,
        sequence_identity=args.nuc_identity,
        coverage=args.nuc_coverage,
        cluster_mode=args.nuc_cluster_mode,
        coverage_mode=args.nuc_coverage_mode,
        spaced_kmer_mode=args.nuc_spaced_kmer_mode,
    )

    params_prot = dict(
        e_value=args.prot_evalue,
        coverage=args.prot_coverage,
        coverage_mode=args.prot_coverage_mode,
        sequence_identity=args.prot_identity,
    )

    return PipelineParameters(
        nuc1=params_nuc1,
        nuc2=params_nuc2,
        prot=params_prot,
    )


def get_clustering_strategy(args: argparse.Namespace) -> ClusteringStrategy:
    if args.clustering_method == "none":
        return None
    elif args.clustering_method == "linclust":
        return LinearClustering(
            distance=args.clustering_distance,
        )
    else:
        return HierarchicalClustering(
            method=args.clustering_method,
            distance=args.clustering_distance,
            precision=args.clustering_precision,
            jobs=args.jobs,
        )


def get_protein_representative(
    mmseqs: MMSeqs,
    input_path: pathlib.Path,
    output_prefix: pathlib.Path,
    fasta_path: pathlib.Path,
) -> None:
    db = Database(mmseqs, input_path.with_suffix(".db"))
    result = Clustering(mmseqs, output_prefix.with_suffix(".db"), db)
    subdb = result.to_subdb(output_prefix.with_name(f"{output_prefix.name}_rep_seq.db"))
    subdb.to_fasta(fasta_path)


def report_completion(
    stopwatch: Stopwatch,
    console: rich.console.Console,
) -> None:
    end_time_str = stopwatch.stop_datetime.strftime("%Y-%m-%d at %H:%M:%S")
    time_str = stopwatch.total_human()

    console.print(f"[bold blue]{'Completed':>12}[/] on {end_time_str}")
    console.print(f"[bold green]{'Total time':>12}[/] {time_str}")


def main(argv: typing.Optional[typing.List[str]] = None) -> int:
    # retrieve argv
    argv = sys.argv[1:] if argv is None else argv

    # build parser and get arguments
    parser = build_parser(argv)
    if not isinstance(argcomplete, ImportError):
        argcomplete.autocomplete(parser)
    args = parser.parse_args(argv)

    with contextlib.ExitStack() as ctx:

        # use a stopwatch to measure time
        stopwatch = ctx.enter_context(Stopwatch())

        # extract MMseqs parameters from command line
        params = get_mmseqs_params(args)
        strategy = get_clustering_strategy(args)

        # use user provided workdir or create a new one in `tempfile`
        if args.workdir is None:
            tempdir = tempfile.TemporaryDirectory(prefix="igua-")
        else:
            tempdir = pathlib.Path(args.workdir)
            tempdir.mkdir(parents=True, exist_ok=True)
            tempdir = tempfile.TemporaryDirectory(prefix="igua-", dir=args.workdir)

        # validate individual file arguments
        if not args.inputs:
            parser.error("No input files given.")

        # acquire temporary folder
        workdir = pathlib.Path(ctx.enter_context(tempdir))

        # prepare progress bar
        progress = ctx.enter_context(
            rich.progress.Progress(
                "",
                rich.progress.SpinnerColumn(),
                *rich.progress.Progress.get_default_columns(),
            )
        )
        progress.console.print(f"[bold blue]{'Started':>12}[/] on {stopwatch.start_datetime.strftime('%Y-%m-%d at %H:%M:%S')}")

        # check mmseqs version (and if an mmseqs binary is available)
        mmseqs = MMSeqs(progress=progress, threads=args.jobs, tempdir=workdir)
        try:
            v = mmseqs.version()
            progress.console.print(f"[bold green]{'Using':>12}[/] [purple]MMSeqs2[/] version {v!r}")
        except RuntimeError:
            progress.console.print(f"[bold red]{'Failed':>12}[/] to locate [purple]MMSeqs2[/] binary")
            return errno.ENOENT

        # create a pipeline with the configuration from the CLI
        pipeline = Pipeline(
            strategy=strategy,
            params=params,
            mmseqs=mmseqs,
            workdir=workdir,
            progress=progress,
            weight=None if args.clustering_weight == "none" else args.clustering_weight,
        )

        # create appropriate dataset handler
        dataset = create_dataset(progress=progress, args=args)

        # run pipeline and retrieve GCFs
        result = pipeline.run(dataset)

        # save GCFs
        result.gcfs.to_csv(args.output, sep="\t", index=False)
        progress.console.print(
            f"[bold green]{'Saved':>12}[/] final GCFs table to {str(args.output)!r}"
        )

        # save compositions
        if args.compositions is not None:
            gcf_representatives = result.gcfs["gcf_representative"].unique().tolist()
            representatives_compositions = anndata.AnnData(
                X=result.compositions[gcf_representatives].X,
                var=result.compositions.var,
                obs=(
                    result.gcfs.set_index("cluster_id")
                        .loc[gcf_representatives, ["gcf_id", "gcf_representative"]]
                        .set_index("gcf_id")
                ),
            )
            representatives_compositions.write(args.compositions)
            progress.console.print(
                f"[bold green]{'Saved':>12}[/] GCF compositions to {str(args.compositions)!r}"
            )

        # save protein features
        if args.features is not None:
            get_protein_representative(
                mmseqs,
                result.proteins_faa,
                workdir.joinpath("step3"),
                args.features,
            )
            progress.console.print(
                f"[bold green]{'Saved':>12}[/] protein features to {str(args.features)!r}"
            )

    # report runtime before exiting
    report_completion(stopwatch, progress.console)

    return 0
