import pathlib

import ipywidgets as widgets
from IPython.display import display

import ipyniivue
from ipyniivue import NiiVue, download_dataset

## Download Formats

DATA_FOLDER = pathlib.Path(ipyniivue.__file__).parent / "images"

local_imgs = [
    "FLAIR.nrrd",
    "cactus.nii.gz",
    "DoG.png",
    "anat_final.FT+tlrc.HEAD",
    "mha.mha",
    "template.mif.gz",
    "trix/fa.mif",
    "dsistudio.src.gz",
    "dsistudio.fib.gz",
    "wm.mgz",
]

files_to_download_local = local_imgs.copy()
for img in local_imgs:
    if img.endswith(".HEAD"):
        # Also need to download the paired .BRIK file
        brik_file = img.replace(".HEAD", ".BRIK")
        files_to_download_local.append(brik_file)

download_dataset(
    api_url="https://niivue.com/demos/images/",
    dest_folder=DATA_FOLDER,
    files=files_to_download_local,
)

# Download images

modalities_imgs = [
    "chris_MRA",
    "chris_PD",
    "chris_t1",
    "chris_t2",
    "CT_Abdo",
    "CT_AVM",
    "CT_Electrodes",
    "CT_Philips",
    "CT_pitch",
    "fmri_pitch",
    "Iguana",
    "mni152",
    "MR_Gd",
    "pcasl",
    "spm152",
    "spmMotor",
    "visiblehuman",
    "rgb_bmp.jpg",
    "gray_bmp.png",
    "HCD1464653.qsdr.fz",
]

# Add the appropriate extensions
files_to_download_modalities = []
for img in modalities_imgs:
    if (
        not img.endswith(".png")
        and not img.endswith(".jpg")
        and not img.endswith(".fz")
    ):
        img_with_ext = img + ".nii.gz"
    else:
        img_with_ext = img
    files_to_download_modalities.append(img_with_ext)

download_dataset(
    api_url="https://niivue.github.io/niivue-demo-images/",
    dest_folder=DATA_FOLDER,
    files=files_to_download_modalities,
)

## Setup NiiVue Instance

nv = NiiVue(
    back_color=(0.7, 0.7, 0.9, 1),
)
nv.opts.is_colorbar = True
nv.set_slice_type("RENDER")
nv.set_clip_plane(0.35, 270, 0)

initial_volume_path = DATA_FOLDER / "mni152.nii.gz"
nv.load_volumes(
    [
        {
            "path": initial_volume_path,
            "colormap": "gray",
            "opacity": 1,
            "visible": True,
        }
    ]
)

## Store image paths

img_paths = {}

# Formats
for img in local_imgs:
    img_path = DATA_FOLDER / img
    img_paths[img] = img_path
    if img.endswith(".HEAD"):
        # Also map the paired .BRIK file
        brik_file = img.replace(".HEAD", ".BRIK")
        brik_path = DATA_FOLDER / brik_file
        img_paths[img + "_paired"] = brik_path

# Modalities
for img in modalities_imgs:
    if (
        not img.endswith(".png")
        and not img.endswith(".jpg")
        and not img.endswith(".fz")
    ):
        img_with_ext = img + ".nii.gz"
    else:
        img_with_ext = img
    img_path = DATA_FOLDER / img_with_ext
    img_paths[img] = img_path


## User interface dropdowns

# Dropdown for Formats
formats_dropdown = widgets.Dropdown(
    options=[(img, img) for img in local_imgs],
    description="Formats:",
    value=None,
)

# Dropdown for Modalities
modalities_dropdown = widgets.Dropdown(
    options=[(img, img) for img in modalities_imgs],
    description="Modalities:",
    value=None,
)

# Dropdown for Colormaps
colormaps = sorted(nv.colormaps())
colormap_dropdown = widgets.Dropdown(
    options=colormaps,
    description="Colormap:",
    value="gray",
)

## Event handlers

def on_format_change(change):
    """Handle selection in Formats Dropdown."""
    selected_img = change["new"]
    if not selected_img:
        return
    img_path = img_paths.get(selected_img)
    if not img_path.exists():
        print(f"Image {img_path} not found.")
        return
    volumeObj = {
        "path": img_path,
        "colormap": "gray",
        "opacity": 1,
        "visible": True,
    }
    if selected_img.endswith(".HEAD"):
        paired_img_path = img_paths.get(selected_img + "_paired")
        if not paired_img_path or not paired_img_path.exists():
            print(f"Paired image {paired_img_path} not found.")
            return
        volumeObj["paired_img_path"] = paired_img_path
    print("Loading format image:", img_path.name)
    nv.load_volumes([volumeObj])
    # Reset Modalities Dropdown
    modalities_dropdown.unobserve(on_modality_change, names="value")
    modalities_dropdown.value = None
    modalities_dropdown.observe(on_modality_change, names="value")
    # Update colormap
    if nv.volumes:
        colormap_dropdown.value = nv.volumes[0].colormap


def on_modality_change(change):
    """Handle selection in Modalities Dropdown."""
    selected_img = change["new"]
    if not selected_img:
        return
    img_path = img_paths.get(selected_img)
    if not img_path.exists():
        print(f"Image {img_path} not found.")
        return
    volumeObj = {
        "path": img_path,
        "colormap": "gray",
        "opacity": 1,
        "visible": True,
    }
    print("Loading modality image:", img_path.name)
    nv.load_volumes([volumeObj])
    # Reset Formats Dropdown
    formats_dropdown.unobserve(on_format_change, names="value")
    formats_dropdown.value = None
    formats_dropdown.observe(on_format_change, names="value")
    # Update colormap
    if nv.volumes:
        colormap_dropdown.value = nv.volumes[0].colormap


def on_colormap_change(change):
    """Handle selection in Colormap Dropdown."""
    selected_cmap = change["new"]
    if nv.volumes:
        nv.set_colormap(nv.volumes[0].id, selected_cmap)

## Observe the dropdowns

formats_dropdown.observe(on_format_change, names="value")
modalities_dropdown.observe(on_modality_change, names="value")
colormap_dropdown.observe(on_colormap_change, names="value")

## Display

controls = widgets.VBox(
    [
        formats_dropdown,
        modalities_dropdown,
        colormap_dropdown,
    ]
)

display(controls, nv)

Downloading FLAIR.nrrd...

Downloading cactus.nii.gz...

Downloading DoG.png...
Downloading anat_final.FT+tlrc.HEAD...

Downloading mha.mha...

Downloading template.mif.gz...

Downloading trix/fa.mif...

Downloading dsistudio.src.gz...

Downloading dsistudio.fib.gz...

Downloading wm.mgz...

Downloading anat_final.FT+tlrc.BRIK...
Dataset downloaded successfully to /home/runner/work/ipyniivue/ipyniivue/src/ipyniivue/images.

Modalities and format¶