#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun 11 10:55:37 2026

@author: begum@mit.edu

Fit Gaussian beam parameters for measurements taken downstream of ZM2
and upstream of ZM3.

This version:
  - reads the spreadsheet from ./data/ next to this .py file
  - groups ZM2 strain gauge values by first digit / integer part
      e.g. 1.28 and 1.35 -> 1.x V
      e.g. 4.45 and 4.50 -> 4.x V
  - averages duplicate measurements at the same point within a voltage group
  - fits x and y beams separately for each voltage group
  - shows one interactive propagation plot
  - radio buttons select which ZM2 setting is displayed
  - extrapolates propagation from FC1 to ZM2
  - uses FC1 as the position reference, x = 0
  - propagates the FC eigenmode from FC1 to ZM2 using Finesse
    without forcing lho.cavFC.qx/qy at FC1.p2.o
  - propagates VOPO target and O4 beams from ZM1 to ZM2,
    applies the ZM2 mirror RoC, then propagates them to FC1
  - plots FC eigenmode, VOPO target, and VOPO O4 beams on the propagation plot
  - computes overlap between FC eigenmode / VOPO target / VOPO O4 and each SG solution at ZM3
  - prints and plots a direct FC eigenmode vs VOPO O4 overlap comparison
  - uses fixed y-limits for all SG settings
  - computes beam ellipticity/aspect ratio
  - prints and displays fitted waist size/location

Coordinate convention:
  s = distance from FC1 along the FC1 -> ZM3 -> ZM2 path [m]
  FC1 is s = 0
  ZM3 is s = ZM3_to_FC1
  ZM2 is s = ZM3_to_FC1 + ZM2_to_ZM3

Spreadsheet convention:
  x_up = distance upstream of ZM3 [m]
  Therefore:
      s = X_ZM3 + x_up

(In Spyder) use this first if the GUI does not respond:
  %matplotlib qt
