from __future__ import annotations

import abc
from typing import List

import numpy as np

from yupi.stats._stats import _kurtosis
from yupi.trajectory import Trajectory

DEFAULT_ZERO_THRESHOLD = 1e-12
"""Default threshold to check if a value is zero."""


class Featurizer(abc.ABC):
    """Base class for featurizers.

    A featurizer is a class that takes a set of trajectories and returns a
    feature matrix. The feature matrix is a 2D numpy array with shape
    (n_trajs, n_features)

    Parameters
    ----------
    zero_threshold : float, optional
        Threshold to check if a value is zero, by default DEFAULT_ZERO_THRESHOLD
    """

    def __init__(self, zero_threshold: float = DEFAULT_ZERO_THRESHOLD):
        self.zero_threshold = zero_threshold

    def is_zero(self, val: float) -> bool:
        """Check if a value is zero using the zero threshold."""
        return abs(val) < self.zero_threshold

    @property
    @abc.abstractmethod
    def count(self) -> int:
        """Number of features generated by this featurizer."""

    @abc.abstractmethod
    def featurize(self, trajs: List[Trajectory]) -> np.ndarray:
        """Featurize a list of trajectories.

        Parameters
        ----------
        trajs : List[Trajectory]
            A list of trajectories.

        Returns
        -------
        np.ndarray
            A numpy array of shape (n_trajs, n_features).
        """

    def __add__(self, other: Featurizer) -> CompoundFeaturizer:
        return self.append(other)

    def append(self, other: Featurizer) -> CompoundFeaturizer:
        """Append another featurizer to this one.

        Parameters
        ----------
        other : Featurizer
            The featurizer to append.

        Returns
        -------
        CompundFeaturizer
            A new featurizer that is the concatenation of this one and the
            other one.
        """
        return CompoundFeaturizer(self, other)


class CompoundFeaturizer(Featurizer):
    """
    Gather multiple featurizers into one.

    Parameters
    ----------
    *featurizers : Featurizer
        The featurizers to gather.
    """

    def __init__(self, *featurizers: Featurizer):
        super().__init__()
        self.featurizers = featurizers
        self._count = sum(f.count for f in featurizers)
        if not self._valid_compound():
            raise ValueError(
                "Featurizers names must be unique. Featurizers used:\n"
                f"{self.featurizers_names()}"
            )

    def _valid_compound(self) -> bool:
        """Check that the compound featurizer is valid."""
        feat_names = self.featurizers_names()
        return len(feat_names) == len(set(feat_names))

    def featurizers_names(self) -> List[str]:
        """Return the names of the featurizers."""
        names = []
        for featurizer in self.featurizers:
            if isinstance(featurizer, CompoundFeaturizer):
                names.extend(featurizer.featurizers_names())
            else:
                names.append(featurizer.__class__.__name__)
        return names

    @property
    def count(self) -> int:
        """Number of features generated by this featurizer."""
        return self._count

    def featurize(self, trajs: List[Trajectory]) -> np.ndarray:
        """Featurize a list of trajectories.

        Parameters
        ----------
        trajs : List[Trajectory]
            A list of trajectories.

        Returns
        -------
        np.ndarray
            A numpy array of shape (n_traj, n_features).
        """
        feats = np.empty((len(trajs), self.count))
        idx = 0
        for featurizer in self.featurizers:
            feats[:, idx : idx + featurizer.count] = featurizer.featurize(trajs)
            idx += featurizer.count
        return feats


class GlobalStatsFeaturizer(Featurizer):
    """Featurizer that extracts all the global features related to an array.

    Statistics used:
        - mean
        - median
        - kurtoisis
        - autocorrelation coefficient
        - min
        - max
        - range
        - std
        - var
        - coeff of var
        - iqr
    """

    @property
    def count(self) -> int:
        return 11

    @abc.abstractmethod
    def _values(self, traj: Trajectory) -> np.ndarray:
        """Get the array to extract the features from.

        Parameters
        ----------
        traj : Trajectory
            The trajectory.

        Returns
        -------
        np.ndarray
            The array to extract the features from.
        """

    def _acc_k(self, values: np.ndarray, k: int, mean: float) -> float:
        N = len(values)  # pylint: disable=invalid-name
        _sum = np.sum((values[1 : N - k] - mean) * (values[1 + k : N] - mean))
        return _sum / N

    def _get_global_stats(self, traj: Trajectory) -> np.ndarray:
        values = self._values(traj)
        mean = np.mean(values)
        c_1 = self._acc_k(values, 1, mean)
        c_0 = self._acc_k(values, 0, mean)
        stats = np.array(
            [
                float(val)
                for val in [
                    mean,
                    np.median(values),
                    _kurtosis(values),
                    0 if self.is_zero(c_0) else c_1 / c_0,
                    np.min(values),
                    np.max(values),
                    np.ptp(values),
                    np.std(values),
                    np.var(values),
                    0 if self.is_zero(mean) else np.std(values) / mean,
                    np.subtract(*np.percentile(values, [75, 25])),
                ]
            ]
        )
        assert len(stats) == self.count
        return stats

    def featurize(self, trajs: List[Trajectory]) -> np.ndarray:
        feats = np.empty((len(trajs), self.count))
        for i, traj in enumerate(trajs):
            feats[i, :] = self._get_global_stats(traj)
        return feats
