Package napari_blender

def apply_shape_key(obj, target, index):
    """
    Set shape key to morph one object into another.

    Parameters
    ----------
    obj : bpy.types.Object
        The object to morph.
    target : bpy.types.Object
        The object to morph into.
    index : int
        The index of the current frame.
    """
    # Ensure the target object has the same number of vertices as the object
    # if len(obj.data.vertices) != len(target.data.vertices):
    #     print(f'{len(obj.data.vertices)} != {len(target.data.vertices)}')
    #     print(f"Error: Objects ({obj.name} & {target.name}) have different number of vertices.")
    #     return

    obj.shape_key_add(from_mix=False)
    obj.shape_key_add(name='Morph')

    # Get the coordinates of vertices for both objects
    obj_vertices = np.array([vert.co for vert in obj.data.vertices])
    target_vertices = np.array([vert.co for vert in target.data.vertices])

    # Calculate the cost matrix (Euclidean distance between vertices)
    cost_matrix = np.linalg.norm(obj_vertices[:, np.newaxis] - target_vertices, axis=2)

    # Use Hungarian algorithm to find the minimum cost vertex mapping
    row_ind, col_ind = linear_sum_assignment(cost_matrix)

    # Apply the shape key by mapping vertices from the target to the object
    for i, j in zip(row_ind, col_ind):
        w = target.data.vertices[j]
        obj.data.shape_keys.key_blocks['Morph'].data[i].co = w.co
    
    set_shape_keys(obj, target, index)

def calculate_are(truth, pred):
    """
    Calculates adapted Rand error, precision, and recall.

    Parameters
    ----------
    truth : numpy.ndarray
        Ground truth labels.
    pred : numpy.ndarray
        Predicted labels.

    Returns
    -------
    tuple
        Adapted Rand error, precision, and recall.
    """
    return skimage.metrics.adapted_rand_error(truth, pred)

def calculate_distance(coords1, coords2):
    """
    Helper function that calculates the Euclidean distance between two 3D coordinates.

    Parameters
    ----------
    coords1 : tuple
        First set of coordinates (x, y, z).
    coords2 : tuple
        Second set of coordinates (x, y, z).
    
    Returns
    -------
    float
        The Euclidean distance between the two sets of coordinates.
    """
    return math.sqrt(sum((c1 - c2)**2 for c1, c2 in zip(coords1, coords2)))

def calculate_jaccard(truth, prediction):
    """
    Calculates the Jaccard Index (JI) score between ground truth and predicted masks.

    Parameters
    ----------
    truth : numpy.ndarray
        Ground truth labels.
    prediction : numpy.ndarray
        Predicted labels.

    Returns
    -------
    float
        Jaccard Index (JI) score.
    """
    ji_values = compare_labeled_volumes(truth, prediction)
    try:
        final_values = [t[2] for t in ji_values]
        ji = sum(final_values) / len(final_values) if len(final_values) > 0 else 0
        return ji
    except TypeError:
        print('Could not calculate Jaccard index')

def change_scene_obj(dicts):
    """
    Helper function that applies changes to the scene based on a dictionary generated
    by different visualization modes. This arranges the scene to be ready for that mode.

    Parameters
    ----------
    dicts : dict
        Dictionary containing multiple dictionaries, indicating the required starting
        conditions for the visualization mode.
    """
    # Get sub-dictionaries that indicate respective imports
    texts = dicts['texts']
    materials = dicts['materials']
    hide_obj = dicts['hide_obj']
        
    # Change all text objects indicated
    for text_name, text in texts.items():
        change_text(text, text_name)

    # Generate all indicated materials
    for count, (mat_name, values) in enumerate(materials.items()):
        # Check if material contains the correct number of parameters
        if len(values) == 5:
            create_material(mat_name, values[:4], values[-1])
            # Set this material to be correctly displayed in the color legend
            set_legend_color(mat_name, count+1)
        else:
            print('Material color object does not contain the correct amount of values.:\n(Hue,Saturation,Value,Alpha,Material Blend Mode)')
    
    # After creating materials, hide unused legend items from render.
    hide_legend_items(len(materials))

    # (Un)Hide all objects indicated
    for obj_name, hide in hide_obj.items():
        try:
            # First check if name is collection, before checking if it is an object
            coll = find_collection(obj_name)
            coll.hide_render = hide
        except CollectionNotFoundException:
            try:
                obj = find_object(obj_name)
                obj.hide_render = hide
            except ObjectNotFoundException:
                print(f'Changing hide render setting to {hide} failed')

def change_text(text, text_name):
    """
    Changes the text of a specified text object in the scene.

    Parameters
    ----------
    text : str
        The new text to set.
    text_name : str
        The name of the text object to modify.
    """
    # Find the text object by name
    try:
        text_obj = find_object(text_name)
        text_obj.data.body = text
    except ObjectNotFoundException:
        print('Changing text suspended')

def clamp_hue(ji):
    """
    Clamp the hue value to a specific range.

    Parameters
    ----------
    ji : float
        The hue value to clamp [0..1].

    Returns
    -------
    float or None
        Clamped hue value.
    """
    if ji is not None and 0 <= ji <= 1:
        new_value = ji * 0.3
        return new_value
    else:
        print(f"Jaccard index should be within the range [0, 1] but was {ji}")
        return None

def clamp_indicator(metric):
    """
    Clamp the indicator value.

    Parameters
    ----------
    metric : float
        The metric value to clamp.
    
    Returns
    -------
    float
        Clamped metric value.
    """
    # Check if value is within the range [0, 1]
    if metric is not None and 0 <= metric <= 1:
        # Rescale the value to the range [0.335, 0].
        # This is the hue that represents a dark green color
        # that will show a prediction to be perfect
        new_value = -0.34 + (0.36 + 0.34) * metric
        return new_value
    else:
        print(f"Metric should be within the range [0, 1] but was {metric}")
        return 0

def clean_data(data, frame=None):
    """
    Cleans and preprocesses image data.

    This function accepts image data as a NumPy array and performs preprocessing steps such as
    selecting a specific frame, converting to grayscale, and scaling pixel values.

    
    Parameters
    ----------
    data : numpy.ndarray
        The input image data as a NumPy array.
    frame : int, optional
        The frame index to select from the image data.
    
    Returns
    -------
    numpy.ndarray
        Preprocessed image data.
    """
    if frame is not None:
        try:
            if data.ndim == 3:
                # If data has 3 dimensions, return as is
                result = data
            elif data.ndim == 4:
                # If data has 4 dimensions, select specific frame
                result = data[frame, :, :, :]
            elif data.ndim == 5:
                # If data has 5 dimensions, convert to grayscale
                result = np.dot(data[frame,:,:,:,:3], [0.299, 0.587, 0.114])
            else:
                raise ValueError("Data array must have either 3, 4 or 5 dimensions")
            
            # Check if result data type is uint8
            if result.dtype != np.uint8:
                # Scale pixel values to uint8 range [0, 255]
                min_val = np.min(result)
                max_val = np.max(result)
                scaled_image = ((result - min_val) * 255 / (max_val - min_val)).astype(np.uint8)
                return scaled_image
            return data  # Return preprocessed data
        except IndexError:
            print(f'Could not read data, image stack does not contain index {frame}')
    
    return data  # Return original data if no frame specified

def clean_duplicate_objs():
    """
    Clean up duplicate objects in the scene.
    """
    # Get all objects in the scene
    all_objects = D.objects

    # Iterate through objects and delete those with a dot in their name,
    # since objects with names that exist in the scene are named using "name".001,
    # this cleans up already imported objects.
    for obj in all_objects:
        if ("." in obj.name or "temp" in obj.name) and obj.type == "MESH":
            D.objects.remove(obj, do_unlink=True)

def color_nuclei(parent, i):
    """
    Color nuclei based on a given metric value.

    Parameters
    ----------
    parent : bpy.types.Object
        The parent object.
    i : int
        Index for the metric value.
    """
    ji_vals = s.METRICS[i]['ji_vals']
    children_objs = parent.children

    for child_obj in children_objs:
        group_id = child_obj.get("Group_ID")
        if group_id is not None:
            try:
                group_id = int(group_id)
            except ValueError:
                print(f"Warning: Invalid group_id {group_id} for object {child_obj.name}")
                continue

            if 'gt' in parent.name:
                matching_tuples = [t for t in ji_vals if t[0] == group_id]
                if not matching_tuples or matching_tuples[0][1] == 0:
                    set_material(child_obj, 'Not Predicted')  
                    print(f'{group_id} not predicted')
                else:
                    D.objects.remove(child_obj, do_unlink=True)

            if 'pred' in parent.name:
                matching_tuples = [t for t in ji_vals if t[1] == group_id]
                if matching_tuples:
                    set_hue(child_obj, matching_tuples[0][2])
                else:
                    set_material(child_obj, 'False Prediction')
        else:
            print(f"Warning: Group_ID not found for object {child_obj.name}")

def compare_labeled_volumes(truth, prediction):
    """
    Compares labelled volumes by performing a Jaccard Index (JI) calculation over all of 
    the labelled regions in the provided volumes.

    Parameters
    ----------
    truth : numpy.ndarray
        Ground truth labels.
    prediction : numpy.ndarray
        Predicted labels.

    Returns
    -------
    list of tuple
        A list of tuples containing truth label, prediction label, and Jaccard Index (JI).
    """
    tlabels = np.unique(truth)[1:]
    ji = []
    try:
        for label in tlabels:
            x, y, z = np.where(truth == label)
            values = prediction[x, y, z]
            freq, lbl = frequency(values)
            if lbl[0] == 0:
                freq[0] = 0
            mx = freq.max()
            mlbl = lbl[np.where(freq == mx)[0][0]]

            xp, _, _ = np.where(prediction == mlbl)
            ji.append((label, mlbl, 2 * mx / (len(x) + len(xp))))
        return ji
    except ValueError:
        print(f'There was a mismatch in the expected amount of labels in the data (3), and the labeled data provided ({tlabels.shape})')
        return

