from itertools import product

import numpy as np
from scipy.integrate import quad, simps as simpson

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


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 energydistribution(e_a, mu, sigma):
        return np.exp(-0.5 * ((mu-e_a) / sigma) ** 2) / (np.sqrt(2 * np.pi) * sigma)

    @staticmethod
    def correlation(t, temperature, *args):
        tau0, e_m, e_b = args

        def integrand(e_a, ti, t0, mu, sigma, te):
            # correlation time would go to inf for higher energies, so we use rate
            return np.exp(-ti*EnergyBarriers.rate(t0, e_a, te)) * EnergyBarriers.energydistribution(e_a, mu, sigma)

        t = np.atleast_1d(t)
        temperature = np.atleast_1d(temperature)

        e_axis = np.linspace(max(0, e_m-50*e_b), e_m+50*e_b, num=5001)

        ret_val = np.array([simpson(integrand(e_axis, o, tau0, e_m, e_b, tt), e_axis)
                            for o in t for tt in temperature])

        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)

        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(omega, temperature):
            ret_val.append(simpson(_integrand_freq_real(e_axis, o, tau0, e_m, e_b, tt), e_axis) -
                           1j * simpson(_integrand_freq_imag(e_axis, o, tau0, e_m, e_b, tt), e_axis))

        return np.array(ret_val)

    @staticmethod
    def specdens(omega, temperature, *args):
        # in contrast to other spectral densities, it's omega and temperature
        tau0, e_m, e_b = args

        def integrand(e_a, w, t0, mu, sigma, t):
            r = EnergyBarriers.rate(t0, e_a, t)
            return r/(r**2 + w**2) * EnergyBarriers.energydistribution(e_a, mu, sigma)

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

        e_axis = np.linspace(max(0, e_m-50*e_b), e_m+50*e_b, num=5001)

        ret_val = np.array([simpson(integrand(e_axis, o, tau0, e_m, e_b, tt), e_axis)
                            for o in omega for tt in temperature])

        return ret_val

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

    @staticmethod
    def logmean(*args):
        return args[1] + args[2] / (kB * args[0])

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


def _integrate_process_imag(args):
    pass


def _integrate_process_real(args):
    omega_i, t, tau0, mu, sigma, temp_j = args
    return quad(_integrand_freq_real(), 0, 10, args=(omega_i, t, tau0, mu, sigma, temp_j))[0]


def _integrate_process_time(args):
    omega_i, t, tau0, mu, sigma, temp_j = args
    return quad(_integrand_time, 0, 10, args=(omega_i, t, tau0, mu, sigma, temp_j))[0]


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


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


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