"""
Train the CLIP classifier head on cached embeddings.

Loads precomputed embeddings (from `scripts/precompute_embeddings.py`),
trains the head architecture from `clip_classifier.py`, and saves a
state_dict that can be loaded by `ClipClassifier.load_head_weights()`.

The head is intentionally tiny (~130k params), so training is laptop-fast
even on CPU — each epoch over 100k cached embeddings takes a few seconds.
That makes it practical to iterate on hyperparameters without re-encoding
the dataset.

Usage
-----
    python scripts/train_head.py \\
        --emb-dir data/embeddings \\
        --out data/checkpoints/head_v1.pt \\
        --epochs 30

Save a Stage 3A metrics report:

    python scripts/train_head.py \\
        --emb-dir data/embeddings \\
        --out data/checkpoints/head_v3a.pt \\
        --report-out data/reports/head_v3a_metrics.json \\
        --eval-split test_augmented
"""
from __future__ import annotations

import argparse
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np


@dataclass(frozen=True)
class SplitEmbeddings:
    x: Any
    y: Any
    paths: np.ndarray
    sources: np.ndarray
    generators: np.ndarray
    model_families: np.ndarray
    augmentations: np.ndarray
    original_paths: np.ndarray


# Mirror src/deepfake_scanner/detectors/clip_classifier.py:78-83 exactly.
# If that architecture changes, this MUST change too — checkpoints won't
# load otherwise.
def _build_head():
    import torch.nn as nn

    return nn.Sequential(
        nn.Linear(512, 256),
        nn.ReLU(),
        nn.Dropout(0.2),
        nn.Linear(256, 2),
    )


def _load_split(emb_dir: Path, name: str):
    import torch

    with np.load(emb_dir / f"{name}.npz") as z:
        embeddings = z["embeddings"].copy()
        labels = z["labels"].copy()
        n = len(labels)
        x = torch.from_numpy(embeddings).float()
        y = torch.from_numpy(labels).long()
        return SplitEmbeddings(
            x=x,
            y=y,
            paths=_optional_str_array(z, "paths", n),
            sources=_optional_str_array(z, "sources", n),
            generators=_optional_str_array(z, "generators", n),
            model_families=_optional_str_array(z, "model_families", n),
            augmentations=_optional_str_array(z, "augmentations", n),
            original_paths=_optional_str_array(z, "original_paths", n),
        )


def _optional_str_array(z, key: str, n: int) -> np.ndarray:
    if key not in z:
        return np.asarray([""] * n)
    values = np.asarray(z[key]).astype(str)
    if len(values) != n:
        raise ValueError(f"{key} has {len(values)} rows but labels has {n}")
    return values


def _accuracy(logits, y) -> float:
    return float((logits.argmax(dim=-1) == y).float().mean().item())


def _confusion(logits, y) -> dict:
    pred = logits.argmax(dim=-1)
    tp = int(((pred == 1) & (y == 1)).sum().item())
    tn = int(((pred == 0) & (y == 0)).sum().item())
    fp = int(((pred == 1) & (y == 0)).sum().item())
    fn = int(((pred == 0) & (y == 1)).sum().item())
    return {"tp": tp, "tn": tn, "fp": fp, "fn": fn}


def _rate(num: int, den: int) -> float:
    return float(num / den) if den else 0.0


def _metrics(logits, y, loss_fn=None) -> dict[str, Any]:
    cm = _confusion(logits, y)
    result: dict[str, Any] = {
        "n": int(len(y)),
        "accuracy": _accuracy(logits, y),
        "confusion": cm,
        "false_positive_rate": _rate(cm["fp"], cm["fp"] + cm["tn"]),
        "false_negative_rate": _rate(cm["fn"], cm["fn"] + cm["tp"]),
    }
    if loss_fn is not None:
        result["loss"] = float(loss_fn(logits, y).item())
    return result


def _group_metrics(logits, y, values: np.ndarray, loss_fn=None) -> dict[str, dict]:
    import torch

    result: dict[str, dict] = {}
    labels = sorted({str(v) for v in values if str(v)})
    for label in labels:
        idx = [i for i, value in enumerate(values) if str(value) == label]
        if not idx:
            continue
        tensor_idx = torch.as_tensor(idx, dtype=torch.long, device=logits.device)
        result[label] = _metrics(
            logits.index_select(0, tensor_idx),
            y.index_select(0, tensor_idx),
            loss_fn,
        )
    return result


def _augmentation_values(split: SplitEmbeddings) -> np.ndarray:
    values = np.asarray(split.augmentations).astype(str)
    return np.asarray([value if value else "clean" for value in values])


def _generator_values(split: SplitEmbeddings) -> np.ndarray:
    values: list[str] = []
    labels = split.y.cpu().numpy()
    for i, label in enumerate(labels):
        if label != 1:
            values.append("")
            continue
        generator = split.generators[i] or split.sources[i]
        values.append(str(generator))
    return np.asarray(values)


