"""
Inscoper API Example Script: Tiling Acquisition
===============================================
This script performs a multi-tile acquisition grid and stitches them into a panorama.

Configuration Prerequisites:
---------------------------
1. Environment Variable: Ensure 'InscoperConfigsPath' is defined in your environment
   and points to your Inscoper configurations folder.
2. Configuration Name: This script assumes a hardware configuration named "camera_and_microscope_and_fluo_source".
   Change it below if your configuration is named differently.
3. Device/Channel Names: Customize device, channel parameters, and stage calibration settings below.
"""
import os
from typing import Dict
from dataclasses import dataclass
import numpy as np
import matplotlib.pyplot as plt
import inscoper_api


@dataclass
class SubDevice:
    name: str = None # a descriptive string for the sub-device
    tag: int = None # a unique identifier within the device for this sub-device
    inscoper_type: int = None # a "purpose" type for the sub-device (what it does, e.g., CUBE, OBJECTIVE, SHUTTER...)

@dataclass
class Device:
    name: str # the name given to the device during the configuration step
    hardware_id: int # an identifier for this device model
    sub_device_map: Dict[int, SubDevice] # the collection of chosen sub-devices for this particular device instance

class CustomImageListener(inscoper_api.ImageListener):
    def __init__(self):
        super().__init__()
        self._received_images = []
        self._image_counter = 0

    def handleNewImage(self, image):
        print(f"Image {self._image_counter} received!")
        self._image_counter += 1
        self._received_images.append(image)

    def getReceivedImages(self):
        return self._received_images

# Macro to all Inscoper configurations
CONFIG_PATH = os.environ['InscoperConfigsPath']

my_bridge = inscoper_api.Bridge()
my_listener = CustomImageListener()
my_bridge.addImageListener(my_listener)

config = my_bridge.readConfigFile(os.sep.join([CONFIG_PATH, "camera_and_microscope_and_fluo_source"]))
my_bridge.loadConfigFile(os.sep.join([CONFIG_PATH, "camera_and_microscope_and_fluo_source"]))
my_bridge.initDevices()

# Dictionary to store the mapping of hardware IDs to our Device objects.
device_map = {}

# Step 0: Iterate over the "Device List" from the configuration. 
# This tells us what specific devices have been added to this system's configuration.
for device in config.getDeviceList():
    driver_config = device.getDriverConfig()
    
    # In this example, we only keep devices controlled by the Inscoper Device Controller.
    if isinstance(driver_config, inscoper_api.InscoperBoxDriverConfig):
        device_name = device.getId() # e.g., "My Microscope" or "My Stage"
        device_id = driver_config.getHardwareId() # The internal hardware identifier

        sub_device_map = {}
        
        # Iterate over all sub-devices configured for this device.
        for sub_device_config in device.getSubDeviceConfigList():
            sub_device_tag = sub_device_config.getTag() # The unique tag will allow us to connect config and description
            sub_device_id = sub_device_config.getId() # Descriptive name of the sub-device
            
            # Create our SubDevice instance and map it by its tag.
            sub_device_map[sub_device_tag] = SubDevice(name=sub_device_id, tag=sub_device_tag)
            
        # Store the complete Device object (with its sub-devices) in our global map.
        device_map[device_id] = Device(name=device_name, hardware_id=device_id, sub_device_map=sub_device_map)

# At this point, device_map contains our devices, but we are missing the "purpose" of each sub-device.
print(device_map)

