"""
train.py -- Part B, Section 3.3 (model training + loss-function ablation).

Two arms of the ablation, identical in every other respect:
    python train.py --loss ce    --epochs 15
    python train.py --loss focal --epochs 15 --focal-alpha 0.25 --focal-gamma 2.0
    python train.py --loss weighted_ce --fg-weight 4.0 --epochs 15   # optional third arm

Model selection: best validation mAP@0.50 (not loss), with early stopping.
"""

from __future__ import annotations

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np
import torch

from common import (DEFAULT_IMG_SIZE, make_grad_scaler, autocast, DET_SPLIT_CSV, METRICS_DIR, PLOTS_DIR,
                    SEED, SEED_TAG, WEIGHTS_DIR, KvasirDetectionDataset, _mpl,
                    average_precision, coco_map, get_device, make_loader,
                    progress, save_json, set_seed)
from model import build_detector, parameter_report, set_roi_loss


def run_tag(args) -> str:
    t = f"det_{args.arch}_{args.loss}"
    if args.loss == "focal":
        t += f"-a{args.focal_alpha}g{args.focal_gamma}"
    if args.tag:
        t += f"_{args.tag}"
    return f"{t}_{SEED_TAG}"


# ---------------------------------------------------------------------------
@torch.no_grad()
def run_inference(model, loader, device):
    """Returns (preds, gts) in the flat dict format the metric helpers expect."""
    model.eval()
    preds, gts = [], []
    for images, targets in progress(loader, desc="infer"):
        images = [im.to(device) for im in images]
        outs = model(images)
        for o, t in zip(outs, targets):
            preds.append({"boxes": o["boxes"].detach().cpu().numpy(),
                          "scores": o["scores"].detach().cpu().numpy(),
                          "labels": o["labels"].detach().cpu().numpy()})
            gts.append({"boxes": t["boxes"].numpy(), "stem": t["stem"]})
    return preds, gts


def warmup_lr(optimizer, base_lr: float, step: int, warmup_steps: int):
    if warmup_steps <= 0 or step >= warmup_steps:
        return
    f = (step + 1) / warmup_steps
    for g in optimizer.param_groups:
        g["lr"] = base_lr * f


# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description="Kvasir-SEG detector training (seed 16)")
    ap.add_argument("--arch", default="fasterrcnn_resnet50_fpn",
                    choices=["fasterrcnn_resnet50_fpn", "retinanet_resnet50_fpn"])
    ap.add_argument("--loss", default="ce", choices=["ce", "focal", "weighted_ce"],
                    help="ROI-head classification loss (the ablation variable)")
    ap.add_argument("--focal-alpha", type=float, default=0.25)
    ap.add_argument("--focal-gamma", type=float, default=2.0)
    ap.add_argument("--fg-weight", type=float, default=4.0)
    ap.add_argument("--epochs", type=int, default=15)
    ap.add_argument("--bs", type=int, default=4)
    ap.add_argument("--lr", type=float, default=5e-3)
    ap.add_argument("--optimizer", default="sgd", choices=["sgd", "adamw"])
    ap.add_argument("--momentum", type=float, default=0.9)
    ap.add_argument("--wd", type=float, default=5e-4)
    ap.add_argument("--img-size", type=int, default=DEFAULT_IMG_SIZE)
    ap.add_argument("--trainable-layers", type=int, default=3,
                    help="ResNet stages left unfrozen (0-5)")
    ap.add_argument("--no-augment", action="store_true")
    ap.add_argument("--no-pretrained", action="store_true")
    ap.add_argument("--warmup-frac", type=float, default=0.1,
                    help="fraction of the first epoch used for linear LR warmup")
    ap.add_argument("--patience", type=int, default=5)
    ap.add_argument("--min-delta", type=float, default=1e-4)
    ap.add_argument("--clip-grad", type=float, default=10.0)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--amp", action="store_true", default=True)
    ap.add_argument("--tag", default="")
    args = ap.parse_args()

    set_seed(SEED)
    device = get_device()
    tag = run_tag(args)
    loss_desc = set_roi_loss(args.loss, alpha=args.focal_alpha,
                             gamma=args.focal_gamma, fg_weight=args.fg_weight)
    print(json.dumps({"run": tag, "seed": SEED, "device": str(device),
                      "loss": loss_desc, **vars(args)}, indent=2))

    ds_tr = KvasirDetectionDataset(DET_SPLIT_CSV, "train", augment=not args.no_augment,
                                   img_size=args.img_size)
    ds_va = KvasirDetectionDataset(DET_SPLIT_CSV, "val", augment=False,
                                   img_size=args.img_size)
    print(f"[data] train={len(ds_tr)} val={len(ds_va)} img_size={args.img_size} "
          f"augment={not args.no_augment}")
    dl_tr = make_loader(ds_tr, args.bs, shuffle=True, num_workers=args.workers)
    dl_va = make_loader(ds_va, max(1, args.bs), shuffle=False, num_workers=args.workers)

    model = build_detector(args.arch, pretrained_backbone=not args.no_pretrained,
                           img_size=args.img_size,
                           trainable_backbone_layers=args.trainable_layers).to(device)
    print(f"[model] {args.arch}: {parameter_report(model)}")

    params = [p for p in model.parameters() if p.requires_grad]
    if args.optimizer == "sgd":
        opt = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.wd)
    else:
        opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.wd)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs,
                                                       eta_min=args.lr * 0.01)
    scaler = make_grad_scaler(enabled=(device.type == "cuda" and args.amp))
    warmup_steps = int(args.warmup_frac * max(1, len(dl_tr)))

    hist, best, bad, gstep = [], -1.0, 0, 0
    ckpt_path = WEIGHTS_DIR / f"{tag}.pt"
    t0 = time.time()

    for ep in range(1, args.epochs + 1):
        model.train()
        agg, n = {}, 0
        for images, targets in progress(dl_tr, desc=f"ep{ep} train"):
            images = [im.to(device) for im in images]
            targets = [{k: (v.to(device) if torch.is_tensor(v) else v)
                        for k, v in t.items()} for t in targets]
            if gstep < warmup_steps:
                warmup_lr(opt, args.lr, gstep, warmup_steps)
            opt.zero_grad(set_to_none=True)
            with autocast(enabled=scaler.is_enabled()):
                losses = model(images, targets)
                total = sum(losses.values())
            if not math.isfinite(float(total.detach())):
                print(f"[warn] non-finite loss at step {gstep}, skipping batch")
                gstep += 1
                continue
            scaler.scale(total).backward()
            if args.clip_grad:
                scaler.unscale_(opt)
                torch.nn.utils.clip_grad_norm_(params, args.clip_grad)
            scaler.step(opt); scaler.update()
            for k, v in losses.items():
                agg[k] = agg.get(k, 0.0) + float(v)
            agg["total"] = agg.get("total", 0.0) + float(total)
            n += 1
            gstep += 1
        train_losses = {k: v / max(1, n) for k, v in agg.items()}

        preds, gts = run_inference(model, dl_va, device)
        ap50 = average_precision(preds, gts, 0.5)["ap"]
        ap75 = average_precision(preds, gts, 0.75)["ap"]
        sched.step()

        row = {"epoch": ep, **{f"loss_{k}": v for k, v in train_losses.items()},
               "val_mAP50": float(ap50), "val_mAP75": float(ap75),
               "lr": opt.param_groups[0]["lr"]}
        hist.append(row)
        print(f"[{tag}] ep {ep:03d} | loss {train_losses.get('total', 0):.4f} "
              f"(cls {train_losses.get('loss_classifier', 0):.4f}, "
              f"box {train_losses.get('loss_box_reg', 0):.4f}) "
              f"| val mAP@50 {ap50:.4f} | mAP@75 {ap75:.4f}")

        if ap50 > best + args.min_delta:
            best, bad = ap50, 0
            torch.save({"model_state": model.state_dict(), "args": vars(args),
                        "epoch": ep, "val_mAP50": float(ap50), "seed": SEED,
                        "tag": tag, "loss_description": loss_desc}, ckpt_path)
            print(f"    -> new best, saved {ckpt_path.name}")
        else:
            bad += 1
            if bad >= args.patience:
                print(f"[{tag}] early stop after {ep} epochs")
                break

    summary = {"tag": tag, "seed": SEED, "best_val_mAP50": float(best),
               "epochs_run": len(hist), "minutes": round((time.time() - t0) / 60, 2),
               "loss_description": loss_desc, "config": vars(args),
               "parameter_report": parameter_report(model), "history": hist,
               "training_recipe": {
                   "architecture": args.arch,
                   "input_resolution": f"{args.img_size}x{args.img_size} (letterbox-free square resize)",
                   "annotation_format": "xyxy absolute pixels + label 1 (polyp), "
                                        "converted from kvasir bbox json",
                   "optimizer": args.optimizer, "base_lr": args.lr,
                   "lr_schedule": f"linear warmup over {warmup_steps} steps, then cosine to "
                                  f"{args.lr * 0.01:.2e}",
                   "batch_size": args.bs, "weight_decay": args.wd,
                   "augmentation": "none" if args.no_augment else
                                   "random h/v flip + brightness/contrast + saturation jitter",
                   "inference_score_threshold": 0.05,
                   "nms_iou_threshold": 0.5,
                   "stopping_rule": f"early stop on val mAP@0.50, patience={args.patience}, "
                                    f"min_delta={args.min_delta}; best checkpoint restored"}}
    save_json(summary, METRICS_DIR / f"train_{tag}.json")

    plt = _mpl()
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    ep_ = [h["epoch"] for h in hist]
    for k in [k for k in hist[0] if k.startswith("loss_")]:
        axes[0].plot(ep_, [h[k] for h in hist], marker="o", ms=3, label=k.replace("loss_", ""))
    axes[0].set_xlabel("epoch"); axes[0].set_ylabel("loss"); axes[0].legend(fontsize=8)
    axes[0].set_title(f"{tag}: training losses")
    axes[1].plot(ep_, [h["val_mAP50"] for h in hist], marker="o", ms=3, label="val mAP@0.50")
    axes[1].plot(ep_, [h["val_mAP75"] for h in hist], marker="s", ms=3, label="val mAP@0.75")
    axes[1].set_xlabel("epoch"); axes[1].legend(); axes[1].grid(alpha=0.3)
    axes[1].set_title(f"{tag}: validation mAP")
    fig.tight_layout(); fig.savefig(PLOTS_DIR / f"curves_{tag}.png", dpi=160); plt.close(fig)
    print(f"[saved] {PLOTS_DIR / f'curves_{tag}.png'}")
    print(f"[done] best val mAP@0.50 = {best:.4f}  |  weights: {ckpt_path}")


if __name__ == "__main__":
    main()