"""
Inscoper API Example Script: Stage Control
==========================================
This script demonstrates how to control a motorized stage (X, Y axis and movement).

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".
   Change it below if your configuration is named differently.
3. Device/Channel Names: Customize device and axis sub-device names below.
"""
import sys
import os
from typing import Dict
from dataclasses import dataclass
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

# Macro to all Inscoper configurations
CONFIG_PATH = os.environ['InscoperConfigsPath']

my_bridge = inscoper_api.Bridge()

config = my_bridge.readConfigFile(os.sep.join([CONFIG_PATH, "camera_and_microscope"]))
my_bridge.loadConfigFile(os.sep.join([CONFIG_PATH, "camera_and_microscope"]))
my_bridge.initDevices()

# Dictionary to store the mapping of hardware IDs to our Device objects.
device_map = {}

# Step 1: 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 2: 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)

# --8<-- [start:main_logic]

# To move the stage we need to interact with the X_AXIS, Y_AXIS, and MOVE_XY sub-devices.
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)

# Check that stage control sub-devices were found in your config.
assert x_axis_sub_device_id is not None,"No X axis position sub-device was found. Check your configuration."
assert y_axis_sub_device_id is not None,"No Y axis position sub-device was found. Check your configuration."
assert move_xy_sub_device_id is not None,"No move action sub_device was found. Check your configuration."

# Read the current values for the x/y positions.
current_x = my_bridge.getIntValue(x_axis_sub_device_id)
current_y = my_bridge.getIntValue(y_axis_sub_device_id)

print(f"X axis position sub-device: {x_axis_sub_device_id.getName()}. Current value: {current_x} nm")
print(f"Y axis position sub-device: {y_axis_sub_device_id.getName()}. Current value: {current_y} nm")
print(f"Move action sub-device: {move_xy_sub_device_id.getName()}")

# Set a new x position and order the stage to move
new_x_value = current_x - 1_000_000 # offset x position by 1 mm
my_bridge.setValue(x_axis_sub_device_id, new_x_value)
my_bridge.setValue(move_xy_sub_device_id)

# Did you see your stage move?

# --8<-- [end:main_logic]

my_bridge.close()