"""Per-class confusion analysis and hard-negative pair detection.

Compares to config.HARD_NEG_PAIRS prior to validate the dataset's known
confusable pairs.
"""
import sys, os
sys.path.insert(0, "/mnt/d/SpiceNet" if os.path.exists("/mnt/d/SpiceNet") else "D:/SpiceNet")

import json
import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report

import config
from src.dataset import get_dataloaders
from src.model import load_checkpoint


@torch.no_grad()
def collect_preds(model, loader, device, mode="fusion"):
    model.eval()
    y_true, y_pred = [], []
    for imgs, tex, col, lbl in loader:
        imgs = imgs.to(device)
        if mode == "fusion":
            tex, col = tex.to(device), col.to(device)
            logits, _ = model.forward_fusion(imgs, tex, col)
        else:
            logits = model.forward_image(imgs)
        y_true.extend(lbl.tolist())
        y_pred.extend(logits.argmax(1).cpu().tolist())
    return np.array(y_true), np.array(y_pred)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, *_ = load_checkpoint(str(config.CHECKPOINT_DIR / "best.pth"), device)

    _, _, test_loader, _, _ = get_dataloaders(multimodal=True)
    y_true, y_pred = collect_preds(model, test_loader, device)

    cm = confusion_matrix(y_true, y_pred, labels=list(range(len(config.CLASSES))))
    print("\nConfusion matrix counts:")
    print("    " + " ".join(f"{c[:4]:>5}" for c in config.CLASSES))
    for i, row in enumerate(cm):
        print(f"{config.CLASSES[i][:4]:>4} " + " ".join(f"{v:>5d}" for v in row))

    # Find all off-diagonal errors (i, j with i != j)
    off = []
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            if i != j and cm[i, j] > 0:
                off.append((cm[i, j], i, j))
    off.sort(reverse=True)

    print(f"\nTotal misclassifications: {sum(c for c,_,_ in off)} / {cm.sum()}")
    print("\nTop confusion pairs (true -> pred):")
    for cnt, i, j in off[:10]:
        marker = " ⚠ PRIOR" if {i, j} in [set(p) for p in config.HARD_NEG_PAIRS] else ""
        print(f"  {cnt:>3d}  {config.CLASSES[i]:>14s} -> {config.CLASSES[j]:<14s}{marker}")

    # Symmetric confusion: cm[i,j] + cm[j,i]
    n = cm.shape[0]
    sym = []
    for i in range(n):
        for j in range(i + 1, n):
            both = cm[i, j] + cm[j, i]
            if both > 0:
                sym.append((both, i, j))
    sym.sort(reverse=True)

    # Compare with the prior HARD_NEG_PAIRS
    prior_set = set(frozenset(p) for p in config.HARD_NEG_PAIRS)
    print("\nPrior hard-negative pairs (from config.HARD_NEG_PAIRS):")
    for i, j in config.HARD_NEG_PAIRS:
        pair_cnt = int(cm[i, j] + cm[j, i])
        print(f"  {config.CLASSES[i]} <-> {config.CLASSES[j]}: {pair_cnt} confusions")

    print("\nMost-confused pairs in practice (top 5):")
    for cnt, i, j in sym[:5]:
        in_prior = frozenset((i, j)) in prior_set
        print(f"  {config.CLASSES[i]} <-> {config.CLASSES[j]}: {cnt}  (in prior: {in_prior})")

    # Save report
    rep = classification_report(y_true, y_pred,
                                target_names=config.CLASSES, digits=4,
                                output_dict=True)
    out = {
        "confusion_matrix": cm.tolist(),
        "classes": config.CLASSES,
        "prior_hard_neg_pairs": config.HARD_NEG_PAIRS,
        "top_confusion_pairs": [(int(c), config.CLASSES[i], config.CLASSES[j]) for c, i, j in off[:10]],
        "symmetric_confusion_pairs": [(int(c), config.CLASSES[i], config.CLASSES[j]) for c, i, j in sym[:10]],
        "classification_report": rep,
    }
    out_path = config.OUTPUT_DIR / "confusion_analysis.json"
    with open(out_path, "w") as f:
        json.dump(out, f, indent=2)
    print(f"\nSaved -> {out_path}")

    # Pretty CM heatmap (counts + percentages)
    fig, ax = plt.subplots(figsize=(11, 9))
    cm_pct = cm.astype(float) / cm.sum(axis=1, keepdims=True) * 100
    im = ax.imshow(cm_pct, cmap="Blues", vmin=0, vmax=100)
    ax.set_xticks(range(len(config.CLASSES)))
    ax.set_yticks(range(len(config.CLASSES)))
    ax.set_xticklabels(config.CLASSES, rotation=45, ha="right")
    ax.set_yticklabels(config.CLASSES)
    ax.set_xlabel("Predicted"); ax.set_ylabel("True")
    ax.set_title("Confusion matrix (strong-aug, % per row)")
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            v = cm[i, j]
            if v > 0:
                color = "white" if cm_pct[i, j] > 50 else "black"
                ax.text(j, i, f"{v}", ha="center", va="center", color=color, fontsize=8)
    plt.colorbar(im, ax=ax, label="% of true class")
    plt.tight_layout()
    plt.savefig(config.OUTPUT_DIR / "confusion_matrix_annotated.png", dpi=150)
    plt.close()
    print(f"Saved -> {config.OUTPUT_DIR / 'confusion_matrix_annotated.png'}")


if __name__ == "__main__":
    main()