"""
⚡ ULTRON Systems // 100% Deterministic Empirical Figure Generator (300 DPI)
Zero synthetic data. Reads directly from:
- logs/pretraining_telemetry_verified.json (855 real logged pre-training steps)
- logs/defensible_tier1_audit.json (Stratified 15-ticker contiguous OOS & Pre-2024 Stress audit,
  exact global cosine similarities, 1m return-magnitude buckets, paired bootstrap CIs,
  mean-ablation + random-channel controls, and non-overlapping state-turnover backtests)
"""
import os
import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

os.makedirs("papers/figures", exist_ok=True)

plt.style.use('seaborn-v0_8-whitegrid' if 'seaborn-v0_8-whitegrid' in plt.style.available else 'default')
mpl.rcParams['font.family'] = 'serif'
mpl.rcParams['font.size'] = 9.5
mpl.rcParams['axes.labelsize'] = 10
mpl.rcParams['axes.titlesize'] = 10.5
mpl.rcParams['xtick.labelsize'] = 8.5
mpl.rcParams['ytick.labelsize'] = 8.5
mpl.rcParams['legend.fontsize'] = 7.8
mpl.rcParams['figure.titlesize'] = 12

with open("logs/pretraining_telemetry_verified.json") as f:
    pre_data = json.load(f)

audit_file = "logs/round2_reviewer_audit.json" if os.path.exists("logs/round2_reviewer_audit.json") else "logs/defensible_tier1_audit.json"
with open(audit_file) as f:
    audit = json.load(f)

horizons = ["1m", "3m", "5m", "15m", "30m", "60m", "1d"]

# =========================================================================
# Figure 1: REAL Empirical Pre-Training Telemetry (855 Logged Steps)
# =========================================================================
print("Rendering Figure 1 from real pretraining_telemetry_verified.json (855 steps)...")
records = pre_data["records"]
steps = np.array([r["step"] for r in records])
t_loss = np.array([r["total_loss"] for r in records])
d_acc = np.array([r["dir_acc_total"] for r in records])

def causal_sma(arr, w=25):
    out = np.zeros_like(arr, dtype=np.float64)
    for i in range(len(arr)):
        lo = max(0, i - w + 1)
        out[i] = np.mean(arr[lo:i+1])
    return out

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.0), dpi=300)

ax1.scatter(steps, t_loss, color='#94A3B8', s=6, alpha=0.45, label='Raw Batch Telemetry (n=855)')
ax1.plot(steps, causal_sma(t_loss, 25), color='#0F766E', lw=1.8, label='25-Log Rolling Mean')
ax1.axvline(5010, color='#64748B', linestyle=':', lw=1.0, label='Epoch 1/2 Boundary (Step 5,010)')
ax1.axvline(10020, color='#475569', linestyle='--', lw=1.0, label='Epoch 2/3 Boundary (Step 10,020)')
ax1.axhline(0.3459, color='#E11D48', linestyle='-.', lw=1.1, label='Empirical Base-Rate Floor (0.5 H_b = 0.3459)')
ax1.set_xlabel('Pre-Training Optimization Step (Steps 320 – 15,030)')
ax1.set_ylabel('Total Joint Loss (Huber + 0.5 BCE)')
ax1.set_title('(a) Empirical Loss Trajectory & Variance Contraction')
ax1.set_ylim(0.315, 0.370)
ax1.legend(loc='upper right', frameon=True)
ax1.grid(True, alpha=0.3)

