"""Computer-Vision block (BONUS): "is this document image a résumé?".

A recruiter often receives a résumé as a **scanned image / PDF page**, not clean
text. This block validates an uploaded document and gates it into the pipeline:
only pages that look like a résumé are admitted; their text is then handed to the
NLP→ML pipeline (see :mod:`src.vision` for the inference + OCR bridge).

Approach — **transfer learning** (the lecture's CV method), compared head-to-head:
1. **Majority baseline** (DummyClassifier) — the number to beat.
2. **ResNet18** (ImageNet) as a *frozen feature extractor* → LogisticRegression.
3. **MobileNet-V3-Small** (ImageNet) frozen features → LogisticRegression.

We freeze the CNN and train only a linear head — fast and reliable on CPU, and a
fair comparison of two backbones. Data: ``jordyvl/rvl_cdip_100_examples_per_class``
(a small, balanced slice of the classic RVL-CDIP document corpus; ``resume`` is one
of its 16 classes). We collapse the 16 classes to a binary **résumé vs other**.

Run:  .venv\\Scripts\\python.exe src/cv_model.py
CV deps (torch/torchvision) live in ``requirements-cv.txt`` — NOT in the Space.
"""
from __future__ import annotations

import json
import sys
from datetime import date
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:  # noqa: BLE001
    pass

ARTIFACTS = ROOT / "artifacts"
EVAL_DIR = ARTIFACTS / "eval"
RANDOM_STATE = 42
N_FOLDS = 5

CV_DATASET = "jordyvl/rvl_cdip_100_examples_per_class"
POSITIVE_NAME = "resume"
BACKBONES = ["resnet18", "mobilenet_v3_small"]
BATCH_SIZE = 32


# --------------------------------------------------------------------------- #
# Backbone (frozen ImageNet feature extractor) + feature extraction
# --------------------------------------------------------------------------- #
def load_backbone(name: str):
    """Return ``(model, preprocess, dim)`` for a frozen torchvision backbone."""
    import torch
    from torchvision import models

    if name == "resnet18":
        weights = models.ResNet18_Weights.IMAGENET1K_V1
        model = models.resnet18(weights=weights)
        model.fc = torch.nn.Identity()
        dim = 512
    elif name == "mobilenet_v3_small":
        weights = models.MobileNet_V3_Small_Weights.IMAGENET1K_V1
        model = models.mobilenet_v3_small(weights=weights)
        model.classifier = torch.nn.Identity()
        dim = 576
    else:
        raise ValueError(f"Unknown backbone: {name}")

    model.eval()
    for p in model.parameters():  # freeze: we only train the linear head
        p.requires_grad_(False)
    return model, weights.transforms(), dim


def extract_features(images, name: str, batch_size: int = BATCH_SIZE) -> np.ndarray:
    """Embed a list of PIL images with a frozen backbone (CPU, no grad)."""
    import torch

    model, preprocess, _ = load_backbone(name)
    feats, batch = [], []

    def _flush():
        if not batch:
            return
        with torch.no_grad():
            x = torch.stack(batch)
            feats.append(model(x).cpu().numpy())
        batch.clear()

    for img in images:
        if img.mode != "RGB":
            img = img.convert("RGB")  # ImageNet transforms expect 3 channels
        batch.append(preprocess(img))
        if len(batch) == batch_size:
            _flush()
    _flush()
    return np.vstack(feats)


# --------------------------------------------------------------------------- #
# Data → binary (résumé vs other)
# --------------------------------------------------------------------------- #
def _load_binary():
    from datasets import load_dataset

    ds = load_dataset(CV_DATASET)
    names = ds["train"].features["label"].names
    pos = names.index(POSITIVE_NAME)

    def to_xy(split: str):
        imgs = list(ds[split]["image"])
        y = ["resume" if lab == pos else "other" for lab in ds[split]["label"]]
        return imgs, y

    tr_i, tr_y = to_xy("train")
    va_i, va_y = to_xy("validation")
    te_i, te_y = to_xy("test")
    # Train on train+validation; evaluate on the held-out test split.
    return (tr_i + va_i, tr_y + va_y), (te_i, te_y), names


def _round(x, n: int = 4):
    return None if x is None else round(float(x), n)


