#  Orakel is a library providing different regression models.
#  Copyright (C) 2025 culpeo <culpeo@dismail.de>

#  This program is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or (at your
#  option) any later version.

#  This program is distributed in the hope that it will be useful, but WITHOUT
#  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
#  FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License
#  for more details.

#  You should have received a copy of the GNU Lesser General Public License
#  along with this program. If not, see https://www.gnu.org/licenses/.

from dataclasses import dataclass

import numpy as np
from scipy.optimize import curve_fit
from scipy.special import lambertw

from orakel.metrics import calculate_aic, calculate_r2


def exponential_decay(x: np.ndarray, a: float, k: float, b: float = 0.0):
    """
    Exponential decay with optional non-zero asymptote.

    y = b + (a - b) * exp(-k * x)

    Parameters
    ----------
    x : np.ndarray
        The independent variable, for example the time.
    a : float
        Initial value at x = 0.
    k : float
        Rate constant.
    b : float, optional
        Lower limit of function (asymptote, by default zero).

    Returns
    -------
    y : np.ndarray
        The dependent variable.
    """
    return b + (a - b) * np.exp(-k * x)


def _specify_exponential_decay(
    x: np.ndarray, y: np.ndarray, non_zero_asymptote: bool = False
):
    """Guess parameters and specify limits for non-linear regression."""
    a = y[0]
    b = y[-1] if non_zero_asymptote else 0.0
    k = (a - b) / x.max() if x.max() != 0 else 1.0
    guessed_params = (a, k, b) if non_zero_asymptote else (a, k)
    lower = (0.0, ) * len(guessed_params)
    upper = (a * 2, np.inf, a) if non_zero_asymptote else (a * 2, np.inf)
    param_limits = (lower, upper)
    return guessed_params, param_limits


def double_exponential_decay(
    x: np.ndarray, a1: float, k1: float, a2: float, k2: float, b: float = 0.0
):
    """
    Double exponential decay with optional non-zero asymptote.

    b + (a1 - a2) * exp(-k1 * x) + (a2 - b) * exp(-k2 * x)

    Parameters
    ----------
    x : np.ndarray
        The independent variable, for example the time.
    a1 : float
        Initial value at x = 0.
    k1 : float
        Rate constant of the first phase of exponential decay.
    a2 : float
        Initial value of the second phase of exponential decay.
    k2 : float
        Rate constant of the second phase of exponential decay.
    b : float, optional
        Lower limit of function (asymptote, by default zero).

    Returns
    -------
    y : np.ndarray
        The dependent variable.
    """
    return b + (a1 - a2) * np.exp(-k1 * x) + (a2 - b) * np.exp(-k2 * x)


def _specify_double_exponential_decay(
    x: np.ndarray, y: np.ndarray, non_zero_asymptote: bool = False
):
    """Guess parameters and specify limits for non-linear regression."""
    a1 = y[0]
    b = y[-1] if non_zero_asymptote else 0.0
    k1 = (a1 - b) / x.max() if x.max() != 0 else 1.0
    a2 = a1 / 2
    k2 = k1 / 10
    if non_zero_asymptote:
        guessed_params = (a1, k1, a2, k2, b)
        upper = (a1 * 2, np.inf, a1, np.inf, a1)
    else:
        guessed_params = (a1, k1, a2, k2)
        upper = (a1 * 2, np.inf, a1, np.inf)
    lower = (0.0, ) * len(guessed_params)
    param_limits = (lower, upper)
    return guessed_params, param_limits


def gompertz(x: np.ndarray, a: float, ku: float, xi: float):
    """
    Gompertz' generalised logistic function.

    y = a * exp(-exp(-(exp(1) * ku * (x - xi) / a)))

    Special case of Richards generalised logistic function with three or four
    parameters and the inflection occurring at 36.79 % of the upper asymptote
    after Benjamin Gompertz. Re-parameterised by Tjørve and Tjørve 2017.

    Parameters
    ----------
    a : float
        Upper limit of function (upper asymptote).
    ku : float
        Absolute maximum growth rate (slope at inflection point).
    xi : float
        Abscissa of inflection point.

    Returns
    -------
    y : np.ndarray
        The dependent variable.

    References
    ----------
    Tjørve KMC, Tjørve E. The use of Gompertz models in growth analyses, and
    new Gompertz-model approach: An addition to the Unified-Richards family.
    PLoS One. 2017;12(6):e0178691. doi:10.1371/journal.pone.0178691
    """
    # a * np.exp(-np.exp(-np.exp(1) * ku * (x - xi)))
    return a * np.exp(-np.exp(-(np.exp(1) * ku * (x - xi) / a)))


