import pathlib
import subprocess

import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt



def run_sims(taus, ste: bool = True, spectrum: bool = False, exec_file: str = './build/rwsim', config_file: str = './config.txt'):
    for tau in taus:
        arguments = [exec_file, config_file]
        if ste:
            arguments += ['--ste']
        if spectrum:
            arguments += ['--spectrum']

        arguments += ['IsotropicAngle']

        with pathlib.Path(config_file).open('a') as f:
            f.write(f'tau={tau}\n')

        subprocess.run(arguments)



def dampening(x: np.ndarray, apod: float) -> np.ndarray:
    return np.exp(-apod * x)


def pulse_attn(freq: np.ndarray, t_pulse: float):
    # 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()


def post_process_ste(taus):
    tevo = np.linspace(1e-6, 60e-6, num=121)

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

        t_mix = raw_data_cc[:, 0]

        fig_cc_raw, ax_cc_raw = plt.subplots()
        fig_ss_raw, ax_ss_raw = plt.subplots()

        ax_cc_raw.semilogx(t_mix, raw_data_cc[:, 1:])
        ax_ss_raw.semilogx(t_mix, raw_data_ss[:, 1:])

        scaled_cc = (raw_data_cc[:, 1:]-raw_data_cc[-1, 1:])/(raw_data_cc[0, 1:]-raw_data_cc[-1, 1:])
        scaled_ss = (raw_data_ss[:, 1:]-raw_data_ss[-1, 1:])/(raw_data_ss[0, 1:]-raw_data_ss[-1, 1:])

        fig_tau, ax_tau = plt.subplots()
        fig_beta, ax_beta= plt.subplots()
        fig_finfty, ax_finfty = plt.subplots()

        tau_cc_fit = []
        beta_cc_fit = []
        finfty_cc_fit = []

        tau_ss_fit = []
        beta_ss_fit = []
        finfty_ss_fit = []

        for j in range(1, raw_data_cc.shape[1]-1):
            p0 = [
                raw_data_cc[0, 1],
                t_mix[np.argmin(np.abs(scaled_cc[:, j]-np.exp(-1)))],
                1,
                0.1
            ]

            res = curve_fit(ste, t_mix, raw_data_cc[:, j+1], p0, bounds=[(0, 0, 0, 0), (np.inf, np.inf, 1, 1)])
            m0, tauc, beta, finfty = res[0]
            # print(f'Cos-Cos-Fit for {tevo[j]}: tau_c = {tauc:.6e}, beta={beta:.4e}, amplitude = {m0: .4e}, f_infty={finfty:.4e}')

            tau_cc_fit.append(res[0][1])
            beta_cc_fit.append(res[0][2])
            finfty_cc_fit.append(res[0][3])

            p0 = [
                raw_data_cc[0, 1],
                t_mix[np.argmin(np.abs(scaled_ss[:, j]-np.exp(-1)))],
                1,
                0.1
            ]

            res = curve_fit(ste, t_mix, raw_data_ss[:, j+1], p0, bounds=[(0, 0, 0, 0), (np.inf, np.inf, 1, 1)])
            m0, tauc, beta, finfty = res[0]
            # print(f'Sin-Sin-Fit for {tevo[j]}: tau_c = {tauc:.6e}, beta={beta:.4e}, amplitude = {m0: .4e}, f_infty={finfty:.4e}')

            tau_ss_fit.append(res[0][1])
            beta_ss_fit.append(res[0][2])
            finfty_ss_fit.append(res[0][3])

        ax_tau.semilogy(tevo[1:], tau_cc_fit, 'C0-')
        ax_tau.axhline(tau)
        ax_beta.plot(tevo[1:], beta_cc_fit, 'C0-')
        ax_finfty.plot(tevo[1:], finfty_cc_fit, 'C0-')

        ax_tau.semilogy(tevo[1:], tau_ss_fit, 'C1-')
        ax_tau.axhline(tau)
        ax_beta.plot(tevo[1:], beta_ss_fit, 'C1-')
        ax_finfty.plot(tevo[1:], finfty_ss_fit, 'C1-')

        plt.show()

        np.savetxt('cc_tauc.dat', list(zip(tevo[1:], tau_cc_fit)))
        np.savetxt('cc_beta.dat', list(zip(tevo[1:], beta_cc_fit)))
        np.savetxt('cc_finfty.dat', list(zip(tevo[1:], finfty_cc_fit)))
        np.savetxt('ss_tauc.dat', list(zip(tevo[1:], tau_ss_fit)))
        np.savetxt('ss_beta.dat', list(zip(tevo[1:], beta_ss_fit)))
        np.savetxt('ss_finfty.dat', list(zip(tevo[1:], finfty_ss_fit)))



    # for i, tau in enumerate(taus):
    #
    #     try:
    #         raw_data_cc = np.loadtxt(f'coscos_tau={tau:.6e}.dat')
    #         raw_data_ss = np.loadtxt(f'sinsin_tau={tau:.6e}.dat')
    #     except OSError:
    #         continue
    #
    #     t_mix = raw_data_cc[:, 0]
    #
    #     plt.semilogx(t_mix, raw_data_cc[:, 1:])
    #     plt.show()
    #
    #     plt.semilogx(t_mix, raw_data_ss[:, 1:])
    #     plt.show()
    #
    #     plt.plot(raw_data_cc[0, 1:])
    #     plt.plot(raw_data_ss[0, 1:])
    #     plt.show()
    #
    #     plt.plot(raw_data_cc[-1, 1:]/raw_data_cc[0, 1:])
    #     plt.plot(raw_data_ss[-1, 1:]/raw_data_ss[0, 1:])
    #     plt.show()


def ste(x, m0, t, beta, finfty):
    return m0 * ((1-finfty) * np.exp(-(x/t)**beta) + finfty)


if __name__ == '__main__':
    tau_values = np.geomspace(1e-2, 1e-6, num=1)   # if num=1, tau is first value

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

    run_sims(tau_values)
    post_process_ste(tau_values)
    # post_process_spectrum(tau_values, lb, pulse_length)