"""
Precision Wavelength Measurement by Fringe Counting
====================================================

Measure the wavelength of a laser by counting interference fringes in a
Michelson interferometer. This script coordinates motor motion and DAQ
acquisition, counts fringes in the recorded signal, and calculates:

    λ = 2 ΔL / N

where ΔL is the mirror travel distance (from motor encoder) and N is
the number of complete fringes (from zero-crossing count).

This is the reference implementation for the complete measurement — like
04_beam_profiler.py in the Gaussian Beams lab. You can run it directly
for your measurements or study the code and adapt it.

Theory:
    When one mirror of a Michelson interferometer moves a distance ΔL,
    the optical path difference changes by 2ΔL (round trip). Each time
    the path difference changes by one wavelength, the detector signal
    completes one full fringe (bright → dark → bright). Therefore:

        N = 2ΔL / λ   →   λ = 2ΔL / N

    The uncertainty from ±1 fringe is:  δλ ≈ λ / N

    Precision improves with distance: at 5 mm travel, ~15,800 fringes
    give δλ ≈ 0.04 nm. At 10 mm, ~31,600 fringes give δλ ≈ 0.02 nm.

Counting Methods:
    This script provides two fringe counting functions:

    count_fringes_simple(voltage):
        Naive zero-crossing detection. Subtracts the mean and counts every
        sign change. This is the starting point — it counts ALL crossings
        including those from environmental vibrations during DAQ buffer
        periods before and after motor motion.

    count_fringes_gated(voltage, sample_rate):
        Envelope-gated hysteresis. Uses a sliding-RMS envelope to identify
        the active signal region (motor moving) and only counts crossings
        there. This reduces bias from buffer noise.

    By default, the measurement series uses the simple method so you can
    observe and investigate the systematic bias before fixing it.

Synchronization:
    The DAQ must start sampling BEFORE the motor starts moving:

        1. Start DAQ acquisition (in a background thread)
        2. Brief delay to ensure DAQ is ready
        3. Start motor move (non-blocking, timeout=0)
        4. Wait for DAQ to finish collecting all samples
        5. Wait for motor to stop

Hardware:
    - Thorlabs KST101 + ZST225B stepper motor
    - NI USB-6009 DAQ
    - Michelson interferometer with HeNe laser
    - Photodetector (connected to DAQ in differential mode)

Software:
    pip install nidaqmx pythonnet numpy matplotlib

First-Time Setup:
    Motor: Configure stage type via KST101 front panel — see the
    Thorlabs Motors resource page for instructions.

    DAQ: Connect photodetector in differential mode (signal to AI0+
    on pin 2, ground to AI0- on pin 3).

Usage:
    python 04_fringe_counting.py

Author: PHYS 4430
"""

import time
import threading
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime

# Physical constants
HENE_WAVELENGTH_NM = 632.8
HENE_WAVELENGTH_M = HENE_WAVELENGTH_NM * 1e-9

# --- Hardware imports (graceful degradation) ---

try:
    import nidaqmx
    from nidaqmx.constants import AcquisitionType, TerminalConfiguration
    NIDAQMX_AVAILABLE = True
except ImportError:
    print("nidaqmx not available — install with: pip install nidaqmx")
    NIDAQMX_AVAILABLE = False

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}")
    THORLABS_AVAILABLE = False


# ---- Motor control functions ----

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

    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 {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)

    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}")
    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."""
    return float(str(device.Position))


def set_velocity(device, velocity_mm_s, acceleration_mm_s2=2.0):
    """Set motor velocity and acceleration."""
    vel_params = device.GetVelocityParams()
    vel_params.MaxVelocity = Decimal(velocity_mm_s)
    vel_params.Acceleration = Decimal(acceleration_mm_s2)
    device.SetVelocityParams(vel_params)


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.
    """
    for _ in range(40):  # up to 2 s
        if device.Status.IsMoving:
            break
        time.sleep(0.05)
    while device.Status.IsMoving:
        time.sleep(0.05)


# ---- DAQ acquisition ----