ax2.scatter(steps, d_acc, color='#93C5FD', s=6, alpha=0.45, label='Raw Batch Accuracy (n=855)')
ax2.plot(steps, causal_sma(d_acc, 25), color='#1D4ED8', lw=1.8, label='25-Log Rolling Mean (53.2% Mean)')
ax2.axvline(5010, color='#64748B', linestyle=':', lw=1.0)
ax2.axvline(10020, color='#475569', linestyle='--', lw=1.0)
ax2.axhline(50.0, color='#E11D48', linestyle='-.', lw=1.1, label='Martingale Null (50.0%)')
ax2.set_xlabel('Pre-Training Optimization Step (Steps 320 – 15,030)')
ax2.set_ylabel('Directional Accuracy (%) across 7 Horizons')
ax2.set_title('(b) Empirical Multi-Horizon Directional Accuracy')
ax2.set_ylim(35.0, 78.0)
ax2.legend(loc='upper right', frameon=True)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
fig.savefig("papers/figures/fig1_pretraining_curves.png", dpi=300, bbox_inches='tight', pad_inches=0.1)
plt.close()

# =========================================================================
# Figure 2: 1m Base-Rate / Flat-Bar Decomposition, Return Buckets & Baselines
# =========================================================================
print("Rendering Figure 2 from 1m return-magnitude buckets & multi-horizon ROC-AUC...")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.0), dpi=300)

# Panel (a): All-Bar ROC-AUC vs Strict Non-Zero ROC-AUC across Horizons (2024 OOS, N=8,044)
ep3_24 = audit["multi_checkpoint_oos_and_stress_audit"]["epoch3"]["2024"]["per_horizon"]
ep3_26 = audit["multi_checkpoint_oos_and_stress_audit"]["epoch3"]["2026"]["per_horizon"]
auc_all_24 = [ep3_24[h]["roc_auc"] for h in horizons]
auc_nz_24 = [ep3_24[h]["nonzero_roc_auc"] for h in horizons]
auc_nz_26 = [ep3_26[h]["nonzero_roc_auc"] for h in horizons]

x = np.arange(len(horizons))
w = 0.26
ax1.bar(x - w, auc_all_24, w, label='2024 All-Bar ROC-AUC', color='#0284C7', edgecolor='black', lw=0.5)
ax1.bar(x, auc_nz_24, w, label='2024 Non-Zero (y≠0) ROC-AUC', color='#0D9488', edgecolor='black', lw=0.5)
ax1.bar(x + w, auc_nz_26, w, label='2026 Non-Zero (y≠0) ROC-AUC', color='#D97706', edgecolor='black', lw=0.5)
ax1.axhline(0.5000, color='#E11D48', linestyle='--', lw=1.2, label='Random Guess / Flat Heuristic (0.5000)')
ax1.set_xticks(x)
ax1.set_xticklabels(horizons)
ax1.set_ylim(0.480, 0.545)
ax1.set_ylabel('Threshold-Independent ROC-AUC')
ax1.set_title('(a) All-Bar vs. Strict Non-Zero (y≠0) ROC-AUC by Horizon')
ax1.legend(loc='upper right', frameon=True)
ax1.grid(axis='y', alpha=0.3)

# Panel (b): 1-Minute Accuracy & ROC-AUC by |y| Return Magnitude Quartile (Q1..Q4, y != 0)
b_keys = ["Q1_Smallest_Move", "Q2_Moderate_Move", "Q3_Medium_Move", "Q4_Largest_Move"]
b_lbls = ["Q1 Micro\n(~0.75 bps)", "Q2 Small\n(~2.1 bps)", "Q3 Medium\n(~4.4 bps)", "Q4 Large\n(~12.4 bps)"]
cal_acc_24 = [ep3_24["1m"]["return_magnitude_buckets"][bk]["calibrated_dir_acc"] for bk in b_keys]
cal_acc_26 = [ep3_26["1m"]["return_magnitude_buckets"][bk]["calibrated_dir_acc"] for bk in b_keys]
b_auc_24 = [ep3_24["1m"]["return_magnitude_buckets"][bk]["roc_auc"] * 100.0 for bk in b_keys]
b_auc_26 = [ep3_26["1m"]["return_magnitude_buckets"][bk]["roc_auc"] * 100.0 for bk in b_keys]

