"""
Motor Velocity Control for Fringe Counting
============================================

Control the Thorlabs ZST225B translation stage at constant velocity.
Use this script to verify motor control before combining with DAQ
acquisition.

What's New Compared to Gaussian Beams:
    In the Gaussian Beams lab, you used MoveTo() to step the motor to
    specific positions and wait for it to arrive. For fringe counting,
    you need the motor to move *continuously* at a constant velocity
    while the DAQ records data. This script shows the velocity-control
    pattern:

        vel_params = device.GetVelocityParams()
        vel_params.MaxVelocity = Decimal(0.5)     # mm/s
        vel_params.Acceleration = Decimal(2.0)     # mm/s^2
        device.SetVelocityParams(vel_params)

    Then MoveTo() with timeout=0 starts the move without blocking,
    so your code can do other things (like acquire data) while the
    motor is in motion.

Hardware:
    - Thorlabs KST101 K-Cube Stepper Motor Controller
    - Thorlabs ZST225B stepper motor actuator
    - USB connection to computer

Software:
    - Thorlabs Kinesis SDK (download from thorlabs.com)
    - pip install pythonnet

First-Time Setup:
    If this is your first time using the motor, you must configure the
    stage type via the KST101 front panel. See the Thorlabs Motors
    resource page for step-by-step instructions. This only needs to be
    done once per controller.

Usage:
    python 03_motor_velocity.py

Author: PHYS 4430
"""

import time

try:
    import clr
    clr.AddReference(
        r"C:\Program Files\Thorlabs\Kinesis"
        r"\Thorlabs.MotionControl.DeviceManagerCLI.dll"
    )
    clr.AddReference(
        r"C:\Program Files\Thorlabs\Kinesis"
        r"\Thorlabs.MotionControl.GenericMotorCLI.dll"
    )
    clr.AddReference(
        r"C:\Program Files\Thorlabs\Kinesis"
        r"\Thorlabs.MotionControl.KCube.StepperMotorCLI.dll"
    )
    from Thorlabs.MotionControl.DeviceManagerCLI import (
        DeviceManagerCLI, DeviceConfiguration
    )
    from Thorlabs.MotionControl.KCube.StepperMotorCLI import KCubeStepper
    from System import Decimal
    THORLABS_AVAILABLE = True
except Exception as e:
    print(f"Thorlabs libraries not available: {e}")
    print("Install Kinesis SDK from thorlabs.com and: pip install pythonnet")
    THORLABS_AVAILABLE = False


def connect_motor(serial_number):
    """
    Connect to a KST101 stepper motor controller.

    This is the same connection sequence used in the Gaussian Beams lab.

    Parameters:
        serial_number: 8-digit serial number from the KST101 LCD

    Returns:
        device: Connected KCubeStepper device object
    """
    if not THORLABS_AVAILABLE:
        raise RuntimeError("Thorlabs libraries not available")

    print(f"Connecting to motor controller {serial_number}...")
    DeviceManagerCLI.BuildDeviceList()

    device = KCubeStepper.CreateKCubeStepper(serial_number)
    device.Connect(serial_number)
    time.sleep(0.5)

    if not device.IsSettingsInitialized():
        device.WaitForSettingsInitialized(5000)

    device.StartPolling(250)
    time.sleep(0.25)
    device.EnableDevice()
    time.sleep(0.5)

    # Load motor configuration
    try:
        device.LoadMotorConfiguration(
            serial_number,
            DeviceConfiguration.DeviceSettingsUseOptionType.UseDeviceSettings
        )
    except Exception:
        device.LoadMotorConfiguration(
            serial_number,
            DeviceConfiguration.DeviceSettingsUseOptionType.UseFileSettings
        )

    info = device.GetDeviceInfo()
    print(f"  Connected: {info.Description}")
    print(f"  Position: {get_position(device):.3f} mm")

    return device


def disconnect_motor(device):
    """Disconnect from the motor controller."""
    device.StopPolling()
    device.Disconnect()
    print("Motor disconnected.")


def get_position(device):
    """
    Get current motor position in mm.

    Returns:
        Position (float, mm)
    """
    return float(str(device.Position))


