"""Fine-tune two pre-trained backbones on the Food-101 subset and compare them.

Models:
    - ResNet18      (richer features, larger)
    - MobileNetV3   (deployment-friendly, faster on CPU)

The best model (by validation top-1 accuracy) is persisted along with the
class list. Designed to run in ~10-20 minutes on a single GPU; falls back
to CPU automatically with a reduced default epoch count.

Usage:
    python -m src.cv.train --epochs 5 --batch 64
"""

from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import models, transforms
from torchvision.datasets import ImageFolder

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

from src.config import (  # noqa: E402
    CV_CLASSES_PATH,
    CV_METRICS_PATH,
    CV_MODEL_PATH,
    MODELS_DIR,
    PROCESSED_DIR,
)


CV_OUT = PROCESSED_DIR / "cv"

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

# Same fixed seed as the ML block (src/ml/train.py) for cross-block consistency.
RANDOM_SEED = 42


def set_seed(seed: int = RANDOM_SEED) -> None:
    """Make training as reproducible as the framework allows."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # warn_only: some torchvision ops have no deterministic CPU kernel; we keep
    # going instead of hard-failing so the run still completes.
    torch.use_deterministic_algorithms(True, warn_only=True)


def _seed_worker(worker_id: int) -> None:
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def build_transforms() -> tuple[transforms.Compose, transforms.Compose]:
    train_tf = transforms.Compose(
        [
            transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            transforms.ToTensor(),
            transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
        ]
    )
    eval_tf = transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
        ]
    )
    return train_tf, eval_tf


def build_model(name: str, num_classes: int) -> nn.Module:
    if name == "resnet18":
        model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
        return model
    if name == "mobilenet_v3_small":
        model = models.mobilenet_v3_small(
            weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1
        )
        in_feat = model.classifier[-1].in_features
        model.classifier[-1] = nn.Linear(in_feat, num_classes)
        return model
    raise ValueError(f"unknown model {name}")


@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader, device: torch.device) -> dict:
    model.eval()
    top1 = top5 = total = 0
    for x, y in loader:
        x = x.to(device)
        y = y.to(device)
        logits = model(x)
        _, pred1 = logits.topk(1, dim=1)
        _, pred5 = logits.topk(min(5, logits.size(1)), dim=1)
        top1 += (pred1.squeeze(1) == y).sum().item()
        top5 += pred5.eq(y.unsqueeze(1)).any(dim=1).sum().item()
        total += y.size(0)
    return {"top1": top1 / max(total, 1), "top5": top5 / max(total, 1), "n": total}


def train_one(
    name: str,
    train_loader: DataLoader,
    val_loader: DataLoader,
    num_classes: int,
    epochs: int,
    device: torch.device,
    lr: float = 3e-4,
) -> tuple[nn.Module, dict]:
    model = build_model(name, num_classes).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)

    best_val = 0.0
    best_state = None
    history: list[dict] = []
    for epoch in range(1, epochs + 1):
        model.train()
        start = time.time()
        running = 0.0
        seen = 0
        for x, y in train_loader:
            x = x.to(device)
            y = y.to(device)
            optimizer.zero_grad()
            logits = model(x)
            loss = criterion(logits, y)
            loss.backward()
            optimizer.step()
            running += loss.item() * x.size(0)
            seen += x.size(0)
        train_loss = running / max(seen, 1)
        val = evaluate(model, val_loader, device)
        history.append(
            {
                "epoch": epoch,
                "train_loss": train_loss,
                "val_top1": val["top1"],
                "val_top5": val["top5"],
                "duration_s": time.time() - start,
            }
        )
        print(
            f"[{name}] epoch {epoch:>2}  loss={train_loss:.3f}  "
            f"val_top1={val['top1']:.3f}  val_top5={val['top5']:.3f}"
        )
        if val["top1"] > best_val:
            best_val = val["top1"]
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}

    if best_state is not None:
        model.load_state_dict(best_state)
    return model, {"name": name, "best_val_top1": best_val, "history": history}


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--epochs", type=int, default=4)
    p.add_argument("--batch", type=int, default=64)
    p.add_argument("--workers", type=int, default=2)
    args = p.parse_args()

    set_seed(RANDOM_SEED)
    print(f"[cv.train] seed: {RANDOM_SEED}")

    train_dir = CV_OUT / "train"
    val_dir = CV_OUT / "val"
    test_dir = CV_OUT / "test"
    if not train_dir.exists():
        raise FileNotFoundError(
            "Training images missing. Run 'python -m src.cv.prepare_data' first."
        )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[cv.train] device: {device}")

    train_tf, eval_tf = build_transforms()
    train_ds = ImageFolder(train_dir, transform=train_tf)
    val_ds = ImageFolder(val_dir, transform=eval_tf)
    test_ds = ImageFolder(test_dir, transform=eval_tf)

    classes = train_ds.classes
    print(f"[cv.train] classes ({len(classes)}): {classes}")

    loader_generator = torch.Generator()
    loader_generator.manual_seed(RANDOM_SEED)
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch,
        shuffle=True,
        num_workers=args.workers,
        worker_init_fn=_seed_worker,
        generator=loader_generator,
    )
    val_loader = DataLoader(
        val_ds, batch_size=args.batch, shuffle=False, num_workers=args.workers
    )
    test_loader = DataLoader(
        test_ds, batch_size=args.batch, shuffle=False, num_workers=args.workers
    )

    results: list[dict] = []
    best_model: nn.Module | None = None
    best_name = ""
    best_metric = -1.0
    for name in ("resnet18", "mobilenet_v3_small"):
        model, hist = train_one(
            name, train_loader, val_loader, len(classes), args.epochs, device
        )
        test_metrics = evaluate(model, test_loader, device)
        record = {**hist, "test": test_metrics}
        results.append(record)
        print(
            f"[cv.train] {name} test_top1={test_metrics['top1']:.3f}"
            f" test_top5={test_metrics['top5']:.3f}"
        )
        if test_metrics["top1"] > best_metric:
            best_metric = test_metrics["top1"]
            best_model = model
            best_name = name

    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    assert best_model is not None
    torch.save(
        {"state_dict": best_model.state_dict(), "model_name": best_name, "classes": classes},
        CV_MODEL_PATH,
    )
    CV_CLASSES_PATH.write_text(json.dumps(classes, indent=2))
    CV_METRICS_PATH.write_text(
        json.dumps({"results": results, "best": best_name}, indent=2)
    )
    print(f"[cv.train] saved best model {best_name} -> {CV_MODEL_PATH}")


if __name__ == "__main__":
    main()