from ctypes import c_double, cast, c_void_p, pointer

import numpy as np
from scipy import special as special, LowLevelCallable
from scipy.integrate import quad

from ..utils import gamma
from nmreval.distributions.helper import HAS_C_FUNCS, diffusion_lib


class Diffusion:
    type = 'Diffusion'
    name = 'Free diffusion'
    equation = r'M_{0}[exp[- \gamma^{2} g^{2} t_{ev}^{2} D(2/3t_{ev}+t_{mix})] * Relax.'
    params = ['M_{0}', 'D', 'T_{Rel.}', r'\beta_{Rel.}', 'g', 't_{ev}/t_{mix}']
    choices = [(r'\gamma', 'nucleus', gamma), ('Varied time', 't_axis', {'t_mix': 'tm', 't_ev': 'te'})]
    bounds = [(0, None), (0, None), (0, None), (0, 1), (0, None), (0, None)]

    @staticmethod
    def func(x, m0, d, trel, brel, g, t2, nucleus=2.67522128e8, t_axis='tm'):
        if t_axis == 'tm':
            tm = x
            tp = t2
            red = np.exp(-(tm/trel)**brel)
        else:
            tm = t2
            tp = x
            # T2 decay happens twice
            red = np.exp(-(tp/trel)**brel) * np.exp(-(tp/trel)**brel)
        q2 = (g * nucleus * tp)**2
        t_eff = (2*tp)/3 + tm

        return m0 * np.exp(-q2 * d * t_eff) * red


class Diffusion1D:
    type = 'Diffusion'
    name = '1D Diffusion'
    equation = r'M_{0} erf(z) / z * Relax\n' \
               r'with z = sqrt([g \gamma t_{ev}]^{2} D [2/3t_{ev} + t_{mix}])'
    params = ['M_{0}', 'D', 'T_{Rel.}', r'\beta_{Rel.}', 'g', 't_{ev}/t_{mix}']
    choices = [(r'\gamma', 'nucleus', gamma), ('Varied time', 't_axis', {'t_mix': 'tm', 't_ev': 'te'})]
    bounds = [(0, None), (0, None), (0, None), (0, 1), (0, None), (0, None)]

    @staticmethod
    def func(x, m0, d, trel, brel, g, t2, nucleus=2.67522128e8, t_axis='tm'):
        if t_axis == 'tm':
            tm = x
            tp = t2
            relax = np.exp(-(tm/trel)**brel)
        else:
            tm = t2
            tp = x
            relax = np.exp(-(tp/trel)**brel) * np.exp(-(tp/t2)**brel)
        # Callaghan eq (6.91)
        z = (tp * nucleus * g * np.sqrt(d * (tm + (2 / 3)*tp)))
        diffs = special.erf(z) / z

        return m0 * diffs * relax


class Diffusion2D(object):
    type = 'Diffusion'
    name = '2D Diffusion'
    equation = r'M_{0} exp(-z) * erfi(z)/z * Relax\n' \
               r'with z = sqrt([g \gamma t_{ev}]^{2} D [t_{mix} + 2/3t_{ev}])'
    params = ['M_{0}', 'D', 'T_{Rel.}', r'\beta_{Rel.}', 'g', 't_{ev}/t_{mix}']
    bounds = [(0, None), (0, None), (0, None), (0, 1), (0, None), (0, None)]
    choices = [(r'\gamma', 'nucleus', gamma), ('Varied time', 't_axis', {'t_mix': 'tm', 't_ev': 'te'})]

    @staticmethod
    def func(x, m0, d, trel, brel, g, t2, nucleus=2.67522128e8, t_axis='tm'):
        if t_axis == 'tm':
            tm = x
            tp = t2
            relax = np.exp(-(tm/trel)**brel)
        else:
            tm = t2
            tp = x
            relax = np.exp(-(tp/trel)**brel) * np.exp(-(tp/trel)**brel)
        q_squared = (g*nucleus*tp)**2
        a = q_squared * (2*tp/3 + tm) * d

        # Callaghan eq (6.90)
        diffs = np.exp(-a) * special.erfi(np.sqrt(a)) / np.sqrt(a)

        return m0*diffs*relax