def _specify_gompertz(x: np.ndarray, y: np.ndarray):
    """Guess parameters and specify limits for non-linear regression."""
    a = y.max()
    ku = (a - y.min()) / x.max() if x.max() != 0 else 1.0
    xi = x[np.argmin(abs(y - a * 0.368))]
    guessed_params = (a, ku, xi)
    param_limits = (0.0, (a * 2, np.inf, x.max()))
    return guessed_params, param_limits


def richards(x: np.ndarray, a: float, d: float, ku: float, xi: float):
    """
    Richards generalised logistic function.

    Unified variant of Richards generalised logistic function with four
    parameters, re-parameterised by Tjørve and Tjørve 2010.

    Parameters
    ----------
    a : float
        Upper limit of function (upper asymptote)
    d : float
        Proportion of upper asymptote at inflection; a * (d ** (1 / (1 - d)))
    ku : float
        Relative maximum growth rate (slope at inflection point); Ku = ku * a
    xi : float
        Abscissa of inflection point

    Returns
    -------
    y : np.ndarray
        The dependent variable.

    References
    ----------
    Tjørve E, Tjørve KMC. A unified approach to the Richards-model family for
    use in growth analyses: Why we need only two model forms. Journal of
    Theoretical Biology. 2010;267(3):417-425. doi:10.1016/j.jtbi.2010.09.008
    """
    exponent = -ku * (x - xi) / d ** (d / (1 - d))
    base = 1 + (d - 1) * np.exp(exponent)
    return a * base ** (1 / (1 - d))


def _specify_richards(x: np.ndarray, y: np.ndarray):
    """Guess parameters and specify limits for non-linear regression."""
    a = y.max()
    d = 1.5
    ku = (a - y.min()) / x.max() / a if x.max() != 0 else 1.0 / a
    xi = x[np.argmin(abs(y - a / 2))]
    guessed_params = (a, d, ku, xi)
    param_limits = (0.0, (a * 2, 2, np.inf, x.max()))
    return guessed_params, param_limits


def michaelis_menten(x: np.ndarray, km: float, vmax: float):
    """
    The Michaelis-Menten equation v = V_max * x / (K_M + x).

    Parameters
    ----------
    x : array-like with one dimension
        Substrate concentration (the independent variable).
    km : float
        Substrate concentration at half-maximal velocity (Michaelis constant).
    vmax : float
        Maximum reaction rate or velocity.

    Returns
    -------
    y : np.ndarray
        Reaction rate for given substrate concentration(s).
    """
    return vmax * x / (km + x)


def _specify_michaelis_menten(x: np.ndarray, y: np.ndarray):
    """Guess parameters and specify limits for non-linear regression."""
    vmax = y.max()
    km = x[np.argmin(abs(y - vmax / 2))]
    guessed_params = (km, vmax)
    param_limits = (1e-6, (x.max(), vmax * 2))
    return guessed_params, param_limits


def schnell_mendoza(t: np.ndarray, s0: float, km: float, vmax: float):
    """
    Explicit form of the Michaelis-Menten equation using Lambert's W function.

    Defined with the substrate s, the substrate concentration at half-maximal
    velocity K_M, and the maximum reaction rate or velocity V_max:

        s = K_M * W((s0 / K_M) * exp((s0 - V_max * t) / K_M))

    Parameters
    ----------
    t : np.ndarray
        Time (the independent variable).
    s0 : float
        Initial substrate concentration (i. e., at t = 0).
    km : float
        Substrate concentration at half-maximal velocity (Michaelis constant).
    vmax : float
        Maximum reaction rate or velocity.

    Returns
    -------
    s : np.ndarray
        Substrate concentration at a given time.

    References
    ----------
    1. Schnell S, Mendoza C. Closed form solution for time-dependent enzyme
       kinetics. Journal of Theoretical Biology. 1997;187(2):207-212.
       doi:10.1006/jtbi.1997.0425
    2. Goudar CT, Harris SK, McInerney MJ, Suflita JM. Progress curve analysis
       for enzyme and microbial kinetic reactions using explicit solutions
       based on the Lambert W function. Journal of Microbiological Methods.
       2004;59(3):317-326. doi:10.1016/j.mimet.2004.06.013
    """
    arg = (s0 / km) * np.exp((s0 - vmax * t) / km)
    return km * lambertw(arg).real


def power_function(x: np.ndarray, a: float, b: float):
    """
    The power function y = a * x ** b.

    Parameters
    ----------
    x : np.ndarray
        The independent variable, for example the time.
    a : float
        Coefficient. Acts as a multiplier.
    b : float
        Power. The exponent of the power function.

    Returns
    -------
    y : np.ndarray
        The dependent variable.
    """
    return a * x ** b