# Step 0: Iterate over the "Device Description List".
# This information acts as the "class definitions" (descriptions.xml) for all possible hardware supported by Inscoper,
# providing metadata like the generic type (purpose) of the sub-devices.
for device_desc in my_bridge.getDeviceDescriptionList():

    driver_desc = device_desc.getDriverDescription()
    
    # Again, we only keep devices controlled by the Inscoper Device Controller.
    if isinstance(driver_desc, inscoper_api.InscoperBoxDriverDescription):
        device_id = driver_desc.getHardwareId()
        
        # If this hardware_id matches one of the devices we found in our configuration...
        if device_id in device_map:
            current_device = device_map[device_id]
            sub_device_map = current_device.sub_device_map
            
            # Look at all possible sub-devices described for this hardware model.
            for sub_device_desc in device_desc.getSubDeviceDescriptionList():
                sub_device_tag = sub_device_desc.getTag()
                sub_device_type = sub_device_desc.getType() # The general purpose of the sub-device (e.g., CUBE, SHUTTER, etc.)
                
                # If this configured device actually uses this specific sub-device...
                if sub_device_tag in sub_device_map:
                    current_sub_device = sub_device_map[sub_device_tag]
                    
                    # Store the human-readable string of the sub-device type 
                    current_sub_device.inscoper_type = inscoper_api.Utils.subDeviceTypeToString(sub_device_type)

objective_turret_sub_device_id = None   # Turret that selects the objective
cube_turret_sub_device_id = None        # Filter-cube turret
tl_power_sub_device_id = None           # Power controller for the transmitted light
tl_shutter_sub_device_id = None         # Shutter inside the microscope for the transmitted light
epi_shutter_sub_device_id = None        # Shutter inside the microscope for the epi-fluorescence
led_blue_power_sub_device_id = None     # Power controller for the blue (440nm) LED
led_blue_shutter_sub_device_id = None   # Shutter inside the light source for the blue (440nm) LED

for device_id, device in device_map.items():
    for sub_device_tag, sub_device in device.sub_device_map.items():
        if sub_device.inscoper_type == "SHUTTER" and "TL" in sub_device.name:
            # Transmitted-light shutter: type is SHUTTER and the name contains "TL"
            tl_shutter_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "SHUTTER" and "IL" in sub_device.name and "Down" in sub_device.name:
            # Epi-fluorescence (incident-light) shutter: type is SHUTTER, path is "IL Down"
            epi_shutter_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "OBJECTIVE":
            # Objective turret: uniquely identified by the "OBJECTIVE" type
            objective_turret_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "CUBE" and "Down" in sub_device.name:
            # Filter-cube turret in the epi (downward) light path
            cube_turret_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif ("Intensity" in sub_device.name or "Power" in sub_device.name) and "TL" in sub_device.name:
            # Transmitted-light intensity: name contains "Intensity" or "Power" AND "TL"
            tl_power_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif ("Intensity" in sub_device.name or "Power" in sub_device.name) and "440" in sub_device.name:
            # 440 nm LED power: name contains "Intensity" or "Power" AND the wavelength "440"
            led_blue_power_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "SHUTTER" and "440" in sub_device.name:
            # 440 nm LED shutter: type is SHUTTER and the wavelength "440" appears in the name
            led_blue_shutter_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)


print(f"Objective Turret: {objective_turret_sub_device_id}")
print(f"Cube Turret: {cube_turret_sub_device_id}")
print(f"TL Power: {tl_power_sub_device_id}")
print(f"TL Shutter: {tl_shutter_sub_device_id}")
print(f"Epi Shutter: {epi_shutter_sub_device_id}")
print(f"LED Blue Power: {led_blue_power_sub_device_id}")
print(f"LED Blue Shutter: {led_blue_shutter_sub_device_id}")


blue_channel_description = {objective_turret_sub_device_id: 1,
             cube_turret_sub_device_id: 2,
             tl_power_sub_device_id: 0,
             tl_shutter_sub_device_id: 0,
             epi_shutter_sub_device_id: 1,
             led_blue_power_sub_device_id: 100,
             led_blue_shutter_sub_device_id: 1}

blue_status = inscoper_api.Status()
for sub_device_id, value in blue_channel_description.items():
    blue_status.addParam(sub_device_id, value)

# --8<-- [start:main_logic]

# ── Step 1: Resolve the sub-device IDs for the stage
x_axis_sub_device_id = None
y_axis_sub_device_id = None
move_xy_sub_device_id = None

for device_id, device in device_map.items():
    for sub_device_tag, sub_device in device.sub_device_map.items():
        if sub_device.inscoper_type == "X_AXIS":
            x_axis_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "Y_AXIS":
            y_axis_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)
        elif sub_device.inscoper_type == "MOVE_AXIS":
            move_xy_sub_device_id = inscoper_api.SubDeviceId(device.name, sub_device.name)

