from __future__ import annotations

from time import perf_counter

import numpy as np
from numpy.random import Generator
from datetime import datetime
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

from .functions import ste
from .parameter import Parameter
from .distributions import BaseDistribution
from .motions import BaseMotion
from .parsing import parse


def run_ste_sim(config_file: str):
    p = parse(config_file)

    rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)

    t_mix = p.ste.t_mix
    t_evo = p.ste.t_evo
    t_echo = p.ste.t_echo

    fig, ax = plt.subplots(2)
    fig2, ax2 = plt.subplots(2)
    fig3, ax3 = plt.subplots(2)

    # outer loop over variables of distribution of correlation times
    for (i, dist_values) in enumerate(p.dist):
        # noinspection PyCallingNonCallable
        dist = p.dist.dist_type(**dist_values, rng=rng)

        # second loop over parameter of motion model
        for (j, motion_values) in enumerate(p.motion):
            # noinspection PyCallingNonCallable
            motion = p.motion.model(delta, eta, **motion_values, rng=rng)

            print(f'\nStart of {dist}, {motion}')

            start = last_print = perf_counter()

            cc = np.zeros((len(t_mix), len(t_evo)))
            ss = np.zeros((len(t_mix), len(t_evo)))

            # inner loop to create trajectories
            for n in range(num_traj):
                phase = make_trajectory(dist, motion, t_max)

                for (k, t_evo_k) in enumerate(t_evo):
                    dephased = phase(t_evo_k)
                    t0 = t_mix + t_evo_k
                    rephased = phase(t0 + t_evo_k + 2*t_echo) + phase(t0) - 2 * phase(t0+t_echo)
                    # print(t_evo_k, t0 + t_evo_k + 2*t_echo, t0)
                    cc[:, k] += np.cos(dephased)*np.cos(rephased)
                    ss[:, k] += np.sin(dephased)*np.sin(rephased)

                last_print = print_step(n, num_traj, start, last_print)

            cc[:, 1:] /= num_traj
            ss[:, 1:] /= num_traj

            fig4, ax4 = plt.subplots()
            ax4.semilogx(t_mix, cc/cc[0, :], '.-')
            fig5, ax5 = plt.subplots()
            ax5.semilogx(t_mix, ss/ss[0, :], '.-')

            for k in range(num_variables):
                p0 = [0.5, 0, dist_values.get('tau', 1), 1]

                p_final = []
                p_final1 = []
                for k, t_evo_k in enumerate(p.ste.t_evo):

                    try:
                        res = curve_fit(ste, t_mix, cc[:, k], p0=p0, bounds=([0, 0, 0, 0], [np.inf, 1, np.inf, 1]))
                        p_final.append(res[0].tolist())
                    except RuntimeError:
                        p_final.append([np.nan, np.nan, np.nan, np.nan])

                    try:
                        res2 = curve_fit(ste, t_mix, ss[:, k], p0=p0, bounds=([0, 0, 0, 0], [np.inf, 1, np.inf, 1]))
                        p_final1.append(res2[0].tolist())
                    except RuntimeError:
                        p_final1.append([np.nan, np.nan, np.nan, np.nan])

                p_final = np.array(p_final)
                p_final1 = np.array(p_final1)
                # ax[0].semilogy(p.ste.t_evo, p_final[:, 0], '.--')
                # ax[1].semilogy(t_evo, p_final1[:, 0], '.--')
                ax[0].plot(t_evo, p_final[:, 1], '.-')
                ax[1].plot(t_evo, p_final1[:, 1], '.-')
                ax2[0].semilogy(t_evo, p_final[:, 2], '.-')
                ax2[1].semilogy(t_evo, p_final1[:, 2], '.-')
                ax3[0].plot(t_evo, p_final[:, 3], '.-')
                ax3[1].plot(t_evo, p_final1[:, 3], '.-')

    plt.show()


def run_spectrum_sim(config_file: str):
    p = parse(config_file)

    rng, num_traj, t_max, delta, eta, num_variables = _prepare_sim(p)

    num_echos = len(p.spec.t_echo)
    reduction_factor = np.zeros((num_variables, num_echos))
    t_echo = p.spec.t_echo
    t_echo_strings = list(map(str, t_echo))

    # outer loop over variables of distribution of correlation times
    for (i, dist_values) in enumerate(p.dist):
        # noinspection PyCallingNonCallable
        dist = p.dist.dist_type(**dist_values, rng=rng)
        # second loop over parameter of motion model
        for (j, motion_values) in enumerate(p.motion):
            # noinspection PyCallingNonCallable
            motion = p.motion.model(delta, eta, **motion_values, rng=rng)
            print(f'\nStart of {dist}, {motion}')

            timesignal = np.zeros((p.spec.num_points, num_echos))

            start = perf_counter()
            last_print = start

            # inner loop to create trajectories
            for n in range(num_traj):
                phase = make_trajectory(dist, motion, t_max)

                for (k, t_echo_k) in enumerate(t_echo):
                    # effect of de-phasing and re-phasing
                    start_amp = -2 * phase(t_echo_k)

                    # start of actual acquisition
                    timesignal[:, k] += np.cos(start_amp + phase(p.spec.t_acq + 2*t_echo_k))
                    reduction_factor[max(p.motion.num_variables, 1)*i+j, k] += np.cos(phase(2*t_echo_k) + start_amp)

                # print(n+1, num_traj, start, last_print)
                last_print = print_step(n+1, num_traj, start, last_print)

            # apply line broadening
            timesignal *= p.spec.dampening[:, None]
            timesignal /= num_traj
            timesignal[0, :] /= 2

            # FT to spectrum
            spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
            spec -= spec[0]
            spec *= p.spec.pulse_attn[:, None]

            save_spectrum_data(timesignal, spec, p, dist, motion, t_echo_strings)

    fig2, ax2 = plt.subplots()
    lines = ax2.semilogx(p.dist.variables['tau'], reduction_factor / num_traj, 'o--')
    ax2.legend(lines, t_echo_strings)
    plt.savefig(f'{dist.name}_{motion.name}_reduction.png')

    plt.show()


