import abc
import contextlib
import copy
import dataclasses
import pathlib
import tempfile
import typing
from typing import Dict, Literal, Mapping, Tuple

import anndata
import rich.console
import rich.progress
import scipy.sparse
import numpy
import pandas

from .dataset.base import BaseDataset
from .dataset.defensefinder import DefenseFinderDataset
from .dataset.fasta_gff import FastaGFFDataset
from .mmseqs import MMSeqs, Database, Clustering
from .clustering import ClusteringStrategy, HierarchicalClustering, LinearClustering, default_strategy
from ._utils import Stopwatch


class _BaseSink(abc.ABC):

    def __init__(self):
        self.stats = []
        self.done = set()

    def add_record(self, name: str, sequence: str, **kwargs) -> bool:
        if name not in self.done:
            self.stats.append({"id": name, "length": len(sequence), **kwargs})
            self.done.add(name)
            return True
        else:
            return False

    def report_statistic(self) -> pandas.DataFrame:
        return pandas.DataFrame(self.stats)


class _FASTASink(_BaseSink):

    def __init__(self, file: typing.TextIO) -> None:
        super().__init__()
        self.file = file

    def add_record(self, name: str, sequence: str, **kwargs) -> None:
        if not super().add_record(name, sequence, **kwargs):
            return False
        self.file.write(">")
        self.file.write(name)
        self.file.write("\n")
        self.file.write(sequence)
        self.file.write("\n")
        return True


@dataclasses.dataclass
class PipelineParameters:
    """The parameters of the IGUA pipeline.

    Attributes:
        nuc1 (`dict`): A dictionary of parameters to pass to MMseqs2 for the
            stage 1 (nucleotide deduplication).
        nuc2 (`dict`): A dictionary of parameters to pass to MMseqs2 for the
            stage 2 (nucleotide clustering).
        prot (`dict`): A dictionary of parameters to pass to MMseqs2 for the
            stage 3 (protein clustering).
        clustering_method (`str`): The name of the clustering method to
            use on the stage 3, either one of the supported hierarchical
            clustering mode, or ``linclust`` to use MMseqs2-derived linear
            clustering.
        clustering_distance (`float`): The distance cutoff to use for
            cluster creation in stage 3.
        clustering_precision (`str`): The precision to use to perform
            hierarchical clustering: one of *half*, *single*, or *double*.
        clustering_weight (`str`): The weights to use for computing pairwise
            distances during stage 3: either *protein* to weigh each
            dimension by the length of the protein it represents, or `None`
            to use uniform weights.

    """
    nuc1: Dict[str, object]
    nuc2: Dict[str, object]
    prot: Dict[str, object]

    @classmethod
    def default(cls) -> "PipelineParameters":
        """Create new default clustering parameters.

        Returns:
            `~igua.pipeline.PipelineParameters`: A new pipeline parameters
            object with all parameters initialized to defaults.

        """
        return cls(
            nuc1=dict(
                e_value=0.001,
                sequence_identity=0.85,
                coverage=1.0,
                cluster_mode=0,
                coverage_mode=1,
                spaced_kmer_mode=0,
            ),
            nuc2=dict(
                e_value=0.001,
                sequence_identity=0.6,
                coverage=0.5,
                cluster_mode=0,
                coverage_mode=0,
                spaced_kmer_mode=0,
            ),
            prot=dict(
                e_value=0.001,
                coverage=0.9,
                coverage_mode=1,
                sequence_identity=0.5,
            ),
        )


@dataclasses.dataclass
class PipelineResult:
    """The results of the IGUA clustering pipeline.

    Attributes:
        gcfs (`pandas.DataFrame`): A dataframe with columns ``cluster_id``,
            ``cluster_length``, ``gcf_id``, ``gcf_representative``,
            ``nucleotide_representative``, ``fragment_representative``
            and ``filename`` summarizing the generated gene cluster families
            at each stage.
        compositions (`anndata.AnnData` or `None`): The protein compositions
            of the stage 2 representative clusters generated by IGUA. May
            be `None` if no clustering method was requested, or if stage 2
            already only produced a singleton.

    """
    gcfs: pandas.DataFrame
    compositions: typing.Optional[anndata.AnnData]


