experiments/data_scaling_study/train.py

"""
Unified trainer for the data-scaling study.

Usage:
    python train.py --model unet         --share 25
    python train.py --model unet         --share 50
    python train.py --model segformer_b0 --share 25
    python train.py --model segformer_b0 --share 50

The 100% runs are not trained here — they are bootstrapped from the existing
checkpoints in pv_panel_models/{unet,vit}_model/. See bootstrap_100.py.

Each run:
  - reads the matching subset_{share}.txt for training filenames
  - validates on the full val set every epoch
  - logs per-epoch metrics (loss, dice, iou, miou, pixel_acc) to logs/{model}_{share}.json
  - saves a single checkpoint at the highest val Dice:
        checkpoints/{model}_{share}_best.pth

Hyperparameters mirror the existing pv_panel_models/{unet,vit}_model trainers:
    Adam, lr=1e-4, ReduceLROnPlateau(mode='max', patience=5, factor=0.5),
    50 epochs, batch_size=16, image_size=128, CombinedLoss(0.5·BCE + 0.5·Dice),
    augmentations: HFlip, VFlip, Rot15.
"""
import argparse
import json
import os
import time
from pathlib import Path

import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm

from dataset import SubsetSolarPanelDataset
from metrics import SegMetrics
from models import MODEL_REGISTRY

THIS_DIR = Path(__file__).resolve().parent
REPO_ROOT = THIS_DIR.parents[1]

TRAIN_IMG = REPO_ROOT / "final_data" / "train" / "images"
TRAIN_MSK = REPO_ROOT / "final_data" / "train" / "masks"
VAL_IMG = REPO_ROOT / "final_data" / "val" / "images"
VAL_MSK = REPO_ROOT / "final_data" / "val" / "masks"

SUBSETS_DIR = THIS_DIR / "subsets"
LOG_DIR = THIS_DIR / "logs"
CKPT_DIR = THIS_DIR / "checkpoints"


