from __future__ import annotations

import re
import copy
from typing import Callable, Iterable, Iterator, Literal, overload, TYPE_CHECKING
import numpy as np
import matplotlib.pyplot as plt
from .utils import (
    append_log,
    plot2,
    concat,
    calc_covars,
    calc_startprob,
    calc_transmat,
    sfHMMAnalysisError,
    optimize_ax,
    optimize_b,
    under_development,
)
from .single_sfhmm import sfHMM1, _S
from .base import sfHMMBase

if TYPE_CHECKING:
    import pandas as pd

__all__ = ["sfHMMn"]


class sfHMMn(sfHMMBase):
    r"""
    Multi-trajectory sfHMM.
    This class shares many attributes in sfHMM1.

    Analysis Results
    ----------------
    gmm_opt : `GMMs` or `DPGMM` object. For more detail, see `sfHMM.gmm`.
        The result of GMM clustering, which has following attributes:
        - wt ... Weights.
        - mu ... Means.
        - sg ... Standard deviations.
        - n_components ... The number of states.
        If AIC/BIC minimization of standard GMM clustering was conducted, the clustering
        results will be stored in 'gmm'.
    n_components : int
        The optimal number of states. Same as 'gmm_opt.n_components'.
    self[i] (accessed by indexing) : sfHMM object
        The i-th sfHMM object. See .\single_sfhmm.py.
    """

    def __init__(
        self,
        data_raw: Iterable[_S] | None = None,
        *,
        sg0: float = -1,
        psf: float = -1,
        krange: int | tuple[int, int] = None,
        model: str = "g",
        name: str = "",
        **kwargs,
    ):
        """
        Parameters
        ----------
        data_raw : iterable, optional
            Datasets to be added.
        sg0 : float, optional
            Parameter used in filtering method. Expected to be 20% of signal change.
            If <= 0, sg0 will be determined automatically.
        psf : float, optional
            Transition probability used in step finding algorithm.
            if 0 < p < 0.5 is not satisfied, the original Kalafut-Visscher's algorithm
            is executed.
        krange : int or (int, int)
            Minimum and maximum number of states to search in GMM clustering. If it is
            integer, then
            it will be interpretted as (krange, krange).
        model: str, by default "g" (= Gaussian)
            Distribution of noise. Gauss and Poisson distribution are available for now.
        name : str, optional
            Name of the object.
        """
        self.n_data = 0
        super().__init__(sg0, psf, krange, model, name, **kwargs)
        self.ylim = [np.inf, -np.inf]
        self._sf_list = []
        self.gmm_opt = None
        data_raw is None or self.appendn(data_raw)

    @overload
    def __getitem__(self, key: int) -> sfHMM1: ...
    @overload
    def __getitem__(self, key: slice) -> sfHMMn: ...

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._sf_list[key]
        elif isinstance(key, slice):
            _k = ":".join([str(key.start), str(key.stop), str(key.step or "")])
            if _k.endswith(":"):
                _k = _k[:-1]
            new = self.__class__(
                sg0=self.sg0,
                psf=self.psf,
                krange=self.krange,
                model=self.model,
                name=f"{self.name}[{_k}]",
            )
            new._sf_list = self._sf_list[key]
            new.ylim = self.ylim
            new.n_data = len(new._sf_list)
            return new
        else:
            raise TypeError(f"Index key must be int or slice, got {type(key)}")

    def __iter__(self) -> Iterator[sfHMM1]:
        return iter(self._sf_list)

    def __add__(self, other: sfHMMn) -> sfHMMn:
        """
        `+` is supported for two sfHMMn objects. a+b makes a new sfHMMn object with
        concatenated list of sfHMM1 objects.
        """
        if self is other:
            new = self.__class__(
                sg0=self.sg0,
                psf=self.psf,
                krage=self.krange,
                model=self.model,
                name=self.name + "+" + other.name,
            )
            new.n_data = self.n_data + other.n_data
            new.ylim = [
                min(self.ylim[0], other.ylim[0]),
                max(self.ylim[1], other.ylim[1]),
            ]
            new._sf_list = copy.deepcopy(self._sf_list) + copy.deepcopy(other._sf_list)
            return new
        else:
            raise TypeError(
                "Unsupported operand type(s) for +: "
                f"{self.__class__.__name__} and {type(other)}."
            )

    @property
    def names(self) -> list[str]:
        return [sf.name for sf in self]

    @append_log
    def append(self, data: _S, name: str = None) -> sfHMMn:
        """
        Append a trajectory as sfHMM object.

        Parameters
        ----------
        data : array
            Data to analyze.
        name : str, optional
            Name of the data.
        """
        if name is None:
            name = self.name + f"[{self.n_data}]"

        sf = sfHMM1(
            data,
            sg0=self.sg0,
            psf=self.psf,
            krange=self.krange,
            model=self.StepClass,
            name=name,
        )
        self.n_data += 1
        self._sf_list.append(sf)
        self.ylim[0] = min(sf.ylim[0], self.ylim[0])
        self.ylim[1] = max(sf.ylim[1], self.ylim[1])
        return self

    def appendn(self, datasets: Iterable[_S]) -> sfHMMn:
        """Append all the data in `datasets`.

        Parameters
        ----------
        datasets : dict, list or iterable objects except for np.ndarray or pd.DataFrame.
            Datasets to be appended. If it is dict, then keys are interpreted as names.
        """
        import pandas as pd

        if isinstance(datasets, dict):
            self.from_dict(datasets)
        elif isinstance(datasets, np.ndarray):
            raise TypeError(
                "Datasets of ndarray is ambiguious. Please use self.append(a) for 1-D "
                "ndarray, or explicitly specify along which axis to iterate by such as "
                "list(a) or np.split(a, axis=1)."
            )
        elif isinstance(datasets, pd.DataFrame):
            self.from_pandas(datasets)
        else:
            for data in datasets:
                self.append(data)

        return self

    @append_log
    def delete(self, indices: int | Iterable[int]) -> None:
        """
        Delete sfHMM1 object(s) from the list.

        Parameters
        ----------
        indices : int or iterable of int
            Indices to delete.
        """
        if isinstance(indices, int):
            indices = [indices]
        indices = sorted(indices)

        for i in reversed(indices):
            self._sf_list.pop(i)

        data_raw_all = self.data_raw
        self.ylim = [data_raw_all.min(), data_raw_all.max()]
        self.n_data -= len(indices)

        return None

    def deleteif(self, filter_func: Callable[[sfHMM1], bool]) -> list[int]:
        """
        Delete sfHMM1 object(s) from the list if certain conditions are
        satisfied.

        Parameters
        ----------
        filter_func : callable
            `filter_func(sf)==True` then delete `sf`.

        Returns
        -------
        list of int
            Deleted indices.
        """
        if not callable(filter_func):
            raise TypeError("`filter_func` must be callable")

        indices = [i for i, sf in enumerate(self) if filter_func(sf)]

        len(indices) > 0 and self.delete(indices)

        return indices

    @append_log
    def pop(self, ind: int) -> sfHMM1:
        """
        Delete one sfHMM1 object and return it.

        Parameters
        ----------
        ind : int
            Indice to pop.
        """
        out = self._sf_list.pop(ind)
        data_raw_all = self.data_raw
        self.ylim = [data_raw_all.min(), data_raw_all.max()]
        self.n_data -= 1
        return out

    def from_dict(self, d: dict[str, _S], like: str = None, regex: str = None):
        """
        Load datasets from dict.

        Parameters
        ----------
        d : dict
            Input data.
        like : str, optional
            If given, key that contains this string is appended.
        regex : regular expression, optional
            If given, key that matches this regular expression is appended.
        """
        for name, data in d.items():
            # only keys matched like or regex requirement are appended.
            if like and like not in name:
                continue
            if regex and not re.match(regex, name):
                continue

            self.append(data, name=str(name))

        if self.n_data == 0:
            raise ValueError(
                "No data appended. Confirm that input dict is in a correct format."
            )

        return self

    def from_pandas(
        self,
        df: pd.DataFrame,
        like: str | None = None,
        regex: str | None = None,
        melt: bool | Literal["infer"] = "infer",
    ) -> sfHMMn:
        """
        Load datasets from pandas.DataFrame.

        Parameters
        ----------
        df : DataFrame
            Input data.
        like : str, optional
            If given, dataset that contains this string is appended.
        regex : regular expression, optional
            If given, dataset that matches this regular expression is appended.
        melt : bool or "infer", optional
            If input DataFrame is melted, which is automatically determined when melt is
            "infer".
        """
        if melt == "infer":
            # Determine if df is melted or not
            col0 = df.iloc[:, 0]
            if df.shape[1] == 2 and len(col0.unique()) < len(col0) // 20:
                return self.from_pandas(df, like, regex, melt=True)
            else:
                return self.from_pandas(df, like, regex, melt=False)

        if melt:
            if df.shape[1] != 2:
                raise ValueError(
                    "For melted DataFrame, it must composed of two columns, with names "
                    "in the first and values in the second."
                )
            name_col, value_col = df.columns
            for name in df[name_col].unique():
                if like and like not in name:
                    continue
                elif regex and not re.match(regex, name):
                    continue
                self.append(df[df[name_col] == name][value_col], name)

        else:
            if like is not None or regex is not None:
                df = df.filter(like=like, regex=regex)
            for name in df:
                data = df[name].dropna()  # delete NaN
                self.append(data, name)  # append data

        if self.n_data == 0:
            raise ValueError(
                "No data appended. Confirm that input DataFrame is in a correct format."
            )

        return self

    def read(
        self,
        path: str,
        sep: str = None,
        encoding: str = None,
        header: int = 0,
        **kwargs,
    ):
        """
        Read a file using appropriate function, and import its data to sfHMMn object.

        Parameters
        ----------
        path : str
            Path to file.
        sep, encoding, header
            Important arguments in pd.read_csv. Default is header=0 rather than "infer"
            because header="infer" usually works in a wrong way.
        **kwargs
            Other keyword arguments that will passed to pd.read_csv or pd.read_excel.
        """
        from .io import read

        read(path, out=self, sep=sep, encoding=encoding, header=header, **kwargs)
        self.source = path
        return self

    @append_log
    def step_finding(self) -> sfHMMn:
        """
        Step finding by extended version of Kalafut-Visscher's algorithm.
        Run independently for each sfHMM object.
        """
        if self.n_data <= 0:
            raise sfHMMAnalysisError("Cannot start analysis before appending data.")

        for sf in self:
            sf.psf = self.psf
            sf.step_finding()
        return self

    @append_log
    def denoising(self) -> sfHMMn:
        """
        Denoising by cutting of the standard deviation of noise to sg0.
        Run independently for each sfHMM object.
        """
        if self.n_data <= 0:
            raise sfHMMAnalysisError("Cannot start analysis before appending data.")

        self._init_sg0()

        for sf in self:
            sf.sg0 = self.sg0
            sf.denoising()

        return self

    @append_log
    def gmmfit(
        self, method: str = "bic", n_init: int = 1, random_state: int = 0
    ) -> sfHMMn:
        """
        Fit the denoised data to Gaussian mixture model.

        Paramters
        ---------
        method : str, optional
            How to determine the optimal number of states. This parameter must be
            'aic', 'bic' or 'Dirichlet'. By default "bic".
        n_init : int, optional
            How many times initialization will be performed in K-means, by default 1.
        random_state : int , optional
            Random seed for K-means initialization., by default 0.
        """
        if self.n_data <= 0:
            raise sfHMMAnalysisError("Cannot start analysis before appending data.")

        self._gmmfit(method, n_init, random_state)

        for sf in self:
            sf.states = self.gmm_opt.predict(np.asarray(sf.step.fit).reshape(-1, 1))
            sf.n_components = self.n_components
        return self

    @append_log
    def hmmfit(self) -> sfHMMn:
        """
        HMM paramter optimization by Forward-Backward algorithm, and state inference by
        Viterbi algorithm.
        """
        if self.n_data <= 0:
            raise sfHMMAnalysisError("Cannot start analysis before appending data.")

        self.data_raw_all = self.data_raw
        self.states_list = [sf.states for sf in self]

        self._set_hmm_params()

        self.fit(self.data_raw_all.reshape(-1, 1), lengths=self.size)

        for sf in self:
            self._copy_params(sf)
            sf.states = sf.predict(np.asarray(sf.data_raw).reshape(-1, 1))
            sf.viterbi = sf.means_[sf.states, 0]
        del self.data_raw_all, self.states_list
        return self

    def run_all_separately(self, plot_every: int = 0) -> sfHMMn:
        """
        Run the function `run_all` for every sfHMM1 object.
        Paramters
        ---------
        plot_every : int
            Every `plot_every` result will be plotted during iteration.
        """
        for i, sf in enumerate(self):
            pl = plot_every > 0 and i % plot_every == 0
            try:
                sf.run_all(plot=pl)
            except Exception as e:
                print(f"{e.__class__.name} during {i}-th trace: {e}")

        return self

    def _set_covars(self):
        self.covars_ = calc_covars(
            self.data_raw_all,
            concat(self.states_list),
            self.n_components,
        )
        return None

    def _set_means(self):
        self.means_ = self.gmm_opt.means_.copy()
        return None

    def _set_startprob(self):
        d0_list = [sf.data_raw[0] for sf in self]
        self.startprob_ = calc_startprob(d0_list, self.gmm_opt)
        return None

    def _set_transmat(self):
        self.transmat_ = calc_transmat(self.states_list, self.n_components)
        return None

    def plot(self):
        if self.gmm_opt is not None:
            self.plot_hist()
        self.plot_traces()
        return None

    def plot_hist(self):
        """
        Plot histogram of data_raw, data_fil and GMM fitting result.
        """
        with plt.style.context(self.__class__.styles):
            plt.figure(figsize=(3, 4))
            plt.suptitle(self.name, fontweight="bold")
            self._hist()
            plt.show()
        return None

    def plot_traces(
        self,
        data: str = "Viterbi path",
        filter_func: Callable | None = None,
        sharex: bool = False,
    ):
        """
        Plot all the trajectories. The figure will look like:
         __ __ __ __
        | 0| 1| 2| 3|
        |__|__|__|__|
        | 4| 5| 6| 7|
        |__|__|__|__|
        | 8| 9|10|11|
        :     :     :

        Parameters
        ----------
        data : str, default is "Viterbi path"
            Which data to plot over the raw data trajectories.
        filter_func : callable or None, optional
            If not None, only sfHMM objects that satisfy filter_func(sf)==True are
            plotted.
        sharex : bool, default is False
            If True, all the subplots will share x-limits.
        """
        c_other = self.colors.get(data, None)

        # index list that satisfies filter_func
        if filter_func is None:
            indices = range(self.n_data)
        else:
            indices = [i for (i, sf) in enumerate(self) if filter_func(sf)]

        n_col = 4
        n_row = (len(indices) - 1) // n_col + 1

        if sharex:
            xlim = [0, max(sf.size for sf in self)]

        with plt.style.context(self.__class__.styles):
            plt.figure(figsize=(n_col * 2.9, n_row * 2.4))

            for i, ind in enumerate(indices):
                sf = self[ind]
                ax = plt.subplot(n_row, n_col, i + 1)

                if data == "Viterbi path":
                    d = sf.viterbi
                elif data == "denoised":
                    d = sf.data_fil
                elif data == "step finding":
                    d = sf.step.fit
                elif data == "none":
                    d = None
                else:
                    raise ValueError(
                        "'data' must be 'step finding', 'denoised', "
                        "'Viterbi path' or 'none'"
                    )

                plot2(
                    sf.data_raw,
                    d,
                    ylim=self.ylim,
                    legend=False,
                    color1=self.colors["raw data"],
                    color=c_other,
                )

                sharex and plt.xlim(xlim)

                plt.text(
                    ax.get_xlim()[1] * 0.98,
                    ax.get_ylim()[1] * 0.98,
                    str(ind),
                    ha="right",
                    va="top",
                    color="gray",
                )

            plt.tight_layout()
            plt.show()
        return None

    def view_in_qt(self, title: str | None = None):
        """
        Open a Qt viewer and plot the results.
        """
        from .viewer import TrajectoryViewer

        if title is None:
            title = self.name

        datalist = []
        for sf in self:
            data = dict()
            sf.data_raw is None or data.update({"raw data": sf.data_raw})
            sf.step is None or data.update({"step finding": sf.step.fit})
            sf.data_fil is None or data.update({"denoised": sf.data_fil})
            sf.viterbi is None or data.update({"Viterbi path": sf.viterbi})
            datalist.append(data)

        viewer = TrajectoryViewer(
            datalist, self.__class__.styles, self.__class__.colors
        )
        viewer.setWindowTitle(title)
        viewer.show()
        return viewer

    def accumulate_transitions(self) -> list[tuple[int, int]]:
        return concat([sf.accumulate_transitions() for sf in self])

    def _accumulate_step_sizes(self) -> np.ndarray:
        if self[0].step is None:
            raise sfHMMAnalysisError("Steps are not detected yet.")
        return np.array(concat([sf.step.step_size_list for sf in self]))

    def _copy_params(self, sf: sfHMM1):
        if self.covariance_type == "spherical":
            sf.covars_ = self.covars_.ravel()
        else:
            sf.covars_ = [[self.covars_[0, 0, 0]]]
        sf.min_covar = self.min_covar
        sf.means_ = self.means_
        sf.startprob_ = self.startprob_
        sf.transmat_ = self.transmat_
        return None

    @under_development
    @append_log
    def align(
        self,
        bounds: tuple[float, float],
        bins: int = 32,
        formula: Literal["ax", "y=ax", "b", "y=x+b"] = "y=ax",
    ) -> list:
        """
        Align step finding results with `formula` transformation. The optimal parameter
        is determined by minimizing normalized mutual information of two step finding
        results: `self[0].step.fit` as the reference and `a * self[i].step.fit + b` as
        the variable.

        (1) y = ax
          __                  __    __
         |      __           |     |  |
        _|    _|  |__  -->  _|    _|  |__

        (2) y = x + b
        _                  _       __
         |__    __          |__  _|  |__
              _|  |__  -->


        Parameters
        ----------
        bounds : tuple of floats
            Bounds of parameter $a$ or $b$, i.e. optimal parameter is searched in the
            range from `bounds[0]` to `bounds[1]`.
        bins : int, default is 32
            Bin number for calculating shannon entropy and mutual information.
        formula: str, default is "y=ax"
            Formulation of transformation.

        """

        if self.n_data < 2:
            raise sfHMMAnalysisError("Cannot align datasets because n_data < 2.")
        if self[0].step is None:
            raise sfHMMAnalysisError("Cannot align datasets before step finding.")

        formula = re.sub(" ", "", formula)
        optimization_func = {
            "ax": optimize_ax,
            "y=ax": optimize_ax,
            "b": optimize_b,
            "y=x+b": optimize_b,
        }[formula]

        for sf in self[1:]:
            result = optimization_func(
                self[0].step.fit,
                sf.step.fit,
                bins=bins,
                range=self.ylim,
                bounds=[bounds],
            )

            a, b = result
            sf.data_raw[:] = a * sf.data_raw + b
            sf.ylim = [a * sf.ylim[0] + b, a * sf.ylim[1] + b]
            sf.step.fit[:] = a * sf.step.fit + b
            sf.step.mu_list[:] = a * sf.step.mu_list + b
            sf.step.step_size_list[:] = a * sf.step.step_size_list
            if sf.data_fil is not None:
                sf.data_fil[:] = a * sf.data_fil + b

        return None

    @property
    def data_raw(self) -> np.ndarray:
        return np.array(concat([sf.data_raw for sf in self]))

    @property
    def data_fil(self) -> np.ndarray:
        return np.array(concat([sf.data_fil for sf in self]))

    @property
    def _sg_list(self) -> np.ndarray:
        return np.array(concat([sf._sg_list for sf in self]))

    @property
    def size(self) -> list[int]:
        return [sf.size for sf in self]
