# required to get ArrayLike to render correctly
from __future__ import annotations

import abc
from typing import Optional, Tuple

import jax.numpy as jnp
import numpy as np
from numpy.typing import ArrayLike, NDArray
from pynapple import Tsd, TsdFrame, TsdTensor

from ..type_casting import support_pynapple
from ..typing import FeatureMatrix
from ._basis import Basis, check_transform_input, min_max_rescale_samples
from ._basis_mixin import AtomicBasisMixin


class RaisedCosineBasisLinear(AtomicBasisMixin, Basis, abc.ABC):
    r"""Represent linearly-spaced raised cosine basis functions.

    This implementation is based on the cosine bumps used by Pillow et al. [1]_
    to uniformly tile the internal points of the domain.

    Parameters
    ----------
    n_basis_funcs :
        The number of basis functions.
    width :
        Width of the raised cosine. Must be $x/2$ for any integer $x\geq 3$ (e.g., 1.5,
        2, 2.5, 3, etc.)
    label :
        The label of the basis, intended to be descriptive of the task variable being processed.
        For example: velocity, position, spike_counts.

    References
    ----------
    .. [1] Pillow, J. W., Paninski, L., Uzzel, V. J., Simoncelli, E. P., & J.,
        C. E. (2005). Prediction and decoding of retinal ganglion cell responses
        with a probabilistic spiking model. Journal of Neuroscience, 25(47),
        11003–11013.

    Raises
    ------
    ValueError
        If width <= 1 or 2*width is not a positive integer. Values that do not match
        this constraint will result in the following, which contribute to aliasing:
        - No overlap between bumps (width < 1).
        - Oscillatory behavior when summing the basis elements (2*width not integer).
    """

    def __init__(
        self,
        n_basis_funcs: int,
        width: float = 2.0,
        label: Optional[str] = "RaisedCosineBasisLinear",
    ) -> None:
        AtomicBasisMixin.__init__(self, n_basis_funcs=n_basis_funcs, label=label)
        Basis.__init__(
            self,
        )
        self._n_input_dimensionality = 1
        self._check_width(width)
        self._width = width
        # for these linear raised-cosine basis functions,
        # the samples must be rescaled to 0 and 1.
        self._rescale_samples = True

    @property
    def width(self):
        """Return width of the raised cosine."""
        return self._width

    @width.setter
    def width(self, width: float):
        self._check_width(width)
        self._width = width

    @staticmethod
    def _check_width(width: float) -> None:
        """Validate the width value.

        Parameters
        ----------
        width :
            The width value to validate.

        Raises
        ------
        ValueError
            If width <= 1 or 2*width is not a positive integer. Values that do not match
            this constraint will result in the following, which contribute to aliasing:
            - No overlap between bumps (width < 1).
            - Oscillatory behavior when summing the basis elements (2*width not integer).
        """
        if width <= 1 or (not np.isclose(width * 2, round(2 * width))):
            raise ValueError(
                f"Invalid raised cosine width. Width must be strictly greater than 1 and "
                f"2*width must be a integer, but got width={width} instead! At other "
                "values, basis sampling will lead to aliasing."
            )

    @support_pynapple(conv_type="numpy")
    @check_transform_input
    def evaluate(  # call these evaluate
        self,
        sample_pts: ArrayLike | Tsd | TsdFrame | TsdTensor,
    ) -> FeatureMatrix:
        """Evaluate the raised cosine basis at the sample points.

        Parameters
        ----------
        sample_pts :
            Spacing for basis functions, holding elements on interval [0, 1].
            `sample_pts` is a n-dimensional (n >= 1) array with first axis being the samples, i.e.
            `sample_pts.shape[0] == n_samples`.

        Raises
        ------
        ValueError
            If the sample provided do not lie in [0,1].

        """
        if self._rescale_samples:
            # note that sample points is converted to NDArray
            # with the decorator.
            # copy is necessary otherwise:
            # basis1 = nmo.basis.RaisedCosineBasisLinear(5)
            # basis2 = nmo.basis.RaisedCosineBasisLog(5)
            # additive_basis = basis1 + basis2
            # additive_basis.evaluate(*([x] * 2)) would modify both inputs
            sample_pts, _ = min_max_rescale_samples(
                jnp.copy(sample_pts), getattr(self, "bounds", None)
            )

        peaks = self._compute_peaks()
        delta = peaks[1] - peaks[0]

        # reshape samples
        shape = sample_pts.shape
        sample_pts = sample_pts.reshape(
            -1,
        )

        # generate a set of shifted cosines, and constrain them to be non-zero
        # over a single period, then enforce the codomain to be [0,1], by adding 1
        # and then multiply by 0.5
        basis_funcs = 0.5 * (
            jnp.cos(
                jnp.clip(
                    np.pi * (sample_pts[:, None] - peaks[None]) / (delta * self.width),
                    -np.pi,
                    np.pi,
                )
            )
            + 1
        )
        basis_funcs = basis_funcs.reshape(*shape, basis_funcs.shape[1])
        return basis_funcs

    def _compute_peaks(self) -> NDArray:
        """
        Compute the location of raised cosine peaks.

        Returns
        -------
            Peak locations of each basis element.
        """
        return np.linspace(0, 1, self.n_basis_funcs)

    def evaluate_on_grid(self, n_samples: int) -> Tuple[NDArray, NDArray]:
        """Evaluate the basis set on a grid of equi-spaced sample points.

        Parameters
        ----------
        n_samples :
            The number of points in the uniformly spaced grid. A higher number of
            samples will result in a more detailed visualization of the basis functions.

        Returns
        -------
        X :
            Array of shape (n_samples,) containing the equi-spaced sample
            points where we've evaluated the basis.
        basis_funcs :
            Raised cosine basis functions, shape (n_samples, n_basis_funcs)
        """
        return super().evaluate_on_grid(n_samples)

    def _check_n_basis_min(self) -> None:
        """Check that the user required enough basis elements.

        Check that the number of basis is at least 2.

        Raises
        ------
        ValueError
            If n_basis_funcs < 2.
        """
        if self.n_basis_funcs < 2:
            raise ValueError(
                f"Object class {self.__class__.__name__} requires >= 2 basis elements. "
                f"{self.n_basis_funcs} basis elements specified instead"
            )


