import numpy as np
from scipy.integrate import quad

from .base import Distribution


class FFHS(Distribution):
    name = 'Intermolecular (FFHS)'
    parameter = None

    @staticmethod
    def distribution(taus, tau0, *args):
        # Distribution over tau, not log(tau), like in other cases
        z = taus / tau0
        return 54 * np.sqrt(z) / (162*z**3 + 18*z**2 - 4*z + 2) / np.pi / tau0

    @staticmethod
    def correlation(t, tau0, *args):
        def integrand(u, tt, tau0):
            return FFHS.distribution(u, tau0) * np.exp(-tt/u) / u

        ret_val = np.array([quad(integrand, 0, np.infty, args=(tt, tau0), epsabs=1e-12, epsrel=1e-12)[0] for tt in t])

        return ret_val

    @staticmethod
    def specdens(omega, tau0, *args):
        def integrand(u, o, tau0):
            return u**4 * tau0 / (81 + 9*u**2 - 2*u**4 + u**6) / (u**4 + (o*tau0)**2)
            # return FFHS.distribution(u, tau0) * u / (1+o**2 * u**2)

        ret_val = np.array([quad(integrand, 0, np.infty, args=(o, tau0), epsabs=1e-12, epsrel=1e-12)[0] for o in omega])

        return ret_val

    @staticmethod
    def susceptibility(omega, tau0, *args):
        def integrand_real(u, o, tau0):
            return FFHS.distribution(u, tau0) / (1+o**2 * u**2)

        def integrand_imag(u, o, tau0):
            return FFHS.distribution(u, tau0) * o*u / (1+o**2 * u**2)

        ret_val = np.array([quad(integrand_real, 0, np.infty, args=(o, tau0),
                                 epsabs=1e-12, epsrel=1e-12)[0] for o in omega], dtype=complex)

        ret_val.imag += np.array([quad(integrand_imag, 0, np.infty, args=(o, tau0),
                                       epsabs=1e-12, epsrel=1e-12)[0] for o in omega])

        return ret_val


# class Bessel(Distribution):
#     name = 'Intermolecular (Bessel)'
#     parameter = None
#
#     @staticmethod
#     def distribution(taus, tau0, *args):
#         x = np.sqrt(tau0 / taus)
#         return special.jv(1.5, x)**2 / tau0 / 2
#
#     @staticmethod
#     def correlation(t, tau0, *args):
#         def integrand(lx, tt):
#             x = np.exp(lx)
#             return Bessel.distribution(x, tau0)*np.exp(-tt/lx)
#
#         logaxis = np.linspace(-20+np.log(tau0), 20+np.log(tau0), num=5001)
#
#         ret_val = np.array([simps(integrand(logaxis, tt), logaxis) for tt in t])
#
#         return ret_val
#
#     @staticmethod
#     def susceptibility(omega, tau0, *args):
#         def integrand(lx, o):
#             x = np.exp(lx)
#             return Bessel.distribution(x, tau0) * 1/(1+1j*omega*x)
#
#         logaxis = np.linspace(-20 + np.log(tau0), 20 + np.log(tau0), num=5001)
#
#         ret_val = np.array([simps(integrand(logaxis, o), logaxis) for o in omega])
#
#         return ret_val
