from ipyniivue import NiiVue
import ipywidgets as widgets
from IPython.display import display
import math

nv = NiiVue(back_color=(0.3, 0.3, 0.3, 1))

meshLHLayersList1 = [
    {
        "path": "../images/BrainMesh_ICBM152.lh.curv",
        "colormap": "gray",
        "cal_min": 0.3,
        "cal_max": 0.5,
        "opacity": 0.7,
    },
    {
        "path": "../images/BrainMesh_ICBM152.lh.motor.mz3",
        "cal_min": 1.64,
        "cal_max": 5,
        "colormap": "warm",
        "colormap_negative": "winter",
        "use_negative_cmap": True,
        "opacity": 0.7,
    },
]

nv.load_meshes(
    [
        {
            "path": "../images/BrainMesh_ICBM152.lh.mz3",
            "layers": meshLHLayersList1,
        }
    ]
)

nv.opts.is_colorbar = True

nv.meshes[0].layers[0].colorbar_visible = False

## Create interactive widgets

# Invert Checkbox
invert_check = widgets.Checkbox(
    value=False,
    description="Invert",
)


def on_invert_change(change):
    """Set mesh layer 1 colormap invert."""
    nv.set_mesh_layer_property(nv.meshes[0].id, 1, "colormap_invert", change["new"])


invert_check.observe(on_invert_change, names="value")

# Curvature Slider
curve_slider = widgets.IntSlider(
    value=7,
    min=0,
    max=10,
    step=1,
    description="Curvature",
    readout=False
)


def on_curve_slider_change(change):
    """Set mesh layer 0 opacity."""
    nv.set_mesh_layer_property(nv.meshes[0].id, 0, "opacity", change["new"] * 0.1)


curve_slider.observe(on_curve_slider_change, names="value")

# Opacity Slider
opacity_slider = widgets.IntSlider(
    value=7,
    min=0,
    max=10,
    step=1,
    description="Opacity",
    readout=False
)


def on_opacity_slider_change(change):
    """Set mesh layer 1 opacity."""
    nv.set_mesh_layer_property(nv.meshes[0].id, 1, "opacity", change["new"] * 0.1)


opacity_slider.observe(on_opacity_slider_change, names="value")

# Save Bitmap Button
save_bmp_button = widgets.Button(
    description="Save Bitmap",
)


def on_save_bmp_clicked(b):
    """Save scene."""
    nv.save_scene("ScreenShot.png")
    with out:
        from IPython.display import clear_output

        clear_output(wait=True)
        print('Scene saved as "ScreenShot.png"')


save_bmp_button.on_click(on_save_bmp_clicked)

# Output widget for messages
out = widgets.Output()

## Create the colormap header buttons

# Get available colormaps
colormaps = nv.colormaps()
colormap_buttons = []


def create_colormap_button(name):
    """Create colormap."""
    btn = widgets.Button(description=name)

    def on_click(b):
        """Set mesh layer colormap."""
        nv.set_mesh_layer_property(nv.meshes[0].id, 1, "colormap", name)

    btn.on_click(on_click)
    return btn


colormap_buttons = [create_colormap_button(name) for name in colormaps]

# Organize colormap buttons in a grid
num_cols = 10
num_rows = math.ceil(len(colormap_buttons) / num_cols)
colormap_grid = []

for i in range(num_rows):
    row_buttons = colormap_buttons[i * num_cols : (i + 1) * num_cols]
    colormap_grid.append(widgets.HBox(row_buttons))

colormap_buttons_widget = widgets.VBox(colormap_grid)

## Display 

# Display header controls
header_controls = widgets.HBox(
    [
        invert_check,
        curve_slider,
        opacity_slider,
        save_bmp_button,
    ]
)

# Display header, image, footer and status 

display(header_controls)
display(nv)
display(colormap_buttons_widget)
display(out)
Mesh colormaps¶