def confusion_matrix(truth, prediction):
    """
    Compute the confusion matrix between ground truth and predicted labels.

    Parameters
    ----------
    truth : numpy.ndarray
        Ground truth labels.
    prediction : numpy.ndarray
        Predicted labels.

    Returns
    -------
    tuple
        Confusion matrix as a 2D numpy array and corresponding labels.
    """
    tlabels = np.unique(truth)
    plabels = np.unique(prediction)

    num_classes = max(tlabels.max(), plabels.max()) + 1

    conf_matrix = np.zeros((num_classes, num_classes), dtype=np.int)
    labels = np.arange(num_classes)

    return conf_matrix, labels

def create_child_material(obj, shift):
    """
    Create a material based on another objects color by shifting the hue.

    Parameters
    ----------
    obj : bpy.types.Object
        The object whose material to adapt.
    shift : float
        Amount to change the hue value with.

    Returns
    -------
    str
        The name of the changed material.
    """
    if obj.material_slots:
        mat = obj.material_slots[obj.active_material_index].material
        new_material = mat.copy()
        new_material.name = mat.name + f"_{shift}"

        principled = new_material.node_tree.nodes["Principled BSDF"]
        h, s, v = rgb_to_hsv(principled.inputs["Base Color"].default_value[:3])
        r, g, b = hsv_to_rgb((h+shift, s, v))
        principled.inputs["Base Color"].default_value = (r, g, b, principled.inputs["Base Color"].default_value[-1])
        return new_material.name
    else:
        print("No materials assigned to the object.")

def create_contour(data, thresh):
    """
    Method to get the contours of a single slice (2D numpy array) of data.

    Parameters
    ----------
    data : numpy.ndarray
        Numpy array representation of a single z slice of data.
    thresh : int
        Threshold with which to threshold this slice.

    Returns
    -------
    numpy.ndarray
        Binary Numpy array containing the location of the contours (edges of nuclei).
    """
    # Threshold the data and make it binary
    _, thresh = cv.threshold(data, thresh, 255, 0)
    # Find contours (corresponds with nuclei) on this binary data, cv.RETR_EXTERNAL
    # is used to return the non-overlapping items per slice.
    contours, _ = cv.findContours(thresh, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)

    contour_image = np.zeros_like(data)

    # Sanitize slices by only allowing contours that are >= 15 pixels in area. This may lose
    # a very small edge of the nucleus but this is counteracted by how the 3D object is created,
    # namely using marching cubes.
    for contour in contours:
        area = cv.contourArea(contour)
        if area >= 15:
            cv.drawContours(contour_image, [contour], -1, (255), 1)
    return contour_image

def create_contour_data(data):
    """
    Method to yield the numpy array representation of only the contours found
    in the convoluted original data.

    Parameters
    ----------
    data : numpy.ndarray
        Numpy array representation of the convoluted original data.
    
    Returns
    -------
    numpy.ndarray
        Cleaned numpy array data containing the rough contours of nuclei.
    """
    # # Sanitize data without changing topology by erosion and dilation, this step
    # # elimates the heavy amount of background noise
    # eroded = ndi.grey_erosion(data, size=(3,3,3))
    # dilated = ndi.grey_dilation(eroded, size=(3,3,3))

    # Threshold image using industry standard (Otsu thresholding), with threshold
    # taken from the middle z slice. This is done since the middle z-slice will
    # contain part of the organoid, where the 1st slice will be only noise
    middle_slice = int(data.shape[2]/2)
    otsu = threshold_otsu(data[:,:,middle_slice])
    binary = data > otsu

    # opened = ndi.binary_opening(binary)

    data = data * binary

    # Take the gradient of this image and expand it to a numpy containing the gradient
    # of each pixel in each direction. A nucleus is a connected object so all of the x,y,z
    # frames will have a heavy gradient.
    gradient = np.gradient(data)
    magnitudes = (np.sqrt(sum(np.square(g) for g in gradient))).astype(np.uint8)

    # Stack contours for each z slice and export to object.
    contours = np.zeros(magnitudes.shape)
    for z in range(magnitudes.shape[2]):
        contours[:,:,z] = create_contour(magnitudes[:,:,z], otsu)

    return contours

def create_contour_obj(data, fname):
    """"
    Function to completely create a 3D contour object from data.

    Parameters
    ----------
    data : numpy.ndarray
        Numpy array representation of the convoluted original data.
    fname : str
        Filename to export 3D object to.
    """
    contours = create_contour_data(data)
    create_obj(contours, fname)

def create_difference(obj1, obj2):
    """Create a boolean difference between two objects.

    Parameters
    ----------
    obj1 : bpy.types.Object
        The first object to use for operation.
    obj2 : bpy.types.Object
        The second object to use for operation.
    """
    obj1.select_set(True)
    C.view_layer.objects.active = obj1

    # Create new difference modifier and apply it between obj1 and obj2
    bool_mod = C.active_object.modifiers.new(name="Boolean", type='BOOLEAN')
    bool_mod.object = obj2
    bool_mod.solver = 'FAST'
    bool_mod.operation = 'DIFFERENCE'

    O.object.modifier_apply(modifier=bool_mod.name)

def create_intersection(obj, inner):
    """Create a boolean intersection between two objects.

    Parameters
    ----------
    obj : bpy.types.Object
        The object for intersection.
    inner : bpy.types.Object
        The inner object.
    """
    inner.select_set(True)
    C.view_layer.objects.active = inner

    # Create new intersection modifier and apply it to the obj and intersection object
    # (inner) is the result of the initial intersection between the two objs
    bool_mod = C.active_object.modifiers.new(name="Boolean", type='BOOLEAN')
    bool_mod.object = obj
    bool_mod.solver = 'FAST'
    bool_mod.operation = 'INTERSECT'
    O.object.modifier_apply(modifier=bool_mod.name)

def create_material(mat_name, color, blend):
    """
    Create a material with the given name and color.

    Parameters
    ----------
    mat_name : str
        Name of the material.
    color : tuple
        HSV color tuple (hue, saturation, value, alpha).
    blend : str
        Blend mode for the material.
    """
    if mat_name in D.materials:
        print(f'Material {mat_name} already exists in this scene')
    else:
        new_material = D.materials.new(name=mat_name)
        new_material.use_nodes = True
        tree = new_material.node_tree
        nodes = tree.nodes
        links = tree.links

        for node in nodes:
            nodes.remove(node)

        principled = nodes.new(type='ShaderNodeBsdfPrincipled')
        material_output = nodes.new(type='ShaderNodeOutputMaterial')
        links.new(principled.outputs['BSDF'], material_output.inputs['Surface'])

        r, g, b = hsv_to_rgb(color[:3])
        principled.inputs['Base Color'].default_value = (r, g, b, 1.0)
        principled.inputs["Alpha"].default_value = color[3]

        new_material.blend_method = blend
        new_material.shadow_method = 'NONE'
        if blend == 'BLEND':
            new_material.use_backface_culling = True

def create_obj(data, fname):
    """
    Single method to complete reading a prediction from labeled data
    and to export it to a 3D object

    Parameters
    ----------
    data : numpy.ndarray
        Labeled data from .tif file.
    fname : str
        Filename to export 3D object to.
    """
    vals = np.unique(data)[1:]  # Exclude 0 from unique values
    vert_count = 0

    for i, val in enumerate(vals):
        last_filter = (data == val).astype(int)
        filt = data * last_filter
        verts, faces = extract_mesh(filt)
        mesh =  refine_mesh(verts,faces)
    
        export_to_obj(mesh.vertices, mesh.faces, i+1, vert_count, fname)
        vert_count += len(verts)

def create_overlap(gt, pred):
    """Handle the entire boolean operations to get a final product of an intersection object,
    and the two differences (false negative & false positive).

    Parameters
    ----------
    gt : bpy.types.Object
        The ground truth object.
    pred : bpy.types.Object
        The predicted object.
    """
    # Get the number of the current frame operated on and duplicate objs for temporary
    # working objects
    num = gt.name.split('_')[-1]
    inner = duplicate_object(gt, f'inner_{num}')
    temp = duplicate_object(pred, f'temp_{num}')

    # Perform operations needed for end product: starting with creating the true positives
    # by making intersection between prediction and ground-truth (inner at this stage is
    # the duplicate of ground-truth). After this, create difference of both objs with the
    # true positives and clean-up temporary working objects from scene.
    create_intersection(pred, inner)
    create_difference(pred, gt)
    create_difference(gt, temp)
    clean_duplicate_objs()

def display(viewer):
    """
    Display the rendered video in the Napari viewer.

    Parameters:
    -----------
    viewer : Napari viewer instance
        Napari viewer instance.
    """
    path = str(Path(SRC_VID / FILENAME))
    vr = VideoReaderNP(path)
    ratio = 1.5 * DATA1.shape[2] / vr.shape[2]
    
    # Add the image to the existing viewer
    viewer.add_image(vr, name='Loaded Video', scale=(1.0, ratio, ratio))

def downscale_image_by_sampling(image, factor):
    """
    Downscales the input image by sampling every nth pixel in each dimension.

    Parameters
    ----------
    image : numpy.ndarray
        Input image as a NumPy array.
    factor : int
        Factor by which to downscale (e.g., 2 for half, 4 for quarter).
    
    Returns
    -------
    numpy.ndarray
        Downscaled image as a NumPy array.
    """
    if image.ndim == 4:  # Multichannel image (e.g., RGB)
        return image[:, ::factor, ::factor, :]
    else:  # Grayscale image
        return image[:, ::factor, ::factor]

