visualize_v3.py

"""
visualize_v3.py — Hatch v3 Comparison Visualizations
======================================================
4-panel figure:
  1. Three-way bar chart: v1 vs v2 vs v3 best configs
  2. Mean F1 heatmap: T_fill × K (best D, T_drain=2, SH=True)
  3. Cup fill trace: hemoglobin (folded) — v2 vs v3 side-by-side
  4. Super-Hatch impact: SH=True vs SH=False across K values
"""

import json
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LinearSegmentedColormap

from training_engine import load_thresholds
from inference_engine import HatchClassifier
from inference_engine_v2 import HatchClassifierV2
from inference_engine_v3 import HatchClassifierV3

# ─── Style ─────────────────────────────────────────────────────────────────────
plt.rcParams.update({
    "figure.facecolor": "#0d1117",
    "axes.facecolor":   "#161b22",
    "axes.edgecolor":   "#30363d",
    "axes.labelcolor":  "#c9d1d9",
    "xtick.color":      "#8b949e",
    "ytick.color":      "#8b949e",
    "text.color":       "#c9d1d9",
    "grid.color":       "#21262d",
    "grid.linewidth":   0.5,
    "font.family":      "monospace",
    "font.size":        10,
})

ACCENT_RED    = "#f85149"
ACCENT_ORANGE = "#f0883e"
ACCENT_GREEN  = "#3fb950"
ACCENT_BLUE   = "#58a6ff"
ACCENT_PURPLE = "#bc8cff"
ACCENT_CYAN   = "#39d353"

# ─── Load data ─────────────────────────────────────────────────────────────────

df_v3 = pd.read_csv("optimization_report_v3.csv")
with open("best_config_v3.json") as f:
    best_v3 = json.load(f)
with open("best_config_v2.json") as f:
    best_v2 = json.load(f)

thresholds, train_means, _ = load_thresholds("hatch_thresholds_W30.json")

v1_best = {
    "W": 40, "K": 6, "score_threshold": 4,
    "f1_folded": 0.3491, "f1_disordered": 0.3551, "mean_f1": 0.3521,
}

print(f"v1  best: mean_f1={v1_best['mean_f1']:.4f}")
print(f"v2  best: N={best_v2['N']} D={best_v2['D']} K={best_v2['K']} "
      f"score≥{best_v2['score_threshold']} mean_f1={best_v2['mean_f1']:.4f}")
print(f"v3  best: T_fill={best_v3['T_fill']} T_drain={best_v3['T_drain']} "
      f"K={best_v3['K']} D={best_v3['D']} SH={best_v3['super_hatch']} "
      f"mean_f1={best_v3['mean_f1']:.4f}")

# ─── Figure layout ─────────────────────────────────────────────────────────────

fig = plt.figure(figsize=(18, 14), facecolor="#0d1117")
fig.suptitle(
    "Hatch Protein Classifier — v1 → v2 → v3 Progression\n"
    "Hysteresis Asymmetric Thresholds + Super-Hatch Hard Reset",
    fontsize=14, color="#e6edf3", fontweight="bold", y=0.98
)

gs = gridspec.GridSpec(2, 2, figure=fig, hspace=0.38, wspace=0.32,
                       left=0.07, right=0.97, top=0.92, bottom=0.06)

# ─── Plot 1: Three-way bar chart ───────────────────────────────────────────────

ax1 = fig.add_subplot(gs[0, 0])

metrics = ["F1 Folded", "F1 Disordered", "Mean F1"]
v1_vals = [v1_best["f1_folded"],  v1_best["f1_disordered"],  v1_best["mean_f1"]]
v2_vals = [best_v2["f1_folded"],  best_v2["f1_disordered"],  best_v2["mean_f1"]]
v3_vals = [best_v3["f1_folded"],  best_v3["f1_disordered"],  best_v3["mean_f1"]]

