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

@author: begum


Fit Gaussian beam parameters from direct beam-profiler measurements,
compute overlaps with target/O4, and make a GIF of the propagation plot.

GIF behavior:
  - cycles through A:L2 = 0 mm, -13 mm, +17 mm
  - each lens setting is shown for 2 seconds
  - each frame shows:
      target
      O4 x
      O4 y
      fitted x beam
      fitted y beam
      measured x/y data points
      +w and -w propagation

Units:
  z values are in mm from ZM1
  wx/wy values are in um
  sigma_wx/sigma_wy values are in um

Here wx/wy are raw profiler sizes, so INPUT_SIZE_IS_DIAMETER = True.
If you manually divide all wx/wy values by 2 before this block,
change INPUT_SIZE_IS_DIAMETER = False.

"""

import os
import tempfile
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

try:
    from PIL import Image
except ImportError:
    raise ImportError(
        "This GIF code needs Pillow. Install it with: pip install pillow"
    )

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

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


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

lambda0 = 1064e-9  # [m]

INPUT_SIZE_IS_DIAMETER = True

# Prevent tiny reported sigmas from over-weighting one point too strongly.
SIGMA_FLOOR_UM = 0.5

# How far upstream of ZM1 to show in propagation plot
Z_UPSTREAM_OF_ZM1 = 0.10  # [m]

# Vertical location of waist text labels, in mm on the plot.
# Put x waist below zero and y waist above zero.
WAIST_LABEL_Y_MM = {
    "x": -0.12,
    "y": 0.12,
}

GIF_OUTPUT_PATH = "beam_propagation_lens_scan.gif"
GIF_FRAME_DURATION_MS = 2000
GIF_DPI = 140


# -------------------------
# Target and O4 VIP beams at ZM1
# -------------------------

AM3_to_ZM1 = 158.2e-3

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


# -------------------------
# Direct beam-profiler measurements
# -------------------------

# A:L2 = 0 mm
z_p13_l0 = 202
wx_p13_l0 = 729.029
wy_p13_l0 = 728.033
sigma_wx_p13_l0 = 0.1794
sigma_wy_p13_l0 = 0.0524

z_p14_l0 = 80 + 240
wx_p14_l0 = 1015.436
wy_p14_l0 = 1031.594
sigma_wx_p14_l0 = 0.1631
sigma_wy_p14_l0 = 0.1546

z_p15_l0 = 160 + 240
wx_p15_l0 = 1239.723
wy_p15_l0 = 1259.941
sigma_wx_p15_l0 = 0.1279
sigma_wy_p15_l0 = 0.1590

z_p16_l0 = 218 + 240
wx_p16_l0 = 1406.296
wy_p16_l0 = 1441.320
sigma_wx_p16_l0 = 0.1376
sigma_wy_p16_l0 = 0.3010

z_p17_l0 = 315 + 240
wx_p17_l0 = 1689.945
wy_p17_l0 = 1749.56
sigma_wx_p17_l0 = 0.3463
sigma_wy_p17_l0 = 0.3902


# A:L2 = -13 mm
z_p21_l13 = 68 + 240
wx_p21_l13 = 946.038
wy_p21_l13 = 959.841
sigma_wx_p21_l13 = 0.7827
sigma_wy_p21_l13 = 1.0425

z_p20_l13 = 100 + 240
wx_p20_l13 = 1037.039
wy_p20_l13 = 1054.128
sigma_wx_p20_l13 = 0.127
sigma_wy_p20_l13 = 0.5848

z_p19_l13 = 176 + 240
wx_p19_l13 = 1260.202
wy_p19_l13 = 1292.843
sigma_wx_p19_l13 = 0.1148
sigma_wy_p19_l13 = 0.5305

z_p18_l13 = 245 + 240
wx_p18_l13 = 1455.835
wy_p18_l13 = 1486.505
sigma_wx_p18_l13 = 82.7851
sigma_wy_p18_l13 = 10.22

z_p17_l13 = 315 + 240
wx_p17_l13 = 1661.984
wy_p17_l13 = 1709.732
sigma_wx_p17_l13 = 0.1614
sigma_wy_p17_l13 = 0.1092


# A:L2 = +17 mm
z_p21_l17 = 68 + 240
wx_p21_l17 = 1035.985
wy_p21_l17 = 1056.579
sigma_wx_p21_l17 = 0.0675
sigma_wy_p21_l17 = 0.1429

z_p20_l17 = 100 + 240
wx_p20_l17 = 1127.876
wy_p20_l17 = 1152.449
sigma_wx_p20_l17 = 0.3735
sigma_wy_p20_l17 = 0.2607

z_p19_l17 = 176 + 240
wx_p19_l17 = 1338.529
wy_p19_l17 = 1362.110
sigma_wx_p19_l17 = 0.1214
sigma_wy_p19_l17 = 0.1618

z_p18_l17 = 245 + 240
wx_p18_l17 = 1543.718
wy_p18_l17 = 1558.867
sigma_wx_p18_l17 = 0.7390
sigma_wy_p18_l17 = 0.2033

z_p17_l17 = 315 + 240
wx_p17_l17 = 1737.392
wy_p17_l17 = 1783.611
sigma_wx_p17_l17 = 0.2439
sigma_wy_p17_l17 = 0.2429


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

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 convert_sigma_um_to_radius_m(sigma_um):
    sigma_m = np.asarray(sigma_um, dtype=float) * 1e-6
    if INPUT_SIZE_IS_DIAMETER:
        return sigma_m / 2
    return sigma_m


def gaussian_beam_w(z, w0, z0):
    """
    Gaussian beam radius w(z).

    z  : position downstream from ZM1 [m]
    w0 : waist radius [m]
    z0 : waist position downstream from ZM1 [m]
    """
    zR = np.pi * w0**2 / lambda0
    return w0 * np.sqrt(1 + ((z - z0) / zR)**2)


def fit_beam(z_mm, w_um, sigma_um=None, label="", use_sigma=True):
    z = np.asarray(z_mm, dtype=float) * 1e-3
    w = convert_size_um_to_radius_m(w_um)

    if use_sigma and sigma_um is not None:
        sigma_um_eff = np.maximum(np.asarray(sigma_um, dtype=float), SIGMA_FLOOR_UM)
        sigma = convert_sigma_um_to_radius_m(sigma_um_eff)
        absolute_sigma = True
    else:
        sigma = None
        absolute_sigma = False

    idx_min = np.argmin(w)
    w0_guess = 0.8 * w[idx_min]
    z0_guess = z[idx_min] - 0.10

    popt, pcov = curve_fit(
        gaussian_beam_w,
        z,
        w,
        p0=[w0_guess, z0_guess],
        sigma=sigma,
        absolute_sigma=absolute_sigma,
        bounds=(
            [1e-6, -5.0],
            [5e-3,  5.0],
        ),
        maxfev=10000,
    )

    w0_fit, z0_fit = popt
    zR_fit = np.pi * w0_fit**2 / lambda0

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

    q_fit = BeamParam(w0=w0_fit, z=z0_fit)

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

    return {
        "label": label,
        "z_m": z,
        "w_m": w,
        "sigma_m": sigma,
        "w0": w0_fit,
        "z0": z0_fit,
        "zR": zR_fit,
        "q": q_fit,
        "pcov": pcov,
        "sigma_w0": sigma_w0,
        "sigma_z0": sigma_z0,
        "rms_residual": rms_residual,
    }


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 w_along_from_waist_location(q, x_array):
    """
    Beam radius along downstream coordinate x from ZM1.

    q.w0 = waist radius
    q.z  = waist location relative to ZM1
    """
    return np.array([
        BeamParam(w0=q.w0, z=x - q.z).w
        for x in x_array
    ])


def ov(q1, q2):
    return 100 * BeamParam.overlap(q1, q2)


# -------------------------
# Assemble data sets
# -------------------------

data_l0 = {
    "label": "A:L2 0 mm",
    "short_label": "0 mm",
    "z_mm": np.array([z_p13_l0, z_p14_l0, z_p15_l0, z_p16_l0, z_p17_l0]),
    "wx_um": np.array([wx_p13_l0, wx_p14_l0, wx_p15_l0, wx_p16_l0, wx_p17_l0]),
    "wy_um": np.array([wy_p13_l0, wy_p14_l0, wy_p15_l0, wy_p16_l0, wy_p17_l0]),
    "sigma_wx_um": np.array([sigma_wx_p13_l0, sigma_wx_p14_l0, sigma_wx_p15_l0, sigma_wx_p16_l0, sigma_wx_p17_l0]),
    "sigma_wy_um": np.array([sigma_wy_p13_l0, sigma_wy_p14_l0, sigma_wy_p15_l0, sigma_wy_p16_l0, sigma_wy_p17_l0]),
}

data_lm13 = {
    "label": "A:L2 -13 mm",
    "short_label": "-13 mm",
    "z_mm": np.array([z_p21_l13, z_p20_l13, z_p19_l13, z_p18_l13, z_p17_l13]),
    "wx_um": np.array([wx_p21_l13, wx_p20_l13, wx_p19_l13, wx_p18_l13, wx_p17_l13]),
    "wy_um": np.array([wy_p21_l13, wy_p20_l13, wy_p19_l13, wy_p18_l13, wy_p17_l13]),
    "sigma_wx_um": np.array([sigma_wx_p21_l13, sigma_wx_p20_l13, sigma_wx_p19_l13, sigma_wx_p18_l13, sigma_wx_p17_l13]),
    "sigma_wy_um": np.array([sigma_wy_p21_l13, sigma_wy_p20_l13, sigma_wy_p19_l13, sigma_wy_p18_l13, sigma_wy_p17_l13]),
}

data_lp17 = {
    "label": "A:L2 +17 mm",
    "short_label": "+17 mm",
    "z_mm": np.array([z_p21_l17, z_p20_l17, z_p19_l17, z_p18_l17, z_p17_l17]),
    "wx_um": np.array([wx_p21_l17, wx_p20_l17, wx_p19_l17, wx_p18_l17, wx_p17_l17]),
    "wy_um": np.array([wy_p21_l17, wy_p20_l17, wy_p19_l17, wy_p18_l17, wy_p17_l17]),
    "sigma_wx_um": np.array([sigma_wx_p21_l17, sigma_wx_p20_l17, sigma_wx_p19_l17, sigma_wx_p18_l17, sigma_wx_p17_l17]),
    "sigma_wy_um": np.array([sigma_wy_p21_l17, sigma_wy_p20_l17, sigma_wy_p19_l17, sigma_wy_p18_l17, sigma_wy_p17_l17]),
}

datasets = [
    ("A:L2 0 mm", data_l0),
    ("A:L2 -13 mm", data_lm13),
    ("A:L2 +17 mm", data_lp17),
]

choice_to_lens = {
    "0 mm": "A:L2 0 mm",
    "-13 mm": "A:L2 -13 mm",
    "+17 mm": "A:L2 +17 mm",
}

lens_to_data = {
    "A:L2 0 mm": data_l0,
    "A:L2 -13 mm": data_lm13,
    "A:L2 +17 mm": data_lp17,
}


# -------------------------
# Fit x/y beams
# -------------------------

fits = {}

for lens_label, data in datasets:
    fits[(lens_label, "x")] = fit_beam(
        data["z_mm"],
        data["wx_um"],
        data["sigma_wx_um"],
        label=f"{lens_label} x",
        use_sigma=True,
    )

    fits[(lens_label, "y")] = fit_beam(
        data["z_mm"],
        data["wy_um"],
        data["sigma_wy_um"],
        label=f"{lens_label} y",
        use_sigma=True,
    )


# -------------------------
# BeamParam objects from weighted fits
# -------------------------

q_fit_l0_x = fits[("A:L2 0 mm", "x")]["q"]
q_fit_l0_y = fits[("A:L2 0 mm", "y")]["q"]

q_fit_lm13_x = fits[("A:L2 -13 mm", "x")]["q"]
q_fit_lm13_y = fits[("A:L2 -13 mm", "y")]["q"]

q_fit_lp17_x = fits[("A:L2 +17 mm", "x")]["q"]
q_fit_lp17_y = fits[("A:L2 +17 mm", "y")]["q"]


# -------------------------
# Print fit results
# -------------------------

fit_rows = []

for lens_label, axis_label in [
    ("A:L2 0 mm", "x"),
    ("A:L2 0 mm", "y"),
    ("A:L2 -13 mm", "x"),
    ("A:L2 -13 mm", "y"),
    ("A:L2 +17 mm", "x"),
    ("A:L2 +17 mm", "y"),
]:
    fit = fits[(lens_label, axis_label)]

    fit_rows.append([
        lens_label,
        axis_label,
        1e3 * fit["w0"],
        1e3 * fit["sigma_w0"],
        1e3 * fit["z0"],
        1e3 * fit["sigma_z0"],
        1e3 * fit["zR"],
        1e6 * fit["rms_residual"],
    ])

print("\nWeighted Gaussian beam fit results:")
print_table(
    fit_rows,
    headers=[
        "lens",
        "axis",
        "w0 [mm]",
        "sigma w0 [mm]",
        "z0 from ZM1 [mm]",
        "sigma z0 [mm]",
        "zR [mm]",
        "RMS residual [um]",
    ],
    floatfmt=".3f",
)


# -------------------------
# Compute overlaps with target and O4
# -------------------------

overlap_rows = []

fit_qs_for_overlap = [
    ("0 mm", q_fit_l0_x, q_fit_l0_y),
    ("-13 mm", q_fit_lm13_x, q_fit_lm13_y),
    ("+17 mm", q_fit_lp17_x, q_fit_lp17_y),
]

for lens_short, qx_fit, qy_fit in fit_qs_for_overlap:
    overlap_rows.append([
        lens_short,
        "Target",
        ov(q_V0_target, qx_fit),
        ov(q_V0_target, qy_fit),
    ])

    overlap_rows.append([
        lens_short,
        "O4",
        ov(q_V0_O4_x, qx_fit),
        ov(q_V0_O4_y, qy_fit),
    ])

print("\nOverlap of fitted beams with target/O4:")
print_table(
    overlap_rows,
    headers=[
        "A:L2",
        "reference",
        "x overlap [%]",
        "y overlap [%]",
    ],
    floatfmt=".2f",
)


# -------------------------
# Plot overlap bar chart
# -------------------------

overlap_labels = [
    "Target\n0 mm",
    "Target\n-13 mm",
    "Target\n+17 mm",
    "O4\n0 mm",
    "O4\n-13 mm",
    "O4\n+17 mm",
]

overlap_fit_plot = np.array([
    [ov(q_V0_target, q_fit_l0_x),    ov(q_V0_target, q_fit_l0_y)],
    [ov(q_V0_target, q_fit_lm13_x),  ov(q_V0_target, q_fit_lm13_y)],
    [ov(q_V0_target, q_fit_lp17_x),  ov(q_V0_target, q_fit_lp17_y)],

    [ov(q_V0_O4_x, q_fit_l0_x),      ov(q_V0_O4_y, q_fit_l0_y)],
    [ov(q_V0_O4_x, q_fit_lm13_x),    ov(q_V0_O4_y, q_fit_lm13_y)],
    [ov(q_V0_O4_x, q_fit_lp17_x),    ov(q_V0_O4_y, q_fit_lp17_y)],
])

component_labels = ["x", "y"]

fig, ax = plt.subplots(figsize=(8.5, 6))

x = np.array([0.0, 0.75, 1.5, 2.75, 3.5, 4.25])
bar_width = 0.26

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

for i, comp in enumerate(component_labels):
    offset = (i - 0.5) * bar_width
    vals = overlap_fit_plot[:, i]

    bars = ax.bar(
        x + offset,
        vals,
        width=bar_width,
        color=bar_colors[comp],
        label=comp,
    )

    for bar, val in zip(bars, vals):
        ax.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.axhline(99.5, color="black", linestyle="--", linewidth=1.3)
ax.axhline(95, color="red", linestyle="--", linewidth=1.3)

separator_x = 2.15
ax.axvline(separator_x, color="0.7", linestyle=":", linewidth=1.2)

ax.text(
    separator_x,
    99.5,
    "99.5%",
    color="black",
    ha="center",
    va="center",
    fontsize=10,
    bbox=dict(facecolor="white", edgecolor="none", alpha=0.8, pad=1.5),
)

ax.text(
    separator_x,
    95,
    "95%",
    color="red",
    ha="center",
    va="center",
    fontsize=10,
    bbox=dict(facecolor="white", edgecolor="none", alpha=0.8, pad=1.5),
)

all_overlap_vals = overlap_fit_plot.reshape(-1)
ymin = max(0, np.floor(np.nanmin(all_overlap_vals)) - 5)
if ymin > 90:
    ymin = 90

ax.set_ylim(ymin, 100.2)
ax.set_xticks(x)
ax.set_xticklabels(overlap_labels)
ax.set_ylabel("Mode overlap [%]")
ax.set_title(
    "Overlap of fitted direct-profiler beams\n"
    "with target and O4 VIP beam profiles"
)

ax.grid(True, axis="y", alpha=0.3)

ax.legend(
    frameon=True,
    loc="upper left",
    bbox_to_anchor=(1.0, 1.0),
)

fig.tight_layout()
plt.show()


# -------------------------
# GIF propagation plot
# -------------------------

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

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

reference_beams = {
    "target": {
        "q": q_V0_target,
        "color": "black",
        "linestyle": "-",
        "linewidth": 2.8,
    },
    "O4 x": {
        "q": q_V0_O4_x,
        "color": "0.35",
        "linestyle": "-.",
        "linewidth": 2.4,
    },
    "O4 y": {
        "q": q_V0_O4_y,
        "color": "0.35",
        "linestyle": "--",
        "linewidth": 2.4,
    },
}

all_z_data_m = np.concatenate([
    data_l0["z_mm"],
    data_lm13["z_mm"],
    data_lp17["z_mm"],
]) * 1e-3

z_min_global = -Z_UPSTREAM_OF_ZM1
z_max_global = np.max(all_z_data_m) + 0.12
z_plot_global = np.linspace(z_min_global, z_max_global, 1200)

gif_choices = ["0 mm", "-13 mm", "+17 mm"]


def draw_selected_lens_for_gif(choice_label, save_path=None):
    """
    Draw selected A:L2 setting and optionally save as PNG.
    """
    fig = plt.figure(figsize=(12.5, 6.9))
    ax = fig.add_axes([0.08, 0.12, 0.68, 0.74])

    lens_label = choice_to_lens[choice_label]
    data = lens_to_data[lens_label]
    z_data = data["z_mm"] * 1e-3

    # Reference target/O4 beams
    for ref_label, ref in reference_beams.items():
        q_ref = ref["q"]
        w_ref = w_along_from_waist_location(q_ref, z_plot_global)

        ax.plot(
            z_plot_global,
            1e3 * w_ref,
            color=ref["color"],
            linestyle=ref["linestyle"],
            linewidth=ref["linewidth"],
            label=ref_label,
        )

        ax.plot(
            z_plot_global,
            -1e3 * w_ref,
            color=ref["color"],
            linestyle=ref["linestyle"],
            linewidth=ref["linewidth"],
        )

        ax.plot(
            q_ref.z,
            0,
            marker="o",
            color=ref["color"],
            markersize=4,
        )

    # Fitted beams and measured points
    for axis_label, w_key, sigma_key in [
        ("x", "wx_um", "sigma_wx_um"),
        ("y", "wy_um", "sigma_wy_um"),
    ]:
        fit = fits[(lens_label, axis_label)]

        w_data = convert_size_um_to_radius_m(data[w_key])
        sigma_data = convert_sigma_um_to_radius_m(
            np.maximum(data[sigma_key], SIGMA_FLOOR_UM)
        )

        w_fit = gaussian_beam_w(z_plot_global, fit["w0"], fit["z0"])

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

        ax.plot(
            z_plot_global,
            -1e3 * w_fit,
            color=plot_colors[axis_label],
            linestyle=plot_linestyles[axis_label],
            linewidth=2.2,
        )

        # Measured data points: +w and -w
        ax.errorbar(
            z_data,
            1e3 * w_data,
            yerr=1e3 * sigma_data,
            fmt="o",
            color=plot_colors[axis_label],
            alpha=0.85,
            capsize=3,
            label=f"{axis_label} data",
        )

        ax.errorbar(
            z_data,
            -1e3 * w_data,
            yerr=1e3 * sigma_data,
            fmt="o",
            color=plot_colors[axis_label],
            alpha=0.85,
            capsize=3,
        )

        # Mark fitted waist
        ax.plot(
            fit["z0"],
            0,
            marker="x",
            color=plot_colors[axis_label],
            markersize=9,
            mew=2,
        )

        # Waist label: x below zero, y above zero
        waist_label_y = WAIST_LABEL_Y_MM[axis_label]

        ax.text(
            fit["z0"],
            waist_label_y,
            f" {axis_label} waist",
            color=plot_colors[axis_label],
            ha="left",
            va="top" if waist_label_y < 0 else "bottom",
            fontsize=8,
        )

    # Mark ZM1
    ax.axvline(
        0.0,
        color="0.35",
        linestyle="-",
        linewidth=1.3,
        alpha=0.9,
    )

    # Mark measurement planes
    for zz in z_data:
        ax.axvline(
            zz,
            color="0.75",
            linestyle=":",
            linewidth=1.0,
            alpha=0.8,
        )

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

    ax.set_xlim(z_min_global, z_max_global)
    ax.set_xlabel("Distance downstream from ZM1 [m]")
    ax.set_ylabel("Beam radius ±w [mm]")

    ax.set_title(
        f"A:L2 = {choice_label}",
        pad=16,
    )

    ax.grid(True, alpha=0.3)

    ymin, ymax = ax.get_ylim()

    ax.text(
        0.0,
        ymin + 0.96 * (ymax - ymin),
        " ZM1",
        ha="left",
        va="top",
        fontsize=9,
        rotation=90,
        color="0.25",
    )

    for zz_mm, zz_m in zip(data["z_mm"], z_data):
        ax.text(
            zz_m,
            ymin + 0.96 * (ymax - ymin),
            f" {zz_mm:.0f} mm",
            ha="left",
            va="top",
            fontsize=7,
            rotation=90,
            color="0.35",
        )

    ax.legend(
        frameon=True,
        fontsize=8,
        loc="upper left",
        bbox_to_anchor=(1.01, 1.0),
    )

    fig.suptitle(
        "Fitted direct-profiler beam propagation\n"
        "with target and O4 VIP beam profiles",
        y=0.995,
    )

    if save_path is not None:
        fig.savefig(save_path, dpi=GIF_DPI)
        plt.close(fig)
    else:
        plt.show()

    return fig


# Preview one static frame
draw_selected_lens_for_gif("0 mm", save_path=None)


# Build GIF
with tempfile.TemporaryDirectory() as tmpdir:
    frame_paths = []

    for i, choice_label in enumerate(gif_choices):
        frame_path = os.path.join(tmpdir, f"frame_{i:02d}_{choice_label.replace('+', 'p').replace('-', 'm').replace(' ', '')}.png")
        draw_selected_lens_for_gif(choice_label, save_path=frame_path)
        frame_paths.append(frame_path)

    frames = [Image.open(path).convert("P", palette=Image.ADAPTIVE) for path in frame_paths]

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

print(f"\nSaved GIF to: {os.path.abspath(GIF_OUTPUT_PATH)}")