#!/usr/bin/env python3
"""
Generate SR analysis result charts for reporting.
Output: docs/figures/SR_analysis/  (PNG files)
"""
import json, os, sys, glob
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
_OUT = os.path.join(_REPO, "docs", "figures", "SR_analysis")
os.makedirs(_OUT, exist_ok=True)

# ---------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------

# Illusion three-scenario comparison (new SR route: phase-state + error-state + absolute)
illusion_results = {
    "0.75L": {"old_v23": 0.908, "new_phase": 0.974, "ppo": 0.972, "model": "d1a3o14_250525_imit_075L_2U_400S", "S": 400},
    "1L":    {"old_v23": 0.962, "new_phase": 0.958, "ppo": 0.973, "model": "d1a3o14_250525_imit_1L_2U_600S", "S": 600},
    "1.5L":  {"old_v23": 0.926, "new_phase": "N/A", "ppo": 0.945, "model": "d1a3o14_250525_imit_15L_2U", "S": 800},
}

# Karman ablation
karman_ablation = {
    "old v23\n(a_lag+da)": 0.901,
    "static->deriv\n(8dim)": 0.745,
    "full-lag->deriv\n(16dim)": 0.619,
    "phase->deriv\n(6dim)": 0.656,
    "phase->abs\n(6dim)": 0.699,
    "phase+mu->abs\n(9dim)": 0.700,
    "expanded->abs\n(10dim)": 0.580,
}

# Karman generalization
karman_gen = {
    "Re25": 0.567, "Re50": 0.582, "Re70": 0.577,
    "Re100": 0.901, "Re150": 0.595, "Re200": 0.793,
    "Re300": 0.541, "Re400": 0.664
}

# Ablation one-step R2
ablation_r2 = {
    "static\n0": 0.321,
    "phase\n6": 0.837,
    "phase+abs\n6": 0.965,
    "full-lag\n16": 0.939,
    "expanded\n10": 0.980,
    "phase+mu\n9": 0.979,
}

# ---------------------------------------------------------------------------
# Color scheme
# ---------------------------------------------------------------------------
C_OLD = "#d62728"    # red
C_NEW_PHASE = "#2ca02c"  # green 
C_PPO = "#1f77b4"    # blue
C_DERIV = "#ff7f0e"  # orange
C_ABS = "#2ca02c"    # green
C_GEN = "#9467bd"    # purple
C_BG = "#f0f0f0"

# ---------------------------------------------------------------------------
# Fig 1: Illusion comparison bar chart
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(10, 5))
labels = list(illusion_results.keys())
x = np.arange(len(labels))
w = 0.25

old_vals = [illusion_results[k]["old_v23"] for k in labels]
new_vals = [illusion_results[k]["new_phase"] for k in labels]
ppo_vals = [illusion_results[k]["ppo"] for k in labels]

# Convert N/A to NaN for plotting
new_vals_plot = [v if isinstance(v, (int, float)) else 0 for v in new_vals]

ax.bar(x - w, old_vals, w, label="Old v23 (动作历史)", color=C_OLD, alpha=0.8)
ax.bar(x, new_vals_plot, w, label="New phase-state (无动作历史)", color=C_NEW_PHASE, alpha=0.8)
ax.bar(x + w, ppo_vals, w, label="PPO 基线", color=C_PPO, alpha=0.6)

# 1.5L label
ax.text(x[2], 0.05, "N/A\n(bang-bang)", ha="center", va="bottom", fontsize=9, color="gray")

ax.set_xticks(x)
ax.set_xticklabels(labels)
ax.set_ylabel("闭环相似度 (DTW)")
ax.set_title("Figure 1: Illusion 三场景 — 新路线 vs 旧版 vs PPO 基线", fontsize=12)
ax.legend(fontsize=9)
ax.set_ylim(0, 1.05)
ax.grid(axis="y", alpha=0.3)
fig.tight_layout()
fig.savefig(os.path.join(_OUT, "fig1_illusion_comparison.png"), dpi=150)
print("Saved fig1_illusion_comparison.png")
plt.close(fig)