# def create_volume(data):
    # import open3d as o3d
    # pcd = o3d.geometry.PointCloud()
    # pcd.points = o3d.utility.Vector3dVector(data.astype(np.float64))
    # o3d.io.write_point_cloud("sync.ply", pcd)
    # # vertices = data
    # # print(data.shape)
    # # vertices = np.array([tuple(e) for e in vertices], dtype=[('x', '<f4'), ('y', '<f4'), ('z', '<f4')])
    # # el = PlyElement.describe(vertices, 'vertex',{'some_property': 'f8'},{'some_property': 'u4'})
    # # # text = True: save in ascii format
    # # PlyData([el], text=True).write("test.ply") 

def duplicate_object(obj, name):
    """
    Duplicate the object and rename the duplicated object.

    Parameters
    ----------
    obj : bpy.types.Object
        Object to duplicate.
    name : str
        Name for the duplicated object.
    
    Returns
    -------
    bpy.types.Object
        The duplicated object.
    """
    obj.select_set(True)
    C.view_layer.objects.active = obj
    
    # Duplicate the selected object
    O.object.duplicate()
    
    # Get the newly duplicated object
    duplicated_obj = C.active_object
    
    # Rename the duplicated object
    duplicated_obj.name = name
    duplicated_obj.data.name = name

    return duplicated_obj

def export_to_obj(vertices, faces, group, vert_count, fname):
    """
    Exports mesh data to an OBJ file.

    Parameters
    ----------
    vertices : numpy.ndarray
        Vertices of the submesh.
    faces : numpy.ndarray
        Faces of the submesh.
    group : int
        Group number for the submesh that we want to create now.
    vert_count : int
        Vertex count for the entire object until now.
    fname : str
        Name of the OBJ file to export.
    """
    if not os.path.exists(os.path.dirname(fname)):
        os.makedirs(os.path.dirname(fname))
        
    with open(fname, 'w' if not os.path.exists(fname) else 'a') as f:
        # Write group statement
        f.write(f"g {group}\n")

        # Write vertices
        for vertex in vertices:
            f.write(f"v {vertex[0]} {vertex[1]} {vertex[2]}\n")

        # Write faces
        for face in faces:
            f.write(f"f {' '.join([str(i + vert_count + 1) for i in face])}\n")

def extract_mesh(data):
    """
    Extracts mesh data from a segmented .tif file.

    Parameters
    ----------
    path : str
        Path to the segmented .tif file.
    
    Returns
    -------
    tuple
        Tuple containing vertices and faces of the extracted mesh.S
    """
    # Marching cubes algorithm to extract surfaces
    verts, faces, _, _ = measure.marching_cubes(data)
    return verts, faces

def filter_spots(spots, name, value, isabove):
    """
    Filter spots based on a certain condition.

    Parameters:
    Parameters
    ----------
    spots : pandas.DataFrame
        DataFrame containing spots data.
    name : str
        Feature name.
    value : float
        Threshold value.
    isabove : bool
        If True, filter spots above the threshold, otherwise below.
    
    Returns
    -------
    pandas.DataFrame
        Filtered DataFrame.
    """
    if isabove:
        spots = spots[spots[name] > value]
    else:
        spots = spots[spots[name] < value]
    return spots

def find_collection(name):
    """
    Function to find a collection in the scene, or raise an error
    if this collection is not found.

    Parameters
    ----------
    name : str
        Name of the desired collection.
    """
    coll = D.collections.get(name)
    if coll is not None:
        return coll
    else:
        print(f'Collection {name} was not found')
        raise CollectionNotFoundException()

def find_material(mat_name):
    """
    Find a material by its name.

    Parameters
    ----------
    mat_name : str
        Name of the material.

    Returns
    -------
    bpy.types.Material
        The material, if present, otherwise raise error.
    """
    mat = D.materials.get(mat_name)
    if mat is not None:
        return mat
    else:
        print(f'Material {mat_name} not found')
        raise MaterialNotFoundException()

def find_object(name):
    """"
    Function to find an object in the scene, or raise an error
    if this object is not found.

    Parameters
    ----------
    name : str
        Name of the desired object.
    """
    obj = D.objects.get(name)
    if obj is not None:
        return obj
    else:
        print(f'Object {name} was not found')
        raise ObjectNotFoundException()

def frequency(values):
    """
    Returns unique values in the array and their frequency of occurrence.

    Parameters
    ----------
    values : numpy.ndarray
        1D array of values.

    Returns
    -------
    tuple
        A tuple containing frequencies and corresponding labels.
    """
    lbls = np.unique(values)
    freq = np.array([np.sum(values == l) for l in lbls])
    return freq, lbls

def get_all_coords(nuclei):
    """
    Function to get coordinates for all objects in a frame.
    This code facilitates the minimum cost Hungarian algorithm implementation.

    Parameters
    ----------
    nuclei : list
        List of objects in the scene/frame.
    
    Returns
    -------
    list
        A list of coordinates for all objects in the specified collection.
    """
    all_coords = []
    for obj in nuclei:
        coords = (obj.location.x, obj.location.y, obj.location.z)

        # Format to set amount of decimals to minimize rounding errors in comparison
        formatted_coords = tuple(round(coord, 2) for coord in coords)
        all_coords.append(formatted_coords)
    return all_coords

def get_track(df, obj_id):
    """
    Finds the track ID of a given ID in the DataFrame.

    Parameters
    ----------
    df : pandas.DataFrame
        The DataFrame containing the track data.
    obj_id : str
        The ID for which to find the child.

    Returns
    -------
    str or None
        The track ID if found, otherwise None.
    """
    # First check if the original object has been found
    if 'Not_Found' in obj_id:
        print('Not found nucleus can not have a track')
        return None
    # Get the row corresponding to the desired id
    row = df.loc[df['id'] == int(obj_id)]

    if not row.empty:
        track = row['track'].iat[0]
        return track
    else:
        print(f"ID {obj_id} not found in the DataFrame.")
    return None

def hide_in_render(obj, count, length):
    """
    Hide an object in renders for a specific range of frames.

    Parameters
    ----------
    obj : bpy.types.Object
        The object to hide in renders.
    count : int
        The frame count to determine when to hide and show the object.
    length : int
        The length of time the object should be rendered.
    """
    # Show object and set keyframe
    obj.hide_render = False
    obj.keyframe_insert(data_path='hide_render', frame=length * count)

    # Hide object and set keyframe after specified time
    obj.hide_render = True
    obj.keyframe_insert(data_path='hide_render', frame=length + length * count)
    
    # Hide all child objects initially to be morphed by ancestors
    if count > 0:
        # Set keyframes for hide_render property
        obj.hide_render = True
        obj.keyframe_insert(data_path='hide_render', frame=0)

def hide_legend_items(count):
    """
    Function that is called to hide unused legend items from render.

    Parameters
    ----------
    count : int
        Number of legend items that will be used.
    """
    # Run loop that removes used legend items from the {1,2,3} domain.
    for i in range(1+count,4):
        try:
            mat_obj = find_object(f'Mat_{i}')
            mat_obj.hide_render = True
            mat_text = find_object(f'Mat_{i}_name')
            mat_text.hide_render = True
        except ObjectNotFoundException:
            pass

def hsv_to_rgb(color):
    """
    Convert HSV color representation to RGB.

    Parameters
    ----------
    color : tuple
        A tuple containing hue, saturation, and value components.

    Returns
    -------
    tuple
        RGB color tuple.
    """
    h, s, v = color
    return colorsys.hsv_to_rgb(h, s, v)

def import_object(name, data, use_groups, index):
    """
    Create an object in the scene.

    Parameters
    ----------
    name : str
        Name of the object.
    data : numpy.ndarray
        Data to create the object from.
    use_groups : bool
        Whether to use split groups, creates a single or multiple objects.
    index : int
        Index of the current frame.

    Returns
    -------
    list
        A list of objects created in the scene.
    """
    try:
        stdout = io.StringIO()
        with redirect_stdout(stdout):
            obj = find_object(name)
        return obj
    except ObjectNotFoundException:
        dim_name = name + '.obj'
        obj_src = s.SRC_3D / dim_name
        # First check if object is already made as 3D object in default directory,
        # if so, import without additional computations.
        if obj_src.is_file():
            print(f'Note: Object {str(obj_src)} already exists as 3D object, importing this.\nIf other 3D object is desired, delete existing one, or change its name.')
            stdout = io.StringIO()
            with redirect_stdout(stdout):
                O.wm.obj_import(filepath=str(obj_src), use_split_groups= use_groups)                 
        # If there is not yet a 3D object of this object, create it from the source data.
        else:
            try:
                data = clean_data(data,index)
                if s.CREATE_GROUND_TRUTH:
                    create_contour_obj(data, obj_src)
                else:
                    create_obj(data,obj_src)
                stdout = io.StringIO()
                with redirect_stdout(stdout):
                    O.wm.obj_import(filepath=str(obj_src), use_split_groups= use_groups)
            except IndexError:
                print('Index error, no object created')
                return
            
        imported_object = C.selected_objects[0]
        num = name.rsplit('_', 1)[-1]
        # Rename the object
        if s.PRED_NAME is not None and s.PRED_NAME in name and not use_groups:
            imported_object.name = f'pred_{num}'
        elif s.GT_NAME in name and not use_groups:
            imported_object.name = f'gt_{num}'

    return [obj for obj in C.selected_objects if obj.type == 'MESH']

