# Prepare Python environment

import scipy.io as sio
from pathlib import Path

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

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]

fname_mag_e1 = data_dir / "sub-fmap_magnitude1.nii.gz"
fname_mag_e1_json = data_dir / "sub-fmap_magnitude1.json"
fname_phase_e1 = data_dir / "sub-fmap_phase1.nii.gz"
fname_phase_e2 = data_dir / "sub-fmap_phase2.nii.gz"
fname_mask = data_dir / "mask.nii.gz"
fname_fmap = data_dir / "fmap.nii.gz"
mag_e1 = nib.load(fname_mag_e1).get_fdata()[30:-30,8:108,30]
phase_e1 = nib.load(fname_phase_e1).get_fdata()[30:-30,8:108,30]
phase_e2 = nib.load(fname_phase_e2).get_fdata()[30:-30,8:108,30]
fmap_hz = nib.load(fname_fmap).get_fdata()[30:-30,8:108,30]
# mag_e1 = nib.load(fname_mag_e1).get_fdata()
# phase_e1 = nib.load(fname_phase_e1).get_fdata()
# phase_e2 = nib.load(fname_phase_e2).get_fdata()
# mask = nib.load(fname_mask).get_fdata()
# fmap_hz = nib.load(fname_fmap).get_fdata()

with open(fname_mag_e1_json, 'r') as json_data:
    data = json.load(json_data)

freq = data["ImagingFrequency"]
fmap_t = fmap_hz / GYRO_BAR_RATIO_H * 1e6
fmap_ppm = fmap_hz / freq

fig = go.Figure()
fig.add_trace(go.Heatmap(z=np.rot90(mag_e1, k=-1),
                         colorscale='gray',
                         colorbar=dict(
                            title="a.u.",
                            titleside="top",
                            tickmode="array"
                         ))
             )
fig.add_trace(go.Heatmap(z=np.rot90(phase_e2 / 4095 * 2 * math.pi - math.pi, k=-1),
                         colorscale='gray',
                         colorbar=dict(
                            title="Rad",
                            titleside="top",
                            tickmode="array",
                            tickvals=[-math.pi, 0, math.pi-0.01],
                            ticktext = [f"-{PI_UNICODE}", 0, f'{PI_UNICODE}']
                        ),
                         visible=False))
fig.add_trace(go.Heatmap(z=np.rot90(fmap_hz, k=-1),
                         colorscale='gray',
                         colorbar=dict(
                            title="Hz",
                            titleside="top",
                            tickmode="array"
                        ),
                         visible=False))
fig.add_trace(go.Heatmap(z=np.rot90(fmap_t, k=-1),
                         colorscale='gray',
                         colorbar=dict(
                            title=f"{MICRO_UNICODE}T",
                            titleside="top",
                            tickmode="array"
                        ),
                         visible=False))
fig.add_trace(go.Heatmap(z=np.rot90(fmap_ppm, k=-1),
                         colorscale='gray',
                         colorbar=dict(
                            title="ppm",
                            titleside="top",
                            tickmode="array"
                        ),
                         visible=False))



x0=0
y0=89
x1=10
y1=99
h = 2
rounded_bottom_left = f' M {x0+h}, {y0} Q {x0}, {y0} {x0}, {y0+h}'#
rounded_top_left = f' L {x0}, {y1-h} Q {x0}, {y1} {x0+h}, {y1}'
rounded_top_right = f' L {x1-h}, {y1} Q {x1}, {y1} {x1}, {y1-h}'
rounded_bottom_right = f' L {x1}, {y0+h} Q {x1}, {y0} {x1-h}, {y0}Z'
path = rounded_bottom_left + rounded_top_left+\
         rounded_top_right+rounded_bottom_right

annotations = ['A', 'B', 'C', 'D', 'E']
shapes = []
for i_shape, annotation in enumerate(annotations):
    shapes.append(dict(type='path',
                       path=path,
                       fillcolor='white',
                       layer='above',
                       line=dict(width=1),
                       label=dict(text=f"<b>{annotation}</b>")
                      )
                 )

fig.add_shape(shapes[0])
# Add dropdown
fig.update_layout(
    title_text="Magnitude",
    title_x=0.5,
    height=500,
    width=600,
    updatemenus=[
        dict(
            buttons=list([
                dict(
                    method="update",
                    args=[{"visible": [True, False, False, False, False]},
                          {'shapes': [shapes[0]], "title": "Magnitude"}],
                    label="Magnitude",
                ),
                dict(
                    method="update",
                    args=[{"visible": [False, True, False, False, False]},
                          {'shapes': [shapes[1]], "title": "Phase"}],
                    label="Phase",
                    
                ),
                dict(
                    method="update",
                    args=[{"visible": [False, False, True, False, False]},
                          {'shapes': [shapes[2]], "title": "B0 Fieldmap (Hz)"}],
                    label="B0 field map (Hz)",
                ),
                dict(
                    method="update",
                    args=[{"visible": [False, False, False, True, False]},
                          {'shapes': [shapes[3]], "title": f"B0 Fieldmap ({MICRO_UNICODE}Tesla)"}],
                    label=f"B0 field map ({MICRO_UNICODE}Tesla)",
                ),
                dict(
                    method="update",
                    args=[{"visible": [False, False, False, False, True]},
                          {'shapes': [shapes[4]], "title": "B0 Fieldmap (ppm)"}],
                    label="B0 field map (ppm)",
                )
            ]),
            direction="down",
            showactive=True,
        )
    ]
)
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.show()