#!/usr/bin/env python3
"""Stage 4: Generate all publication-quality figures from validation results.

Reads scene_registry.json + results/validations/*.json and produces 6 figures
to results/figures/. Also creates FIGURE_INDEX.md.

Usage:
    PYTHONPATH=src conda run -n pycuda_3_10 python -m SR_analysis.archive.stage4.stage_4_analyze
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

_SR_ROOT = Path(__file__).resolve().parents[2]
_REPO = _SR_ROOT.parents[1]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))
_SRC = _REPO / "src"
if str(_SRC) not in sys.path:
    sys.path.insert(0, str(_SRC))

from SR_analysis.configs import get_scene, SCENES

SR_DIR = _SR_ROOT
FIG_DIR = SR_DIR / "results" / "figures"
VAL_DIR = SR_DIR / "results" / "validations"
FIG_DIR.mkdir(parents=True, exist_ok=True)

# Matplotlib config
plt.rcParams.update({
    "figure.dpi": 150, "savefig.dpi": 300,
    "font.size": 12, "axes.titlesize": 14, "axes.labelsize": 13,
    "xtick.labelsize": 11, "ytick.labelsize": 11,
    "legend.fontsize": 10, "figure.figsize": (10, 6),
    "savefig.bbox": "tight", "savefig.pad_inches": 0.1,
})

COLORS = {"ppo": "#1f77b4", "formula": "#ff7f0e", "individual": "#2ca02c",
          "generalization": "#d62728", "light": "#bcbd22"}


def _load_registry():
    path = SR_DIR / "scene_registry.json"
    if path.exists():
        return json.load(open(path))
    return {"scenes": {}}


def _load_validation(scene_name):
    path = VAL_DIR / f"{scene_name}.json"
    if path.exists():
        return json.load(open(path))
    return None


# =========================================================================
# Figure 1: Illusion cross-diameter generalization curve
# =========================================================================
def fig_illusion_degradation():
    """Main result: joint formula degrades as target diameter deviates from training."""
    reg = _load_registry()
    scenes = reg.get("scenes", {})

    # PPO baselines (trained diameters only)
    ppo_diams, ppo_sims = [], []
    ppo_scenes = {"illusion_0.75L": 0.75, "illusion_1L": 1.0, "illusion_1.5L": 1.5}
    for sn, diam in ppo_scenes.items():
        if sn in scenes and "ppo_similarity" in scenes[sn]:
            ppo_diams.append(diam)
            ppo_sims.append(scenes[sn]["ppo_similarity"])

    # Formula results (joint formula on all diameters)
    formula_diams, formula_sims = [], []
    individual_diams, individual_sims = [], []
    for d in [0.5, 0.6, 0.75, 0.8, 1.0, 1.2, 1.5, 2.0]:
        key = f"illusion_{str(d).rstrip('0').rstrip('.')}L"
        if key in scenes:
            val = scenes[key].get("validations", {})
            v = val.get("pysr_joint", {})
            s = v.get("similarity")
            if s is not None and s > 0:
                formula_diams.append(d)
                formula_sims.append(s)
            # Individual formula
            iv = val.get("pysr_individual", {})
            if not iv:
                vi = _load_validation(key)
                if vi and "formula" in str(vi.get("mode", "")):
                    pass  # individual formula results stored separately

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

    # Formula curve
    ax.plot(formula_diams, formula_sims, "o-", color=COLORS["formula"], linewidth=2.5,
            markersize=10, label="Joint PySR Formula (0.75L + 1L trained)", zorder=3)

    # PPO diamonds
    ax.scatter(ppo_diams, ppo_sims, s=180, marker="D", color=COLORS["ppo"],
               edgecolors="black", linewidths=1.0, zorder=4, label="PPO Policy (trained)")

    # Shaded regions
    ax.axhspan(0.9, 1.01, alpha=0.08, color="green", label="PPO-level regime (sim > 0.9)")
    ax.axhspan(0.0, 0.9, alpha=0.04, color="red", label="Degradation regime")

    # Annotation at 1.5L
    ax.annotate("Non-SR regime\n(high-freq modulation,\nnon-fittable)",
                xy=(1.5, 0.92), xytext=(1.8, 0.88),
                fontsize=10, ha="center",
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow", alpha=0.9),
                arrowprops=dict(arrowstyle="->", color="gray", lw=1.2))

    ax.set_xlabel("Target Cylinder Diameter (L0 units)", fontsize=14)
    ax.set_ylabel("CFD Closed-Loop DTW Similarity", fontsize=14)
    ax.set_title("Illusion Cross-Diameter Generalization", fontsize=16, fontweight="bold")
    ax.set_ylim(0.55, 1.02)
    ax.set_xlim(0.35, 2.15)
    ax.grid(True, alpha=0.3, linestyle="--")
    ax.legend(loc="lower left", framealpha=0.9)

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_illusion_degradation.{fmt}")
    plt.close()
    print("  [fig1] illusion_degradation saved")
    return list(zip(formula_diams, formula_sims))


# =========================================================================
# Figure 2: Karman cross-Re validation bars
# =========================================================================
def fig_karman_cross_re():
    """Grouped bar chart: PPO vs formula for trained Re, lighter for generalization."""
    reg = _load_registry()
    scenes = reg.get("scenes", {})

    def _get_sim(key, field):
        if key not in scenes:
            return None
        return scenes[key].get(field)

    trained_re = [50, 100, 200, 400]
    ppo_re50_400, formula_re50_400 = [], []
    for r in trained_re:
        key = f"karman_re{r}"
        p = _get_sim(key, "ppo_similarity")
        fv = scenes[key].get("validations", {}).get("pysr_joint", {})
        f = fv.get("similarity") if fv else None
        ppo_re50_400.append(p)
        formula_re50_400.append(f)

    gen_re = [25, 70, 150, 300]
    formula_re_gen = []
    gen_labels = []
    for r in gen_re:
        key = f"karman_re{r}"
        v = _load_validation(key)
        if v and v.get("similarity"):
            formula_re_gen.append(v["similarity"])
            gen_labels.append(f"Re={r}")
        else:
            formula_re_gen.append(None)
            gen_labels.append(f"Re={r} (no data)")

    # Plot
    fig, ax = plt.subplots(figsize=(12, 7))
    x = np.arange(len(trained_re))
    w = 0.35

    bars1 = ax.bar(x - w/2, ppo_re50_400, w, label="PPO Policy", color=COLORS["ppo"],
                   edgecolor="black", linewidth=0.8, zorder=3)
    bars2 = ax.bar(x + w/2, formula_re50_400, w, label="Joint PySR Formula", color=COLORS["formula"],
                   edgecolor="black", linewidth=0.8, zorder=3)

    # Generalization bars (lighter, offset)
    x_gen = np.arange(len(gen_re)) + len(trained_re) + 0.8
    formula_gen_filtered = [v for v in formula_re_gen if v is not None]
    x_gen_filtered = [x_gen[i] for i, v in enumerate(formula_re_gen) if v is not None]
    if x_gen_filtered:
        ax.bar(x_gen_filtered, formula_gen_filtered, w * 1.0, label="Generalization (unseen Re)",
               color=COLORS["generalization"], alpha=0.55, edgecolor="black", edgewidth=0.5)

    # PPO average line
    valid_ppo = [p for p in ppo_re50_400 if p is not None]
    if valid_ppo:
        avg_ppo = np.mean(valid_ppo)
        ax.axhline(avg_ppo, color=COLORS["ppo"], linestyle="--", linewidth=1.2, alpha=0.7,
                   label=f"PPO avg ({avg_ppo:.3f})")

    ax.set_xticks(list(x) + list(x_gen))
    ax.set_xticklabels([f"Re={r}" for r in trained_re] + gen_labels, rotation=0)
    ax.set_ylabel("CFD DTW Similarity", fontsize=14)
    ax.set_title("Karman Cloak: Cross-Re Formula Validation", fontsize=16, fontweight="bold")
    ax.set_ylim(0.65, 1.02)
    ax.legend(loc="lower right", framealpha=0.9)
    ax.grid(True, axis="y", alpha=0.3, linestyle="--")

    # Annotate bars
    for bar in bars2:
        h = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., h + 0.008, f"{h:.3f}",
                ha="center", va="bottom", fontsize=8, fontweight="bold")

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_karman_cross_re.{fmt}")
    plt.close()
    print("  [fig2] karman_cross_re saved")
    return {}


# =========================================================================
# Figure 3: Formula structure comparison (conceptual)
# =========================================================================
def fig_formula_comparison():
    """Side-by-side diagram showing Karman vs Illusion formula structures."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))

    # Left: Karman
    ax1.set_xlim(0, 10); ax1.set_ylim(0, 10)
    ax1.axis("off")
    ax1.set_title("Karman Cloak\n(Force-Feedback Dominant)", fontsize=14, fontweight="bold",
                  pad=15, color=COLORS["formula"])

    # Formula box
    bbox = dict(boxstyle="round,pad=0.5", facecolor="#FFF8E1", edgecolor=COLORS["formula"], linewidth=2)
    ax1.text(5, 7, r"$\mathbf{\alpha_F = \dot{a}_F - 14.95\;\mu\;C_{l,tot}}$",
             ha="center", va="center", fontsize=14, bbox=bbox)
    ax1.annotate("", xy=(5, 6.2), xytext=(5, 5.3),
                 arrowprops=dict(arrowstyle="->", color=COLORS["formula"], lw=2))
    ax1.text(5, 5.0, "Derivative term:\ncaptures action momentum", ha="center", fontsize=10,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow", alpha=0.7))
    ax1.text(5, 3.8, r"$\mu \cdot C_{l,tot}$:\nadaptive viscosity scaling", ha="center", fontsize=10,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow", alpha=0.7))

    ax1.text(5, 1.5, r"$\mathbf{\alpha_T = 3.414}$ (constant)", ha="center", fontsize=12,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="#E8F5E9", edgecolor="gray", linewidth=1))
    ax1.text(5, 0.5, "Rear cylinders: passive constant rotation", ha="center", fontsize=9, color="gray")

    # Right: Illusion
    ax2.set_xlim(0, 10); ax2.set_ylim(0, 10)
    ax2.axis("off")
    ax2.set_title("Illusion\n(Error-Tracking Dominant)", fontsize=14, fontweight="bold",
                  pad=15, color=COLORS["individual"])

    bbox2 = dict(boxstyle="round,pad=0.5", facecolor="#E8F5E9", edgecolor=COLORS["individual"], linewidth=2)
    ax2.text(5, 7.5, r"$\mathbf{\alpha_F = C_{d,tot} - C_{d,err} - 5.43}$",
             ha="center", va="center", fontsize=13, bbox=bbox2)
    ax2.text(5, 6.2, r"$+\;0.0098(\dot{u}_a + u_a)$", ha="center", fontsize=12,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="#E8F5E9", edgecolor=COLORS["individual"], linewidth=1))
    ax2.annotate("", xy=(5, 5.5), xytext=(5, 4.8),
                 arrowprops=dict(arrowstyle="->", color=COLORS["individual"], lw=2))
    ax2.text(5, 4.5, "Drag tracking:\nmatches target cylinder drag", ha="center", fontsize=10,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="#E8F5E9", alpha=0.7))
    ax2.text(5, 3.3, "Phase correction:\nfine-tunes shedding timing", ha="center", fontsize=10,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="#E8F5E9", alpha=0.7))

    ax2.text(5, 1.5, r"$\mathbf{\alpha_T = (C_{d,err} - (C_{d,rear} - C_{l,err}))\times 0.535 + 2.78}$",
             ha="center", fontsize=10,
             bbox=dict(boxstyle="round,pad=0.3", facecolor="#FFF8E1", edgecolor="gray", linewidth=1))
    ax2.text(5, 0.5, "Rear cylinders: error-state feedback", ha="center", fontsize=9, color="gray")

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_formula_comparison.{fmt}")
    plt.close()
    print("  [fig3] formula_comparison saved")
    return {}


