"""
train.py — train the U-Net on RadioMapSeer coverage maps.

v1: 2-channel input (buildings, transmitter).
v2: add --distance for a 3rd channel (distance-to-Tx) that fixes the square cutoff.

Examples:
  # v1
  python src\\train.py --epochs 30 --batch-size 32 --num-workers 8
  # v2 (separate checkpoint dir so v1 stays intact)
  python src\\train.py --distance --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v2
"""

import argparse
import csv
import time
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from dataset import RadioMapSeerDataset
from model import UNet

torch.backends.cudnn.benchmark = True   # autotune convs (input size is fixed)


def evaluate(net, loader, criterion, device, amp):
    net.eval()
    total, n = 0.0, 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
                loss = criterion(net(x), y)
            total += loss.item() * x.size(0)
            n += x.size(0)
    return total / n


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--epochs", type=int, default=30)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--lr", type=float, default=1e-3)
    p.add_argument("--img-size", type=int, default=256)
    p.add_argument("--base", type=int, default=64)
    p.add_argument("--max-samples", type=int, default=None)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--ckpt-dir", default="checkpoints")
    p.add_argument("--distance", action="store_true",
                   help="v2: add a distance-to-transmitter input channel")
    p.add_argument("--compile", action="store_true",
                   help="wrap model in torch.compile (optional; flaky on Windows)")
    args = p.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    amp = (device == "cuda")
    print(f"Device: {device} | AMP: {amp} | distance: {args.distance} | compile: {args.compile}")

    train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size,
                                   max_samples=args.max_samples, add_distance=args.distance)
    val_max = None if args.max_samples is None else max(20, args.max_samples // 5)
    val_ds = RadioMapSeerDataset(split="val", img_size=args.img_size,
                                 max_samples=val_max, add_distance=args.distance)
    print(f"Train: {len(train_ds)} | Val: {len(val_ds)}")

    dl_kwargs = dict(batch_size=args.batch_size, num_workers=args.num_workers,
                     pin_memory=True, persistent_workers=(args.num_workers > 0))
    train_loader = DataLoader(train_ds, shuffle=True, **dl_kwargs)
    val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs)

    in_ch = 3 if args.distance else 2
    model = UNet(in_channels=in_ch, out_channels=1, base=args.base).to(device)
    net = torch.compile(model) if args.compile else model
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    scaler = torch.amp.GradScaler(enabled=amp)

    ckpt_dir = Path(args.ckpt_dir)
    ckpt_dir.mkdir(exist_ok=True)
    best_val = float("inf")

    history_path = ckpt_dir / "history.csv"
    with open(history_path, "w", newline="") as f:
        csv.writer(f).writerow(["epoch", "train_mse", "val_mse"])

    for epoch in range(1, args.epochs + 1):
        net.train()
        t0 = time.time()
        running, n = 0.0, 0
        for i, (x, y) in enumerate(train_loader, 1):
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
                loss = criterion(net(x), y)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            running += loss.item() * x.size(0)
            n += x.size(0)
            if i % 50 == 0:
                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}")

        train_loss = running / n
        val_loss = evaluate(net, val_loader, criterion, device, amp)
        dt = time.time() - t0
        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.5f} | "
              f"val MSE {val_loss:.5f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")
        with open(history_path, "a", newline="") as f:
            csv.writer(f).writerow([epoch, f"{train_loss:.6f}", f"{val_loss:.6f}"])

        if val_loss < best_val:
            best_val = val_loss
            torch.save({"epoch": epoch, "model_state": model.state_dict(),
                        "val_loss": val_loss, "args": vars(args)},
                       ckpt_dir / "best.pt")
            print(f"           -> saved best.pt (val MSE {val_loss:.5f})")

        if device == "cuda":
            print(f"           peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
            torch.cuda.reset_peak_memory_stats()

    print(f"\nDone. Best val MSE: {best_val:.5f}  (RMSE {best_val**0.5:.5f})")


if __name__ == "__main__":
    main()