from multiprocessing import Pool, cpu_count
from itertools import product

import numpy as np

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

from .base import Distribution


__all__ = ['LogGaussian']


class LogGaussian(Distribution):
    name = 'Log-Gaussian'
    parameter = [r'\sigma']
    bounds = [(0, 10)]

    @staticmethod
    def distribution(tau, tau0, sigma: float):
        return np.exp(-0.5*(np.log(tau/tau0)/sigma)**2)/np.sqrt(2*np.pi)/sigma

    @staticmethod
    def correlation(t, tau0, sigma: float):
        _t = np.atleast_1d(t)
        _tau = np.atleast_1d(tau0)

        pool = Pool(processes=min(cpu_count(), 4))
        integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_t, _tau)]

        with np.errstate(divide='ignore'):
            res = np.array(pool.map(_integrate_process_time, integration_ranges))
        ret_val = res.reshape((_t.shape[0], _tau.shape[0]))

        return ret_val.squeeze()

    @staticmethod
    def susceptibility(omega, tau0, sigma: float):
        _omega = np.atleast_1d(omega)
        _tau = np.atleast_1d(tau0)

        pool = Pool(processes=min(cpu_count(), 4))
        integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_omega, _tau)]

        with np.errstate(divide='ignore'):
            res_real = np.array(pool.map(_integrate_process_imag, integration_ranges))
            res_imag = np.array(pool.map(_integrate_process_real, integration_ranges))
        ret_val = (res_real+1j*res_imag).reshape((_omega.shape[0], _tau.shape[0]))

        return ret_val.squeeze()

    @staticmethod
    def specdens(omega, tau0, sigma):
        _omega = np.atleast_1d(omega)
        _tau = np.atleast_1d(tau0)

        pool = Pool(processes=min(cpu_count(), 4))
        integration_ranges = [(omega_i, tau_j, sigma) for (omega_i, tau_j) in product(_omega, _tau)]

        with np.errstate(divide='ignore'):
            res = np.array(pool.map(_integrate_process_imag, integration_ranges))
        ret_val = res.reshape((_omega.shape[0], _tau.shape[0]))

        ret_val /= _omega[:, None]

        return ret_val.squeeze()

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


def _integrate_process_imag(args):
    omega_i, tau_j, sigma = args
    area = quad(_integrand_freq_imag_high, 0, 50, args=(omega_i, tau_j, sigma))[0]
    area += quad(_integrand_freq_imag_low, -50, 0, args=(omega_i, tau_j, sigma))[0]

    return area


def _integrate_process_real(args):
    omega_i, tau_j, sigma = args
    area = quad(_integrand_freq_real_high, 0, 50, args=(omega_i, tau_j, sigma))[0]
    area += quad(_integrand_freq_real_low, -50, 0, args=(omega_i, tau_j, sigma))[0]

    return area


def _integrate_process_time(args):
    omega_i, tau_j, sigma = args
    return quad(_integrand_time, -50, 50, args=(omega_i, tau_j, sigma))[0]


def _integrand_time(u, t, tau, sigma):
    uu = np.exp(u)
    return LogGaussian.distribution(uu, tau, sigma) * np.exp(-t/uu)


# integrands
def _integrand_freq_imag_low(u, omega, tau, sigma):
    # integrand
    uu = np.exp(u)
    return LogGaussian.distribution(uu, tau, sigma) * omega * uu / (1 + (omega*uu)**2)


def _integrand_freq_imag_high(u, omega, tau, sigma):
    uu = np.exp(-u)
    return LogGaussian.distribution(1/uu, tau, sigma) * omega * uu / (uu**2 + omega**2)


def _integrand_freq_real_low(u, omega, tau, sigma):
    uu = np.exp(u)
    return LogGaussian.distribution(uu, tau, sigma) / (1 + (omega*uu)**2)


def _integrand_freq_real_high(u, omega, tau, sigma):
    uu = np.exp(-2*u)
    return LogGaussian.distribution(np.exp(u), tau, sigma) * uu / (uu + omega**2)