x = np.arange(len(metrics))
w = 0.25
bars1 = ax1.bar(x - w,     v1_vals, w, label="v1 (full-forgiveness)", color=ACCENT_RED,    alpha=0.85)
bars2 = ax1.bar(x,         v2_vals, w, label="v2 (consensus drain)",  color=ACCENT_ORANGE, alpha=0.85)
bars3 = ax1.bar(x + w,     v3_vals, w, label="v3 (hysteresis+SH)",    color=ACCENT_GREEN,  alpha=0.85)

for bars, vals, col in [(bars1, v1_vals, ACCENT_RED),
                         (bars2, v2_vals, ACCENT_ORANGE),
                         (bars3, v3_vals, ACCENT_GREEN)]:
    for bar, val in zip(bars, vals):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.008,
                 f"{val:.3f}", ha="center", va="bottom", fontsize=8, color=col)

ax1.set_xticks(x)
ax1.set_xticklabels(metrics, fontsize=10)
ax1.set_ylim(0, 0.85)
ax1.set_ylabel("Score", fontsize=9)
ax1.set_title("v1 → v2 → v3 Best Configuration", color="#e6edf3", fontsize=11)
ax1.legend(fontsize=8, loc="upper left")
ax1.axhline(0.5, color="#30363d", linestyle="--", linewidth=0.8)
ax1.grid(axis="y", alpha=0.3)

# Annotate total improvement
v1_mf1 = v1_best["mean_f1"]
v3_mf1 = best_v3["mean_f1"]
delta_total = v3_mf1 - v1_mf1
ax1.annotate(
    f"+{delta_total:.3f} total\nvs v1",
    xy=(2 + w, v3_mf1),
    xytext=(2 + w + 0.25, v3_mf1 + 0.05),
    fontsize=9, color=ACCENT_GREEN,
    arrowprops=dict(arrowstyle="->", color=ACCENT_GREEN, lw=1.2),
)

# ─── Plot 2: Mean F1 Heatmap (T_fill × K, best D, T_drain=2, SH=True) ─────────

ax2 = fig.add_subplot(gs[0, 1])

best_D_v3 = best_v3["D"]
sub = df_v3[
    (df_v3["T_drain"] == 2) &
    (df_v3["D"] == best_D_v3) &
    (df_v3["super_hatch"] == True)
]
pivot = sub.pivot(index="T_fill", columns="K", values="mean_f1")

cmap = LinearSegmentedColormap.from_list(
    "hatch3", ["#161b22", "#1f6feb", "#58a6ff", "#3fb950", "#39d353"], N=256
)
im = ax2.imshow(pivot.values, cmap=cmap, aspect="auto",
                vmin=df_v3["mean_f1"].min(), vmax=df_v3["mean_f1"].max())

ax2.set_xticks(range(len(pivot.columns)))
ax2.set_xticklabels([f"K={k}" for k in pivot.columns], fontsize=9)
ax2.set_yticks(range(len(pivot.index)))
ax2.set_yticklabels([f"T_fill={t}" for t in pivot.index], fontsize=9)
ax2.set_title(f"Mean F1 Heatmap (T_drain=2, D={best_D_v3}, SH=True)", color="#e6edf3", fontsize=11)
ax2.set_xlabel("Cup Depth K", fontsize=9)
ax2.set_ylabel("Fill Threshold T_fill", fontsize=9)

for i in range(len(pivot.index)):
    for j in range(len(pivot.columns)):
        val = pivot.values[i, j]
        ax2.text(j, i, f"{val:.3f}", ha="center", va="center",
                 fontsize=9, color="#e6edf3" if val > 0.50 else "#8b949e")

# Mark best cell
best_tf = best_v3["T_fill"]
best_k  = best_v3["K"]
if best_tf in pivot.index.values and best_k in pivot.columns.values:
    bi = list(pivot.index).index(best_tf)
    bj = list(pivot.columns).index(best_k)
    ax2.add_patch(plt.Rectangle((bj - 0.5, bi - 0.5), 1, 1,
                                  fill=False, edgecolor="#f0883e", lw=2.5))
    ax2.text(bj, bi - 0.38, "★ best", ha="center", va="top",
             fontsize=8, color=ACCENT_ORANGE)