def acquire_data(daq_device, daq_channel, sample_rate, duration_s):
    """
    Acquire photodetector signal using timed acquisition.

    Uses differential mode and FINITE acquisition.

    Parameters:
        daq_device: DAQ device name (e.g., "Dev1")
        daq_channel: Analog input channel (e.g., "ai0")
        sample_rate: Sample rate (Hz)
        duration_s: Acquisition duration (s)

    Returns:
        time_array: Time values (s)
        voltage_array: Voltage values (V)
        actual_rate: Actual sample rate achieved by DAQ (Hz)
    """
    num_samples = int(sample_rate * duration_s)

    with nidaqmx.Task() as task:
        task.ai_channels.add_ai_voltage_chan(
            f"{daq_device}/{daq_channel}",
            terminal_config=TerminalConfiguration.DIFF,
            min_val=-10.0, max_val=10.0
        )
        task.timing.cfg_samp_clk_timing(
            rate=sample_rate,
            sample_mode=AcquisitionType.FINITE,
            samps_per_chan=num_samples
        )
        actual_rate = task.timing.samp_clk_rate
        if abs(actual_rate - sample_rate) / sample_rate > 0.01:
            print(f"  WARNING: DAQ coerced rate from {sample_rate} to "
                  f"{actual_rate:.1f} Hz "
                  f"({(actual_rate - sample_rate)/sample_rate*100:+.1f}%)")
        data = task.read(
            number_of_samples_per_channel=num_samples,
            timeout=duration_s + 5.0
        )

    time_array = np.arange(num_samples) / actual_rate
    voltage_array = np.array(data)
    return time_array, voltage_array, actual_rate


# ---- Signal analysis ----

def count_fringes_simple(voltage):
    """
    Count interference fringes using simple zero-crossing detection.

    Subtracts the mean and counts every sign change. Each complete fringe
    produces two zero-crossings (positive-going and negative-going), so
    N_fringes = N_crossings / 2.

    This method counts ALL zero-crossings in the signal, including those
    from environmental vibrations during buffer periods when the motor
    is not moving. See the lab guide for a discussion of how this affects
    the measurement and how to investigate it.

    Parameters:
        voltage: Voltage array (V)

    Returns:
        n_fringes: Number of complete fringes (float)
        n_crossings: Raw number of zero-crossings
    """
    centered = voltage - np.mean(voltage)
    signs = np.sign(centered)
    n_crossings = int(np.sum(np.abs(np.diff(signs)) > 0))
    n_fringes = n_crossings / 2
    return n_fringes, n_crossings


def count_fringes_gated(voltage, sample_rate, gate_fraction=0.25,
                  crossing_fraction=0.3, window_ms=10):
    """
    Count interference fringes using envelope-gated hysteresis.

    Real DAQ data includes quiet periods before and after motor motion
    where environmental vibrations cause spurious zero-crossings. This
    function computes a sliding-RMS envelope to identify the active
    signal region (motor moving), then counts crossings with hysteresis
    only where the fringe signal is strong.

    Parameters:
        voltage: Voltage array (V)
        sample_rate: DAQ sample rate (Hz)
        gate_fraction: Envelope threshold as fraction of peak RMS (default 0.25)
        crossing_fraction: Hysteresis threshold as fraction of active std (default 0.3)
        window_ms: RMS averaging window in ms (default 10)

    Returns:
        n_fringes: Number of complete fringes (float)
        n_crossings: Raw number of threshold crossings
    """
    centered = voltage - np.mean(voltage)

    # Sliding RMS envelope to find where fringes are present
    window_samples = max(int(sample_rate * window_ms / 1000), 1)
    kernel = np.ones(window_samples) / window_samples
    rms_envelope = np.sqrt(np.convolve(centered ** 2, kernel, mode='same'))

    # Gate: only count where envelope exceeds fraction of peak
    gate_threshold = gate_fraction * np.max(rms_envelope)
    active = rms_envelope > gate_threshold

    # Hysteresis threshold based on active-region amplitude
    if np.any(active):
        threshold = crossing_fraction * np.std(centered[active])
    else:
        threshold = crossing_fraction * np.std(centered)

    # Count crossings only in gated region
    n_crossings = 0
    high = False
    for i in range(len(centered)):
        if not active[i]:
            continue
        v = centered[i]
        if not high and v > threshold:
            high = True
            n_crossings += 1
        elif high and v < -threshold:
            high = False
            n_crossings += 1

    n_fringes = n_crossings / 2
    return n_fringes, n_crossings