def _specify_power_function(x: np.ndarray, y: np.ndarray):
    """Guess parameters and specify limits for non-linear regression."""
    x1, x2 = np.percentile(x, [5.0, 95.0])
    y1, y2 = np.percentile(y, [5.0, 95.0])
    b = (np.log(y2) - np.log(y1)) / (np.log(x2) - np.log(x1))
    a = np.exp(np.log(y1) - b * np.log(x1))
    guessed_params = (a, b)
    param_limits = ((0.0, -5.0), (a * 5, 5.0))
    return guessed_params, param_limits


@dataclass
class ModelSpecification:
    """Specification for a non-linear model."""
    function: callable
    auto_params: callable

MODELS = {
    'double_exponential_decay': ModelSpecification(
        function=double_exponential_decay,
        auto_params=_specify_double_exponential_decay
    ),
    'exponential_decay': ModelSpecification(
        function=exponential_decay,
        auto_params=_specify_exponential_decay
    ),
    'gompertz': ModelSpecification(
        function=gompertz,
        auto_params=_specify_gompertz
    ),
    'michaelis_menten': ModelSpecification(
        function=michaelis_menten,
        auto_params=_specify_michaelis_menten
    ),
    'power_function': ModelSpecification(
        function=power_function,
        auto_params=_specify_power_function
    ),
    'richards': ModelSpecification(
        function=richards,
        auto_params=_specify_richards
    )
}


class NonlinearRegression:
    """
    Perform regression analysis with the selected non-linear model.

    Parameters
    ----------
    model_name : str
        The regression model to be used.
    **model_kwargs : additional, model-specific arguments
        Optional arguments for the selected model (e. g., flag for model of
        exponential decay non_zero_asymptote=True).

    Raises
    ------
    ValueError
        Raises an error if the selected model is not available.

    Returns
    -------
    self : object
    """

    def __init__(self, model_name: str, **model_kwargs):
        if model_name not in MODELS:
            raise ValueError(f'Select a model from: {list(MODELS.keys())}.')
        self._model = MODELS[model_name]
        self._model_kwargs = model_kwargs
        self.params_ = None
        self.scores_ = None

    def fit(self, x: np.ndarray, y: np.ndarray):
        """
        Estimate parameters of the selected model.

        Parameters
        ----------
        x : array-like with one dimension
            Independent variable (e. g., time)
        y : array-like of shape (n_samples, n_points)
            Dependent variable (e. g., measured signal)

        Returns
        -------
        self : object
            Regression model describing relationship between x and y with the
            optimal parameters and scores as attributes:

            params_ : np.ndarray of shape (n_samples, n_params)
            scores_ : np.ndarray of shape (n_samples, 2)
                Scores of model and estimation: coefficient of determination
                and Akaike information criterion (AIC).
        """
        if x.ndim != 1:
            raise ValueError('x must be a one-dimensional array.')
        if y.ndim == 1:
            y = y[np.newaxis, :]

        n_samples, n_points = y.shape
        if n_points != x.size:
            raise ValueError(f'y has {n_points} points but x has {x.size}.')

        # Cache values before loop
        model = self._model.function
        model_kwargs = self._model_kwargs
        specify_model = self._model.auto_params

        # Determine number of model parameters
        guessed_params, _ = specify_model(x, y[0], **model_kwargs)
        n_params = len(guessed_params)

        params = np.empty((n_samples, n_params), dtype=float)
        scores = np.empty((n_samples, 2), dtype=float)

        for i in range(n_samples):
            yi = y[i]

            # Fit model based on initial estimate and model-specific limits
            guessed_params, param_limits = specify_model(x, yi, **model_kwargs)
            try:
                opt_params, _ = curve_fit(
                    model,
                    x,
                    yi,
                    p0=guessed_params,
                    bounds=param_limits
                )
            except ValueError:
                opt_params = (np.nan, ) * n_params
            params[i] = opt_params
            y_pred = model(x, *opt_params)
            scores[i] = (
                calculate_r2(yi, y_pred), calculate_aic(yi, y_pred, n_params)
            )

        self.params_ = params
        self.scores_ = scores

        return self

    def predict(self, x: np.ndarray):
        """
        Predict y for a given x using the trained model.

        Parameters
        ----------
        x : array-like with one dimension
            Independent variable (e. g., time)

        Returns
        -------
        y : np.ndarray of shape (n_samples, n_points)
            Estimated dependent variable (e. g., measured signal)
        """
        if self.params_ is None:
            raise RuntimeError('Model has not been fitted yet.')
        if x.ndim != 1:
            raise ValueError('x must be an one-dimensional array.')

        params = (
            self.params_[:, p][:, None] for p in range(self.params_.shape[1])
        )
        return self._model.function(x, *params)