plt.colorbar(im, ax=ax2, label="Mean F1", fraction=0.046, pad=0.04)

# ─── Plot 3: Cup Fill Trace — v2 vs v3 on hemoglobin ──────────────────────────

ax3 = fig.add_subplot(gs[1, 0])

TRACE_K = 20  # large K to see full dynamics without early overflow

clf_v2 = HatchClassifierV2(
    thresholds, train_means,
    K=TRACE_K, score_threshold=4, window_size=30, N=2, D=3
)
clf_v3 = HatchClassifierV3(
    thresholds, train_means,
    K=TRACE_K, T_fill=3, T_drain=2, N=2, D=3,
    window_size=30, super_hatch=True
)

# Hemoglobin alpha — a folded protein that v2 still struggles with
hemo_seq = (
    "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALT"
    "NAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR"
)

t2 = clf_v2.trace(hemo_seq)
t3 = clf_v3.trace(hemo_seq)

pos2 = t2["positions"]
pos3 = t3["positions"]

ax3.plot(pos2, t2["entropy_levels"], color=ACCENT_ORANGE, linewidth=1.8,
         label=f"v2 (score≥4, K={TRACE_K})", alpha=0.9)
ax3.plot(pos3, t3["entropy_levels"], color=ACCENT_GREEN, linewidth=1.8,
         label=f"v3 (T_fill=3 T_drain=2, K={TRACE_K})", alpha=0.9)

# Mark v3 Super-Hatch events
sh_pos  = [pos3[i] for i, s in enumerate(t3["super_hatch_events"]) if s]
sh_lvls = [t3["entropy_levels"][i] for i, s in enumerate(t3["super_hatch_events"]) if s]
if sh_pos:
    ax3.scatter(sh_pos, sh_lvls, marker="*", color=ACCENT_PURPLE, s=120,
                zorder=6, label=f"v3 Super-Hatch ({len(sh_pos)} events)", alpha=0.9)

# Mark v3 grace zone events
grace_pos  = [pos3[i] for i, g in enumerate(t3["grace_events"]) if g]
grace_lvls = [t3["entropy_levels"][i] for i, g in enumerate(t3["grace_events"]) if g]
if grace_pos:
    ax3.scatter(grace_pos, grace_lvls, marker="o", color=ACCENT_CYAN, s=30,
                zorder=5, label=f"v3 Grace zone ({len(grace_pos)} events)", alpha=0.7)

# Mark v2 drain events
drain2_pos  = [pos2[i] for i, d in enumerate(t2["drain_events"]) if d]
drain2_lvls = [t2["entropy_levels"][i] for i, d in enumerate(t2["drain_events"]) if d]
if drain2_pos:
    ax3.scatter(drain2_pos, drain2_lvls, marker="v", color=ACCENT_BLUE, s=30,
                zorder=5, label=f"v2 drain ({len(drain2_pos)} events)", alpha=0.6)

ax3.axhline(TRACE_K, color="#8b949e", linestyle="--", linewidth=0.8, alpha=0.5,
            label=f"Overflow threshold (K={TRACE_K})")

ax3.set_xlabel("Sequence Position (residue)", fontsize=9)
ax3.set_ylabel("Entropy Level (cup fill)", fontsize=9)
ax3.set_title("Cup Fill Trace: Hemoglobin Alpha (FOLDED)\nv2 Consensus vs v3 Hysteresis+Super-Hatch",
              color="#e6edf3", fontsize=11)
ax3.legend(fontsize=7, loc="upper right", ncol=1)
ax3.grid(alpha=0.3)