class Pipeline:
    """The IGUA multi-stage clustering pipeline.
    """

    def __init__(
        self,
        strategy: ClusteringStrategy = default_strategy(),
        params: typing.Optional[PipelineParameters] = None,
        *,
        prefix: str = "GCF",
        jobs: int = 1,
        weight: Literal["protein", None] = "protein",
        mmseqs: typing.Optional[MMSeqs] = None,
        progress: typing.Optional[rich.progress.Progress] = None,
        workdir: typing.Optional[pathlib.Path] = None,
    ):
        """Create a new pipeline.

        Arguments:
            strategy (`ClusteringStrategy` or `None`): The clustering
                strategy to use to form the final clusters using the
                protein compositions in the stage 3 of the pipeline.
                Pass `None` to disable the stage 3 altogether.
            params (`PipelineParameters` or `None`): An object with
                parameters to pass to MMseqs2 for each stages of the
                pipeline.
            prefix (`str`): The prefix to use for the GCF identifiers.
            jobs (`int`): The number of threads to use with MMseqs2 and
                the distance computation in stage 3.
            weight (`str` or `None`): How to weight dimensions in the
                protein composition matrix. Use `protein` to weigh by
                protein length, or `None` to disable weighting.
            mmseqs (`~igua.mmseqs.MMseqs`): The MMSeqs2 driver to use
                to conduct linear clustering in stages 1 and 2 and
                protein clustering in stage 3.
            progress (`~rich.progress.Progress` or `None`): A `rich`
                progress object to use for reporting progress.
            workdir (`pathlib.Path` or `None`): The path to a folder
                to use for temporary files. If `None` given, generates
                a temporary folder with `tempfile.TemporaryDirectory`
                on each `~Pipeline.run` invokation.

        """
        self.jobs = jobs
        self.params = params or PipelineParameters.default()
        self.workdir = None if workdir is None else pathlib.Path(workdir)
        self.prefix = prefix
        self.mmseqs = mmseqs

        self.progress = progress
        if progress is None:
            self.console = rich.console.Console(quiet=True)
        else:
            self.console = progress.console

        if weight is not None and weight != "protein":
            raise ValueError(f"invalid weight: {weight!r}")
        self.weight = weight

        if strategy is None:
            self.clustering_strategy = None
        else:
            if not isinstance(strategy, ClusteringStrategy):
                raise TypeError(f"expected ClusteringStrategy, got {type(strategy).__name__}")
            self.clustering_strategy = copy.deepcopy(strategy)
            self.clustering_strategy.jobs = self.jobs

    # ---

    def _extract_clusters_to_file(
        self,
        dataset: BaseDataset,
        output: pathlib.Path,
    ) -> pandas.DataFrame:
        """Extract the dataset clusters to the given file.

        Arguments:
            dataset (`~igua.dataset.base.BaseDataset`): The dataset
                containing the gene clusters to extract.
            output (`pathlib.Path`): The path to the output file to
                write the gene clusters sequence to.

        Returns:
            `pandas.DataFrame`: A table with columns ``cluster_id`` and
            ``cluster_length`` reporting the ID and lengths of the
            extracted clusters.

        """
        # record processed sequences to flag duplicates
        done = set()
        n_duplicates = 0
        # extract raw sequences
        self.console.print(f"[bold blue]{'Loading':>12}[/] input clusters")
        with output.open("w") as dst:
            sink = _FASTASink(dst)
            for cluster in dataset.extract_clusters(self.progress):
                if cluster.id not in done:
                    done.add(cluster.id)
                    sink.add_record(cluster.id, cluster.sequence, filename=cluster.source)
                else:
                    n_duplicates += 1
            input_sequences = (
                sink.report_statistic()
                    .rename(columns={"id": "cluster_id", "length": "cluster_length"})
                    .set_index("cluster_id")
            )
        if n_duplicates > 0:
            self.console.print(
                f"[bold yellow]{'Warning':>12}[/] {n_duplicates} duplicate "
                "clusters found in inputs"
            )

        self.console.print(
            f"[bold green]{'Loaded':>12}[/] {len(input_sequences)} "
            "clusters to process"
        )
        return input_sequences

    def _extract_proteins_to_file(
        self,
        dataset: BaseDataset,
        clusters: typing.Collection[str],
        output: pathlib.Path,
    ) -> pandas.DataFrame:
        """Extract the dataset proteins to the given file.

        Arguments:
            dataset (`~igua.dataset.BaseDataset`): The dataset containing
                the gene clusters to extract and their proteins.
            clusters (`~collections.abc.Collection` of `str`): A collection
                of clusters from which to extract proteins (typically the
                clusters selected as representatives in the previous step).
            output (`pathlib.Path`): The path to the output file to write
                proteins to.

        Returns:
            `pandas.DataFrame`: A table with columns ``protein_id``,
            ``protein_length`` and ``cluster_id`` reporting the ID, length
            and source cluster of the extracted proteins.

        """
        self.console.print(
            f"[bold blue]{'Extracting':>12}[/] protein sequences from "
            "representative clusters"
        )
        with output.open("w") as dst:
            sink = _FASTASink(dst)
            for protein in dataset.extract_proteins(self.progress, clusters):
                sink.add_record(protein.id, protein.sequence, cluster_id=protein.cluster_id)
            protein_sizes = (
                sink.report_statistic()
                    .rename(columns={"id": "protein_id", "length": "protein_length"})
                    .set_index("protein_id")
                    .astype({"protein_length": "int32"})
            )
        return protein_sizes

    def _make_compositions(
        self,
        protein_clusters: pandas.DataFrame,
        representatives: typing.Dict[str, int],
        protein_representatives: typing.Dict[str, int],
    ) -> anndata.AnnData:
        """Build a compositional matrix from the given protein clusters.

        Arguments:
            protein_clusters (`pandas.DataFrame`): The clustering of
                proteins returned by MMSeqs.
            representatives (`dict` of `str` to `int`): The index mapping
                gene cluster IDs to their index in the compositional table.
            protein_representatives (`dict` of `str` to `int`): The index
                mapping protein IDs to their index in the compositional
                table.

        Returns:
            `~anndata.AnnData`: An annotated compositional matrix where
            each cluster is represented as a combination of protein
            counts.

        """
        self.console.print(
            f"[bold blue]{'Building':>12}[/] protein compositional array"
        )

        # index the proteins by protein ID
        lengths = protein_clusters.set_index("protein_id")["protein_length"]

        # create the empty compositional matrix
        compositions = scipy.sparse.dok_matrix(
            (len(representatives), len(protein_representatives)),
            dtype=numpy.int32
        )

        # Build compositional matrix using protein lengths as weights
        task = self.progress.add_task(
            description=f"[bold blue]{'Working':>9}[/]",
            total=len(protein_clusters)
        )
        for row in self.progress.track(protein_clusters.itertuples(), task_id=task):
            cluster_index = representatives[row.cluster_id]
            prot_index = protein_representatives[row.protein_representative]
            compositions[cluster_index, prot_index] += 1 #lengths.loc[row.protein_representative]
        self.progress.remove_task(task)

        # Build the annotated data matrix
        sorted_representatives = sorted(representatives, key=representatives.__getitem__)
        sorted_protein_representatives = sorted(
            protein_representatives, key=protein_representatives.__getitem__
        )
        return anndata.AnnData(
            X=compositions.tocsr(),
            obs=pandas.DataFrame(
                index=pandas.Index(sorted_representatives, name="cluster_id")
            ),
            var=pandas.DataFrame(
                index=pandas.Index(sorted_protein_representatives, name="protein_id"),
                data=dict(size=lengths.loc[sorted_protein_representatives]),
            ),
        )

    def _run_nucleotide_deduplication(
        self,
        workdir: pathlib.Path,
        clusters_fna: pathlib.Path,
        mmseqs: MMSeqs,
    ) -> Tuple[Clustering, pandas.DataFrame]:
        """Run the nucleotide deduplication stage of the pipeline.

        Arguments:
            workdir (`pathlib.Path`): The path to the temporary directory
                to use for processing data.
            clusters_fna (`pathlib.Path`): The path to a two-line FASTA
                file containing the gene clusters to process.

        Returns:
            (`~igua.mmseqs.Clustering`, `pandas.DataFrame`): The clustering
            file created by MMseqs, and a dataframe containing a summary of
            the result.

        """
        # create initial sequence database
        self.console.print(
            f"[bold blue]{'Starting':>12}[/] nucleotide deduplication step "
            "with [purple]MMSeqs2[/]"
        )
        with Stopwatch() as stopwatch:
            db = Database.create(mmseqs, clusters_fna)
            step1 = db.cluster(workdir / "step1.db", **self.params.nuc1)
            gcfs1 = step1.to_dataframe(columns=["fragment_representative", "cluster_id"]).sort_values("cluster_id")  # type: ignore
        self.console.print(
            f"[bold blue]{'Finished':>12}[/] nucleotide deduplication step "
            f"in [bold cyan]{stopwatch.total_human()}[/]"
        )
        self.console.print(
            f"[bold green]{'Reduced':>12}[/] {len(gcfs1)} clusters to "
            f"{gcfs1.fragment_representative.nunique()} complete representatives"
        )
        return step1, gcfs1

    def _run_nucleotide_clustering(
        self,
        workdir: pathlib.Path,
        step1: Clustering,
    ) -> pandas.DataFrame:
        """Run the nucleotide clustering stage of the pipeline.

        Arguments:
            workdir (`pathlib.Path`): The path to the temporary directory
                to use for processing data.
            step1 (`~igua.mmseqs.Clustering`): The MMseqs clustering file
                created at the previous stage.

        Returns:
            `pandas.DataFrame`: The result of the clustering on the
            representative sequences of the ``step1`` clustering.

        """
        # cluster sequences
        self.console.print(
            f"[bold blue]{'Starting':>12}[/] nucleotide clustering step with [purple]MMSeqs2[/]"
        )
        with Stopwatch() as stopwatch:
            repdb = step1.to_subdb(workdir / "step1.rep_seq.db")
            step2 = repdb.cluster(workdir / "step2.db", **self.params.nuc2)
            gcfs2 = step2.to_dataframe(columns=["nucleotide_representative", "fragment_representative"]).sort_values("fragment_representative")  # type: ignore
        self.console.print(
            f"[bold blue]{'Finished':>12}[/] nucleotide clustering step "
            f"in [bold cyan]{stopwatch.total_human()}[/]"
        )
        self.console.print(
            f"[bold green]{'Reduced':>12}[/] {len(gcfs2)} clusters to "
            f"{len(gcfs2.nucleotide_representative.unique())} nucleotide "
            "representatives"
        )
        return gcfs2

    def _extract_representatives(
        self,
        gcfs2: pandas.DataFrame
    ) -> Dict[str, int]:
        """Extract the representative clusters in ``gcfs2`` into a index.

        Arguments:
            gcfs2 (`pandas.DataFrame`): The clustering results of the
                second stage.

        Returns:
            `dict` of `str` to `int`: A dictionary mapping every unique
            representative cluster in ``gcfs2`` to a single index.

        """
        # load representatives
        self.console.print(
            f"[bold blue]{'Extracting':>12}[/] representative clusters"
        )
        representatives = {
            x: i
            for i, x in enumerate(sorted(gcfs2["nucleotide_representative"].unique()))
        }
        self.console.print(
            f"[bold green]{'Found':>12}[/] {len(representatives)} nucleotide "
            "representative clusters"
        )
        return representatives

    def _run_protein_clustering(
        self,
        workdir: pathlib.Path,
        dataset: BaseDataset,
        representatives: typing.Collection[str],
        mmseqs: MMSeqs,
    ) -> anndata.AnnData:
        """Run protein clustering and build a compositional matrix.

        Arguments:
            workdir (`pathlib.Path`): The path to the temporary directory
                to use for processing data.
            dataset (`~igua.dataset.BaseDataset`): The dataset containing
                the gene clusters to extract and their proteins.
            representatives (`collections.abc.Collection` of `str`): A
                collection of gene cluster IDs to extract proteins from.

        Returns:
            `anndata.AnnData`: A compositional matrix derived from the
            protein clustering. See `ClusteringPipeline._make_compositions`
            for more information.

        """
        # extract proteins and record sizes
        proteins_faa = workdir / "proteins.faa"
        protein_sizes = self._extract_proteins_to_file(dataset, representatives, output=proteins_faa)
        if not proteins_faa.exists() or proteins_faa.stat().st_size == 0:
            self.console.print(
                f"[bold yellow]{'Warning':>12}[/] No proteins extracted from input dataset"
            )

        # cluster proteins
        self.console.print(
            f"[bold blue]{'Starting':>12}[/] protein clustering step with [purple]MMSeqs2[/]"
        )
        with Stopwatch() as stopwatch:
            prot_db = Database.create(mmseqs, proteins_faa)
            prot_result = prot_db.cluster(workdir / "step3.db", **self.params.prot)
            prot_clusters = prot_result.to_dataframe(
                columns=["protein_representative", "protein_id"]
            )
            # record protein lengths and source cluster
            prot_clusters = pandas.merge(
                prot_clusters,
                protein_sizes[["cluster_id", "protein_length"]],
                on="protein_id"
            )
            # extract protein representatives
            protein_representatives = {
                x: i
                for i, x in enumerate(
                    sorted(prot_clusters["protein_representative"].unique())
                )
            }
        self.console.print(
            f"[bold blue]{'Finished':>12}[/] protein clustering step "
            f"in [bold cyan]{stopwatch.total_human()}[/]"
        )
        self.console.print(
            f"[bold green]{'Found':>12}[/] {len(protein_representatives)} "
            f"protein representatives for {len(prot_clusters)} proteins"
        )
        return self._make_compositions(
            prot_clusters,
            representatives,
            protein_representatives,
        )

    def _hierarchical_clustering(
        self,
        compositions: anndata.AnnData,
    ) -> pandas.DataFrame:
        """Perform hierarchical (or linear) clustering on the given data.

        Arguments:
            compositions (`anndata.AnnData`): A compositional matrix
                obtained from the protein clustering.

        Returns:
            `pandas.DataFrame`: A table mapping each row of ``compositions``
            to a GCF with arbitrary identifiers.

        """
        assert self.clustering_strategy is not None
        # use weights unless disabled
        if self.weight == "protein":
            weights = compositions.var["size"].values
        elif self.weight is None:
            weights = None
        else:
            raise ValueError(f"invalid clustering weight: {self.weight!r}")
        # run clustering based on protein array membership
        self.console.print(
            f"[bold blue]{'Clustering':>12}[/] gene clusters using "
            f"{self.clustering_strategy.method} linkage"
        )
        with Stopwatch() as stopwatch:
            flat = self.clustering_strategy.cluster(compositions.X, weights)
            # build GCFs based on flat clustering
            gcfs3 = pandas.DataFrame(
                {
                    "gcf_id": [f"{self.prefix}{i:07}" for i in flat],
                    "nucleotide_representative": compositions.obs_names,
                }
            )
        self.console.print(
            f"[bold blue]{'Finished':>12}[/] hierarchical clustering step "
            f"in [bold cyan]{stopwatch.total_human()}[/]"
        )
        return gcfs3

    def _hierarchical_clustering_fallback(
        self,
        representatives: Mapping[str, int],
    ) -> pandas.DataFrame:
        self.console.print(
            f"[bold yellow]{'Skipping':>12}[/] hierarchical clustering step"
        )
        compositions = None
        # proteins_faa = None
        sorted_representatives = sorted(
            representatives,
            key=representatives.__getitem__
        )
        gcfs3 = pandas.DataFrame(
            {
                "gcf_id": [
                    f"{self.prefix}{i+1:07}"
                    for i in range(len(sorted_representatives))
                ],
                "nucleotide_representative": sorted_representatives,
            }
        )
        return gcfs3

    def _join_results(
        self,
        gcfs1: pandas.DataFrame,
        gcfs2: pandas.DataFrame,
        gcfs3: pandas.DataFrame,
        input_sequences: pandas.DataFrame,
    ) -> pandas.DataFrame:
        """Join the results of the different stages to a single dataframe.
        """
        #
        self.console.print(
            f"[bold green]{'Built':>12}[/] {len(gcfs3.gcf_id.unique())} "
            f"GCFs from {len(input_sequences)} initial clusters"
        )

        # build GCFs based on flat clustering
        gcf3_representatives = (
            pandas.merge(
                gcfs3,
                input_sequences["cluster_length"],
                left_on="nucleotide_representative",
                right_index=True,
            )
            .sort_values("cluster_length")
            .drop_duplicates("gcf_id", keep="last")
            .set_index("gcf_id")
        )
        gcfs3 = pandas.merge(
            gcfs3,
            gcf3_representatives["nucleotide_representative"].rename(
                "gcf_representative"
            ),
            left_on="gcf_id",
            right_index=True,
        )

        # build final GCF table
        gcfs = pandas.merge(
            pandas.merge(
                pandas.merge(gcfs1, gcfs2, on="fragment_representative"),
                gcfs3,
                on="nucleotide_representative",
            ),
            input_sequences,
            left_on="cluster_id",
            right_index=True,
        )

        # sort results and drop unused columns
        gcfs.sort_values(["gcf_id", "cluster_length"], inplace=True)
        gcfs = gcfs[
            [
                "cluster_id",
                "cluster_length",
                "gcf_id",
                "gcf_representative",
                "nucleotide_representative",
                "fragment_representative",
                "filename",
            ]
        ]
        return gcfs

    # ---

    def run(
        self,
        dataset: BaseDataset,
    ):
        """Run the clustering pipeline on the given dataset.

        Arguments:
            dataset (`~igua.dataset.base.BaseDataset`): A dataset containing
                the gene clusters to cluster into families.

        Returns:
            `~igua.pipeline.PipelineResult`: The results of the pipeline.
            See class-level documentation for more information.

        """
        # create a context to use the user-provided temporary directory
        # or create one if none was given
        if self.workdir is None:
            workdir_context = tempfile.TemporaryDirectory(prefix="igua-")
        else:
            workdir_context = contextlib.nullcontext(self.workdir)

        with workdir_context as workdir:
            workdir = pathlib.Path(workdir)
            # use the user-provided MMseqs or create a default one
            mmseqs = self.mmseqs or MMSeqs(
                progress=self.progress,
                threads=self.jobs,
                tempdir=workdir
            )
            # extract input nucleotides
            clusters_fna = workdir / "clusters.fna"
            input_sequences = self._extract_clusters_to_file(dataset, clusters_fna)
            # run first two steps of clustering and get nucleotide representatives
            step1, gcfs1 = self._run_nucleotide_deduplication(workdir, clusters_fna, mmseqs)
            gcfs2 = self._run_nucleotide_clustering(workdir, step1)
            representatives = self._extract_representatives(gcfs2)
            # run protein clustering if enabled and if there are enough
            # gene clusters left
            if self.clustering_strategy is not None and len(representatives) > 1:
                compositions = self._run_protein_clustering(
                    workdir,
                    dataset,
                    representatives,
                    mmseqs,
                )
                gcfs3 = self._hierarchical_clustering(
                    compositions,
                )
            else:
                compositions = None
                gcfs3 = self._hierarchical_clustering_fallback(
                    representatives
                )

        # generate result table
        gcfs = self._join_results(gcfs1, gcfs2, gcfs3, input_sequences)
        return PipelineResult(gcfs=gcfs, compositions=compositions)
