"""
Ticket 6 — Evaluation & Confusion Matrix Report (ATA classifier).

Runs the trained model against the TEST split (never used for training or
checkpoint selection -- that was val) and produces:
  - overall accuracy / macro-F1 (comparable to the PRD's success metric)
  - per-class precision/recall/F1 (exposes which classes are actually weak,
    which a single headline number hides)
  - a confusion matrix image
  - a metrics JSON, for the Model Performance Dashboard (Ticket 16) to read later
"""
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    precision_recall_fscore_support,
)
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_DIR = "models/ata_classifier"
TEST_CSV = "data/processed/test.csv"
MAX_LENGTH = 256
OUT_DIR = "reports/ata_classifier"


def main():
    Path(OUT_DIR).mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR, local_files_only=True)
    model.eval()

    with open(f"{MODEL_DIR}/label_map.json") as f:
        label_map = json.load(f)
    id2label = {int(k): v for k, v in label_map["id2label"].items()}
    label2id = label_map["label2id"]

    df = pd.read_csv(TEST_CSV)
    df = df[df["ata_chapter_label"].notna()].reset_index(drop=True)
    print(f"Test examples (ATA-labeled): {len(df)}")

    true_labels = df["ata_chapter_label"].astype(str).tolist()
    y_true = [label2id[t] for t in true_labels]

    y_pred = []
    with torch.no_grad():
        for text in df["narrative_text"].tolist():
            inputs = tokenizer(text, truncation=True, max_length=MAX_LENGTH, return_tensors="pt")
            logits = model(**inputs).logits
            pred_id = int(torch.argmax(logits, dim=1).item())
            y_pred.append(pred_id)

    acc = accuracy_score(y_true, y_pred)
    precision, recall, f1, support = precision_recall_fscore_support(
        y_true, y_pred, average="macro", zero_division=0
    )

    print(f"\n=== TEST SET RESULTS (never used in training or checkpoint selection) ===")
    print(f"Accuracy: {acc:.1%}   (PRD target: >=80%)")
    print(f"Macro F1: {f1:.3f}")
    print(f"Macro Precision: {precision:.3f}   Macro Recall: {recall:.3f}")

    labels_present = sorted(set(y_true) | set(y_pred))
    target_names = [id2label[i] for i in labels_present]
    report_dict = classification_report(
        y_true, y_pred, labels=labels_present, target_names=target_names,
        zero_division=0, output_dict=True
    )
    print("\nPer-class report:")
    print(classification_report(y_true, y_pred, labels=labels_present, target_names=target_names, zero_division=0))

    cm = confusion_matrix(y_true, y_pred, labels=labels_present)
    fig, ax = plt.subplots(figsize=(10, 9))
    im = ax.imshow(cm, cmap="Blues")
    ax.set_xticks(range(len(target_names)))
    ax.set_yticks(range(len(target_names)))
    ax.set_xticklabels(target_names, rotation=90)
    ax.set_yticklabels(target_names)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title("ATA Chapter Classifier — Confusion Matrix (Test Set)")
    for i in range(len(target_names)):
        for j in range(len(target_names)):
            if cm[i, j] > 0:
                ax.text(j, i, str(cm[i, j]), ha="center", va="center",
                         color="white" if cm[i, j] > cm.max() / 2 else "black", fontsize=7)
    fig.colorbar(im)
    fig.tight_layout()
    fig.savefig(f"{OUT_DIR}/confusion_matrix.png", dpi=150)
    print(f"\nSaved confusion matrix to {OUT_DIR}/confusion_matrix.png")

    metrics_out = {
        "test_accuracy": acc,
        "test_macro_f1": f1,
        "test_macro_precision": precision,
        "test_macro_recall": recall,
        "n_test_examples": len(df),
        "per_class": report_dict,
        "prd_target_accuracy": 0.80,
        "meets_prd_target": acc >= 0.80,
    }
    with open(f"{OUT_DIR}/metrics.json", "w") as f:
        json.dump(metrics_out, f, indent=2)
    print(f"Saved metrics JSON to {OUT_DIR}/metrics.json")


if __name__ == "__main__":
    main()