"""
Fringe Counting Simulation
===========================

Simulate fringe counting in a Michelson interferometer to verify your
understanding of precision wavelength measurement before working with
real data.

This script is a companion to the Week 2 prelab. It generates a synthetic
photodetector signal for a moving mirror, counts zero-crossings to determine
the number of fringes, and calculates the laser wavelength using:

    λ = 2 ΔL / N

where ΔL is the mirror travel distance and N is the number of complete
fringes. The uncertainty from fringe counting alone is:

    δλ ≈ λ / N   (from ±1 fringe uncertainty)

No hardware is required — this is pure simulation.

Key Concepts:
    - Fringe frequency:  f = 2v / λ  (from mirror velocity and wavelength)
    - A complete fringe is one full cycle of constructive → destructive →
      constructive interference, corresponding to λ/2 of mirror travel
    - Zero-crossings of the AC signal: N_fringes = N_crossings / 2
    - Longer travel distance → more fringes → better precision

Usage:
    python 01_fringe_simulation.py

Author: PHYS 4430
"""

import numpy as np
import matplotlib.pyplot as plt


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


def generate_fringes(distance_m, wavelength_m=HENE_WAVELENGTH_M,
                     sample_rate=10000, velocity_m_s=0.5e-3,
                     noise_amplitude=0.0):
    """
    Generate a synthetic photodetector signal for a Michelson interferometer.

    The signal is a cosine at the fringe frequency f = 2v/λ, representing
    the AC component of the photodetector output during constant-velocity
    mirror motion.

    Parameters:
        distance_m: Mirror travel distance (m)
        wavelength_m: Light wavelength (m), default HeNe 632.8 nm
        sample_rate: Sample rate (Hz)
        velocity_m_s: Mirror velocity (m/s)
        noise_amplitude: RMS amplitude of Gaussian noise to add (V)

    Returns:
        t: Time array (s)
        voltage: Simulated photodetector voltage (V, AC component)
    """
    duration_s = distance_m / velocity_m_s
    t = np.arange(0, duration_s, 1 / sample_rate)

    fringe_freq = 2 * velocity_m_s / wavelength_m
    voltage = np.cos(2 * np.pi * fringe_freq * t)

    if noise_amplitude > 0:
        voltage += np.random.normal(0, noise_amplitude, len(t))

    return t, voltage


def count_zero_crossings(voltage):
    """
    Count zero-crossings in a voltage signal.

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

    Note: This method counts ALL zero-crossings, including those from
    noise or environmental vibrations. For a discussion of how this
    affects real hardware measurements, see Part 3 of this script.

    Parameters:
        voltage: Voltage array (V)

    Returns:
        n_crossings: Number of zero-crossings (divide by 2 for fringes)
    """
    centered = voltage - np.mean(voltage)
    signs = np.sign(centered)
    n_crossings = int(np.sum(np.abs(np.diff(signs)) > 0))
    return 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