xb = np.arange(len(b_lbls))
wb = 0.20
ax2.bar(xb - 1.5*wb, cal_acc_24, wb, label='2024 Calibrated Acc (%)', color='#1E3A8A', edgecolor='black', lw=0.5)
ax2.bar(xb - 0.5*wb, b_auc_24, wb, label='2024 Bucket ROC-AUC (×100)', color='#38BDF8', edgecolor='black', lw=0.5)
ax2.bar(xb + 0.5*wb, cal_acc_26, wb, label='2026 Calibrated Acc (%)', color='#0F766E', edgecolor='black', lw=0.5)
ax2.bar(xb + 1.5*wb, b_auc_26, wb, label='2026 Bucket ROC-AUC (×100)', color='#F59E0B', edgecolor='black', lw=0.5)
ax2.axhline(50.0, color='#E11D48', linestyle='--', lw=1.2, label='50.0% Null Benchmark')
ax2.set_xticks(xb)
ax2.set_xticklabels(b_lbls, fontsize=8)
ax2.set_ylim(49.0, 54.8)
ax2.set_ylabel('Non-Zero Bucket Accuracy (%) / AUC (×100)')
ax2.set_title('(b) 1-Minute Non-Zero Discrimination by |Return| Quartile')
ax2.legend(loc='upper right', frameon=True, fontsize=7.2)
ax2.grid(axis='y', alpha=0.3)

plt.tight_layout()
fig.savefig("papers/figures/fig2_oos_directional_accuracy.png", dpi=300, bbox_inches='tight', pad_inches=0.1)
plt.close()

# =========================================================================
# Figure 3: Frobenius Drift vs RMS Update Arithmetic & Pre-2024 Stress Forgetting
# =========================================================================
print("Rendering Figure 3 from Frobenius drift / Adam RMS arithmetic & Pre-2024 stress audit...")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.0), dpi=300)

modules = ["Return\nHead", "Depthwise\nConv", "SwiGLU\nFFN", "Conformer\nBackbone", "Dir.\nHead", "GQA\nAttn", "Inception\nTokenizer"]
mod_keys = ["return_regression_head", "sub_depthwise_conv", "sub_swiglu_ffn", "conformer_backbone_all", "directional_class_head", "sub_gqa_attention", "inception_tokenizer"]

wd = audit["weight_drift_and_cosine_audit"]
drift_wf24 = [wd["epoch3_to_wf2024"]["by_module"][k]["rel_frobenius_pct"] for k in mod_keys]
drift_wf25 = [wd["epoch3_to_wf2025"]["by_module"][k]["rel_frobenius_pct"] for k in mod_keys]
rms_dw_wf24 = [wd["epoch3_to_wf2024"]["by_module"][k]["rms_update_magnitude"] * 1e4 for k in mod_keys]
rms_w_base = [wd["epoch3_to_wf2024"]["by_module"][k]["rms_weight_norm"] * 50.0 for k in mod_keys]

xm = np.arange(len(modules))
wm = 0.25
ax1.bar(xm - wm, drift_wf24, wm, label='Rel. Drift Ep3→WF24 (%)', color='#0284C7', edgecolor='black', lw=0.5)
ax1.bar(xm, drift_wf25, wm, label='Rel. Drift Ep3→WF25 (%)', color='#0D9488', edgecolor='black', lw=0.5)
ax1.bar(xm + wm, rms_dw_wf24, wm, label='Abs. RMS Update ||ΔW||_RMS (×10⁴)', color='#F59E0B', edgecolor='black', lw=0.5)
ax1.plot(xm, rms_w_base, marker='D', color='#E11D48', lw=1.5, linestyle='--', label='Base Weight Norm ||W||_RMS (×50)')
ax1.set_xticks(xm)
ax1.set_xticklabels(modules, fontsize=7.5)
ax1.set_ylabel('Drift (%) / Scaled RMS Magnitude')
ax1.set_title('(a) Plasticity Hierarchy: Relative Drift vs. Absolute RMS Update')
ax1.legend(loc='upper left', frameon=True, fontsize=7.2)
ax1.grid(True, alpha=0.3)