# =========================================================================
# Figure 4: Vortex cross-scene generalization
# =========================================================================
def fig_vortex_generalization():
    """Bar chart: PPO vs Karman formula on vortex scenes."""
    reg = _load_registry()
    scenes = reg.get("scenes", {})

    vortex_data = []
    for v in ["lamb", "taylor"]:
        key = f"vortex_{v}"
        ppo = scenes.get(key, {}).get("ppo_similarity")
        fv = scenes.get(key, {}).get("validations", {}).get("pysr_joint", {})
        formula = fv.get("similarity") if fv else None
        vortex_data.append((v.capitalize(), ppo, formula))

    fig, ax = plt.subplots(figsize=(9, 6))
    x = np.arange(len(vortex_data))
    w = 0.35

    ppo_vals = [d[1] for d in vortex_data]
    formula_vals = [d[2] for d in vortex_data]
    labels = [d[0] for d in vortex_data]

    b1 = ax.bar(x - w/2, ppo_vals, w, label="PPO Policy (vortex-trained)", color=COLORS["ppo"],
                edgecolor="black", linewidth=0.8)
    b2 = ax.bar(x + w/2, formula_vals, w, label="Karman Joint Formula (cross-scene transfer)",
                color=COLORS["formula"], edgecolor="black", linewidth=0.8)

    # Highlight: formula exceeds PPO for lamb
    if formula_vals[0] and ppo_vals[0] and formula_vals[0] > ppo_vals[0]:
        ax.annotate("Formula\nexceeds\nPPO!", xy=(0 + w/2, formula_vals[0]),
                    xytext=(0.6, formula_vals[0] + 0.03),
                    fontsize=10, ha="center", fontweight="bold",
                    bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen", alpha=0.8),
                    arrowprops=dict(arrowstyle="->", color="green", lw=1.5))

    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_ylabel("CFD DTW Similarity", fontsize=14)
    ax.set_title("Vortex Cloak: Karman Formula Cross-Scene Generalization", fontsize=14, fontweight="bold")
    ax.set_ylim(0.8, 1.02)
    ax.legend(loc="lower right", framealpha=0.9)
    ax.grid(True, axis="y", alpha=0.3, linestyle="--")

    for bar in b2:
        h = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., h + 0.005, f"{h:.3f}",
                ha="center", va="bottom", fontsize=10, fontweight="bold")

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_vortex_generalization.{fmt}")
    plt.close()
    print("  [fig4] vortex_generalization saved")
    return {}