def calculate_wavelength(distance_m, n_fringes):
    """
    Calculate wavelength from mirror travel and fringe count.

    Uses λ = 2 ΔL / N, with uncertainty δλ = λ / N from ±1 fringe.

    Parameters:
        distance_m: Mirror travel distance (m)
        n_fringes: Number of complete fringes

    Returns:
        wavelength_nm: Calculated wavelength (nm)
        uncertainty_nm: Uncertainty from ±1 fringe (nm)
    """
    if n_fringes == 0:
        return float('inf'), float('inf')

    wavelength_m = 2 * distance_m / n_fringes
    wavelength_nm = wavelength_m * 1e9
    uncertainty_nm = wavelength_nm / n_fringes
    return wavelength_nm, uncertainty_nm


# ---- Coordinated scan ----

def run_scan(device, daq_device="Dev1", daq_channel="ai0",
             velocity_mm_s=0.5, travel_mm=5.0, sample_rate=10000):
    """
    Run a coordinated motor + DAQ scan for fringe counting.

    The DAQ starts first (in a background thread), then the motor starts
    moving. Motor start and end positions are recorded from the encoder.

    Parameters:
        device: Connected KCubeStepper motor device
        daq_device: DAQ device name
        daq_channel: Analog input channel
        velocity_mm_s: Mirror velocity (mm/s)
        travel_mm: Total travel distance (mm)
        sample_rate: DAQ sample rate (Hz)

    Returns:
        time_array: Time values (s)
        voltage_array: Voltage values (V)
        start_pos_mm: Motor start position (mm)
        end_pos_mm: Motor end position (mm)
        actual_rate: Actual DAQ sample rate (Hz)
    """
    scan_duration = travel_mm / velocity_mm_s
    daq_duration = scan_duration + 1.0  # buffer

    fringe_freq = 2 * velocity_mm_s * 1e-3 / HENE_WAVELENGTH_M

    print(f"\n  Scan parameters:")
    print(f"    Velocity:         {velocity_mm_s} mm/s")
    print(f"    Travel:           {travel_mm} mm")
    print(f"    Scan duration:    {scan_duration:.1f} s")
    print(f"    Expected fringes: ~{2 * travel_mm * 1e-3 / HENE_WAVELENGTH_M:.0f}")
    print(f"    Fringe frequency: {fringe_freq:.0f} Hz")
    print(f"    Sample rate:      {sample_rate} Hz")

    if fringe_freq > sample_rate / 2 * 0.8:
        print(f"    WARNING: Fringe frequency near Nyquist limit!")

    # Set motor velocity
    set_velocity(device, velocity_mm_s)

    # Record start position from encoder
    start_pos = get_position(device)
    end_target = start_pos + travel_mm
    print(f"    Start position:   {start_pos:.4f} mm")
    print(f"    End target:       {end_target:.4f} mm")

    # --- Synchronized acquisition ---
    result = {'time': None, 'voltage': None, 'rate': None, 'error': None}

    def daq_thread_fn():
        try:
            t, v, r = acquire_data(
                daq_device, daq_channel, sample_rate, daq_duration
            )
            result['time'] = t
            result['voltage'] = v
            result['rate'] = r
        except Exception as e:
            result['error'] = str(e)

    # Step 1: Start DAQ in background thread
    print("\n  Starting DAQ acquisition...")
    daq_thread = threading.Thread(target=daq_thread_fn)
    daq_thread.start()

    # Step 2: Brief delay to ensure DAQ is sampling
    time.sleep(0.2)

    # Step 3: Start motor move (non-blocking)
    print("  Starting motor...")
    device.MoveTo(Decimal(end_target), 0)

    # Step 4: Wait for DAQ to finish
    daq_thread.join()

    # Step 5: Wait for motor to finish
    wait_for_move(device)

    # Record end position from encoder
    end_pos = get_position(device)
    print(f"  Scan complete.")
    print(f"    End position:     {end_pos:.4f} mm")
    print(f"    Actual travel:    {end_pos - start_pos:.4f} mm")

    if result['error']:
        raise RuntimeError(f"DAQ acquisition failed: {result['error']}")

    return (result['time'], result['voltage'],
            start_pos, end_pos, result['rate'])


# ---- Measurement series ----

