# Prepare Python environment
import math
import json
import nibabel as nib
import numpy as np
from numpy.fft import ifftn, fftn, ifft, fftshift, ifftshift
import os
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy.signal import butter, lfilter, freqz, filtfilt
from scipy.io import loadmat
import warnings
PI_UNICODE = "\U0001D70B"
CHI_UNICODE = "\U0001D712"
MICRO_UNICODE = "\u00B5"
GYRO_BAR_RATIO_H = 42.6e6  # [Hz/T]

def calc_dk(gx, gy, dt):
    dkx = GYRO_BAR_RATIO_H * gx * dt
    dky = GYRO_BAR_RATIO_H * gy * dt
    return (dkx, dky)

gy_bad = 100e-6

end_time = 0.0912
n_times = 913
t = np.linspace(0, end_time, n_times)
dt = end_time / n_times
nx = 64
k = np.zeros([n_times, 2])
k_distorted = np.zeros([n_times, 2])
for it in range(1, n_times):
    
    if it <= 20:
        gx = -40e-3
        gy = -40e-3
    else:
        n_steps = 138
        i = (it - 20) % n_steps
        if i <= 0:
            gx = 0
            gy = 25e-3
        elif i <= nx:
            gx = 25e-3
            gy = 0
        elif i <= nx + 5:
            gx = 0
            gy = 25e-3
        elif i <= (2*nx) + 5:
            gx = -25e-3
            gy = 0
        elif i <= n_steps:
            gx = 0
            gy = 25e-3

    dkx, dky = calc_dk(gx, gy, dt)
    kx = k[it - 1, 0] + dkx
    ky = k[it - 1, 1] + dky
    k[it, :] = [kx, ky]
    dkx, dky = calc_dk(gx, gy + gy_bad, dt)
    kx = k_distorted[it - 1, 0] + dkx
    ky = k_distorted[it - 1, 1] + dky
    k_distorted[it, :] = [kx, ky]


fig = go.Figure()
fig.add_trace(go.Scatter(x=k[:, 0], y=k[:, 1],
                     mode='lines',
                     line=dict(color='#636EFA'),
                     name='Theoretical trajectory'))
fig.add_trace(go.Scatter(x=k_distorted[:, 0], y=k_distorted[:, 1],
                     mode='lines',
                     line=dict(color='#fa6363'),
                     name='Inhomogeneous trajectory'))
frames = [dict(
    data=[go.Scatter(x=k[:2*i, 0], y=k[:2*i, 1],
                     mode='lines',
                     line=dict(color='#636EFA'),
                     name='Theoretical trajectory'),
          go.Scatter(x=k_distorted[:2*i, 0], y=k_distorted[:2*i, 1],
                     mode='lines',
                     line=dict(color='#fa6363'),
                     name='Inhomogeneous trajectory')],
    name=str(i),
    traces=[0,1]) for i in range(int(n_times/2))]
fig.frames = frames

fig.update_xaxes(range=[-3500, 3500])
fig.update_yaxes(range=[-3700, 3700])
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.update_layout(title="K-space Trajectory",
                  title_x=0.5,
                  height=600,
                  width=700,
                  updatemenus=[dict(
                                type='buttons',
                                buttons=[dict(label='Play',
                                              method='animate',
                                              args=[None, dict(frame=dict(duration=15, redraw=False), transition=dict(duration=15), fromcurrent=True, mode='immediate')]),
                                         dict(label='Pause',
                                              method='animate',
                                              args=[[None], dict(frame=dict(duration=0, redraw=False), mode='immediate')])
                                        ])])
fig.show()