print(f"X Axis: {x_axis_sub_device_id}")
print(f"Y Axis: {y_axis_sub_device_id}")
print(f"Move XY: {move_xy_sub_device_id}")


# ── Step 2: Start the camera in the correct mode
start_cam_status = inscoper_api.Status()
start_cam_status.addTag("MyCamera")
start_cam_status.addTag("SOFTWARE_TRIGGER")

my_bridge.loadSequence(0, "StartCamera", [start_cam_status])
my_bridge.runSequence(0, False)

# ── Step 3: Set the blue channel
my_bridge.loadSequence(1, "DeviceUpdate", [blue_status])
my_bridge.runSequence(1, False)


# ── Step 4: Wait for the user to adjust the X and Y position on the microscope
input("Please manually adjust the X and Y position on the microscope. Once done, press Enter to continue...")


# ── Step 5: Get current stage positions and define the tiling grid
assert x_axis_sub_device_id is not None, "No X axis sub-device was found. Check your configuration."
assert y_axis_sub_device_id is not None, "No Y axis sub-device was found. Check your configuration."

current_x = my_bridge.getLongValue(x_axis_sub_device_id)
current_y = my_bridge.getLongValue(y_axis_sub_device_id)
print(f"Current X position: {current_x} nm")
print(f"Current Y position: {current_y} nm")

# Control parameters for tiling acquisition
tiling_size = 3          # Tiling grid dimension (e.g. 3 for a 3x3 grid)
overlap = 0.10           # Overlap fraction between adjacent tiles (10%)
pixel_size_um = 0.2166     # Calibration: Pixel size in microns (microns per pixel)

# Stage calibration parameters (aligns stage movements with camera sensor)
stage_calib_mirror = True
stage_calib_scale = 0.99942
stage_calib_angle = 4.711666431527656  # in radians (~270 degrees)

# Get camera dimensions to compute the field of view (FOV)
my_camera = my_bridge.getCamera("MyCamera")
camera_width = my_camera.getImageWidth()
camera_height = my_camera.getImageHeight()
print(f"Camera frame dimensions: {camera_width}x{camera_height} pixels")

# Calculate FOV in nanometers
fov_x_nm = camera_width * pixel_size_um * 1000
fov_y_nm = camera_height * pixel_size_um * 1000
print(f"FOV size: {fov_x_nm/1000:.2f} um x {fov_y_nm/1000:.2f} um")

# Calculate step sizes in nanometers, taking overlap into account
step_x = fov_x_nm * (1.0 - overlap)
step_y = fov_y_nm * (1.0 - overlap)

# Read the step constraints from the API
try:
    x_step = my_bridge.getStepValue(x_axis_sub_device_id)
    if x_step <= 0:
        x_step = 1
except Exception:
    x_step = 1
    
try:
    y_step = my_bridge.getStepValue(y_axis_sub_device_id)
    if y_step <= 0:
        y_step = 1
except Exception:
    y_step = 1

print(f"X Axis Step Constraint: {x_step} nm")
print(f"Y Axis Step Constraint: {y_step} nm")

tiling_status = []
tile_grid_coords = []