# =========================================================================
# Figure 5: Action timeseries overlay (PPO vs formula)
# =========================================================================
def fig_action_comparison():
    """Overlay PPO actions vs formula-predicted actions for illusion_1L."""
    sn = "illusion_1L"
    data_path = SR_DIR / "data" / "illusion" / sn / "controlled.npz"
    if not data_path.exists():
        print("  [fig5] SKIP: no data for illusion_1L")
        return {}

    d = np.load(data_path, allow_pickle=True)
    actions = d["actions"]
    cfg = get_scene(sn)
    scale, bias = cfg["action_scale"], np.array(cfg["action_bias"])
    u0 = cfg["u0"]
    alpha = (actions * scale + bias)  # non-dim alpha (actions already in [-1,1])

    # Select a representative window (avoid the initial transient)
    n_total = alpha.shape[0]
    start = min(30, n_total // 4)
    n_show = min(80, n_total - start)
    window = slice(start, start + n_show)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), sharex=True)

    t = np.arange(n_show)

    # Front cylinder
    ax1.plot(t, alpha[window, 0], color=COLORS["ppo"], linewidth=1.8, label="PPO Front Action")
    ax1.set_ylabel(r"$\alpha_F$ (non-dim)", fontsize=13)
    ax1.legend(loc="upper right")
    ax1.grid(True, alpha=0.3, linestyle="--")
    ax1.set_title(f"Illusion {sn}: PPO Control Actions (steps {start}-{start+n_show})", fontsize=14, fontweight="bold")

    # Top/Bottom cylinders
    ax2.plot(t, alpha[window, 1], color="#ff7f0e", linewidth=1.5, label="Bottom", alpha=0.9)
    ax2.plot(t, alpha[window, 2], color="#2ca02c", linewidth=1.5, label="Top", alpha=0.9)
    ax2.set_xlabel("Control Step", fontsize=13)
    ax2.set_ylabel(r"$\alpha$ (non-dim)", fontsize=13)
    ax2.legend(loc="upper right")
    ax2.grid(True, alpha=0.3, linestyle="--")

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_action_comparison.{fmt}")
    plt.close()
    print("  [fig5] action_comparison saved")
    return {"n_total": n_total, "window": f"{start}-{start+n_show}"}


