"""Reproduce Fig 3 (diameter offsets, 4 modes at df = +2 / -2 nm) from the
regenerated 6930-image simulated set, using the Apr-1 helpers unchanged.

Replicates process_file()'s diameter path exactly:
    load -> auto_roi -> normalize[0,1] -> sum(axis=0) -> measure_diameter_all_modes

Run from the VibeTube-0 root (or anywhere):
    python outputs/sim/60kV/.../Diameter_Analysis/reproduce_fig3.py
Outputs (written next to this script):
    diameter_offsets_4modes_df_pm2.csv, fig3_a_df+2nm_underfocus.png,
    fig3_b_df-2nm_overfocus.png
"""
import sys, csv, re
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = Path(__file__).resolve()
VIBETUBE_ROOT = next((p for p in HERE.parents if (p / "analysis" / "chirality_helpers.py").exists()), HERE.parent)
SIM = VIBETUBE_ROOT / "data" / "simulated" / "6-0-to-25-25"                     # input images
OUT = VIBETUBE_ROOT / "results" / "simulated" / "6-0-to-25-25" / "Diameter_Analysis"  # output
OUT.mkdir(parents=True, exist_ok=True)
sys.path.insert(0, str(VIBETUBE_ROOT / "analysis"))
import chirality_helpers as h                   # noqa: E402
SCALE = 0.05  # nm/px (0.5 A/px)
MODES = ["d_min", "d_max", "d_ave", "d_inf"]
COL = {"d_min": "#1f77b4", "d_max": "#ff7f0e", "d_ave": "#2ca02c", "d_inf": "#d62728"}
name_re = re.compile(r"tem_image_\((\d+),(\d+)\)_df([+-][\d.]+)nm\.png")

# ground-truth geometric diameter per (n,m,df)
dcal = {}
with open(SIM / "summary.csv", newline="") as f:
    for row in csv.DictReader(f):
        dcal[(int(row["n"]), int(row["m"]), float(row["defocus_nm"]))] = float(row["diameter_nm"])

rows = []
for df_tag, df_val in [("+2.0", 2.0), ("-2.0", -2.0)]:
    for p in sorted(SIM.glob(f"tem_image_*_df{df_tag}nm.png")):
        m = name_re.match(p.name)
        n, mm = int(m.group(1)), int(m.group(2))
        img, _ = h.load_image_data(str(p))
        if img is None:
            continue
        roi = h.auto_roi(img)
        rng = img.max() - img.min()
        if rng < 1e-10:
            continue
        proj = np.sum(((img - img.min()) / rng)[roi[0]:roi[1], roi[2]:roi[3]], axis=0)
        dm = h.measure_diameter_all_modes(proj, SCALE)
        if dm is None:
            continue
        rows.append({"n": n, "m": mm, "defocus_nm": df_val,
                     "d_cal": dcal.get((n, mm, df_val), np.nan), **{k: dm[k] for k in MODES}})
    print(f"df={df_val:+.0f} nm: measured {sum(1 for r in rows if r['defocus_nm']==df_val)} tubes")

with open(OUT / "diameter_offsets_4modes_df_pm2.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["n", "m", "defocus_nm", "d_cal"] + MODES)
    w.writeheader(); w.writerows(rows)

print("\nMean offset (measured - d_cal), nm:")
print(f"{'focus':>10} " + " ".join(f"{k:>8}" for k in MODES))
for df_val, label in [(2.0, "under+2"), (-2.0, "over -2")]:
    sub = [r for r in rows if r["defocus_nm"] == df_val and np.isfinite(r["d_cal"])]
    offs = {k: np.mean([r[k] - r["d_cal"] for r in sub]) for k in MODES}
    print(f"{label:>10} " + " ".join(f"{offs[k]:+8.3f}" for k in MODES))

for df_val, tag in [(2.0, "a_df+2nm_underfocus"), (-2.0, "b_df-2nm_overfocus")]:
    sub = [r for r in rows if r["defocus_nm"] == df_val and np.isfinite(r["d_cal"])]
    xc = np.array([r["d_cal"] for r in sub])
    fig, ax = plt.subplots(figsize=(5.2, 5.0))
    lim = [0, max(3.6, xc.max() + 0.2)]
    ax.plot(lim, lim, "k--", lw=1, alpha=0.6, label="y = x")
    for k in MODES:
        ax.scatter(xc, [r[k] for r in sub], s=14, c=COL[k], alpha=0.7, label=k, edgecolors="none")
    ax.set_xlim(lim); ax.set_ylim(lim); ax.set_aspect("equal")
    ax.set_xlabel("Geometric diameter  d$_{cal}$  (nm)")
    ax.set_ylabel("Measured diameter (nm)")
    ax.set_title(f"Fig 3{tag[0]} - diameter modes, df = {df_val:+.0f} nm")
    ax.legend(frameon=False, fontsize=9, loc="upper left")
    fig.tight_layout(); fig.savefig(OUT / f"fig3_{tag}.png", dpi=200); plt.close(fig)
print(f"\nWrote CSV + 2 plots to {OUT}")
