from itertools import product
from ctypes import c_double, cast, pointer, c_void_p

import numpy as np
from scipy import LowLevelCallable

from scipy.integrate import quad
try:
    from scipy.integrate import simps as simpson
except ImportError:
    from scipy.integrate import simpson

from .base import Distribution
from ..lib.utils import ArrayLike
from ..utils.constants import kB

from .helper import HAS_C_FUNCS, lib


# noinspection PyMethodOverriding
class EnergyBarriers(Distribution):
    name = 'Energy barriers'
    parameter = [r'\tau_{0}', r'E_{m}', r'\Delta E']

    @staticmethod
    def distribution(tau, temperature, *args):
        t0, em, sigma = args
        m = np.exp(em / (kB * temperature)) * t0
        return temperature * np.exp(-0.5 * (np.log(tau/m)*kB*temperature/sigma)**2) / np.sqrt(2*np.pi)/sigma

    @staticmethod
    def rate(t0, e_a, te):
        return np.exp(-e_a / (kB * te)) / t0

    @staticmethod
    def energy_distribution(e_a, mu, sigma):
        return np.exp(-0.5 * ((mu-e_a) / sigma) ** 2) / (np.sqrt(2 * np.pi) * sigma)

    @staticmethod
    def correlation(t: ArrayLike, temperature: ArrayLike, tau0: float, e_m: float, e_b: float) -> ArrayLike:
        t = np.atleast_1d(t)
        temperature = np.atleast_1d(temperature)

        if HAS_C_FUNCS:
            ret_val = _integrate_c(lib.energyDistCorrelation, t, temperature, tau0, e_m, e_b)
        else:
            ret_val = _integrate_py(_integrand_time, t, temperature, tau0, e_m, e_b)

        return ret_val

    @staticmethod
    def specdens(omega: ArrayLike, temperature: ArrayLike, tau0: float, e_m: float, e_b: float) -> ArrayLike:
        omega = np.atleast_1d(omega)
        temperature = np.atleast_1d(temperature)

        if HAS_C_FUNCS:
            ret_val = _integrate_c(lib.energyDist_SD, omega, temperature, tau0, e_m, e_b)
        else:
            ret_val = _integrate_py(_integrand_sd, omega, temperature, tau0, e_m, e_b)

        return ret_val

    @staticmethod
    def susceptibility(omega: ArrayLike, temperature: ArrayLike, tau0: float, e_m: float, e_b: float) -> ArrayLike:
        # in contrast to other spectral densities, it's omega and temperature

        omega = np.atleast_1d(omega)
        temperature = np.atleast_1d(temperature)

        if HAS_C_FUNCS:
            ret_val = _integrate_c(lib.energyDistSuscReal, omega, temperature, tau0, e_m, e_b) + \
                      1j * _integrate_c(lib.energyDistSuscImag, omega, temperature, tau0, e_m, e_b)
        else:
            ret_val = _integrate_py(_integrand_susc_real, omega, temperature, tau0, e_m, e_b) + \
                      1j * _integrate_py(_integrand_susc_imag, omega, temperature, tau0, e_m, e_b)

        return ret_val

    @staticmethod
    def mean(temperature, tau0, ea):
        return tau0*np.exp(ea/(kB*temperature))

    @staticmethod
    def logmean(temperature, tau0, ea):
        return tau0 + ea / (kB * temperature)

    @staticmethod
    def max(*args):
        return args[1] * np.exp(args[2] / (kB * args[0]))


# helper functions
def _integrate_c(func, omega: np.ndarray, temperature: np.ndarray, tau0: float, e_m: float, e_b: float) -> np.ndarray:
    res = []
    for o, t in product(omega, temperature):
        c = (c_double * 5)(o, tau0, e_m, e_b, t)
        user_data = cast(pointer(c), c_void_p)
        area = quad(LowLevelCallable(func, user_data), 0, np.inf, epsabs=1e-13)[0]

        res.append(area)

    ret_val = np.array(res).reshape(omega.shape[0], temperature.shape[0])

    return ret_val.squeeze()


def _integrate_py(func, axis, temp, tau0, e_m, e_b):
    x = np.atleast_1d(axis)
    temperature = np.atleast_1d(temp)

    e_axis = np.linspace(max(0., e_m - 50*e_b), e_m + 50*e_b, num=5001)
    ret_val = []
    for o, tt in product(x, temperature):
        ret_val.append(simpson(y=func(e_axis, o, tau0, e_m, e_b, tt), x=e_axis))

    ret_val = np.array(ret_val).reshape(x.shape[0], temperature.shape[0])

    return ret_val.squeeze()


# python integrands
def _integrand_sd(u, omega, tau0, mu, sigma, temp):
    r = EnergyBarriers.rate(tau0, u, temp)
    return r / (r**2 + omega**2) * EnergyBarriers.energy_distribution(u, mu, sigma)


def _integrand_susc_real(u, omega, tau0, mu, sigma, temp):
    r = EnergyBarriers.rate(tau0, u, temp)
    return 1 / (r**2 + omega**2) * EnergyBarriers.energy_distribution(u, mu, sigma)


def _integrand_susc_imag(u, omega, tau0, mu, sigma, temp):
    rate = EnergyBarriers.rate(tau0, u, temp)
    return omega * rate / (rate**2 + omega**2) * EnergyBarriers.energy_distribution(u, mu, sigma)


def _integrand_time(u, t, tau0, mu, sigma, temp):
    rate = EnergyBarriers.rate(tau0, u, temp)
    return EnergyBarriers.energy_distribution(u, mu, sigma) * np.exp(-t*rate)