# Panel (b): Pre-2024 Historical Stress Regimes & Forward OOS ROC-AUC (Epoch 3 vs WF-2024 vs WF-2025)
cohorts_lbl = ["2020-03\nCOVID Crash", "2022-06\nRate Shock", "2023-12\nPre-Anchor", "Pre-2024\nCombined", "2024\nOOS", "2025\nOOS", "2026\nForward"]
c_keys = ["stress_2020_03", "stress_2022_06", "stress_2023_12", "pre2024_stress", "2024", "2025", "2026"]

mc = audit["multi_checkpoint_oos_and_stress_audit"]
auc_ep3 = [mc["epoch3"][ck]["overall_roc_auc"] for ck in c_keys]
auc_wf24 = [mc["wf2024"][ck]["overall_roc_auc"] for ck in c_keys]
auc_wf25 = [mc["wf2025"][ck]["overall_roc_auc"] for ck in c_keys]

xc = np.arange(len(cohorts_lbl))
wc = 0.25
ax2.bar(xc - wc, auc_ep3, wc, label='Frozen Epoch 3 (2010–2023)', color='#1E3A8A', edgecolor='black', lw=0.5)
ax2.bar(xc, auc_wf24, wc, label='Online WF-2024 (+2024)', color='#0EA5E9', edgecolor='black', lw=0.5)
ax2.bar(xc + wc, auc_wf25, wc, label='Online WF-2025 (+2024–25)', color='#10B981', edgecolor='black', lw=0.5)
ax2.axhline(0.5000, color='#E11D48', linestyle='--', lw=1.1, label='Random Null (0.5000)')
ax2.set_xticks(xc)
ax2.set_xticklabels(cohorts_lbl, fontsize=7.5)
ax2.set_ylim(0.500, 0.522)
ax2.set_ylabel('7-Horizon Mean ROC-AUC')
ax2.set_title('(b) Backward Stress Forgetting & Forward Transfer Audit')
ax2.legend(loc='upper right', frameon=True, fontsize=7.2)
ax2.grid(axis='y', alpha=0.3)

plt.tight_layout()
fig.savefig("papers/figures/fig3_walk_forward_comparison.png", dpi=300, bbox_inches='tight', pad_inches=0.1)
plt.close()

# =========================================================================
# Figure 4: Controlled Mechanistic Kernel Ablation (Mean-Ablation, Controls & Denominator Effect)
# =========================================================================
print("Rendering Figure 4 from controlled mean-ablation, random-channel controls & RMS logit shift...")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.0), dpi=300)

ab = audit["mechanistic_kernel_ablation_controlled"]
ab_keys = ["MeanAblate_K3", "MeanAblate_K7", "MeanAblate_K15", "MeanAblate_K31", "MeanAblate_K63", "MeanAblate_K127", "Random_341ch_Control", "MeanAblate_Micro_K3_K7", "MeanAblate_Meso_K15_K31", "MeanAblate_Macro_K63_K127", "Random_682ch_Control"]
ab_lbls = ["K=3", "K=7", "K=15", "K=31", "K=63", "K=127", "Rand\n341ch", "Micro\n(3,7)", "Meso\n(15,31)", "Macro\n(63,127)", "Rand\n682ch"]

cent_mean = [ab["mean_ablations"][k]["mean_centered_logit_l2_pct"] for k in ab_keys]
cent_zero = [ab["zero_ablations"][k.replace("MeanAblate_", "ZeroAblate_")]["mean_centered_logit_l2_pct"] for k in ab_keys]

