# Copyright (c) VOLVICON.
#
# This file is provided solely for educational purposes in connection with
# the Volvicon software and related documentation.
#
# You may use and modify this file for personal or internal purposes when
# working with Volvicon products. Any reproduction, distribution, or use
# of this file outside the Volvicon ecosystem is strictly prohibited
# without prior written authorization from VOLVICON.

"""
Voxel Data Helpers Tutorial.

This tutorial demonstrates Python helper functions available on flat voxel
data containers such as api.MaskUint8, api.MaskUint16, api.VolumeUint8,
api.VolumeInt16, api.VolumeUint16, and api.VolumeFloat32.

Prerequisites:
- Volvicon application must be running
- A mask named 'Mask' should exist in the project
- NumPy must be available for the to_numpy() and update_from_numpy() examples
"""

import ScriptingApi as api


app = api.Application()
mask_operations = app.get_mask_operations()

# Retrieve mask voxels as an 8-bit flat voxel buffer.
mask_uint8 : api.MaskUint8 = mask_operations.get_mask_uint8('Mask')

if not mask_uint8.data:
    print("Mask 'Mask' is empty or was not found. Create a mask named 'Mask' before running this tutorial.")
else:
    # =============================================================================
    # Basic Shape and Count Helpers
    # =============================================================================
    width, height, depth = mask_uint8.voxel_dimensions()
    print(f"Dimensions: width={width}, height={height}, depth={depth}")

    total_voxels = mask_uint8.voxel_count()
    occupied_voxels = mask_uint8.count_nonzero()
    print(f"Total voxels: {total_voxels}")
    print(f"Non-zero voxels: {occupied_voxels}")

    # =============================================================================
    # Coordinate Access Helpers
    # =============================================================================
    x = min(width - 1, width // 2)
    y = min(height - 1, height // 2)
    z = min(depth - 1, depth // 2)

    center_flat_index = mask_uint8.flat_index(x, y, z)
    center_value = mask_uint8.value_at(x, y, z)
    print(f"Voxel ({x}, {y}, {z}) is stored at flat index {center_flat_index}")
    print(f"Voxel ({x}, {y}, {z}) value: {center_value}")

    # set_value_at() updates the in-memory object. This example works on a copy
    # and does not write the modified data back to the project.
    edited_mask_uint8 : api.MaskUint8 = api.MaskUint8()
    edited_mask_uint8.data = list(mask_uint8.data)
    edited_mask_uint8.dimensions = list(mask_uint8.dimensions)
    edited_mask_uint8.spacing = list(mask_uint8.spacing)
    edited_mask_uint8.origin = list(mask_uint8.origin)
    edited_mask_uint8.set_value_at(x, y, z, center_value)

    # To write edited voxels back to the project, call this intentionally:
    # mask_operations.set_mask_uint8('Mask', edited_mask_uint8)

    # =============================================================================
    # Row and Slice Iteration Helpers
    # =============================================================================
    first_slice_row_count = 0
    first_slice_nonzero = 0
    for row in mask_uint8.iter_rows(0):
        first_slice_row_count += 1
        first_slice_nonzero += sum(1 for value in row if value)

    print(f"Rows in first Z slice: {first_slice_row_count}")
    print(f"Non-zero voxels in first Z slice: {first_slice_nonzero}")

    first_slice_rows = mask_uint8.rows(0)
    print(f"First slice materialized as rows: {len(first_slice_rows)} rows")

    slice_count = 0
    for voxel_slice in mask_uint8.iter_slices():
        slice_count += 1
        if slice_count == 1:
            print(f"First materialized slice has {len(voxel_slice)} rows")

    all_slices = mask_uint8.slices()
    print(f"All slices materialized: {len(all_slices)} slices")

    # =============================================================================
    # NumPy Conversion Helpers
    # =============================================================================
    try:
        mask_array = mask_uint8.to_numpy(True)
        print(f"NumPy shape: {mask_array.shape}")
        print(f"NumPy non-zero voxels: {int((mask_array != 0).sum())}")

        # update_from_numpy() accepts arrays with shape (depth, height, width)
        # and updates both dimensions and flat row-major data.
        rebuilt_mask_uint8 : api.MaskUint8 = api.MaskUint8()
        rebuilt_mask_uint8.update_from_numpy(mask_array)
        rebuilt_mask_uint8.spacing = list(mask_uint8.spacing)
        rebuilt_mask_uint8.origin = list(mask_uint8.origin)

        print(f"Rebuilt dimensions: {rebuilt_mask_uint8.voxel_dimensions()}")

        # Write back only when you intend to replace compatible project data:
        # mask_operations.set_mask_uint8('Mask', rebuilt_mask_uint8)
    except ImportError:
        print("NumPy is not available, skipping to_numpy() and update_from_numpy() examples.")

    # For multi-label masks, the same helper methods are available on MaskUint16.
    # Note the multi-label getter is get_mask_uint_16 (with an underscore before 16).
    # mask_uint16 : api.MaskUint16 = mask_operations.get_mask_uint_16('Mask')
    # print(mask_uint16.voxel_dimensions())

    # =============================================================================
    # The same helpers work on volume voxel data
    # =============================================================================
    # Volume voxel data uses the typed getters on volume_operations:
    #   get_volume_uint8, get_volume_int_16, get_volume_uint_16, get_volume_float_32
    # (the bit count carries an underscore for the 16-/32-bit variants).
    volume_operations = app.get_volume_operations()
    volume_names = app.get_all_volume_names()
    if volume_names:
        # Get the voxel (scalar) data type of the volume as enum.
        # voxel_data_type : api.VoxelDataType = volume_operations.get_voxel_data_type(volume_names[0])

        # Retrieve volume voxel data with the typed getters. The returned object type depends on the getter used:

        # Get the volume voxel data as uint8 (unsigned 8-bit integers).
        # volume_uint8 : api.VolumeUint8 = volume_operations.get_volume_uint8(volume_names[0])

        # Get the volume voxel data as int16 (signed 16-bit integers).
        # volume_int_16 : api.VolumeInt16 = volume_operations.get_volume_int_16(volume_names[0])

        # Get the volume voxel data as uint16 (unsigned 16-bit integers).
        # volume_uint_16 : api.VolumeUint16 = volume_operations.get_volume_uint_16(volume_names[0])

        # Get the volume voxel data as float32 (32-bit floating point).
        volume_float_32 : api.VolumeFloat32 = volume_operations.get_volume_float_32(volume_names[0])

        print(f"Volume dimensions: {volume_float_32.voxel_dimensions()}")
        try:
            volume_array = volume_float_32.to_numpy(True)   # shape (depth, height, width)
            print(f"Volume NumPy shape: {volume_array.shape}")
            print(f"Volume mean intensity: {float(volume_array.mean()):.3f}")

            # Write modified voxels back with the matching typed setter:

            # Set the volume voxel data back as uint8 (unsigned 8-bit integers).
            # volume_operations.set_volume_uint8(volume_names[0], volume_uint8)

            # Set the volume voxel data back as int16 (signed 16-bit integers).
            # volume_operations.set_volume_int_16(volume_names[0], volume_int_16)

            # Set the volume voxel data back as uint16 (unsigned 16-bit integers).
            # volume_operations.set_volume_uint_16(volume_names[0], volume_uint_16)

            # Set the volume voxel data back as float32 (32-bit floating point).
            # volume_operations.set_volume_float_32(volume_names[0], volume_float_32)
        except ImportError:
            print("NumPy is not available, skipping the volume to_numpy() example.")