import numpy
import numpy as np
from matplotlib import pyplot



# parameter for spectrum simulations
lb = 2e3
pulse_length = 2e-6


def dampening(x: np.ndarray, apod: float) -> np.ndarray:
    """
    Calculate additional dampening to account e.g. for field inhomogeneities.

    :param x: Time axis in seconds
    :param apod: Dampening factor in 1/seconds
    :return: Exponential dampening
    """

    return np.exp(-apod * x)


def pulse_attn(freq: np.ndarray, t_pulse: float) -> np.ndarray:
    """
    Calculate attenuation of signal to account for finite pulse lengths.

    See Schmitt-Rohr/Spieß, eq. 2.126 for more information.

    :param freq: Frequency axis in Hz
    :param t_pulse: Assumed pulse length in s
    :return: Attenuation factor.
    """
    # cf. Schmitt-Rohr/Spieß eq. 2.126; omega_1 * t_p = pi/2
    pi_half_squared = np.pi**2 / 4
    omega = 2 * np.pi * freq

    numerator = np.sin(np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 2))
    denominator = np.sqrt(pi_half_squared + omega**2 * t_pulse**2 / 4)

    return np.pi * numerator / denominator / 2


def post_process_spectrum(taus, apod, tpulse):
    reduction_factor = np.zeros((taus.size, 5)) # hard-coded t_echo :(

    for i, tau in enumerate(taus):
        try:
            raw_data = np.loadtxt(f'fid_tau={tau:.6e}.dat')
        except OSError:
            continue

        t = raw_data[:, 0]
        timesignal = raw_data[:, 1:]

        timesignal *= dampening(t, apod)[:, None]
        timesignal[0, :] /= 2

        # FT to spectrum
        freq = np.fft.fftshift(np.fft.fftfreq(t.size, d=1e-6))
        spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
        spec *= pulse_attn(freq, t_pulse=tpulse)[:, None]

        reduction_factor[i, :] = 2*timesignal[0, :]

        plt.plot(freq, spec)
        plt.show()

    plt.semilogx(taus, reduction_factor, '.')
    plt.show()