xa = np.arange(len(ab_lbls))
wa = 0.36
ax1.bar(xa - wa/2, cent_mean, wa, label='In-Distribution Mean-Ablation (Centered L2 %)', color='#4F46E5', edgecolor='black', lw=0.5)
ax1.bar(xa + wa/2, cent_zero, wa, label='Zero-Clamping Ablation (Centered L2 %)', color='#0EA5E9', edgecolor='black', lw=0.5)
ax1.axhline(ab["mean_ablations"]["Random_341ch_Control"]["mean_centered_logit_l2_pct"], color='#E11D48', linestyle=':', lw=1.2, label='Random 341-Ch Control (20.61% ± 0.90%)')
ax1.set_xticks(xa)
ax1.set_xticklabels(ab_lbls, fontsize=7.2)
ax1.set_ylabel('Centered Logit Shift ||Δ(z - z̄)||_2 / ||z - z̄||_2 (%)')
ax1.set_title('(a) Centered Logit Sensitivity vs. Random-Channel Controls')
ax1.legend(loc='upper left', frameon=True, fontsize=7.2)
ax1.grid(axis='y', alpha=0.3)

# Panel (b): Resolving the Denominator Effect across Horizons (Uncentered Raw L2 vs Centered L2 vs Abs RMS)
k7_ph = ab["mean_ablations"]["MeanAblate_K7"]["per_horizon"]
meso_ph = ab["mean_ablations"]["MeanAblate_Meso_K15_K31"]["per_horizon"]
raw_l2_k7 = [k7_ph[h]["raw_logit_l2_pct"] for h in horizons]
cent_l2_k7 = [k7_ph[h]["centered_logit_l2_pct"] for h in horizons]
rms_k7_scaled = [k7_ph[h]["abs_logit_rms_shift"] * 1000.0 for h in horizons]
base_rms_scaled = [ab["baseline"]["per_horizon"][h]["logit_rms"] * 200.0 for h in horizons]

xh = np.arange(len(horizons))
ax2.plot(xh, raw_l2_k7, marker='o', lw=2.0, color='#D97706', label='Uncentered ||Δz||_2 / ||z||_2 (%) [Peaks at 15m]')
ax2.plot(xh, cent_l2_k7, marker='s', lw=2.0, color='#4F46E5', label='Centered ||Δ(z-z̄)||_2 / ||z-z̄||_2 (%) [Variance-Norm]')
ax2.plot(xh, rms_k7_scaled, marker='^', lw=2.0, color='#0F766E', label='Denominator-Free RMS(Δz) (×1000) [Monotonic]')
ax2.plot(xh, base_rms_scaled, marker='x', lw=1.4, linestyle='--', color='#64748B', label='Baseline Denom. ||z||_RMS (×200) [Min at 15m]')
ax2.set_xticks(xh)
ax2.set_xticklabels(horizons)
ax2.set_ylabel('Perturbation (%) / Scaled RMS Units')
ax2.set_title('(b) Decomposition of the 15m Logit Denominator Effect (K=7)')
ax2.legend(loc='upper right', frameon=True, fontsize=7.0)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
fig.savefig("papers/figures/fig4_kernel_ablation.png", dpi=300, bbox_inches='tight', pad_inches=0.1)
plt.close()

# =========================================================================
# Figure 5: Non-Overlapping Contiguous Backtest & Frozen 2024 Policy OOS
# =========================================================================
print("Rendering Figure 5 from non-overlapping state-turnover backtest & frozen 2024 policy...")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.0), dpi=300)

# Panel (a): Fixed Hurdle (tau=1.0sigma) Walk-Forward Gross Return & State-Hold Breakeven Across 2024, 2025, 2026
wf_tau1 = audit["frozen_2024_policy_oos_evaluation"]["walk_forward_weights"]["fixed_tau_1sigma_all_horizons"]
h_plot = ["1m", "3m", "5m", "15m", "30m", "60m"]
be_24 = [wf_tau1[h]["2024"]["breakeven_one_way_bps_state_hold"] for h in h_plot]
be_25 = [wf_tau1[h]["2025"]["breakeven_one_way_bps_state_hold"] for h in h_plot]
be_26 = [wf_tau1[h]["2026"]["breakeven_one_way_bps_state_hold"] for h in h_plot]

