from itertools import product

import numpy as np

from .correlationfuncs import GammaCD
from ..math.mittagleffler import mlf


class StimEcho:
    type = 'Stimulated Echo'
    name = 'Stimulated Echo'
    equation = r'M_{0} * [(1-F_{\infty})C(-(x/\tau)^{\alpha}) + F_{\infty}]' \
               r'R[-(x/T_{1})^{\beta}]'
    params = ['M_{0}', r'F_{\infty}', r'\tau', r'\alpha', 'T_{1}', r'\beta']
    bounds = [(0, None), (0, 1), (0, None), (0, 1), (0, None), (0, 1)]
    choices = [('C', 'ctype', {'KWW': 'kww', 'MLF': 'mlf', 'CD': 'cd'}),
               ('R', 'rtype', {'KWW': 'kww', 'MLF': 'mlf', 'CD': 'cd'})]

    @staticmethod
    def func(x, m0, finf, tau, alpha, t1, beta, ctype='kww', rtype='kww'):
        if ctype == 'kww':
            corr_f = np.exp(-(x / tau)**alpha)
        elif ctype == 'mlf':
            corr_f = mlf(-(x/tau)**alpha, alpha)
        elif ctype == 'cd':
            corr_f = GammaCD.func(x, 1., tau, alpha)
        else:
            raise ValueError(f'Unknown ctype {ctype}, use "kww", "mlf" or "cd"')

        if rtype == 'kww':
            rel_f = np.exp(-(x / t1)**beta)
        elif rtype == 'mlf':
            rel_f = mlf(-(x/t1)**beta, beta)
        elif rtype == 'cd':
            rel_f = GammaCD.func(x, 1., t1, beta)
        else:
            raise ValueError(f'Unknown rtype {rtype}, use "kww", "mlf" or "cd"')

        return m0 * ((1 - finf) * corr_f + finf) * rel_f


class StimEchoZ0:
    name = 'Initial magnetization'
    type = 'Stimulated Echo'
    params = ['A', r'\Deltat', r'\delta', r'\eta']
    equation = r'A(1 \pm \langle cos[2\omega_{Q}(x+\Deltat)]\rangle)'
    choices = [('Mode', 'mode', {'Zeeman': 'cos', 'Spin-Alignment': 'sin'})]

    @staticmethod
    def func(x, a, t_ev, delta, eta, mode='cos'):
        x = np.asarray(x)

        theta = np.linspace(0, np.pi/2)
        phi = np.linspace(0, 2*np.pi)
        angles = np.array(list(product(theta, phi)))
        dt = theta[1] - theta[0]
        dp = phi[1] - phi[0]

        omega = np.pi*delta*(3*np.cos(angles[:, 0])**2 - 1 - eta * np.sin(angles[:, 0])**2 * np.cos(2*angles[:, 1]))
        wt = np.sin(angles[:, 0])

        ret_val = (np.cos(2 * np.outer(omega, x-t_ev)) * wt[:, None]).sum(axis=0) * dt * dp / 2 / np.pi

        if mode == 'cos':
            return a*(1 + ret_val) / 2
        elif mode == 'sin':
            return a*(1 - ret_val) / 2
        else:
            raise ValueError