def set_velocity(device, velocity_mm_s, acceleration_mm_s2=2.0):
    """
    Set motor velocity and acceleration parameters.

    This is what's new compared to Gaussian Beams — instead of just
    calling MoveTo(), you first configure how fast the motor moves.

    Parameters:
        device: Connected KCubeStepper device
        velocity_mm_s: Maximum velocity (mm/s)
        acceleration_mm_s2: Acceleration (mm/s^2)
    """
    vel_params = device.GetVelocityParams()
    vel_params.MaxVelocity = Decimal(velocity_mm_s)
    vel_params.Acceleration = Decimal(acceleration_mm_s2)
    device.SetVelocityParams(vel_params)
    print(f"  Velocity set to {velocity_mm_s} mm/s "
          f"(accel: {acceleration_mm_s2} mm/s^2)")


def move_constant_velocity(device, travel_mm):
    """
    Command a constant-velocity move over a specified distance.

    Uses MoveTo() with timeout=0 so the call returns immediately
    (non-blocking). This lets your code do other things while the
    motor is moving — like acquiring data from the DAQ.

    Parameters:
        device: Connected KCubeStepper device
        travel_mm: Distance to travel (mm, positive = forward)

    Returns:
        start_pos: Position before move (mm)
        end_target: Commanded end position (mm)
    """
    start_pos = get_position(device)
    end_target = start_pos + travel_mm

    print(f"  Moving from {start_pos:.3f} to {end_target:.3f} mm "
          f"({travel_mm:.2f} mm travel)...")

    # timeout=0 makes this non-blocking — the function returns
    # immediately while the motor continues to move
    device.MoveTo(Decimal(end_target), 0)

    return start_pos, end_target


def wait_for_move(device):
    """Wait for the motor to finish moving.

    After a non-blocking MoveTo(..., 0), the motor needs a moment before
    IsMoving becomes True. We wait up to 2 s for motion to start, then
    poll until it stops.
    """
    # Wait for motor to actually start moving
    for _ in range(40):  # up to 2 s
        if device.Status.IsMoving:
            break
        time.sleep(0.05)
    # Now wait for it to finish
    while device.Status.IsMoving:
        time.sleep(0.05)


def main():
    """Demonstrate constant-velocity motor control."""
    print("=" * 55)
    print("  MOTOR VELOCITY CONTROL FOR FRINGE COUNTING")
    print("=" * 55)

    if not THORLABS_AVAILABLE:
        print("\nThorlabs libraries not available.")
        print("Install Kinesis SDK and pythonnet.")
        return

    # Get serial number
    serial = input("\nMotor serial number: ").strip()
    if not serial:
        print("Serial number required (8-digit number on KST101 LCD).")
        return

    device = None
    try:
        device = connect_motor(serial)

        # Home the motor
        home = input("\nHome the motor? (y/n, default y): ").strip()
        if home.lower() != 'n':
            print("Homing...")
            device.Home(60000)
            while device.Status.IsHoming:
                time.sleep(0.1)
            print(f"  Homed at {get_position(device):.3f} mm")

        # Set velocity
        vel = input("\nVelocity in mm/s (default 0.5): ").strip()
        velocity = float(vel) if vel else 0.5
        set_velocity(device, velocity)

        # Move
        dist = input("Travel distance in mm (default 5.0): ").strip()
        travel = float(dist) if dist else 5.0

        expected_time = travel / velocity
        print(f"\n  Expected travel time: {expected_time:.1f} s")

        input("Press Enter to start move...")

        # Time the move
        start_time = time.time()
        start_pos, end_target = move_constant_velocity(device, travel)

        # Wait for completion
        wait_for_move(device)
        elapsed = time.time() - start_time

        # Report results
        end_pos = get_position(device)
        actual_travel = end_pos - start_pos
        actual_velocity = actual_travel / elapsed if elapsed > 0 else 0

        print(f"\n  Results:")
        print(f"    Start position:    {start_pos:.3f} mm")
        print(f"    End position:      {end_pos:.3f} mm")
        print(f"    Actual travel:     {actual_travel:.3f} mm")
        print(f"    Elapsed time:      {elapsed:.2f} s")
        print(f"    Average velocity:  {actual_velocity:.3f} mm/s")
        print(f"    Commanded:         {velocity:.3f} mm/s")

        if velocity > 0:
            error_pct = (actual_velocity - velocity) / velocity * 100
            print(f"    Velocity error:    {error_pct:+.1f}%")

    except Exception as e:
        print(f"\nError: {e}")
        import traceback
        traceback.print_exc()

    finally:
        if device is not None:
            disconnect_motor(device)

    print("\nDone!")


if __name__ == "__main__":
    main()