# Prepare Python environment

import scipy.io as sio
from pathlib import Path

data_dir = Path("../../../data/05-B0/data")

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]
# Note: Field was reduced a lot to be able to show the sinusoid
# Note: *2 after lowpass filter is because this is a single coil (sin instead of e^(-ix)) and demodulating by multiplying a sinusoid creates a 1/2 difference. In practice, since we have both x and y components, we can recover the full signal instead of doing X2.
GYRO_BAR_RATIO_H = 42.6e6  # [Hz/T]
b0 = 0.000002  # [T]
T2 = 0.3  # s
y_0_cst = 100
fs = 10000

f_larmor = b0 * GYRO_BAR_RATIO_H
t = np.linspace(0, 1, fs + 1)  # 1 second

def butter_lowpass(cutoff, fs, order=5):
    return butter(order, cutoff, fs=fs, btype='low', analog=False)

def butter_lowpass_filter(data, cutoff, fs, order=5):
    b, a = butter_lowpass(cutoff, fs, order=order)
    y = filtfilt(b, a, data, method='gust')
    return y

# Lab frame
y_0 = y_0_cst * np.sin(2 * math.pi * f_larmor * t)
exp = np.exp(-t/T2)
y = y_0 * exp / y_0_cst
temp = y * (y_0 / y_0_cst)
y_demod = butter_lowpass_filter(temp, f_larmor, fs, order=5) * 2

fig = go.Figure()
fig.add_scatter(x=t, y=y, name="FID")
fig.add_scatter(x=t, y=y_demod, name="Demodulated FID")
fig.add_scatter(x=t, y=exp, name="T2 decay")

# 2 isochromats
y_1amp = y_0_cst / 10
y_1 = y_1amp * np.sin(2 * math.pi * (f_larmor + 10) * t)
y = (y_0 + y_1) * exp / (y_0_cst + y_1amp)
temp = y * (y_0 / y_0_cst)
y_demod = butter_lowpass_filter(temp, f_larmor, fs, order=5) * 2

fig.add_scatter(x=t, y=y, name="FID", visible=False)
fig.add_scatter(x=t, y=y_demod, name="Demodulated FID", visible=False)
fig.add_scatter(x=t, y=exp, name="T2 decay", visible=False)

# Multiple isochromats
n_freqs = 100
fid = y_0 * exp
y_sum_demod = butter_lowpass_filter(fid * (y_0 / y_0_cst), f_larmor, fs, order=5) * 2
for i in range(n_freqs):
    amp = 1
    freq_offset = 10
    scale = freq_offset/n_freqs
    mid = n_freqs // 2
    y_1 = amp * np.sin(2 * math.pi * (f_larmor + scale*(mid - i)) * t) * exp
    fid += y_1
    y_demod = butter_lowpass_filter(y_1 * (y_0 / y_0_cst), f_larmor, fs, order=5) * 2
    y_sum_demod += y_demod

y_demod_scaled = y_sum_demod / (y_0_cst + (n_freqs * amp))
fid_scaled = fid / (y_0_cst + (n_freqs * amp))

fig.add_scatter(x=t, y=fid_scaled, name="FID", visible=False)
fig.add_scatter(x=t, y=y_demod_scaled, name="Demodulated FID", visible=False)
fig.add_scatter(x=t, y=exp, name="T2 decay", visible=False)
fig.update_traces(marker=dict(size=3))
fig.update_layout(
    title="Single species",
    title_x=0.5,
    updatemenus=[
        dict(
            buttons=list([
                dict(
                    args=[{"visible": [True, True, True, False, False, False, False, False, False]},
                          {"title": "Single species"}],
                    label="Single species",
                    method="update"
                ),
                dict(
                    args=[{"visible": [False, False, False, True, True, True, False, False, False]},
                          {"title": "Two species"}],
                    label="Two species",
                    method="update"
                ),
                dict(
                    args=[{"visible": [False, False, False, False, False, False, True, True, True]},
                          {"title": "Multiple Species"}],
                    label="Multiple Species",
                    method="update"
                )
            ]),
            direction="down",
            showactive=True,

        ),
    ])
fig.show()