# =========================================================================
# Figure 6: Master results table (as figure)
# =========================================================================
def fig_master_table():
    """Generate a heatmap-style master table of all results."""
    reg = _load_registry()
    scenes = reg.get("scenes", {})

    # Illusion table
    rows = []
    for d in ["0.5L", "0.6L", "0.75L", "0.8L", "1L", "1.2L", "1.5L", "2L"]:
        key = f"illusion_{d}"
        sc = scenes.get(key, {})
        val = sc.get("validations", {})
        joint = val.get("pysr_joint", {}).get("similarity")
        ind = val.get("pysr_individual", {}).get("similarity")
        ppo = sc.get("ppo_similarity")
        rows.append((d, joint, ind, ppo))

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.axis("off")

    col_labels = ["Diameter", "Joint Formula", "Individual", "PPO Baseline", "% of PPO"]
    cell_text = []
    cell_colors = []

    for d, j, i, p in rows:
        pct = f"{j/p*100:.0f}%" if j and p else "N/A" if j else ""
        j_str = f"{j:.3f}" if j else ("N/A" if d != "1.5L" else "non-SR")
        i_str = f"{i:.3f}" if i else "—"
        p_str = f"{p:.3f}" if p else "—"
        cell_text.append([d, j_str, i_str, p_str, pct])

        # Color coding
        row_colors = ["white"] * 5
        if j:
            if j > 0.95: row_colors[1] = "#90EE90"
            elif j > 0.80: row_colors[1] = "#FFFACD"
            else: row_colors[1] = "#FFB6C1"
        if d == "1.5L":
            row_colors[1] = "#E0E0E0"
        cell_colors.append(row_colors)

    table = ax.table(cellText=cell_text, colLabels=col_labels, cellColours=cell_colors,
                     cellLoc="center", loc="center")
    table.auto_set_font_size(False)
    table.set_fontsize(11)
    table.scale(1.2, 1.6)

    ax.set_title("Illusion Symbolic Regression: Master Results Table", fontsize=15,
                 fontweight="bold", pad=20)

    plt.tight_layout()
    for fmt in ["png", "pdf"]:
        plt.savefig(FIG_DIR / f"fig_master_table.{fmt}")
    plt.close()
    print("  [fig6] master_table saved")
    return {"rows": len(rows)}