class AnisotropicDiffusion(object):
    type = 'Diffusion'
    name = 'Anisotropic Diffusion'
    equation = r'M_{0} exp[-(g \gamma t_{p})^{2} D_{\perp}] * erf(z) / z * Relax\n' \
               r'with z = sqrt([g \gamma t_{ev}]^{2} [D_{\para} - D_{\perp}] [2/3t_{ev} + t_{mix}])'
    params = ['M_{0}', r'D_{\perp}', r'D_{\para}', r'T_{Rel.}', r'\beta_{Rel.}', 'g', 't_{ev}/t_{mix}']
    choices = [(r'\gamma', 'nucleus', gamma), ('Varied time', 't_axis', {'t_mix': 'tm', 't_ev': 'te'})]
    bounds = [(0, None), (0, None), (0, None), (0, None), (0, 1), (0, None), (0, None)]

    @staticmethod
    def func(x, m0, d_perp, d_par, trel, brel, g, t2, nucleus=2.67522128e8, t_axis='tm'):
        if t_axis == 'tm':
            tm = x
            tp = t2
            relax = np.exp(-(tm/trel)**brel)
        else:
            tm = t2
            tp = x
            relax = np.exp(-(tp/trel)**brel)*np.exp(-(tp/trel)**brel)

        q = g * nucleus * tp
        t = 2 * tp / 3 + tm
        # Callaghan eq (6.89)
        if HAS_C_FUNCS:
            # divide by 2 to normalize by integral sin(x), x=0..pi
            diffusion_decay = AnisotropicDiffusion._integrate_c(q, t, d_perp, d_par) / 2
        else:
            z = np.sqrt(q**2 * (d_par - d_perp) * t)
            diffusion_decay = np.exp(-q**2 * t * d_perp) * special.erf(z) / z

        return m0 * diffusion_decay * relax

    @staticmethod
    def _integrate_c(q, t, d_perp, d_par) -> np.ndarray:
        diffusion_decay = np.zeros_like(t)

        for (i, t_i) in enumerate(t):
            c = (c_double * 4)(q, t_i, d_perp, d_par)
            user_data = cast(pointer(c), c_void_p)

            diffusion_decay[i] = quad(LowLevelCallable(diffusion_lib.anistropicDiffusion, user_data), 0, np.pi, epsabs=1e-13)[0]

        return diffusion_decay


class Peschier:
    name = 'Diffusion + Cross-Relaxation'
    type = 'Diffusion'
    equation = r'Diffusion with cross-relax f(ast) \rightarrow s(low)'
    params = ['M_{0}', 'D', 'T_{1f}', 'T_{1s}', 'k', 'p_{f}', 't_{ev}', 'g']
    bounds = [(0, None), (0, None), (0, None), (0, None), (0, None), (0, None)]
    choices = [(r'\gamma', 'nucleus', gamma)]

    @staticmethod
    def func(x, m0, d, t1f, t1s, k, pf, tp, g, nucleus: float = 2.67522128e8):
        q = nucleus*g*tp

        r1s, r1f = 1 / t1s, 1 / t1f
        kf, ks = (1-pf)*k, pf*k
        a_plus = 0.5 * (d*q*q + kf + ks + r1f + r1s + np.sqrt((d*q*q + kf + r1f - ks - r1s)**2 + 4*kf*ks))
        a_minu = 0.5 * (d*q*q + kf + ks + r1f + r1s - np.sqrt((d*q*q + kf + r1f - ks - r1s)**2 + 4*kf*ks))

        return m0 * np.exp(-(2/3)*tp*d*q**2)*((a_plus - ks - r1s) / (a_plus - a_minu) * np.exp(-a_plus * x) -
                                              (a_minu - ks - r1s) / (a_plus - a_minu) * np.exp(-a_minu * x))


class DiffusionGradients:
    name = 'Divided Gradients'
    type = 'Diffusion'
    equation = r'M_{0}exp[-\gamma^{2} (g_{1}^{2} - g_{2}^{2}) t_{ev}^{2} D(2/3t_{ev}+t_{mix})]'
    params = ['M_{0}', 'D', 'g_{1}', 'g_{2}', 't_{ev}/t_{mix}']
    choices = [(r'\gamma', 'nucleus', gamma), ('Varied time', 't_axis', {'t_mix': 'tm', 't_ev': 'te'})]
    bounds = [(0, None), (0, None), (0, None), (0, None), (0, None)]

    @staticmethod
    def func(x, m0, d, g1, g2, t2, nucleus=2.67522128e8, t_axis='tm'):
        if t_axis == 'tm':
            tm = x
            tp = t2
        else:
            tm = t2
            tp = x
        q2 = (g1**2 - g2**2) * (nucleus * tp)**2
        t_eff = (2 * tp) / 3 + tm

        return m0 * np.exp(-q2 * d * t_eff)
