'''GRE simulation.'''

import logging

import numpy as np

from .utils import ernst

def spoiled_gre(T1, T2star, TR, TE, alpha=None, M0=1):
    '''Spoiled, steady state GRE contrast simulation.

    Magnitude image simulations.

    Parameters
    ----------
    T1 : array_like
        longitudinal exponential decay time constant (in sec).
    T2star : array_like
        Effective transverse exponential decay time constant (in sec).
    TR : float
        repetition time (in sec).
    TE : float
        echo time (in sec).
    alpha : array_like, optional
        flip angle (in rad).
    M0 : array_like, optional
        proton density.

    Returns
    -------
    S : array_like
        Simulated magnitude spoiled GRE image.

    Notes
    -----
    alpha=None will use the Ernst angle.

    Assuming a longitudinal steady-state and perfect spoiling. Note
    that dependence is on T2* rather than T2 because SE/STE formation
    is suppressed by spoiling and the signal is generated by gradient
    refocusing of an FID.

    Implements equation [14.8] from [1]_.

    References
    ----------
    .. [1] Notes from Bernstein, M. A., King, K. F., & Zhou, X. J.
           (2004). Handbook of MRI pulse sequences. Elsevier.
    '''

    if alpha is None:
        logging.warning(
            'Ernst angle will computed for every T1 value, this may '
            'or may not be what you wanted')
        alpha = ernst(TR, T1)

    # Make sure we don't divide by zero
    idx1 = np.nonzero(T1)
    idx2 = np.nonzero(T2star)
    E1 = np.zeros(T1.shape)
    E1[idx1] = np.exp(-TR/T1[idx1])
    E2 = np.zeros(T2star.shape)
    E2[idx2] = np.exp(-TE/T2star[idx2])

    # Do the thing:
    return M0*np.sin(alpha)*(1 - E1)*E2/(1 - np.cos(alpha)*E1)
