import numpy as np
from tifffile import imread
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.patches as mpatches
import matplotlib.cbook as cbook
from matplotlib_scalebar.scalebar import ScaleBar
from ipywidgets import interact, IntSlider

# Load the TIFF file into a NumPy array
label_volume = imread('data/Li_and_Cu_segmented_SaveImage11-19.align.view.labels_3dtiff.tif')

# Define RGB colors for each label (0 = background, 1 = other, 2 = Cu, 3 = Li)
label_colors = {
    0: [255, 255, 255],         # white for background
    1: [255, 255, 0],     # yellow for other label (if exists)
    2: [255, 194, 10],   # dark magenta for Current Collector
    3: [12, 123, 220]        # dark green for Li 
}

# Create an empty RGB volume
rgb_volume = np.zeros(label_volume.shape + (3,), dtype=np.uint8)

# Map labels to RGB
for label, color in label_colors.items():
    rgb_volume[label_volume == label] = color

# Check the dimensions of the array
#print(tiff_array)

def show_slice(data, slice_num):
        fig, ax = plt.subplots(layout= 'constrained')
        ax.imshow(data[slice_num, :, :])
        plt.title(f'slice number: {slice_num}')
        ax.axis('off')
        CC = mpatches.Patch(color= '#ffc20A', label='Current Collector')
        Li = mpatches.Patch(color= '#0c7bdc', label='Lithium')
        ax.legend(handles=[CC, Li], loc='upper left')
        scalebar = ScaleBar(200, "nm", length_fraction=0.1, location="lower right")
        ax.set_aspect(1)
        ax.add_artist(scalebar)
        plt.show()

@interact(slice_num=IntSlider(min=0, max=rgb_volume.shape[0]-1, step=1, description='Slice Number'))
def explore_slices(slice_num):
    show_slice(rgb_volume, slice_num)