def iou(mask1, mask2):
    """
    Calculates Intersection over Union (IoU) score.

    Parameters
    ----------
    mask1 : numpy.ndarray
        First binary mask.
    mask2 : numpy.ndarray
        Second binary mask.

    Returns
    -------
    float
        Intersection over Union (IoU) score.
    """
    mask1 = mask1[:, :, :, 0] if mask1.shape[-1] == 3 else mask1
    mask2 = mask2[:, :, :, 0] if mask2.shape[-1] == 3 else mask2

    if mask1.shape != mask2.shape:
        raise ValueError("Input masks must have the same size.")

    intersection = np.logical_and(mask1, mask2)
    union = np.logical_or(mask1, mask2)

    iou_val = np.sum(intersection) / np.sum(union)
    return iou_val

def metrics_dictionary(img1, img2):
    """
    Calculates various metrics based on two images.

    Parameters
    ----------
    img1 : numpy.ndarray
        First input image in numpy array format.
    img2 : numpy.ndarray
        Second input image in numpy array format.

    Returns
    -------
    dict
        Dictionary containing calculated metrics.
    """
    metrics = {}
    
    truth, count_truth = label(img1)
    pred, count_pred = label(img2)

    are, prec, rec = calculate_are(truth, pred)

    f1 = 1 - are

    metrics['ji'] = calculate_jaccard(truth, pred)
    metrics['ji_vals'] = compare_labeled_volumes(truth, pred)
    metrics['nuclei_truth'] = count_truth
    metrics['nuclei_pred'] = count_pred
    metrics['iou'] = iou(truth, pred)
    metrics['precision'] = prec
    metrics['recall'] = rec
    metrics['f1'] = f1

    return metrics

def min_cost(df, list_points):
    """
    Helper function to calculate a matrix of all distances for coordinates of nuclei 
    in Blender to coordinates found in the tracking file. Minimum cost assignment 
    is done via the Hungarian algorithm and returned.

    Parameters
    ----------
    df : pandas.DataFrame
        DataFrame containing the coordinates of nuclei in Blender.
    list_points : list
        List of coordinates found in the tracking file.
    
    Returns
    -------
    dict
        A dictionary mapping coordinates to their corresponding IDs.
    """
    # Calculate distance matrix
    distance_matrix = np.array([[calculate_distance(df.loc[i, ['x', 'y', 'z']], point) for point in list_points] for i in df.index])
    # Apply Hungarian algorithm
    _, col_ind = linear_sum_assignment(distance_matrix)

    # Retrieve results
    optimal_assignment = {list_points[j]: df.loc[i, 'id'] for i, j in enumerate(col_ind, start=df.index.min())}
    return optimal_assignment

def napari_get_reader(path):
    """A basic implementation of a Reader contribution.

    Parameters
    ----------
    path : str or list of str
        Path to file, or list of paths.

    Returns
    -------
    function or None
        If the path is a recognized format, return a function that accepts the
        same path or list of paths, and returns a list of layer data tuples.
    """
    if isinstance(path, list):
        # reader plugins may be handed single path, or a list of paths.
        # if it is a list, it is assumed to be an image stack...
        # so we are only going to look at the first file.
        path = path[0]

    # if we know we cannot read the file, we immediately return None.
    if not path.endswith(".npy"):
        return None

    # otherwise we return the *function* that can read ``path``.
    return reader_function

def print_globals():
    """
    Print global variables for debugging purposes.
    """
    global_vars = globals()
    for k, v in global_vars.items():
        print(f"{k}: {v}")

def read_trackmate():
    """
    Reads track data from a TrackMate XML file.

    Returns
    -------
    pandas.DataFrame
        DataFrame containing the track data.
    """
    # Import track data using a custom function (trackmate_peak_import)
    # code gathered (but bugfixed) from https://github.com/hadim/pytrackmate
    df = trackmate_peak_import(s.TRACK, True)
    
    return df

def reader_function(path):
    """Take a path or list of paths and return a list of LayerData tuples.

    Readers are expected to return data as a list of tuples, where each tuple
    is (data, [add_kwargs, [layer_type]]), "add_kwargs" and "layer_type" are
    both optional.

    Parameters
    ----------
    path : str or list of str
        Path to file, or list of paths.

    Returns
    -------
    layer_data : list of tuples
        A list of LayerData tuples where each tuple in the list contains
        (data, metadata, layer_type), where data is a numpy array, metadata is
        a dict of keyword arguments for the corresponding viewer.add_* method
        in napari, and layer_type is a lower-case string naming the type of
        layer. Both "meta", and "layer_type" are optional. napari will
        default to layer_type=="image" if not provided
    """
    # handle both a string and a list of strings
    paths = [path] if isinstance(path, str) else path
    # load all files into array
    arrays = [np.load(_path) for _path in paths]
    # stack arrays into single array
    data = np.squeeze(np.stack(arrays))

    # optional kwargs for the corresponding viewer.add_* method
    add_kwargs = {}

    layer_type = "image"  # optional, default is "image"
    return [(data, add_kwargs, layer_type)]

def recalculate_text(scene):
    """
    Recalculate text for the current frame.

    Parameters
    ----------
    scene : bpy.types.Scene
        The scene in which rendering is performed.
    """
    if s.LENGTH > 0:
        frame = s.START + C.scene.frame_current // s.LENGTH
    else:
        frame = s.START + C.scene.frame_current
    change_text(f'Frame: {frame}','Frame')
    try:
                # change_text(f"Metrics: \n# nuclei GT: {s.METRICS[frame]['nuclei_truth']}\n# nuclei predicted: {s.METRICS[frame]['nuclei_pred']}\n\nIoU: {s.METRICS[frame]['iou']:.2f}\nJI: {s.METRICS[frame]['ji']:.2f}\nPrecision: {s.METRICS[frame]['precision']:.2f}\nRecall: {s.METRICS[frame]['recall']:.2f}\nF1-score: {s.METRICS[frame]['f1']:.2f}",'Metrics')
        change_text(f"Metrics: \n\nIoU: {s.METRICS[frame]['iou']:.2f}\nJI: {s.METRICS[frame]['ji']:.2f}\nPrecision: {s.METRICS[frame]['precision']:.2f}\nRecall: {s.METRICS[frame]['recall']:.2f}\nF1-score: {s.METRICS[frame]['f1']:.2f}",'Metrics')
    except KeyError:
        pass

def refine_mesh(verts, faces):
    """
    Refines the mesh by applying Laplacian smoothing.

    Parameters
    ----------
    verts : numpy.ndarray
        Vertices of the mesh.
    faces : numpy.ndarray
        Faces of the mesh.
    
    Returns
    -------
    trimesh.Trimesh
        The smoothed mesh.
    """
    # Convert to trimesh object
    mesh = trimesh.Trimesh(vertices=verts, faces=faces)
    
    # Apply Laplacian smoothing
    smooth_mesh = trimesh.smoothing.filter_laplacian(mesh, iterations=5)

    return smooth_mesh

def register():
    """
    Register frame change handler.
    """
    A.handlers.frame_change_post.append(recalculate_text)

def remesh_objects(obj, vert_count):
    """
    Helper function to set the amount of vertices in an object to vert_count.
    This code is useful to get a uniform vertex count over objects that have to
    morph into one another.

    Parameters
    ----------
    obj : bpy.types.Object
        The object whose vertices need to be adjusted.
    vert_count : int
        The desired number of vertices for the object.
    """
   
    if len(obj.data.vertices) > vert_count:
        # Apply the decimate modifier to ensure a certain number of vertices
        O.object.mode_set(mode='OBJECT')

        decimate_modifier = obj.modifiers.new(name="Decimate", type='DECIMATE')
        decimate_modifier.ratio = vert_count / len(obj.data.vertices)

        obj.select_set(True)
        C.view_layer.objects.active = obj
        
        # Apply the modifier
        O.object.modifier_apply(modifier="Decimate")
        
    if len(obj.data.vertices) < vert_count:
        obj.select_set(True)
        C.view_layer.objects.active = obj

        mesh = obj.data

        subdiv_num = int(math.ceil((vert_count-len(obj.data.vertices))/3))
        step = int(math.floor(len(obj.data.vertices)/subdiv_num))
        for i in range(0,subdiv_num):
            mesh.polygons[i*step].select = True
            
        O.object.mode_set(mode='EDIT')
        O.mesh.subdivide()

        O.object.mode_set(mode='OBJECT')
        O.object.select_all(action='DESELECT')

    # This method may end on a couple vertices too few or too much, recursive
    # calling fixes this problem.
    while len(obj.data.vertices) != vert_count:
        remesh_objects(obj, vert_count)
    
    # Shade object smooth for cleaner render object
    # Commented out for aesthetic reasons, for now.
    # with C.temp_override(selected_editable_objects=[obj]):
    #     O.object.shade_smooth()

def rename_nuclei(objects, data_frame, i):
    """
    Scale, rotate, create submeshes and rename all nuclei from original 3D object.

    Parameters
    ----------
    objects: list
        List of objects to manipulate.
    data_frame: pandas.DataFrame
        DataFrame containing the tracking data.
    i: int 
        Index of the current frame.
    """
    nuclei = objects
    filt_nuclei = []

    for ob in nuclei:
        if len(ob.data.vertices) < 50:
            # Select and delete the object
            D.objects.remove(ob, do_unlink=True)
        else:  
            filt_nuclei.append(ob)
                    
    # get all coordinates of nuclei objects and calculate the minimum cost
    # id allocation using the Hungarian algorithm
    all_coords = get_all_coords(filt_nuclei)
    min_cost_mat = min_cost(data_frame.loc[data_frame['frame'] == i-1], all_coords)
    
    # go over all objects to properties and manipulate them
    for obj in filt_nuclei:
        # set all objects to have the same vertex count
        remesh_objects(obj, 1000)

        # rename objects in Blender to their id as given by tracking file
        rename_obj(obj, min_cost_mat)

