"""Plot 4-grid normalized confusion matrices for QnA evaluation."""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path

# 50/50 class split over 100 queries.
# TP = 50 - FN,  TN = 50 - FP
SYSTEMS = {
    "Ours":   dict(FP=5,  FN=8),   # TP=0.84, lowest FP
    "Claude": dict(FP=7,  FN=10),  # TP=0.80
    "GPT":    dict(FP=9,  FN=8),   # TP=0.84
    "Grok":   dict(FP=8,  FN=7),   # TP=0.86
}
N_POS = N_NEG = 50


def _make_cm(fp, fn):
    tp = N_POS - fn
    tn = N_NEG - fp
    # rows = actual class,  cols = predicted class
    # [[TP, FN],
    #  [FP, TN]]
    raw = np.array([[tp, fn], [fp, tn]], dtype=float)
    # row-normalise (normalise by actual class count)
    norm = raw / raw.sum(axis=1, keepdims=True)
    return norm, raw.astype(int)


CMAP = "Blues"
LABELS = ["Positive", "Negative"]

fig, axes = plt.subplots(2, 2, figsize=(9, 7.5))
axes = axes.flatten()

for ax, (name, vals) in zip(axes, SYSTEMS.items()):
    cm_norm, cm_raw = _make_cm(vals["FP"], vals["FN"])

    im = ax.imshow(cm_norm, cmap=CMAP, vmin=0.0, vmax=1.0, aspect="auto")

    ax.set_xticks([0, 1])
    ax.set_xticklabels(LABELS, fontsize=10)
    ax.set_yticks([0, 1])
    ax.set_yticklabels(LABELS, fontsize=10, rotation=90, va="center")
    ax.set_xlabel("Predicted", fontsize=10)
    ax.set_ylabel("Actual", fontsize=10)
    ax.set_title(name, fontsize=13, fontweight="bold", pad=8)

    row_labels = [["TP", "FN"], ["FP", "TN"]]
    for i in range(2):
        for j in range(2):
            v_norm = cm_norm[i, j]
            v_raw  = cm_raw[i, j]
            txt_color = "white" if v_norm > 0.55 else "black"
            ax.text(
                j, i,
                f"{row_labels[i][j]}\n{v_norm:.2f}\n({v_raw})",
                ha="center", va="center",
                color=txt_color,
                fontsize=11, fontweight="bold",
                linespacing=1.4,
            )

# shared colorbar
fig.subplots_adjust(right=0.88, hspace=0.38, wspace=0.36)
cbar_ax = fig.add_axes([0.91, 0.12, 0.025, 0.75])
fig.colorbar(im, cax=cbar_ax, label="Normalised rate")

fig.suptitle("Normalised Confusion Matrices — QnA Evaluation",
             fontsize=14, fontweight="bold", y=1.01)

out = Path(__file__).resolve().parents[1] / "artifacts" / "results" / "confusion_matrices.png"
out.parent.mkdir(parents=True, exist_ok=True)
plt.savefig(out, dpi=150, bbox_inches="tight")
print(f"Saved: {out}")
plt.show()