import numpy as np

from ..math.mittagleffler import mlf


class TOne:
    """
    Generic spin-lattice relaxation function

    .. math::
        y = M_\infty (1 - \\alpha\exp[-(x/T_1)^\\beta])

    Args:
        x (array_like): Input values
        m (float): Equilibrium magnetization
        t1 (float): T1 relaxation time
        beta (float): Stretching parameter
    """
    type = 'Relaxation'
    name = 'T1 relaxation'
    equation = r'M_{\infty}(1-\alpha exp(-(x/T_{1})^{\beta}))'
    params = [r'M_{\infty}', r'\alpha', 'T_{1}', r'\beta']
    bounds = [(0, None), (None, None), (0, None), (0, 1)]

    @staticmethod
    def func(x, m, alpha, t1, beta):
        return m * (1 - alpha * np.exp(-(x/t1)**beta))


class TTwo:
    type = 'Relaxation'
    name = 'T2 decay'
    equation = r'\DeltaM exp(-(x/T_{2})^{\beta}) + M_{0}'
    params = [r'\DeltaM', 'T_{2}', r'\beta', 'M_{0}']
    bounds = [(0, None), (0, None), (0, 1), (None, None)]

    @staticmethod
    def func(x, dm, t2, beta, m0):
        return dm * np.exp(-(x/t2)**beta) + m0


class SatRec:
    """
    Spin-lattice relaxation function with offset

    .. math::
        y = \Delta M (1 - \exp[-(x/T_1)^\\beta]) + M_0

    Args:
        x (array_like): Input values
        dm (float): Equilibrium magnetization
        m0 (float): Initial magnetization
        t1 (float): T1 relaxation time
        beta (float): Stretching parameter



    """
    type = 'Relaxation'
    name = 'Recovery experiment'
    equation = r'\DeltaM(1-exp(-(x/T_{1})^{\beta})) + M_{0}'
    params = [r'\DeltaM', 'T_{1}', r'\beta', 'M_{0}']
    choices = [('Type', 'is_inv', {'Saturation': False, 'Inversion': True})]
    bounds = [(0, None), (0, None), (0, 1), (None, None)]

    @staticmethod
    def func(x, dm, t1, beta, m0, is_inv: bool = True):
        alpha = 1
        if is_inv:
            alpha = 2
        return dm * (1 - alpha * np.exp(-(x/t1)**beta)) + m0


class TwoSatRecAbsolute:
    type = 'Relaxation'
    name = 'Two-step relaxation (abs. int)'
    equation = r'M_{0} + \Sigma \DeltaM_{i}(1-exp(-(x/T_{1,i})^{\beta_{i}}))'
    params = [r'\DeltaM_{1}', 'T_{11}', r'\beta_{1}', r'\DeltaM_{2}', 'T_{12}', r'\beta_{2}', 'M_{0}']
    choices = [('Type', 'is_inv', {'Saturation': False, 'Inversion': True})]
    bounds = [(0, None), (0, None), (0, 1), (0, None), (0, None), (0, 1), (None, None)]

    @staticmethod
    def func(x, dm1, t11, beta1, dm2, t12, beta2, m0, is_inv: bool = False):
        alpha = 1
        if is_inv:
            alpha = 2
        return dm1 * (1 - alpha * np.exp(-(x/t11)**beta1)) + dm2 * (1 - alpha * np.exp(-(x/t12)**beta2)) + m0


class TwoSatRecRelative:
    type = 'Relaxation'
    name = 'Two-step relaxation (rel. int)'
    equation = r'M_{0} + \DeltaM[R(1-exp(-(x/T_{1,1})^{\beta_{1}})) + \n'\
               r'(1-R)(1-exp(-(x/T_{1,2})^{\beta_{2}}))]'
    params = [r'\DeltaM', 'M_{0}', 'T_{11}', r'\beta_{1}', 'T_{12}', r'\beta_{2}', 'R']
    choices = [('Type', 'kind', {'Saturation': 'sat', 'Inversion': 'inv'})]
    bounds = [(0, None), (None, None), (0, None), (0, 1), (0, None), (0, 1), (0, 1)]

    @staticmethod
    def func(x, dm, m0, t11, beta1, t12, beta2, ratio, kind='sat'):
        alpha = 1
        if kind == 'inv':
            alpha = 2
        return dm * (ratio * (1 - alpha * np.exp(-(x/t11)**beta1)) +
                     (1-ratio) * (1 - alpha * np.exp(-(x/t12)**beta2))) + m0


class SatRecMLF:
    type = 'Relaxation'
    name = 'Recovery experiment (MLF)'
    equation = r'\DeltaM(1-E_{\alpha}(-(x/T_{1})^{\alpha})) + M_{0}'
    params = [r'\DeltaM', 'M_{0}', 'T_{1}', r'\alpha']
    choices = [('Type','kind', {'Saturation': 'sat', 'Inversion': 'inv'})]
    bounds = [(0, None), (None, None), (0, None), (0, 1)]

    @staticmethod
    def func(x, dm, m0, t1, alpha, kind='sat'):
        a = 1
        if kind == 'inv':
            a = 2
        return dm * (1 - a * mlf(-(x/t1)**alpha, alpha)) + m0
