import numpy as np
import pandas as pd
from lmfit import fit_report, minimize, Parameters
import plotly.graph_objects as go


class Cosinor:
    """
    Class for Cosinor analysis. Cosinor analysis fits a cosine curve to a time series
    to summarize a rhythmic pattern using a small set of interpretable parameters.

    It answers questions like:
        - What is the typical baseline level of the signal?  -> **Mesor**
        - How large is the rhythm around that baseline?      -> **Amplitude**
        - When does the peak occur within the cycle?         -> **Acrophase**
        - What cycle length best matches the data?           -> **Period**

    Model
    -----
    This implementation uses a standard 1-harmonic cosinor model:

        y(x) = Mesor + Amplitude * cos(2π/Period * x + Acrophase)

    where `x`is time expressed in sample units (e.g., minutes since start),
    not necessarily clock time.

    Typical usage
    -------------
    >>> model = Cosinor()
    >>> result = model.fit(series)
    >>> mesor = result.params['Mesor'].value
    >>> amplitude = result.params['Amplitude'].value
    >>> period = result.params['Period'].value
    >>> acrophase = result.params['Acrophase'].value
    >>> best = model.best_fit(series, result.params)
    >>> fig = model.plot(series, result.params)

    References
    ----------
    This code is derived from the original implementation in pyActigraphy, distributed under the BSD 3-Clause License.
    Original author: Grégory Hammad (gregory.hammad@uliege.be).

    [1] Hammad, G., Reyt, M., Beliy, N., Baillet, M., Deantoni, M., Lesoinne, A., Muto, V., & Schmidt, C. (2021).
    pyActigraphy: Open-source python package for actigraphy data visualization and analysis.
    PLoS Computational Biology, 17(10), 1009514–1009535. https://doi.org/10.1371/journal.pcbi.1009514

    [2] Hammad, G., Wulff, K., Skene, D. J., Münch, M., & Spitschan, M. (2024). Open-Source Python Module for the
    Analysis of Personalized Light Exposure Data from Wearable Light Loggers and Dosimeters.
    LEUKOS, 20(4), 380–389. https://doi.org/10.1080/15502724.2023.2296863

    [3] Cornelissen, G. (2014). Cosinor-based rhythmometry. Theoretical Biology and Medical Modelling, 11(1), 16.
    https://doi.org/10.1186/1742-4682-11-16

    """
    @staticmethod
    def _cosinor(x, params):
        """
        1-harmonic cosine function

        Parameters
        ----------
        x : np.ndarray
            Time values expressed in sample units (e.g., minutes since start).
        params : lmfit.Parameters
            Must contain: 'Amplitude', 'Acrophase', 'Period', 'Mesor'.

        Returns
        -------
        np.ndarray
            Model values evaluated at `x`.
        """
        A = params['Amplitude']
        phi = params['Acrophase']
        T = params['Period']
        M = params['Mesor']
        return M + A * np.cos(2 * np.pi / T * x + phi)

    @staticmethod
    def _residual(params, x, data, fit_func):
        """
        Residuals to be minimized by the optimizer.

        Residuals are simply: observed - model.

        Parameters
        ----------
        params : lmfit.Parameters
            Candidate parameter values during optimization.
        x : np.ndarray
            Time values (in sample units).
        data : np.ndarray
            Observed data values.
        fit_func : callable
            Model function such as `_cosinor`.

        Returns
        -------
        np.ndarray
            Residual vector used by the optimizer.

        """
        model = fit_func(x, params)
        return data - model

    def __init__(self,fit_params=None):
        """
        Create a Cosinor model with default or user-provided initial parameters.

        Parameters
        ----------
        fit_params : lmfit.Parameters, optional
            Initial guesses and bounds for the fit. If None, sensible defaults
            are created.
        """
        self.fit_func = self.__class__._cosinor
        self.fit_obj_func = self.__class__._residual

        if fit_params is None:
            fit_params = Parameters()
            # Default parameters for the cosinor fit function
            fit_params.add('Amplitude', value=50, min=0)
            fit_params.add('Acrophase', value=np.pi, min=0, max=2*np.pi)
            fit_params.add('Period', value=1440, min=0)
            fit_params.add('Mesor', value=50, min=0)
        self.fit_initial_params = fit_params

    # ----------------------------
    # Time handling
    # ----------------------------

    @staticmethod
    def _convert_timestamp_to_index(data):
        """Convert timestamps"""
        # Define the x range by converting timestamps to indices, in order to
        # deal with time series with irregular index.
        return ((data.index - data.index[0]) / data.index.freq).values

    def fit(
        self,
        data,
        params=None,
        method='leastsq',
        nan_policy='raise',
        reduce_fcn=None,
        verbose=False
    ):
        """
        Fit the actigraphy data using a cosinor function.

        Parameters
        ----------
        data : pandas.Series
            Input time series.
        params: instance of Parameters [1]_, optional.
            Initial fit parameters. If None, use the default parameters.
            Default is None.
        method: str, optional
            Name of the fitting method to use [1]_.
            Default is 'leastsq'.
        nan_policy: str, optional
            Specifies action if the objective function returns NaN values.
            One of:

            * 'raise': a ValueError is raised
            * 'propagate': the values returned from userfcn are un-altered
            * 'omit': non-finite values are filtered

            Default is 'raise'.
        reduce_fcn: str, optional
            Function to convert a residual array to a scalar value for the
            scalar minimizers. Optional values are:

            * 'None' : sum of squares of residual
            * 'negentropy' : neg entropy, using normal distribution
            * 'neglogcauchy': neg log likelihood, using Cauchy distribution

            Default is None.
        verbose: bool, optional
            If set to True, display fit informations.
            Default is False.

        Returns
        -------
        fit_results : MinimizerResult
            Fit results.

        References
        ----------
        This code is derived from the original implementation in pyActigraphy, distributed under the BSD 3-Clause License.
        Original author: Grégory Hammad (gregory.hammad@uliege.be).

        [1] Hammad, G., Reyt, M., Beliy, N., Baillet, M., Deantoni, M., Lesoinne, A., Muto, V., & Schmidt, C. (2021).
        pyActigraphy: Open-source python package for actigraphy data visualization and analysis.
        PLoS Computational Biology, 17(10), 1009514–1009535. https://doi.org/10.1371/journal.pcbi.1009514

        [2] Hammad, G., Wulff, K., Skene, D. J., Münch, M., & Spitschan, M. (2024). Open-Source Python Module for the
        Analysis of Personalized Light Exposure Data from Wearable Light Loggers and Dosimeters.
        LEUKOS, 20(4), 380–389. https://doi.org/10.1080/15502724.2023.2296863

        [3] Non-Linear Least-Squares Minimization and Curve-Fitting for Python.
        https://lmfit.github.io/lmfit-py/index.html
        """

        # Define the x range by converting timestamps to indices, in order to
        # deal with time series with irregular index.
        x = self._convert_timestamp_to_index(data)

        # Minimize residuals
        fit_results = minimize(
            self.fit_obj_func,
            self.fit_initial_params if params is None else params,
            method=method,
            args=(x, data.values, self.fit_func),
            nan_policy=nan_policy,
            reduce_fcn=reduce_fcn
        )
        # Print fit parameters if verbose
        if verbose:
            print(fit_report(fit_results))

        return fit_results

    def best_fit(self, data, params):
        """
        Best fit function of the data.

        Parameters
        ----------
        data : pandas.Series
            Originally fitted time series.
        params: instance of Parameters [1]_
            Best fit parameters.

        Returns
        -------
        bestfit_data : pandas.Series
            Time series of the best fit data.

        References
        ----------
        This code is derived from the original implementation in pyActigraphy, distributed under the BSD 3-Clause License.
        Original author: Grégory Hammad (gregory.hammad@uliege.be).

        [1] Hammad, G., Reyt, M., Beliy, N., Baillet, M., Deantoni, M., Lesoinne, A., Muto, V., & Schmidt, C. (2021).
        pyActigraphy: Open-source python package for actigraphy data visualization and analysis.
        PLoS Computational Biology, 17(10), 1009514–1009535. https://doi.org/10.1371/journal.pcbi.1009514

        [2] Hammad, G., Wulff, K., Skene, D. J., Münch, M., & Spitschan, M. (2024). Open-Source Python Module for the
        Analysis of Personalized Light Exposure Data from Wearable Light Loggers and Dosimeters.
        LEUKOS, 20(4), 380–389. https://doi.org/10.1080/15502724.2023.2296863

        [3] Non-Linear Least-Squares Minimization and Curve-Fitting for Python.
        https://lmfit.github.io/lmfit-py/index.html
        """
        # Define the x range by converting timestamps to indices, in order to
        # deal with time series with irregular index.
        x = self._convert_timestamp_to_index(data)
        y = self.fit_func(x, params)
        return pd.Series(index=data.index, data=y)

    def plot(self, data, best_params):
        layout = go.Layout(
            title="Cosinor",
            xaxis=dict(title="Date time"),
            yaxis=dict(title="Counts/period"),
            showlegend=False
        )

        fig = go.Figure(
            data=[
                go.Scatter(x=data.index.astype(str),
                           y=data, name='Raw data'),
                go.Scatter(x=self.best_fit(data, best_params).index.astype(str),
                           y=self.best_fit(data, best_params),
                           name='Best fit')
            ], layout=layout,
        )
        return fig