def run_epoch(model, loader, criterion, optimizer, device, train: bool):
    model.train(mode=train)
    metrics = SegMetrics()
    total_loss = 0.0
    n_batches = 0

    desc = "Train" if train else "Val"
    ctx = torch.enable_grad() if train else torch.no_grad()
    with ctx:
        for images, masks in tqdm(loader, desc=desc, leave=False):
            images = images.to(device, non_blocking=True)
            masks = masks.to(device, non_blocking=True)

            if train:
                optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            if train:
                loss.backward()
                optimizer.step()

            total_loss += loss.item()
            n_batches += 1
            metrics.update(outputs.detach(), masks)

    avg_loss = total_loss / max(n_batches, 1)
    return avg_loss, metrics.compute()


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True, choices=list(MODEL_REGISTRY.keys()))
    p.add_argument("--share", required=True, type=int, choices=[25, 50])
    p.add_argument("--epochs", type=int, default=50)
    p.add_argument("--batch-size", type=int, default=16)
    p.add_argument("--image-size", type=int, default=128)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def main():
    args = parse_args()

    torch.manual_seed(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[run] model={args.model}  share={args.share}%  device={device}")

    LOG_DIR.mkdir(parents=True, exist_ok=True)
    CKPT_DIR.mkdir(parents=True, exist_ok=True)

    subset_file = SUBSETS_DIR / f"subset_{args.share}.txt"
    if not subset_file.is_file():
        raise FileNotFoundError(
            f"{subset_file} not found. Run subsets/make_subsets.py first."
        )

    train_set = SubsetSolarPanelDataset(
        TRAIN_IMG, TRAIN_MSK,
        file_list=subset_file,
        image_size=args.image_size,
        augment=True,
    )
    val_set = SubsetSolarPanelDataset(
        VAL_IMG, VAL_MSK,
        file_list=None,
        image_size=args.image_size,
        augment=False,
    )

    print(f"[data] train={len(train_set)}  val={len(val_set)}")

    train_loader = DataLoader(
        train_set, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, pin_memory=True,
    )
    val_loader = DataLoader(
        val_set, batch_size=args.batch_size, shuffle=False,
        num_workers=args.num_workers, pin_memory=True,
    )

    model_fn = MODEL_REGISTRY[args.model]
    model, criterion = model_fn()
    model = model.to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[model] {args.model}  params={n_params:,}")

    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="max", patience=5, factor=0.5
    )

    history = {
        "model": args.model,
        "share": args.share,
        "n_train": len(train_set),
        "n_val": len(val_set),
        "epochs": [],
    }

    best_dice = -1.0
    best_epoch = -1
    best_path = CKPT_DIR / f"{args.model}_{args.share}_best.pth"
    log_path = LOG_DIR / f"{args.model}_{args.share}.json"

    def _fmt(seconds: float) -> str:
        seconds = int(round(seconds))
        h, rem = divmod(seconds, 3600)
        m, s = divmod(rem, 60)
        return f"{h:d}:{m:02d}:{s:02d}" if h else f"{m:d}:{s:02d}"

    t0 = time.time()
    history["start_time_iso"] = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime(t0))

    for epoch in range(args.epochs):
        print(f"\nEpoch {epoch + 1}/{args.epochs}")
        epoch_t0 = time.time()

        train_t0 = time.time()
        train_loss, train_m = run_epoch(model, train_loader, criterion, optimizer, device, train=True)
        train_seconds = time.time() - train_t0

        val_t0 = time.time()
        val_loss, val_m = run_epoch(model, val_loader, criterion, optimizer, device, train=False)
        val_seconds = time.time() - val_t0

        scheduler.step(val_m["dice"])

        epoch_seconds = time.time() - epoch_t0
        elapsed = time.time() - t0
        avg_per_epoch = elapsed / (epoch + 1)
        eta = avg_per_epoch * (args.epochs - epoch - 1)

        epoch_record = {
            "epoch": epoch + 1,
            "lr": optimizer.param_groups[0]["lr"],
            "train_loss": train_loss,
            "val_loss": val_loss,
            **{f"train_{k}": v for k, v in train_m.items()},
            **{f"val_{k}": v for k, v in val_m.items()},
            "epoch_seconds": epoch_seconds,
            "train_seconds": train_seconds,
            "val_seconds": val_seconds,
        }
        history["epochs"].append(epoch_record)

        print(
            f"  train  loss={train_loss:.4f}  dice={train_m['dice']:.4f}  "
            f"iou={train_m['iou']:.4f}  miou={train_m['miou']:.4f}  "
            f"pixel_acc={train_m['pixel_acc']:.4f}"
        )
        print(
            f"  val    loss={val_loss:.4f}  dice={val_m['dice']:.4f}  "
            f"iou={val_m['iou']:.4f}  miou={val_m['miou']:.4f}  "
            f"pixel_acc={val_m['pixel_acc']:.4f}"
        )
        print(
            f"  time   epoch={_fmt(epoch_seconds)}  "
            f"(train={_fmt(train_seconds)}  val={_fmt(val_seconds)})  "
            f"elapsed={_fmt(elapsed)}  ETA={_fmt(eta)}"
        )

        # Save logs every epoch (resilient to crashes / SIGINT).
        with open(log_path, "w") as f:
            json.dump(history, f, indent=2)

        if val_m["dice"] > best_dice:
            best_dice = val_m["dice"]
            best_epoch = epoch + 1
            torch.save({
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "val_metrics": val_m,
                "model_name": args.model,
                "share": args.share,
            }, best_path)
            print(f"  ↳ new best (dice={best_dice:.4f}) → {best_path.name}")

    total_seconds = time.time() - t0
    history["best_epoch"] = best_epoch
    history["best_val_dice"] = best_dice
    history["wall_clock_seconds"] = total_seconds
    history["end_time_iso"] = time.strftime("%Y-%m-%dT%H:%M:%S")
    with open(log_path, "w") as f:
        json.dump(history, f, indent=2)

    print(f"\n[done] best epoch {best_epoch} (dice={best_dice:.4f})")
    print(f"       wall   {_fmt(total_seconds)}  ({total_seconds:.1f} s)")
    print(f"       best   → {best_path}")
    print(f"       log    → {log_path}")


if __name__ == "__main__":
    main()