# ---------------------------------------------------------------------------
# Fig 2: Karman ablation bar chart
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(12, 4.5))
keys = list(karman_ablation.keys())
vals = list(karman_ablation.values())
colors = []
for k in keys:
    if "old" in k.lower(): colors.append(C_OLD)
    elif "abs" in k: colors.append(C_ABS)
    elif "deriv" in k: colors.append(C_DERIV)
    else: colors.append("#7f7f7f")

bars = ax.bar(range(len(keys)), vals, color=colors, alpha=0.85)
ax.axhline(y=0.901, color=C_OLD, linestyle="--", alpha=0.5, label="Old v23 baseline (0.901)")
ax.set_xticks(range(len(keys)))
ax.set_xticklabels(keys, fontsize=8, rotation=20, ha="right")
ax.set_ylabel("闭环相似度")
ax.set_title("Figure 2: Karman re100 消融实验 — 输入/输出形式对比", fontsize=12)
ax.legend(fontsize=9)
ax.set_ylim(0, 1.0)
ax.grid(axis="y", alpha=0.3)

# Add red dashed line at 0.699 highlighting best new route
ax.axhline(y=0.699, color=C_ABS, linestyle=":", alpha=0.5)
ax.text(5.5, 0.705, "phase+abs 最佳: 0.699", fontsize=8, color=C_ABS)

fig.tight_layout()
fig.savefig(os.path.join(_OUT, "fig2_karman_ablation.png"), dpi=150)
print("Saved fig2_karman_ablation.png")
plt.close(fig)

# ---------------------------------------------------------------------------
# Fig 3: Karman generalization across Re
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(10, 4.5))
re_vals = sorted(karman_gen.keys(), key=lambda s: int(s.replace("Re","")))
sims = [karman_gen[k] for k in re_vals]

ax.plot(range(len(re_vals)), sims, "o-", color=C_GEN, linewidth=2, markersize=8)
# Mark training Re
train_re = [0, 3, 5, 7]  # indices of Re50/100/200/400
for i in train_re:
    ax.plot(i, sims[i], "o", color=C_OLD, markersize=12, markeredgecolor="black", markeredgewidth=1.5)
ax.set_xticks(range(len(re_vals)))
ax.set_xticklabels(re_vals)
ax.set_ylabel("闭环相似度")
ax.set_title("Figure 3: Karman 跨 Re 泛化 (旧 v23 模型)", fontsize=12)
ax.set_xlabel("Reynolds Number (2D reference)")
ax.axhline(y=0.5, color="gray", linestyle=":", alpha=0.5)
ax.grid(axis="y", alpha=0.3)
ax.set_ylim(0, 1.0)

# Annotations
ax.annotate("训练 Re", xy=(1.8, 0.92), fontsize=9, color=C_OLD)
ax.annotate("泛化 Re\n(未见过的)", xy=(4.5, 0.55), fontsize=9, color=C_GEN)

fig.tight_layout()
fig.savefig(os.path.join(_OUT, "fig3_karman_generalization.png"), dpi=150)
print("Saved fig3_karman_generalization.png")
plt.close(fig)

# ---------------------------------------------------------------------------
# Fig 4: One-step R2 vs closed-loop scatter (diagnostic)
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(7, 5.5))

points = [
    ("old v23", 0.996, 0.901, C_OLD),
    ("static→deriv", 0.321, 0.745, C_DERIV),
    ("full-lag→deriv", 0.939, 0.619, "#7f7f7f"),
    ("phase→deriv", 0.837, 0.656, C_DERIV),
    ("phase→abs", 0.965, 0.699, C_ABS),
    ("expanded→abs", 0.980, 0.580, "#7f7f7f"),
    ("phase+mu→abs", 0.979, 0.700, C_ABS),
]
for name, r2, sim, color in points:
    ax.scatter(r2, sim, s=100, color=color, zorder=5)
    ax.annotate(name, (r2, sim), textcoords="offset points", xytext=(5, 5), fontsize=8)

