"""Compute final metrics, confusion matrix, threshold sweep and a failure
gallery for the saved CV model.

Outputs:
    models/cv_confusion_matrix.json
    models/cv_metrics.json   (test_overall, per_class, threshold_curve)
    docs/screenshots/cv_threshold_curve.png
    docs/screenshots/cv_failures.png

Usage:
    python -m src.cv.evaluate
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import torch
from sklearn.metrics import classification_report, confusion_matrix
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.config import (  # noqa: E402
    CV_METRICS_PATH,
    CV_MODEL_PATH,
    MODELS_DIR,
    PROCESSED_DIR,
    SCREENSHOTS_DIR,
)
from src.cv.train import build_model, build_transforms  # noqa: E402


CONFUSION_PATH = MODELS_DIR / "cv_confusion_matrix.json"
THRESHOLD_PLOT = SCREENSHOTS_DIR / "cv_threshold_curve.png"
FAILURE_PLOT = SCREENSHOTS_DIR / "cv_failures.png"

THRESHOLDS = [round(0.30 + 0.05 * i, 2) for i in range(14)]  # 0.30 .. 0.95
PRECISION_TARGET = 0.95


def _apply_temperature(logits: torch.Tensor, temperature: float | None) -> torch.Tensor:
    if temperature and temperature > 0:
        return logits / temperature
    return logits


def threshold_sweep(
    confidences: np.ndarray, correct: np.ndarray
) -> tuple[list[dict], float | None]:
    """Precision/recall/coverage of the 'accept (to_guest)' decision per threshold.

    precision = accepted_correct / accepted
    coverage  = accepted / total
    recall    = accepted_correct / total_correct
    """
    total = len(confidences)
    total_correct = int(correct.sum())
    curve: list[dict] = []
    recommended: float | None = None
    for t in THRESHOLDS:
        accepted_mask = confidences >= t
        accepted = int(accepted_mask.sum())
        accepted_correct = int(correct[accepted_mask].sum())
        precision = accepted_correct / accepted if accepted else 1.0
        coverage = accepted / total if total else 0.0
        recall = accepted_correct / total_correct if total_correct else 0.0
        curve.append(
            {
                "threshold": t,
                "precision": round(precision, 4),
                "recall": round(recall, 4),
                "coverage": round(coverage, 4),
                "accepted": accepted,
            }
        )
        if recommended is None and precision >= PRECISION_TARGET and accepted > 0:
            recommended = t
    return curve, recommended


def plot_threshold_curve(curve: list[dict], recommended: float | None) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    ts = [c["threshold"] for c in curve]
    prec = [c["precision"] for c in curve]
    rec = [c["recall"] for c in curve]
    cov = [c["coverage"] for c in curve]

    plt.figure(figsize=(8, 5))
    plt.plot(ts, prec, marker="o", label="Precision (accepted correct)")
    plt.plot(ts, rec, marker="s", label="Recall (of all correct)")
    plt.plot(ts, cov, marker="^", label="Coverage (auto-passed)")
    plt.axhline(PRECISION_TARGET, color="gray", linestyle="--", linewidth=1,
                label=f"Precision target {PRECISION_TARGET}")
    if recommended is not None:
        plt.axvline(recommended, color="red", linestyle=":", linewidth=1.5,
                    label=f"Recommended t={recommended}")
    plt.xlabel("Top-1 confidence threshold")
    plt.ylabel("Score")
    plt.title("Pass-decision threshold sweep")
    plt.legend(loc="lower left", fontsize=8)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    SCREENSHOTS_DIR.mkdir(parents=True, exist_ok=True)
    plt.savefig(THRESHOLD_PLOT, dpi=120)
    plt.close()
    print(f"[cv.evaluate] wrote {THRESHOLD_PLOT}")


def plot_failures(failures: list[dict], classes: list[str]) -> None:
    if not failures:
        print("[cv.evaluate] no misclassifications - skipping failure gallery")
        return

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from PIL import Image

    n = min(8, len(failures))
    cols, rows = 4, 2
    fig, axes = plt.subplots(rows, cols, figsize=(14, 7))
    for ax in axes.flat:
        ax.axis("off")
    for ax, f in zip(axes.flat, failures[:n]):
        img = Image.open(f["path"]).convert("RGB")
        ax.imshow(img)
        ax.set_title(
            f"true: {f['true']}\npred: {f['pred']} ({f['confidence']*100:.0f}%)",
            fontsize=9,
        )
    fig.suptitle("Most confident misclassifications (test set)", fontsize=12)
    plt.tight_layout(rect=(0, 0, 1, 0.96))
    SCREENSHOTS_DIR.mkdir(parents=True, exist_ok=True)
    plt.savefig(FAILURE_PLOT, dpi=120)
    plt.close()
    print(f"[cv.evaluate] wrote {FAILURE_PLOT}")


@torch.no_grad()
def main() -> None:
    if not CV_MODEL_PATH.exists():
        raise FileNotFoundError(
            "CV model missing. Train first with 'python -m src.cv.train'."
        )
    checkpoint = torch.load(CV_MODEL_PATH, map_location="cpu")
    classes = checkpoint["classes"]
    temperature = checkpoint.get("temperature")
    model = build_model(checkpoint["model_name"], len(classes))
    model.load_state_dict(checkpoint["state_dict"])
    model.eval()

    _, eval_tf = build_transforms()
    test_ds = ImageFolder(PROCESSED_DIR / "cv" / "test", transform=eval_tf)
    samples = test_ds.samples  # list of (path, class_idx)
    loader = DataLoader(test_ds, batch_size=64, shuffle=False, num_workers=2)

    y_true: list[int] = []
    y_pred: list[int] = []
    confidences: list[float] = []
    top5_hits = 0
    for x, y in loader:
        logits = _apply_temperature(model(x), temperature)
        probs = torch.softmax(logits, dim=1)
        conf1, pred1 = probs.topk(1, dim=1)
        _, pred5 = logits.topk(min(5, logits.size(1)), dim=1)
        y_true.extend(y.tolist())
        y_pred.extend(pred1.squeeze(1).tolist())
        confidences.extend(conf1.squeeze(1).tolist())
        top5_hits += pred5.eq(y.unsqueeze(1)).any(dim=1).sum().item()

    y_true_arr = np.array(y_true)
    y_pred_arr = np.array(y_pred)
    conf_arr = np.array(confidences)
    correct = (y_true_arr == y_pred_arr).astype(int)

    top1 = float(correct.mean())
    top5 = top5_hits / max(len(y_true), 1)
    print(f"[cv.evaluate] test top1={top1:.3f}  top5={top5:.3f}  n={len(y_true)}"
          f"  (temperature={temperature})")

    cm = confusion_matrix(y_true, y_pred, labels=list(range(len(classes))))
    report = classification_report(
        y_true, y_pred, target_names=classes, output_dict=True, zero_division=0
    )

    CONFUSION_PATH.write_text(
        json.dumps(
            {"classes": classes, "matrix": cm.tolist(), "top1": top1, "top5": top5},
            indent=2,
        )
    )
    print(f"[cv.evaluate] wrote {CONFUSION_PATH}")

    # Q3 threshold sweep
    curve, recommended = threshold_sweep(conf_arr, correct)
    print(f"[cv.evaluate] recommended threshold (precision>={PRECISION_TARGET}): "
          f"{recommended}")
    plot_threshold_curve(curve, recommended)

    # Q4 failure gallery: most confident wrong predictions
    failures = [
        {
            "path": samples[i][0],
            "true": classes[y_true[i]],
            "pred": classes[y_pred[i]],
            "confidence": float(conf_arr[i]),
        }
        for i in range(len(y_true))
        if not correct[i]
    ]
    failures.sort(key=lambda f: f["confidence"], reverse=True)
    plot_failures(failures, classes)

    existing = json.loads(CV_METRICS_PATH.read_text()) if CV_METRICS_PATH.exists() else {}
    existing["test_overall"] = {"top1": top1, "top5": top5, "n": len(y_true)}
    existing["per_class"] = {
        cls: {"precision": v["precision"], "recall": v["recall"], "f1": v["f1-score"]}
        for cls, v in report.items()
        if cls in classes
    }
    existing["threshold_curve"] = {
        "precision_target": PRECISION_TARGET,
        "recommended_threshold": recommended,
        "points": curve,
    }
    existing["top_failures"] = [
        {"true": f["true"], "pred": f["pred"], "confidence": round(f["confidence"], 3)}
        for f in failures[:8]
    ]
    CV_METRICS_PATH.write_text(json.dumps(existing, indent=2))
    print(f"[cv.evaluate] updated {CV_METRICS_PATH}")


if __name__ == "__main__":
    main()