"""End-to-end demonstration of the he11lib reconstruction pipeline.

Simulates a gyrotron beam that is mostly the LG_00 fundamental mode with a
small admixture of LG_01, viewed by a thermal camera at four distances from
the output window. The camera has an unknown transverse offset/pointing and
adds sensor noise; the target also has some thermal-diffusion blur that we
correct for. We then reconstruct the mode purity, beam center/pointing, and
plot the diagnostics.

Run with:

    python examples/full_pipeline_example.py
"""

from __future__ import annotations

import matplotlib.pyplot as plt

from he11lib import (
    BeamReconstructor,
    DiffusionDeconvolver,
    LGBasis,
    SyntheticBeamGenerator,
    plot_center_trace,
    plot_mode_purity,
    plot_residuals,
)

# --- Known reference beam parameters (from the gyrotron/mode-converter design) ---
W0 = 5e-3          # reference waist radius, meters
Z0 = 0.5           # reference waist location, meters from the output window
WAVELENGTH = 1.76e-3  # radiation wavelength, meters (e.g. a 170 GHz gyrotron)

# --- Ground truth for the synthetic beam (unknown to the reconstructor) ---
TRUE_COEFFICIENTS = {(0, 0): 0.95 + 0j, (0, 1): 0.25 + 0.05j}
TRUE_CENTER = (0.4e-3, -0.3e-3)       # beam offset from the camera's optical axis
TRUE_POINTING_DEG = 0.15              # beam pointing (tilt) angle
CAMERA_VIEWING_ANGLE_DEG = 5.0        # oblique camera viewing angle (known)
CAMERA_PIXEL_SCALE = 4e-4             # meters/pixel (known calibration)
IMAGE_SHAPE = (81, 81)
# Measurement plane distances, meters. Kept within roughly +/-2 Rayleigh
# ranges of z0 so the (widening) beam stays well within the camera frame --
# planes much farther out would be clipped by the finite frame, which
# degrades the fit.
Z_LIST = [0.4, 0.45, 0.55, 0.6]

# --- Target thermal-diffusion blur (known target material properties) ---
THERMAL_DIFFUSIVITY = 1e-6  # m^2/s
DWELL_TIME = 0.2            # s


def main() -> None:
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    generator = SyntheticBeamGenerator(
        basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=CAMERA_PIXEL_SCALE
    )

    planes = generator.generate(
        coefficients=TRUE_COEFFICIENTS,
        z_list=Z_LIST,
        center=TRUE_CENTER,
        pointing_angle_deg=TRUE_POINTING_DEG,
        viewing_angle_deg=CAMERA_VIEWING_ANGLE_DEG,
        noise_std=2e-4,
        seed=42,
    )

    # Apply the same thermal-diffusion blur a real target would exhibit.
    blur_deconvolver = DiffusionDeconvolver(
        thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME
    )
    for plane in planes:
        plane.flux = blur_deconvolver.blur(plane.flux, plane.pixel_scale)

    # The ground truth only has order-0 and order-1 content, so a max_order
    # of 1 is enough for automatic mode-set growth to find it; growing much
    # further would start fitting deconvolution/noise artifacts as spurious
    # higher-order modes.
    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        max_order=1,
        deconvolver=blur_deconvolver,
    )
    result = reconstructor.reconstruct(planes)

    print("Mode purity table (power fraction, phase [rad]):")
    for mode, (fraction, phase) in sorted(
        result.purity.items(), key=lambda item: -item[1][0]
    ):
        print(f"  LG_{mode[0]},{mode[1]}: {fraction:6.3%}  (phase {phase:+.3f} rad)")

    print(f"\nFitted pointing angle: {result.pointing_angle_deg:.4f} deg")
    print("Fitted beam center per plane (m):")
    for plane, (cx, cy) in zip(planes, result.centers):
        print(f"  z={plane.z:.2f} m -> ({cx:.3e}, {cy:.3e})")

    print(f"\nUsed phase-retrieval fallback: {result.used_phase_retrieval}")

    plot_mode_purity(result)
    plot_center_trace(planes, result)
    plot_residuals(planes, result)
    plt.show()


if __name__ == "__main__":
    main()