ax.set_xlabel("One-step R²")
ax.set_ylabel("CFD 闭环相似度")
ax.set_title("Figure 4: Karman re100 — One-step R² 与闭环不一致性", fontsize=12)
ax.grid(alpha=0.3)

# Upper-left region = good closed-loop, bad one-step (static-deriv)
# Upper-right region = good both (old v23)
# Lower-right region = good one-step, bad closed-loop (most new methods)
ax.annotate("稳健欠拟合", xy=(0.15, 0.85), fontsize=9, color="gray", fontstyle="italic")
ax.annotate("分布偏移\n(训练好, 闭环差)", xy=(0.75, 0.45), fontsize=9, color="gray", fontstyle="italic")

fig.tight_layout()
fig.savefig(os.path.join(_OUT, "fig4_r2_vs_closedloop.png"), dpi=150)
print("Saved fig4_r2_vs_closedloop.png")
plt.close(fig)

# ---------------------------------------------------------------------------
# Fig 5: Phase-state feature coefficients (Illusion 1L)
# ---------------------------------------------------------------------------
# Load the SINDy results for illusion 1L phase-state
try:
    with open(os.path.join(_REPO, "src/SR_analysis/sindy/illusion/sindy_results_deriv.json")) as f:
        sr = json.load(f)
    per = sr["per_scene"]["illusion_1L"]
    fn_f = per["feature_names_front"]
    coef_f = per["front"]["best_coef"][:len(fn_f)]
    
    fig, ax = plt.subplots(figsize=(8, 4.5))
    # Sort by |coef|
    pairs = sorted(zip(fn_f, coef_f), key=lambda p: -abs(p[1]))
    names = [p[0] for p in pairs]
    vals = [p[1] for p in pairs]
    colors_bar = [C_NEW_PHASE if v > 0 else C_OLD for v in vals]
    ax.barh(range(len(names)), vals, color=colors_bar, alpha=0.8)
    ax.set_yticks(range(len(names)))
    ax.set_yticklabels(names)
    ax.axvline(x=0, color="black", linewidth=0.5)
    ax.set_xlabel("系数值")
    ax.set_title("Figure 5: Illusion 1L Front — Phase-state 特征系数", fontsize=12)
    ax.grid(axis="x", alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(_OUT, "fig5_illusion_coefficients.png"), dpi=150)
    print("Saved fig5_illusion_coefficients.png")
    plt.close(fig)
except Exception as e:
    print(f"fig5 skipped: {e}")

# ---------------------------------------------------------------------------
# Fig 6: Summary timeline / roadmap
# ---------------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(10, 3.5))
phases = [
    ("Phase 0\nBug Audit", "2026-06-12\n12 bugs\nfixed", 0.8),
    ("Phase 1\nTarget fix", "2026-06-13\nIllusion target\ninfo added", 0.85),
    ("Phase 2\nAblation", "2026-06-14\nPhase-state\nvalidated", 0.90),
    ("Phase 2b\nOutput mode", "2026-06-15\nPhase+abs\nIllusion 0.97", 0.95),
    ("Phase 3\nIllusion SR", "Next\nPySR on\n0.75L/1L", 0.7),
]
y_pos = 1
for i, (label, desc, conf) in enumerate(phases):
    color = plt.cm.RdYlGn(conf)
    ax.barh(y_pos, 1, left=i, height=0.5, color=color, alpha=0.8)
    ax.text(i + 0.5, y_pos, label, ha="center", va="center", fontsize=8, fontweight="bold")
    ax.text(i + 0.5, y_pos - 0.3, desc, ha="center", va="top", fontsize=6, color="gray")
ax.set_xlim(0, len(phases))
ax.set_ylim(0, 2)
ax.axis("off")
ax.set_title("Figure 6: 研究进展路线图", fontsize=12)
fig.tight_layout()
fig.savefig(os.path.join(_OUT, "fig6_roadmap.png"), dpi=150)
print("Saved fig6_roadmap.png")
plt.close(fig)

print(f"\nAll figures saved to {_OUT}/")