"""
Evaluate saved ModernBERT model on the held-out test split.
=========================================================
Loads the saved test split from splits/test.csv and evaluates
the model on it.

Usage:
    python eval_bert.py --model modernbert/best --splits-dir splits
"""

import csv
import argparse
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix

from config import CLASSES, BINARY_CLASSES

csv.field_size_limit(10_000_000)

MAX_TOKENS = 8192
BATCH_SIZE = 4


# ── Inference ─────────────────────────────────────────────────────────────────

def run_inference(model, tokenizer, device, texts, id2label):
    all_preds, all_confs, all_probs = [], [], []
    for i in range(0, len(texts), BATCH_SIZE):
        batch = texts[i : i + BATCH_SIZE]
        enc = tokenizer(
            batch,
            truncation=True,
            max_length=MAX_TOKENS,
            padding="longest",
            return_tensors="pt",
        ).to(device)
        with torch.no_grad():
            logits = model(**enc).logits
        probs = torch.softmax(logits, dim=-1).float().cpu().numpy()
        for prob in probs:
            best_idx = prob.argmax()
            all_preds.append(id2label[best_idx])
            all_confs.append(prob[best_idx])
            all_probs.append(prob)
        if (i // BATCH_SIZE) % 50 == 0:
            print(f"  {i + len(batch):,} / {len(texts):,}", end="\r")
    print()
    return all_preds, all_confs, all_probs


# ── Main ──────────────────────────────────────────────────────────────────────

def evaluate(model_dir, splits_dir):
    # Load saved test split
    splits_dir = Path(splits_dir)
    X_test, y_test = [], []
    with open(splits_dir / "test.csv") as f:
        for row in csv.DictReader(f):
            X_test.append(row["text"])
            y_test.append(row["label"])

    # Auto-detect binary vs multiclass from labels in data
    all_labels_set = set(y_test)
    if all_labels_set.issubset(set(BINARY_CLASSES)):
        classes = BINARY_CLASSES
        print("Detected binary classification mode")
    else:
        classes = CLASSES
        print("Detected multiclass classification mode")

    id2label = {i: cls for i, cls in enumerate(classes)}

    counts = defaultdict(int)
    for l in y_test:
        counts[l] += 1
    print(f"Test set: {len(X_test):,} examples")
    for cls in classes:
        print(f"  {cls:<25}: {counts[cls]:>5}")

    # Load model
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")
    print(f"Loading model from {model_dir}...")
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_dir, dtype=torch.bfloat16
    ).to(device)
    model.eval()

    # Run inference
    print(f"\nRunning inference on {len(X_test):,} test examples...")
    pred_names, confidences, _ = run_inference(model, tokenizer, device, X_test, id2label)

    # Report
    acc = accuracy_score(y_test, pred_names)
    report = classification_report(y_test, pred_names, labels=classes, zero_division=0)

    print(f"\n--- Test Set Results (acc={acc:.3f}) ---")
    print(report)

    # Confusion matrix
    cm = confusion_matrix(y_test, pred_names, labels=classes)
    print("Confusion matrix (rows=true, cols=pred):")
    header = "".join(f"{c[:8]:>10}" for c in classes)
    print(f"{'':>25}{header}")
    for i, cls in enumerate(classes):
        row = "".join(f"{cm[i,j]:>10}" for j in range(len(classes)))
        print(f"  {cls:<25}{row}")

    # Per-confidence breakdown
    print("\nMean confidence by true class:")
    class_confs = defaultdict(list)
    for true, conf in zip(y_test, confidences):
        class_confs[true].append(conf)
    for cls in classes:
        if class_confs[cls]:
            print(f"  {cls:<25}: {np.mean(class_confs[cls]):.3f}")

    # Save misclassified examples
    misclassified_path = Path(model_dir) / "misclassified.csv"
    with open(misclassified_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["true_label", "pred_label", "confidence", "text"])
        for text, true, pred, conf in zip(X_test, y_test, pred_names, confidences):
            if true != pred:
                writer.writerow([true, pred, f"{conf:.4f}", text])
    print(f"\nMisclassified examples saved to {misclassified_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model",      default="modernbert/best")
    parser.add_argument("--splits-dir", default="splits")
    args = parser.parse_args()

    evaluate(args.model, args.splits_dir)