src/m2a_vision/train.py

"""
train.py — Training Script for M2a Vision Models
==================================================
Module: M2a Vision DL | MicroPlastiNet Pipeline
Author: MicroPlastiNet Team

USAGE
-----
# Classifier training (recommended first step):
python train.py --task classify --data_dir data/synthetic --epochs 10 --batch_size 32

# Detector training:
python train.py --task detect --data_dir data/synthetic --epochs 10 --batch_size 8

# Full run with all options:
python train.py \\
    --task classify \\
    --data_dir data/synthetic \\
    --epochs 20 \\
    --batch_size 32 \\
    --lr 1e-3 \\
    --checkpoint_dir checkpoints/ \\
    --log_dir runs/exp1 \\
    --freeze_backbone \\
    --amp

TRAINING STRATEGY
-----------------
Phase 1 (freeze backbone, epochs 1-5):
  Warm up only the classifier head with high LR.
  Prevents destroying pretrained features on a small dataset.

Phase 2 (unfreeze, epochs 6+):
  Fine-tune entire EfficientNet-B0 with cosine-annealed LR.

Mixed precision (--amp): Uses torch.cuda.amp for ~30% speedup on GPU.
  Falls back to float32 silently on CPU (no error, just slower).

EXPECTED ACCURACY (synthetic data — indicative only)
-----------------------------------------------------
  Classifier val accuracy after 10 epochs: ~65-75%
  Classifier val F1-macro: ~0.60-0.70
  Detector mAP@0.5 after 10 epochs: ~0.45-0.60

  With real Kaggle data (when available):
    Classifier: 75-85% | Detector mAP@0.5: 65-78%
    With UV fluorescence (MP-Set): +5-10% on both metrics

  Field-grade (camera alone): 60-70% accuracy
  Lab-grade (UV fluorescence augmentation): ~85% accuracy
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Dict, Tuple, Optional

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from torch.cuda.amp import GradScaler, autocast

# ── Local imports ──────────────────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))
from dataset import (
    get_detection_loaders, get_classification_loaders,
    SHAPE_CLASSES, generate_dataset
)
from model import (
    build_detector, build_classifier, YOLOLoss,
    NUM_CLASSES, TinyYOLO, MPClassifier
)


# ─────────────────────── Metrics Helpers ────────────────────────────────────

def classification_metrics(
    all_preds: list, all_labels: list
) -> Dict[str, float]:
    """
    Compute accuracy and macro-F1 for multi-class classification.

    Parameters
    ----------
    all_preds  : List of predicted class indices.
    all_labels : List of ground-truth class indices.

    Returns
    -------
    Dict with 'accuracy' and 'f1_macro'.
    """
    from sklearn.metrics import accuracy_score, f1_score
    acc = accuracy_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds, average="macro", zero_division=0)
    return {"accuracy": acc, "f1_macro": f1}


def simple_map_estimate(
    total_correct: int, total_pred: int, total_gt: int
) -> float:
    """
    Simplified mAP proxy: F1 of object detection (precision × recall harmonic mean).
    A crude but fast estimate during training; use evaluate.py for proper mAP.

    Parameters
    ----------
    total_correct : True positive detections.
    total_pred    : Total predictions made.
    total_gt      : Total ground-truth objects.

    Returns
    -------
    F1-score proxy for detection quality.
    """
    precision = total_correct / (total_pred + 1e-6)
    recall = total_correct / (total_gt + 1e-6)
    f1 = 2 * precision * recall / (precision + recall + 1e-6)
    return f1


# ─────────────────────── Classifier Training ────────────────────────────────

def train_classifier(args, device: torch.device) -> Dict:
    """
    Full training loop for MPClassifier (EfficientNet-B0).

    Returns dict with final metrics.
    """
    print("\n=== Training MPClassifier (EfficientNet-B0) ===")

    train_loader, val_loader = get_classification_loaders(
        data_dir=args.data_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        img_size=224,
    )
    print(f"Train samples: {len(train_loader.dataset)} | "
          f"Val samples: {len(val_loader.dataset)}")

    model = build_classifier(
        num_classes=NUM_CLASSES,
        pretrained=True,
        freeze_backbone=args.freeze_backbone,
    ).to(device)

    # ── Optimizer & Scheduler ────────────────────────────────────────────
    # Separate param groups: lower LR for backbone, higher for head
    backbone_params = list(model.features.parameters())
    head_params = list(model.classifier.parameters())
    optimizer = optim.AdamW([
        {"params": backbone_params, "lr": args.lr * 0.1},
        {"params": head_params,     "lr": args.lr},
    ], weight_decay=1e-4)

    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs, eta_min=1e-6)

    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
    scaler = GradScaler(enabled=(args.amp and device.type == "cuda"))

    writer = SummaryWriter(log_dir=os.path.join(args.log_dir, "classify"))
    os.makedirs(args.checkpoint_dir, exist_ok=True)

    best_acc = 0.0
    history = []

    for epoch in range(1, args.epochs + 1):
        # ── Training phase ───────────────────────────────────────────────
        model.train()

        # Unfreeze backbone after warm-up phase (default: after epoch 3)
        if epoch == args.unfreeze_epoch and args.freeze_backbone:
            print(f"  [Epoch {epoch}] Unfreezing backbone")
            for param in model.features.parameters():
                param.requires_grad = True

        train_loss = 0.0
        t0 = time.time()

        for batch_i, (images, labels) in enumerate(train_loader):
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)

            with autocast(enabled=(args.amp and device.type == "cuda")):
                logits = model(images)
                loss = criterion(logits, labels)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()

            train_loss += loss.item()

            if (batch_i + 1) % max(1, len(train_loader) // 4) == 0:
                print(f"  Epoch {epoch}/{args.epochs} "
                      f"[{batch_i+1}/{len(train_loader)}] "
                      f"loss={train_loss/(batch_i+1):.4f}")

        scheduler.step()
        avg_train_loss = train_loss / len(train_loader)

        # ── Validation phase ─────────────────────────────────────────────
        model.eval()
        val_loss = 0.0
        all_preds = []
        all_labels_list = []

        with torch.no_grad():
            for images, labels in val_loader:
                images = images.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                logits = model(images)
                loss = criterion(logits, labels)
                val_loss += loss.item()
                preds = logits.argmax(dim=1)
                all_preds.extend(preds.cpu().numpy().tolist())
                all_labels_list.extend(labels.cpu().numpy().tolist())

        avg_val_loss = val_loss / len(val_loader)
        metrics = classification_metrics(all_preds, all_labels_list)
        val_acc = metrics["accuracy"]
        val_f1 = metrics["f1_macro"]
        elapsed = time.time() - t0

        print(f"Epoch {epoch:02d}/{args.epochs:02d} | "
              f"train_loss={avg_train_loss:.4f} | "
              f"val_loss={avg_val_loss:.4f} | "
              f"val_acc={val_acc:.4f} | "
              f"val_F1={val_f1:.4f} | "
              f"time={elapsed:.1f}s")

        # TensorBoard
        writer.add_scalar("Loss/train", avg_train_loss, epoch)
        writer.add_scalar("Loss/val", avg_val_loss, epoch)
        writer.add_scalar("Metrics/accuracy", val_acc, epoch)
        writer.add_scalar("Metrics/f1_macro", val_f1, epoch)
        writer.add_scalar("LR", optimizer.param_groups[1]["lr"], epoch)

        history.append({
            "epoch": epoch, "train_loss": avg_train_loss,
            "val_loss": avg_val_loss, "val_acc": val_acc, "val_f1": val_f1
        })

        # ── Checkpoint ───────────────────────────────────────────────────
        if val_acc > best_acc:
            best_acc = val_acc
            ckpt_path = os.path.join(
                args.checkpoint_dir, "best_classifier.pt")
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_acc": val_acc,
                "val_f1": val_f1,
                "train_loss": avg_train_loss,
                "classes": SHAPE_CLASSES,
                "architecture": "EfficientNet-B0",
            }, ckpt_path)
            print(f"  ✓ Saved best checkpoint (val_acc={val_acc:.4f})")

    # Save last checkpoint
    torch.save({
        "epoch": args.epochs,
        "model_state_dict": model.state_dict(),
        "val_acc": val_acc,
        "val_f1": val_f1,
        "classes": SHAPE_CLASSES,
        "architecture": "EfficientNet-B0",
    }, os.path.join(args.checkpoint_dir, "last_classifier.pt"))

    writer.close()

    final_metrics = {
        "task": "classify",
        "best_val_acc": best_acc,
        "final_val_acc": history[-1]["val_acc"],
        "final_val_f1": history[-1]["val_f1"],
        "epochs": args.epochs,
        "history": history,
    }
    print(f"\nBest val accuracy: {best_acc:.4f}")
    return final_metrics


# ─────────────────────── Detector Training ──────────────────────────────────

def train_detector(args, device: torch.device) -> Dict:
    """
    Full training loop for TinyYOLO detector.

    Returns dict with final metrics.
    """
    print("\n=== Training TinyYOLO Detector ===")

    train_loader, val_loader = get_detection_loaders(
        data_dir=args.data_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        img_size=416,
    )
    print(f"Train batches: {len(train_loader)} | Val batches: {len(val_loader)}")

    model = build_detector(num_classes=NUM_CLASSES).to(device)
    criterion = YOLOLoss(num_classes=NUM_CLASSES)

    optimizer = optim.SGD(
        model.parameters(), lr=args.lr,
        momentum=0.937, weight_decay=5e-4, nesterov=True)

    scheduler = optim.lr_scheduler.OneCycleLR(
        optimizer,
        max_lr=args.lr,
        steps_per_epoch=len(train_loader),
        epochs=args.epochs,
        pct_start=0.1,
    )

    scaler = GradScaler(enabled=(args.amp and device.type == "cuda"))
    writer = SummaryWriter(log_dir=os.path.join(args.log_dir, "detect"))
    os.makedirs(args.checkpoint_dir, exist_ok=True)

    best_map = 0.0
    history = []

    for epoch in range(1, args.epochs + 1):
        model.train()
        epoch_loss = 0.0
        t0 = time.time()

        for batch_i, batch in enumerate(train_loader):
            images = batch["image"].to(device, non_blocking=True)
            boxes = [b.to(device) for b in batch["boxes"]]
            labels = [l.to(device) for l in batch["labels"]]

            optimizer.zero_grad(set_to_none=True)

            with autocast(enabled=(args.amp and device.type == "cuda")):
                preds = model(images)
                loss, loss_comps = criterion(preds, boxes, labels)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

            epoch_loss += loss.item()

        avg_loss = epoch_loss / len(train_loader)

        # ── Validation: fast mAP proxy ───────────────────────────────────
        model.eval()
        tp_count, pred_count, gt_count = 0, 0, 0
        val_loss = 0.0

        with torch.no_grad():
            for batch in val_loader:
                images = batch["image"].to(device)
                boxes = [b.to(device) for b in batch["boxes"]]
                labels = [l.to(device) for l in batch["labels"]]
                preds = model(images)
                loss, _ = criterion(preds, boxes, labels)
                val_loss += loss.item()
                for lbl in labels:
                    gt_count += lbl.shape[0]
                # Simple detection count (objects with high confidence)
                for pred_scale in preds:
                    conf = torch.sigmoid(pred_scale[..., 4])
                    pred_count += (conf > 0.5).sum().item()
                    tp_count += min((conf > 0.5).sum().item(), gt_count)

        avg_val_loss = val_loss / max(1, len(val_loader))
        map_proxy = simple_map_estimate(tp_count, pred_count, gt_count)
        elapsed = time.time() - t0

        print(f"Epoch {epoch:02d}/{args.epochs:02d} | "
              f"train_loss={avg_loss:.4f} | "
              f"val_loss={avg_val_loss:.4f} | "
              f"mAP_proxy={map_proxy:.4f} | "
              f"time={elapsed:.1f}s")

        writer.add_scalar("Loss/train", avg_loss, epoch)
        writer.add_scalar("Loss/val", avg_val_loss, epoch)
        writer.add_scalar("Metrics/mAP_proxy", map_proxy, epoch)

        history.append({
            "epoch": epoch, "train_loss": avg_loss,
            "val_loss": avg_val_loss, "map_proxy": map_proxy
        })

        if map_proxy > best_map:
            best_map = map_proxy
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "map_proxy": map_proxy,
                "train_loss": avg_loss,
                "classes": SHAPE_CLASSES,
                "architecture": "TinyYOLO",
            }, os.path.join(args.checkpoint_dir, "best_detector.pt"))
            print(f"  ✓ Saved best detector checkpoint (mAP_proxy={map_proxy:.4f})")

    torch.save({
        "epoch": args.epochs,
        "model_state_dict": model.state_dict(),
        "map_proxy": map_proxy,
        "classes": SHAPE_CLASSES,
        "architecture": "TinyYOLO",
    }, os.path.join(args.checkpoint_dir, "last_detector.pt"))

    writer.close()

    return {
        "task": "detect",
        "best_map_proxy": best_map,
        "final_map_proxy": history[-1]["map_proxy"],
        "epochs": args.epochs,
        "history": history,
    }


# ─────────────────────────────── CLI ────────────────────────────────────────

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train M2a Vision models for microplastic detection",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--task", choices=["classify", "detect", "both"],
                        default="classify",
                        help="Which model to train")
    parser.add_argument("--data_dir", default="data/synthetic",
                        help="Root dataset directory (train/ and val/ subdirs)")
    parser.add_argument("--generate_data", action="store_true",
                        help="Generate synthetic dataset before training")
    parser.add_argument("--n_train", type=int, default=2000)
    parser.add_argument("--n_val", type=int, default=500)
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--num_workers", type=int, default=0)
    parser.add_argument("--checkpoint_dir", default="checkpoints",
                        help="Directory to save model checkpoints")
    parser.add_argument("--log_dir", default="runs",
                        help="TensorBoard log directory")
    parser.add_argument("--amp", action="store_true",
                        help="Enable mixed precision (GPU only)")
    parser.add_argument("--freeze_backbone", action="store_true",
                        help="Freeze EfficientNet backbone for warm-up")
    parser.add_argument("--unfreeze_epoch", type=int, default=3,
                        help="Epoch to unfreeze backbone (if --freeze_backbone)")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output_json", default=None,
                        help="Save final metrics JSON to this path")
    return parser.parse_args()


def main():
    args = parse_args()

    # Reproducibility
    torch.manual_seed(args.seed)
    import random, numpy as np
    random.seed(args.seed)
    np.random.seed(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    print(f"Task:   {args.task}")
    print(f"Data:   {args.data_dir}")
    print(f"Epochs: {args.epochs} | Batch: {args.batch_size} | LR: {args.lr}")

    # ── Optionally generate data ─────────────────────────────────────────
    if args.generate_data or not Path(args.data_dir).exists():
        print(f"\nGenerating synthetic dataset → {args.data_dir}")
        generate_dataset(
            out_dir=args.data_dir,
            n_train=args.n_train,
            n_val=args.n_val,
            seed=args.seed,
        )

    results = {}

    if args.task in ("classify", "both"):
        clf_metrics = train_classifier(args, device)
        results["classifier"] = clf_metrics

    if args.task in ("detect", "both"):
        det_metrics = train_detector(args, device)
        results["detector"] = det_metrics

    print("\n=== Training Complete ===")
    print(json.dumps({k: {kk: vv for kk, vv in v.items() if kk != "history"}
                      for k, v in results.items()}, indent=2))

    if args.output_json:
        os.makedirs(os.path.dirname(args.output_json) or ".", exist_ok=True)
        with open(args.output_json, "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nMetrics saved to {args.output_json}")

    return results


if __name__ == "__main__":
    main()