scripts/train_unified.py

"""
Unified MVTec training loop with tqdm, AMP, early stopping, and resumable checkpoints.
"""

import os
import sys
import time
from typing import Any, Dict, Optional, Tuple

import torch
from torch.amp import GradScaler, autocast
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm.auto import tqdm

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

from models.model import OmniAL
from scripts.evaluate import evaluate_dataset
from utils.checkpointing import (
    append_epoch_metrics,
    ensure_run_dir,
    find_latest_checkpoint,
    load_checkpoint,
    restore_scheduler,
    save_checkpoint,
    save_metrics_json,
)
from utils.metrics import OmniALLoss


def build_model_and_optimizer(cfg: dict, device: str = None):
    device = device or cfg["train"]["device"]
    device = device if torch.cuda.is_available() else "cpu"

    model = OmniAL(cfg).to(device)
    tcfg = cfg["train"]
    optimizer = Adam(
        model.parameters(),
        lr=float(tcfg.get("lr", 1e-4)),
        weight_decay=float(tcfg.get("weight_decay", 1e-4)),
    )
    criterion = OmniALLoss(cfg)
    return model, optimizer, criterion, device


def build_scheduler(optimizer, cfg: dict):
    """ReduceLROnPlateau on the eval monitor (mode=max). Used by unified and separate."""
    tcfg = cfg.get("train", {})
    if str(tcfg.get("scheduler", "plateau")).lower() in ("none", "off", "false"):
        return None
    return ReduceLROnPlateau(
        optimizer,
        mode=tcfg.get("scheduler_mode", "max"),
        factor=float(tcfg.get("scheduler_factor", 0.5)),
        patience=int(tcfg.get("scheduler_patience", 3)),
        min_lr=float(tcfg.get("scheduler_min_lr", 1e-7)),
    )


def run_one_epoch(
    model,
    dataloader,
    optimizer,
    criterion,
    device: str,
    train: bool = True,
    use_amp: bool = False,
    scaler: Optional[GradScaler] = None,
    desc: str = "train",
) -> Dict[str, float]:
    model.train() if train else model.eval()
    epoch_losses = {
        "total": 0.0,
        "l2_img": 0.0,
        "ssim_img": 0.0,
        "l2_jnd": 0.0,
        "ssim_jnd": 0.0,
        "focal": 0.0,
    }
    n_batches = 0
    use_amp = bool(use_amp and device != "cpu" and torch.cuda.is_available())
    amp_device = "cuda" if (device != "cpu" and torch.cuda.is_available()) else "cpu"
    if scaler is None:
        scaler = GradScaler("cuda", enabled=use_amp) if amp_device == "cuda" else GradScaler(enabled=False)

    context = torch.enable_grad() if train else torch.no_grad()
    pbar = tqdm(dataloader, desc=desc, leave=False, dynamic_ncols=True)

    with context:
        for batch in pbar:
            anomaly_img = batch["anomaly_img"].to(device, non_blocking=True)
            targets = {
                "image": batch["image"].to(device, non_blocking=True),
                "jnd": batch["jnd"].to(device, non_blocking=True),
                "severity": batch["severity"].to(device, non_blocking=True),
            }

            with autocast(amp_device, enabled=use_amp):
                outputs = model(anomaly_img)
                losses = criterion(outputs, targets)

            if train:
                loss = losses["total"]
                if not torch.isfinite(loss):
                    print(f"[WARN] non-finite loss={loss.item()}; skipping step", flush=True)
                    optimizer.zero_grad(set_to_none=True)
                    continue
                optimizer.zero_grad(set_to_none=True)
                scaler.scale(loss).backward()
                # Unscale then clip to reduce AMP NaNs
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                scaler.step(optimizer)
                scaler.update()

            for k in epoch_losses:
                v = losses[k]
                epoch_losses[k] += v.item() if torch.is_tensor(v) else v
            n_batches += 1
            pbar.set_postfix(loss=f"{losses['total'].item():.4f}")

            del anomaly_img, targets, outputs, losses

    for k in epoch_losses:
        epoch_losses[k] /= max(n_batches, 1)
    return epoch_losses


def _monitor_value(eval_results: Dict[str, Any], monitor: str) -> float:
    if not eval_results or "average" not in eval_results:
        return float("-inf")
    avg = eval_results["average"]
    key_map = {
        "avg_I-AUROC": "I-AUROC",
        "avg_P-AUROC": "P-AUROC",
        "avg_P-AP": "P-AP",
        "I-AUROC": "I-AUROC",
        "P-AUROC": "P-AUROC",
        "P-AP": "P-AP",
    }
    k = key_map.get(monitor, monitor)
    return float(avg.get(k, float("-inf")))