def run_measurement_series(device, distances_mm, n_repeats=3,
                           daq_device="Dev1", daq_channel="ai0",
                           velocity_mm_s=0.5, sample_rate=10000,
                           counting_method="simple"):
    """
    Run fringe counting measurements at multiple distances.

    Homes between runs. Returns the motor to the start position after
    each measurement.

    Parameters:
        device: Connected KCubeStepper motor device
        distances_mm: List of travel distances (mm)
        n_repeats: Number of repeat measurements at each distance
        daq_device: DAQ device name
        daq_channel: Analog input channel
        velocity_mm_s: Mirror velocity (mm/s)
        sample_rate: DAQ sample rate (Hz)
        counting_method: "simple" (default) or "gated". Simple counts all
            zero-crossings; gated uses envelope detection to exclude buffer
            periods. Start with "simple" to see the raw measurement, then
            try "gated" to see how algorithmic improvements affect results.

    Returns:
        results: List of dicts with measurement results
    """
    results = []

    for dist_mm in distances_mm:
        for rep in range(n_repeats):
            print(f"\n{'='*55}")
            print(f"  Distance: {dist_mm} mm, Run {rep + 1}/{n_repeats}")
            print(f"{'='*55}")

            # Move to start position (near home)
            set_velocity(device, 1.0)  # fast travel to start
            device.MoveTo(Decimal(1.0), 60000)  # start at 1 mm
            wait_for_move(device)
            time.sleep(0.5)

            # Run measurement scan
            t, v, start_pos, end_pos, actual_rate = run_scan(
                device, daq_device=daq_device, daq_channel=daq_channel,
                velocity_mm_s=velocity_mm_s, travel_mm=dist_mm,
                sample_rate=sample_rate
            )

            # Compute distance from encoder
            actual_travel_mm = end_pos - start_pos
            actual_travel_m = actual_travel_mm * 1e-3

            # Count fringes
            if counting_method == "gated":
                n_fringes, n_crossings = count_fringes_gated(v, actual_rate)
            else:
                n_fringes, n_crossings = count_fringes_simple(v)

            # Calculate wavelength
            wavelength_nm, uncertainty_nm = calculate_wavelength(
                actual_travel_m, n_fringes
            )
            error_nm = wavelength_nm - HENE_WAVELENGTH_NM

            result = {
                'distance_mm': dist_mm,
                'repeat': rep + 1,
                'start_pos_mm': start_pos,
                'end_pos_mm': end_pos,
                'actual_travel_mm': actual_travel_mm,
                'n_fringes': n_fringes,
                'n_crossings': n_crossings,
                'wavelength_nm': wavelength_nm,
                'uncertainty_nm': uncertainty_nm,
                'error_nm': error_nm,
            }
            results.append(result)

            print(f"\n  Result:")
            print(f"    Encoder travel:   {actual_travel_mm:.4f} mm")
            print(f"    Fringes:          {n_fringes:.0f}")
            print(f"    Wavelength:       {wavelength_nm:.3f} nm")
            print(f"    Uncertainty:      ±{uncertainty_nm:.4f} nm")
            print(f"    Error:            {error_nm:+.3f} nm")

    return results


# ---- Plotting ----

