# Prepare Python environment

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" 

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 dipole_kernel(b0_dir, voxel_size, n_voxels):
    """ Create a dipole kernel
    dipole kernel: (3*cos(theta)**2 - 1) / (4*pi*r**3)
                => (3*r**2*cos(theta)**2 - r**2) / (4*pi*r**5)
                => (3*b0_dir**2 - r**2) / (4*pi*r**2**2.5)

        Function inspired and derived from: https://onlinelibrary.wiley.com/doi/10.1002/mrm.28716
    """
    eps = 0.00001
    x, y, z = np.meshgrid(range(round(-n_voxels[0]/2+0.5), round(n_voxels[0]/2+0.5)), range(round(-n_voxels[1]/2+0.5), round(n_voxels[1]/2+0.5)), range(round(-n_voxels[2]/2+0.5), round(n_voxels[2]/2+0.5)), indexing='ij')

    x = x * voxel_size[0] + eps
    y = y * voxel_size[1] + eps
    z = z * voxel_size[2] + eps

    r2 = (x**2 + y**2 + z**2)

    d = np.prod(voxel_size) * ( 3 * ((x*b0_dir[0] + y*b0_dir[1] + z*b0_dir[2])**2) - r2 ) / (4 * math.pi * r2**2.5)

    d[np.isnan(d)] = eps
    D = np.real(fftshift(fftn(ifftshift(d))))

    mid_voxel = n_voxels[0]//2
    return d[n_voxels[1]//2], D[n_voxels[1]//2]

b0_dir = (0, 0, 1)
voxel_size = np.array((1, 1, 1)) * 1e-3
n_voxels = (201,201,201)
d, D = dipole_kernel(b0_dir, voxel_size, n_voxels)

fig = go.Figure()
fig = make_subplots(rows=1, cols=2, shared_xaxes=False, horizontal_spacing=0.13, vertical_spacing = 0.12, subplot_titles=("Dipole Kernel (d)", "Dipole Kernel (D)"), specs=[[{"type": "Heatmap"}, {"type": "Heatmap"}]])
fig.add_trace(go.Heatmap(z=d, colorscale='gray', showscale=False, zmin=-1e-6, zmax=1e-6))
fig.add_trace(go.Heatmap(z=D, colorscale='gray', showscale=False), 1, 2)
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.update_layout(
    height=400,
    width=750)
fig.show()