def make_trajectory(
        dist: BaseDistribution,
        motion: BaseMotion,
        t_max: float,
        t_passed: float = 0.,
        init_phase: float = 0.
):

    # set initial orientations and correlation times
    motion.start()
    dist.start()

    # number of jumps that are simulated at once
    chunks = min(int(0.51 * t_max / dist.tau_jump), 100_000) + 1
    # print(chunks)

    t = [np.array([t_passed])]
    phase = [np.array([init_phase])]
    # omega = [np.array([0])]
    while t_passed < t_max:
        # frequencies between jumps
        current_omega = motion.jump(size=chunks)
        # times at a particular position
        t_wait = dist.wait(size=chunks)

        accumulated_phase = np.cumsum(t_wait * current_omega) + phase[-1][-1]
        phase.append(accumulated_phase)
        # omega.append(current_omega)

        t_wait = np.cumsum(t_wait) + t_passed
        t_passed = t_wait[-1]
        t.append(t_wait)

    t = np.concatenate(t)
    phase = np.concatenate(phase)
    # omega = np.concatenate(omega)

    # fig_test, ax_test = plt.subplots()
    # ax_test.plot(t, phase, 'x-')

    # np.savetxt('trajectory.dat', np.c_[t, phase, omega])

    # convenient interpolation to get phase at arbitrary times
    phase_interpol = interp1d(t, phase)

    return phase_interpol


def _prepare_sim(parameter: Parameter) -> tuple[Generator, int, float, float, float, int]:
    # collect variables that are common to spectra and stimulated echo

    # random number generator
    rng = np.random.default_rng(parameter.sim.seed)

    # number of random walkers
    num_traj = parameter.sim.num_walker

    # length of trajectories
    t_max = parameter.sim.t_max

    # parameter for omega_q
    delta, eta = parameter.molecule.delta, parameter.molecule.eta

    num_variables = parameter.dist.num_variables * parameter.motion.num_variables

    return rng, num_traj, t_max, delta, eta, num_variables


def print_step(n: int, num_traj: int, start: float, last_print: float) -> float:
    step_time = perf_counter()
    dt = step_time - last_print
    if dt > 10 or n == num_traj:
        date = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f'{date} - step {n} of {num_traj}', end=' - ')

        elapsed = step_time - start
        total = num_traj * elapsed / n
        print(f'expected total: {total:.2f}s - elapsed: {elapsed:.2f}s - remaining: {total - elapsed:.2f}s')
        if dt > 10:
            last_print = step_time

    return last_print


def make_filename(dist: BaseDistribution, motion: BaseMotion) -> str:
    filename = f'{dist}_{motion}'
    filename = filename.replace(' ', '_')
    filename = filename.replace('.', 'p')

    return filename


def save_spectrum_data(
        timesignal: np.ndarray,
        spectrum: np.ndarray,
        param: Parameter,
        dist: BaseDistribution,
        motion: BaseMotion,
        echo_strings: list[str]
):
    filename = make_filename(dist, motion)

    header = param.totext(sim=True, spec=True)
    header += '\nx\t' + '\t'.join(echo_strings)

    np.savetxt(filename + '_timesignal.dat', np.c_[param.spec.t_acq, timesignal], header=header)
    np.savetxt(filename + '_spectrum.dat', np.c_[param.spec.freq, spectrum], header=header)

    fig, ax = plt.subplots()
    lines = ax.plot(param.spec.freq, spectrum)
    ax.set_title(f'{dist}, {motion}')
    ax.legend(lines, echo_strings)
    plt.savefig(filename + '_spectrum.png')

    fig1, ax1 = plt.subplots()
    lines = ax1.plot(param.spec.t_acq, timesignal)
    ax1.set_title(f'{dist}, {motion}')
    ax1.legend(lines, echo_strings)
    plt.savefig(filename + '_timesignal.png')

    plt.show()