# --------------------------------------------------------------------------- #
# Train + compare + persist
# --------------------------------------------------------------------------- #
def train_cv_classifier() -> dict:
    import joblib
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import seaborn as sns
    from sklearn.dummy import DummyClassifier
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import (
        accuracy_score,
        classification_report,
        confusion_matrix,
        f1_score,
    )
    from sklearn.model_selection import StratifiedKFold, cross_validate

    print("Loading document images (résumé vs other) ...")
    (train_imgs, ytr), (test_imgs, yte), class_names = _load_binary()
    n_pos_tr = sum(v == "resume" for v in ytr)
    n_pos_te = sum(v == "resume" for v in yte)
    print(f"  train={len(ytr)} (résumé={n_pos_tr}) | test={len(yte)} (résumé={n_pos_te})")
    print(f"  source: {CV_DATASET} (16 classes -> binary)")

    labels = ["other", "resume"]
    cv = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=RANDOM_STATE)

    # Baseline: predict the majority class regardless of pixels.
    base = DummyClassifier(strategy="most_frequent").fit(np.zeros((len(ytr), 1)), ytr)
    bpred = base.predict(np.zeros((len(yte), 1)))
    results = {
        "baseline_majority": {
            "test_accuracy": _round(accuracy_score(yte, bpred)),
            "test_f1_macro": _round(f1_score(yte, bpred, average="macro")),
        }
    }

    best = {"name": None, "f1": -1.0, "clf": None, "Xte": None, "pred": None}
    print(f"\nComparing {len(BACKBONES)} frozen backbones "
          f"(+ baseline), {N_FOLDS}-fold CV on the linear head ...\n")
    for name in BACKBONES:
        print(f"  [{name}] extracting features ...")
        Xtr = extract_features(train_imgs, name)
        Xte = extract_features(test_imgs, name)
        clf = LogisticRegression(max_iter=2000, class_weight="balanced",
                                 C=1.0, random_state=RANDOM_STATE)
        cvres = cross_validate(clf, Xtr, ytr, cv=cv,
                               scoring=["accuracy", "f1_macro"], n_jobs=-1)
        clf.fit(Xtr, ytr)
        pred = clf.predict(Xte)
        rec = {
            "feature_dim": int(Xtr.shape[1]),
            "cv_accuracy_mean": _round(cvres["test_accuracy"].mean()),
            "cv_f1_macro_mean": _round(cvres["test_f1_macro"].mean()),
            "cv_f1_macro_std": _round(cvres["test_f1_macro"].std()),
            "test_accuracy": _round(accuracy_score(yte, pred)),
            "test_f1_macro": _round(f1_score(yte, pred, average="macro")),
        }
        results[name] = rec
        print(f"      CV macro-F1={rec['cv_f1_macro_mean']:.3f} | "
              f"test acc={rec['test_accuracy']:.3f} macro-F1={rec['test_f1_macro']:.3f}")
        if rec["test_f1_macro"] > best["f1"]:
            best.update(name=name, f1=rec["test_f1_macro"], clf=clf, Xte=Xte, pred=pred)

    print(f"\nBest backbone by test macro-F1: {best['name']}")
    print("\nPer-class report (best on test):")
    print(classification_report(yte, best["pred"], labels=labels, digits=3, zero_division=0))

    # --- figures ----------------------------------------------------------- #
    EVAL_DIR.mkdir(parents=True, exist_ok=True)
    sns.set_theme(style="whitegrid")

    cm = confusion_matrix(yte, best["pred"], labels=labels)
    fig, ax = plt.subplots(figsize=(5.5, 4.8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Greens", cbar=False,
                xticklabels=labels, yticklabels=labels, ax=ax)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title(f"CV résumé detector — confusion ({best['name']}, test)")
    fig.tight_layout()
    fig.savefig(EVAL_DIR / "cv_confusion_matrix.png", dpi=120)
    plt.close(fig)

    names_cmp = [n for n in ["baseline_majority", *BACKBONES]]
    f1s = [results[n]["test_f1_macro"] for n in names_cmp]
    fig, ax = plt.subplots(figsize=(6.5, 4.2))
    ax.bar(names_cmp, f1s, color=["grey", "steelblue", "seagreen"])
    for i, v in enumerate(f1s):
        ax.text(i, v + 0.01, f"{v:.2f}", ha="center")
    ax.set_ylim(0, 1.05)
    ax.set_ylabel("Test macro-F1")
    ax.set_title("CV block — model comparison (résumé vs other)")
    plt.xticks(rotation=15, ha="right")
    fig.tight_layout()
    fig.savefig(EVAL_DIR / "cv_model_comparison.png", dpi=120)
    plt.close(fig)

    # --- persist artifacts ------------------------------------------------- #
    ARTIFACTS.mkdir(exist_ok=True)
    bundle = {
        "backbone": best["name"],
        "clf": best["clf"],
        "classes": list(best["clf"].classes_),
        "positive_class": "resume",
        "img_mode": "RGB",
    }
    joblib.dump(bundle, ARTIFACTS / "cv_model.joblib")

    metrics = {
        "task": "document_image_resume_vs_other",
        "data_source": CV_DATASET,
        "n_train": len(ytr), "n_test": len(yte),
        "n_resume_train": int(n_pos_tr), "n_resume_test": int(n_pos_te),
        "best_backbone": best["name"],
        "models": results,
        "test_confusion_matrix": {"labels": labels, "matrix": cm.tolist()},
        "method": "frozen ImageNet backbone -> LogisticRegression (transfer learning)",
    }
    (ARTIFACTS / "cv_metrics.json").write_text(
        json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8")
    meta = {
        "task": "document_image_resume_vs_other",
        "backbone": best["name"],
        "classes": list(best["clf"].classes_),
        "positive_class": "resume",
        "data_source": CV_DATASET,
        "trained_at": date.today().isoformat(),
    }
    (ARTIFACTS / "cv_metadata.json").write_text(
        json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")

    print(f"\nSaved -> {ARTIFACTS / 'cv_model.joblib'}")
    print(f"Saved -> {ARTIFACTS / 'cv_metrics.json'}")
    print(f"Saved -> {EVAL_DIR / 'cv_confusion_matrix.png'} (+ cv_model_comparison.png)")
    return metrics


if __name__ == "__main__":
    train_cv_classifier()
    print("\nDONE.")