#!/usr/bin/env python3
"""
Normalized confusion matrices + accuracy analysis for all four systems
vs ground truth from query_translation_eval.csv.
"""

from __future__ import annotations
import csv, re
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

ROOT = Path(__file__).resolve().parents[1]

# ── Load ground truth ────────────────────────────────────────────────────────
def load_ground_truth() -> list[str]:
    gt = []
    with (ROOT / "datasets" / "query_translation_eval.csv").open(encoding="utf-8") as f:
        for row in csv.DictReader(f):
            gt.append(row["answer"].strip().upper())
    return gt   # 100 entries, "T" or "F"

# ── Load baseline trans_query_eval.csv (end-to-end via baseline UPPAAL model) ──
def load_raw_qna(path: Path) -> list[str]:
    preds = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        m = re.search(r":\s*([TF])\s*$", line, re.IGNORECASE)
        if m:
            preds.append(m.group(1).upper())
    return preds

# ── Load our e2e results ──────────────────────────────────────────────────────
def load_ours() -> list[str]:
    preds = []
    with (ROOT / "artifacts" / "e2e_eval" / "e2e_query_eval.csv").open(encoding="utf-8") as f:
        for row in csv.DictReader(f):
            v = row["verdict"].strip().upper()
            preds.append(v if v in ("T", "F") else "?")
    return preds

# ── Confusion matrix (normalized by true label) ───────────────────────────────
def confusion(gt: list[str], pred: list[str]):
    """Returns 2x2 numpy array normalized by row (true label).
       Rows = actual [T, F], Cols = predicted [T, F]
    """
    cm = np.zeros((2, 2), dtype=float)
    for g, p in zip(gt, pred):
        if g not in ("T", "F") or p not in ("T", "F"):
            continue
        ri = 0 if g == "T" else 1
        ci = 0 if p == "T" else 1
        cm[ri, ci] += 1
    # normalize by row
    row_sums = cm.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1
    return cm / row_sums, cm   # normalized, raw

def accuracy(gt, pred):
    correct = sum(g == p for g, p in zip(gt, pred) if g in ("T","F") and p in ("T","F"))
    total   = sum(1 for g in gt if g in ("T","F"))
    return correct / total if total else 0.0

def precision_recall_f1(gt, pred):
    tp = sum(1 for g,p in zip(gt,pred) if g=="T" and p=="T")
    fp = sum(1 for g,p in zip(gt,pred) if g=="F" and p=="T")
    fn = sum(1 for g,p in zip(gt,pred) if g=="T" and p=="F")
    tn = sum(1 for g,p in zip(gt,pred) if g=="F" and p=="F")
    prec = tp/(tp+fp) if (tp+fp) else 0.0
    rec  = tp/(tp+fn) if (tp+fn) else 0.0
    f1   = 2*prec*rec/(prec+rec) if (prec+rec) else 0.0
    return prec, rec, f1, tp, fp, fn, tn

# ── Main ──────────────────────────────────────────────────────────────────────
gt = load_ground_truth()

systems = {
    "Ours\n(pipeline)":  load_ours(),
    "Claude\n(direct)":  load_raw_qna(ROOT / "artifacts/baselines/claude/raw_qna.txt"),
    "GPT\n(direct)":     load_raw_qna(ROOT / "artifacts/baselines/gpt/raw_qna.txt"),
    "Grok\n(direct)":    load_raw_qna(ROOT / "artifacts/baselines/grok/raw_qna.txt"),
}

# ── Print stats table ─────────────────────────────────────────────────────────
print(f"\n{'System':<22} {'Acc':>5}  {'Prec':>5}  {'Rec':>5}  {'F1':>5}  {'TP':>3} {'TN':>3} {'FP':>3} {'FN':>3}")
print("-" * 70)
stats = {}
for name, pred in systems.items():
    acc = accuracy(gt, pred)
    prec, rec, f1, tp, fp, fn, tn = precision_recall_f1(gt, pred)
    label = name.replace("\n", " ")
    print(f"{label:<22} {acc:.3f}  {prec:.3f}  {rec:.3f}  {f1:.3f}  {tp:>3} {tn:>3} {fp:>3} {fn:>3}")
    stats[name] = (acc, prec, rec, f1)

# ── Plot: 1 row of 4 normalized confusion matrices ───────────────────────────
fig, axes = plt.subplots(1, 4, figsize=(14, 4))
fig.suptitle("Normalized Confusion Matrices — QA Verdict Prediction\n(rows = actual, cols = predicted; normalized by true label)",
             fontsize=12, y=1.02)

labels = ["T (True)", "F (False)"]
cmap = "Blues"

for ax, (name, pred) in zip(axes, systems.items()):
    cm_norm, cm_raw = confusion(gt, pred)
    acc = stats[name][0]

    im = ax.imshow(cm_norm, vmin=0, vmax=1, cmap=cmap)

    # Cell annotations: normalized value + raw count
    for i in range(2):
        for j in range(2):
            val  = cm_norm[i, j]
            raw  = int(cm_raw[i, j])
            color = "white" if val > 0.6 else "black"
            ax.text(j, i, f"{val:.2f}\n({raw})", ha="center", va="center",
                    fontsize=11, color=color, fontweight="bold")

    ax.set_xticks([0, 1]); ax.set_yticks([0, 1])
    ax.set_xticklabels(labels, fontsize=9)
    ax.set_yticklabels(labels, fontsize=9)
    ax.set_xlabel("Predicted", fontsize=10)
    ax.set_ylabel("Actual", fontsize=10)
    ax.set_title(f"{name}\nAcc = {acc:.2f}", fontsize=10)

plt.colorbar(im, ax=axes[-1], fraction=0.046, pad=0.04, label="Proportion")
plt.tight_layout()
out = ROOT / "docs" / "confusion_matrices.png"
plt.savefig(out, dpi=150, bbox_inches="tight")
print(f"\nSaved -> {out}")
plt.show()