xp = np.arange(len(h_plot))
wp = 0.25
ax1.bar(xp - wp, be_24, wp, label='2024 OOS State-Hold Breakeven (bps)', color='#0284C7', edgecolor='black', lw=0.5)
ax1.bar(xp, be_25, wp, label='2025 OOS State-Hold Breakeven (bps)', color='#0D9488', edgecolor='black', lw=0.5)
ax1.bar(xp + wp, be_26, wp, label='2026 OOS State-Hold Breakeven (bps)', color='#7C3AED', edgecolor='black', lw=0.5)
ax1.axhline(0.0, color='black', lw=0.9)
ax1.axhline(1.5, color='#E11D48', linestyle='--', lw=1.2, label='Maker / Midpoint Cost Hurdle (1.5 bps)')
ax1.set_xticks(xp)
ax1.set_xticklabels(h_plot)
ax1.set_ylabel('State-Hold Breakeven One-Way Cost (bps)')
ax1.set_title('(a) Non-Overlapping Breakeven Cost (Fixed τ=1.0σ, Walk-Forward)')
ax1.legend(loc='lower left', frameon=True, fontsize=7.2)
ax1.grid(axis='y', alpha=0.3)

# Panel (b): Frozen 2024 Selected Policy (60m, tau=0.0sigma) vs Fixed 3m (tau=1.0sigma) Across Cost Grid
cost_bps = [0.0, 1.5, 3.0, 5.0]
cost_keys = ["0.0bps_oneway", "1.5bps_oneway", "3.0bps_oneway", "5.0bps_oneway"]

fpol_wf = audit["frozen_2024_policy_oos_evaluation"]["walk_forward_weights"]["frozen_2024_optimal_policy"]["yearly_performance"]
sr_f24 = [fpol_wf["2024"]["cost_grid"][ck]["annualized_sharpe_state_hold"] for ck in cost_keys]
sr_f25 = [fpol_wf["2025"]["cost_grid"][ck]["annualized_sharpe_state_hold"] for ck in cost_keys]
sr_f26 = [fpol_wf["2026"]["cost_grid"][ck]["annualized_sharpe_state_hold"] for ck in cost_keys]
sr_3m_26 = [wf_tau1["3m"]["2026"]["cost_grid"][ck]["annualized_sharpe_state_hold"] for ck in cost_keys]

ax2.plot(cost_bps, sr_f24, marker='o', lw=2.0, color='#0284C7', label='Frozen Policy on 2024 (Selection: SR=+4.40 at 0bp, +2.04 at 1.5bp)')
ax2.plot(cost_bps, sr_f25, marker='s', lw=2.0, color='#0D9488', label='Frozen Policy Applied Untouched to 2025 (SR=+1.01 at 0bp)')
ax2.plot(cost_bps, sr_f26, marker='^', lw=2.0, color='#D97706', label='Frozen Policy Applied Untouched to 2026 (SR=-2.01 at 0bp)')
ax2.plot(cost_bps, sr_3m_26, marker='D', lw=1.8, linestyle='--', color='#7C3AED', label='2026 3m Demeaned (τ=1.0σ: SR=+3.56 at 0bp)')
ax2.axhline(0.0, color='black', linestyle='--', lw=1.1, label='Zero Net Sharpe Boundary')
ax2.set_xlabel('One-Way Execution Cost: Spread + Fee + Slippage (bps)')
ax2.set_ylabel('Annualized Net Sharpe Ratio (Non-Overlapping)')
ax2.set_title('(b) Frozen 2024 Policy Untouched OOS Transfer (2024–2026)')
ax2.legend(loc='lower left', frameon=True, fontsize=7.0)
ax2.grid(True, alpha=0.3)

plt.tight_layout()
fig.savefig("papers/figures/fig5_slippage_friction_curve.png", dpi=300, bbox_inches='tight', pad_inches=0.1)
plt.close()

print("✅ All 5 upgraded deterministic empirical figures rendered at 300 DPI!")