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


def spgr(pd, r1, r2s=None, mt=None, transmit=None, receive=None, gfactor=None,
         te=0, tr=25e-3, fa=20, sigma=None, device=None):
    """Simulate data generated by a Spoiled Gradient-Echo (SPGR/FLASH) sequence.

    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. Mandatory if any `te > 0`.
    mt : tensor_like, optional
        MTsat. Mandatory if any `mtpulse == True`.

    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
    -------------------
    te : float, default=0
        Echo time, in sec
    tr : float default=2.5e-3
        Repetition time, in sec
    fa : float, default=20
        Flip angle, in deg

    Noise
    -----
    sigma : float, optional
        Standard-deviation of the sampled Rician noise (no sampling if `None`)
    Returns
    -------
    sim : tensor
        Simulated SPGR image

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

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

    if receive is not None:
        pd = pd * receive
    del receive
    pd = pd * fa.sin()
    fa = fa.cos()

    e1, r1 = r1.mul(tr).neg_().exp(), None
    signal = pd * (1 - e1)

    if mt is not None:
        omt = mt.neg().add_(1)
        signal *= omt
        signal /= (1 - fa * omt * e1)
        del omt
    else:
        signal /= (1 - fa * e1)

    if r2s is not None:
        e2, r2s = r2s.mul(te).neg_().exp(), None
        signal *= e2
        del e2

    # noise
    signal = add_noise(signal, std=sigma)
    return signal



def stack_maps(*prm):
    prm = [p for p in prm if p is not None]
    return torch.stack(prm)


def unstack_maps(prm, has_r2s, has_mt):
    pd = prm[0]
    r1 = prm[1]
    r2s = prm[2] if has_r2s else None
    mt = prm[-1] if has_mt else None
    return pd, r1, r2s, mt


def exp_maps(pd, r1, r2s, mt, inplace=False):
    if not inplace:
        pd = pd.clone()
        r1 = r1.clone()
        if r2s is not None:
            r2s = r2s.clone()
        if mt is not None:
            mt = mt.clone()
    pd = pd.exp_()
    r1 = r1.exp_()
    if r2s is not None:
        r2s = r2s.exp_()
    if mt is not None:
        mt = mt.neg_().exp_().add_(1).reciprocal_()
    return pd, r1, r2s, mt


def exp_maps_(pd, r1, r2s, mt):
    return exp_maps(pd, r1, r2s, mt, inplace=True)


def safelog(x, eps=1e-42):
    return x.clamp_min(eps).log_()
