# 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]


# Load cylinder (Y 64)
susc = nib.load(data_dir / "field_simulations" / "cylinder" / "Chi.nii.gz").get_fdata()
fmap_hz_all = nib.load(data_dir / "field_simulations" / "cylinder" / "fmap_hz.nii.gz").get_fdata()
local_field_cyl = nib.load(data_dir  / "field_simulations" / "cylinder" / "local_field.nii.gz").get_fdata()

# Load brain (Z 210)
susc_brain = nib.load(data_dir / "field_simulations" / "brain" / "chi_masked.nii.gz").get_fdata()
fmap_hz_brain_all = nib.load(data_dir / "field_simulations" / "brain" / "fmap_masked.nii.gz").get_fdata()
local_field_brain = nib.load(data_dir / "field_simulations" / "brain" / "local_field.nii.gz").get_fdata()

fig = make_subplots(rows=1, cols=3, shared_xaxes=False, horizontal_spacing=0.13, vertical_spacing = 0.12, subplot_titles=(f"Susceptibility distribution ({CHI_UNICODE})", "Simulated B0 map", "Simulated B0 map<br>no background field"), specs=[[{"type": "Heatmap"}, {"type": "Heatmap"}, {"type": "Heatmap"}]])
fig.add_trace(go.Heatmap(z=susc, colorscale='gray', colorbar_x=1/3 - 0.09, colorbar=dict(title="ppm", titleside="top")), 1, 1)
fig.add_trace(go.Heatmap(z=fmap_hz_all, colorscale='gray', colorbar_x=2/3 - 0.045, colorbar=dict(title="Hz", titleside="top")), 1, 2)
fig.add_trace(go.Heatmap(z=local_field_cyl, colorscale='gray', colorbar_x=1-0.004, colorbar=dict(title="Hz", titleside="top")), 1, 3)
fig.add_trace(go.Heatmap(z=np.rot90(susc_brain, k=-1), colorscale='gray', colorbar_x=1/3 - 0.09, zmin=-0.5, zmax=0.5, colorbar=dict(title="ppm", titleside="top"), visible=False), 1, 1)
fig.add_trace(go.Heatmap(z=np.rot90(fmap_hz_brain_all, k=-1), colorscale='gray', colorbar_x=2/3 - 0.045, zmin=1100, zmax=2300, colorbar=dict(title="Hz", titleside="top"), visible=False), 1, 2)
fig.add_trace(go.Heatmap(z=np.rot90(local_field_brain, k=-1), colorscale='gray', zmin=-4, zmax=4, colorbar_x=1-0.004, colorbar=dict(title="Hz", titleside="top"), visible=False), 1, 3)

### Create buttons for drop down menu
labels = ["Cylinders", "Brain"]
buttons = []
for i, label in enumerate(labels):
    if label == "Cylinders":
        visibility = [True, True, True, False, False, False]
    else:
        visibility = [False, False, False, True, True, True]
    button = dict(
                 label =  label,
                 method = 'update',
                 args = [{'visible': visibility},
                         {'title': label}])
    buttons.append(button)

updatemenus = list([
    dict(active=0,
         x=0.5,
         y=-0.15,
         buttons=buttons,
         showactive=True,
    )
])

fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.update_layout({"height": 350, "width": 750},
                  title_text="B0 maps from susceptibility simulations",
                  title_x=0.5,
                  updatemenus=updatemenus,
                  showlegend=False
                 )
fig.show()