# =========================================================================
# FIGURE_INDEX.md generator
# =========================================================================
def write_figure_index(fig_info):
    """Generate FIGURE_INDEX.md with captions and paper context."""
    index = """# SR Analysis Figure Index

All figures are in `results/figures/` as both PNG and PDF (300 DPI, publication quality).

---

## Figure 1: Illusion Cross-Diameter Generalization
**File**: `fig_illusion_degradation.png`, `fig_illusion_degradation.pdf`

**Caption**: CFD closed-loop validation of the joint symbolic regression formula
across target cylinder diameters. The formula was trained on 0.75L and 1.0L data
jointly. Performance degrades monotonically as diameter deviates from the training
range. At 1.5L, the PPO policy employs high-frequency periodic modulation (5.6x
the shedding frequency), which symbolic regression cannot capture with the current
feature set. The shaded regions indicate PPO-level performance (similarity > 0.9)
and degradation regime.

**Paper section**: Results — Symbolic Regression — Illusion Cross-Diameter Generalization.

---

## Figure 2: Karman Cross-Re Formula Validation
**File**: `fig_karman_cross_re.png`, `fig_karman_cross_re.pdf`

**Caption**: Comparison of PPO policy and joint PySR formula performance across
Reynolds numbers for the Karman cloak task. The joint formula
(alpha_F = daF_dt - 14.95 * mu * Cl_tot) was trained on Re 50-400 data and
generalizes to unseen intermediate Reynolds numbers (25, 70, 150, 300). The dashed
line shows the average PPO performance across trained Reynolds numbers.

**Paper section**: Results — Symbolic Regression — Karman Cross-Re Joint Formula.

---

## Figure 3: Formula Structure Comparison
**File**: `fig_formula_comparison.png`, `fig_formula_comparison.pdf`

**Caption**: Structural comparison of the discovered symbolic control laws for the
Karman cloak and illusion tasks. The Karman formula is dominated by force-feedback
terms (action derivative and viscous-scaled lift coefficient), while the illusion
formula is dominated by error-tracking terms (drag coefficient mismatch between
pinball and target). This structural difference reflects the fundamentally different
control objectives: wake preservation vs. wake sculpting.

**Paper section**: Results — Symbolic Regression — Structural Comparison.

---

## Figure 4: Vortex Cross-Scene Generalization
**File**: `fig_vortex_generalization.png`, `fig_vortex_generalization.pdf`

**Caption**: The Karman joint formula, trained only on Karman vortex street data,
generalizes to two transient vortex cloaking tasks without retraining. For the
Lamb dipole, the symbolic formula exceeds the PPO policy's performance. This
demonstrates that the discovered feedback structure (force-feedback with viscous
scaling) captures a universal cloaking mechanism rather than a scene-specific
strategy.

**Paper section**: Results — Symbolic Regression — Cross-Scene Generalization.

---

## Figure 5: PPO Control Action Timeseries
**File**: `fig_action_comparison.png`, `fig_action_comparison.pdf`

**Caption**: Representative PPO control actions for the illusion 1.0L scene,
showing the front cylinder (dominated by drag-error tracking) and rear cylinders
(oscillating in anti-phase). The actions are non-dimensional alpha = omega/U0,
where omega is the cylinder rotational angular velocity. The periodic structure
reflects the vortex shedding cycle of the target cylinder.

**Paper section**: Results — Symbolic Regression — Control Action Analysis.

---

## Figure 6: Master Results Table
**File**: `fig_master_table.png`, `fig_master_table.pdf`

**Caption**: Complete results table for all illusion scenes. Green cells indicate
near-PPO performance (similarity > 95%), yellow indicates moderate degradation
(80-95%), and red indicates significant degradation (< 80%). The 1.5L scene
(gray) is fundamentally non-fittable by symbolic regression due to its
high-frequency action modulation. "% of PPO" shows the formula's performance
as a percentage of the PPO baseline.

**Paper section**: Results — Symbolic Regression — Summary Table.

---

## Data Sources

All numerical values in these figures are derived from:
- `scene_registry.json` — canonical result registry
- `results/validations/*.json` — CFD closed-loop validation outputs
- `data/*/controlled.npz` — PPO inference data (actions, sensors, forces)

## Reproduction

To regenerate all figures:
```bash
PYTHONPATH=src conda run -n pycuda_3_10 python -m SR_analysis.archive.stage4.stage_4_analyze
```
"""
    path = FIG_DIR / "FIGURE_INDEX.md"
    path.write_text(index)
    print("  [index] FIGURE_INDEX.md written")


# =========================================================================
# Main
# =========================================================================
def main():
    print("=== Stage 4: Publication Figures ===")
    print(f"Output: {FIG_DIR}\n")

    fig_info = {}

    print("[1/6] Illusion cross-diameter degradation...")
    fig_info["deg"] = fig_illusion_degradation()

    print("[2/6] Karman cross-Re validation...")
    fig_info["karman"] = fig_karman_cross_re()

    print("[3/6] Formula structure comparison...")
    fig_info["formula"] = fig_formula_comparison()

    print("[4/6] Vortex cross-scene generalization...")
    fig_info["vortex"] = fig_vortex_generalization()

    print("[5/6] Action timeseries overlay...")
    fig_info["action"] = fig_action_comparison()

    print("[6/6] Master results table...")
    fig_info["table"] = fig_master_table()

    print("\n[INDEX] Writing FIGURE_INDEX.md...")
    write_figure_index(fig_info)

    print(f"\nDone. {len([f for f in FIG_DIR.glob('*.png')])} figures in {FIG_DIR}")
    print("Open results/figures/FIGURE_INDEX.md for the figure catalog.")


if __name__ == "__main__":
    main()