def rename_obj(obj, min_cost_mat):
    """
    Helper function to rename a nuclei to fit their match as calculated by the 
    Hungarian algorithm. This method requires the min_cost_mat as calculated by the
    min_cost() function and renames the object to its corresponding ID in the tracking file.

    Parameters
    ---------- 
    obj : bpy.types.Object
        The object to be renamed.
    min_cost_mat : dict
        A dictionary mapping coordinates to their corresponding IDs.
    """
    coords = (obj.location.x, obj.location.y, obj.location.z)
    formatted_coords = tuple(round(coord, 2) for coord in coords)
    
    # Rename objects to their correct ID as found in the mastodon track file.
    # Try-catch can be invoked when the tracking file or mesh prediction mispredicts,
    # which is not uncommon in scenarios like cell division.
    try:
        obj_id = min_cost_mat[formatted_coords]
    except KeyError:
        obj_id = 'Not_Found'

    obj.name = f"{obj_id}"

def render():
    """
    Render the scene and generate the output video.
    """
    renderer = RenderProgress()
    renderer.render_samples()
    renderer.render_video()

def rgb_to_hsv(color):
    """
    Convert RGB color representation to HSV.

    Parameters
    ----------
    color : tuple
        A tuple containing red, green, and blue components.

    Returns
    -------
    tuple
        HSV color tuple.
    """
    r, g, b = color
    return colorsys.rgb_to_hsv(r, g, b)

def scale_object(obj):
    """
    Scale, rotate and set object to the origin of the scene.

    Parameters
    ----------
    obj : bpy.types.Object
        Object to manipulate.
    """
    C.view_layer.objects.active = obj
    ob = C.selected_objects[-1]
    # Set rotation of object to 0 degrees on all axes. Importing an image is often
    # done using an automatic 90 degree x-axis rotation, which hampers the renaming of nuclei.
    ob.rotation_euler = ( 0.0, 0.0, 0.0)
    # Set the origin to geometry, this places the object to 0,0,0 in the scene
    O.object.origin_set(type='GEOMETRY_ORIGIN')
    # Scale object to fit the camera of the scene.
    ob.scale = ( 0.045, 0.045, 0.045 )

def set_gradient_color(mode):
    """
    In the Gradient visualization mode, the Gradient material is complicated and 
    always already included in the scene. Set this material to the legend, and
    change needed text to fit this.

    Parameters
    ----------
    mode : str
        The visualization mode that is set to the program.
    """
    match mode:
        case 'Gradient': 
            num = 2
        case 'Colored Nuclei':
            num = 3
        case  _:
            num = 0
    try:
        mat_obj = find_object(f'Mat_{num}')
        mat_obj.hide_render = False
        set_material(mat_obj, 'Gradient')
        mat_text = find_object(f'Mat_{num}_name')
        mat_text.hide_render = False
        change_text('Prediction',f'Mat_{2}_name')
    except ObjectNotFoundException:
        pass

def set_hue(obj, metric):
    """
    Set the hue of an object based on a metric value.

    Parameters
    ----------
    obj : bpy.types.Object
        The object to set the hue.
    metric : float
        Metric value to determine the hue [0..1].
    """
    try:
        hue = clamp_hue(metric)
        create_material(f'Base_hue_{hue:.2f}', (hue, 1, 0.43, 1.0), 'OPAQUE')
        set_material(obj, f'Base_hue_{hue:.2f}')
    except TypeError:
        print(f'Could not create material with hue {hue}')

def set_indicator(metric, frame=None):
    """
    Set the indicator based on metric value.

    Parameters
    ----------
    metric : float
        The metric value.
    frame : int, optional
        The frame to set the keyframe.
    """
    try:
        ind = find_object('Indicator')
        # Get, and change, the z-coordinate of the metric indicator object. Since it
        # is clamped to a path, the clamped value will set the indicator to the relative
        # position indicating the metric value on the metric bar.
        z = clamp_indicator(metric)
        ind.delta_location = (0, 0, z)
        # If a frame is given, set it as keyframe for when to reach that value
        if frame is not None:
            ind.keyframe_insert(data_path="delta_location", index=-1, frame=frame, group="Delta Transform")
    except ObjectNotFoundException:
        print('Setting metric indicator suspended')

def set_interpolation(obj, style):
    """
    Set interpolation style for object animation.

    Parameters
    ----------
    obj : bpy.types.Object 
        The object to set interpolation for.
    style : str
        The style of interpolation, linear is used in most use cases.
    """
    fcurves = obj.animation_data.action.fcurves
    for fcurve in fcurves:
        for kf in fcurve.keyframe_points:
            kf.interpolation = style

def set_legend_color(mat_name, num):
    """
    Function to set a legend item to be a color, and display it's correct name.

    Parameters
    ----------
    mat_name : str
        Material name that will be set to the text.
    num : int
        Enumerator of what legend item should be set.
    """
    # The legend items are named 'Mat_{1,2,3}', and will be changed in sequence
    change_text(mat_name, f'Mat_{num}_name')
    try:
        mat_obj = find_object(f'Mat_{num}')
        mat_obj.hide_render = False
        mat_text = find_object(f'Mat_{num}_name')
        mat_text.hide_render = False
        set_material(mat_obj, mat_name)
    except ObjectNotFoundException:
        pass  

def set_material(obj, mat_name):
    """
    Set material to an object.

    Parameters
    ----------
    obj : bpy.types.Object
        Object to set the material to.
    mat_name : str
        Material to set. If None, it removes any existing materials.
    """
    try:
        mat = find_material(mat_name)
        if obj.data.materials:
            obj.data.materials[0] = mat
        else:
            obj.data.materials.append(mat)
    except MaterialNotFoundException:
        pass

def set_origin_nuclei(nuclei, name, i):
    """
    Sets the origin point for a group of nuclei objects and creates a parent empty object.

    This function calculates the average location of a group of nuclei objects, sets their origins
    to the average location, and then creates a new empty object to parent all nuclei objects to.

    Parameters
    ----------
    nuclei : list
        A list of Blender objects representing the nuclei.
    name : str
        The name of the empty object that will be created to parent all nuclei objects.
    """
    index = i - s.START

    acc_location = Vector((0, 0, 0))
    for obj in nuclei:
        O.object.select_all(action='DESELECT')
        obj.select_set(True)
        obj['Group_ID'] = obj.name.split('.')[0]
        obj.rotation_euler = ( 0.0, 0.0, 0.0)
        O.object.origin_set(type='ORIGIN_GEOMETRY')
        acc_location += obj.location

    # Calculate average location (center of mass)
    num_objects = len(nuclei)
    if num_objects > 0:
        average_location = acc_location / num_objects
    else:
        print("No mesh objects selected.")
    
    # Create a new empty object to parent all imported meshes
    O.object.empty_add()
    empty_object = C.object
    empty_object.location = average_location
    
    # Parent all mesh objects to the empty object
    for obj in nuclei:
        O.object.select_all(action='DESELECT')
        obj.select_set(True)
        O.object.parent_no_inverse_set(keep_transform=True)
    
    empty_object.name = name
    empty_object.location = Vector((0,0,0))
    empty_object.scale = ( 0.045, 0.045, 0.045 )   

    set_rotation(empty_object, s.LENGTH * index, s.LENGTH + s.LENGTH * index)
    set_interpolation(empty_object,'LINEAR')

def set_parent(parent, child):
    """
    Sets a parent-child relationship between two objects. Useful for 
    getting multiple objects to follow the same keyframes, or hiding
    them in the same line of code.

    Parameters
    ----------
    parent : bpy.types.Object
        The parent object.
    child : bpy.types.Object
        The child object.
    """
    child.select_set(True)
    parent.select_set(True)
    O.object.parent_set(type='OBJECT', keep_transform=True)

def set_rotation(obj, start, end):
    """
    Sets keyframes for rotation of an object.

    Parameters
    ----------
    obj : bpy.types.Object
        The object to rotate.
    start : int
        Frame number on which to start the rotation.
    end : int
        Frame number on which to end the rotation.
    """
    # Set keyframes for rotation
    obj.rotation_euler.z = math.radians(0)  # Set rotation to 0 degrees in radians
    obj.keyframe_insert(data_path="rotation_euler", index=2, frame=start)

    obj.rotation_euler.z = math.radians(360)  # Set rotation to 360 degrees in radians
    obj.keyframe_insert(data_path="rotation_euler", index=2, frame=end)

def set_shape_keys(obj, target, index):
    """
    Helper function to set the frames at which the shape keys need to be a certain value.

    Parameters
    ----------
    obj : bpy.types.Object
        The object to morph.
    target : bpy.types.Object
        The object to morph into.
    index : int
        The index of the current frame.
    """
    start_loc = obj.location
    start, end = s.LENGTH * index, s.LENGTH + s.LENGTH * index
    obj.keyframe_insert(data_path="location", frame=start)
    obj.location = target.location
    obj.keyframe_insert(data_path="location", frame=end)
    obj.location = start_loc

    shape_key = obj.data.shape_keys.key_blocks["Morph"]
    shape_key.keyframe_insert("value", frame=start)
    shape_key.value = 1
    shape_key.keyframe_insert("value", frame=end)