# Annotate final cup levels
final_v2 = t2["entropy_levels"][-1] if t2["entropy_levels"] else 0
final_v3 = t3["entropy_levels"][-1] if t3["entropy_levels"] else 0
ax3.annotate(f"v2 final: {final_v2}",
             xy=(pos2[-1], final_v2), xytext=(pos2[-1] - 20, final_v2 + 1.5),
             fontsize=8, color=ACCENT_ORANGE,
             arrowprops=dict(arrowstyle="->", color=ACCENT_ORANGE, lw=0.8))
ax3.annotate(f"v3 final: {final_v3}",
             xy=(pos3[-1], final_v3), xytext=(pos3[-1] - 20, final_v3 - 2.5),
             fontsize=8, color=ACCENT_GREEN,
             arrowprops=dict(arrowstyle="->", color=ACCENT_GREEN, lw=0.8))

# ─── Plot 4: Super-Hatch Impact — SH=True vs SH=False across K ────────────────

ax4 = fig.add_subplot(gs[1, 1])

# For T_fill=3, T_drain=2, D=3 — compare SH=True vs SH=False across K
sub_sh_on  = df_v3[(df_v3["T_fill"] == 3) & (df_v3["T_drain"] == 2) &
                    (df_v3["D"] == 3) & (df_v3["super_hatch"] == True)].sort_values("K")
sub_sh_off = df_v3[(df_v3["T_fill"] == 3) & (df_v3["T_drain"] == 2) &
                    (df_v3["D"] == 3) & (df_v3["super_hatch"] == False)].sort_values("K")

k_vals = sub_sh_on["K"].values

ax4.plot(k_vals, sub_sh_on["mean_f1"],   color=ACCENT_PURPLE, linewidth=2.0,
         marker="o", markersize=7, label="SH=True (Super-Hatch ON)")
ax4.plot(k_vals, sub_sh_off["mean_f1"],  color="#8b949e",     linewidth=2.0,
         marker="s", markersize=7, linestyle="--", label="SH=False (Super-Hatch OFF)")

ax4.fill_between(k_vals,
                  sub_sh_on["mean_f1"],
                  sub_sh_off["mean_f1"],
                  alpha=0.15, color=ACCENT_PURPLE, label="SH delta")

# Also show F1 components for SH=True
ax4.plot(k_vals, sub_sh_on["f1_folded"],     color=ACCENT_GREEN, linewidth=1.2,
         marker="^", markersize=5, linestyle=":", alpha=0.8, label="SH=True F1 Folded")
ax4.plot(k_vals, sub_sh_on["f1_disordered"], color=ACCENT_ORANGE, linewidth=1.2,
         marker="v", markersize=5, linestyle=":", alpha=0.8, label="SH=True F1 Disordered")

# Annotate the delta at each K
for k, on, off in zip(k_vals, sub_sh_on["mean_f1"], sub_sh_off["mean_f1"]):
    delta = on - off
    if abs(delta) > 0.0005:
        ax4.annotate(f"+{delta:.4f}", xy=(k, (on + off) / 2),
                     xytext=(k + 0.1, (on + off) / 2 + 0.002),
                     fontsize=7, color=ACCENT_PURPLE)

# Mark v2 baseline
ax4.axhline(best_v2["mean_f1"], color=ACCENT_ORANGE, linestyle="--",
            linewidth=1.0, alpha=0.7, label=f"v2 baseline ({best_v2['mean_f1']:.3f})")

ax4.set_xlabel("Cup Depth K", fontsize=9)
ax4.set_ylabel("F1 Score", fontsize=9)
ax4.set_title("Super-Hatch Impact: SH=True vs SH=False\n(T_fill=3, T_drain=2, D=3, N=2, W=30)",
              color="#e6edf3", fontsize=11)
ax4.legend(fontsize=7, loc="lower right", ncol=1)
ax4.grid(alpha=0.3)
ax4.set_xticks(k_vals)

# ─── Save ──────────────────────────────────────────────────────────────────────

plt.savefig("hatch_v3_analysis.png", dpi=150, bbox_inches="tight",
            facecolor="#0d1117")
print("Saved: hatch_v3_analysis.png")
plt.close()