import ctypes
from multiprocessing import Pool, cpu_count
from itertools import product
from typing import Callable

import numpy as np
from scipy import LowLevelCallable
from scipy.special import erf

from nmreval.lib.utils import ArrayLike

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

from nmreval.distributions.helper import HAS_C_FUNCS, lib
from nmreval.distributions.base import Distribution


__all__ = ['LogGaussian']


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

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

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

        if HAS_C_FUNCS:
            res = _integrate_correlation_c(_t, _tau, sigma)
        else:
            res = _integration_parallel(_t, _tau, sigma, _integrate_process_time)

        return res.squeeze()

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

        if HAS_C_FUNCS:
            res_real = _integrate_susc_real_c(_omega, _tau, sigma)
            res_imag = _integrate_susc_imag_c(_omega, _tau, sigma)
        else:
            res_real = _integration_parallel(_omega, _tau, sigma, _integrate_process_imag)
            res_imag = _integration_parallel(_omega, _tau, sigma, _integrate_process_real)

        return (res_real + 1j * res_imag).squeeze()

    @staticmethod
    def specdens(omega: ArrayLike, tau: ArrayLike, sigma: float) -> np.ndarray:
        _omega = np.atleast_1d(omega)
        _tau = np.atleast_1d(tau)

        if HAS_C_FUNCS:
            ret_val = _integrate_susc_imag_c(_omega, _tau, sigma)
        else:
            ret_val = _integration_parallel(_omega, _tau, sigma, _integrate_process_imag)

        ret_val /= _omega[:, None]
        ret_val[_omega == 0, :] = _tau[None, :] * np.exp(sigma**2 / 2)

        return ret_val.squeeze()

    @staticmethod
    def mean(tau, sigma):
        return tau*np.exp(sigma**2 / 2)


def _integration_parallel(x: np.ndarray, tau: np.ndarray, sigma: float, func: Callable) -> np.ndarray:
    pool = Pool(processes=min(cpu_count(), 4))
    integration_ranges = [(x_i, tau_j, sigma) for (x_i, tau_j) in product(x, tau)]

    with np.errstate(divide='ignore'):
        res = pool.map(func, integration_ranges)

    res = np.array(res).reshape((x.shape[0], tau.shape[0]))

    return res


def _integrate_susc_imag_c(omega: np.ndarray, tau: np.ndarray, sigma: float) -> np.ndarray:
    return _integrate_susc_c(lib.logGaussian_imag_low, lib.logGaussian_imag_high, omega, tau, sigma)


def _integrate_susc_real_c(omega: np.ndarray, tau: np.ndarray, sigma: float) -> np.ndarray:
    return _integrate_susc_c(lib.logGaussian_real_low, lib.logGaussian_real_high, omega, tau, sigma)


def _integrate_susc_c(lowfunc, highfunc, omega, tau, sigma):
    res = []

    for o, t in product(omega, tau):
        c = (ctypes.c_double * 3)(o, t, sigma)
        user_data = ctypes.cast(ctypes.pointer(c), ctypes.c_void_p)

        area = 0
        for (func, limits) in [(highfunc, (0, np.inf)), (lowfunc, (-np.infty, 0))]:
            epsabs = 1e-12
            while epsabs > 1e-25:
                a = quad(LowLevelCallable(func, user_data), *limits, epsabs=epsabs, epsrel=1e-12, full_output=1)
                if a[2]['last'] > 2 or a[0] < 1e-48:
                    break
                epsabs /= 10.

            area += a[0]

        res.append(area)

    res = np.asanyarray(res).reshape((omega.shape[0], tau.shape[0]))

    return res


def _integrate_process_imag(omega, tau, sigma):
    area = quad(_integrand_freq_imag_high, 0, 50, args=(omega, tau, sigma), epsabs=1e-12, epsrel=1e-12)[0]
    area += quad(_integrand_freq_imag_low, -50, 0, args=(omega, tau, sigma), epsabs=1e-12, epsrel=1e-12)[0]

    return area


def _integrate_process_real(omega: float, tau: float, sigma: float):
    area = quad(_integrand_freq_real_high, 0, 50, args=(omega, tau, sigma))[0]
    area += quad(_integrand_freq_real_low, -50, 0, args=(omega, tau, sigma))[0]

    return area


def _integrate_correlation_c(t, tau, sigma):
    res = []

    for t_i, tau_i in product(t, tau):
        c = (ctypes.c_double * 3)(t_i, tau_i, sigma)
        user_data = ctypes.cast(ctypes.pointer(c), ctypes.c_void_p)

        area = quad(LowLevelCallable(lib.logGaussianCorrelation, user_data), -np.infty, np.infty)[0]

        res.append(area)

    res = np.asanyarray(res).reshape((t.shape[0], tau.shape[0]))

    return res


def _integrate_process_time(omega, tau, sigma):
    return quad(_integrand_time, -50, 50, args=(omega, tau, sigma), epsabs=1e-12, epsrel=1e-12)[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)