def _write_latest_per_category(run_dir: str, epoch: int, per_cat: Dict[str, Any], average: Dict[str, float]) -> None:
    """Overwrite a human-readable table for the latest epoch only."""
    path = os.path.join(run_dir, "metrics", "latest_per_category.txt")
    lines = [f"Epoch {epoch} per-category results", "-" * 56]
    for cat in sorted(per_cat.keys()):
        m = per_cat[cat]
        lines.append(
            f"{cat:15s} | I-AUROC={m['I-AUROC']:.3f}  P-AUROC={m['P-AUROC']:.3f}  P-AP={m['P-AP']:.3f}"
        )
    if average:
        lines.append("-" * 56)
        lines.append(
            f"{'AVERAGE':15s} | I-AUROC={average.get('I-AUROC', 0):.3f}  "
            f"P-AUROC={average.get('P-AUROC', 0):.3f}  P-AP={average.get('P-AP', 0):.3f}"
        )
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")


def train_unified(
    model,
    train_loader,
    optimizer,
    criterion,
    device: str,
    cfg: dict,
    run_dir: str,
    start_epoch: int = 1,
    history: Optional[Dict[str, Any]] = None,
    scaler: Optional[GradScaler] = None,
    scheduler=None,
    best_metric: float = float("-inf"),
    hf_repo_id: Optional[str] = None,
    hf_token: Optional[str] = None,
    hf_run_name: Optional[str] = None,
) -> Tuple[Dict[str, Any], int, float]:
    """
    Train unified OmniAL on MVTec with per-epoch eval, checkpointing, and early stopping.
    Returns (history, last_epoch, best_metric).
    """
    run_dir = str(ensure_run_dir(run_dir))
    ckpt_dir = os.path.join(run_dir, "checkpoints")
    max_epochs = int(cfg["train"].get("max_epochs", cfg["train"].get("epochs", 400)))
    patience = int(cfg["train"].get("early_stopping_patience", 10))
    monitor = cfg["train"].get("early_stopping_monitor", "avg_I-AUROC")
    save_every = int(cfg["logging"].get("save_every", 2))
    hf_push_every = int(cfg["logging"].get("hf_push_every", 0))
    use_amp = bool(cfg["train"].get("use_amp", True))
    dataset_name = cfg["train"].get("dataset_name", "mvtec")

    if scaler is None:
        amp_on = use_amp and device != "cpu" and torch.cuda.is_available()
        scaler = GradScaler("cuda", enabled=amp_on) if amp_on else GradScaler(enabled=False)
    if scheduler is None:
        scheduler = build_scheduler(optimizer, cfg)

    if history is None:
        history = {
            "train": [],
            "eval": [],
            "per_category": [],
        }

    epochs_without_improvement = 0
    last_per_category = {}

    print(
        f"Run dir: {run_dir} | epochs {start_epoch}-{max_epochs} | "
        f"batch={train_loader.batch_size} | amp={use_amp} | patience={patience} | "
        f"monitor={monitor} | wd={cfg['train'].get('weight_decay', 1e-4)} | "
        f"lr={optimizer.param_groups[0]['lr']}"
    )
    print("Per-category metrics are written to metrics/ (not printed each epoch).")

    for epoch in range(start_epoch, max_epochs + 1):
        t0 = time.time()
        train_losses = run_one_epoch(
            model, train_loader, optimizer, criterion, device,
            train=True, use_amp=use_amp, scaler=scaler,
            desc=f"Epoch {epoch}/{max_epochs} train",
        )
        history["train"].append({"epoch": epoch, **train_losses})

        # Quiet eval: no duplicate per-category spam in the training log
        eval_results = evaluate_dataset(
            model, cfg, dataset_name=dataset_name, device=device, verbose=False
        )
        monitor_val = _monitor_value(eval_results, monitor)
        elapsed = time.time() - t0
        if scheduler is not None:
            scheduler.step(monitor_val)
        current_lr = optimizer.param_groups[0]["lr"]

        per_cat = {k: v for k, v in eval_results.items() if k != "average"}
        last_per_category = per_cat
        epoch_record = {
            "epoch": epoch,
            "train": train_losses,
            "eval_average": eval_results.get("average", {}),
            "eval_per_category": per_cat,
            "monitor": monitor,
            "monitor_value": monitor_val,
            "lr": current_lr,
            "elapsed_sec": round(elapsed, 2),
        }
        history["eval"].append(epoch_record["eval_average"])
        history["per_category"].append({
            "epoch": epoch,
            "categories": per_cat,
        })
        append_epoch_metrics(run_dir, epoch_record)
        save_metrics_json(run_dir, history)
        # Overwrite a single "latest" table each epoch for easy inspection
        _write_latest_per_category(run_dir, epoch, per_cat, epoch_record["eval_average"])

        # One-line epoch history only (stays in output for all epochs)
        print(
            f"[Epoch {epoch}/{max_epochs}] "
            f"train_total={train_losses['total']:.4f} | "
            f"I-AUROC={epoch_record['eval_average'].get('I-AUROC', 0):.4f} | "
            f"P-AUROC={epoch_record['eval_average'].get('P-AUROC', 0):.4f} | "
            f"P-AP={epoch_record['eval_average'].get('P-AP', 0):.4f} | "
            f"{monitor}={monitor_val:.4f} | lr={current_lr:.2e} | {elapsed:.1f}s",
            flush=True,
        )

        improved = monitor_val > best_metric
        if improved:
            best_metric = monitor_val
            epochs_without_improvement = 0
            save_checkpoint(
                os.path.join(ckpt_dir, "best.pt"),
                epoch, model, optimizer, scaler, history, best_metric, cfg, tag="best",
                scheduler=scheduler,
            )
            print(f"  -> new best {monitor}={best_metric:.4f} (saved best.pt)", flush=True)
        else:
            epochs_without_improvement += 1
            print(f"  -> no improvement ({epochs_without_improvement}/{patience})", flush=True)
        save_checkpoint(
            os.path.join(ckpt_dir, "last.pt"),
            epoch, model, optimizer, scaler, history, best_metric, cfg, tag="last",
            scheduler=scheduler,
        )

        if epoch % save_every == 0:
            ep_path = os.path.join(ckpt_dir, f"epoch_{epoch:04d}.pt")
            save_checkpoint(
                ep_path, epoch, model, optimizer, scaler, history, best_metric, cfg,
                tag=f"epoch_{epoch}", scheduler=scheduler,
            )
            print(f"  -> checkpoint saved: {ep_path}")

        if (
            hf_repo_id
            and hf_push_every > 0
            and epoch % hf_push_every == 0
        ):
            try:
                from utils.hf_sync import push_folder_to_hub
                path_in_repo = f"runs/{hf_run_name or os.path.basename(run_dir)}"
                url = push_folder_to_hub(
                    run_dir, hf_repo_id, path_in_repo=path_in_repo,
                    token=hf_token,
                    commit_message=f"Checkpoint epoch {epoch}",
                )
                print(f"  -> HF push OK: {url}")
            except Exception as e:
                print(f"  -> HF push failed: {e}")

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if epochs_without_improvement >= patience:
            print(f"Early stopping at epoch {epoch} (no {monitor} improvement for {patience} epochs).")
            break

    # Print final per-category table once (keeps epoch one-liners clean above)
    if last_per_category:
        print("\nFinal per-category results:")
        for cat, m in sorted(last_per_category.items()):
            print(f"  {cat:15s} I-AUROC={m['I-AUROC']:.3f}  P-AUROC={m['P-AUROC']:.3f}  P-AP={m['P-AP']:.3f}")

    return history, epoch, best_metric


def resume_from_run(
    model,
    optimizer,
    scaler,
    device: str,
    run_dir: str,
    checkpoint_path: Optional[str] = None,
    scheduler=None,
) -> Tuple[int, Dict[str, Any], float]:
    """Load last/best checkpoint if present. Returns (start_epoch, history, best_metric)."""
    path = checkpoint_path or find_latest_checkpoint(run_dir)
    if not path or not os.path.isfile(path):
        return 1, None, float("-inf")

    ckpt = load_checkpoint(path, model, optimizer, scaler, device=device)
    restore_scheduler(scheduler, ckpt)
    start_epoch = int(ckpt["epoch"]) + 1
    history = ckpt.get("history")
    best_metric = float(ckpt.get("best_metric", float("-inf")))
    lr = optimizer.param_groups[0]["lr"]
    print(f"Resumed from {path} -> start_epoch={start_epoch}, best_metric={best_metric:.4f}, lr={lr:.2e}")
    return start_epoch, history, best_metric