# Prepare Python environment
import numpy as np
import plotly.express as px
import os
import nibabel as nib
import numpy as np
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import math
import json

import scipy.io as sio
from pathlib import Path

from contextlib import contextmanager
import sys, os
from pathlib import Path

@contextmanager
def suppress_stdout():
    with open(os.devnull, "w") as devnull:
        old_stdout = sys.stdout
        sys.stdout = devnull
        try:  
            yield
        finally:
            sys.stdout = old_stdout

import os
from pathlib import Path

def find_myst_yml_directories(start_dir=None):
    """
    Recursively search for directories containing myst.yml file.
    
    Args:
        start_dir (str or Path): Starting directory (defaults to current directory)
    
    Returns:
        list: List of full paths to directories containing myst.yml
    """
    if start_dir is None:
        start_dir = Path.cwd()
    else:
        start_dir = Path(start_dir)
    
    myst_dirs = []
    
    def _search_directory(current_dir):
        # Check if myst.yml exists in current directory
        myst_file = current_dir / "myst.yml"
        if myst_file.exists():
            myst_dirs.append(str(current_dir.resolve()))
            # Don't search subdirectories if we found myst.yml here
            return
        
        # Recursively search all subdirectories
        for item in current_dir.iterdir():
            if item.is_dir():
                try:
                    _search_directory(item)
                except (PermissionError, OSError):
                    # Skip directories we can't access
                    continue
    
    _search_directory(start_dir)
    return myst_dirs

def find_myst_yml_directories_upwards(start_dir=None):
    """
    Search for myst.yml in current directory, if not found go to parent and repeat.
    
    Args:
        start_dir (str or Path): Starting directory (defaults to current directory)
    
    Returns:
        str or None: Full path of directory containing myst.yml, or None if not found
    """
    if start_dir is None:
        current_dir = Path.cwd()
    else:
        current_dir = Path(start_dir)
    
    # Keep going up until we reach the filesystem root
    while current_dir != current_dir.parent:  # Stop at root
        myst_file = current_dir / "myst.yml"
        if myst_file.exists():
            return str(current_dir.resolve())
        
        # Move to parent directory
        current_dir = current_dir.parent
    
    return None

with suppress_stdout():
    repo_path = Path(find_myst_yml_directories_upwards())
    print(repo_path)
    data_req_path = repo_path / "binder" / "data_requirement.json"
    data_path = repo_path / "data"
    dataset_path = data_path / "qmrlab-mooc"
    
data_dir = dataset_path / "05-B0" / "05-B0" / "data"/ "fmap" 

PI_UNICODE = "\U0001D70B"
fname_mag_e1 = data_dir / "sub-fmap_magnitude1.nii.gz"
fname_phase_e1 = data_dir / "sub-fmap_phase1.nii.gz"
fname_phase_e1_json = data_dir / "sub-fmap_phase1.json"
fname_phase_e2 = data_dir / "sub-fmap_phase2.nii.gz"
fname_mask = data_dir / "mask.nii.gz"
fname_fmap = data_dir / "fmap.nii.gz"

nii_mag_e1 = nib.load(fname_mag_e1)
nii_phase_e1 = nib.load(fname_phase_e1)
nii_phase_e2 = nib.load(fname_phase_e2)
nii_mask = nib.load(fname_mask)
nii_fmap = nib.load(fname_fmap)

# Phase evolution though different echo times
mask = nii_mask.get_fdata()[30:-30,8:105,30]
fmap = nii_fmap.get_fdata()[30:-30,8:105,30] * mask  # [Hz]
phase1 = (nii_phase_e1.get_fdata()[30:-30,8:105,30] / 4095 * 2 * math.pi - math.pi) * mask

with open(fname_phase_e1_json, 'r') as json_data1:
    data1 = json.load(json_data1)
    
echo_time1 = data1['EchoTime']
phase0 = phase1 - (echo_time1 * (fmap * 2 * math.pi))
zmin = -math.pi
zmax = math.pi

steps = 31
last_echo_time = 0.03
echo_times = np.linspace(0.0, last_echo_time, steps)
fig = go.Figure()
for i_echo, echo_time in enumerate(echo_times):
    phase = phase0 + (fmap * echo_time * 2 * math.pi)
    phase = np.angle(np.exp(1j*phase))
    if i_echo >= len(echo_times) - 1:
        fig.add_trace(go.Heatmap(z=np.rot90(phase, k=-1), visible=True, coloraxis = "coloraxis"))
    else:
        fig.add_trace(go.Heatmap(z=np.rot90(phase, k=-1), visible=False, coloraxis = "coloraxis"))

fig.update_layout(coloraxis = {'colorscale':'gray'},
                 coloraxis_cmin=zmin, coloraxis_cmax=zmax)
fig.update_coloraxes(
    colorbar=dict(title="Rad",
                  titleside="top",
                  tickmode="array",
                  tickvals=[-math.pi, 0, math.pi-0.01],
                  ticktext = [f"-{PI_UNICODE}", 0, f'{PI_UNICODE}']))

echo_times_str = [f"{time:.2}" for time in echo_times]
steps = []
for i in range(len(fig.data)):
    step = dict(
        method="update",
        label=echo_times_str[i],
        args=[{"visible": [False] * len(fig.data)}],  # layout attribute
    )
    step["args"][0]["visible"][i] = True  # Toggle i'th trace to "visible"
    steps.append(step)

sliders = [dict(
    active=30,
    currentvalue={"prefix": "Echo Time (s): "},
    steps=steps
)]

fig.update_layout(
    sliders=sliders
)

fig.update_layout(
    title=dict(text="Phase at different echo times", x=0.5)
)
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.update_layout({"height": 550, "width": 500})
fig.show()