import json
import ipywidgets as widgets
from IPython.display import display
from ipyniivue import ColorMap, NiiVue, ShowRender, SliceType

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

nv.set_radiological_convention(False)
nv.opts.multiplanar_show_render = ShowRender.ALWAYS
nv.set_slice_type(SliceType.MULTIPLANAR)

nv.load_volumes([{"path": "../images/FLAIR.nii.gz"}])

## Set up Callbacks

drawing_loaded = False

location_label = widgets.Label(value="")

@nv.on_location_change
def handle_location_change(location):
    """Update the location label when the crosshair location changes."""
    location_label.value = f"Location: {location['string']}"

@nv.on_image_loaded
def handle_image_loaded(volume):
    """Reset drawing settings when a new image is loaded."""
    global drawing_loaded
    if not drawing_loaded:
        drawing_loaded = True
        nv.load_drawing("../images/lesion.nii.gz")
    nv.set_drawing_enabled(False)
    draw_pen.value = -1  # Set the drawPen dropdown to 'Off'

## Create GUI Controls

# Dropdown for Draw Pen
draw_pen = widgets.Dropdown(
    options=[
        ("Off", -1),
        ("Erase", 0),
        ("Red", 1),
        ("Green", 2),
        ("Blue", 3),
        ("Filled Erase", 8),
        ("Filled Red", 9),
        ("Filled Green", 10),
        ("Filled Blue", 11),
        ("Erase Selected Cluster", 12),
    ],
    value=-1,
    description="Draw color:",
)

# Movement Buttons
left_button = widgets.Button(description="Left")
right_button = widgets.Button(description="Right")
posterior_button = widgets.Button(description="Posterior")
anterior_button = widgets.Button(description="Anterior")
inferior_button = widgets.Button(description="Inferior")
superior_button = widgets.Button(description="Superior")

# Other Buttons
save_button = widgets.Button(description="Save Drawing")
undo_button = widgets.Button(description="Undo")
growcut_button = widgets.Button(description="Grow Cut")

# Draw Opacity Slider
draw_opacity = widgets.IntSlider(
    value=80,
    min=0,
    max=100,
    step=1,
    description="Drawing Opacity",
    readout=False,
    style={"description_width": "initial"},
)

# Checkboxes
fill_pen_overwrites_checkbox = widgets.Checkbox(
    value=True,
    description="Fill pen overwrites",
)
radiological_checkbox = widgets.Checkbox(
    value=False,
    description="Radiological",
)
world_space_checkbox = widgets.Checkbox(
    value=False,
    description="World space",
)
linear_interpolation_checkbox = widgets.Checkbox(
    value=True,
    description="Linear Interpolation",
)
highdpi_checkbox = widgets.Checkbox(
    value=True,
    description="HighDPI",
)

# Textarea for Custom Colormap
text_value = """{
    "R": [0, 255, 22, 127],
    "G": [0, 20, 192, 187],
    "B": [0, 152, 80, 255],
    "labels": ["clear", "pink", "lime", "sky"]
}"""
num_lines = text_value.count("\n") + 1  # Adding 1 to include the last line
script_text = widgets.Textarea(
    value=text_value,
    description="Colormap",
    layout=widgets.Layout(width="60%"),
    rows=num_lines,  # Set rows to the number of lines in your text
)

# Apply Button for Custom Colormap
custom_button = widgets.Button(description="Apply")

## Define Event Handlers

def on_draw_opacity_change(change):
    """Handle changes in drawing opacity."""
    nv.draw_opacity = change["new"] / 100

draw_opacity.observe(on_draw_opacity_change, names="value")

def on_draw_pen_change(change):
    """Handle changes in draw pen selection."""
    mode = int(change["new"])
    nv.set_drawing_enabled(mode >= 0)
    if 0 <= mode <= 11:
        nv.set_pen_value(mode & 7, mode > 7)
    if mode == 12:
        # Erase selected cluster
        nv.set_pen_value(-0.0, False)

draw_pen.observe(on_draw_pen_change, names="value")

def on_left_click(b):
    """Move left."""
    nv.move_crosshair_in_vox(-1, 0, 0)

def on_right_click(b):
    """Move right."""
    nv.move_crosshair_in_vox(1, 0, 0)

def on_posterior_click(b):
    """Posterior."""
    nv.move_crosshair_in_vox(0, -1, 0)

def on_anterior_click(b):
    """Anterior."""
    nv.move_crosshair_in_vox(0, 1, 0)

def on_inferior_click(b):
    """Inferior."""
    nv.move_crosshair_in_vox(0, 0, -1)

def on_superior_click(b):
    """Superior."""
    nv.move_crosshair_in_vox(0, 0, 1)

left_button.on_click(on_left_click)
right_button.on_click(on_right_click)
posterior_button.on_click(on_posterior_click)
anterior_button.on_click(on_anterior_click)
inferior_button.on_click(on_inferior_click)
superior_button.on_click(on_superior_click)

def on_undo_click(b):
    """Undo drawing action."""
    nv.draw_undo()

undo_button.on_click(on_undo_click)

def on_growcut_click(b):
    """Draw grow cut."""
    nv.draw_grow_cut()

growcut_button.on_click(on_growcut_click)

def on_save_click(b):
    """Save drawing."""
    nv.save_image("test.nii", is_save_drawing=True)

save_button.on_click(on_save_click)

def on_fill_pen_overwrites_change(change):
    """Draw fill overwrites option."""
    nv.draw_fill_overwrites = change["new"]

fill_pen_overwrites_checkbox.observe(on_fill_pen_overwrites_change, names="value")

def on_radiological_change(change):
    """Set radiological convention."""
    nv.set_radiological_convention(change["new"])

radiological_checkbox.observe(on_radiological_change, names="value")

def on_world_space_change(change):
    """Set slice mm."""
    nv.set_slice_mm(change["new"])

world_space_checkbox.observe(on_world_space_change, names="value")

def on_linear_interpolation_change(change):
    """Set interpolation."""
    nv.set_interpolation(not change["new"])

linear_interpolation_checkbox.observe(on_linear_interpolation_change, names="value")

def on_highdpi_change(change):
    """Set high resolution capable."""
    nv.set_high_resolution_capable(change["new"])

highdpi_checkbox.observe(on_highdpi_change, names="value")

def on_custom_button_click(b):
    """Set draw colormap."""
    try:
        val = script_text.value
        cmap = ColorMap(**json.loads(val))
        nv.set_draw_colormap(cmap)
    except Exception as e:
        print(f"Error applying custom colormap: {e}")

custom_button.on_click(on_custom_button_click)

## Setup controls arrangement

controls_row1 = widgets.HBox([draw_pen])

movement_buttons = widgets.HBox(
    [
        left_button,
        right_button,
        posterior_button,
        anterior_button,
        inferior_button,
        superior_button,
    ]
)

other_buttons = widgets.HBox([save_button, undo_button, growcut_button])

checkboxes1 = widgets.HBox(
    [
        fill_pen_overwrites_checkbox,
        radiological_checkbox,
        world_space_checkbox,
    ]
)

checkboxes2 = widgets.HBox(
    [
        linear_interpolation_checkbox,
        highdpi_checkbox,
    ]
)

controls = widgets.VBox(
    [
        controls_row1,
        movement_buttons,
        other_buttons,
        draw_opacity,
        checkboxes1,
        checkboxes2,
        script_text,
        custom_button,
        location_label,
    ]
)

## Display all

display(controls, nv)
Drawing¶
