import numpy as np
import matplotlib.pyplot as plt
from hmmlearn.utils import normalize
from scipy import special
from sfHMM._sfhmm_ext import forward, backward, viterbi, compute_log_xi_sum
from sfHMM.base import sfHMMBase


class sfHMMmotorBase(sfHMMBase):
    """Base class for sfHMM for motor stepping.

    This base class enables sparse transition probability matrix aiming at analyzing
    motor stepping trajectories. The attribute `transmat_` is generated from
    `transmat_kernel` every time it is called. Also, during M-step transmat_kernel is
    updated.
    """

    transmat_kernel: "np.ndarray"

    def __init__(
        self,
        sg0: float = -1,
        psf: float = -1,
        krange: "int | tuple[int, int] | None" = None,
        model: str = "g",
        name: str = "",
        max_stride: int = 2,
    ):
        super().__init__(
            sg0=sg0,
            psf=psf,
            krange=krange,
            model=model,
            name=name,
            covariance_type="tied",
        )
        self.max_stride = max_stride

    @property
    def transmat_(self):
        transmat = np.zeros((self.n_components, self.n_components))
        for i, p in enumerate(self.transmat_kernel):
            transmat += np.eye(self.n_components, k=i - self.max_stride) * p

        normalize(transmat, axis=1)
        return transmat

    def gmmfit(
        self,
        method: str = "Dirichlet",
        n_init: int = 1,
        random_state: int = 0,
        estimation: str = "fast",
    ):
        """
        estimation : str, optional
            How to estimate krange from step finding result.
            - "fast" ... narrower range but faster.
            - "safe" ... wider range but safer.
            or estimated krange can be directly given.
        """
        if self.krange is None:
            self._estimate_krange(estimation)
        return super().gmmfit(method, n_init, random_state)

    def tdp(self, **kwargs):
        dy_step = self._accumulate_step_sizes()
        dy_vit = np.array(
            [
                self.means_[sy] - self.means_[sx]
                for sx, sy in self.accumulate_transitions()
            ]
        )

        with plt.style.context(self.__class__.styles):
            if dy_step.size == 0:
                print("no step found")
                return None
            elif dy_vit.size > 0:
                xmin = min(dy_step.min(), dy_vit.min())
                xmax = max(dy_step.max(), dy_vit.max())
            else:
                xmin = dy_step.min()
                xmax = dy_step.max()
            plt.figure(figsize=(6, 4.4))
            # plot step sizes using step finding result
            plt.subplot(2, 1, 1)
            kw = dict(
                bins=int((self.max_stride * 2 + 1) * 5),
                color=self.__class__.colors["step finding"],
            )
            kw.update(kwargs)
            plt.hist(dy_step, **kw)
            plt.xlim(xmin, xmax)
            # plot step sizes using Viterbi path
            plt.subplot(2, 1, 2)
            kw = dict(
                bins=int((self.max_stride * 2 + 1) * 5),
                color=self.__class__.colors["Viterbi path"],
            )
            kw.update(kwargs)
            plt.hist(dy_vit, **kw)
            plt.xlim(xmin, xmax)

            plt.xlabel("step size")
            plt.show()
        return None

    def _init_sg0(self, p: float = 50):
        """Initialize 'sg0' if sg0 is negative."""
        return super()._init_sg0(p=p)

    def _check(self):
        self.startprob_ = np.asarray(self.startprob_)
        if len(self.startprob_) != self.n_components:
            raise ValueError("startprob_ must have length n_components")
        if not np.allclose(self.startprob_.sum(), 1.0):
            raise ValueError(
                f"startprob_ must sum to 1.0 (got {self.startprob_.sum():.4f})"
            )

        if not np.allclose(self.transmat_kernel.sum(), 1.0):
            raise ValueError(
                "transmat_kernel must sum to 1.0" f" (got {self.transmat_kernel.sum()})"
            )

    def _do_mstep(self, stats):
        if "s" in self.params:
            startprob_ = np.maximum(self.startprob_prior - 1 + stats["start"], 0)
            self.startprob_ = np.where(self.startprob_ == 0, 0, startprob_)
            normalize(self.startprob_)
        if "t" in self.params:
            transmat_ = np.maximum(self.transmat_prior - 1 + stats["trans"], 0)
            transmat_ = np.where(self.transmat_ == 0, 0, transmat_)
            for i in range(len(self.transmat_kernel)):
                self.transmat_kernel[i] = np.sum(
                    np.diag(transmat_, k=i - self.max_stride)
                )
            normalize(self.transmat_kernel)

        # GaussianHMM._do_mstep must be copied here.
        means_prior = self.means_prior
        means_weight = self.means_weight

        denom = stats["post"][:, None]
        if "m" in self.params:
            self.means_ = (means_weight * means_prior + stats["obs"]) / (
                means_weight + denom
            )

        if "c" in self.params:
            covars_prior = self.covars_prior
            covars_weight = self.covars_weight
            meandiff = self.means_ - means_prior

        if self.covariance_type in ("spherical", "diag"):
            c_n = (
                means_weight * meandiff**2
                + stats["obs**2"]
                - 2 * self.means_ * stats["obs"]
                + self.means_**2 * denom
            )
            c_d = max(covars_weight - 1, 0) + denom
            self._covars_ = (covars_prior + c_n) / np.maximum(c_d, 1e-5)
            if self.covariance_type == "spherical":
                self._covars_ = np.tile(
                    self._covars_.mean(1)[:, None], (1, self._covars_.shape[1])
                )
        elif self.covariance_type in ("tied", "full"):
            c_n = np.empty((self.n_components, self.n_features, self.n_features))
            for c in range(self.n_components):
                obsmean = np.outer(stats["obs"][c], self.means_[c])
                c_n[c] = (
                    means_weight * np.outer(meandiff[c], meandiff[c])
                    + stats["obs*obs.T"][c]
                    - obsmean
                    - obsmean.T
                    + np.outer(self.means_[c], self.means_[c]) * stats["post"][c]
                )
            cvweight = max(covars_weight - self.n_features, 0)
            if self.covariance_type == "tied":
                self._covars_ = (covars_prior + c_n.sum(axis=0)) / (
                    cvweight + stats["post"].sum()
                )
            elif self.covariance_type == "full":
                self._covars_ = (covars_prior + c_n) / (
                    cvweight + stats["post"][:, None, None]
                )

    def _do_viterbi_pass(self, framelogprob: np.ndarray):
        n_samples, n_components = framelogprob.shape
        state_sequence, logprob = viterbi(
            n_samples,
            n_components,
            log_mask_zero(self.startprob_),
            log_mask_zero(self.transmat_kernel),
            framelogprob,
            self.max_stride,
        )
        return logprob, state_sequence

    def _do_forward_pass(self, framelogprob: np.ndarray):
        n_samples, n_components = framelogprob.shape
        fwdlattice = forward(
            n_samples,
            n_components,
            log_mask_zero(self.startprob_),
            log_mask_zero(self.transmat_kernel),
            framelogprob,
            self.max_stride,
        )
        with np.errstate(under="ignore"):
            return special.logsumexp(fwdlattice[-1]), fwdlattice

    def _do_backward_pass(self, framelogprob: np.ndarray):
        n_samples, n_components = framelogprob.shape
        bwdlattice = backward(
            n_samples,
            n_components,
            log_mask_zero(self.transmat_kernel),
            framelogprob,
            self.max_stride,
        )
        return bwdlattice

    def _accumulate_sufficient_statistics(
        self, stats, obs, framelogprob, posteriors, fwdlattice, bwdlattice
    ):
        stats["nobs"] += 1
        if "s" in self.params:
            stats["start"] += posteriors[0]
        if "t" in self.params:
            n_samples, n_components = framelogprob.shape
            # when the sample is of length 1, it contains no transitions
            # so there is no reason to update our trans. matrix estimate
            if n_samples <= 1:
                return

            log_xi_sum = compute_log_xi_sum(
                n_samples,
                n_components,
                fwdlattice,
                log_mask_zero(self.transmat_kernel),
                bwdlattice,
                framelogprob,
                self.max_stride,
            )
            with np.errstate(under="ignore"):
                stats["trans"] += np.exp(log_xi_sum)

        if "m" in self.params or "c" in self.params:
            stats["post"] += posteriors.sum(axis=0)
            stats["obs"] += np.dot(posteriors.T, obs)

        if "c" in self.params:
            if self.covariance_type in ("spherical", "diag"):
                stats["obs**2"] += np.dot(posteriors.T, obs**2)
            elif self.covariance_type in ("tied", "full"):
                # posteriors: (nt, nc); obs: (nt, nf); obs: (nt, nf)
                # -> (nc, nf, nf)
                stats["obs*obs.T"] += np.einsum("ij,ik,il->jkl", posteriors, obs, obs)

    def _get_n_fit_scalars_per_param(self):
        nc = self.n_components
        nf = self.n_features
        return {
            "s": nc - 1,
            "t": self.max_stride * 2 + 1,
            "m": nc * nf,
            "c": {
                "spherical": nc,
                "diag": nc * nf,
                "full": nc * nf * (nf + 1) // 2,
                "tied": nf * (nf + 1) // 2,
            }[self.covariance_type],
        }


def log_mask_zero(a):
    with np.errstate(divide="ignore"):
        return np.log(a)