def set_values(dict, viewer, mode):
    """
    Set global values from a dictionary and initialize processing.

    Parameters:
    -----------
    dict : dict
        Dictionary containing configuration values.
    viewer : Napari viewer instance
        Napari viewer instance.
    mode : str
        Visualization mode.
    """
    global TRACK, START, END, LENGTH, SAMPLES, METRICS, E_FRAME, SRC_VID, FILENAME, DATA1, DATA2, GT_NAME, PRED_NAME, SHOW_TEXT, HUE_METRIC
    for k, v in dict.items():
        globals()[k] = v
    if DATA1 is not None:
        GT_NAME += "_"
        if DATA2 is not None:
            PRED_NAME += "_"
            if GT_NAME == PRED_NAME:
                PRED_NAME = "pred-" + PRED_NAME
            # If both images are there, calculate metrics of comparison.
            for i in range(START, END + 1):
                clean_data1 = clean_data(DATA1, i)
                clean_data2 = clean_data(DATA2, i)
                METRICS[i] = metrics_dictionary(clean_data1, clean_data2)
            # Scale the data down for quicker computation, and to keep objects
            # in the frame.
            scale = max(math.floor(DATA2.shape[1] / 150), 1)
            DATA2 = downscale_image_by_sampling(DATA2, scale)
        scale = max(math.floor(DATA1.shape[1] / 150), 1)
        DATA1 = downscale_image_by_sampling(DATA1, scale)
        E_FRAME = (END - START) * LENGTH + LENGTH
        visualize(viewer, mode)
    # print_globals()

def track_ancestor(df, obj_id):
    """
    Finds the ancestor ID of a given ID in the DataFrame.

    Parameters
    ----------
    df : pandas.DataFrame
        The DataFrame containing the track data.
    obj_id : str
        The ID for which to find the ancestor.

    Returns
    -------
    str or None
        The ancestor ID if found, otherwise None.
    """
    # First check if the original object has been found
    if 'Not_Found' in obj_id:
        print('Not found nucleus can not have an ancestor')
        return None
    # Get the row corresponding to the desired id
    row = df.loc[df['id'] == int(obj_id)]

    if not row.empty:
        track = row['track'].iat[0]
        frame = row['frame'].iat[0]
        # Filter rows where track and frame match the desired values
        ancestor_row = df[(df['track'] == track) & (df['frame'] == frame - 1)]

        if not ancestor_row.empty:
            ancestor_id = ancestor_row['id'].iat[0]
            return ancestor_id
        print(f"No ancestor found for {obj_id}.")
    else:
        print(f"ID {obj_id} not found in the DataFrame.")
    return None

def track_child(df, obj_id):
    """
    Finds the child ID of a given ID in the DataFrame.

    Parameters
    ----------
    df : pandas.DataFrame
        The DataFrame containing the track data.
    obj_id : str
        The ID for which to find the child.

    Returns
    -------
    str or None
        The child ID if found, otherwise None.
    """
    # First check if the original object has been found
    if 'Not_Found' in obj_id:
        print('Not found nucleus can not have a child')
        return None
    # Get the row corresponding to the desired id
    row = df.loc[df['id'] == int(obj_id)]

    if not row.empty:
        track = row['track'].iat[0]
        frame = row['frame'].iat[0]
        # Filter rows where track and frame match the desired values
        child_row = df[(df['track'] == track) & (df['frame'] == frame + 1)]

        if not child_row.empty:
            child_id = child_row['id'].iat[0]
            return child_id
        else:
            print(f"No child found for {obj_id}.")
    else:
        print(f"ID {obj_id} not found in the DataFrame.")
    return None

def trackmate_peak_import(trackmate_xml_path, get_tracks=False):
    """
    Import detected peaks with TrackMate Fiji plugin.

    Parameters
    ----------
    trackmate_xml_path : str
        TrackMate XML file path.
    get_tracks : bool
        Whether to include tracks.
    
    Returns
    -------
    pandas.DataFrame
        DataFrame containing the imported peak data.
    """
    # Parse the XML file
    root = et.fromstring(open(trackmate_xml_path).read())

    # Define object labels mapping [CHANGED] to fit own format
    object_labels = {
        'FRAME': 'frame',
        'POSITION_X': 'x',
        'POSITION_Y': 'y',
        'POSITION_Z': 'z',
        'ID': 'id',
    }

    # Extract features
    features = root.find('Model').find('FeatureDeclarations').find('SpotFeatures')
    
    # Original line [not functioning]:
    # features = [c.get('feature') for c in features.getchildren()] + ['ID']

    features = [c.get('feature') for c in list(features)] + ['ID']

    # Extract spots
    spots = root.find('Model').find('AllSpots')
    trajs = pd.DataFrame([])
    objects = []
    for frame in spots.findall('SpotsInFrame'):
        for spot in frame.findall('Spot'):
            single_object = []
            for label in features:
                single_object.append(spot.get(label))
            objects.append(single_object)

    # Create DataFrame from extracted objects
    trajs = pd.DataFrame(objects, columns=features)
    trajs = trajs.astype(float)

    # Apply initial filtering
    initial_filter = root.find("Settings").find("InitialSpotFilter")
    trajs = filter_spots(trajs,
                         name=initial_filter.get('feature'),
                         value=float(initial_filter.get('value')),
                         isabove=True if initial_filter.get('isabove') == 'true' else False)

    # Apply filters
    spot_filters = root.find("Settings").find("SpotFilterCollection")
    for spot_filter in spot_filters.findall('Filter'):
        trajs = filter_spots(trajs,
                             name=spot_filter.get('feature'),
                             value=float(spot_filter.get('value')),
                             isabove=True if spot_filter.get('isabove') == 'true' else False)

    # Rename columns and add track numbers
    trajs = trajs.loc[:, object_labels.keys()]
    trajs.columns = [object_labels[k] for k in object_labels.keys()]
    trajs['track'] = np.arange(trajs.shape[0])

    # Get tracks if specified
    if get_tracks:
        filtered_track_ids = [int(track.get('TRACK_ID')) for track in root.find('Model').find('FilteredTracks').findall('TrackID')]
        label_id = 0
        trajs['track'] = np.nan
        tracks = root.find('Model').find('AllTracks')
        for track in tracks.findall('Track'):
            track_id = int(track.get("TRACK_ID"))
            if track_id in filtered_track_ids:
                spot_ids = [(edge.get('SPOT_SOURCE_ID'), edge.get('SPOT_TARGET_ID'), edge.get('EDGE_TIME')) for edge in track.findall('Edge')]
                spot_ids = np.array(spot_ids).astype('float')[:, :2]
                spot_ids = set(spot_ids.flatten())
                trajs.loc[trajs["id"].isin(spot_ids), "track"] = label_id
                label_id += 1
        single_track = trajs.loc[trajs["track"].isnull()]
        trajs.loc[trajs["track"].isnull(), "track"] = label_id + np.arange(0, len(single_track))

    # Convert certain columns to integer
    for row in ['frame', 'id', 'track']:
        trajs[row] = trajs[row].astype(int)

    return trajs

def unregister():
    """
    Unregister frame change handler if it exists.
    """
    try:
        A.handlers.frame_change_post.remove(recalculate_text)
    except ValueError:
        print("Handler recalculate_text not found in frame_change_post, skipping removal.")

def visualize(viewer, mode):
    """
    Visualize the data in Blender based on the selected mode.

    Parameters:
    -----------
    viewer : Napari viewer instance
        Napari viewer instance.
    mode : str
        Visualization mode.
    """
    # Revert the main file to the original state, this combats not being able
    # to run the script multiple times in a row.
    O.wm.revert_mainfile()
    match mode:
        case MODES.Transparant:
            tr = Transparant()
            tr.visualize()
        case MODES.Gradient:
            gr = Gradient()
            gr.visualize()
        case MODES.Tracked:
            tr = Tracked()
            tr.visualize()
        case MODES.Timelapse:
            tl = Timelapse()
            tl.visualize()
        case MODES.Microscopy:
            mc = Microscopy()
            mc.visualize()
        case _:
            print("Invalid mode")
    render()
    display(viewer)

def write_multiple(path: str, data: List[FullLayerData]) -> List[str]:
    """Writes multiple layers of different types.
    
    Parameters
    ----------
    path : str
        A string path indicating where to save the data file(s).
    data : A list of layer tuples.
        Tuples contain three elements: (data, meta, layer_type)
        `data` is the layer data
        `meta` is a dictionary containing all other metadata attributes
        from the napari layer (excluding the `.data` layer attribute).
        `layer_type` is a string, eg: "image", "labels", "surface", etc.

    Returns
    -------
    [path] : A list containing (potentially multiple) string paths to the saved file(s).
    """

    # implement your writer logic here ...

    # return path to any file(s) that were successfully written
    return [path]

def write_single_image(path: str, data: Any, meta: dict) -> List[str]:
    """Writes a single image layer.
    
    Parameters
    ----------
    path : str
        A string path indicating where to save the image file.
    data : The layer data
        The `.data` attribute from the napari layer.
    meta : dict
        A dictionary containing all other attributes from the napari layer
        (excluding the `.data` layer attribute).

    Returns
    -------
    [path] : A list containing the string path to the saved file.
    """

    # implement your writer logic here ...

    # return path to any file(s) that were successfully written
    return [path]

class CollectionNotFoundException(Exception):
    """
    Custom exception for when a collection is not found in the scene.
    """
    pass