"""

import os
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from matplotlib.widgets import RadioButtons, CheckButtons
from matplotlib.transforms import blended_transform_factory
from PIL import Image

try:
    from finesse import BeamParam
except ImportError:
    from finesse.gaussian import BeamParam

try:
    import tabulate
    HAVE_TABULATE = True
except ImportError:
    HAVE_TABULATE = False


# -------------------------
# Optional Finesse / LIGO imports for target beam propagation
# -------------------------

try:
    import git
    from finesse_ligo.factory import ALIGOFactory
    HAVE_LIGO_FACTORY = True
except ImportError:
    HAVE_LIGO_FACTORY = False


# -------------------------
# User settings
# -------------------------

lambda0 = 1064e-9  # [m]

# Spreadsheet path: ./data/HAM7ZM2ZM3Data_10June2026.xlsx
try:
    SCRIPT_DIR = Path(__file__).resolve().parent
except NameError:
    SCRIPT_DIR = Path.cwd()

DATA_PATH = SCRIPT_DIR / "data" / "HAM7ZM2ZM3Data_10June2026.xlsx"

# PSAM / ZM2 RoC calibration file.
# Put this file in ./data/ next to the beam-size spreadsheet.
# If it is not found there, the script will also check the script directory itself.
PSAM_CALIBRATION_FILENAME = "Actuation_Curves_ZM2_SN1.xlsx"
PSAM_CALIBRATION_PATH = SCRIPT_DIR / "data" / PSAM_CALIBRATION_FILENAME

print(f"Reading data from: {DATA_PATH}")
print(f"Reading PSAM calibration from: {PSAM_CALIBRATION_PATH}")

# The spreadsheet gives 2w in um.
INPUT_SIZE_IS_DIAMETER = True

# Measurement uncertainties are not in the spreadsheet.
USE_SIGMA = False

# Build and plot the target / FC eigenmode using Finesse.
# This requires running from inside the ligo-commissioning-modeling repo.
BUILD_TARGET_FROM_FINESSE = True
IFO = "lho"

# VOPO target/O4 beam values from the ZM1 analysis
# These are BeamParam objects referred to ZM1, using the same convention as the ZM1 code:
#   q.z = waist location relative to ZM1 along the ZM1 -> ZM2 path.
BUILD_VOPO_REFERENCES = True

# PSAM projection settings.
# New calibration from today's note:
#   SG 1.x  -> PZT 0 V
#   SG 3.1x -> PZT 90 V
# Therefore the measured 3.x beam is treated as the current 90 V state.
BUILD_PSAM_PROJECTION = True
PSAM_PZT_MIN = 0.0
PSAM_PZT_MAX = 90.0
PSAM_CURRENT_FIT_GROUP_LABEL = "3.x V"
PSAM_CURRENT_PZT_VOLTAGE = 90.0
PSAM_CURRENT_ROC_SCAN = "increasing"

# ZM2 RoC for the 3.15 V strain-gauge setting
# ZM2_ROC_3p15 = 0.86194  # [m]
ZM2_ROC_3p15 = 0.85230  # [m]

# Explicit starting-state assumptions for PSAM hysteresis projections.
# Each "family" in the projected-overlap figure uses one of these starting
# assumptions, then projects along the available increasing/decreasing scan
# RoC values from the calibration table.
PSAM_START_CASES = [
    {
        "panel": "3.x V baseline",
        "start_group_label": "3.x V",
        "start_branch": "decreasing",
        "start_pzt_v": 60.0,
        "start_roc_m": 0.85025,
    },
]

# Mirror sign convention for applying the ZM2 RoC.
# This implements: 1/q_after = 1/q_before - 2/R.
# If later we find the RoC sign convention is opposite, flip this to -1.
ZM2_ROC_SIGN = +1

# Optical layout distances
AM3_to_ZM1 = 158.2e-3
ZM1_to_ZM2 = 1498.6e-3
ZM2_to_ZM3 = 1821.5e-3
ZM3_to_FC1 = 1000.3e-3

# VOPO target/O4 beam definitions at ZM1, copied from the ZM1 analysis.
q_V0_target = BeamParam(
    w0=0.2027e-3,
    z=185.2e-3 - AM3_to_ZM1,
)

q_V0_O4_x = BeamParam(
    w0=0.2176e-3,
    z=171.5e-3 - AM3_to_ZM1,
)

q_V0_O4_y = BeamParam(
    w0=0.2104e-3,
    z=171.2e-3 - AM3_to_ZM1,
)

# New plotting coordinate:
# s = distance from FC1 along FC1 -> ZM3 -> ZM2
X_FC1 = 0.0
X_ZM3 = ZM3_to_FC1
X_ZM2 = ZM3_to_FC1 + ZM2_to_ZM3

# Waist label vertical positions in mm.
# Put them farther away from the horizontal zero line for readability.
WAIST_LABEL_Y_MM = {
    "x": -0.8,
    "y": 0.8,
}


# -------------------------
# Helper functions
# -------------------------

def print_table(rows, headers, floatfmt=".3f"):
    if HAVE_TABULATE:
        print(tabulate.tabulate(rows, headers=headers, floatfmt=floatfmt))
    else:
        print(headers)
        for row in rows:
            print(row)


def convert_size_um_to_radius_m(size_um):
    size_m = np.asarray(size_um, dtype=float) * 1e-6

    if INPUT_SIZE_IS_DIAMETER:
        return size_m / 2

    return size_m


def gaussian_beam_w(x_up, w0, x0_up):
    """
    Gaussian beam radius w(x_up).

    x_up : distance upstream of ZM3 [m]
    w0   : waist radius [m]
    x0_up: waist location in upstream-of-ZM3 coordinate [m]
    """
    zR = np.pi * w0**2 / lambda0
    return w0 * np.sqrt(1 + ((x_up - x0_up) / zR)**2)


def w_from_q_at_ZM3_along_FC1_to_ZM2(q_at_ZM3, s_array):
    """
    Evaluate a BeamParam referred to ZM3 along the FC1 -> ZM2 plotting coordinate.

    q_at_ZM3 is assumed to be oriented in the FC1 -> ZM2 propagation direction.
    s = distance from FC1
    ZM3 is at s = X_ZM3

    Moving from ZM3 by ds changes the beam-coordinate z as:
        z -> q.z + ds
    """
    return np.array([
        BeamParam(w0=q_at_ZM3.w0, z=q_at_ZM3.z + (s - X_ZM3)).w
        for s in s_array
    ])


def q_complex_from_beamparam(q):
    """
    Convert a BeamParam into complex q = z + i zR.
    """
    zR = np.pi * q.w0**2 / lambda0
    return complex(q.z, zR)


def beamparam_from_q_complex(q_complex):
    """
    Convert complex q = z + i zR into a BeamParam.
    """
    zR = np.imag(q_complex)

    if zR <= 0:
        raise ValueError(
            f"Cannot create BeamParam from q with non-positive zR: q = {q_complex}"
        )

    w0 = np.sqrt(lambda0 * zR / np.pi)
    z = np.real(q_complex)

    return BeamParam(w0=w0, z=z)


def apply_zm2_mirror_roc(q_before_zm2, roc):
    """
    Apply ZM2 mirror curvature as a thin-lens equivalent.

    Convention used here:
        1/q_after = 1/q_before - 2/R

    where R = ZM2_ROC_SIGN * roc.
    """
    R_eff = ZM2_ROC_SIGN * roc
    q_after = 1 / (1 / q_before_zm2 - 2 / R_eff)

    return q_after


def reverse_q_orientation(q):
    """
    Reverse propagation direction at the same physical plane.

    If q = z + i zR in one direction, then in the opposite direction:
        q_rev = -z + i zR
    """
    return BeamParam(w0=q.w0, z=-q.z)


def q_fit_return_at_ZM3(group_label, axis_label):
    """
    Fitted SG beam at ZM3, oriented in the physical return direction:
        ZM2 -> ZM3 -> FC1

    The fitted x0_up is the waist location upstream of ZM3.
    For the return direction, the waist is downstream in the beam coordinate,
    so q.z at ZM3 is +x0_up.
    """
    fit = fits[(group_label, axis_label)]

    return BeamParam(
        w0=fit["w0"],
        z=+fit["x0_up"],
    )


def propagate_zm1_beam_through_zm2(q_at_zm1, roc):
    """
    Propagate a ZM1-referred VOPO beam to ZM2 and apply the ZM2 RoC.

    Input q_at_zm1 uses the ZM1 code convention:
        q.z = waist location relative to ZM1 along the ZM1 -> ZM2 path.

    Therefore the actual beam coordinate at ZM2 before reflection is:
        z_at_ZM2 = ZM1_to_ZM2 - q.z

    Returns:
        BeamParam just after ZM2, oriented in the physical return direction
        ZM2 -> ZM3 -> FC1.
    """
    q_before_ZM2 = BeamParam(
        w0=q_at_zm1.w0,
        z=ZM1_to_ZM2 - q_at_zm1.z,
    )

    q_after_complex = apply_zm2_mirror_roc(
        q_complex_from_beamparam(q_before_ZM2),
        roc,
    )

    return beamparam_from_q_complex(q_after_complex)


def q_return_at_ZM3_from_q_after_ZM2(q_after_ZM2):
    """
    Propagate a return-direction q from ZM2 to ZM3.

    q_after_ZM2 is oriented ZM2 -> ZM3 -> FC1.
    """
    return BeamParam(
        w0=q_after_ZM2.w0,
        z=q_after_ZM2.z + ZM2_to_ZM3,
    )


def propagate_zm1_beam_through_zm2_to_zm3(q_at_zm1, roc):
    """
    Compatibility wrapper.

    Returns a BeamParam at ZM3 oriented in the FC1 -> ZM2 direction,
    because the plotting code for FC eigenmode and fitted SG beams uses that
    convention.
    """
    q_after_ZM2_return = propagate_zm1_beam_through_zm2(q_at_zm1, roc)
    q_at_ZM3_return = q_return_at_ZM3_from_q_after_ZM2(q_after_ZM2_return)

    return reverse_q_orientation(q_at_ZM3_return)


def w_from_return_beam_after_ZM2(q_after_ZM2, s_array):
    """
    Evaluate a beam that is physically propagating ZM2 -> ZM3 -> FC1.

    Plot coordinate:
        s = 0 at FC1
        s increases toward ZM2

    Distance propagated from ZM2 to position s is:
        d = X_ZM2 - s
    """
    return np.array([
        BeamParam(w0=q_after_ZM2.w0, z=q_after_ZM2.z + (X_ZM2 - s)).w
        for s in s_array
    ])


def waist_s_from_return_beam_after_ZM2(q_after_ZM2):
    """
    Waist location in the FC1-referenced plot coordinate for a return beam.
    """
    return X_ZM2 + q_after_ZM2.z


def aspect_ratio_from_wx_wy(wx, wy):
    """
    Beam ellipticity/aspect ratio.

    Returns max(wx, wy) / min(wx, wy).

    A circular beam has aspect_ratio = 1.
    """
    wx = float(wx)
    wy = float(wy)

    return max(wx, wy) / min(wx, wy)


def ellipticity_percent_from_wx_wy(wx, wy):
    """
    Percent ellipticity relative to a circular beam.

    0% means circular.
    10% means the larger radius is 10% larger than the smaller radius.
    """
    return 100 * (aspect_ratio_from_wx_wy(wx, wy) - 1)


def fit_beam_three_points(x_up_m, w_um, label=""):
    """
    Fit Gaussian beam radius using measurement points.
    """
    x = np.asarray(x_up_m, dtype=float)
    w = convert_size_um_to_radius_m(w_um)

    idx_min = np.argmin(w)

    # Try a few initial guesses and keep the best fit.
    # This helps avoid weird local minima with only three points.
    w0_guesses = [
        0.35 * w[idx_min],
        0.50 * w[idx_min],
        0.70 * w[idx_min],
        0.90 * w[idx_min],
    ]

    x0_guesses = [
        np.min(x) - 0.30,
        np.min(x),
        np.mean(x),
        np.max(x),
        np.max(x) + 0.30,
    ]

    best = None

    for w0_guess in w0_guesses:
        for x0_guess in x0_guesses:
            try:
                popt, pcov = curve_fit(
                    gaussian_beam_w,
                    x,
                    w,
                    p0=[w0_guess, x0_guess],
                    bounds=(
                        [1e-6, -10.0],
                        [10e-3,  10.0],
                    ),
                    maxfev=20000,
                )

                w0_fit, x0_fit = popt
                residual = w - gaussian_beam_w(x, w0_fit, x0_fit)
                rss = np.sum(residual**2)

                if best is None or rss < best["rss"]:
                    best = {
                        "popt": popt,
                        "pcov": pcov,
                        "rss": rss,
                    }

            except Exception:
                pass

    if best is None:
        raise RuntimeError(f"Fit failed for {label}")

    w0_fit, x0_fit = best["popt"]
    pcov = best["pcov"]

    zR_fit = np.pi * w0_fit**2 / lambda0

    try:
        perr = np.sqrt(np.diag(pcov))
        sigma_w0 = perr[0]
        sigma_x0 = perr[1]
    except Exception:
        sigma_w0 = np.nan
        sigma_x0 = np.nan

    residual = w - gaussian_beam_w(x, w0_fit, x0_fit)
    rms_residual = np.sqrt(np.mean(residual**2))

    # BeamParam referred to ZM3.
    # Positive x0_fit means waist upstream of ZM3,
    # so q.z at ZM3 is negative.
    q_at_ZM3 = BeamParam(w0=w0_fit, z=-x0_fit)

    # Waist position in FC1-referenced coordinate.
    waist_from_FC1 = X_ZM3 + x0_fit

    return {
        "label": label,
        "x_up_m": x,
        "w_m": w,
        "w0": w0_fit,
        "x0_up": x0_fit,
        "waist_from_FC1": waist_from_FC1,
        "zR": zR_fit,
        "q_at_ZM3": q_at_ZM3,
        "pcov": pcov,
        "sigma_w0": sigma_w0,
        "sigma_x0_up": sigma_x0,
        "rms_residual": rms_residual,
        "residual_m": residual,
    }


def point_sort_key(point_name):
    """
    Sort p22, p23, p24 by their numeric suffix.
    """
    text = str(point_name).lower().strip()
    digits = "".join(ch for ch in text if ch.isdigit())

    if digits:
        return int(digits)

    return 9999


def voltage_group_key(voltage):
    """
    Group by first digit / integer part.
    Examples:
        1.28 -> 1
        1.35 -> 1
        3.15 -> 3
        4.45 -> 4
        6.08 -> 6
    """
    return int(np.floor(float(voltage)))


def safe_get_node_q(node):
    """
    Try a few ways of reading the traced q at a Finesse node.

    This is only for diagnostics; not all Finesse node objects expose qx/qy
    directly in all versions.
    """
    out = {}

    for attr in ["qx", "qy", "q"]:
        try:
            out[attr] = getattr(node, attr)
        except Exception as exc:
            out[attr] = f"unavailable: {exc}"

    return out


def print_fc_mode_diagnostics(lho):
    """
    Diagnostic prints suggested by Kevin Kuns.

    The important point is that lho.cavFC.qx/qy are defined at the cavity
    source/reference node, which is not necessarily FC1.p2.o. Therefore we do
    not force qx_in/qy_in at FC1.p2.o. Instead, we let Finesse use the traced
    eigenmode dependency for that node.
    """
    print("\nFC eigenmode diagnostic:")

    try:
        print("  lho.cavFC.source =", lho.cavFC.source)
    except Exception as exc:
        print("  Could not print lho.cavFC.source:", exc)

    try:
        dep = lho.trace_forest.find_dependency_from_node(lho.FC1.p2.o)
        print("  dependency from FC1.p2.o =", dep)
    except Exception as exc:
        print("  Could not evaluate dependency from FC1.p2.o:", exc)

    print("\n  Local node-q diagnostic, if exposed by this Finesse version:")

    node_list = [
        ("FC1.p1.i", lho.FC1.p1.i),
        ("FC1.p1.o", lho.FC1.p1.o),
        ("FC1.p2.i", lho.FC1.p2.i),
        ("FC1.p2.o", lho.FC1.p2.o),
        ("FC2.p1.i", lho.FC2.p1.i),
        ("FC2.p1.o", lho.FC2.p1.o),
    ]

    for name, node in node_list:
        print(f"    {name}: {safe_get_node_q(node)}")


def build_target_beam_from_finesse():
    """
    Build the FC eigenmode propagated from FC1 to ZM2.

    Returns:
        qx_target_at_ZM3, qy_target_at_ZM3

    These BeamParams are referred to ZM3 and oriented in the FC1 -> ZM2 direction.

    Important convention change:
        Do NOT pass qx_in=lho.cavFC.qx or qy_in=lho.cavFC.qy here.
        lho.cavFC.qx/qy are at lho.cavFC.source, which is not guaranteed to be
        FC1.p2.o. Calling propagate_beam_astig without qx_in/qy_in lets Finesse
        use the traced eigenmode/dependency for FC1.p2.o.
    """
    if not HAVE_LIGO_FACTORY:
        raise ImportError(
            "Could not import git and/or finesse_ligo.factory.ALIGOFactory. "
            "Cannot build target beam from Finesse."
        )

    class SQZFactory(ALIGOFactory):
        pass

    repo = git.Repo(".", search_parent_directories=True)
    root_dir = repo.git.rev_parse("--show-toplevel")

    factory = SQZFactory(
        os.path.join(
            root_dir,
            IFO.upper(),
            "yaml",
            f"{IFO}_O4.yaml",
        )
    )

    factory.options.SQZ.add = True
    factory.options.SQZ.add_VIP = True

    lho = factory.make()

    print_fc_mode_diagnostics(lho)

    # Let Finesse propagate the q specified by the trace/eigenmode dependency
    # at FC1.p2.o. Do not force lho.cavFC.qx/qy at this node.
    mode_FC1_to_ZM2 = lho.propagate_beam_astig(
        lho.FC1.p2.o,
        lho.ZM2.p1.o,
    )

    qx_target_at_ZM3 = mode_FC1_to_ZM2.qx(at=lho.ZM3.p1.o)
    qy_target_at_ZM3 = mode_FC1_to_ZM2.qy(at=lho.ZM3.p1.o)

    return qx_target_at_ZM3, qy_target_at_ZM3


# -------------------------
# Build target / FC eigenmode beam from Finesse
# -------------------------

TARGET_BEAM_AVAILABLE = False
qx_target_at_ZM3 = None
qy_target_at_ZM3 = None

if BUILD_TARGET_FROM_FINESSE:
    try:
        qx_target_at_ZM3, qy_target_at_ZM3 = build_target_beam_from_finesse()
        TARGET_BEAM_AVAILABLE = True

        print("\nBuilt target / FC eigenmode from Finesse:")
        print(f"  qx target at ZM3: w0 = {1e3*qx_target_at_ZM3.w0:.4f} mm, z = {qx_target_at_ZM3.z:.4f} m")
        print(f"  qy target at ZM3: w0 = {1e3*qy_target_at_ZM3.w0:.4f} mm, z = {qy_target_at_ZM3.z:.4f} m")

    except Exception as exc:
        TARGET_BEAM_AVAILABLE = False
        print("\nWARNING: Could not build target beam from Finesse.")
        print(f"Reason: {exc}")
        print("Continuing without target beam / target overlaps.")


# -------------------------
# Build VOPO target/O4 references through ZM2 RoC
# -------------------------

VOPO_REFERENCES_AVAILABLE = False

# Return-direction beams just after ZM2.
q_vopo_target_after_ZM2_return = None
q_vopo_o4_x_after_ZM2_return = None
q_vopo_o4_y_after_ZM2_return = None

# Same physical beams referred to ZM3 in the physical return direction:
#   ZM2 -> ZM3 -> FC1
q_vopo_target_at_ZM3_return = None
q_vopo_o4_x_at_ZM3_return = None
q_vopo_o4_y_at_ZM3_return = None

# Same physical beams referred to ZM3 but reversed into the FC1 -> ZM2
# direction, useful for plotting with w_from_q_at_ZM3_along_FC1_to_ZM2.
q_vopo_target_at_ZM3 = None
q_vopo_o4_x_at_ZM3 = None
q_vopo_o4_y_at_ZM3 = None

if BUILD_VOPO_REFERENCES:
    try:
        q_vopo_target_after_ZM2_return = propagate_zm1_beam_through_zm2(
            q_V0_target,
            ZM2_ROC_3p15,
        )

        q_vopo_o4_x_after_ZM2_return = propagate_zm1_beam_through_zm2(
            q_V0_O4_x,
            ZM2_ROC_3p15,
        )

        q_vopo_o4_y_after_ZM2_return = propagate_zm1_beam_through_zm2(
            q_V0_O4_y,
            ZM2_ROC_3p15,
        )

        q_vopo_target_at_ZM3_return = q_return_at_ZM3_from_q_after_ZM2(
            q_vopo_target_after_ZM2_return
        )

        q_vopo_o4_x_at_ZM3_return = q_return_at_ZM3_from_q_after_ZM2(
            q_vopo_o4_x_after_ZM2_return
        )

        q_vopo_o4_y_at_ZM3_return = q_return_at_ZM3_from_q_after_ZM2(
            q_vopo_o4_y_after_ZM2_return
        )

        q_vopo_target_at_ZM3 = reverse_q_orientation(q_vopo_target_at_ZM3_return)
        q_vopo_o4_x_at_ZM3 = reverse_q_orientation(q_vopo_o4_x_at_ZM3_return)
        q_vopo_o4_y_at_ZM3 = reverse_q_orientation(q_vopo_o4_y_at_ZM3_return)

        VOPO_REFERENCES_AVAILABLE = True

        print("\nBuilt VOPO target/O4 references from ZM1 through ZM2 RoC:")
        print(f"  ZM2 RoC used: {ZM2_ROC_SIGN * ZM2_ROC_3p15:.5f} m")
        print("  Return-direction q's at ZM3, for overlap with fitted SG return q:")
        print(f"    VOPO target: w0 = {1e3*q_vopo_target_at_ZM3_return.w0:.4f} mm, z = {q_vopo_target_at_ZM3_return.z:.4f} m")
        print(f"    VOPO O4 x:  w0 = {1e3*q_vopo_o4_x_at_ZM3_return.w0:.4f} mm, z = {q_vopo_o4_x_at_ZM3_return.z:.4f} m")
        print(f"    VOPO O4 y:  w0 = {1e3*q_vopo_o4_y_at_ZM3_return.w0:.4f} mm, z = {q_vopo_o4_y_at_ZM3_return.z:.4f} m")
        print("  FC1 -> ZM2 q's at ZM3, for plotting:")
        print(f"    VOPO target: w0 = {1e3*q_vopo_target_at_ZM3.w0:.4f} mm, z = {q_vopo_target_at_ZM3.z:.4f} m")
        print(f"    VOPO O4 x:  w0 = {1e3*q_vopo_o4_x_at_ZM3.w0:.4f} mm, z = {q_vopo_o4_x_at_ZM3.z:.4f} m")
        print(f"    VOPO O4 y:  w0 = {1e3*q_vopo_o4_y_at_ZM3.w0:.4f} mm, z = {q_vopo_o4_y_at_ZM3.z:.4f} m")

    except Exception as exc:
        VOPO_REFERENCES_AVAILABLE = False
        print("\nWARNING: Could not build VOPO target/O4 references through ZM2 RoC.")
        print(f"Reason: {exc}")
        print("Continuing without VOPO target/O4 overlaps.")


# -------------------------
# Load and clean spreadsheet
# -------------------------

df_raw = pd.read_excel(DATA_PATH)

# Expected columns:
#   point
#   Z(mm) upstream of ZM3
#   2w(um) X
#   2w(um) Y
#   ZM2 Strain Gauge (V)
#
# Use the first five columns, so the exact column names do not matter too much.
df = df_raw.iloc[:, :5].copy()

df.columns = [
    "point",
    "z_upstream_mm",
    "diam_x_um",
    "diam_y_um",
    "voltage",
]

df["point"] = (
    df["point"]
    .astype(str)
    .str.strip()
    .str.lower()
)

for col in ["z_upstream_mm", "diam_x_um", "diam_y_um", "voltage"]:
    df[col] = pd.to_numeric(df[col], errors="coerce")

df = df.dropna(
    subset=[
        "point",
        "z_upstream_mm",
        "diam_x_um",
        "diam_y_um",
        "voltage",
    ]
).copy()

df["x_up_m"] = df["z_upstream_mm"] * 1e-3
df["s_from_FC1_m"] = X_ZM3 + df["x_up_m"]

df["voltage_group"] = df["voltage"].apply(voltage_group_key)
df["voltage_group_label"] = df["voltage_group"].apply(lambda v: f"{v}.x V")

point_order = sorted(df["point"].unique(), key=point_sort_key)

print("\nLoaded data:")
print(df[[
    "point",
    "z_upstream_mm",
    "s_from_FC1_m",
    "diam_x_um",
    "diam_y_um",
    "voltage",
    "voltage_group_label",
]])


# -------------------------
# Group data by voltage first digit
# -------------------------
#
# If multiple rows exist for the same point and same voltage group,
# average them. This implements the assumption:
#   1.2 and 1.3 are the same setting, etc.

grouped_df = (
    df
    .groupby(["voltage_group", "voltage_group_label", "point"], as_index=False)
    .agg(
        z_upstream_mm=("z_upstream_mm", "mean"),
        x_up_m=("x_up_m", "mean"),
        s_from_FC1_m=("s_from_FC1_m", "mean"),
        diam_x_um=("diam_x_um", "mean"),
        diam_y_um=("diam_y_um", "mean"),
        voltage_mean=("voltage", "mean"),
        voltage_min=("voltage", "min"),
        voltage_max=("voltage", "max"),
        n_rows=("voltage", "size"),
    )
)

print("\nGrouped data, after averaging duplicate point/group rows:")
print(grouped_df[[
    "voltage_group_label",
    "point",
    "z_upstream_mm",
    "s_from_FC1_m",
    "diam_x_um",
    "diam_y_um",
    "voltage_mean",
    "voltage_min",
    "voltage_max",
    "n_rows",
]])


# -------------------------
# Fit x/y beams for each complete voltage group
# -------------------------

fits = {}
fit_group_labels = []
grouped_data = {}

for voltage_group in sorted(grouped_df["voltage_group"].unique()):
    group_label = f"{voltage_group}.x V"

    group_data = grouped_df[grouped_df["voltage_group"] == voltage_group].copy()
    group_data["point_sort"] = group_data["point"].apply(point_sort_key)
    group_data = group_data.sort_values("point_sort").reset_index(drop=True)

    points_present = list(group_data["point"].values)

    missing_points = [
        point for point in point_order
        if point not in points_present
    ]

    if len(group_data) < 3:
        print(f"\nSkipping {group_label}: only {len(group_data)} points found.")
        continue

    if len(missing_points) > 0:
        print(f"\nWARNING for {group_label}: missing points {missing_points}")
        print("Fitting with available points only.")

    grouped_data[group_label] = group_data
    fit_group_labels.append(group_label)

    fits[(group_label, "x")] = fit_beam_three_points(
        group_data["x_up_m"].values,
        group_data["diam_x_um"].values,
        label=f"{group_label} x",
    )

    fits[(group_label, "y")] = fit_beam_three_points(
        group_data["x_up_m"].values,
        group_data["diam_y_um"].values,
        label=f"{group_label} y",
    )


# -------------------------
# Print fit table
# -------------------------

fit_rows = []

for group_label in fit_group_labels:
    group_data = grouped_data[group_label]
    voltage_mean = np.mean(group_data["voltage_mean"].values)

    for axis_label in ["x", "y"]:
        fit = fits[(group_label, axis_label)]

        fit_rows.append([
            group_label,
            voltage_mean,
            axis_label,
            1e3 * fit["w0"],
            1e3 * fit["sigma_w0"],
            1e3 * fit["waist_from_FC1"],
            1e3 * fit["x0_up"],
            1e3 * fit["sigma_x0_up"],
            1e3 * fit["zR"],
            1e6 * fit["rms_residual"],
            fit["q_at_ZM3"].z,
        ])

print("\nGaussian beam fit results:")
print_table(
    fit_rows,
    headers=[
        "ZM2 SG group",
        "mean V",
        "axis",
        "w0 [mm]",
        "sigma w0 [mm]",
        "waist from FC1 [mm]",
        "waist upstream of ZM3 [mm]",
        "sigma waist [mm]",
        "zR [mm]",
        "RMS residual [um]",
        "q_at_ZM3.z [m]",
    ],
    floatfmt=".4f",
)


# -------------------------
# Print compact waist solution table
# -------------------------

waist_rows = []

for group_label in fit_group_labels:
    group_data = grouped_data[group_label]
    voltage_mean = np.mean(group_data["voltage_mean"].values)

    fx = fits[(group_label, "x")]
    fy = fits[(group_label, "y")]

    waist_rows.append([
        group_label,
        voltage_mean,
        1e3 * fx["w0"],
        1e3 * fx["waist_from_FC1"],
        1e3 * fy["w0"],
        1e3 * fy["waist_from_FC1"],
    ])

print("\nCompact waist solutions, using FC1 as position reference:")
print_table(
    waist_rows,
    headers=[
        "ZM2 SG group",
        "mean V",
        "x w0 [mm]",
        "x waist from FC1 [mm]",
        "y w0 [mm]",
        "y waist from FC1 [mm]",
    ],
    floatfmt=".3f",
)


# -------------------------
# Create named q objects for later use
# -------------------------
#
# These are BeamParam objects referred to ZM3.
# The keys are like "1.x V", "3.x V", etc.

q_fit_by_group = {}

for group_label in fit_group_labels:
    q_fit_by_group[group_label] = {
        "x": fits[(group_label, "x")]["q_at_ZM3"],
        "y": fits[(group_label, "y")]["q_at_ZM3"],
    }


# -------------------------
# Target overlap with each SG solution
# -------------------------

def compute_target_overlap_for_group(group_label):
    """
    Compute overlap between target / FC eigenmode and fitted SG solution at ZM3.
    """
    if not TARGET_BEAM_AVAILABLE:
        return {
            "x_overlap": np.nan,
            "y_overlap": np.nan,
            "mean_overlap": np.nan,
        }

    fx = fits[(group_label, "x")]
    fy = fits[(group_label, "y")]

    x_overlap = BeamParam.overlap(qx_target_at_ZM3, fx["q_at_ZM3"])
    y_overlap = BeamParam.overlap(qy_target_at_ZM3, fy["q_at_ZM3"])

    mean_overlap = 0.5 * (x_overlap + y_overlap)

    return {
        "x_overlap": x_overlap,
        "y_overlap": y_overlap,
        "mean_overlap": mean_overlap,
    }


target_overlap_rows = []

if TARGET_BEAM_AVAILABLE:
    for group_label in fit_group_labels:
        group_data = grouped_data[group_label]
        voltage_mean = np.mean(group_data["voltage_mean"].values)

        overlaps = compute_target_overlap_for_group(group_label)

        target_overlap_rows.append([
            group_label,
            voltage_mean,
            100 * overlaps["x_overlap"],
            100 * overlaps["y_overlap"],
            100 * overlaps["mean_overlap"],
        ])

    print("\nOverlap with target / FC eigenmode at ZM3:")
    print_table(
        target_overlap_rows,
        headers=[
            "ZM2 SG group",
            "mean V",
            "x overlap [%]",
            "y overlap [%]",
            "mean overlap [%]",
        ],
        floatfmt=".3f",
    )




# -------------------------
# Direct FC eigenmode overlap table
# -------------------------

fc_overlap_by_group = {}

if TARGET_BEAM_AVAILABLE:
    for group_label in fit_group_labels:
        fc_overlap_by_group[group_label] = compute_target_overlap_for_group(group_label)


# -------------------------
# VOPO target/O4 overlap with each SG solution
# -------------------------

def compute_vopo_overlap_for_group(group_label):
    """
    Compute overlap between fitted SG solution and VOPO target/O4 beams at ZM3.

    This uses the physical return direction for both beams:
        ZM2 -> ZM3 -> FC1

    This avoids the sign-convention mistake where two beams can look close in
    radius but have a poor q-overlap because their q.z values are written in
    opposite propagation directions.
    """
    if not VOPO_REFERENCES_AVAILABLE:
        return {
            "target_x_overlap": np.nan,
            "target_y_overlap": np.nan,
            "target_mean_overlap": np.nan,
            "o4_x_overlap": np.nan,
            "o4_y_overlap": np.nan,
            "o4_mean_overlap": np.nan,
        }

    qx_fit_return = q_fit_return_at_ZM3(group_label, "x")
    qy_fit_return = q_fit_return_at_ZM3(group_label, "y")

    # VOPO target is circular/single-q, so compare it against both fitted axes.
    target_x_overlap = BeamParam.overlap(
        q_vopo_target_at_ZM3_return,
        qx_fit_return,
    )

    target_y_overlap = BeamParam.overlap(
        q_vopo_target_at_ZM3_return,
        qy_fit_return,
    )

    target_mean_overlap = 0.5 * (target_x_overlap + target_y_overlap)

    # VOPO O4 has x/y q's.
    o4_x_overlap = BeamParam.overlap(
        q_vopo_o4_x_at_ZM3_return,
        qx_fit_return,
    )

    o4_y_overlap = BeamParam.overlap(
        q_vopo_o4_y_at_ZM3_return,
        qy_fit_return,
    )

    o4_mean_overlap = 0.5 * (o4_x_overlap + o4_y_overlap)

    return {
        "target_x_overlap": target_x_overlap,
        "target_y_overlap": target_y_overlap,
        "target_mean_overlap": target_mean_overlap,
        "o4_x_overlap": o4_x_overlap,
        "o4_y_overlap": o4_y_overlap,
        "o4_mean_overlap": o4_mean_overlap,
    }


def q_return_at_s_from_ZM3(q_at_ZM3_return, s_position):
    """
    Propagate a return-direction q from ZM3 to a plane in the FC1-referenced coordinate.

    q_at_ZM3_return is oriented in the physical return direction:
        ZM2 -> ZM3 -> FC1

    Plot coordinate:
        s = 0 at FC1
        s = X_ZM3 at ZM3
        s = X_ZM2 at ZM2

    The signed propagation distance from ZM3 to plane s in the return direction is:
        d = X_ZM3 - s
    """
    return BeamParam(
        w0=q_at_ZM3_return.w0,
        z=q_at_ZM3_return.z + (X_ZM3 - s_position),
    )


def compute_vopo_overlap_at_planes_for_group(group_label):
    """
    Compute VOPO target/O4 overlap with one SG setting at ZM2, ZM3, and FC1.

    For correct q conventions, the three overlap values should be identical
    to numerical precision because this is common free-space propagation.
    """
    if not VOPO_REFERENCES_AVAILABLE:
        return {}

    qx_fit_ZM3 = q_fit_return_at_ZM3(group_label, "x")
    qy_fit_ZM3 = q_fit_return_at_ZM3(group_label, "y")

    planes = {
        "ZM2": X_ZM2,
        "ZM3": X_ZM3,
        "FC1": X_FC1,
    }

    out = {}

    for plane_name, s_plane in planes.items():
        qx_fit_plane = q_return_at_s_from_ZM3(qx_fit_ZM3, s_plane)
        qy_fit_plane = q_return_at_s_from_ZM3(qy_fit_ZM3, s_plane)

        q_target_plane = q_return_at_s_from_ZM3(
            q_vopo_target_at_ZM3_return,
            s_plane,
        )

        q_o4_x_plane = q_return_at_s_from_ZM3(
            q_vopo_o4_x_at_ZM3_return,
            s_plane,
        )

        q_o4_y_plane = q_return_at_s_from_ZM3(
            q_vopo_o4_y_at_ZM3_return,
            s_plane,
        )

        target_x = BeamParam.overlap(q_target_plane, qx_fit_plane)
        target_y = BeamParam.overlap(q_target_plane, qy_fit_plane)
        target_mean = 0.5 * (target_x + target_y)

        o4_x = BeamParam.overlap(q_o4_x_plane, qx_fit_plane)
        o4_y = BeamParam.overlap(q_o4_y_plane, qy_fit_plane)
        o4_mean = 0.5 * (o4_x + o4_y)

        out[plane_name] = {
            "target_x_overlap": target_x,
            "target_y_overlap": target_y,
            "target_mean_overlap": target_mean,
            "o4_x_overlap": o4_x,
            "o4_y_overlap": o4_y,
            "o4_mean_overlap": o4_mean,
        }

    return out


vopo_overlap_rows = []

if VOPO_REFERENCES_AVAILABLE:
    for group_label in fit_group_labels:
        group_data = grouped_data[group_label]
        voltage_mean = np.mean(group_data["voltage_mean"].values)

        overlaps = compute_vopo_overlap_for_group(group_label)

        vopo_overlap_rows.append([
            group_label,
            voltage_mean,
            100 * overlaps["target_x_overlap"],
            100 * overlaps["target_y_overlap"],
            100 * overlaps["target_mean_overlap"],
            100 * overlaps["o4_x_overlap"],
            100 * overlaps["o4_y_overlap"],
            100 * overlaps["o4_mean_overlap"],
        ])

    print("\nOverlap with VOPO target/O4 propagated through ZM2 RoC at ZM3:")
    print_table(
        vopo_overlap_rows,
        headers=[
            "ZM2 SG group",
            "mean V",
            "VOPO target x [%]",
            "VOPO target y [%]",
            "VOPO target mean [%]",
            "VOPO O4 x [%]",
            "VOPO O4 y [%]",
            "VOPO O4 mean [%]",
        ],
        floatfmt=".3f",
    )

    # Common-propagation diagnostic: ZM2, ZM3, and FC1 should agree.
    plane_overlap_rows = []
    plane_spread_rows = []

    for group_label in fit_group_labels:
        plane_overlaps = compute_vopo_overlap_at_planes_for_group(group_label)

        target_means = []
        o4_means = []

        for plane_name in ["ZM2", "ZM3", "FC1"]:
            vals = plane_overlaps[plane_name]

            target_means.append(100 * vals["target_mean_overlap"])
            o4_means.append(100 * vals["o4_mean_overlap"])

            plane_overlap_rows.append([
                group_label,
                plane_name,
                100 * vals["target_x_overlap"],
                100 * vals["target_y_overlap"],
                100 * vals["target_mean_overlap"],
                100 * vals["o4_x_overlap"],
                100 * vals["o4_y_overlap"],
                100 * vals["o4_mean_overlap"],
            ])

        plane_spread_rows.append([
            group_label,
            max(target_means) - min(target_means),
            max(o4_means) - min(o4_means),
        ])

    print("\nVOPO overlap common-propagation diagnostic at ZM2, ZM3, and FC1:")
    print_table(
        plane_overlap_rows,
        headers=[
            "ZM2 SG group",
            "plane",
            "target x [%]",
            "target y [%]",
            "target mean [%]",
            "O4 x [%]",
            "O4 y [%]",
            "O4 mean [%]",
        ],
        floatfmt=".6f",
    )

    print("\nOverlap spread across ZM2/ZM3/FC1; should be ~0 if q conventions are consistent:")
    print_table(
        plane_spread_rows,
        headers=[
            "ZM2 SG group",
            "target mean spread [percentage points]",
            "O4 mean spread [percentage points]",
        ],
        floatfmt=".6e",
    )

    if "3.x V" in fit_group_labels:
        print("\nSign-convention diagnostic for 3.x V against VOPO O4:")
        for axis_label, q_ref_return in [
            ("x", q_vopo_o4_x_at_ZM3_return),
            ("y", q_vopo_o4_y_at_ZM3_return),
        ]:
            fit = fits[("3.x V", axis_label)]

            q_fit_wrong_sign = BeamParam(
                w0=fit["w0"],
                z=-fit["x0_up"],
            )

            q_fit_right_sign = BeamParam(
                w0=fit["w0"],
                z=+fit["x0_up"],
            )

            ov_wrong = BeamParam.overlap(q_ref_return, q_fit_wrong_sign)
            ov_right = BeamParam.overlap(q_ref_return, q_fit_right_sign)

            print(
                f"  {axis_label}: "
                f"wrong sign overlap = {100*ov_wrong:.2f}%, "
                f"right sign overlap = {100*ov_right:.2f}%"
            )




# -------------------------
# Direct comparison: FC eigenmode vs VOPO O4 overlaps
# -------------------------

fc_vs_o4_overlap_rows = []

if TARGET_BEAM_AVAILABLE and VOPO_REFERENCES_AVAILABLE:
    for group_label in fit_group_labels:
        group_data = grouped_data[group_label]
        voltage_mean = np.mean(group_data["voltage_mean"].values)

        fc_overlaps = compute_target_overlap_for_group(group_label)
        o4_overlaps = compute_vopo_overlap_for_group(group_label)

        fc_vs_o4_overlap_rows.append([
            group_label,
            voltage_mean,
            100 * fc_overlaps["x_overlap"],
            100 * fc_overlaps["y_overlap"],
            100 * fc_overlaps["mean_overlap"],
            100 * o4_overlaps["o4_x_overlap"],
            100 * o4_overlaps["o4_y_overlap"],
            100 * o4_overlaps["o4_mean_overlap"],
            100 * o4_overlaps["o4_mean_overlap"] - 100 * fc_overlaps["mean_overlap"],
        ])

    print("\nDirect overlap comparison: fitted SG beams vs FC eigenmode and VOPO O4:")
    print_table(
        fc_vs_o4_overlap_rows,
        headers=[
            "ZM2 SG group",
            "mean V",
            "FC x [%]",
            "FC y [%]",
            "FC mean [%]",
            "O4 x [%]",
            "O4 y [%]",
            "O4 mean [%]",
            "O4 mean - FC mean [pp]",
        ],
        floatfmt=".3f",
    )



# -------------------------
# PSAM calibration and projected-overlap estimate
# -------------------------

def resolve_optional_data_path(path):
    """
    Resolve a data path.

    First try the requested path, usually SCRIPT_DIR / "data" / filename.
    If not present, also try SCRIPT_DIR / filename.
    """
    path = Path(path)

    if path.exists():
        return path

    alt = SCRIPT_DIR / path.name

    if alt.exists():
        return alt

    return path


def read_psam_roc_calibration(path):
    """
    Read the ZM2 PSAM calibration spreadsheet.

    The spreadsheet contains two sections:
        INCREASING PZT VOLTAGE
        DECREASING PZT VOLTAGE

    We intentionally ignore the old strain-gauge column and keep only:
        PZT voltage [V]
        RoC [m]

    Returns a dict with keys "increasing" and "decreasing".
    """
    path = resolve_optional_data_path(path)

    if not path.exists():
        raise FileNotFoundError(
            f"Could not find PSAM calibration file: {path}\n"
            f"Expected it at ./data/{PSAM_CALIBRATION_FILENAME} or next to this script."
        )

    raw = pd.read_excel(path, sheet_name=0, header=None)

    sections = {}
    current_section = None

    for _, row in raw.iterrows():
        first = row.iloc[0]

        if isinstance(first, str):
            text = first.strip().lower()

            if "increasing pzt voltage" in text:
                current_section = "increasing"
                sections[current_section] = []
                continue

            if "decreasing pzt voltage" in text:
                current_section = "decreasing"
                sections[current_section] = []
                continue

        if current_section is None:
            continue

        pzt = pd.to_numeric(row.iloc[0], errors="coerce")
        roc = pd.to_numeric(row.iloc[4], errors="coerce")

        if np.isfinite(pzt) and np.isfinite(roc):
            sections[current_section].append(
                {
                    "scan": current_section,
                    "pzt_voltage_v": float(pzt),
                    "roc_m": float(roc),
                }
            )

    out = {}

    for scan_name, rows in sections.items():
        scan_df = pd.DataFrame(rows)

        if len(scan_df) == 0:
            out[scan_name] = pd.DataFrame(
                columns=["scan", "pzt_voltage_v", "roc_m"]
            )
            continue

        scan_df = scan_df.drop_duplicates(
            subset=["pzt_voltage_v"],
            keep="first",
        )

        scan_df = scan_df.sort_values("pzt_voltage_v").reset_index(drop=True)

        out[scan_name] = scan_df

    for required_scan in ["increasing", "decreasing"]:
        if required_scan not in out:
            out[required_scan] = pd.DataFrame(
                columns=["scan", "pzt_voltage_v", "roc_m"]
            )

    return out


def filter_psam_scan_to_available_range(scan_df):
    """
    Keep only available PZT voltage points in the requested range.

    No interpolation is done here; we only select points present in the Excel file.
    """
    if len(scan_df) == 0:
        return scan_df.copy()

    return (
        scan_df[
            (scan_df["pzt_voltage_v"] >= PSAM_PZT_MIN)
            & (scan_df["pzt_voltage_v"] <= PSAM_PZT_MAX)
        ]
        .copy()
        .sort_values("pzt_voltage_v")
        .reset_index(drop=True)
    )


def make_available_psam_roc_table(psam_curves):
    """
    Build a compact table with one row per available PZT voltage.
    """
    inc = filter_psam_scan_to_available_range(psam_curves["increasing"])
    dec = filter_psam_scan_to_available_range(psam_curves["decreasing"])

    inc_small = inc[["pzt_voltage_v", "roc_m"]].rename(
        columns={"roc_m": "roc_increasing_m"}
    )

    dec_small = dec[["pzt_voltage_v", "roc_m"]].rename(
        columns={"roc_m": "roc_decreasing_m"}
    )

    table = pd.merge(
        inc_small,
        dec_small,
        on="pzt_voltage_v",
        how="outer",
    )

    table = table.sort_values("pzt_voltage_v").reset_index(drop=True)

    return table


def undo_zm2_mirror_roc(q_after_ZM2_return, roc):
    """
    Undo the ZM2 mirror curvature used to produce q_after_ZM2_return.

    Forward convention in this script:
        1/q_after = 1/q_before - 2/R

    Therefore:
        1/q_before = 1/q_after + 2/R

    q_after_ZM2_return is the measured beam just after ZM2, oriented in the
    physical return direction ZM2 -> ZM3 -> FC1.
    """
    R_eff = ZM2_ROC_SIGN * roc
    q_after_complex = q_complex_from_beamparam(q_after_ZM2_return)

    q_before_complex = 1 / (1 / q_after_complex + 2 / R_eff)

    return beamparam_from_q_complex(q_before_complex)


def apply_zm2_roc_to_before_beam(q_before_ZM2, roc):
    """
    Apply a new ZM2 RoC to the reconstructed incoming beam at ZM2.

    Returns a return-direction BeamParam just after ZM2.
    """
    q_after_complex = apply_zm2_mirror_roc(
        q_complex_from_beamparam(q_before_ZM2),
        roc,
    )

    return beamparam_from_q_complex(q_after_complex)


def current_fit_group_for_psam_projection():
    """
    Choose the measured beam used as the starting point for PSAM projection.

    By default this is 3.x V, because the new calibration says 3.1x corresponds
    to PZT 90 V. If the exact label is missing, use the first label starting
    with '3.'.
    """
    if PSAM_CURRENT_FIT_GROUP_LABEL in fit_group_labels:
        return PSAM_CURRENT_FIT_GROUP_LABEL

    for label in fit_group_labels:
        if str(label).startswith("3."):
            print(
                f"\nWARNING: {PSAM_CURRENT_FIT_GROUP_LABEL} not found; "
                f"using {label} for PSAM projection."
            )
            return label

    raise RuntimeError(
        "Could not find a 3.x V fit group to use as the PSAM projection baseline."
    )


def current_roc_for_psam_projection(psam_curves):
    """
    Get the current RoC used for undoing the measured 3.x beam.

    New calibration says the 3.1x measurement corresponds to PZT 90 V. The user
    specifically asked to undo the RoC due to the 3.1x SG increasing-scan value,
    so we use the increasing-scan RoC at 90 V.
    """
    scan_df = psam_curves[PSAM_CURRENT_ROC_SCAN]

    if len(scan_df) == 0:
        raise RuntimeError(f"No PSAM calibration rows for scan {PSAM_CURRENT_ROC_SCAN}")

    pzt_vals = np.asarray(scan_df["pzt_voltage_v"], dtype=float)
    idx = np.argmin(np.abs(pzt_vals - PSAM_CURRENT_PZT_VOLTAGE))

    chosen_pzt = float(scan_df.iloc[idx]["pzt_voltage_v"])
    chosen_roc = float(scan_df.iloc[idx]["roc_m"])

    if abs(chosen_pzt - PSAM_CURRENT_PZT_VOLTAGE) > 1e-9:
        print(
            f"\nWARNING: no exact PZT={PSAM_CURRENT_PZT_VOLTAGE:.1f} V point "
            f"in {PSAM_CURRENT_ROC_SCAN} scan. Using nearest available "
            f"PZT={chosen_pzt:.1f} V."
        )

    return chosen_roc, chosen_pzt


def project_fit_to_new_psam_roc(group_label, axis_label, current_roc, new_roc):
    """
    Project a measured fit to a new ZM2 RoC.

    Procedure:
        fitted measured q at ZM3, return direction
        -> propagate back to ZM2
        -> undo current ZM2 RoC
        -> apply new ZM2 RoC
        -> propagate forward to ZM3

    Returns:
        q_projected_at_ZM3_return
    """
    q_current_ZM3_return = q_fit_return_at_ZM3(group_label, axis_label)

    # Back-propagate from ZM3 to ZM2 in the return-direction q convention.
    q_current_ZM2_return = q_return_at_s_from_ZM3(
        q_current_ZM3_return,
        X_ZM2,
    )

    # Reconstruct the incoming beam just before the current ZM2 reflection.
    q_before_ZM2 = undo_zm2_mirror_roc(
        q_current_ZM2_return,
        current_roc,
    )

    # Apply the candidate RoC.
    q_new_after_ZM2_return = apply_zm2_roc_to_before_beam(
        q_before_ZM2,
        new_roc,
    )

    # Propagate the candidate reflected beam to ZM3.
    q_new_ZM3_return = q_return_at_ZM3_from_q_after_ZM2(
        q_new_after_ZM2_return
    )

    return q_new_ZM3_return


def compute_projected_psam_overlap_for_roc(group_label, current_roc, new_roc):
    """
    Compute projected x/y/mean overlap with the FC eigenmode for one candidate RoC.
    """
    if not TARGET_BEAM_AVAILABLE:
        return {
            "x_overlap": np.nan,
            "y_overlap": np.nan,
            "mean_overlap": np.nan,
        }

    qx_proj_return = project_fit_to_new_psam_roc(
        group_label,
        "x",
        current_roc,
        new_roc,
    )

    qy_proj_return = project_fit_to_new_psam_roc(
        group_label,
        "y",
        current_roc,
        new_roc,
    )

    # FC eigenmode q's are represented in the FC1 -> ZM2 direction.
    # Reverse the projected return-direction beams before comparing.
    qx_proj_fc1_to_zm2 = reverse_q_orientation(qx_proj_return)
    qy_proj_fc1_to_zm2 = reverse_q_orientation(qy_proj_return)

    x_overlap = BeamParam.overlap(qx_target_at_ZM3, qx_proj_fc1_to_zm2)
    y_overlap = BeamParam.overlap(qy_target_at_ZM3, qy_proj_fc1_to_zm2)

    return {
        "x_overlap": x_overlap,
        "y_overlap": y_overlap,
        "mean_overlap": 0.5 * (x_overlap + y_overlap),
    }


def build_psam_projected_overlap_table(psam_curves, current_group_label, current_roc):
    """
    Build projected overlap table for all available PZT/RoC values in 0--90 V.
    """
    rows = []

    for scan_name in ["increasing", "decreasing"]:
        scan_df = filter_psam_scan_to_available_range(psam_curves[scan_name])

        for _, row in scan_df.iterrows():
            pzt_voltage = float(row["pzt_voltage_v"])
            roc = float(row["roc_m"])

            overlaps = compute_projected_psam_overlap_for_roc(
                current_group_label,
                current_roc,
                roc,
            )

            rows.append(
                {
                    "scan": scan_name,
                    "pzt_voltage_v": pzt_voltage,
                    "roc_m": roc,
                    "overlap_x": overlaps["x_overlap"],
                    "overlap_y": overlaps["y_overlap"],
                    "overlap_mean": overlaps["mean_overlap"],
                    "overlap_x_percent": 100 * overlaps["x_overlap"],
                    "overlap_y_percent": 100 * overlaps["y_overlap"],
                    "overlap_mean_percent": 100 * overlaps["mean_overlap"],
                }
            )

    table = pd.DataFrame(rows)

    if len(table) > 0:
        table = table.sort_values(["scan", "pzt_voltage_v"]).reset_index(drop=True)

    return table


def build_psam_projected_overlap_table_multistart(psam_curves, start_cases):
    """
    Build one combined projection table for multiple starting assumptions.

    Each starting assumption is a combination of:
      - fitted measured group label (for example 3.x V or 1.x V)
      - assumed starting branch on the hysteresis curve (increasing/decreasing)
      - assumed starting PZT voltage
      - assumed starting RoC used to undo the present ZM2 curvature

    For each starting assumption, project the measured fitted beam onto both the
    increasing and decreasing PSAM calibration curves.
    """
    rows = []

    for case in start_cases:
        group_label = case["start_group_label"]
        start_branch = case["start_branch"]
        start_pzt_v = float(case["start_pzt_v"])
        start_roc = float(case["start_roc_m"])
        panel = case["panel"]

        if group_label not in fit_group_labels:
            print(
                f"\nWARNING: start-group {group_label} not available in fitted groups; "
                "skipping this PSAM projection case."
            )
            continue

        for scan_name in ["increasing", "decreasing"]:
            scan_df = filter_psam_scan_to_available_range(psam_curves[scan_name])

            for _, row in scan_df.iterrows():
                pzt_voltage = float(row["pzt_voltage_v"])
                roc = float(row["roc_m"])

                overlaps = compute_projected_psam_overlap_for_roc(
                    group_label,
                    start_roc,
                    roc,
                )

                rows.append(
                    {
                        "panel": panel,
                        "start_group_label": group_label,
                        "start_branch": start_branch,
                        "start_pzt_v": start_pzt_v,
                        "start_roc_m": start_roc,
                        "candidate_scan": scan_name,
                        "pzt_voltage_v": pzt_voltage,
                        "roc_m": roc,
                        "overlap_x": overlaps["x_overlap"],
                        "overlap_y": overlaps["y_overlap"],
                        "overlap_x_percent": 100 * overlaps["x_overlap"],
                        "overlap_y_percent": 100 * overlaps["y_overlap"],
                    }
                )

    table = pd.DataFrame(rows)

    if len(table) > 0:
        table = table.sort_values(
            ["panel", "start_group_label", "start_branch", "candidate_scan", "pzt_voltage_v"]
        ).reset_index(drop=True)

    return table


def print_psam_multistart_projection_table(multistart_table):
    """Print the expanded PSAM multistart projection table."""
    if len(multistart_table) == 0:
        print("\nNo multistart PSAM projection table was built.")
        return

    print("\nExpanded projected overlap table (x/y only, no mean):")
    rows = []
    for _, row in multistart_table.iterrows():
        rows.append([
            row["panel"],
            row["start_group_label"],
            row["start_branch"],
            row["start_pzt_v"],
            row["start_roc_m"],
            row["candidate_scan"],
            row["pzt_voltage_v"],
            row["roc_m"],
            row["overlap_x_percent"],
            row["overlap_y_percent"],
        ])

    print_table(
        rows,
        headers=[
            "panel",
            "start group",
            "start branch",
            "start PZT [V]",
            "start RoC [m]",
            "candidate scan",
            "PZT [V]",
            "RoC [m]",
            "x overlap [%]",
            "y overlap [%]",
        ],
        floatfmt=".4f",
    )


def print_psam_start_cases():
    """Print the explicit starting assumptions being used."""
    print("\nPSAM projection start cases used in version 30:")
    rows = []
    for case in PSAM_START_CASES:
        rows.append([
            case["panel"],
            case["start_group_label"],
            case["start_branch"],
            case["start_pzt_v"],
            case["start_roc_m"],
        ])

    print_table(
        rows,
        headers=["panel", "start group", "start branch", "start PZT [V]", "start RoC [m]"],
        floatfmt=".5f",
    )


def print_psam_tables(available_table, projected_table, current_group_label, current_roc, current_pzt):
    """
    Print PSAM calibration and projected-overlap tables.
    """
    print("\nAvailable PSAM RoC table, using only available Excel points from 0--90 V:")

    available_rows = []

    for _, row in available_table.iterrows():
        available_rows.append([
            row["pzt_voltage_v"],
            row.get("roc_increasing_m", np.nan),
            row.get("roc_decreasing_m", np.nan),
        ])

    print_table(
        available_rows,
        headers=[
            "PZT [V]",
            "RoC increasing [m]",
            "RoC decreasing [m]",
        ],
        floatfmt=".5f",
    )

    print(
        f"\nPSAM projection baseline: fitted {current_group_label} beam, "
        f"treated as PZT={current_pzt:.1f} V, "
        f"current undo RoC={current_roc:.5f} m from the {PSAM_CURRENT_ROC_SCAN} scan."
    )

    if TARGET_BEAM_AVAILABLE:
        projected_rows = []

        for _, row in projected_table.iterrows():
            projected_rows.append([
                row["scan"],
                row["pzt_voltage_v"],
                row["roc_m"],
                row["overlap_x_percent"],
                row["overlap_y_percent"],
                row["overlap_mean_percent"],
            ])

        print("\nProjected overlap with FC eigenmode after changing ZM2 RoC:")
        print_table(
            projected_rows,
            headers=[
                "scan",
                "PZT [V]",
                "RoC [m]",
                "x overlap [%]",
                "y overlap [%]",
                "mean overlap [%]",
            ],
            floatfmt=".4f",
        )

        if len(projected_table) > 0:
            best_idx = projected_table["overlap_mean_percent"].idxmax()
            best = projected_table.loc[best_idx]
            print(
                "\nBest projected mean overlap in available 0--90 V points:"
                f"\n  scan = {best['scan']}"
                f"\n  PZT = {best['pzt_voltage_v']:.1f} V"
                f"\n  RoC = {best['roc_m']:.5f} m"
                f"\n  x/y/mean overlap = "
                f"{best['overlap_x_percent']:.3f}% / "
                f"{best['overlap_y_percent']:.3f}% / "
                f"{best['overlap_mean_percent']:.3f}%"
            )

    else:
        print("\nProjected overlap table unavailable because FC eigenmode was not built.")


def plot_psam_roc_and_projection(available_table, projected_table):
    """
    Version 30 plotting.

    Keep the RoC-vs-PZT plot, and show a single projected-overlap figure.

    Projected-overlap panel:
        use the measured 3.x V beam, assuming the starting branch is decreasing
        with start point ~60 V and RoC = 0.85025 m.

        From that starting beam, project onto both the increasing and
        decreasing PSAM calibration curves.

    Only x and y overlaps are shown; no mean curve is shown.
    Solid = x overlap, dashed = y overlap.
    Upward markers = candidate increasing scan.
    Downward markers = candidate decreasing scan.
    """
    if len(available_table) > 0:
        fig_roc, ax_roc = plt.subplots(figsize=(7.4, 5.0))

        if "roc_increasing_m" in available_table:
            inc = available_table.dropna(subset=["roc_increasing_m"])
            ax_roc.plot(
                inc["pzt_voltage_v"],
                inc["roc_increasing_m"],
                marker="^",
                linestyle="-",
                linewidth=2.0,
                label="increasing scan",
            )

        if "roc_decreasing_m" in available_table:
            dec = available_table.dropna(subset=["roc_decreasing_m"])
            ax_roc.plot(
                dec["pzt_voltage_v"],
                dec["roc_decreasing_m"],
                marker="v",
                linestyle="--",
                linewidth=2.0,
                label="decreasing scan",
            )

        ax_roc.axvline(
            60.0,
            color="0.35",
            linestyle=":",
            linewidth=1.2,
            label="3.x start, dec branch",
        )

        ax_roc.set_xlabel("PZT voltage [V]")
        ax_roc.set_ylabel("ZM2 RoC [m]")
        ax_roc.set_title("Available ZM2 PSAM RoC values, 0--90 V")
        ax_roc.grid(True, alpha=0.3)
        ax_roc.legend(frameon=True)
        fig_roc.tight_layout()

        roc_png = SCRIPT_DIR / "psam_available_roc_vs_pzt_0_to_90V.png"
        fig_roc.savefig(roc_png, dpi=200, bbox_inches="tight")
        print(f"Saved PSAM RoC-vs-PZT plot to: {roc_png}")

    if len(projected_table) > 0 and TARGET_BEAM_AVAILABLE:
        fig_proj, ax = plt.subplots(figsize=(10.2, 5.9))

        panel_name = "3.x V baseline"
        sub_panel = projected_table[projected_table["panel"] == panel_name].copy()

        start_cases_panel = [
            case for case in PSAM_START_CASES
            if case["panel"] == panel_name
        ]

        # Since only the decreasing-start case is retained, color encodes the
        # candidate hysteresis branch; marker direction also encodes it.
        family_styles = {
            "increasing": {"color": "tab:blue", "marker": "^"},
            "decreasing": {"color": "tab:red", "marker": "v"},
        }

        for case in start_cases_panel:
            start_branch = case["start_branch"]
            start_roc = float(case["start_roc_m"])
            start_pzt = float(case["start_pzt_v"])

            for candidate_scan in ["increasing", "decreasing"]:
                sub = sub_panel[
                    (sub_panel["start_branch"] == start_branch)
                    & (sub_panel["candidate_scan"] == candidate_scan)
                ].sort_values("pzt_voltage_v")

                if len(sub) == 0:
                    continue

                style = family_styles[candidate_scan]
                family_label = (
                    f"start dec ({start_pzt:.0f} V, R={start_roc:.5f} m) "
                    f"→ {candidate_scan[:3]}"
                )

                ax.plot(
                    sub["pzt_voltage_v"],
                    sub["overlap_x_percent"],
                    color=style["color"],
                    marker=style["marker"],
                    linestyle="-",
                    linewidth=2.1,
                    markersize=6.8,
                    label=family_label + " x",
                )

                ax.plot(
                    sub["pzt_voltage_v"],
                    sub["overlap_y_percent"],
                    color=style["color"],
                    marker=style["marker"],
                    linestyle="--",
                    linewidth=2.1,
                    markersize=6.8,
                    label=family_label + " y",
                )

        for case in start_cases_panel:
            ax.axvline(
                float(case["start_pzt_v"]),
                color="0.55",
                linestyle=":",
                linewidth=1.1,
                alpha=0.9,
            )

        ax.set_xlabel("PZT voltage [V]")
        ax.set_ylabel("Projected overlap with FC eigenmode [%]")
        ax.set_title(
            "Projected FC overlap using 3.1x measured beam\n"
            "starting branch assumed decreasing"
        )
        ax.grid(True, alpha=0.3)
        ax.legend(frameon=True, fontsize=8.2, ncols=2, loc="best")

        fig_proj.suptitle(
            "Projected FC mode overlap after changing ZM2 RoC\n"
            "Solid = x overlap, dashed = y overlap; upward markers = increasing scan, downward markers = decreasing scan",
            fontsize=13,
            y=1.02,
        )
        fig_proj.tight_layout()

        proj_png = SCRIPT_DIR / "psam_projected_fc_overlap_vs_pzt_version30.png"
        fig_proj.savefig(proj_png, dpi=200, bbox_inches="tight")
        print(f"Saved projected FC-overlap plot to: {proj_png}")


PSAM_CALIBRATION_AVAILABLE = False
PSAM_AVAILABLE_ROC_TABLE = pd.DataFrame()
PSAM_PROJECTED_OVERLAP_TABLE = pd.DataFrame()
PSAM_MULTISTART_PROJECTED_OVERLAP_TABLE = pd.DataFrame()
PSAM_PRE_ZM2_RECONSTRUCTION_TABLE = pd.DataFrame()
PSAM_PRE_ZM2_RECONSTRUCTION_PAIRWISE_TABLE = pd.DataFrame()

if BUILD_PSAM_PROJECTION:
    try:
        psam_curves = read_psam_roc_calibration(PSAM_CALIBRATION_PATH)
        PSAM_AVAILABLE_ROC_TABLE = make_available_psam_roc_table(psam_curves)

        current_group_label = current_fit_group_for_psam_projection()
        current_roc, current_pzt = current_roc_for_psam_projection(psam_curves)

        # This should normally match ZM2_ROC_3p15. Print a warning if it does not.
        if abs(current_roc - ZM2_ROC_3p15) > 1e-6:
            print(
                f"\nWARNING: ZM2_ROC_3p15 = {ZM2_ROC_3p15:.5f} m, "
                f"but calibration {PSAM_CURRENT_ROC_SCAN} scan at "
                f"PZT={current_pzt:.1f} V gives {current_roc:.5f} m. "
                "PSAM projection will use the calibration-table value."
            )

        PSAM_PROJECTED_OVERLAP_TABLE = build_psam_projected_overlap_table(
            psam_curves,
            current_group_label,
            current_roc,
        )

        PSAM_CALIBRATION_AVAILABLE = True

        print("\nNOTE: The following legacy single-baseline table is kept only for comparison.")
        print("      The plotted v29 projection uses PSAM_START_CASES below, not this legacy 90 V increasing baseline.")

        print_psam_tables(
            PSAM_AVAILABLE_ROC_TABLE,
            PSAM_PROJECTED_OVERLAP_TABLE,
            current_group_label,
            current_roc,
            current_pzt,
        )

        print_psam_start_cases()
        PSAM_MULTISTART_PROJECTED_OVERLAP_TABLE = build_psam_projected_overlap_table_multistart(
            psam_curves,
            PSAM_START_CASES,
        )
        print_psam_multistart_projection_table(PSAM_MULTISTART_PROJECTED_OVERLAP_TABLE)

        # Save the tables next to the script for easy copying into alog/slides.
        available_csv = SCRIPT_DIR / "psam_available_roc_table_0_to_90V.csv"
        projected_csv = SCRIPT_DIR / "psam_projected_fc_overlap_0_to_90V.csv"
        projected_multistart_csv = SCRIPT_DIR / "psam_projected_fc_overlap_multistart_v30.csv"
        PSAM_AVAILABLE_ROC_TABLE.to_csv(available_csv, index=False)
        PSAM_PROJECTED_OVERLAP_TABLE.to_csv(projected_csv, index=False)
        PSAM_MULTISTART_PROJECTED_OVERLAP_TABLE.to_csv(projected_multistart_csv, index=False)
        print(f"\nSaved PSAM RoC table to: {available_csv}")
        print(f"Saved projected overlap table to: {projected_csv}")
        print(f"Saved multistart projected overlap table to: {projected_multistart_csv}")

        # PSAM plots are made near the end of the script, right before the final
        # plt.show().  This avoids backend issues where figures made before the
        # interactive GUI can be hidden behind/blocking windows.

    except Exception as exc:
        PSAM_CALIBRATION_AVAILABLE = False
        print("\nWARNING: Could not build PSAM projection table/plots.")
        print(f"Reason: {exc}")

# -------------------------
# Single-plot interactive propagation setup
# -------------------------

plot_colors = {
    "x": "tab:blue",
    "y": "tab:orange",
}

plot_linestyles = {
    "x": "-",
    "y": "--",
}

target_plot_styles = {
    "x": {
        # Grey dash-dot curve: this is the dashed FC eigenmode line.
        "color": "0.45",
        "linestyle": "-.",
        "linewidth": 2.8,
        "label": "FC eigenmode x",
    },
    "y": {
        "color": "black",
        "linestyle": "-",
        "linewidth": 2.8,
        "label": "FC eigenmode y",
    },
}

vopo_plot_styles = {
    "target": {
        "color": "tab:green",
        "linestyle": "-",
        "linewidth": 2.4,
        "label": "VOPO target",
    },
    "o4_x": {
        "color": "tab:red",
        "linestyle": "-.",
        "linewidth": 2.1,
        "label": "VOPO O4 x",
    },
    "o4_y": {
        "color": "magenta",
        "linestyle": "--",
        "linewidth": 2.1,
        "label": "VOPO O4 y",
    },
}

# Common plotting coordinate:
# extrapolate all the way from FC1 to ZM2.
s_min_global = X_FC1
s_max_global = X_ZM2
s_plot_global = np.linspace(s_min_global, s_max_global, 1800)

# Convert plotting coordinate back to the fit coordinate.
x_up_plot_global = s_plot_global - X_ZM3

# Zoom-panel range in the FC1-referenced coordinate.
# This is the right-side inset/zoom plot.
ZOOM_X_MIN = 2.0
ZOOM_X_MAX = 2.4

# Curve visibility controls.
# The labels here are used by the CheckButtons widget.
CURVE_OPTIONS = [
    ("FC x", "fc_x", True),
    ("FC y", "fc_y", True),
    ("VOPO target", "vopo_target", True),
    ("VOPO O4 x", "vopo_o4_x", True),
    ("VOPO O4 y", "vopo_o4_y", True),
    ("SG x fit/data", "sg_x", True),
    ("SG y fit/data", "sg_y", True),
]

CURVE_VISIBLE = {
    key: default
    for _, key, default in CURVE_OPTIONS
}

CURVE_LABEL_TO_KEY = {
    label: key
    for label, key, _ in CURVE_OPTIONS
}


def curve_is_visible(key):
    return CURVE_VISIBLE.get(key, True)


def mark_return_beam_waist(ax, q_after_ZM2, label, color, y_offset_mm):
    """
    Mark the waist of a return-propagating beam.

    q_after_ZM2 is oriented in the physical direction:
        ZM2 -> ZM3 -> FC1
    """
    waist_s = waist_s_from_return_beam_after_ZM2(q_after_ZM2)

    if s_min_global - 0.20 <= waist_s <= s_max_global + 0.20:
        ax.plot(
            waist_s,
            0,
            marker="D",
            markersize=5,
            color=color,
            zorder=10,
        )

        ax.text(
            waist_s,
            y_offset_mm,
            f" {label} waist",
            color=color,
            ha="left",
            va="bottom" if y_offset_mm > 0 else "top",
            fontsize=8,
            fontweight="bold",
            clip_on=True,
        )


def compute_ellipticity_summary_for_group(group_label):
    """
    Compute fitted beam ellipticity at FC1, ZM3, ZM2,
    and averaged across the measured points.
    """
    fx = fits[(group_label, "x")]
    fy = fits[(group_label, "y")]

    positions_s = {
        "FC1": X_FC1,
        "ZM3": X_ZM3,
        "ZM2": X_ZM2,
    }

    out = {}

    for name, spos in positions_s.items():
        x_up_pos = spos - X_ZM3

        wx = gaussian_beam_w(x_up_pos, fx["w0"], fx["x0_up"])
        wy = gaussian_beam_w(x_up_pos, fy["w0"], fy["x0_up"])

        out[f"wx_{name}"] = wx
        out[f"wy_{name}"] = wy
        out[f"aspect_{name}"] = aspect_ratio_from_wx_wy(wx, wy)
        out[f"ellipticity_percent_{name}"] = ellipticity_percent_from_wx_wy(wx, wy)

    group_data = grouped_data[group_label]
    x_meas = group_data["x_up_m"].values

    wx_meas_fit = gaussian_beam_w(x_meas, fx["w0"], fx["x0_up"])
    wy_meas_fit = gaussian_beam_w(x_meas, fy["w0"], fy["x0_up"])

    aspect_meas = np.array([
        aspect_ratio_from_wx_wy(wx, wy)
        for wx, wy in zip(wx_meas_fit, wy_meas_fit)
    ])

    out["aspect_meas_mean"] = np.mean(aspect_meas)
    out["ellipticity_percent_meas_mean"] = 100 * (out["aspect_meas_mean"] - 1)

    return out


def compute_global_ylim_mm():
    """
    Compute a fixed y-limit for all SG settings, including the target beam.
    """
    ymax_mm = 0.0

    for group_label in fit_group_labels:
        group_data = grouped_data[group_label]

        for axis_label, diam_col in [
            ("x", "diam_x_um"),
            ("y", "diam_y_um"),
        ]:
            fit = fits[(group_label, axis_label)]

            w_fit = gaussian_beam_w(
                x_up_plot_global,
                fit["w0"],
                fit["x0_up"],
            )

            w_data = convert_size_um_to_radius_m(group_data[diam_col].values)

            ymax_mm = max(
                ymax_mm,
                np.nanmax(1e3 * w_fit),
                np.nanmax(1e3 * w_data),
            )

    if TARGET_BEAM_AVAILABLE:
        w_target_x = w_from_q_at_ZM3_along_FC1_to_ZM2(
            qx_target_at_ZM3,
            s_plot_global,
        )

        w_target_y = w_from_q_at_ZM3_along_FC1_to_ZM2(
            qy_target_at_ZM3,
            s_plot_global,
        )

        ymax_mm = max(
            ymax_mm,
            np.nanmax(1e3 * w_target_x),
            np.nanmax(1e3 * w_target_y),
        )

    if VOPO_REFERENCES_AVAILABLE:
        w_vopo_target = w_from_return_beam_after_ZM2(
            q_vopo_target_after_ZM2_return,
            s_plot_global,
        )

        w_vopo_o4_x = w_from_return_beam_after_ZM2(
            q_vopo_o4_x_after_ZM2_return,
            s_plot_global,
        )

        w_vopo_o4_y = w_from_return_beam_after_ZM2(
            q_vopo_o4_y_after_ZM2_return,
            s_plot_global,
        )

        ymax_mm = max(
            ymax_mm,
            np.nanmax(1e3 * w_vopo_target),
            np.nanmax(1e3 * w_vopo_o4_x),
            np.nanmax(1e3 * w_vopo_o4_y),
        )

    # Add headroom for labels and text boxes.
    return 1.18 * ymax_mm


GLOBAL_YLIM_MM = compute_global_ylim_mm()


# -------------------------
# Print ellipticity summary
# -------------------------

ellipticity_rows = []

for group_label in fit_group_labels:
    group_data = grouped_data[group_label]
    voltage_mean = np.mean(group_data["voltage_mean"].values)

    ell = compute_ellipticity_summary_for_group(group_label)

    ellipticity_rows.append([
        group_label,
        voltage_mean,
        ell["aspect_FC1"],
        ell["ellipticity_percent_FC1"],
        ell["aspect_ZM3"],
        ell["ellipticity_percent_ZM3"],
        ell["aspect_ZM2"],
        ell["ellipticity_percent_ZM2"],
        ell["aspect_meas_mean"],
        ell["ellipticity_percent_meas_mean"],
    ])

print("\nBeam ellipticity summary from fitted x/y propagation:")
print_table(
    ellipticity_rows,
    headers=[
        "ZM2 SG group",
        "mean V",
        "aspect FC1",
        "ellip FC1 [%]",
        "aspect ZM3",
        "ellip ZM3 [%]",
        "aspect ZM2",
        "ellip ZM2 [%]",
        "mean aspect at meas.",
        "mean ellip at meas. [%]",
    ],
    floatfmt=".3f",
)


# -------------------------
# Single-plot interactive propagation GUI
# -------------------------


def plot_plus_minus_curve(ax, x, w_mm, color, linestyle, linewidth, label=None, alpha=None):
    """
    Plot +w and -w envelopes.
    """
    ax.plot(
        x,
        w_mm,
        color=color,
        linestyle=linestyle,
        linewidth=linewidth,
        label=label,
        alpha=alpha,
    )

    ax.plot(
        x,
        -w_mm,
        color=color,
        linestyle=linestyle,
        linewidth=linewidth,
        alpha=alpha,
    )


def mark_waist_on_zoom(ax, waist_s, label, color, y_offset_mm=None, marker="D"):
    """
    Mark a waist location on the zoom panel only.

    The waist designations are handled by the zoom-panel legend, rather than
    writing text next to every marker. The y_offset_mm argument is kept for
    backwards compatibility with the existing calls, but is not used.
    """
    if ZOOM_X_MIN - 0.05 <= waist_s <= ZOOM_X_MAX + 0.05:
        ax.plot(
            waist_s,
            0,
            marker=marker,
            markersize=5,
            color=color,
            linestyle="None",
            label=label,
            zorder=20,
        )


def mark_forward_beam_waist_from_ZM3_on_zoom(ax, q_at_ZM3, label, color, y_offset_mm):
    """
    Mark waist for a beam whose q at ZM3 is oriented FC1 -> ZM2.
    """
    waist_s = X_ZM3 - q_at_ZM3.z
    mark_waist_on_zoom(ax, waist_s, label, color, y_offset_mm, marker="s")


def mark_return_beam_waist_after_ZM2_on_zoom(ax, q_after_ZM2, label, color, y_offset_mm):
    """
    Mark waist for a beam physically propagating ZM2 -> ZM3 -> FC1.
    """
    waist_s = waist_s_from_return_beam_after_ZM2(q_after_ZM2)
    mark_waist_on_zoom(ax, waist_s, label, color, y_offset_mm, marker="D")


def mark_sg_fit_waist_on_zoom(ax, group_label, axis_label, label, color, y_offset_mm):
    """
    Mark fitted SG waist on the zoom panel only.
    """
    fit = fits[(group_label, axis_label)]
    mark_waist_on_zoom(
        ax,
        fit["waist_from_FC1"],
        label,
        color,
        y_offset_mm,
        marker="x",
    )


def draw_zoom_from_main_axis(ax_main, ax_zoom, group_label):
    """
    Copy the visible main propagation plot into a right-side zoom panel.

    Waist locations are deliberately shown only here, not on the big plot.
    """
    ax_zoom.clear()

    y_abs_max = 0.0

    for line in ax_main.get_lines():
        xdata = np.asarray(line.get_xdata(), dtype=float)
        ydata = np.asarray(line.get_ydata(), dtype=float)

        if len(xdata) == 0 or len(ydata) == 0:
            continue

        ax_zoom.plot(
            xdata,
            ydata,
            color=line.get_color(),
            linestyle=line.get_linestyle(),
            linewidth=line.get_linewidth(),
            marker=line.get_marker(),
            markersize=line.get_markersize(),
            markeredgewidth=line.get_markeredgewidth(),
            alpha=line.get_alpha(),
        )

        # Estimate the zoom y-limits from non-vertical, non-horizontal beam/data lines.
        finite = np.isfinite(xdata) & np.isfinite(ydata)
        if np.count_nonzero(finite) < 2:
            continue

        xfinite = xdata[finite]
        yfinite = ydata[finite]

        # Skip mostly vertical/horizontal guide lines for y-limit estimation.
        if np.nanmax(xfinite) - np.nanmin(xfinite) < 1e-9:
            continue

        if np.nanmax(yfinite) - np.nanmin(yfinite) < 1e-9:
            continue

        in_zoom = (xfinite >= ZOOM_X_MIN) & (xfinite <= ZOOM_X_MAX)

        if np.any(in_zoom):
            y_abs_max = max(y_abs_max, np.nanmax(np.abs(yfinite[in_zoom])))

    ax_zoom.set_xlim(ZOOM_X_MIN, ZOOM_X_MAX)

    if y_abs_max > 0:
        ax_zoom.set_ylim(-1.12 * y_abs_max, 1.12 * y_abs_max)
    else:
        ax_zoom.set_ylim(ax_main.get_ylim())

    # Waist markers: shown on zoom only.
    if TARGET_BEAM_AVAILABLE:
        if curve_is_visible("fc_x"):
            mark_forward_beam_waist_from_ZM3_on_zoom(
                ax_zoom,
                qx_target_at_ZM3,
                "FC x waist",
                target_plot_styles["x"]["color"],
                -0.55,
            )

        if curve_is_visible("fc_y"):
            mark_forward_beam_waist_from_ZM3_on_zoom(
                ax_zoom,
                qy_target_at_ZM3,
                "FC y waist",
                target_plot_styles["y"]["color"],
                -0.75,
            )

    if VOPO_REFERENCES_AVAILABLE:
        if curve_is_visible("vopo_target"):
            mark_return_beam_waist_after_ZM2_on_zoom(
                ax_zoom,
                q_vopo_target_after_ZM2_return,
                "VOPO target waist",
                vopo_plot_styles["target"]["color"],
                0.55,
            )

        if curve_is_visible("vopo_o4_x"):
            mark_return_beam_waist_after_ZM2_on_zoom(
                ax_zoom,
                q_vopo_o4_x_after_ZM2_return,
                "O4 x waist",
                vopo_plot_styles["o4_x"]["color"],
                0.75,
            )

        if curve_is_visible("vopo_o4_y"):
            mark_return_beam_waist_after_ZM2_on_zoom(
                ax_zoom,
                q_vopo_o4_y_after_ZM2_return,
                "O4 y waist",
                vopo_plot_styles["o4_y"]["color"],
                0.95,
            )

    if curve_is_visible("sg_x"):
        mark_sg_fit_waist_on_zoom(
            ax_zoom,
            group_label,
            "x",
            "SG x waist",
            plot_colors["x"],
            -0.95,
        )

    if curve_is_visible("sg_y"):
        mark_sg_fit_waist_on_zoom(
            ax_zoom,
            group_label,
            "y",
            "SG y waist",
            plot_colors["y"],
            1.15,
        )

    ax_zoom.axhline(0, color="0.75", linewidth=0.8)
    ax_zoom.grid(True, alpha=0.3)

    # Legend for waist markers only. The copied propagation curves do not carry
    # labels in the zoom panel, so this legend stays focused on waist locations.
    waist_handles, waist_labels = ax_zoom.get_legend_handles_labels()
    if len(waist_handles) > 0:
        ax_zoom.legend(
            waist_handles,
            waist_labels,
            title="Waists",
            frameon=True,
            fontsize=6.2,
            title_fontsize=6.8,
            loc="upper right",
            borderpad=0.35,
            labelspacing=0.25,
            handlelength=1.0,
            handletextpad=0.4,
        )

    ax_zoom.set_title("Zoom: 2.0–2.4 m", fontsize=10, loc="left")
    ax_zoom.set_xlabel("Distance from FC1 [m]", fontsize=9)
    ax_zoom.set_ylabel("±w [mm]", fontsize=9)
    ax_zoom.tick_params(axis="both", labelsize=8)


def draw_selected_group(ax, group_label, ax_zoom=None):
    """
    Draw selected ZM2 setting on a single axis.
    """
    ax.clear()

    group_data = grouped_data[group_label]

    mean_voltage = np.mean(group_data["voltage_mean"].values)
    min_voltage = np.min(group_data["voltage_min"].values)
    max_voltage = np.max(group_data["voltage_max"].values)

    s_data = group_data["s_from_FC1_m"].values

    # Target / FC eigenmode beam
    if TARGET_BEAM_AVAILABLE:
        w_target_x = w_from_q_at_ZM3_along_FC1_to_ZM2(
            qx_target_at_ZM3,
            s_plot_global,
        )

        w_target_y = w_from_q_at_ZM3_along_FC1_to_ZM2(
            qy_target_at_ZM3,
            s_plot_global,
        )

        if curve_is_visible("fc_x"):
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_target_x,
                color=target_plot_styles["x"]["color"],
                linestyle=target_plot_styles["x"]["linestyle"],
                linewidth=target_plot_styles["x"]["linewidth"],
                label=target_plot_styles["x"]["label"],
            )

        if curve_is_visible("fc_y"):
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_target_y,
                color=target_plot_styles["y"]["color"],
                linestyle=target_plot_styles["y"]["linestyle"],
                linewidth=target_plot_styles["y"]["linewidth"],
                label=target_plot_styles["y"]["label"],
            )

    # VOPO target/O4 propagated from ZM1, through ZM2 RoC, to FC1
    if VOPO_REFERENCES_AVAILABLE:
        w_vopo_target = w_from_return_beam_after_ZM2(
            q_vopo_target_after_ZM2_return,
            s_plot_global,
        )

        w_vopo_o4_x = w_from_return_beam_after_ZM2(
            q_vopo_o4_x_after_ZM2_return,
            s_plot_global,
        )

        w_vopo_o4_y = w_from_return_beam_after_ZM2(
            q_vopo_o4_y_after_ZM2_return,
            s_plot_global,
        )

        if curve_is_visible("vopo_target"):
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_vopo_target,
                color=vopo_plot_styles["target"]["color"],
                linestyle=vopo_plot_styles["target"]["linestyle"],
                linewidth=vopo_plot_styles["target"]["linewidth"],
                label=vopo_plot_styles["target"]["label"],
            )

        if curve_is_visible("vopo_o4_x"):
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_vopo_o4_x,
                color=vopo_plot_styles["o4_x"]["color"],
                linestyle=vopo_plot_styles["o4_x"]["linestyle"],
                linewidth=vopo_plot_styles["o4_x"]["linewidth"],
                label=vopo_plot_styles["o4_x"]["label"],
            )

        if curve_is_visible("vopo_o4_y"):
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_vopo_o4_y,
                color=vopo_plot_styles["o4_y"]["color"],
                linestyle=vopo_plot_styles["o4_y"]["linestyle"],
                linewidth=vopo_plot_styles["o4_y"]["linewidth"],
                label=vopo_plot_styles["o4_y"]["label"],
            )

    # Fitted beams and measured points
    for axis_label, diam_col in [
        ("x", "diam_x_um"),
        ("y", "diam_y_um"),
    ]:
        fit = fits[(group_label, axis_label)]

        w_data = convert_size_um_to_radius_m(group_data[diam_col].values)

        w_fit = gaussian_beam_w(
            x_up_plot_global,
            fit["w0"],
            fit["x0_up"],
        )

        curve_key = f"sg_{axis_label}"

        if curve_is_visible(curve_key):
            # Fit propagation: +w and -w
            plot_plus_minus_curve(
                ax,
                s_plot_global,
                1e3 * w_fit,
                color=plot_colors[axis_label],
                linestyle=plot_linestyles[axis_label],
                linewidth=2.2,
                label=f"{axis_label} fit",
            )

            # Measured data: +w and -w
            ax.plot(
                s_data,
                1e3 * w_data,
                "o",
                color=plot_colors[axis_label],
                alpha=0.9,
                label=f"{axis_label} data",
            )

            ax.plot(
                s_data,
                -1e3 * w_data,
                "o",
                color=plot_colors[axis_label],
                alpha=0.9,
            )

    # Mark main optics
    optic_markers = [
        ("FC1", X_FC1, "0.25", "--"),
        ("ZM3", X_ZM3, "0.20", "-"),
        ("ZM2", X_ZM2, "0.25", "--"),
    ]

    for optic_label, optic_x, optic_color, optic_ls in optic_markers:
        ax.axvline(
            optic_x,
            color=optic_color,
            linestyle=optic_ls,
            linewidth=1.4,
            alpha=0.95,
        )

    # Mark measurement planes
    for _, row in group_data.iterrows():
        xpos = row["s_from_FC1_m"]

        ax.axvline(
            xpos,
            color="0.75",
            linestyle=":",
            linewidth=1.0,
            alpha=0.8,
        )

    ax.axhline(0, color="0.75", linewidth=1.0)

    ax.set_xlim(s_min_global, s_max_global)

    # Fixed y-limits for all SG settings
    ax.set_ylim(-GLOBAL_YLIM_MM, GLOBAL_YLIM_MM)

    ax.set_xlabel("Distance from FC1 along FC1 → ZM3 → ZM2 path [m]")
    ax.set_ylabel("Beam radius ±w [mm]")

    # No ordinary axis title; custom title is placed above the plot.
    ax.set_title("")

    ax.grid(True, alpha=0.3)

    # Main optics labels above the plotting area
    optics_label_transform = blended_transform_factory(ax.transData, ax.transAxes)

    for optic_label, optic_x, optic_color, optic_ls in optic_markers:
        ax.text(
            optic_x,
            1.025,
            optic_label,
            transform=optics_label_transform,
            ha="center",
            va="bottom",
            fontsize=15,
            color=optic_color,
            fontweight="bold",
            clip_on=False,
        )

    # SG setting title above the plot, below the main figure title
    sg_title = (
        f"ZM2 strain gauge setting: {group_label} "
        f"(mean {mean_voltage:.2f} V, range {min_voltage:.2f}–{max_voltage:.2f} V)"
    )

    ax.text(
        0.0,
        1.105,
        sg_title,
        transform=ax.transAxes,
        ha="left",
        va="bottom",
        fontsize=12,
        fontweight="bold",
        color="black",
        clip_on=False,
    )

    ymin, ymax = ax.get_ylim()

    # Label measurement planes inside the top of the plot
    for _, row in group_data.iterrows():
        xpos = row["s_from_FC1_m"]
        label = row["point"]

        ax.text(
            xpos,
            ymin + 0.96 * (ymax - ymin),
            f" {label}",
            ha="left",
            va="top",
            fontsize=8,
            rotation=90,
            color="0.35",
        )

    # Fit quality
    fx = fits[(group_label, "x")]
    fy = fits[(group_label, "y")]

    # Ellipticity
    ell = compute_ellipticity_summary_for_group(group_label)

    # Target overlap
    target_overlap = compute_target_overlap_for_group(group_label)

    if TARGET_BEAM_AVAILABLE:
        target_overlap_text = (
            f"FC overlap x = {100 * target_overlap['x_overlap']:.2f}%\n"
            f"FC overlap y = {100 * target_overlap['y_overlap']:.2f}%\n"
            f"FC overlap mean = {100 * target_overlap['mean_overlap']:.2f}%\n"
        )
    else:
        target_overlap_text = (
            "FC overlap unavailable\n"
        )

    vopo_overlap = compute_vopo_overlap_for_group(group_label)

    if VOPO_REFERENCES_AVAILABLE:
        vopo_overlap_text = (
            f"VOPO target x/y/mean = "
            f"{100 * vopo_overlap['target_x_overlap']:.1f}/"
            f"{100 * vopo_overlap['target_y_overlap']:.1f}/"
            f"{100 * vopo_overlap['target_mean_overlap']:.1f}%\n"
            f"VOPO O4 x/y/mean = "
            f"{100 * vopo_overlap['o4_x_overlap']:.1f}/"
            f"{100 * vopo_overlap['o4_y_overlap']:.1f}/"
            f"{100 * vopo_overlap['o4_mean_overlap']:.1f}%\n"
        )
    else:
        vopo_overlap_text = (
            "VOPO overlap unavailable\n"
        )

    info_text = (
        f"ZM2 SG: {group_label}\n"
        f"mean {mean_voltage:.2f} V, range {min_voltage:.2f}–{max_voltage:.2f} V\n"
        f"x w0 = {1e3 * fx['w0']:.3f} mm\n"
        f"x waist = {1e3 * fx['waist_from_FC1']:.1f} mm from FC1\n"
        f"y w0 = {1e3 * fy['w0']:.3f} mm\n"
        f"y waist = {1e3 * fy['waist_from_FC1']:.1f} mm from FC1\n"
        f"x RMS residual = {1e6 * fx['rms_residual']:.1f} µm\n"
        f"y RMS residual = {1e6 * fy['rms_residual']:.1f} µm\n"
        f"{target_overlap_text}"
        f"{vopo_overlap_text}"
        f"ellip. at FC1 = {ell['ellipticity_percent_FC1']:.1f}%\n"
        f"ellip. at ZM3 = {ell['ellipticity_percent_ZM3']:.1f}%\n"
        f"ellip. at ZM2 = {ell['ellipticity_percent_ZM2']:.1f}%"
    )

    # Move the waist / residual / overlap / ellipticity text box to the middle-right figure column.
    old_info_text = getattr(ax.figure, "_beam_prop_info_text", None)

    if old_info_text is not None:
        old_info_text.remove()

    ax.figure._beam_prop_info_text = ax.figure.text(
        0.82,
        0.98,
        info_text,
        ha="left",
        va="top",
        fontsize=7.6,
        bbox=dict(
            facecolor="white",
            edgecolor="0.6",
            alpha=0.88,
            pad=5,
        ),
    )

    # Move legend underneath the curve checkboxes.
    old_legend = getattr(ax.figure, "_beam_prop_legend", None)

    if old_legend is not None:
        old_legend.remove()

    handles, labels = ax.get_legend_handles_labels()

    ax.figure._beam_prop_legend = ax.figure.legend(
        handles,
        labels,
        title="Visible curves",
        frameon=True,
        fontsize=6.8,
        title_fontsize=7.5,
        loc="upper left",
        bbox_to_anchor=(0.64, 0.3),
    )

    if ax_zoom is not None:
        draw_zoom_from_main_axis(ax, ax_zoom, group_label)

# Initial selection
initial_group_label = fit_group_labels[0]

fig_gui = plt.figure(figsize=(17.2, 7.4))

# Main plot axis
# Leave room on the right for:
#   - smaller SG selector
#   - curve checkboxes
#   - legend / text
#   - zoom panel
ax_main = fig_gui.add_axes([0.055, 0.10, 0.56, 0.66])

# Right-side zoom panel, placed to the right of the text/legend column.
ax_zoom = fig_gui.add_axes([0.785, 0.18, 0.20, 0.48])

# Radio button axis: smaller SG selector.
rax_height = min(0.050 * len(fit_group_labels), 0.22)
rax = fig_gui.add_axes([0.640, 0.720, 0.090, rax_height])

radio = RadioButtons(
    rax,
    tuple(fit_group_labels),
    active=0,
)

rax.set_title("ZM2 SG", fontsize=8)
for label in radio.labels:
    label.set_fontsize(7.5)

# Curve visibility checkboxes.
check_labels = [label for label, _, _ in CURVE_OPTIONS]
check_states = [CURVE_VISIBLE[key] for _, key, _ in CURVE_OPTIONS]

check_ax_height = min(0.030 * len(check_labels) + 0.035, 0.26)
check_ax = fig_gui.add_axes([0.640, 0.35, 0.08, check_ax_height])
curve_checks = CheckButtons(
    check_ax,
    check_labels,
    check_states,
)

check_ax.set_title("Show curves", fontsize=8)
for label in curve_checks.labels:
    label.set_fontsize(7.0)

# Initial draw
draw_selected_group(ax_main, initial_group_label, ax_zoom=ax_zoom)

fig_gui.suptitle(
    "Interactive Gaussian beam propagation fits\n"
    "FC1-referenced propagation from FC1 to ZM2",
    x=0.055,
    y=0.965,
    ha="left",
)


def current_group_label():
    return radio.value_selected


def redraw_gui():
    draw_selected_group(ax_main, current_group_label(), ax_zoom=ax_zoom)
    fig_gui.canvas.draw_idle()


def update_group(label):
    redraw_gui()


def update_curve_visibility(label):
    key = CURVE_LABEL_TO_KEY[label]
    CURVE_VISIBLE[key] = not CURVE_VISIBLE[key]
    redraw_gui()


radio.on_clicked(update_group)
curve_checks.on_clicked(update_curve_visibility)

# -------------------------
# Save animated GIF of the propagation plot
# -------------------------
#
# This saves the complete GUI figure while cycling through the ZM2 strain-gauge
# settings. The current checkbox visibility state is respected.
#
# The GIF is saved next to this script as:
#   HAM7_ZM2_strain_gauge_scan.gif

SAVE_PROPAGATION_GIF = True
GIF_FILENAME = SCRIPT_DIR / "HAM7_ZM2_strain_gauge_scan.gif"
GIF_FRAME_DURATION_MS = 2000  # 2 seconds per strain-gauge setting
GIF_LOOP_FOREVER = 0          # 0 means loop forever


def figure_to_pil_image(fig):
    """
    Convert a matplotlib figure canvas to a PIL RGB image.
    """
    fig.canvas.draw()

    # buffer_rgba is more robust than tostring_rgb across matplotlib versions.
    rgba = np.asarray(fig.canvas.buffer_rgba())

    return Image.fromarray(rgba).convert("RGB")


def save_strain_gauge_gif(fig, radio_widget, group_labels, gif_path,
                          frame_duration_ms=2000, loop=0):
    """
    Cycle through SG settings using the existing radio buttons,
    capture each updated figure state, and save as a GIF.
    """
    frames = []

    print("\nCreating propagation GIF...")
    print(f"  Output: {gif_path}")
    print(f"  Frames: {len(group_labels)}")
    print(f"  Duration per frame: {frame_duration_ms} ms")

    # Remember the current GUI selection so we can restore it after saving.
    original_label = radio_widget.value_selected

    for i, label in enumerate(group_labels):
        print(f"  Rendering frame {i + 1}/{len(group_labels)} : {label}")

        # Trigger the same callback path used by clicking the GUI radio buttons.
        radio_widget.set_active(i)

        # Let matplotlib process the redraw before capture.
        plt.pause(0.10)

        frame = figure_to_pil_image(fig)
        frames.append(frame.copy())

    if len(frames) == 0:
        raise RuntimeError("No GIF frames were generated.")

    frames[0].save(
        gif_path,
        save_all=True,
        append_images=frames[1:],
        duration=frame_duration_ms,
        loop=loop,
    )

    print(f"Saved GIF to: {gif_path}")

    # Restore the original selected strain-gauge setting in the GUI.
    if original_label in group_labels:
        original_index = list(group_labels).index(original_label)
        radio_widget.set_active(original_index)
        plt.pause(0.05)


if SAVE_PROPAGATION_GIF:
    save_strain_gauge_gif(
        fig=fig_gui,
        radio_widget=radio,
        group_labels=fit_group_labels,
        gif_path=GIF_FILENAME,
        frame_duration_ms=GIF_FRAME_DURATION_MS,
        loop=GIF_LOOP_FOREVER,
    )

# Do not block here; all figures, including PSAM plots, are shown together at the end.
# plt.show()

# %%
# Additional summary plots: overlaps with VOPO target and O4 references

if VOPO_REFERENCES_AVAILABLE:
    sg_labels = []
    vopo_target_x = []
    vopo_target_y = []
    vopo_target_mean = []
    vopo_o4_x = []
    vopo_o4_y = []
    vopo_o4_mean = []

    for group_label in fit_group_labels:
        overlaps = compute_vopo_overlap_for_group(group_label)

        sg_labels.append(group_label)
        vopo_target_x.append(100 * overlaps["target_x_overlap"])
        vopo_target_y.append(100 * overlaps["target_y_overlap"])
        vopo_target_mean.append(100 * overlaps["target_mean_overlap"])
        vopo_o4_x.append(100 * overlaps["o4_x_overlap"])
        vopo_o4_y.append(100 * overlaps["o4_y_overlap"])
        vopo_o4_mean.append(100 * overlaps["o4_mean_overlap"])

    sg_labels = np.array(sg_labels)
    x = np.arange(len(sg_labels))
    bar_width = 0.24

    # Plot 1: VOPO target overlaps
    fig_target, ax_target = plt.subplots(figsize=(7.5, 5.5))

    bars_x = ax_target.bar(
        x - bar_width,
        vopo_target_x,
        width=bar_width,
        label="x",
    )

    bars_y = ax_target.bar(
        x,
        vopo_target_y,
        width=bar_width,
        label="y",
    )

    bars_mean = ax_target.bar(
        x + bar_width,
        vopo_target_mean,
        width=bar_width,
        label="mean",
    )

    for bars in [bars_x, bars_y, bars_mean]:
        for bar in bars:
            val = bar.get_height()
            ax_target.text(
                bar.get_x() + bar.get_width() / 2,
                val - 0.35,
                f"{val:.2f}%",
                ha="center",
                va="top",
                fontsize=9,
                rotation=90,
                color="white",
                fontweight="bold",
            )

    ax_target.axhline(99.5, color="black", linestyle="--", linewidth=1.2)
    ax_target.axhline(95.0, color="red", linestyle="--", linewidth=1.2)

    ax_target.set_xticks(x)
    ax_target.set_xticklabels(sg_labels)
    ax_target.set_xlabel("ZM2 strain gauge group")
    ax_target.set_ylabel("Mode overlap with VOPO target [%]")
    ax_target.set_title(
        "Overlap of fitted SG beams with VOPO target\n"
        f"propagated through ZM2 RoC = {ZM2_ROC_SIGN * ZM2_ROC_3p15:.5f} m"
    )
    ax_target.grid(True, axis="y", alpha=0.3)
    ax_target.legend(frameon=True)

    fig_target.tight_layout()

    # Plot 2: VOPO O4 overlaps
    fig_o4, ax_o4 = plt.subplots(figsize=(7.5, 5.5))

    bars_x = ax_o4.bar(
        x - bar_width,
        vopo_o4_x,
        width=bar_width,
        label="x",
    )

    bars_y = ax_o4.bar(
        x,
        vopo_o4_y,
        width=bar_width,
        label="y",
    )

    bars_mean = ax_o4.bar(
        x + bar_width,
        vopo_o4_mean,
        width=bar_width,
        label="mean",
    )

    for bars in [bars_x, bars_y, bars_mean]:
        for bar in bars:
            val = bar.get_height()
            ax_o4.text(
                bar.get_x() + bar.get_width() / 2,
                val - 0.35,
                f"{val:.2f}%",
                ha="center",
                va="top",
                fontsize=9,
                rotation=90,
                color="white",
                fontweight="bold",
            )

    ax_o4.axhline(99.5, color="black", linestyle="--", linewidth=1.2)
    ax_o4.axhline(95.0, color="red", linestyle="--", linewidth=1.2)

    ax_o4.set_xticks(x)
    ax_o4.set_xticklabels(sg_labels)
    ax_o4.set_xlabel("ZM2 strain gauge group")
    ax_o4.set_ylabel("Mode overlap with VOPO O4 [%]")
    ax_o4.set_title(
        "Overlap of fitted SG beams with VOPO O4 beam\n"
        f"propagated through ZM2 RoC = {ZM2_ROC_SIGN * ZM2_ROC_3p15:.5f} m"
    )
    ax_o4.grid(True, axis="y", alpha=0.3)
    ax_o4.legend(frameon=True)

    fig_o4.tight_layout()

    # Plot 3: direct FC eigenmode vs VOPO O4 mean-overlap comparison
    if TARGET_BEAM_AVAILABLE:
        fc_mean = []
        fc_x = []
        fc_y = []

        for group_label in fit_group_labels:
            fc_overlaps = compute_target_overlap_for_group(group_label)

            fc_x.append(100 * fc_overlaps["x_overlap"])
            fc_y.append(100 * fc_overlaps["y_overlap"])
            fc_mean.append(100 * fc_overlaps["mean_overlap"])

        fig_fc_o4, ax_fc_o4 = plt.subplots(figsize=(8.2, 5.5))

        compare_bar_width = 0.32

        bars_fc = ax_fc_o4.bar(
            x - compare_bar_width / 2,
            fc_mean,
            width=compare_bar_width,
            label="FC eigenmode mean",
            color="0.55",
        )

        bars_o4 = ax_fc_o4.bar(
            x + compare_bar_width / 2,
            vopo_o4_mean,
            width=compare_bar_width,
            label="VOPO O4 mean",
            color="tab:red",
        )

        for bars in [bars_fc, bars_o4]:
            for bar in bars:
                val = bar.get_height()
                ax_fc_o4.text(
                    bar.get_x() + bar.get_width() / 2,
                    val - 0.35,
                    f"{val:.2f}%",
                    ha="center",
                    va="top",
                    fontsize=9,
                    rotation=90,
                    color="white",
                    fontweight="bold",
                )

        ax_fc_o4.axhline(99.5, color="black", linestyle="--", linewidth=1.2)
        ax_fc_o4.axhline(95.0, color="red", linestyle="--", linewidth=1.2)

        ax_fc_o4.set_xticks(x)
        ax_fc_o4.set_xticklabels(sg_labels)
        ax_fc_o4.set_xlabel("ZM2 strain gauge group")
        ax_fc_o4.set_ylabel("Mean mode overlap [%]")
        ax_fc_o4.set_title(
            "Mean overlap of fitted SG beams\n"
            "with FC eigenmode and VOPO O4 reference"
        )
        ax_fc_o4.grid(True, axis="y", alpha=0.3)
        ax_fc_o4.legend(frameon=True)

        fig_fc_o4.tight_layout()

        # Optional detailed x/y plot for debugging astigmatism.
        fig_fc_o4_xy, ax_fc_o4_xy = plt.subplots(figsize=(9.0, 5.8))

        detailed_bar_width = 0.13
        detailed_offsets = np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5]) * detailed_bar_width

        detailed_series = [
            ("FC x", fc_x, "0.55"),
            ("FC y", fc_y, "0.25"),
            ("FC mean", fc_mean, "0.75"),
            ("O4 x", vopo_o4_x, "tab:red"),
            ("O4 y", vopo_o4_y, "magenta"),
            ("O4 mean", vopo_o4_mean, "tab:pink"),
        ]

        for offset, (label, vals, color) in zip(detailed_offsets, detailed_series):
            ax_fc_o4_xy.bar(
                x + offset,
                vals,
                width=detailed_bar_width,
                label=label,
                color=color,
            )

        ax_fc_o4_xy.axhline(99.5, color="black", linestyle="--", linewidth=1.2)
        ax_fc_o4_xy.axhline(95.0, color="red", linestyle="--", linewidth=1.2)

        ax_fc_o4_xy.set_xticks(x)
        ax_fc_o4_xy.set_xticklabels(sg_labels)
        ax_fc_o4_xy.set_xlabel("ZM2 strain gauge group")
        ax_fc_o4_xy.set_ylabel("Mode overlap [%]")
        ax_fc_o4_xy.set_title(
            "Axis-by-axis overlap of fitted SG beams\n"
            "with FC eigenmode and VOPO O4 reference"
        )
        ax_fc_o4_xy.grid(True, axis="y", alpha=0.3)
        ax_fc_o4_xy.legend(frameon=True, ncols=3, fontsize=8)

        fig_fc_o4_xy.tight_layout()

    
# %%
# PSAM summary plots: available RoC and projected FC overlap vs PZT voltage
#
# These are deliberately made at the end, immediately before the final
# plt.show(), so they appear reliably together with the other summary figures.

if PSAM_CALIBRATION_AVAILABLE:
    print("\nMaking PSAM summary plots at end of script...")
    plot_psam_roc_and_projection(
        PSAM_AVAILABLE_ROC_TABLE,
        PSAM_MULTISTART_PROJECTED_OVERLAP_TABLE,
    )

else:
    print("\nSkipping PSAM summary plots because PSAM calibration was not available.")
    
    

# =============================================================================
# FC1 mode-space plots:
#   (1) two-panel x/y version
#   (2) single-panel combined x+y version
#
# This is the FC1 analog of the familiar SRM/OMC mode-space plot:
#   x-axis = beam radius at FC1 [mm]
#   y-axis = defocus at FC1 = Re(1/q) = 1/R [1/m]
#   contours = constant mismatch to the FC eigenmode [%]
# =============================================================================

# -------------------------
# Helpers for FC1 mode-space coordinates
# -------------------------

def q_fc1_from_q_at_zm3_fcdir(q_at_zm3_fcdir):
    """
    Convert a BeamParam referred to ZM3 in the FC1 -> ZM2 direction
    into a BeamParam referred to FC1 in the same propagation direction.

    The global coordinate is:
        s = 0 at FC1
        s = X_ZM3 at ZM3

    Moving from ZM3 to FC1 changes q.z by:
        z -> z + (X_FC1 - X_ZM3)
    """
    return BeamParam(
        w0=q_at_zm3_fcdir.w0,
        z=q_at_zm3_fcdir.z + (X_FC1 - X_ZM3),
    )


def fc1_radius_defocus_from_q_at_zm3_fcdir(q_at_zm3_fcdir):
    """
    Return:
        w_fc1_m      beam radius at FC1 [m]
        defocus_fc1  Re(1/q) at FC1 [1/m]
        q_fc1        BeamParam referred to FC1, FC1 -> ZM2 direction
    """
    q_fc1 = q_fc1_from_q_at_zm3_fcdir(q_at_zm3_fcdir)
    q_fc1_complex = q_complex_from_beamparam(q_fc1)

    w_fc1_m = BeamParam(w0=q_fc1.w0, z=q_fc1.z).w
    defocus_fc1 = np.real(1.0 / q_fc1_complex)

    return w_fc1_m, defocus_fc1, q_fc1


def beamparam_from_w_defocus_at_fc1(w_fc1_m, defocus_fc1):
    """
    Construct a Gaussian beam at FC1 from:
        w at FC1
        defocus = Re(1/q) = 1/R

    using:
        1/q = 1/R - i * lambda / (pi w^2)
    """
    inv_q = defocus_fc1 - 1j * lambda0 / (np.pi * w_fc1_m**2)
    q_complex = 1.0 / inv_q
    return beamparam_from_q_complex(q_complex)


def mismatch_percent(q1, q2):
    """
    Return mismatch in percent = 100 * (1 - overlap).
    """
    return 100.0 * (1.0 - BeamParam.overlap(q1, q2))


def drop_object_columns_for_csv(df_in):
    """
    Drop non-scalar columns before writing CSV.
    """
    if len(df_in) == 0:
        return df_in

    df_out = df_in.copy()

    for col in list(df_out.columns):
        if df_out[col].dtype == "object":
            # Keep string-like columns, drop arbitrary BeamParam objects.
            non_null = df_out[col].dropna()
            if len(non_null) > 0 and not isinstance(non_null.iloc[0], str):
                df_out = df_out.drop(columns=[col])

    return df_out


# -------------------------
# Collect measured / projected / reference points at FC1
# -------------------------

def build_fc1_modespace_tables():
    """
    Build tables of:
      - measured SG points at FC1
      - projected PSAM points at FC1
      - reference points at FC1

    All q's are expressed in the FC1 -> ZM2 direction before extracting
    beam radius and defocus.
    """
    measured_rows = []
    projected_rows = []
    reference_rows = []

    if not TARGET_BEAM_AVAILABLE:
        return (
            pd.DataFrame(measured_rows),
            pd.DataFrame(projected_rows),
            pd.DataFrame(reference_rows),
        )

    # --- FC eigenmode target at FC1
    for axis_label, q_target_at_zm3 in [("x", qx_target_at_ZM3), ("y", qy_target_at_ZM3)]:
        w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
            q_target_at_zm3
        )

        reference_rows.append({
            "family": "FC target",
            "condition": "FC eigenmode",
            "axis": axis_label,
            "w_fc1_mm": 1e3 * w_m,
            "defocus_fc1_inv_m": defocus,
            "overlap_percent": 100.0,
            "q_fc1": q_fc1,
        })

    # --- Measured SG fits
    for group_label in fit_group_labels:
        for axis_label in ["x", "y"]:
            q_meas_return = q_fit_return_at_ZM3(group_label, axis_label)
            q_meas_fcdir = reverse_q_orientation(q_meas_return)

            w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
                q_meas_fcdir
            )

            q_target = qx_target_at_ZM3 if axis_label == "x" else qy_target_at_ZM3
            ov = 100.0 * BeamParam.overlap(q_meas_fcdir, q_target)

            measured_rows.append({
                "family": "Measured SG",
                "condition": group_label,
                "axis": axis_label,
                "w_fc1_mm": 1e3 * w_m,
                "defocus_fc1_inv_m": defocus,
                "overlap_percent": ov,
                "q_fc1": q_fc1,
            })

    # --- Projected PSAM points from the retained start cases
    #
    # Reload the calibration table here so this block is self-contained
    # and does not depend on a local variable from the earlier try block.
    if len(PSAM_START_CASES) > 0:
        try:
            psam_curves_fc1 = read_psam_roc_calibration(PSAM_CALIBRATION_PATH)

            for case in PSAM_START_CASES:
                start_group_label = case["start_group_label"]
                start_branch = case["start_branch"]
                start_roc = float(case["start_roc_m"])
                start_pzt = float(case["start_pzt_v"])

                if start_group_label not in fit_group_labels:
                    print(
                        f"\nWARNING: FC1 mode-space plot skipping missing group "
                        f"{start_group_label}"
                    )
                    continue

                for scan_name in ["increasing", "decreasing"]:
                    scan_df = filter_psam_scan_to_available_range(
                        psam_curves_fc1[scan_name]
                    )

                    for _, row in scan_df.iterrows():
                        pzt_voltage = float(row["pzt_voltage_v"])
                        roc = float(row["roc_m"])

                        for axis_label in ["x", "y"]:
                            q_proj_return = project_fit_to_new_psam_roc(
                                start_group_label,
                                axis_label,
                                start_roc,
                                roc,
                            )

                            q_proj_fcdir = reverse_q_orientation(q_proj_return)

                            w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
                                q_proj_fcdir
                            )

                            q_target = qx_target_at_ZM3 if axis_label == "x" else qy_target_at_ZM3
                            ov = 100.0 * BeamParam.overlap(q_proj_fcdir, q_target)

                            projected_rows.append({
                                "family": "Projected PSAM",
                                "condition": f"{scan_name} scan",
                                "start_group_label": start_group_label,
                                "start_branch": start_branch,
                                "start_pzt_v": start_pzt,
                                "start_roc_m": start_roc,
                                "axis": axis_label,
                                "pzt_voltage_v": pzt_voltage,
                                "roc_m": roc,
                                "w_fc1_mm": 1e3 * w_m,
                                "defocus_fc1_inv_m": defocus,
                                "overlap_percent": ov,
                                "q_fc1": q_fc1,
                            })

        except Exception as exc:
            print("\nWARNING: Could not build FC1 projected PSAM mode-space points.")
            print(f"Reason: {exc}")

    # --- VOPO target and O4 references, if available
    if VOPO_REFERENCES_AVAILABLE:
        # VOPO target: not astigmatic in this representation, but include it
        # in both panels for easy comparison.
        w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
            q_vopo_target_at_ZM3
        )

        for axis_label, q_target in [("x", qx_target_at_ZM3), ("y", qy_target_at_ZM3)]:
            reference_rows.append({
                "family": "VOPO target",
                "condition": "VOPO target",
                "axis": axis_label,
                "w_fc1_mm": 1e3 * w_m,
                "defocus_fc1_inv_m": defocus,
                "overlap_percent": 100.0 * BeamParam.overlap(q_vopo_target_at_ZM3, q_target),
                "q_fc1": q_fc1,
            })

        # O4 x
        w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
            q_vopo_o4_x_at_ZM3
        )
        reference_rows.append({
            "family": "O4 nominal",
            "condition": "O4 nominal",
            "axis": "x",
            "w_fc1_mm": 1e3 * w_m,
            "defocus_fc1_inv_m": defocus,
            "overlap_percent": 100.0 * BeamParam.overlap(q_vopo_o4_x_at_ZM3, qx_target_at_ZM3),
            "q_fc1": q_fc1,
        })

        # O4 y
        w_m, defocus, q_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(
            q_vopo_o4_y_at_ZM3
        )
        reference_rows.append({
            "family": "O4 nominal",
            "condition": "O4 nominal",
            "axis": "y",
            "w_fc1_mm": 1e3 * w_m,
            "defocus_fc1_inv_m": defocus,
            "overlap_percent": 100.0 * BeamParam.overlap(q_vopo_o4_y_at_ZM3, qy_target_at_ZM3),
            "q_fc1": q_fc1,
        })

    measured_df = pd.DataFrame(measured_rows)
    projected_df = pd.DataFrame(projected_rows)
    reference_df = pd.DataFrame(reference_rows)

    return measured_df, projected_df, reference_df


def get_modespace_limits(values, frac_pad=0.12, min_pad=0.02):
    """
    Get padded plotting limits.
    """
    values = np.asarray(values, dtype=float)
    values = values[np.isfinite(values)]

    if len(values) == 0:
        return -1.0, 1.0

    vmin = np.nanmin(values)
    vmax = np.nanmax(values)

    if np.isclose(vmin, vmax):
        pad = max(min_pad, 0.05 * abs(vmin) if vmin != 0 else min_pad)
    else:
        pad = max(min_pad, frac_pad * (vmax - vmin))

    return vmin - pad, vmax + pad


def build_overlap_grid_single_target(
    q_target_fc1,
    w_mm_lim,
    defocus_lim,
    nw=260,
    nd=260,
):
    """
    Build overlap grid relative to a single target q at FC1.

    The returned grid is overlap in percent, not mismatch.
    """
    w_vals_mm = np.linspace(w_mm_lim[0], w_mm_lim[1], nw)
    d_vals = np.linspace(defocus_lim[0], defocus_lim[1], nd)

    overlap = np.full((len(d_vals), len(w_vals_mm)), np.nan)

    for i, dval in enumerate(d_vals):
        for j, wmm in enumerate(w_vals_mm):
            try:
                q_trial = beamparam_from_w_defocus_at_fc1(1e-3 * wmm, dval)
                overlap[i, j] = 100.0 * BeamParam.overlap(q_trial, q_target_fc1)
            except Exception:
                overlap[i, j] = np.nan

    W, D = np.meshgrid(w_vals_mm, d_vals)
    return W, D, overlap


def build_overlap_grid_average_xy(
    q_target_x_fc1,
    q_target_y_fc1,
    w_mm_lim,
    defocus_lim,
    nw=260,
    nd=260,
):
    """
    Build overlap grid using the average overlap to x and y FC targets.
    This is used only for the single-panel combined view.
    """
    w_vals_mm = np.linspace(w_mm_lim[0], w_mm_lim[1], nw)
    d_vals = np.linspace(defocus_lim[0], defocus_lim[1], nd)

    overlap = np.full((len(d_vals), len(w_vals_mm)), np.nan)

    for i, dval in enumerate(d_vals):
        for j, wmm in enumerate(w_vals_mm):
            try:
                q_trial = beamparam_from_w_defocus_at_fc1(1e-3 * wmm, dval)
                ox = 100.0 * BeamParam.overlap(q_trial, q_target_x_fc1)
                oy = 100.0 * BeamParam.overlap(q_trial, q_target_y_fc1)
                overlap[i, j] = 0.5 * (ox + oy)
            except Exception:
                overlap[i, j] = np.nan

    W, D = np.meshgrid(w_vals_mm, d_vals)
    return W, D, overlap


def fc1_axis_limits_from_points(msub, psub, rsub, frac_pad_xy=1):
    """
    Build tight plot limits from the actual plotted points.

    Only measured points, projected PSAM points, and the FC target are used.
    VOPO target and O4 nominal points are intentionally not included.
    """
    all_w = []
    all_d = []

    if len(msub) > 0:
        all_w.extend(msub["w_fc1_mm"].tolist())
        all_d.extend(msub["defocus_fc1_inv_m"].tolist())

    if len(psub) > 0:
        all_w.extend(psub["w_fc1_mm"].tolist())
        all_d.extend(psub["defocus_fc1_inv_m"].tolist())

    fc_target = rsub[rsub["family"] == "FC target"].copy()
    if len(fc_target) > 0:
        all_w.extend(fc_target["w_fc1_mm"].tolist())
        all_d.extend(fc_target["defocus_fc1_inv_m"].tolist())

    w_lim = get_modespace_limits(np.array(all_w), frac_pad=frac_pad_xy, min_pad=0.015)
    d_lim = get_modespace_limits(np.array(all_d), frac_pad=0.035, min_pad=0.004)

    return w_lim, d_lim


# -------------------------
# Two-panel plot
# -------------------------

def plot_fc1_modespace_two_panel(measured_df, projected_df, reference_df):
    """
    Two-panel analog of the SRM-output style plot:
      left  = x / horizontal
      right = y / vertical

    Contours are overlap with the corresponding FC eigenmode [%].
    """
    if not TARGET_BEAM_AVAILABLE:
        print("\nSkipping FC1 two-panel mode-space plot because TARGET_BEAM_AVAILABLE=False")
        return

    _, _, qx_target_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(qx_target_at_ZM3)
    _, _, qy_target_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(qy_target_at_ZM3)

    fig, axes = plt.subplots(1, 2, figsize=(15, 6), sharey=False)

    contour_levels = [20, 30, 40, 50, 60, 70, 80.0, 85.0, 90.0, 95.0, 97.0, 98.0, 99.0, 99.5, 99.8]

    group_colors = {
        label: plt.cm.tab10(i % 10)
        for i, label in enumerate(fit_group_labels)
    }

    scan_styles = {
        "increasing": {"color": "darkturquoise", "marker": "^", "linestyle": "-"},
        "decreasing": {"color": "deeppink", "marker": "v", "linestyle": "-"},
    }

    for ax, axis_label, q_target_fc1 in zip(
        axes,
        ["x", "y"],
        [qx_target_fc1, qy_target_fc1],
    ):
        msub = measured_df[measured_df["axis"] == axis_label].copy()
        psub = projected_df[projected_df["axis"] == axis_label].copy()
        rsub = reference_df[reference_df["axis"] == axis_label].copy()

        w_lim, d_lim = fc1_axis_limits_from_points(msub, psub, rsub)

        W, D, O = build_overlap_grid_single_target(
            q_target_fc1,
            w_lim,
            d_lim,
        )

        finite_O = O[np.isfinite(O)]
        levels_in_range = [
            lev for lev in contour_levels
            if len(finite_O) > 0 and np.nanmin(finite_O) <= lev <= np.nanmax(finite_O)
        ]

        if len(levels_in_range) > 0:
            cs = ax.contour(
                W,
                D,
                O,
                levels=levels_in_range,
                linewidths=1.6,
            )
            ax.clabel(cs, inline=True, fontsize=9, fmt="%.1f")

        # FC target only. Do not show VOPO target or O4 nominal on this plot.
        fc_target_row = rsub[rsub["family"] == "FC target"]
        if len(fc_target_row) > 0:
            row = fc_target_row.iloc[0]
            ax.scatter(
                row["w_fc1_mm"],
                row["defocus_fc1_inv_m"],
                marker="*",
                s=185,
                color="black",
                edgecolor="white",
                linewidth=0.9,
                label=f"FC target {axis_label}",
                zorder=7,
            )

        # Measured SG points
        for group_label in fit_group_labels:
            subg = msub[msub["condition"] == group_label]
            if len(subg) == 0:
                continue

            row = subg.iloc[0]
            ax.scatter(
                row["w_fc1_mm"],
                row["defocus_fc1_inv_m"],
                marker="o",
                s=78,
                color=group_colors[group_label],
                edgecolor="black",
                linewidth=0.7,
                label=f"{group_label} measured",
                zorder=6,
            )

            ax.annotate(
                f"{group_label}\n{row['overlap_percent']:.1f}%",
                (row["w_fc1_mm"], row["defocus_fc1_inv_m"]),
                xytext=(6, 6),
                textcoords="offset points",
                fontsize=8.4,
                bbox=dict(facecolor="white", edgecolor="none", alpha=0.75, pad=1.2),
            )

        # Projected PSAM curves
        if len(psub) > 0:
            for scan_name in ["increasing", "decreasing"]:
                scan_rows = psub[psub["condition"] == f"{scan_name} scan"].copy()

                if len(scan_rows) == 0:
                    continue

                scan_rows = scan_rows.sort_values("pzt_voltage_v")
                style = scan_styles[scan_name]

                ax.plot(
                    scan_rows["w_fc1_mm"],
                    scan_rows["defocus_fc1_inv_m"],
                    color=style["color"],
                    marker=style["marker"],
                    linestyle=style["linestyle"],
                    linewidth=2.0,
                    markersize=6.0,
                    label=f"Projected {scan_name}",
                    zorder=4,
                )

        ax.set_xlim(w_lim)
        ax.set_ylim(d_lim)

        ax.set_xlabel("Beam radius at FC1 [mm]")
        ax.set_ylabel("Defocus at FC1 [1/m]")
        ax.set_title(
            f"{axis_label.upper()} axis at FC1\nContours: overlap with FC eigenmode [%]"
        )
        ax.grid(True, alpha=0.35)

        handles, labels = ax.get_legend_handles_labels()
        unique = dict(zip(labels, handles))
        ax.legend(
            unique.values(),
            unique.keys(),
            loc="best",
            fontsize=8.2,
            frameon=True,
        )

    fig.suptitle(
        "FC1 mode-space plot: measured and projected beams\n"
        "Two-panel version: x and y treated separately",
        fontsize=14,
        y=0.975,
    )
    fig.tight_layout(rect=[0, 0, 1, 0.94])

    out_png = SCRIPT_DIR / "fc1_modespace_two_panel.png"
    fig.savefig(out_png, dpi=220, bbox_inches="tight")
    print(f"Saved two-panel FC1 mode-space plot to: {out_png}")


# -------------------------
# Single-panel combined plot
# -------------------------

def plot_fc1_modespace_single_panel(measured_df, projected_df, reference_df):
    """
    Single-panel combined version:
      - same condition uses same color
      - x points use triangles
      - y points use circles
      - contours use average overlap to x and y FC targets
    """
    if not TARGET_BEAM_AVAILABLE:
        print("\nSkipping FC1 single-panel mode-space plot because TARGET_BEAM_AVAILABLE=False")
        return

    _, _, qx_target_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(qx_target_at_ZM3)
    _, _, qy_target_fc1 = fc1_radius_defocus_from_q_at_zm3_fcdir(qy_target_at_ZM3)

    fig, ax = plt.subplots(figsize=(11, 7))

    # Tight limits from measured/projected points + FC target only.
    fc_target = reference_df[reference_df["family"] == "FC target"].copy()

    all_w = []
    all_d = []

    if len(measured_df) > 0:
        all_w.extend(measured_df["w_fc1_mm"].tolist())
        all_d.extend(measured_df["defocus_fc1_inv_m"].tolist())

    if len(projected_df) > 0:
        all_w.extend(projected_df["w_fc1_mm"].tolist())
        all_d.extend(projected_df["defocus_fc1_inv_m"].tolist())

    if len(fc_target) > 0:
        all_w.extend(fc_target["w_fc1_mm"].tolist())
        all_d.extend(fc_target["defocus_fc1_inv_m"].tolist())

    w_lim = get_modespace_limits(np.array(all_w), frac_pad=1, min_pad=0.035)
    d_lim = get_modespace_limits(np.array(all_d), frac_pad=0.035, min_pad=0.004)

    contour_levels = [20, 30, 40, 50, 60, 70, 80.0, 85.0, 90.0, 95.0, 97.0, 98.0, 99.0, 99.5, 99.8]

    W, D, O = build_overlap_grid_average_xy(
        qx_target_fc1,
        qy_target_fc1,
        w_lim,
        d_lim,
    )

    finite_O = O[np.isfinite(O)]
    levels_in_range = [
        lev for lev in contour_levels
        if len(finite_O) > 0 and np.nanmin(finite_O) <= lev <= np.nanmax(finite_O)
    ]

    if len(levels_in_range) > 0:
        cs = ax.contour(
            W,
            D,
            O,
            levels=levels_in_range,
            linewidths=1.6,
        )
        ax.clabel(cs, inline=True, fontsize=9, fmt="%.1f")

    axis_markers = {
        "x": "^",
        "y": "o",
    }

    group_colors = {
        label: plt.cm.tab10(i % 10)
        for i, label in enumerate(fit_group_labels)
    }

    projected_scan_colors = {
        "increasing scan": "darkturquoise",
        "decreasing scan": "deeppink",
    }

    # FC eigenmode only. Do not show VOPO target or O4 nominal.
    for _, row in fc_target.iterrows():
        ax.scatter(
            row["w_fc1_mm"],
            row["defocus_fc1_inv_m"],
            marker=axis_markers[row["axis"]],
            s=125,
            color="black",
            edgecolor="white",
            linewidth=0.9,
            label=f"FC target {row['axis']}",
            zorder=7,
        )

    # Measured SG groups
    for group_label in fit_group_labels:
        subg = measured_df[measured_df["condition"] == group_label].copy()
        if len(subg) == 0:
            continue

        color = group_colors[group_label]

        for _, row in subg.iterrows():
            ax.scatter(
                row["w_fc1_mm"],
                row["defocus_fc1_inv_m"],
                marker=axis_markers[row["axis"]],
                s=82,
                color=color,
                edgecolor="black",
                linewidth=0.7,
                label=f"{group_label} {row['axis']}",
                zorder=6,
            )

            ax.annotate(
                f"{group_label} {row['axis']}\n{row['overlap_percent']:.1f}%",
                (row["w_fc1_mm"], row["defocus_fc1_inv_m"]),
                xytext=(6, 5),
                textcoords="offset points",
                fontsize=8.1,
                bbox=dict(facecolor="white", edgecolor="none", alpha=0.75, pad=1.0),
            )

    # Projected PSAM curves
    if len(projected_df) > 0:
        for scan_name in ["increasing scan", "decreasing scan"]:
            scan_color = projected_scan_colors[scan_name]
            subscan = projected_df[projected_df["condition"] == scan_name].copy()

            for axis_label in ["x", "y"]:
                subaxis = subscan[subscan["axis"] == axis_label].copy()

                if len(subaxis) == 0:
                    continue

                subaxis = subaxis.sort_values("pzt_voltage_v")

                ax.plot(
                    subaxis["w_fc1_mm"],
                    subaxis["defocus_fc1_inv_m"],
                    color=scan_color,
                    marker=axis_markers[axis_label],
                    linestyle="-" if axis_label == "x" else "--",
                    linewidth=2.0,
                    markersize=5.8,
                    label=f"{scan_name} {axis_label}",
                    zorder=4,
                )

    ax.set_xlim(w_lim)
    ax.set_ylim(d_lim)

    ax.set_xlabel("Beam radius at FC1 [mm]")
    ax.set_ylabel("Defocus at FC1 [1/m]")
    ax.set_title(
        "FC1 mode-space plot: combined x and y view\n"
        "Contours show average overlap with x/y FC eigenmodes [%]"
    )
    ax.grid(True, alpha=0.35)

    handles, labels = ax.get_legend_handles_labels()
    unique = dict(zip(labels, handles))
    ax.legend(
        unique.values(),
        unique.keys(),
        loc="best",
        fontsize=8.0,
        frameon=True,
        ncols=2,
    )

    fig.tight_layout(rect=[0, 0, 1, 0.97])

    out_png = SCRIPT_DIR / "fc1_modespace_single_panel_xy.png"
    fig.savefig(out_png, dpi=220, bbox_inches="tight")
    print(f"Saved single-panel FC1 mode-space plot to: {out_png}")


# -------------------------
# Build and plot FC1 mode-space figures
# -------------------------

if TARGET_BEAM_AVAILABLE:
    fc1_measured_df, fc1_projected_df, fc1_reference_df = build_fc1_modespace_tables()

    fc1_measured_csv = SCRIPT_DIR / "fc1_modespace_measured_points.csv"
    fc1_projected_csv = SCRIPT_DIR / "fc1_modespace_projected_points.csv"
    fc1_reference_csv = SCRIPT_DIR / "fc1_modespace_reference_points.csv"

    drop_object_columns_for_csv(fc1_measured_df).to_csv(fc1_measured_csv, index=False)
    drop_object_columns_for_csv(fc1_projected_df).to_csv(fc1_projected_csv, index=False)
    drop_object_columns_for_csv(fc1_reference_df).to_csv(fc1_reference_csv, index=False)

    print(f"\nSaved FC1 mode-space measured table to: {fc1_measured_csv}")
    print(f"Saved FC1 mode-space projected table to: {fc1_projected_csv}")
    print(f"Saved FC1 mode-space reference table to: {fc1_reference_csv}")

    plot_fc1_modespace_two_panel(
        fc1_measured_df,
        fc1_projected_df,
        fc1_reference_df,
    )

    plot_fc1_modespace_single_panel(
        fc1_measured_df,
        fc1_projected_df,
        fc1_reference_df,
    )
else:
    print("\nSkipping FC1 mode-space plots because TARGET_BEAM_AVAILABLE=False")



# ============================================================
# True top-view FC overlap map:
# beam radius vs wavefront curvature, color = overlap with FC
#
# This is not an interpolation of the measured/projected points.
# Instead, it builds a dense grid of hypothetical Gaussian beams
# parameterized by:
#     w  = beam radius at ZM3
#     C  = wavefront curvature at ZM3 = 1/R
#
# For each (w, C), it constructs the corresponding Gaussian q:
#     1/q = C - i lambda/(pi w^2)
#
# Then it computes the overlap with the FC eigenmode.
#
# The measured SG points and projected PSAM points are overlaid
# on top of the true overlap map.
# ============================================================

# ------------------------------------------------------------
# Settings
# ------------------------------------------------------------

PSAM_CAL_PATH = SCRIPT_DIR / "data" / "Actuation_Curves_ZM2_SN1.xlsx"
if not PSAM_CAL_PATH.exists():
    # fallback for this chat environment
    PSAM_CAL_PATH = Path("/mnt/data/Actuation_Curves_ZM2_SN1.xlsx")

PZT_MIN_V = 0
PZT_MAX_V = 90

# New calibration:
#   SG 1.x  <-> 0 V
#   SG 3.1x <-> 90 V
#
# For the top-view map, use the same preferred baseline as the v29
# projected-overlap plot: measured 3.x V beam on the decreasing branch.
CURRENT_MEASURED_GROUP = "3.x V"
CURRENT_TOPO_START_ROC_M = 0.85025
CURRENT_TOPO_START_BRANCH = "decreasing"
CURRENT_TOPO_START_PZT_V = 60.0

TOPOMAP_SAVE_PATH = SCRIPT_DIR / "psam_fc_overlap_true_top_view_map.png"
TOPOMAP_DATA_SAVE_PATH = SCRIPT_DIR / "psam_fc_overlap_true_top_view_points.csv"


# ------------------------------------------------------------
# Helper functions for q, beam size, and wavefront curvature
# ------------------------------------------------------------

def q_complex_from_beamparam_topomap(q):
    """
    Convert BeamParam -> complex q = z + i zR.
    """
    zR = np.pi * q.w0**2 / lambda0
    return q.z + 1j * zR


def beamparam_from_q_complex_topomap(q_complex):
    """
    Convert complex q = z + i zR -> BeamParam.
    """
    z = np.real(q_complex)
    zR = np.imag(q_complex)

    if zR <= 0:
        raise ValueError(
            f"Cannot create BeamParam from non-positive zR: q = {q_complex}"
        )

    w0 = np.sqrt(lambda0 * zR / np.pi)

    return BeamParam(w0=w0, z=z)


def q_from_beam_radius_and_wavefront_curvature(w_mm, curvature_inv_m):
    """
    Construct a BeamParam at the reference plane from beam radius and
    wavefront curvature.

    Parameters
    ----------
    w_mm : float or array-like
        Beam radius at the reference plane, in mm.
    curvature_inv_m : float or array-like
        Wavefront curvature C = 1/R at the reference plane, in 1/m.

    Formula
    -------
        1/q = 1/R - i lambda/(pi w^2)
            = C   - i lambda/(pi w^2)

    The returned BeamParam is represented by q = z + i zR.
    """
    w_m = np.asarray(w_mm, dtype=float) * 1e-3
    C = np.asarray(curvature_inv_m, dtype=float)

    inv_q = C - 1j * lambda0 / (np.pi * w_m**2)
    q_complex = 1.0 / inv_q

    if np.ndim(q_complex) == 0:
        return beamparam_from_q_complex_topomap(complex(q_complex))

    return q_complex


def beam_radius_mm_at_ref(q):
    """
    Beam radius at the reference plane, in mm.
    """
    return 1e3 * BeamParam(w0=q.w0, z=q.z).w


def wavefront_curvature_inv_m_at_ref(q):
    """
    Wavefront curvature C = 1/R at the reference plane, in 1/m.

    Since 1/q = 1/R - i lambda/(pi w^2), the wavefront curvature is
    simply Re(1/q).
    """
    q_complex = q_complex_from_beamparam_topomap(q)
    return float(np.real(1.0 / q_complex))


def wavefront_radius_m_at_ref(q):
    """
    Wavefront radius of curvature R at the reference plane, in m.
    This is only used for diagnostic table columns.
    """
    C = wavefront_curvature_inv_m_at_ref(q)

    if np.isclose(C, 0.0):
        return np.inf

    return 1.0 / C


def overlap_percent_topomap(q1, q2):
    """
    BeamParam overlap in percent.
    """
    return 100.0 * BeamParam.overlap(q1, q2)


def load_psam_actuation_curves_topomap(path):
    """
    Parse the PSAM Excel sheet and return clean increasing/decreasing tables.

    Columns:
        pzt_V
        roc_m

    The old strain-gauge column is ignored.
    """
    raw = pd.read_excel(path, sheet_name="ZM2_SN1_20210823", header=None)

    # Increasing scan rows in the provided file
    inc = raw.iloc[7:28, [0, 4]].copy()
    inc.columns = ["pzt_V", "roc_m"]

    # Decreasing scan rows in the provided file
    dec = raw.iloc[30:51, [0, 4]].copy()
    dec.columns = ["pzt_V", "roc_m"]

    for df_ in [inc, dec]:
        df_["pzt_V"] = pd.to_numeric(df_["pzt_V"], errors="coerce")
        df_["roc_m"] = pd.to_numeric(df_["roc_m"], errors="coerce")
        df_.dropna(inplace=True)
        df_.reset_index(drop=True, inplace=True)

    # Restrict to available PZT points in 0--90 V.
    # No interpolation is done.
    inc = inc[(inc["pzt_V"] >= PZT_MIN_V) & (inc["pzt_V"] <= PZT_MAX_V)].copy()
    dec = dec[(dec["pzt_V"] >= PZT_MIN_V) & (dec["pzt_V"] <= PZT_MAX_V)].copy()

    inc.sort_values("pzt_V", inplace=True)
    dec.sort_values("pzt_V", inplace=True)

    inc.reset_index(drop=True, inplace=True)
    dec.reset_index(drop=True, inplace=True)

    return inc, dec


def make_available_voltage_roc_table_topomap(df_inc, df_dec):
    """
    Merge increasing/decreasing scans into one table keyed by PZT voltage.
    """
    table = pd.merge(
        df_inc.rename(columns={"roc_m": "roc_inc_m"}),
        df_dec.rename(columns={"roc_m": "roc_dec_m"}),
        on="pzt_V",
        how="outer",
    ).sort_values("pzt_V").reset_index(drop=True)

    return table


def apply_mirror_roc_topomap(q_before, roc):
    """
    Reflect from a curved mirror with RoC = roc.

    Convention:
        1/q_after = 1/q_before - 2/R
    """
    qc_before = q_complex_from_beamparam_topomap(q_before)
    qc_after = 1.0 / (1.0 / qc_before - 2.0 / roc)
    return beamparam_from_q_complex_topomap(qc_after)


def undo_mirror_roc_topomap(q_after, roc):
    """
    Undo reflective mirror curvature.

    If
        1/q_after = 1/q_before - 2/R
    then
        1/q_before = 1/q_after + 2/R
    """
    qc_after = q_complex_from_beamparam_topomap(q_after)
    qc_before = 1.0 / (1.0 / qc_after + 2.0 / roc)
    return beamparam_from_q_complex_topomap(qc_before)


def q_after_ZM2_from_return_q_at_ZM3_topomap(q_return_at_ZM3):
    """
    Convert return-direction q at ZM3 into q just after ZM2 reflection.

    Return direction:
        ZM2 -> ZM3 -> FC1

    Along propagation from ZM2 to ZM3, z increases by ZM2_to_ZM3.
    Therefore:
        z_after_ZM2 = z_at_ZM3 - ZM2_to_ZM3
    """
    return BeamParam(
        w0=q_return_at_ZM3.w0,
        z=q_return_at_ZM3.z - ZM2_to_ZM3,
    )


def q_return_at_ZM3_from_q_after_ZM2_topomap(q_after_ZM2):
    """
    Propagate q from just after ZM2 reflection to ZM3
    along return direction ZM2 -> ZM3 -> FC1.
    """
    return BeamParam(
        w0=q_after_ZM2.w0,
        z=q_after_ZM2.z + ZM2_to_ZM3,
    )


def fc_q_at_ZM3_return(axis_label):
    """
    FC eigenmode at ZM3, written in the same return-direction convention
    as the measured/projected beams.

    qx_target_at_ZM3 and qy_target_at_ZM3 are oriented FC1 -> ZM2.
    For the top-view map, we use return direction ZM2 -> ZM3 -> FC1.
    """
    if axis_label == "x":
        return reverse_q_orientation(qx_target_at_ZM3)

    if axis_label == "y":
        return reverse_q_orientation(qy_target_at_ZM3)

    raise ValueError(f"Unknown axis label: {axis_label}")


def build_top_view_point_table():
    """
    Build a point table with measured and projected beams.

    Coordinates:
        w_mm               = beam radius at ZM3 [mm]
        curvature_inv_m    = wavefront curvature C = 1/R at ZM3 [1/m]
        R_m                = wavefront radius of curvature at ZM3 [m]

    Values:
        overlap_pct        = overlap with FC eigenmode [%]
    """
    if not TARGET_BEAM_AVAILABLE:
        raise RuntimeError(
            "Cannot build overlap map because FC eigenmode is unavailable."
        )

    if CURRENT_MEASURED_GROUP not in fit_group_labels:
        raise RuntimeError(
            f"{CURRENT_MEASURED_GROUP} not found in fit_group_labels = {fit_group_labels}"
        )

    df_inc, df_dec = load_psam_actuation_curves_topomap(PSAM_CAL_PATH)
    df_available = make_available_voltage_roc_table_topomap(df_inc, df_dec)

    print("\nAvailable PZT / RoC table for top-view map (0--90 V):")
    print(df_available)

    current_topo_roc = float(CURRENT_TOPO_START_ROC_M)

    print(f"\nTop-view map baseline:")
    print(f"  measured group = {CURRENT_MEASURED_GROUP}")
    print(
        f"  undo current RoC using {CURRENT_TOPO_START_BRANCH}-scan "
        f"{CURRENT_TOPO_START_PZT_V:.0f} V value = {current_topo_roc:.6f} m"
    )

    rows = []

    # --------------------------------------------------------
    # Measured SG points
    # --------------------------------------------------------
    for group_label in fit_group_labels:
        for axis_label in ["x", "y"]:
            q_fc = fc_q_at_ZM3_return(axis_label)
            q_meas = q_fit_return_at_ZM3(group_label, axis_label)

            w_mm = beam_radius_mm_at_ref(q_meas)
            C_inv_m = wavefront_curvature_inv_m_at_ref(q_meas)
            R_m = wavefront_radius_m_at_ref(q_meas)
            ovlp = overlap_percent_topomap(q_fc, q_meas)

            rows.append({
                "kind": "measured",
                "scan": "measured",
                "label": group_label,
                "pzt_V": np.nan,
                "axis": axis_label,
                "roc_m": np.nan,
                "w_mm": w_mm,
                "curvature_inv_m": C_inv_m,
                "R_m": R_m,
                "overlap_pct": ovlp,
            })

    # --------------------------------------------------------
    # Projected points from candidate PZT / RoC values
    # --------------------------------------------------------
    for axis_label in ["x", "y"]:
        q_fc = fc_q_at_ZM3_return(axis_label)

        q_current_at_ZM3 = q_fit_return_at_ZM3(
            CURRENT_MEASURED_GROUP,
            axis_label,
        )

        q_current_after_ZM2 = q_after_ZM2_from_return_q_at_ZM3_topomap(
            q_current_at_ZM3
        )

        q_before_ZM2 = undo_mirror_roc_topomap(
            q_current_after_ZM2,
            current_topo_roc,
        )

        for scan_label, scan_df in [("increasing", df_inc), ("decreasing", df_dec)]:
            for _, row in scan_df.iterrows():
                pzt_V = float(row["pzt_V"])
                roc_new = float(row["roc_m"])

                q_after_new = apply_mirror_roc_topomap(q_before_ZM2, roc_new)
                q_proj_at_ZM3 = q_return_at_ZM3_from_q_after_ZM2_topomap(
                    q_after_new
                )

                w_mm = beam_radius_mm_at_ref(q_proj_at_ZM3)
                C_inv_m = wavefront_curvature_inv_m_at_ref(q_proj_at_ZM3)
                R_m = wavefront_radius_m_at_ref(q_proj_at_ZM3)
                ovlp = overlap_percent_topomap(q_fc, q_proj_at_ZM3)

                rows.append({
                    "kind": "projected",
                    "scan": scan_label,
                    "label": f"{scan_label[:3]} {pzt_V:.0f} V",
                    "pzt_V": pzt_V,
                    "axis": axis_label,
                    "roc_m": roc_new,
                    "w_mm": w_mm,
                    "curvature_inv_m": C_inv_m,
                    "R_m": R_m,
                    "overlap_pct": ovlp,
                })

    df_points = pd.DataFrame(rows)

    return df_available, df_points


def make_overlap_grid_for_axis(df_points_axis, axis_label, n_grid=350):
    """
    Make a dense true overlap grid in (w, curvature) space.

    This is the key correction compared with the previous version:
    the color map is NOT interpolated from measured/projected points.
    It is computed directly from the Gaussian mode-overlap formula at each
    point on a regular (w, C) grid.
    """
    q_fc = fc_q_at_ZM3_return(axis_label)

    x_points = df_points_axis["w_mm"].to_numpy(dtype=float)
    y_points = df_points_axis["curvature_inv_m"].to_numpy(dtype=float)

    x_min, x_max = np.nanmin(x_points), np.nanmax(x_points)
    y_min, y_max = np.nanmin(y_points), np.nanmax(y_points)

    x_pad = 0.12 * (x_max - x_min) if x_max > x_min else 0.05 * x_min
    y_pad = 0.18 * (y_max - y_min) if y_max > y_min else 0.05 * max(abs(y_min), 1e-3)

    x_grid = np.linspace(x_min - x_pad, x_max + x_pad, n_grid)
    y_grid = np.linspace(y_min - y_pad, y_max + y_pad, n_grid)

    W_grid, C_grid = np.meshgrid(x_grid, y_grid)

    overlap_grid = np.empty_like(W_grid, dtype=float)

    # Construct q and compute overlap point-by-point.
    # This is not vectorized because BeamParam.overlap expects BeamParam objects.
    for ii in range(W_grid.shape[0]):
        for jj in range(W_grid.shape[1]):
            q_test = q_from_beam_radius_and_wavefront_curvature(
                W_grid[ii, jj],
                C_grid[ii, jj],
            )

            overlap_grid[ii, jj] = overlap_percent_topomap(q_fc, q_test)

    return W_grid, C_grid, overlap_grid


def annotate_overlap_points(ax, df_sub, color, xytext, fontsize, bold=False):
    """
    Annotate point overlaps in percent.
    """
    for _, row in df_sub.iterrows():
        ax.annotate(
            f"{row['overlap_pct']:.1f}%",
            (row["w_mm"], row["curvature_inv_m"]),
            xytext=xytext,
            textcoords="offset points",
            fontsize=fontsize,
            color=color,
            fontweight="bold" if bold else "normal",
            bbox=dict(
                facecolor="white",
                edgecolor="none",
                alpha=0.72,
                pad=0.4,
            ),
        )


def add_true_top_view_overlap_panel(ax, df_points_axis, axis_label, axis_title,
                                    vmin=None, vmax=None):
    """
    Make one true top-view overlap map.

    Axes:
        x = beam radius at ZM3 [mm]
        y = wavefront curvature C = 1/R at ZM3 [1/m]

    Color:
        overlap with FC eigenmode [%]

    No contour lines are used; the view is entirely top-down.
    """
    W_grid, C_grid, overlap_grid = make_overlap_grid_for_axis(
        df_points_axis,
        axis_label,
    )

    im = ax.pcolormesh(
        W_grid,
        C_grid,
        overlap_grid,
        shading="auto",
        cmap="viridis",
        vmin=vmin,
        vmax=vmax,
    )

    # Split into measured / projected subsets.
    df_meas = df_points_axis[df_points_axis["kind"] == "measured"]

    df_inc = df_points_axis[
        (df_points_axis["kind"] == "projected")
        & (df_points_axis["scan"] == "increasing")
    ]

    df_dec = df_points_axis[
        (df_points_axis["kind"] == "projected")
        & (df_points_axis["scan"] == "decreasing")
    ]

    # --------------------------------------------------------
    # Plot projected points first.
    # Increasing scan: upward triangle.
    # Decreasing scan: downward triangle.
    # --------------------------------------------------------
    ax.scatter(
        df_inc["w_mm"],
        df_inc["curvature_inv_m"],
        s=76,
        marker="^",
        facecolor="tab:blue",
        edgecolor="black",
        linewidth=0.85,
        zorder=5,
        label="projected increasing",
    )

    ax.scatter(
        df_dec["w_mm"],
        df_dec["curvature_inv_m"],
        s=76,
        marker="v",
        facecolor="tab:red",
        edgecolor="black",
        linewidth=0.85,
        zorder=5,
        label="projected decreasing",
    )

    # --------------------------------------------------------
    # Plot measured points last so they stay on top.
    # --------------------------------------------------------
    ax.scatter(
        df_meas["w_mm"],
        df_meas["curvature_inv_m"],
        s=130,
        marker="o",
        facecolor="white",
        edgecolor="black",
        linewidth=1.55,
        zorder=8,
        label="measured SG",
    )

    # --------------------------------------------------------
    # Annotate projected increasing points BELOW the curve/points.
    # --------------------------------------------------------
    for _, row in df_inc.iterrows():
        ax.annotate(
            f"{row['overlap_pct']:.1f}%",
            (row["w_mm"], row["curvature_inv_m"]),
            xytext=(5, -13),
            textcoords="offset points",
            fontsize=7.8,
            color="tab:blue",
            ha="left",
            va="top",
            bbox=dict(
                facecolor="white",
                edgecolor="none",
                alpha=0.74,
                pad=0.40,
            ),
            zorder=7,
        )

    # --------------------------------------------------------
    # Annotate projected decreasing points ABOVE the curve/points.
    # --------------------------------------------------------
    for _, row in df_dec.iterrows():
        ax.annotate(
            f"{row['overlap_pct']:.1f}%",
            (row["w_mm"], row["curvature_inv_m"]),
            xytext=(5, 9),
            textcoords="offset points",
            fontsize=7.8,
            color="tab:red",
            ha="left",
            va="bottom",
            bbox=dict(
                facecolor="white",
                edgecolor="none",
                alpha=0.74,
                pad=0.40,
            ),
            zorder=7,
        )

    # --------------------------------------------------------
    # Annotate measured points last, with larger font.
    # --------------------------------------------------------
    for _, row in df_meas.iterrows():
        ax.annotate(
            f"{row['overlap_pct']:.1f}%",
            (row["w_mm"], row["curvature_inv_m"]),
            xytext=(8, 8),
            textcoords="offset points",
            fontsize=11.0,
            color="black",
            fontweight="bold",
            ha="left",
            va="bottom",
            bbox=dict(
                facecolor="white",
                edgecolor="black",
                alpha=0.88,
                pad=0.60,
            ),
            zorder=9,
        )

    ax.set_xlabel("Beam radius at ZM3 [mm]")
    ax.set_ylabel(r"Wavefront curvature at ZM3, $1/R$ [1/m]")
    ax.set_title(axis_title)
    ax.grid(False)

    # Legend order: measured first, then increasing, then decreasing.
    handles, labels = ax.get_legend_handles_labels()
    desired_order = [
        "measured SG",
        "projected increasing",
        "projected decreasing",
    ]

    ordered_handles = []
    ordered_labels = []

    for wanted in desired_order:
        for handle, label in zip(handles, labels):
            if label == wanted:
                ordered_handles.append(handle)
                ordered_labels.append(label)

    ax.legend(
        ordered_handles,
        ordered_labels,
        frameon=True,
        fontsize=8,
        loc="best",
    )

    return im


# ------------------------------------------------------------
# Build tables and true top-view overlap maps
# ------------------------------------------------------------

if TARGET_BEAM_AVAILABLE:
    df_available_roc, df_topo_points = build_top_view_point_table()

    print("\nTrue top-view overlap point table:")
    print(df_topo_points)

    df_topo_points.to_csv(TOPOMAP_DATA_SAVE_PATH, index=False)
    print(f"\nSaved top-view point table to: {TOPOMAP_DATA_SAVE_PATH}")

    # Use one shared color scale for x and y.
    # Since the grid is evaluated directly, compute the scale from the point
    # values but keep a little room at the top/bottom.
    point_min = np.nanmin(df_topo_points["overlap_pct"].values)
    point_max = np.nanmax(df_topo_points["overlap_pct"].values)
    color_pad = 0.05 * (point_max - point_min) if point_max > point_min else 1.0

    global_vmin = max(0.0, point_min - color_pad)
    global_vmax = min(100.0, point_max + color_pad)

    fig_topo, (ax_topo_x, ax_topo_y) = plt.subplots(
        1,
        2,
        figsize=(15.5, 6.2),
        constrained_layout=True,
    )

    im_x = add_true_top_view_overlap_panel(
        ax_topo_x,
        df_topo_points[df_topo_points["axis"] == "x"],
        axis_label="x",
        axis_title="X-axis overlap map",
        vmin=global_vmin,
        vmax=global_vmax,
    )

    im_y = add_true_top_view_overlap_panel(
        ax_topo_y,
        df_topo_points[df_topo_points["axis"] == "y"],
        axis_label="y",
        axis_title="Y-axis overlap map",
        vmin=global_vmin,
        vmax=global_vmax,
    )

    cbar = fig_topo.colorbar(im_y, ax=[ax_topo_x, ax_topo_y], shrink=0.95)
    cbar.set_label("Overlap with FC eigenmode [%]")

    fig_topo.suptitle(
        "True top-view FC mode-overlap map\n"
        r"color = overlap computed directly over $(w,\,1/R)$ space",
        fontsize=14,
    )

    fig_topo.savefig(TOPOMAP_SAVE_PATH, dpi=220, bbox_inches="tight")
    print(f"Saved true top-view overlap figure to: {TOPOMAP_SAVE_PATH}")

else:
    print("\nSkipping true top-view overlap map because FC eigenmode was not available.")


plt.show()