def main():
    """Demonstrate fringe counting at multiple distances."""
    print("=" * 60)
    print("  FRINGE COUNTING SIMULATION")
    print("  Precision wavelength measurement: λ = 2ΔL / N")
    print("=" * 60)

    # --- Part 1: Multi-distance demonstration ---
    print("\n--- Part 1: Wavelength measurement at different distances ---")
    print(f"\nTrue wavelength: {HENE_WAVELENGTH_NM} nm")

    distances_mm = [1, 2, 5, 10]
    velocity_mm_s = 0.5
    sample_rate = 10000

    results = []

    print(f"\n{'Distance':>10s}  {'Fringes':>8s}  {'λ (nm)':>10s}  "
          f"{'δλ (nm)':>10s}  {'Error (nm)':>10s}")
    print("-" * 60)

    for dist_mm in distances_mm:
        dist_m = dist_mm * 1e-3
        vel_m_s = velocity_mm_s * 1e-3

        t, voltage = generate_fringes(
            dist_m, sample_rate=sample_rate, velocity_m_s=vel_m_s
        )

        n_crossings = count_zero_crossings(voltage)
        n_fringes = n_crossings / 2
        wavelength_nm, uncertainty_nm = calculate_wavelength(dist_m, n_fringes)
        error_nm = wavelength_nm - HENE_WAVELENGTH_NM

        results.append({
            'distance_mm': dist_mm,
            'n_fringes': n_fringes,
            'wavelength_nm': wavelength_nm,
            'uncertainty_nm': uncertainty_nm,
            'error_nm': error_nm,
        })

        print(f"{dist_mm:>8d} mm  {n_fringes:>8.0f}  {wavelength_nm:>10.3f}  "
              f"{uncertainty_nm:>10.4f}  {error_nm:>+10.3f}")

    # --- Part 2: Noise sensitivity ---
    print("\n\n--- Part 2: Noise sensitivity (5 mm distance) ---")
    print("\nHow does noise affect fringe counting accuracy?")

    dist_m = 5e-3
    vel_m_s = 0.5e-3

    # Signal amplitude is 1.0 V (peak of cosine)
    noise_levels = [0, 0.05, 0.1, 0.2, 0.5, 1.0]

    print(f"\n{'Noise (V)':>10s}  {'SNR':>6s}  {'Fringes':>8s}  "
          f"{'λ (nm)':>10s}  {'Error (nm)':>10s}")
    print("-" * 55)

    for noise in noise_levels:
        snr = 1.0 / noise if noise > 0 else float('inf')

        # Average 3 trials to show typical behavior
        wavelengths = []
        for _ in range(3):
            t, voltage = generate_fringes(
                dist_m, sample_rate=sample_rate, velocity_m_s=vel_m_s,
                noise_amplitude=noise
            )
            n_crossings = count_zero_crossings(voltage)
            n_fringes = n_crossings / 2
            wl, _ = calculate_wavelength(dist_m, n_fringes)
            wavelengths.append(wl)

        mean_wl = np.mean(wavelengths)
        error = mean_wl - HENE_WAVELENGTH_NM

        snr_str = f"{snr:.0f}" if snr < 1000 else "inf"
        print(f"{noise:>10.2f}  {snr_str:>6s}  {n_fringes:>8.0f}  "
              f"{mean_wl:>10.3f}  {error:>+10.3f}")

    # --- Part 3: Buffer noise simulation ---
    print("\n\n--- Part 3: Buffer noise simulation ---")
    print("\nReal DAQ data includes ~1 s of recording before and after motor")
    print("motion. Environmental vibrations during these 'buffer' periods")
    print("produce extra zero-crossings that bias the fringe count.\n")

    buffer_duration_s = 0.5  # buffer before and after motor motion

    # Generate clean fringe signal during motor motion
    t_fringes, v_fringes = generate_fringes(
        dist_m, sample_rate=sample_rate, velocity_m_s=vel_m_s
    )
    motion_duration = len(t_fringes) / sample_rate

    # Generate buffer noise (environmental vibration, no motor)
    n_buffer = int(buffer_duration_s * sample_rate)
    np.random.seed(42)  # reproducible

    # Pre-buffer: low-frequency vibration (~15 Hz building vibration)
    t_buf = np.arange(n_buffer) / sample_rate
    vibration_freq = 15  # Hz
    pre_buffer = 0.3 * np.sin(2 * np.pi * vibration_freq * t_buf
                               + np.random.uniform(0, 2 * np.pi))
    pre_buffer += 0.1 * np.random.normal(0, 1, n_buffer)

    # Post-buffer: similar but different phase
    post_buffer = 0.2 * np.sin(2 * np.pi * vibration_freq * t_buf
                                + np.random.uniform(0, 2 * np.pi))
    post_buffer += 0.1 * np.random.normal(0, 1, n_buffer)

    # Concatenate: buffer + fringes + buffer
    full_signal = np.concatenate([pre_buffer, v_fringes, post_buffer])
    full_time = np.arange(len(full_signal)) / sample_rate

    # Count with simple method on the full signal vs. clean only
    n_crossings_full = count_zero_crossings(full_signal)
    n_fringes_full = n_crossings_full / 2
    wl_full, _ = calculate_wavelength(dist_m, n_fringes_full)

    n_crossings_clean = count_zero_crossings(v_fringes)
    n_fringes_clean = n_crossings_clean / 2
    wl_clean, _ = calculate_wavelength(dist_m, n_fringes_clean)

    n_crossings_pre = count_zero_crossings(pre_buffer)
    n_crossings_post = count_zero_crossings(post_buffer)

    print(f"  Clean signal:   {n_fringes_clean:.0f} fringes -> "
          f"lambda = {wl_clean:.3f} nm "
          f"(error: {wl_clean - HENE_WAVELENGTH_NM:+.3f} nm)")
    print(f"  Pre-buffer:     {n_crossings_pre} extra crossings")
    print(f"  Post-buffer:    {n_crossings_post} extra crossings")
    print(f"  With buffers:   {n_fringes_full:.0f} fringes -> "
          f"lambda = {wl_full:.3f} nm "
          f"(error: {wl_full - HENE_WAVELENGTH_NM:+.3f} nm)")
    extra_fringes = n_fringes_full - n_fringes_clean
    print(f"\n  Extra fringes from buffers: {extra_fringes:.0f}")
    print(f"  Wavelength biased LOW because N is too large")

    # Show distance dependence of the bias
    print(f"\n  How does this bias scale with distance?")
    print(f"  (Assuming {extra_fringes:.0f} extra fringes at every distance)\n")
    print(f"  {'Distance':>10s}  {'True N':>8s}  {'Biased N':>10s}  "
          f"{'lambda biased':>14s}  {'Error':>10s}")
    print(f"  {'-' * 60}")
    for d_mm in [1, 2, 5, 10]:
        d_m = d_mm * 1e-3
        true_n = 2 * d_m / HENE_WAVELENGTH_M
        biased_n = true_n + extra_fringes
        wl_biased = 2 * d_m / biased_n * 1e9
        error = wl_biased - HENE_WAVELENGTH_NM
        print(f"  {d_mm:>8d} mm  {true_n:>8.0f}  {biased_n:>10.0f}  "
              f"{wl_biased:>10.3f} nm   {error:>+10.3f}")
    print(f"\n  The bias is LARGEST at short distances and shrinks at long")
    print(f"  distances -- a hallmark of a constant additive offset in N.")

    # --- Part 4: Plots ---
    print("\n\nGenerating plots...")

    fig, axes = plt.subplots(2, 2, figsize=(12, 9))

    # Top left: example fringe signal (5 mm, first 5 ms)
    ax = axes[0, 0]
    t, voltage = generate_fringes(5e-3, sample_rate=sample_rate,
                                  velocity_m_s=vel_m_s)
    mask = t * 1000 <= 5  # first 5 ms
    ax.plot(t[mask] * 1000, voltage[mask], 'b-', linewidth=0.8)
    ax.set_xlabel('Time (ms)')
    ax.set_ylabel('Voltage (V)')
    ax.set_title('Simulated fringe signal (first 5 ms)')
    ax.grid(True, alpha=0.3)

    # Top right: fringe signal with noise
    ax = axes[0, 1]
    t_noisy, v_noisy = generate_fringes(5e-3, sample_rate=sample_rate,
                                        velocity_m_s=vel_m_s,
                                        noise_amplitude=0.2)
    mask = t_noisy * 1000 <= 5
    ax.plot(t_noisy[mask] * 1000, v_noisy[mask], 'b-', linewidth=0.8)
    ax.set_xlabel('Time (ms)')
    ax.set_ylabel('Voltage (V)')
    ax.set_title('Fringe signal with noise (SNR = 5)')
    ax.grid(True, alpha=0.3)

    # Bottom left: precision vs distance
    ax = axes[1, 0]
    dists = [r['distance_mm'] for r in results]
    uncerts = [r['uncertainty_nm'] for r in results]
    ax.plot(dists, uncerts, 'bo-', markersize=8)
    ax.set_xlabel('Travel distance (mm)')
    ax.set_ylabel('Uncertainty δλ (nm)')
    ax.set_title('Precision improves with distance')
    ax.set_yscale('log')
    ax.grid(True, alpha=0.3)

    # Bottom right: buffer noise visualization
    ax = axes[1, 1]
    ax.plot(full_time * 1000, full_signal, 'b-', linewidth=0.3)
    ax.axvspan(0, buffer_duration_s * 1000, alpha=0.15, color='red',
               label='Buffer (no motor)')
    ax.axvspan((buffer_duration_s + motion_duration) * 1000,
               full_time[-1] * 1000, alpha=0.15, color='red')
    ax.set_xlabel('Time (ms)')
    ax.set_ylabel('Voltage (V)')
    ax.set_title('Signal with buffer periods (first 20 ms)')
    ax.set_xlim(0, 20)
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig("fringe_simulation.png", dpi=150, bbox_inches='tight')
    print("Plot saved to: fringe_simulation.png")
    plt.show()

    print("\nDone! Compare these results to your prelab calculations.")


if __name__ == "__main__":
    main()