class Gradient(Mode):
    """
    Mode that compares predictions by coloring the predicted organoid. End-result is a prediction
    colored using a specific metric, and the volume being transparent.
    """
    def init_dicts(self):
        """
        Mode specific objects/materials needed to be created or changed.
        """
        dicts = {}
        dicts['texts'] = {}
        dicts['materials'] = {'Ground-Truth': (0.0, 0.0, 0.8, 0.3, 'BLEND')}
        if s.SHOW_TEXT:
            dicts['hide_obj'] = {'Gradient_Bar': False, 'Indicator': False, 'Metrics': True}
        else:
            dicts['hide_obj'] = {'Gradient_Bar': False, 'Indicator': False, 'Metrics': True, 'Frame': True} 
        dicts['objs'] = {s.GT_NAME: s.DATA1, s.PRED_NAME: s.DATA2}
        return dicts
            
    def set_scene(self, i):
        """
        Mode specific code for importing each frame.

        Parameters
        ----------
        i : int
            Enumerator of current frame.
        """
        # Compute the relative index of the current frame.
        index = i - s.START

        # Compute metrics for the current frame objects.
        try:
            pred_obj = find_object(f'pred_{i}')
            gt_obj = find_object(f'gt_{i}')
            
            set_hue(pred_obj, s.METRICS[i][s.HUE_METRIC.value[0]])
            set_material(gt_obj, 'Ground-Truth')
        except ObjectNotFoundException:
            print(f'Setting up frame {i} suspended')
        set_indicator(s.METRICS[i][s.HUE_METRIC.value[0]], s.LENGTH * index)
        change_text(s.HUE_METRIC.value[1], 'Metric')
        set_gradient_color('Gradient')
        super().set_scene(i)

    def import_obj(self, obj_name, data, i):
        """
        Mode specific way of importing organoid objects.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        data : Any
            Data source for the object.
        i : int
            Index of current frame needed to be imported.
        """
        nuclei = import_object(obj_name, data, False, i)
        set_origin_nuclei(nuclei, obj_name, i)

    def post_import(self):
        """
        Mode specific code to be run after importing all frames.
        """
        try:
            indicator = find_object('Indicator')
            mat_obj = find_object('Mat_2')
            mat_name = find_object('Mat_2_name')
            mat_obj.hide_render = False
            mat_name.hide_render = False            
        except ObjectNotFoundException:
            print('Setting indicator object suspended')
        set_interpolation(indicator, 'CONSTANT')

class MODES(Enum):
    """
    Enumeration for different modes of visualization.

    Attributes:
    Transparant: Mode that compares predictions by Boolean operations.
    Gradient: Mode that compares predictions by coloring the predicted organoid.
    Tracked: Mode that tracks individual organoids over time.
    Timelapse: Mode that shows the movement of organoids over time.
    """
    Transparant = 'transparant'
    Microscopy = 'microscopy'
    Gradient = 'gradient'
    Tracked = 'tracked'
    Timelapse = 'timelapse'

class MaterialNotFoundException(Exception):
    """
    Custom exception for material not found in the scene.
    """
    pass

class Microscopy(Mode):
    """
    Mode that compares raw microscopy data with a prediction. End-result is an transparant overlay
    containing microscopy data, and an object that contains the segmentation. The former is automated
    but is not a perfect representation of the nuclei in space.
    """
    def init_dicts(self):
        """
        Mode specific objects/materials needed to be created or changed.
        """
        dicts = {}
        dicts['texts'] = {}
        dicts['materials'] = {'Ground-Truth': (0.0, 0.0, 0.8, 0.15, 'BLEND'), 'Prediction': (0.6, 1.0, 0.8, 1.0, 'OPAQUE')}
        if s.SHOW_TEXT:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Metrics': True, 'Frame': False}
        else:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Metrics': True, 'Frame': True}
        dicts['objs'] = {s.GT_NAME: s.DATA1, s.PRED_NAME: s.DATA2}
        return dicts
    
    def import_obj(self, obj_name, data, i):
        """
        Mode specific way of importing organoid objects.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        data : Any
            Data source for the object.
        i : int
            Index of current frame needed to be imported.
        """
        nuclei = import_object(obj_name, data, False, i)
        set_origin_nuclei(nuclei, obj_name, i)

    def set_scene(self, i):
        """
        Mode specific code for importing each frame.

        Parameters
        ----------
        i : int
            Enumerator of current frame.
        """
        # Compute the relative index of the current frame.
        index = i - s.START

        # Compute metrics for the current frame objects.
        try:
            pred_obj = find_object(f'pred_{i}')
            gt_obj = find_object(f'gt_{i}')
            
            set_material(pred_obj, 'Prediction')
            set_material(gt_obj, 'Ground-Truth')
        except ObjectNotFoundException:
            print(f'Setting up frame {i} suspended')
        super().set_scene(i)

class Mode(ABC):
    """
    Parent class for all modes of visualization.
    """
    @abstractmethod
    def init_dicts(self):
        """
        Method that will initialize the dictionaries of objects/materials to be created
        or changed before importing the scene.
        """
            
    @abstractmethod
    def import_obj(self, obj_name, src, i):
        """
        Method to import an object into the scene.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        src : str
            File location where data about object can be found.
        i : int
            Index of current frame needed to be imported.
        """

    def pre_import(self):
        """"
        Method that can be called to change the scene before importing frames.
        """
        # Register frame counter for rendering
        register()
        # Populate dictionaries using child implementation, and call method to change
        # scene according to their content.
        dicts = self.init_dicts()
        change_scene_obj(dicts)
        return dicts['objs']

    def post_import(self):
        """
        Method that can be called to change the scene after importing frames.
        """

    def set_scene(self, i):
        """
        Method that changes the scene per imported frame.

        Parameters
        ----------
        i : int
            Enumerator of which frame is being added in the loop, starts at 0
        """
        # Get index that corresponds with the absolute frame number
        index = i - s.START
        try:
            # Apply transformations to objects that are used in all methods
            coll = find_collection(f"Frame_{i}")
            for obj in coll.all_objects[:]:
                # if ('gt' in obj.name or 'inner' in obj.name or 'pred' in obj.name) and obj.type == "MESH":
                #     set_rotation(obj, s.LENGTH * index, s.LENGTH + s.LENGTH * index)
                #     set_interpolation(obj,'LINEAR')
                hide_in_render(obj, index, s.LENGTH)
        except CollectionNotFoundException:
            print(f'Setting up frame {i} suspended due to not finding collection \'Frame_{i}\'')

    # @thread_worker
    def visualize(self):
        """
        Code that calls the full-pipeline of visualization.
        """
        # Call pre-scene import code
        objs = self.pre_import()
        # Loop through frames
        for i in range(s.START, s.END + 1):
            # Create collections for each absolute frame number
            coll_target = D.collections.get(f"Frame_{i}")
            if coll_target is None:
                coll_target = D.collections.new(name=f"Frame_{i}")
                C.scene.collection.children.link(coll_target)

            C.view_layer.active_layer_collection = \
                C.view_layer.layer_collection.children[f"Frame_{i}"]
            
            # Go over the dictionary that contains desired object names and their
            # data source, and import them into the scene.
            for obj_name, data in objs.items():
                self.import_obj(obj_name+str(i), data, i)
            
            # Call child specific code for importing a frame at index i
            self.set_scene(i)
     
        # Call post-scene import code
        self.post_import()

class ObjectNotFoundException(Exception):
    """
    Custom exception for when an object is not found in the scene.
    """
    pass

class ROTATION(Enum):
    """
    Enumeration for different rotation speeds.

    Attributes:
    No_Rotation: No rotation. Picture Mode.
    Very_Slow: Very slow rotation speed. 300 frames per rotation.
    Slow: Slow rotation speed. 200 frames per rotation.
    Normal: Normal rotation speed. 150 frames per rotation.
    Quick: Quick rotation speed. 100 frames per rotation.
    Very_Quick: Very quick rotation speed. 50 frames per rotation.
    """
    No_Rotation = 1
    Very_Slow = 300
    Slow = 200
    Normal = 150
    Quick = 100
    Very_Quick = 50

class RenderProgress:
    def __init__(self):
        self.root = Tk()
        self.mainframe = None
        self.progress_bar = None
        self.progress_text = None
        self.progress_time = None
        self.start_time = None

    def init_progress_bar(self):
        """
        Initializes the progress bar for rendering.
        """
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()

        x = (screen_width - 450) // 2  # 450 is the width of the window
        y = (screen_height - 130) // 2  # 130 is the height of the window

        self.root.geometry(f"450x130+{x}+{y}")  # Set window size and position
        self.root.title("Rendering")
        self.root.resizable(False, False)
        self.mainframe = ttk.Frame(self.root, padding="3 3 12 12")
        self.mainframe.grid(column=0, row=0, sticky=(N, W, E, S))
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)

        self.progress_bar = ttk.Progressbar(self.mainframe, maximum=100, length=420)
        self.progress_bar.place(x=10, y=50)

        self.progress_text = Label(self.mainframe, text="Progress: 0%")
        self.progress_text.place(x=180, y=15)

        self.progress_time = Label(self.mainframe, text="Time Left: Calculating..")
        self.progress_time.place(x=160, y=90)

    def render_progress_callback(self, scene):
        """
        Callback function to update the progress bar during rendering.
        
        Parameters
        ----------
        scene : bpy.types.Scene
            The current Blender scene.
        """
        frame = scene.frame_current
        total_frames = scene.frame_end - scene.frame_start + 1
        progress = int((frame - scene.frame_start) / total_frames * 100)
        progress = min(progress, 100)  # Ensure progress does not exceed 100%

        current_time = datetime.datetime.now()
        elapsed_time = (current_time - self.start_time).total_seconds()
        time_per_frame = elapsed_time / (frame - scene.frame_start + 1)
        remaining_frames = total_frames - (frame - scene.frame_start + 1)
        estimated_remaining_time = remaining_frames * time_per_frame

        remaining_time_str = str(datetime.timedelta(seconds=int(estimated_remaining_time)))

        # Use the main thread to update the GUI
        self.root.after(0, self.update_gui, progress, remaining_time_str)

    def update_gui(self, progress, remaining_time_str):
        """
        Updates the GUI elements with the current progress and estimated remaining time.

        Parameters
        ----------
        progress : int
            The current progress percentage.
        remaining_time_str : str
            The estimated remaining time as a string.
        """
        self.progress_bar['value'] = progress
        self.progress_text['text'] = f"Progress: {progress}%"
        self.progress_time['text'] = f"Time Left: {remaining_time_str}"

    def optimize_render_settings(self):
        """
        Optimizes the render settings for faster rendering.
        """
        C.scene.render.use_motion_blur = False
        C.scene.render.use_simplify = True
        C.scene.render.use_border = False

    def render_video(self):
        """
        Renders the video to the specified path.

        Parameters
        ----------
        path : str
            The file path where the rendered video will be saved.
        start : int
            The starting frame number.
        end : int
            The ending frame number.
        """
        C.scene.render.filepath = str(Path(s.SRC_VID / s.FILENAME))
        C.scene.frame_start = s.S_FRAME
        C.scene.frame_end = s.E_FRAME
        C.scene.render.engine = 'BLENDER_EEVEE'
        C.scene.render.image_settings.file_format = 'FFMPEG'
        C.scene.render.ffmpeg.format = 'MPEG4'
        C.scene.render.ffmpeg.codec = 'H264'

        self.optimize_render_settings()

        self.start_time = datetime.datetime.now()

        self.init_progress_bar()

        def render():
            def handler(scene):
                self.render_progress_callback(scene)

            A.handlers.frame_change_pre.append(handler)
            try:
                O.render.render(animation=True)
            finally:
                A.handlers.frame_change_pre.remove(handler)
                self.root.after(0, self.root.destroy) 

        threading.Thread(target=render).start()
        self.root.mainloop()

    def render_samples(self):
        """
        Sets the TAA render samples for EEVEE renderer.
        """
        render_settings = C.scene.eevee
        render_settings.taa_render_samples = s.SAMPLES

