main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Auto-run script for the in-situ XRD experiment sequence.

This module coordinates the UR5e robot, screw automation, angle changer,
and the Anton Paar XRD instrument. It contains helper routines for
waiting for the stage, submitting scans, and the high-level sequencing
logic that orchestrates the full experiment.

Refactor highlights:
- Added type annotations and improved docstrings.
- Safer controller connection checks (avoid AttributeError on missing attrs).
- `perform_xrd_scan` returns a boolean success flag.
"""

import os, sys, time, json

from AngleChanger.AngleChanger_Driver import AngleChangerController
from SampleLoader.ScrewAutomater_Driver import ScrewAutomationController
from config import *
from Logger.Logger import TaskLogger

# --- Import Controller Classes from their respective driver files ---
from Adaptive_insituXRD.AntonPaarXRD_Driver import (
    AntonPaarXRDController,
    ApiAdvancedScanRequestDto,
    ApiSimpleScanRequestDto,
    ApiSetTemperatureRequestDto,
    NonAmbientAttributesDto,
    ScanRequestDto,
    WaitNodeDto,
    SampleAlignmentNodeDto,
    SampleAlignmentDto,
)
from robot_arm_ur5e.UR5e_Driver import UR5eController


# ==============================================================================
# --- Helper Function for XRD Operation ---
# ==============================================================================

def wait_until_idle(client: AntonPaarXRDController,
                    file_path: str, 
                    device_name: str,
                    logger: TaskLogger = None):
    while True:
        try:
            status = client.get_status()
            # instrument_status = status.instrument_status
            # stage_status = status.stage_status
            # current_temp = status.stage_temperature
            # logger.info(device_name, f"Idle: {instrument_status}, Status: {stage_status}, Temp: {current_temp}")

            if os.path.isfile(file_path):
                logger.info(device_name, f"Checking for file: {file_path}:{os.path.isfile(file_path)}")
                logger.info(device_name, "Instrument is stable.")
                client.stop()
                logger.info(device_name, "Instrument is turned off.")
                break
        except Exception as e:
            logger.error(device_name, f"Status check error: {e}")
        time.sleep(2)

def wait_for_sampleAlignment(client:AntonPaarXRDController,
                             dir_path: str, 
                            device_name: str,
                            temperature: int,
                            logger: TaskLogger = None):
    """Wait for the sample alignment to complete."""
    logger.info(device_name, "Waiting for sample alignment to complete...")
    while True:
        try:
            status = client.get_status()

            file_list = os.listdir(dir_path) # 해당 폴더 내의 파일명 리스트 [3, 7]
            found=False

            for filename in file_list:
                if filename.endswith('.scn') and "sampleAlignment_ZScan-NA#" in filename and str(temperature) in filename:
                    logger.info(device_name, f"Sample alignment completed: {filename}")
                    found=True
            
            if found:
                logger.info(device_name, "Instrument is aligned.")
                client.stop()
                break
            
            time.sleep(2)

        except Exception as e:
            logger.error(device_name, f"Error during sample alignment wait: {e}")
        time.sleep(2)

def wait_for_target_temperature(
    client: AntonPaarXRDController,
    device_name: str,
    target_temp: float,
    tolerance: float = 5.0,
    logger: TaskLogger = None,
) -> None:
    """Block until the stage temperature is within `tolerance` of `target_temp`.

    This routine polls the instrument status periodically and will stop
    the instrument once the target is reached. It raises on instrument
    errors.

    Args:
        client: AntonPaarXRDController instance.
        device_name: human-readable device name for logging.
        target_temp: target temperature in degrees Celsius.
        tolerance: allowed deviation from target_temp before considering
            the stage "ready".
        logger: TaskLogger instance (required). If None, a TaskLogger
            with the device name will be created.
    """
    if logger is None:
        logger = TaskLogger(device_name)

    logger.info(device_name, f"Waiting for stage to reach {target_temp}°C (±{tolerance}°C)...")

    poll_interval = 15.0  # seconds between status polls

    while True:
        try:
            status = client.get_status()
            current_temp = status.stage_temperature
            stage_status = status.stage_status

            temp_is_stable = abs(current_temp - target_temp) <= tolerance

            if temp_is_stable:
                logger.info(device_name, f"Target temperature {target_temp}°C reached (measured {current_temp}°C). Stopping instrument.")
                client.stop()
                time.sleep(5)
                return

            if status.instrument_status == "Error" or stage_status == "Error":
                raise RuntimeError(f"Instrument error: {status.last_error_message}")

            logger.debug(device_name, f"Current temp: {current_temp}°C; stage status: {stage_status}; waiting...")
            time.sleep(poll_interval)

        except Exception as e:
            logger.error(device_name, f"Error while waiting for stage readiness: {e}")
            raise

def perform_xrd_scan(
    xrd_client: AntonPaarXRDController,
    xrd_scan_mode: str,
    logger: TaskLogger,
    timestamp: str = time.strftime("%Y%m%d_%H%M%S"),
    experiment_config: dict = None
) -> bool:
    """Perform the configured temperature + scan sequence.

    Returns True on success, False on failure. The function attempts a
    safe shutdown on errors and performs a final cooldown + RT scan in
    the finally block.
    """
    device_name = xrd_client.device_name
    logger.info(device_name, "Preparing to perform in-situ XRD scan(s)...")

    success = True
    # Load the configuration file once outside the try block.
    # with open(experiment_config_path, 'r') as file:
    #     config_data = json.load(file)
    config_data = experiment_config

    try:
        # Pre-check: ensure instrument is idle before starting
        logger.info(device_name, "Starting pre-experiment checks...")
        status = xrd_client.get_status()
        logger.info(device_name, f"Initial instrument status: {status.instrument_status}, stage temp: {status.stage_temperature}°C")

        if status.instrument_status != "Idle":
            logger.warning(device_name, "Instrument not idle; attempting to stop existing activity...")
            xrd_client.stop()
            time.sleep(5)
            status = xrd_client.get_status()
            if status.instrument_status != "Idle":
                raise RuntimeError("Failed to set instrument to Idle state")

        # Main loop over temperature steps
        logger.info(device_name, "Beginning in-situ experiment steps...")

        experiment_name = config_data["EXPERIMENT_NAME"]
        user_name = config_data["USER_NAME"]
        base_storage_path = config_data["BASE_STORAGE_PATH"]

        total_experiment_name = f"{experiment_name}_{timestamp}"
        current_storage_path = f"{base_storage_path}/{user_name}/{total_experiment_name}"

        try:
            os.makedirs(current_storage_path, exist_ok=True)
            main_logger.info(program_name, f"Data for this run will be saved in: {current_storage_path}")
        except OSError as e:
            main_logger.error(f"Failed to create storage directory: {e}")
            sys.exit(1)

        if xrd_scan_mode.lower() == "dynamic":
            logger.info(device_name, "Dynamic scan mode selected.")
            raise NotImplementedError("Dynamic scan mode is not implemented in this version.")
        
        elif xrd_scan_mode.lower() == "static":
            EXPERIMENTS_TO_RUN_LIST = config_data["EXPERIMENTS_TO_RUN_LIST"]
            logger.info(device_name, f"Static scan mode selected - running {len(EXPERIMENTS_TO_RUN_LIST)} configured steps.")

            for idx, experiment in enumerate(EXPERIMENTS_TO_RUN_LIST):
                # time.sleep(4)  # brief pause between steps
                logger.info(device_name, f"--- Starting step {idx + 1}/{len(EXPERIMENTS_TO_RUN_LIST)} ---")

                # --- 1. Conditionally create DTOs for all parameters for the current step ---
                set_temp_humidity = experiment.get("set_temperature_humidity", {})
                coupled_scan = experiment.get("coupled_scan", {})
                sample_alignment_data = experiment.get("sample_alignment", {})
                wait_data = experiment.get("wait", {})

                non_ambient_attributes = None
                if set_temp_humidity:
                    temperature_dict_list=[]
                    for temperature in set_temp_humidity.get("TEMPERATURE_STEPS", []):
                        temperature_dict_list.append({
                            "temperatureCelsius": temperature,
                            "relativeHumidity": 30
                        })
                    non_ambient_attributes = NonAmbientAttributesDto(
                        series_definition_mode=0,
                        set_points=temperature_dict_list,
                        heating_or_cooling_rate=set_temp_humidity.get("HEATING_RATES", 150.0),
                        temperature_hit_tolerance=set_temp_humidity.get("TEMPERATURE_HEAT_TOLERANCE", 5.0),
                        use_auto_height_compensation=set_temp_humidity.get("USE_AUTO_HEIGHT_COMPENSATION", False),
                    )

                sample_alignment = None
                if sample_alignment_data:
                    sample_alignment = SampleAlignmentNodeDto(
                        sampleAlignments=[
                            SampleAlignmentDto(
                                sample_alignment_type=sample_alignment_data.get("SAMPLE_ALIGNMENT_TYPE"),
                                sample_alignment_iterations=sample_alignment_data.get("SAMPLE_ALIGNMENT_ITERATIONS"),
                                sample_alignment_start_type=sample_alignment_data.get("SAMPLE_ALIGNMENT_START_TYPE"),
                            )
                        ]
                    )
                
                wait_node = None
                if wait_data and wait_data.get("WAIT_TIME_SECONDS"):
                    wait_node = WaitNodeDto(wait_time_seconds=wait_data.get("WAIT_TIME_SECONDS"))

                current_scan_request = None
                if coupled_scan: # Only create a scan request if scan parameters exist
                    scan_name = experiment_name + f"_Step{idx + 1}_Scan"
                    current_scan_request = ScanRequestDto(
                        is_start_position_relative=False,
                        two_theta_offset_from=coupled_scan.get("TWO_THETA_OFFSET_FROM"),
                        two_theta_offset_to=coupled_scan.get("TWO_THETA_OFFSET_TO"),
                        step_size=coupled_scan.get("STEP_SIZE"),
                        time_per_step=coupled_scan.get("TIME_PER_STEP"),
                        name=scan_name
                    )

                # --- 2. Call different APIs based on whether scan information is present ---
                # This logic is crucial because the API endpoint for scanning requires a 'scan' node,
                # while non-scan operations (like setting temp) do not.
                
                if current_scan_request:
                    # If scan information exists, call start_advanced_scan.
                    logger.info(device_name, "This step includes a scan. Using start_advanced_scan.")
                    advanced_request = ApiAdvancedScanRequestDto(
                        # save_non_ambient_log_files=True,
                        non_ambient_attributes=non_ambient_attributes,
                        sample_alignment=sample_alignment,
                        wait=wait_node,
                        scan=current_scan_request,
                        experiment_name=experiment_name,
                        storage_path=current_storage_path,
                    )
                    target_temp_str = f" to {set_temp_humidity.get('TEMPERATURE_STEPS', ['N/A'])}°C" if non_ambient_attributes else ""
                    logger.info(device_name, f"Submitting job: Set temperature{target_temp_str}, then scan.")
                    result = xrd_client.start_advanced_scan(advanced_request, blocking=True)
                    
                    temperature_list=set_temp_humidity.get("TEMPERATURE_STEPS", [])
                    scan_number=len(temperature_list)
                    file_path = f"{current_storage_path}/{scan_name}-NA#{scan_number}_{float(temperature_list[-1])}C.hdf5"
                    
                    wait_until_idle(
                        xrd_client,
                        file_path,
                        device_name,
                        logger=logger
                    )
                    logger.info(device_name, f"Completed scan step {idx + 1}: {result}")

                if bool(current_scan_request) == False and bool(non_ambient_attributes) == True:
                    if bool(sample_alignment) == True:
                        # If there's no scan info, but there are temperature or alignment settings, call set_temperature.
                        logger.info(device_name, "This step is for temperature/alignment only. Using set_temperature.")
                        temp_request = ApiSetTemperatureRequestDto(
                            non_ambient_attributes=non_ambient_attributes,
                            wait=wait_node,
                            sample_alignment=sample_alignment,
                            experiment_name=experiment_name,
                            storage_path=current_storage_path,
                        )
                        logger.info(device_name, "Submitting job for temperature setting and sample alignment.")
                        result = xrd_client.set_temperature(temp_request)
                        logger.info(device_name, f"Completed temperature/alignment step {idx + 1}: {result}")
                        
                        # xrd_client.stop()
                        wait_for_sampleAlignment(
                            xrd_client,
                            current_storage_path,
                            device_name,
                            set_temp_humidity.get("TEMPERATURE_STEPS", [])[0],
                            logger=logger
                        )
                    else:
                        logger.info(device_name, "This step is for temperature only. Using set_temperature.")
                        temp_request = ApiSetTemperatureRequestDto(
                            non_ambient_attributes=non_ambient_attributes,
                            wait=wait_node,
                            experiment_name=experiment_name,
                            storage_path=current_storage_path,
                        )
                        logger.info(device_name, "Submitting job for temperature setting.")
                        result = xrd_client.set_temperature(temp_request)
                        logger.info(device_name, f"Completed temperature step {idx + 1}: {result}")

                        # xrd_client.stop()
                        wait_for_target_temperature(
                            xrd_client, 
                            device_name, 
                            set_temp_humidity.get("TEMPERATURE_STEPS", [])[0], 
                            tolerance=5.0, 
                            logger=logger
                        )

                else:
                    # Skip if the step contains no actionable information.
                    logger.warning(device_name, f"Step {idx + 1} is empty and will be skipped.")

                logger.info(device_name, "All configured experiment steps finished")
                
        logger.info(device_name, "All configured experiment steps finished")

    except Exception as e:
        success = False
        logger.error(device_name, f"Experiment error: {e}")
        logger.info(device_name, "Attempting safe shutdown of instrument")
        try:
            xrd_client.stop()
        except Exception:
            logger.warning(device_name, "Failed to stop instrument cleanly")

    finally:
        # The finally block can be adjusted with the same conditional logic if needed.
        # For simplicity, this version keeps the final actions separate and robust.
        logger.info(device_name, f"Finalizing: returning to room temperature (25°C) and performing final scan")
        try:
            EXPERIMENTS_TO_RUN_FINAL = config_data.get("EXPERIMENTS_TO_RUN_FINAL", {})
            if not EXPERIMENTS_TO_RUN_FINAL:
                logger.info(device_name, "No final experiment step configured. Skipping.")
                return success

            final_temp_config = EXPERIMENTS_TO_RUN_FINAL.get("set_temperature_humidity", {})
            final_scan_config = EXPERIMENTS_TO_RUN_FINAL.get("coupled_scan", {})
            final_sample_alignment_data = EXPERIMENTS_TO_RUN_FINAL.get("sample_alignment", {})
            final_wait_config = EXPERIMENTS_TO_RUN_FINAL.get("wait", {})

            # 1. Set temperature to room temperature.
            temp_request = ApiSetTemperatureRequestDto(
                non_ambient_attributes=NonAmbientAttributesDto(
                    series_definition_mode=0,
                    set_points=final_temp_config.get("TEMPERATURE_STEPS", [25.0]),
                    heating_or_cooling_rate=final_temp_config.get("HEATING_RATES", 150.0),
                ),
                sample_alignment = SampleAlignmentNodeDto(
                        sampleAlignments=[
                            SampleAlignmentDto(
                                sample_alignment_type=final_sample_alignment_data.get("SAMPLE_ALIGNMENT_TYPE"),
                                sample_alignment_iterations=final_sample_alignment_data.get("SAMPLE_ALIGNMENT_ITERATIONS"),
                                sample_alignment_start_type=final_sample_alignment_data.get("SAMPLE_ALIGNMENT_START_TYPE"),
                            )
                        ]
                    ),
                wait=WaitNodeDto(wait_time_seconds=final_wait_config.get("WAIT_TIME_SECONDS", 0)),
                experiment_name=experiment_name,
                storage_path=current_storage_path
            )
            xrd_client.set_temperature(temp_request)
            wait_for_target_temperature(
                xrd_client, 
                device_name, 
                final_temp_config.get("TEMPERATURE_STEPS", [25.0])[0], 
                tolerance=5.0, 
                logger=logger
            )
            
            # 2. Perform the final scan at room temperature if configured.
            if final_scan_config:
                final_scan_name = f"{experiment_name}_Final_RT_Scan"
                final_scan_parameters = ScanRequestDto(
                    name=final_scan_name,
                    is_start_position_relative=False,
                    two_theta_offset_from=final_scan_config.get("TWO_THETA_OFFSET_FROM"),
                    two_theta_offset_to=final_scan_config.get("TWO_THETA_OFFSET_TO"),
                    step_size=final_scan_config.get("STEP_SIZE"),
                    time_per_step=final_scan_config.get("TIME_PER_STEP"),
                )
                final_scan_request = ApiAdvancedScanRequestDto(
                    scan=final_scan_parameters, 
                    experiment_name=final_scan_name, 
                    storage_path=current_storage_path
                )
                logger.info(device_name, f"Starting final RT scan: {final_scan_name}")
                res = xrd_client.start_simple_scan(final_scan_request, blocking=True)
                logger.info(device_name, f"Final RT scan complete: {res}")

                file_path = f"{current_storage_path}/{final_scan_name}.hdf5"

                # xrd_client.stop()
                wait_until_idle(
                    xrd_client,
                    file_path,
                    device_name,
                    logger=logger
                )

        except Exception as e:
            logger.error(device_name, f"Error during cooldown/final scan: {e}")
            try:
                xrd_client.stop()
            except Exception:
                logger.warning(device_name, "Failed to stop instrument during final shutdown")

    return success


# ==============================================================================
# --- Main Execution Logic ---
# ==============================================================================

def run_full_sequence(program_name:str, 
                      location_index: int, 
                      logger: TaskLogger, 
                      experiment_config: dict,
                      xrd_scan_mode:str="static", 
                      timestamp: str = time.strftime("%Y%m%d_%H%M%S"), 
                      ) -> None:
    """Executes the entire automated experiment sequence for a given crucible location."""
    # --- Initialize all controllers ---
    logger.info(program_name,"--- Initializing all hardware controllers ---")
    
    try:
        # --- Define dynamic program names based on location_index ---
        prog_crucible_to_holder = f"/motionplanning/crucible_{location_index}_crucibleholder.urp"
        prog_holder_to_crucible = f"/motionplanning/crucibleholder_crucible_{location_index}.urp"

        success = True

        # --- Start of the Action Sequence ---
        logger.info(program_name,f"--- STARTING AUTOMATED SEQUENCE FOR CRUCIBLE AT LOCATION {location_index} ---")

        # Step 1: [UR5e] crucible_{location_index} to crucibleholder
        if success:
            logger.info(program_name,f"Step 1/11: [UR5e] Moving sample carrier from crucible {location_index} to crucible holder...")
            success = robot.run_urp_program(prog_crucible_to_holder)

        # Step 2: [UR5e] crucibleholder to angleChanger
        if success:
            logger.info(program_name,"Step 2/11: [UR5e] Moving crucible holder to angle changer...")
            success = robot.run_urp_program("/motionplanning/crucibleholder_anglechanger.urp")

        # Step 3: [AngleChanger] move up
        if success:
            logger.info(program_name,"Step 3/11: [AngleChanger] Moving up...")
            success = angle_changer.move_up()

        # Step 4: [UR5e] angleChanger to insituXRD
        if success:
            logger.info(program_name,"Step 4/11: [UR5e] Moving from angle changer to XRD chamber...")
            success = robot.run_urp_program("/motionplanning/anglechanger_insituXRD.urp")

        # Step 5: [ScrewAutomater] loadSample
        if success:
            logger.info(program_name,"Step 5/11: [ScrewAutomater] Loading sample...")
            success = screw_automater.loadSample()

        # Step 6: [insituXRD] run XRD operation
        if success:
            logger.info(program_name,"Step 6/11: [insituXRD] Starting in-situ XRD measurement...")
            success = perform_xrd_scan(anton_paar_xrd, xrd_scan_mode, logger, timestamp=timestamp, experiment_config_path=experiment_config)

        # Step 7: [ScrewAutomater] unloadSample
        if success:
            logger.info(program_name,"Step 7/11: [ScrewAutomater] Unloading sample...")
            success = screw_automater.unloadSample()

        # Step 8: [UR5e] insituXRD to angleChanger
        if success:
            logger.info(program_name,"Step 8/11: [UR5e] Moving from XRD chamber to angle changer...")
            success = robot.run_urp_program("/motionplanning/insituXRD_anglechanger.urp")

        # Step 9: [AngleChanger] move down
        if success:
            logger.info(program_name,"Step 9/11: [AngleChanger] Moving down...")
            success = angle_changer.move_down()

        # Step 10: [UR5e] angleChanger to crucibleHolder
        if success:
            logger.info(program_name,"Step 10/11: [UR5e] Moving from angle changer to crucible holder...")
            success = robot.run_urp_program("/motionplanning/anglechanger_crucibleholder.urp")

        # Step 11: [UR5e] crucibleholder to crucible_{location_index}
        if success:
            logger.info(program_name,"Step 11/11: [UR5e] Returning sample carrier to crucible...")
            success = robot.run_urp_program(prog_holder_to_crucible)

        # --- End of Sequence ---
        if not success:
            raise RuntimeError("A step in the sequence failed. Check logs for details.")
        
        logger.info(program_name,"--- MAIN SEQUENCE COMPLETED SUCCESSFULLY ---")

    except Exception as e:
        logger.error(program_name,f"An unexpected critical error occurred: {e}")
        raise Exception(f"An unexpected critical error occurred: {e}")


def load_sample(program_name:str, 
                location_index: int, 
                logger: TaskLogger, 
                robot: UR5eController,
                angle_changer: AngleChangerController,
                screw_automater: ScrewAutomationController
                ) -> None:
    """Executes the load sequence for a given crucible index."""
    # --- Initialize all controllers ---
    logger.info(program_name,"--- Initializing all hardware controllers ---")
    
    try:
        # --- Define dynamic program names based on location_index ---
        prog_crucible_to_holder = f"/motionplanning/crucible_{location_index}_crucibleholder.urp"
        prog_holder_to_crucible = f"/motionplanning/crucibleholder_crucible_{location_index}.urp"

        success = True

        # --- Start of the Action Sequence ---
        logger.info(program_name,f"--- STARTING LOAD SEQUENCE FOR CRUCIBLE AT LOCATION {location_index} ---")

        # Step 1: [UR5e] crucible_{location_index} to crucibleholder
        if success:
            logger.info(program_name,f"Step 1/11: [UR5e] Moving sample carrier from crucible {location_index} to crucible holder...")
            success = robot.run_urp_program(prog_crucible_to_holder)

        # Step 2: [UR5e] crucibleholder to angleChanger
        if success:
            logger.info(program_name,"Step 2/11: [UR5e] Moving crucible holder to angle changer...")
            success = robot.run_urp_program("/motionplanning/crucibleholder_anglechanger.urp")

        # Step 3: [AngleChanger] move up
        if success:
            logger.info(program_name,"Step 3/11: [AngleChanger] Moving up...")
            success = angle_changer.move_up()

        # Step 4: [UR5e] angleChanger to insituXRD
        if success:
            logger.info(program_name,"Step 4/11: [UR5e] Moving from angle changer to XRD chamber...")
            success = robot.run_urp_program("/motionplanning/anglechanger_insituXRD.urp")

        # Step 5: [ScrewAutomater] loadSample
        if success:
            logger.info(program_name,"Step 5/11: [ScrewAutomater] Loading sample...")
            success = screw_automater.loadSample()

        # --- End of Sequence ---
        if not success:
            raise RuntimeError("A step in the sequence failed. Check logs for details.")
        
        logger.info(program_name,"--- LOAD SEQUENCE COMPLETED SUCCESSFULLY ---")

    except Exception as e:
        logger.error(program_name,f"An unexpected critical error occurred: {e}")
        raise Exception(f"An unexpected critical error occurred: {e}")


def unload_sample(program_name:str, 
                    location_index: int, 
                    logger: TaskLogger, 
                    robot:UR5eController, 
                    angle_changer: AngleChangerController, 
                    screw_automater: ScrewAutomationController) -> None:
    """Executes the entire unload sequence for a given crucible location."""
    # --- Initialize all controllers ---
    logger.info(program_name,"--- Initializing all hardware controllers ---")
    
    try:
        # --- Define dynamic program names based on location_index ---
        prog_crucible_to_holder = f"/motionplanning/crucible_{location_index}_crucibleholder.urp"
        prog_holder_to_crucible = f"/motionplanning/crucibleholder_crucible_{location_index}.urp"

        success = True

        # --- Start of the Action Sequence ---
        logger.info(program_name,f"--- STARTING UNLOAD SEQUENCE FOR CRUCIBLE AT LOCATION {location_index} ---")

        # Step 7: [ScrewAutomater] unloadSample
        if success:
            logger.info(program_name,"Step 7/11: [ScrewAutomater] Unloading sample...")
            success = screw_automater.unloadSample()

        # Step 8: [UR5e] insituXRD to angleChanger
        if success:
            logger.info(program_name,"Step 8/11: [UR5e] Moving from XRD chamber to angle changer...")
            success = robot.run_urp_program("/motionplanning/insituXRD_anglechanger.urp")

        # Step 9: [AngleChanger] move down
        if success:
            logger.info(program_name,"Step 9/11: [AngleChanger] Moving down...")
            success = angle_changer.move_down()

        # Step 10: [UR5e] angleChanger to crucibleHolder
        if success:
            logger.info(program_name,"Step 10/11: [UR5e] Moving from angle changer to crucible holder...")
            success = robot.run_urp_program("/motionplanning/anglechanger_crucibleholder.urp")

        # Step 11: [UR5e] crucibleholder to crucible_{location_index}
        if success:
            logger.info(program_name,"Step 11/11: [UR5e] Returning sample carrier to crucible...")
            success = robot.run_urp_program(prog_holder_to_crucible)

        # --- End of Sequence ---
        if not success:
            raise RuntimeError("A step in the sequence failed. Check logs for details.")
        
        logger.info(program_name,"--- UNLOAD SEQUENCE COMPLETED SUCCESSFULLY ---")

    except Exception as e:
        logger.error(program_name,f"An unexpected critical error occurred: {e}")
        raise Exception(f"An unexpected critical error occurred: {e}")


def run_insituXRD(program_name:str, 
                      logger: TaskLogger, 
                      xrd_scan_mode:str="static", 
                      timestamp: str = time.strftime("%Y%m%d_%H%M%S"), 
                      experiment_config_path: str="XRD_template/config.json",
                      anton_paar_xrd: AntonPaarXRDController = None) -> None:
    """Executes the entire automated in-situ XRD sequence for a given crucible location."""
    # --- Initialize all controllers ---
    logger.info(program_name,"--- Initializing all hardware controllers ---")
    
    try:
        success = True

        # --- Start of the Action Sequence ---
        logger.info(program_name,f"--- STARTING in-situ XRD SEQUENCE ---")

        # Step 6: [insituXRD] run XRD operation
        if success:
            logger.info(program_name,"Step 6/11: [insituXRD] Starting in-situ XRD measurement...")
            success = perform_xrd_scan(anton_paar_xrd, xrd_scan_mode, logger, timestamp=timestamp, experiment_config_path=experiment_config_path)

        # --- End of Sequence ---
        if not success:
            raise RuntimeError("A step in the sequence failed. Check logs for details.")
        
        logger.info(program_name,"--- in-situ XRD SEQUENCE COMPLETED SUCCESSFULLY ---")

    except Exception as e:
        logger.error(program_name,f"An unexpected critical error occurred: {e}")
        raise Exception(f"An unexpected critical error occurred: {e}")


if __name__ == "__main__":
    program_name="Lab_Demo_Main"
    main_logger = TaskLogger(program_name, logLevel="INFO", modeType="real")
    
    # --- Dynamically create a timestamped folder for this run's data ---
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    # # --- Get crucible location from command-line argument ---
    # if len(sys.argv) > 1:
    #     try:
    #         crucible_idx = int(sys.argv[1])
    #         experiment_name = sys.argv[2] # experiment_name="Automated_InSitu_XRD_Run"
    #         main_logger.info(program_name, f"Using crucible location from command-line argument: {crucible_idx}")
    #     except ValueError:
    #         main_logger.warning(f"Invalid command-line argument '{sys.argv[1]}'. Using default location.")
    #         crucible_idx = DEFAULT_CRUCIBLE_LOCATION
    # else:
    #     crucible_idx = DEFAULT_CRUCIBLE_LOCATION
    #     main_logger.info(program_name, f"Using default crucible location: {crucible_idx}")

    # --- Run the main sequence ---
    # for i in range(4): # Loop through crucible locations 1 to 3

    robot = UR5eController("UR5e", ROBOT_HOST, main_logger)
    angle_changer = AngleChangerController("AngleChanger", ANGLE_CHANGER_PORT, baudrate=9600, logger=main_logger)
    screw_automater = ScrewAutomationController("ScrewAutomater", SCREW_TMCM_PORT, SCREW_LINEAR_RAIL_PORT, main_logger)
    anton_paar_xrd = AntonPaarXRDController("AntonPaarXRD", XRD_BASE_URL, main_logger)

    # -- Verify all connections are established --
    robot_ok = getattr(robot, "dashboard", None) is not None
    if not robot_ok:
        main_logger.error(program_name, "UR5e Robot connection failed.")
        raise Exception("UR5e Robot connection failed.")
    screw_ok = (screw_automater.tmcm_obj.interface and screw_automater.linear_rail_obj.board) is not None
    if not screw_ok:
        main_logger.error(program_name, "Screw Automation connection failed.")
        raise Exception("Screw Automation connection failed.")
    angle_ok = bool(getattr(getattr(angle_changer, "ser", None), "is_open", False))
    if not angle_ok:
        main_logger.error(program_name, "Angle Changer connection failed.")
        raise Exception("Angle Changer connection failed.")
    xrd_ok = anton_paar_xrd.is_connected
    if not xrd_ok:
        main_logger.error(program_name, "XRD connection failed.")
        raise Exception("XRD connection failed.")
    
    run_full_sequence(
        program_name=program_name,
        location_index=0,
        logger=main_logger,
        xrd_scan_mode="static",
        timestamp=timestamp,
        experiment_config_path="XRD_template/LLZO_1100C_v1.json"
    )

    # need to build another programm