from time import time
from numpy.typing import ArrayLike
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt


# spectral parameter
delta = 161e3  # in Hz
eta = 0
lb = 5e3  # in Hz

# correlation time
tau = np.logspace(-8, -1, num=15)  # in s

# acquisition parameter
acq_length = 4096
dt = 1e-6  # in s
t_echo = [5e-6, 10e-6, 20e-6, 50e-6, 100e-6, 200e-6]  # all in s

# derived parameter
t_acq = np.arange(acq_length) * dt
t_max = acq_length*dt + 2*max(t_echo)
dampening = np.exp(-lb * t_acq)

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

# number of random walkers
num_traj = 1


def omega_q(delta_: float, eta_: float, theta_: ArrayLike, phi_: ArrayLike) -> np.ndarray:
    cos_theta = np.cos(theta_)
    sin_theta = np.sin(theta_)
    return 2 * np.pi * delta_ * (3 * cos_theta * cos_theta - 1 + eta_ * sin_theta*sin_theta * np.cos(2*phi_))


def rotate(x_in: float, y_in: float, z_in: float, a: float) -> tuple[float, float]:
    # rotation by tetrahedral angle is a given, only second angle is free parameter
    beta = 109.45 * np.pi / 180.
    cos_beta = np.cos(beta)
    sin_beta = np.sin(beta)

    cos_alpha = np.cos(a)
    sin_alpha = np.sin(a)

    scale = np.sqrt(1 - z_in * z_in) + 1e-12

    x = x_in * cos_beta + sin_beta / scale * (x_in * z_in * cos_alpha - y_in * sin_alpha)
    y = y_in * cos_beta + sin_beta / scale * (y_in * z_in * cos_alpha + x_in * sin_alpha)
    z = z_in * cos_beta - scale * sin_beta*cos_alpha
    z = max(-1, min(1, z))

    return np.arccos(z), np.arctan2(y, x)


def new_tau(size=1) -> np.ndarray:
    return rng.exponential(tau_i, size=size)


reduction_factor = np.zeros((len(tau), len(t_echo)))


for (n, tau_i) in enumerate(tau):
    print(f'\nStart for tau={tau_i}')

    timesignal = np.zeros((acq_length, len(t_echo)))

    start = time()
    expected_jumps = round(t_max/tau_i)
    if expected_jumps > 1e7:
        print(f'Too many jumps to process, Skip {tau_i}s')
        continue

    chunks = int(0.6 * t_max / tau_i) + 1
    print(f'Chunk size for trajectories: {chunks}')

    for i in range(num_traj):
        # draw orientation
        z_theta, z_phi, z_alpha = rng.random(3)
        theta = np.arccos(1 - 2 * z_theta)
        phi = 2 * np.pi * z_phi
        alpha = 2 * np.pi * z_alpha

        # orientation in cartesian coordinates
        x_start = np.sin(theta) * np.cos(phi)
        y_start = np.sin(theta) * np.sin(phi)
        z_start = np.cos(theta)

        # calculate orientation of tetrahedral edges and their frequencies
        orientation = np.zeros(4)
        orientation[0] = omega_q(delta, eta, theta, phi)
        orientation[1] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha))
        orientation[2] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha + 2 * np.pi / 3))
        orientation[3] = omega_q(delta, eta, *rotate(x_start, y_start, z_start, alpha + 4 * np.pi / 3))

        t_passed = 0
        t = [0]
        phase = [0]
        accumulated_phase = 0
        start_position = rng.choice([0, 1, 2, 3])

        while t_passed < t_max:
            # orientation until the next jump
            jumps = rng.choice([1, 2, 3], size=chunks) + start_position
            jumps = np.cumsum(jumps)
            jumps %= 4

            current_omega = orientation[jumps]
            # current_omega = rng.choice(orientation, size=chunks)

            # time to next jump
            t_wait = new_tau(size=chunks)
            accumulated_phase = np.cumsum(t_wait*current_omega) + phase[-1]

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

            phase.extend(accumulated_phase.tolist())

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

        for j, t_echo_j in enumerate(t_echo):
            # effect of de-phasing and re-phasing
            start_amp = -2 * phase_interpol(t_echo_j)

            # start of actual acquisition
            timesignal[:, j] += np.cos(start_amp + phase_interpol(t_acq + 2*t_echo_j)) * dampening
            reduction_factor[n, j] += np.cos(phase_interpol(2*t_echo_j) + start_amp)

        if (i+1) % 200 == 0:
            elapsed = time()-start
            print(f'Step {i+1} of {num_traj}', end=' - ')
            total = num_traj * elapsed / (i+1)
            print(f'elapsed: {elapsed:.2f}s - total: {total:.2f}s - remaining: {total-elapsed:.2f}s')

    timesignal /= num_traj

    # FT to spectrum
    freq = np.fft.fftshift(np.fft.fftfreq(acq_length, dt))
    spec = np.fft.fftshift(np.fft.fft(timesignal, axis=0), axes=0).real
    spec -= spec[0]
    # spec /= np.max(spec, axis=0)

    t_echo_strings = list(map(str, t_echo))

    # plot spectra
    fig, ax = plt.subplots()
    lines = ax.plot(freq, spec)
    ax.set_title(f'tau = {tau_i}s')
    ax.legend(lines, t_echo_strings)
    # plt.savefig(f'{tau_i}.png')

    # save time signals and spectra
    # np.savetxt(f'spectrum_{tau_i}.dat', np.c_[freq, spec], header='f\t' + '\t'.join(t_echo_strings))
    # np.savetxt(f'timesignal_{tau_i}.dat', np.c_[t_acq, timesignal], header='t\t' + '\t'.join(t_echo_strings))

fig2, ax2 = plt.subplots()
ax2.semilogx(tau, reduction_factor / num_traj, 'o--')

plt.show()