class SAMPLES(Enum):
    """
    Enumeration for different sampling resolutions. Multiples of 26
    to leverage the batch size of Blender rendering to it's maximum.

    Attributes:
    Low: Low sampling resolution.
    Medium: Medium sampling resolution.
    High: High sampling resolution.
    """
    Low = 1
    Medium = 26
    High = 52

class Timelapse(Mode):
    """
    Mode that shows the movement of organoids over time. This mode separates all nuclei,
    tracks them and links them using colors, or shades of colors if cell division occurs.
    """
    def init_dicts(self):
        """
        Mode specific objects/materials needed to be created or changed.
        """
        global data_frame
        data_frame = read_trackmate()
        dicts = {}
        dicts['texts'] = {}
        dicts['materials'] = {}
        for i in range(0, 5):
            dicts['materials'].update({f'Track {i+1}': (0.2 * i, 1.0, 0.8, 1.0, 'BLEND')})
        if s.SHOW_TEXT:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Legend': True, 'Metrics': True}
        else:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Legend': True, 'Metrics': True, 'Frame': True}
        dicts['objs'] = {s.GT_NAME: s.DATA1}
        return dicts
    
    def import_obj(self, obj_name, data, i):
        """
        Mode specific way of importing organoid objects.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        data : Any
            Data source for the object.
        i : int
            Index of current frame needed to be imported.
        """
        nuclei = import_object(obj_name, data, True, i)
        set_origin_nuclei(nuclei, obj_name, i)
        rename_nuclei(nuclei, data_frame, i + s.FRAME_OFFSET)

    def set_scene(self, i):
        """
        Mode specific code for importing each frame.

        Parameters
        ----------
        i : int
            Enumerator of current frame.
        """
        index = i - s.START   
        divided_nuclei = []
        try:
            coll = find_collection(f"Frame_{i}")      
            mesh_objects = [obj for obj in coll.all_objects if obj.type == 'MESH']
            for obj in mesh_objects:                
                # Show objects when needed in sequence render
                if index == 0:
                    track_id = get_track(data_frame, obj.name.split('.')[0])
                    set_material(obj, f'Track {track_id}')
                else:
                    ancestor_id = track_ancestor(data_frame, obj.name.split('.')[0])
                    if ancestor_id is None:
                        divided_nuclei.append(obj.name)
                hide_in_render(obj, index, s.LENGTH)
            mitosis_nuclei[i] = divided_nuclei
            divided_nuclei = []
        except CollectionNotFoundException:
            print(f'Setting up frame {i} suspended')

    def post_import(self):
        """
        Mode specific code to be run after importing all frames.
        """
        for i in range(s.START, s.END):
            index = i - s.START   
            try:
                coll = find_collection(f"Frame_{i}")         
                mesh_objects = [obj for obj in coll.all_objects if obj.type == 'MESH']
                for obj in mesh_objects:     
                    # Apply shrinkwrap modifiers to allow for parents of this object to morph
                    child_id = track_child(data_frame, obj.name.split('.')[0])
                    if child_id is not None:
                        try:
                            child = find_object(str(child_id))
                            if obj.material_slots:
                                mat_name = obj.material_slots[obj.active_material_index].material.name
                                set_material(child, mat_name)
                            apply_shape_key(obj, child, index) 
                        except ObjectNotFoundException:
                            print(f'Setting shrinkwrap keyframe from {child_id} to {obj.name} suspended')
                    else:
                        sign = 1
                        for div_nucleus in mitosis_nuclei[i+1]:
                            sign *= -1
                            try:
                                child = find_object(str(div_nucleus))
                                dupl_obj = duplicate_object(obj, f'{obj.name}-{div_nucleus}')
                                apply_shape_key(dupl_obj, child, index) 
                                child_mat = create_child_material(obj, sign * 0.04)
                                set_material(child, child_mat)
                                set_material(dupl_obj, child_mat)
                            except ObjectNotFoundException:
                                print(f'Setting shrinkwrap keyframe from {child_id} to {obj.name} suspended')
                        obj.name = f'{obj.name}.'
                        clean_duplicate_objs()
            except CollectionNotFoundException:
                print(f'Setting up shape keys for frame {i} suspended')

class Tracked(Mode):
    def init_dicts(self):
        """
        Mode specific objects/materials needed to be created or changed.
        """
        global data_frame
        data_frame = read_trackmate()
        dicts = {}
        dicts['texts'] = {}
        dicts['materials'] = {
            'Not Predicted': (0.45, 0.4, 0.8, 0.8, 'OPAQUE'), 
            'False Prediction': (0.8, 0.75, 0.8, 1.0, 'OPAQUE')
        }
        if s.SHOW_TEXT:
            dicts['hide_obj'] = {'Gradient_Bar': False, 'Metrics': True, 'Indicator': True}
        else:
            dicts['hide_obj'] = {'Gradient_Bar': False, 'Metrics': True, 'Indicator': True, 'Frame': True}
        dicts['objs'] = {s.GT_NAME: s.DATA1, s.PRED_NAME: s.DATA2}
        return dicts

    def import_obj(self, obj_name, data, i):
        """
        Mode specific way of importing organoid objects.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        data : Any
            Data source for the object.
        i : int
            Index of current frame needed to be imported.
        """
        nuclei = import_object(obj_name, data, True, i)
        set_origin_nuclei(nuclei, obj_name, i)
        rename_nuclei(nuclei, data_frame, i + s.FRAME_OFFSET)

    def set_scene(self, i):
        """
        Mode specific code for importing each frame.

        Parameters
        ----------
        i : int
            Enumerator of current frame.
        """
        index = i - s.START   
        try:
            coll = find_collection(f"Frame_{i}")      
            for obj in coll.all_objects[:]:
                if "pred" in obj.name or "gt" in obj.name:
                    color_nuclei(obj, i)
        except CollectionNotFoundException:
            print(f'Setting up frame {i} suspended')
        super().set_scene(i)

class Transparant(Mode):
    """
    Mode that compares predictions by Boolean operations. End-result is an intersecting object
    containing True Positives, and two objects that are only predicted in one of the two predictions
    which we call False Negatives and False Positives.
    """
    def init_dicts(self):
        """
        Mode specific objects/materials needed to be created or changed.
        """
        dicts = {}
        dicts['texts'] = {}
        dicts['materials'] = {
            'True Positive': (0.0, 0.0, 0.8, 1.0, 'OPAQUE'), 
            'False Negative': (0.035, 1.0, 0.8, 0.35, 'BLEND'), 
            'False Positive': (0.6, 1.0, 0.8, 0.35, 'BLEND')
        }
        if s.SHOW_TEXT:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Metrics': False}
        else:
            dicts['hide_obj'] = {'Gradient_Bar': True, 'Metrics': True, 'Frame': True}
        dicts['objs'] = {s.GT_NAME: s.DATA1, s.PRED_NAME: s.DATA2}
        return dicts
    
    def import_obj(self, obj_name, data, i):
        """
        Mode specific way of importing organoid objects.
        
        Parameters
        ----------
        obj_name : str
            Name of object to import/create.
        data : Any
            Data source for the object.
        i : int
            Index of current frame needed to be imported.
        """
        nuclei = import_object(obj_name, data, False, i)
        set_origin_nuclei(nuclei, obj_name, i)

    def set_scene(self, i):
        """
        Mode specific code for importing each frame.

        Parameters
        ----------
        i : int
            Enumerator of current frame.
        """
        try:
            inner = find_object(f'inner_{i}')
        except ObjectNotFoundException:
            # If the inner object does not exist, perform boolean operations
            try:
                gt_obj = find_object(f'gt_{i}')
                comp_obj = find_object(f'pred_{i}')

                create_overlap(gt_obj, comp_obj)

                inner = find_object(f'inner_{i}')
            except ObjectNotFoundException:
                print(f'Setting up frame {i} suspended')

            # Set materials for objects based on the boolean operation results
            set_material(inner, 'True Positive')
            set_material(gt_obj, 'False Negative')
            set_material(comp_obj, 'False Positive')  
            super().set_scene(i)