# Generate a grid centered on (current_x, current_y)
for j in range(tiling_size):
    # Use a snake pattern to minimize stage travel distance
    cols = range(tiling_size - 1, -1, -1) if j % 2 != 0 else range(tiling_size)
    for i in cols:
        # Compute tile offset in camera coordinate system (in nanometers)
        offset_cam_x = (i - (tiling_size - 1) / 2.0) * step_x
        offset_cam_y = (j - (tiling_size - 1) / 2.0) * step_y
        
        # Apply mirroring to camera Y coordinate if calibrated
        offset_cam_y_prime = -offset_cam_y if stage_calib_mirror else offset_cam_y
        
        # Rotate and scale offsets to get stage coordinate offsets (nanometers)
        cos_a = np.cos(stage_calib_angle)
        sin_a = np.sin(stage_calib_angle)
        offset_stage_x = stage_calib_scale * (offset_cam_x * cos_a - offset_cam_y_prime * sin_a)
        offset_stage_y = stage_calib_scale * (offset_cam_x * sin_a + offset_cam_y_prime * cos_a)
        
        # Round to the nearest multiple of the step constraint
        target_x = int(round((current_x + offset_stage_x) / x_step)) * int(x_step)
        target_y = int(round((current_y + offset_stage_y) / y_step)) * int(y_step)
        
        tile_status = inscoper_api.Status()
        for sub_device_id, value in blue_channel_description.items():
            tile_status.addParam(sub_device_id, value)
        
        tile_status.addParam(x_axis_sub_device_id, target_x)
        tile_status.addParam(y_axis_sub_device_id, target_y)
        
        tile_status.addTag("MyCamera")
        tile_status.addTag("SOFTWARE_TRIGGER")
        
        tiling_status.append(tile_status)
        
        # Map loop indices (j, i) to display coordinates (row, col)
        tile_grid_coords.append((j, i))

# ── Step 6: Load and run the acquisition
my_bridge.loadSequence(2, "Acquisition", tiling_status)
my_bridge.runSequence(2, False)

# ── Step 7: Retrieve and display the images
received_images = my_listener.getReceivedImages()

if received_images:
    N = len(received_images)
    print(f"Received {N} images.")
    
    # Get image properties from the first received image
    first_image = received_images[0]
    W = first_image.getWidth()
    H = first_image.getHeight()
    pixel_size_bytes = first_image.getPixelSize()
    bit_depth = first_image.getBitDepth()
    bytes_per_value = bit_depth // 8
    values_per_pixel = pixel_size_bytes // bytes_per_value
    
    # Calculate step size in pixels for the panorama mapping
    step_x_px = int(round(W * (1.0 - overlap)))
    step_y_px = int(round(H * (1.0 - overlap)))
    
    # Total size of the panorama
    pano_width = (tiling_size - 1) * step_x_px + W
    pano_height = (tiling_size - 1) * step_y_px + H
    
    print(f"Stitching {N} tiles into a {pano_width}x{pano_height} pixel panorama...")
    
    # Initialize panorama and weight matrices (assuming single-channel / grayscale)
    pano = np.zeros((pano_height, pano_width), dtype=np.float32)
    weights = np.zeros((pano_height, pano_width), dtype=np.float32)
    
    for idx, image in enumerate(received_images):
        if idx >= len(tile_grid_coords):
            break
        j, i = tile_grid_coords[idx]
        print(f"Processing tile {idx + 1}/{N} at grid position ({j}, {i})...")
        
        image_data = image.getImageData()
        img = np.reshape(image_data, (H, W, values_per_pixel))
        img_channel = img[:, :, 0].astype(np.float32)
        
        # Place in panorama
        start_x = i * step_x_px
        start_y = j * step_y_px
        end_x = start_x + W
        end_y = start_y + H
        
        pano[start_y:end_y, start_x:end_x] += img_channel
        weights[start_y:end_y, start_x:end_x] += 1.0
        
    # Normalize and display the panorama
    print("Normalizing overlap regions...")
    # Avoid division by zero in case of unmapped pixels
    mask = weights > 0
    pano[mask] /= weights[mask]
    
    # Calculate the total field of view (FOV) in microns
    total_fov_x_um = pano_width * pixel_size_um
    total_fov_y_um = pano_height * pixel_size_um
    
    print("Displaying final stitched panorama...")
    fig_pano, ax_pano = plt.subplots(figsize=(8, 8))
    ax_pano.imshow(pano, cmap="gray")
    ax_pano.set_title(f"Convalaria: {total_fov_x_um:.1f}x{total_fov_y_um:.1f} um, Overlap: {overlap*100:.0f}%")
    ax_pano.axis("off")
    plt.tight_layout()
    #plt.savefig("tiling.png", bbox_inches="tight", dpi=300, transparent=True)
    plt.show()
else:
    print("No images received to plot.")

# --8<-- [end:main_logic]

my_bridge.close()