def plot_single_run(time_array, voltage_array, sample_rate,
                    travel_mm, wavelength_nm, n_fringes,
                    save_path=None):
    """
    Plot time-domain signal for a single measurement run.

    Shows both the full trace (useful for identifying buffer noise at
    the start/end of acquisition) and a zoomed view of individual fringes.

    Parameters:
        time_array: Time values (s)
        voltage_array: Voltage values (V)
        sample_rate: Sample rate (Hz)
        travel_mm: Actual travel distance (mm)
        wavelength_nm: Calculated wavelength (nm)
        n_fringes: Number of fringes counted
        save_path: If provided, save figure to this path
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

    t_ms = time_array * 1000

    # Full trace — look for buffer noise at start/end
    ax1.plot(t_ms, voltage_array, 'b-', linewidth=0.5)
    ax1.set_title('Full signal — look for activity before/after motor motion')
    ax1.set_xlabel('Time (ms)')
    ax1.set_ylabel('Voltage (V)')
    ax1.grid(True, alpha=0.3)

    # Zoomed view of fringes
    if t_ms[-1] > 20:
        # Show a 10 ms window from the middle of the trace
        mid = t_ms[-1] / 2
        mask = (t_ms >= mid - 5) & (t_ms <= mid + 5)
        ax2.plot(t_ms[mask], voltage_array[mask], 'b-', linewidth=0.5)
        ax2.set_title('Zoomed view (10 ms from mid-scan)')
    else:
        ax2.plot(t_ms, voltage_array, 'b-', linewidth=0.5)
        ax2.set_title('Fringe signal')
    ax2.set_xlabel('Time (ms)')
    ax2.set_ylabel('Voltage (V)')
    ax2.grid(True, alpha=0.3)

    fig.suptitle(f'λ = 2ΔL/N = 2×{travel_mm:.3f} mm / {n_fringes:.0f} '
                 f'= {wavelength_nm:.2f} nm', fontsize=12, fontweight='bold')
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Plot saved to: {save_path}")

    plt.show()


def plot_measurement_series(results, save_path=None):
    """
    Plot wavelength vs distance with error bars from repeated measurements.

    Parameters:
        results: List of result dicts from run_measurement_series
        save_path: If provided, save figure to this path
    """
    # Group results by distance
    distances = sorted(set(r['distance_mm'] for r in results))

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    means = []
    stds = []
    predicted_uncerts = []

    for dist in distances:
        wls = [r['wavelength_nm'] for r in results
               if r['distance_mm'] == dist]
        mean_wl = np.mean(wls)
        std_wl = np.std(wls, ddof=1) if len(wls) > 1 else 0
        means.append(mean_wl)
        stds.append(std_wl)

        # Predicted uncertainty from ±1 fringe
        n_fringes_expected = 2 * dist * 1e-3 / HENE_WAVELENGTH_M
        predicted_uncerts.append(HENE_WAVELENGTH_NM / n_fringes_expected)

    # Left: wavelength vs distance with error bars
    ax1.errorbar(distances, means, yerr=stds, fmt='bo-', capsize=5,
                 markersize=8, label='Measured')
    ax1.axhline(HENE_WAVELENGTH_NM, color='r', linestyle='--', alpha=0.7,
                label=f'HeNe: {HENE_WAVELENGTH_NM} nm')
    ax1.set_xlabel('Travel distance (mm)')
    ax1.set_ylabel('Wavelength (nm)')
    ax1.set_title('Measured wavelength vs distance')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # Right: scatter vs predicted uncertainty
    ax2.plot(distances, stds, 'bo-', markersize=8, label='Observed scatter (std)')
    ax2.plot(distances, predicted_uncerts, 'r--', label='Predicted δλ = λ/N')
    ax2.set_xlabel('Travel distance (mm)')
    ax2.set_ylabel('Uncertainty (nm)')
    ax2.set_title('Precision vs distance')
    ax2.set_yscale('log')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"Plot saved to: {save_path}")

    plt.show()


# ---- Main ----

def main():
    """Interactive precision wavelength measurement."""
    print("=" * 60)
    print("  PRECISION WAVELENGTH MEASUREMENT BY FRINGE COUNTING")
    print("  λ = 2ΔL / N")
    print("=" * 60)

    # Check dependencies
    if not NIDAQMX_AVAILABLE:
        print("\nNI-DAQmx not available. Install: pip install nidaqmx")
        return
    if not THORLABS_AVAILABLE:
        print("\nThorlabs libraries not available.")
        print("Install Kinesis SDK and: pip install pythonnet")
        return

    # Get user input
    serial = input("\nMotor serial number: ").strip()
    if not serial:
        print("Serial number required.")
        return

    daq_device = input("DAQ device name (default Dev1): ").strip()
    daq_device = daq_device if daq_device else "Dev1"

    vel = input("Scan velocity in mm/s (default 0.5): ").strip()
    velocity = float(vel) if vel else 0.5

    rate = input("Sample rate in Hz (default 10000): ").strip()
    sample_rate = int(rate) if rate else 10000

    # Connect and run
    device = None
    try:
        device = connect_motor(serial)

        # Home
        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")

        # --- Phase 1: Test scan ---
        print("\n" + "=" * 55)
        print("  PHASE 1: Test Scan")
        print("=" * 55)
        print("\nRunning a short test scan to verify hardware...")

        input("Press Enter when ready (motor should be able to move 1 mm)...")

        t, v, start, end, actual_rate = run_scan(
            device, daq_device=daq_device,
            velocity_mm_s=velocity, travel_mm=1.0,
            sample_rate=sample_rate
        )

        n_fringes, n_crossings = count_fringes_simple(v)
        travel_m = (end - start) * 1e-3
        wl, _ = calculate_wavelength(travel_m, n_fringes)
        print(f"\n  Test scan result:")
        print(f"    Fringes:    {n_fringes:.0f}")
        print(f"    Wavelength: {wl:.1f} nm (expect ~632.8 nm)")
        print(f"    Pk-pk:      {np.max(v) - np.min(v):.4f} V")

        if n_fringes < 100:
            cont = input("\nVery few fringes detected — check alignment. "
                         "Continue anyway? (y/n): ")
            if cont.lower() != 'y':
                return

        # --- Phase 2: Measurement runs ---
        print("\n" + "=" * 55)
        print("  PHASE 2: Measurement Series")
        print("=" * 55)

        mode = input("\nMeasurement mode:\n"
                     "  1. Single distance (quick test)\n"
                     "  2. Multi-distance series (full characterization)\n"
                     "Choice (default 2): ").strip()

        if mode == '1':
            dist = input("Travel distance in mm (default 5.0): ").strip()
            travel = float(dist) if dist else 5.0
            reps = input("Number of repeats (default 3): ").strip()
            n_repeats = int(reps) if reps else 3

            distances_mm = [travel]
        else:
            distances_mm = [1, 2, 5, 10]
            n_repeats = 3
            print(f"\nWill measure at {distances_mm} mm, {n_repeats} repeats each")
            print(f"Total scans: {len(distances_mm) * n_repeats}")

        method = input("\nCounting method:\n"
                       "  1. Simple zero-crossing (default — recommended first)\n"
                       "  2. Envelope-gated (advanced)\n"
                       "Choice (default 1): ").strip()
        counting_method = "gated" if method == '2' else "simple"

        input("\nPress Enter to start measurement series...")

        results = run_measurement_series(
            device, distances_mm, n_repeats=n_repeats,
            daq_device=daq_device, daq_channel="ai0",
            velocity_mm_s=velocity, sample_rate=sample_rate,
            counting_method=counting_method
        )

        # --- Summary ---
        print("\n" + "=" * 60)
        print("  RESULTS SUMMARY")
        print("=" * 60)

        print(f"\n{'Dist (mm)':>10s}  {'Run':>4s}  {'Fringes':>8s}  "
              f"{'λ (nm)':>10s}  {'δλ (nm)':>10s}  {'Error':>8s}")
        print("-" * 60)
        for r in results:
            print(f"{r['distance_mm']:>8d}    {r['repeat']:>3d}  "
                  f"{r['n_fringes']:>8.0f}  {r['wavelength_nm']:>10.3f}  "
                  f"{r['uncertainty_nm']:>10.4f}  "
                  f"{r['error_nm']:>+7.3f}")

        # Group statistics
        print(f"\n{'Dist (mm)':>10s}  {'Mean λ (nm)':>12s}  "
              f"{'Std (nm)':>10s}  {'Pred δλ':>10s}")
        print("-" * 50)
        for dist in sorted(set(r['distance_mm'] for r in results)):
            wls = [r['wavelength_nm'] for r in results
                   if r['distance_mm'] == dist]
            mean_wl = np.mean(wls)
            std_wl = np.std(wls, ddof=1) if len(wls) > 1 else 0
            pred = HENE_WAVELENGTH_NM / (2 * dist * 1e-3 / HENE_WAVELENGTH_M)
            print(f"{dist:>8d}    {mean_wl:>12.3f}  "
                  f"{std_wl:>10.4f}  {pred:>10.4f}")

        # Save results
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_file = f"fringe_counting_{timestamp}.csv"
        with open(csv_file, 'w') as f:
            f.write("Distance (mm),Run,Fringes,Wavelength (nm),"
                    "Uncertainty (nm),Error (nm)\n")
            for r in results:
                f.write(f"{r['distance_mm']},{r['repeat']},"
                        f"{r['n_fringes']:.0f},{r['wavelength_nm']:.4f},"
                        f"{r['uncertainty_nm']:.5f},{r['error_nm']:.4f}\n")
        print(f"\nResults saved to: {csv_file}")

        # Plot
        if len(set(r['distance_mm'] for r in results)) > 1:
            plot_measurement_series(
                results,
                save_path=f"fringe_counting_{timestamp}.png"
            )

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

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

    print("\nMeasurement complete!")


if __name__ == "__main__":
    main()
