import torch
from nitorch.core import utils, constants
from .noise import add_noise


def mprage(pd, r1, r2s=None, transmit=None, receive=None, gfactor=None,
           tr=2.3, ti=0.9, tx=None, te=None, fa=9, n=160, eff=0.96,
           sigma=None, device=None):
    """Simulate data generated by a (simplified) MP-RAGE sequence.

    Default parameters mimic the ADNI-3 protocol on 3T Siemens scanners.
    Our Implementation is based on the MP2RAGE paper, where the sequence
    is stripped from the second GRE readout block.

    Tissue parameters
    -----------------
    pd : tensor_like
        Proton density
    r1 : tensor_like
        Longitudinal relaxation rate, in 1/sec
    r2s : tensor_like, optional
        Transverse relaxation rate, in 1/sec.
        If not provided, T2*-bias is not included.

    Fields
    ------
    transmit : tensor_like, optional
        Transmit B1 field
    receive : tensor_like, optional
        Receive B1 field
    gfactor : tensor_like, optional
        G-factor map.
        If provided and `sigma` is not `None`, the g-factor map is used
        to sample non-stationary noise.

    Sequence parameters
    -------------------
    tr : float default=2.3
        Repetition time, in sec.
        (Time between two inversion pulses)
    ti : float, default=0.9
        Inversion time, in sec.
        (Time between inversion pulse and middle of the echo train)
    tx : float, default=2*te or 6e-3
        Excitation repetition time, in sec
        (Time between two excitation pulses within the echo train)
    te : float, default=tx/2
        Echo time, in sec
    fa : float, default=9
        Flip angle, in deg
    n : int, default=160
        Number of excitation pulses (= phase encoding steps) per train.
    eff : float, default=0.96
        Efficiency of the inversion pulse.

    Noise
    -----
    sigma : float, optional
        Standard-deviation of the sampled Rician noise (no sampling if `None`)

    Returns
    -------
    sim : tensor
        Simulated MPRAGE image

    References
    ----------
    ..[1] "MP2RAGE, a self bias-field corrected sequence for improved
        segmentation and T1-mapping at high field."
        Marques JP, Kober T, Krueger G, van der Zwaag W, Van de Moortele PF, Gruetter R.
        Neuroimage. 2010 Jan 15;49(2):1271-81.
        doi: 10.1016/j.neuroimage.2009.10.002

    """

    pd, r1, r2s, transmit, receive, gfactor \
        = utils.to_max_backend(pd, r1, r2s, transmit, receive, gfactor)
    pd, r1, r2s, transmit, receive, gfactor \
        = utils.to(pd, r1, r2s, transmit, receive, gfactor, device=device)
    backend = utils.backend(pd)

    if tx is None and te is None:
        tx = 6e-3
    tx = tx or 2*te                 # Time between excitation pulses
    te = te or tx/2                 # Echo time
    fa = fa * constants.pi / 180    # Flip angle of GRE block
    n = n or min(pd.shape)          # Number of readouts (PE steps) per loop
    tr1 = n*tx                      # GRE block
    tp = ti - tr1/2                 # Preparation time
    td = tr - ti - tr1/2            # Recovery time
    m = n // 2                      # Middle of echo train

    if transmit is not None:
        fa = transmit * fa
    del transmit
    fa = torch.as_tensor(fa, **backend)

    # precompute exponential terms
    ex = r1.mul(-tx).exp()
    ep = r1.mul(-tp).exp()
    ed = r1.mul(-td).exp()
    e1 = r1.mul(-tr).exp()
    c = fa.cos()

    # steady state
    s = (1 - ep) * (c*ex).pow(n)
    s = s + (1 - ex) * (1 - (c*ex).pow(n)) / (1 - c*ex)
    s = s * ed + (1 - ed)
    s = s * pd / (1 + eff * c.pow(n) * e1)

    # IR component
    s = -eff * s * ep / pd + (1 - ep)
    s = s * (c*ex).pow(m-1)
    s = s + (1 - ex) * (1 - (c*ex).pow(m-1)) / (1 - c*ex)
    s = s * fa.sin()
    s = s.abs()

    # Modulation (PD, B1-, R2*)
    if receive is not None:
        pd = pd * receive
    del receive

    s = s * pd

    if r2s is not None:
        e2 = r2s.mul(-te).exp_()
        s = s * e2
    del r2s

    # noise
    s = add_noise(s, std=sigma, gfactor=gfactor)
    return s