class RaisedCosineBasisLog(RaisedCosineBasisLinear, abc.ABC):
    r"""Represent log-spaced raised cosine basis functions.

    Similar to ``RaisedCosineBasisLinear`` but the basis functions are log-spaced.
    This implementation is based on the cosine bumps used by Pillow et al. [1]_
    to uniformly tile the internal points of the domain.

    Parameters
    ----------
    n_basis_funcs :
        The number of basis functions.
    width :
        Width of the raised cosine. Must be $x/2$ for any integer $x\geq 3$ (e.g., 1.5,
        2, 2.5, 3, etc.)
    time_scaling :
        Non-negative hyper-parameter controlling the logarithmic stretch magnitude, with
        larger values resulting in more stretching. As this approaches 0, the
        transformation becomes linear.
    enforce_decay_to_zero:
        If set to True, the algorithm first constructs a basis with ``n_basis_funcs + ceil(width)`` elements
        and subsequently trims off the extra basis elements. This ensures that the final basis element
        decays to 0.
    label :
        The label of the basis, intended to be descriptive of the task variable being processed.
        For example: velocity, position, spike_counts.

    References
    ----------
    .. [1] Pillow, J. W., Paninski, L., Uzzel, V. J., Simoncelli, E. P., & J.,
       C. E. (2005). Prediction and decoding of retinal ganglion cell responses
       with a probabilistic spiking model. Journal of Neuroscience, 25(47),
       11003–11013.

    """

    def __init__(
        self,
        n_basis_funcs: int,
        width: float = 2.0,
        time_scaling: float = None,
        enforce_decay_to_zero: bool = True,
        label: Optional[str] = "RaisedCosineBasisLog",
    ) -> None:
        super().__init__(
            n_basis_funcs,
            width=width,
            label=label,
        )
        # The samples are scaled appropriately in the self._transform_samples which scales
        # and applies the log-stretch, no additional transform is needed.
        self._rescale_samples = False
        if time_scaling is None:
            time_scaling = 50.0

        self.time_scaling = time_scaling
        self.enforce_decay_to_zero = enforce_decay_to_zero

    @property
    def time_scaling(self):
        """Time scaling.

        Time scaling parameter regulating the logarithmic stretch of the basis.
        """
        return self._time_scaling

    @time_scaling.setter
    def time_scaling(self, time_scaling):
        """Setter property for time_scaling."""
        self._check_time_scaling(time_scaling)
        self._time_scaling = time_scaling

    @staticmethod
    def _check_time_scaling(time_scaling: float) -> None:
        if time_scaling <= 0:
            raise ValueError(
                f"Only strictly positive time_scaling are allowed, {time_scaling} provided instead."
            )

    def _transform_samples(
        self,
        sample_pts: ArrayLike,
    ) -> NDArray:
        """
        Map the sample domain to log-space.

        Parameters
        ----------
        sample_pts :
            Sample points used for evaluating the splines,
            shape (n_samples, ).

        Returns
        -------
            Transformed version of the sample points that matches the Raised Cosine basis domain,
            shape (n_samples, ).
        """
        # rescale to [0,1]
        # copy is necessary to avoid unwanted rescaling in additive/multiplicative basis.
        sample_pts, _ = min_max_rescale_samples(
            jnp.copy(sample_pts), getattr(self, "bounds", None)
        )
        # This log-stretching of the sample axis has the following effect:
        # - as the time_scaling tends to 0, the points will be linearly spaced across the whole domain.
        # - as the time_scaling tends to inf, basis will be small and dense around 0 and
        # progressively larger and less dense towards 1.
        log_spaced_pts = jnp.log(self.time_scaling * sample_pts + 1) / jnp.log(
            self.time_scaling + 1
        )
        return log_spaced_pts

    def _compute_peaks(self) -> NDArray:
        """
        Peak location of each log-spaced cosine basis element.

        Compute the peak location for the log-spaced raised cosine basis.
        Enforcing that the last basis decays to zero is equivalent to
        setting the last peak to a value smaller than 1.

        Returns
        -------
            Peak locations of each basis element.

        """
        if self.enforce_decay_to_zero:
            # compute the last peak location such that the last
            # basis element decays to zero at the last sample.
            last_peak = 1 - self.width / (self.n_basis_funcs + self.width - 1)
        else:
            last_peak = 1
        return np.linspace(0, last_peak, self.n_basis_funcs)

    @support_pynapple(conv_type="numpy")
    @check_transform_input
    def evaluate(
        self,
        sample_pts: ArrayLike | Tsd | TsdFrame | TsdTensor,
    ) -> FeatureMatrix:
        """Evaluate the log-spaced raised cosine basis at the sample points.

        Parameters
        ----------
        sample_pts :
            Spacing for basis functions. Samples will be rescaled to the interval [0, 1].
            `sample_pts` is a n-dimensional (n >= 1) array with first axis being the samples, i.e.
            `sample_pts.shape[0] == n_samples`.

        Returns
        -------
        basis_funcs :
            Log-raised cosine basis functions, shape (n_samples, n_basis_funcs).

        Raises
        ------
        ValueError
            If the sample provided do not lie in [0,1].
        """
        return super().evaluate(self._transform_samples(sample_pts))
