from __future__ import annotations

from abc import ABC

import numpy as np
try:
    from scipy.integrate import simpson
except ImportError:
    from scipy.integrate import simps as simpson
from scipy.special import gammaln, gamma as scipy_gamma

from nmreval.distributions.base import Distribution
from nmreval.math.logfourier import logft


class AbstractGG(Distribution, ABC):
    """
    Base class for general gamma functions
    """

    @classmethod
    def correlation(cls, t, tau0, *args):
        tt = np.atleast_1d(t)
        taus, ln_tau = AbstractGG._prepare_integration(tau0)
        g_tau = cls.distribution(taus, tau0, *args)
        ret_val = np.array([simpson(np.exp(-t_i/taus) * g_tau, x=ln_tau) for t_i in tt]).squeeze()

        return ret_val

    @classmethod
    def susceptibility(cls, omega, tau0, *args):
        r"""
        Calculate spectral density \int G(ln(tau) tau/(1+(w*tau)^2) dln(tau)
        """
        _w = np.atleast_1d(omega)
        _t = np.atleast_1d(tau0)
        ret_val = np.zeros((_w.size, _t.size), dtype=complex)

        for (i, tau_i) in enumerate(_t):
            taus, ln_tau = AbstractGG._prepare_integration(tau_i)
            g_tau = cls.distribution(taus, tau_i, *args)

            ret_val[:, i].real = np.array([simpson(g_tau * taus / (1 + (w_i*taus)**2), x=ln_tau) for w_i in _w])
            ret_val[:, i].imag = np.array([simpson(g_tau * w_i * taus / (1 + (w_i*taus)**2), x=ln_tau) for w_i in _w])

        return ret_val.squeeze()

    @classmethod
    def specdens(cls, omega: float | np.ndarray, tau0: float | np.ndarray, *args) -> float | np.ndarray:
        r"""
        Calculate spectral density \int G(ln(tau) tau/(1+(w*tau)^2) dln(tau)
        """
        _w = np.atleast_1d(omega)
        _t = np.atleast_1d(tau0)
        ret_val = np.zeros((_w.size, _t.size))

        for (i, tau_i) in enumerate(_t):
            taus, ln_tau = AbstractGG._prepare_integration(tau_i)
            g_tau = cls.distribution(taus, tau_i, *args)

            ret_val[:, i] = np.array([simpson(g_tau * taus / (1 + (w_i*taus)**2), x=ln_tau) for w_i in _w])

        return ret_val.squeeze()

    @staticmethod
    def _prepare_integration(
            tau0: float,
            limits: tuple[int, int] = (20, 20),
            num_steps: int = 4001,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Create array of correlation times for integration over ln(tau)
        Args:
            tau0 (float): center of correlation times
            limits (tuple of int): Define limits ln(tau0)-limits[0] and ln(tau0)+limits[1] of array. Default is (20, 20)
            num_steps (int): length of array. Default is 4001

        Returns:
            array of taus and array of ln(tau)
        """

        ln_tau0 = np.log(tau0)
        ln_tau = np.linspace(ln_tau0-limits[0], ln_tau0+limits[1], num=num_steps)

        return np.exp(ln_tau), ln_tau


# noinspection PyMethodOverriding
class GGAlpha(AbstractGG):
    name = r'General Gamma (alpha)'
    parameter = [r'\alpha', r'\beta']

    @staticmethod
    def distribution(taus: float | np.ndarray, tau: float, alpha: float, beta: float) -> float | np.ndarray:
        b_to_a = beta / alpha
        inv_norm = alpha * np.exp(-gammaln(b_to_a) + b_to_a * np.log(b_to_a))
        t_to_t0 = taus / tau
        ret_val = np.exp(-b_to_a * t_to_t0**alpha) * t_to_t0**beta

        return ret_val * inv_norm

    @staticmethod
    def mean(tau: float | np.ndarray, alpha: float, beta: float) -> float | np.ndarray:
        beta_by_alpha = beta/alpha
        inv_alpha = 1/alpha
        return tau * (1/beta_by_alpha)**inv_alpha * scipy_gamma(beta_by_alpha + inv_alpha) / scipy_gamma(beta_by_alpha)


# noinspection PyMethodOverriding
class GGAlphaEW(AbstractGG):
    name = r'General Gamma (alpha + EW)'
    parameter = [r'\alpha', r'\beta', r'\sigma', r'\gamma']

    @staticmethod
    def distribution(tau: float | np.ndarray, tau0: float,
                     alpha: float, beta: float, sigma: float, gamma: float) -> float | np.ndarray:

        if gamma == beta or sigma == 0:
            return GGAlpha.distribution(tau, tau0, alpha, beta)

        b_to_a = beta / alpha
        ln_b_to_a = np.log(b_to_a)
        g_to_a = gamma / alpha
        t_to_t0 = tau / tau0

        norm = (np.exp(gammaln(b_to_a)) + sigma**(gamma-beta) * np.exp(gammaln(g_to_a) + (b_to_a-g_to_a)*ln_b_to_a)) / \
            np.exp(b_to_a*ln_b_to_a) / alpha

        ret_val = np.exp(-b_to_a * t_to_t0**alpha) * t_to_t0**beta * (1 + (t_to_t0*sigma)**(gamma-beta))

        return ret_val / norm


# noinspection PyMethodOverriding
class GGBeta(AbstractGG):
    name = r'General Gamma (beta)'
    parameter = ['a', 'b']

    @staticmethod
    def distribution(tau: float | np.ndarray, tau0: float, a: float, b: float) -> float | np.ndarray:
        norm = a * (1+b) * np.sin(np.pi*b/(1+b)) * b**(b/(1+b)) / np.pi
        ret_val = b * (tau/tau0)**a + (tau/tau0)**(-a*b)

        return norm / ret_val


class GGAlphaEWBeta(Distribution):
    name = r'General Gamma (WW (\alpha + EW), \beta)'
    parameter = [
        r'\tau_{\alpha}', r'\alpha', r'\beta', r'\sigma', r'\gamma',
        r'\tau_{\beta}', 'a', 'b', 'R'
    ]

    @staticmethod
    def distribution(tau, tau0, *args):
        pass

    @staticmethod
    def correlation(t, *args):
        tau_alpha, alpha, beta, sigma, gamma, tau_beta, a, b, r = args
        if r == 0:
            return GGAlphaEW.correlation(t, tau_alpha, alpha, beta, sigma, gamma)
        elif r == 1:
            return GGBeta.correlation(t, tau_beta, a, b)

        gg_alpha = GGAlphaEW.correlation(t, tau_alpha, alpha, beta, sigma, gamma)
        gg_beta = GGBeta.correlation(t, tau_beta, a, b)

        return ((1-r)*gg_beta + r) * gg_alpha

    @staticmethod
    def susceptibility(omega, *args):
        t_axis, corr = GGAlphaEWBeta._correlation_for_ft(omega, *args)

        ft = logft(t_axis, corr, new_x=omega, mode='complex')[1]

        return 1 + omega * ft.imag + 1j * omega * ft.real

    @staticmethod
    def specdens(omega, *args):
        t_axis, corr = GGAlphaEWBeta._correlation_for_ft(omega, *args)

        return logft(t_axis, corr, new_x=omega, mode='real')[1]

    @staticmethod
    def _correlation_for_ft(omega, *args):
        # time axis for correlation function
        additional_decades = 5
        lower_limit = np.log10(1 / np.max(omega)) - additional_decades
        upper_limit = np.log10(1 / np.min(omega)) + additional_decades
        num_steps = int(len(omega) * additional_decades)
        t_axis = np.logspace(lower_limit, upper_limit, num=num_steps)

        return t_axis, GGAlphaEWBeta.correlation(t_axis, *args)