def _split_report(name: str, split: SplitEmbeddings, logits, loss_fn) -> dict[str, Any]:
    report: dict[str, Any] = {"overall": _metrics(logits, split.y, loss_fn)}

    by_source = _group_metrics(logits, split.y, split.sources, loss_fn)
    if by_source:
        report["by_source"] = by_source

    by_generator = _group_metrics(logits, split.y, _generator_values(split), loss_fn)
    if by_generator:
        report["by_generator"] = by_generator

    by_model_family = _group_metrics(logits, split.y, split.model_families, loss_fn)
    if by_model_family:
        report["by_model_family"] = by_model_family

    if any(str(value) for value in split.augmentations):
        report["by_augmentation"] = _group_metrics(
            logits,
            split.y,
            _augmentation_values(split),
            loss_fn,
        )

    print(
        f"\n{name}:  loss={report['overall']['loss']:.4f}  "
        f"acc={report['overall']['accuracy']:.4f}"
    )
    cm = report["overall"]["confusion"]
    print(
        f"  confusion matrix: tp={cm['tp']}  tn={cm['tn']}  "
        f"fp={cm['fp']}  fn={cm['fn']}"
    )
    for group_name in [
        "by_generator",
        "by_source",
        "by_model_family",
        "by_augmentation",
    ]:
        if group_name not in report:
            continue
        print(f"  {group_name}:")
        for label, metrics in report[group_name].items():
            print(
                f"    {label}: n={metrics['n']} "
                f"acc={metrics['accuracy']:.4f} "
                f"fpr={metrics['false_positive_rate']:.4f} "
                f"fnr={metrics['false_negative_rate']:.4f}"
            )

    return report


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--emb-dir", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--patience", type=int, default=5,
                        help="Early-stop after N epochs without val-loss improvement")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--eval-split",
        action="append",
        default=[],
        help=(
            "Additional embedding split name to evaluate, without .npz. "
            "Example: --eval-split test_augmented"
        ),
    )
    parser.add_argument(
        "--report-out",
        type=Path,
        default=None,
        help="Optional JSON path for overall and grouped metrics",
    )
    args = parser.parse_args()

    import torch
    import torch.nn as nn

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    print(f"Loading embeddings from {args.emb_dir}...")
    train = _load_split(args.emb_dir, "train")
    val = _load_split(args.emb_dir, "val")
    test = _load_split(args.emb_dir, "test")
    extra_evals = {
        name: _load_split(args.emb_dir, name)
        for name in args.eval_split
    }
    x_train, y_train = train.x, train.y
    x_val, y_val = val.x, val.y
    print(f"  train: {train.x.shape}, val: {val.x.shape}, test: {test.x.shape}")
    for name, split in extra_evals.items():
        print(f"  {name}: {split.x.shape}")

    cls_counts = torch.bincount(y_train, minlength=2)
    print(f"  train class counts: authentic={cls_counts[0].item()}  "
          f"ai_generated={cls_counts[1].item()}")

    head = _build_head()
    optimizer = torch.optim.Adam(
        head.parameters(), lr=args.lr, weight_decay=args.weight_decay,
    )
    loss_fn = nn.CrossEntropyLoss()

    best_val_loss = float("inf")
    best_state = None
    epochs_no_improve = 0

    for epoch in range(1, args.epochs + 1):
        t0 = time.time()
        head.train()
        perm = torch.randperm(len(x_train))
        train_losses: list[float] = []
        for i in range(0, len(x_train), args.batch_size):
            idx = perm[i : i + args.batch_size]
            xb, yb = x_train[idx], y_train[idx]
            optimizer.zero_grad()
            logits = head(xb)
            loss = loss_fn(logits, yb)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        head.eval()
        with torch.no_grad():
            val_logits = head(x_val)
            val_loss = loss_fn(val_logits, y_val).item()
            val_acc = _accuracy(val_logits, y_val)

        train_loss = float(np.mean(train_losses))
        dt = time.time() - t0
        print(
            f"  epoch {epoch:3d}  train_loss={train_loss:.4f}  "
            f"val_loss={val_loss:.4f}  val_acc={val_acc:.4f}  ({dt:.1f}s)"
        )

        if val_loss < best_val_loss - 1e-4:
            best_val_loss = val_loss
            best_state = {k: v.clone() for k, v in head.state_dict().items()}
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= args.patience:
                print(
                    f"  early stop at epoch {epoch} "
                    f"(no val-loss improvement for {args.patience} epochs)"
                )
                break

    if best_state is not None:
        head.load_state_dict(best_state)

    head.eval()
    with torch.no_grad():
        validation_logits = head(val.x)
        eval_logits = {
            "test": head(test.x),
            **{name: head(split.x) for name, split in extra_evals.items()},
        }

    report: dict[str, Any] = {
        "train": {
            "n": int(len(train.y)),
            "class_counts": {
                "authentic": int(cls_counts[0].item()),
                "ai_generated": int(cls_counts[1].item()),
            },
        },
        "validation": _split_report("validation", val, validation_logits, loss_fn),
        "evaluation": {
            "test": _split_report("test", test, eval_logits["test"], loss_fn),
        },
    }
    for name, split in extra_evals.items():
        report["evaluation"][name] = _split_report(
            name,
            split,
            eval_logits[name],
            loss_fn,
        )

    args.out.parent.mkdir(parents=True, exist_ok=True)
    torch.save(head.state_dict(), args.out)
    print(f"\nsaved head to {args.out}")
    if args.report_out is not None:
        args.report_out.parent.mkdir(parents=True, exist_ok=True)
        with args.report_out.open("w", encoding="utf-8") as fh:
            json.dump(report, fh, indent=2, sort_keys=True)
        print(f"saved metrics report to {args.report_out}")
    print("To use this checkpoint at inference time:")
    print("  from deepfake_scanner.detectors.clip_classifier import ClipClassifier")
    print("  c = ClipClassifier()")
    print(f"  c.load_head_weights('{args.out}')")


if __name__ == "__main__":
    main()