"""
Phase 2 — Foundation model fine-tuning for fundus classification.

Backbones added:
  * RETFound (MAE-pretrained on 1.6M fundus images; SOTA on most fundus benchmarks)
    weights: https://github.com/rmaphoh/RETFound_MAE
  * DINOv2-Large (general-purpose strong self-supervised features)
  * Swin-Base (timm)

Two-regime fine-tuning:
  1. linear-probe (head only) for 20 epochs   -> stable feature extraction baseline
  2. full fine-tune at LR 1e-5 for 10 epochs  -> task-specific adaptation
"""

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

import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torch.cuda.amp import autocast, GradScaler
from torchvision import transforms
from PIL import Image
import cv2

from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, average_precision_score


# Re-use building blocks from v2 (CLAHE etc.) by inlining to keep this self-contained.
class CLAHEPreprocess:
    def __init__(self, clip_limit=2.0, tile=(8, 8)):
        self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile)
    def __call__(self, img):
        arr = np.array(img.convert("RGB"))
        lab = cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)
        lab[..., 0] = self.clahe.apply(lab[..., 0])
        return Image.fromarray(cv2.cvtColor(lab, cv2.COLOR_LAB2RGB))


class ImageListDataset(Dataset):
    def __init__(self, samples, transform):
        self.samples = samples; self.transform = transform
    def __len__(self): return len(self.samples)
    def __getitem__(self, i):
        p, l = self.samples[i]
        return self.transform(Image.open(p).convert("RGB")), int(l)


def make_transforms(image_size, train, mean, std):
    pre = [CLAHEPreprocess()]
    if train:
        return transforms.Compose(pre + [
            transforms.Resize((image_size + 32, image_size + 32)),
            transforms.RandomResizedCrop(image_size, scale=(0.75, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.RandAugment(num_ops=2, magnitude=7),
            transforms.ColorJitter(0.15, 0.15, 0.1),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
            transforms.RandomErasing(p=0.2, scale=(0.02, 0.1)),
        ])
    return transforms.Compose(pre + [
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])


# ------------------------- backbones -------------------------

def build_dinov2_large(num_classes):
    """DINOv2-L/14: 1024-dim CLS features."""
    backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
    class M(nn.Module):
        def __init__(self):
            super().__init__()
            self.backbone = backbone
            self.head = nn.Linear(1024, num_classes)
            # Materialize parameter lists (avoid generator exhaustion).
            self._head_params = list(self.head.parameters())
            self._backbone_params = list(self.backbone.parameters())
        def forward(self, x):
            f = self.backbone(x)  # CLS token, [B, 1024]
            return self.head(f)
        def trainable_groups(self):
            return [
                {"params": self._head_params,     "lr": 1e-3, "linear_probe": True},
                {"params": self._backbone_params, "lr": 1e-5, "linear_probe": False},
            ]
    return M(), 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]


def build_swin_base(num_classes):
    import timm
    model = timm.create_model("swin_base_patch4_window7_224", pretrained=True, num_classes=num_classes)
    head_params = list(model.head.parameters()) if hasattr(model, "head") else []
    other_params = [p for n, p in model.named_parameters() if not n.startswith("head")]
    class M(nn.Module):
        def __init__(self):
            super().__init__(); self.m = model
        def forward(self, x): return self.m(x)
        def trainable_groups(self):
            return [
                {"params": head_params,  "lr": 1e-3, "linear_probe": True},
                {"params": other_params, "lr": 1e-5, "linear_probe": False},
            ]
    return M(), 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]


def build_retfound(num_classes, weights_path):
    """RETFound ViT-Large/16, MAE-pretrained on fundus images.
    Loads weights from a local checkpoint downloaded from rmaphoh/RETFound_MAE."""
    import timm
    # RETFound is a vanilla MAE ViT-L/16 with patch 16, image 224.
    model = timm.create_model("vit_large_patch16_224", pretrained=False, num_classes=num_classes,
                              drop_path_rate=0.2, global_pool="token")
    if weights_path and os.path.exists(weights_path):
        ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
        state = ckpt.get("model", ckpt.get("state_dict", ckpt))
        # RETFound checkpoints have 'pos_embed' etc; we drop classifier head keys
        state = {k: v for k, v in state.items()
                 if not k.startswith("head.") and not k.startswith("fc_norm.")}
        missing, unexp = model.load_state_dict(state, strict=False)
        print(f"  RETFound loaded: {len(state)} keys, missing={len(missing)}, unexpected={len(unexp)}")
    else:
        print(f"  WARNING: RETFound weights not found at {weights_path}; using random init for backbone (will perform poorly)")
    head_params = list(model.head.parameters()) + list(model.fc_norm.parameters())
    other_params = [p for n, p in model.named_parameters()
                    if not n.startswith("head") and not n.startswith("fc_norm")]
    class M(nn.Module):
        def __init__(self): super().__init__(); self.m = model
        def forward(self, x): return self.m(x)
        def trainable_groups(self):
            return [
                {"params": head_params,  "lr": 1e-3, "linear_probe": True},
                {"params": other_params, "lr": 1e-5, "linear_probe": False},
            ]
    return M(), 224, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]


# ------------------------- train -------------------------

def expected_calibration_error(probs, labels, n_bins=15):
    conf = probs.max(1); pred = probs.argmax(1); correct = (pred == labels).astype(float)
    bins = np.linspace(0, 1, n_bins+1); ece = 0.0
    for i in range(n_bins):
        m = (conf > bins[i]) & (conf <= bins[i+1])
        if m.sum(): ece += m.mean() * abs(correct[m].mean() - conf[m].mean())
    return float(ece)


def bootstrap_ci(labels, preds, metric_fn, n=1000, seed=42):
    rng = np.random.default_rng(seed); N = len(labels); vals = []
    for _ in range(n):
        idx = rng.integers(0, N, N)
        try: vals.append(metric_fn(labels[idx], preds[idx]))
        except Exception: pass
    vals = np.array(vals)
    return float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))


@torch.no_grad()
def tta_predict(model, x, device):
    model.eval(); B, C, H, W = x.shape
    crop = int(H * 0.9); out = None; n = 0
    views = [x, torch.flip(x, dims=[3])]
    for (y, xc) in [(0, 0), (0, W-crop), (H-crop, 0), (H-crop, W-crop)]:
        c = x[:, :, y:y+crop, xc:xc+crop]
        c = F.interpolate(c, size=(H, W), mode="bilinear", align_corners=False)
        views.append(c)
    for v in views:
        p = F.softmax(model(v.to(device)), dim=1)
        out = p if out is None else out + p; n += 1
    return (out/n).cpu().numpy()


@torch.no_grad()
def evaluate(model, loader, device, num_classes, use_tta=False):
    model.eval(); ps, ls = [], []
    for x, y in loader:
        if use_tta: p = tta_predict(model, x, device)
        else:
            p = F.softmax(model(x.to(device)), dim=1).cpu().numpy()
        ps.append(p); ls.append(y.numpy())
    probs = np.concatenate(ps); labels = np.concatenate(ls); preds = probs.argmax(1)
    acc = accuracy_score(labels, preds)
    p, r, f1, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
    try: roc = roc_auc_score(labels, probs, multi_class="ovr", average="macro", labels=list(range(num_classes)))
    except Exception: roc = float("nan")
    try: pr_auc = average_precision_score(F.one_hot(torch.tensor(labels), num_classes).numpy(), probs, average="macro")
    except Exception: pr_auc = float("nan")
    return {"acc": acc, "precision": p, "recall": r, "f1": f1,
            "roc_auc": roc, "pr_auc": pr_auc, "ece": expected_calibration_error(probs, labels),
            "labels": labels.tolist(), "preds": preds.tolist(), "probs": probs.tolist()}


def mixup(x, y, alpha, nc):
    lam = np.random.beta(alpha, alpha)
    i = torch.randperm(x.size(0), device=x.device)
    x = lam*x + (1-lam)*x[i]
    yoh = F.one_hot(y, nc).float()
    return x, lam*yoh + (1-lam)*yoh[i]


def train_foundation(name, build_fn, samples_tr, samples_va, num_classes, device,
                     batch_size, workers, lp_epochs, ft_epochs, patience, label):
    model, image_size, mean, std = build_fn()
    model = model.to(device)

    tf_tr = make_transforms(image_size, train=True,  mean=mean, std=std)
    tf_va = make_transforms(image_size, train=False, mean=mean, std=std)

    ds_tr = ImageListDataset(samples_tr, tf_tr); ds_va = ImageListDataset(samples_va, tf_va)
    labels_arr = np.array([s[1] for s in samples_tr])
    cw = 1.0 / np.maximum(np.bincount(labels_arr, minlength=num_classes), 1)
    sw = cw[labels_arr]
    sampler = WeightedRandomSampler(sw.tolist(), num_samples=len(sw), replacement=True)
    dl_tr = DataLoader(ds_tr, batch_size=batch_size, sampler=sampler, num_workers=workers, pin_memory=True, drop_last=True)
    dl_va = DataLoader(ds_va, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=True)

    groups = model.trainable_groups()
    head_group = next(g for g in groups if g.get("linear_probe"))
    backbone_group = next(g for g in groups if not g.get("linear_probe"))

    scaler = GradScaler()
    best_f1 = -1; best_state = None; bad = 0; history = []

    # ---- Stage 1: linear probe (freeze backbone) ----
    for p in backbone_group["params"]: p.requires_grad = False
    opt = torch.optim.AdamW([{"params": head_group["params"], "lr": head_group["lr"]}], weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=lp_epochs)
    for ep in range(lp_epochs):
        model.train(); t0 = time.time(); loss_sum, n = 0.0, 0
        for x, y in dl_tr:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            if np.random.rand() < 0.3:
                x, ysoft = mixup(x, y, 0.2, num_classes); soft = True
            else: ysoft = y; soft = False
            opt.zero_grad(set_to_none=True)
            with autocast():
                out = model(x)
                loss = -(ysoft * F.log_softmax(out, 1)).sum(1).mean() if soft else F.cross_entropy(out, ysoft)
            scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
            loss_sum += loss.item()*x.size(0); n += x.size(0)
        sched.step()
        v = evaluate(model, dl_va, device, num_classes)
        history.append({"phase": "lp", "epoch": ep, "loss": loss_sum/n, "val_acc": v["acc"], "val_f1": v["f1"]})
        print(f"[{label} LP] ep {ep+1}/{lp_epochs}  loss {loss_sum/n:.4f}  val_acc {v['acc']*100:5.2f}  val_f1 {v['f1']*100:5.2f}  ({time.time()-t0:.0f}s)", flush=True)
        if v["f1"] > best_f1 + 1e-4:
            best_f1 = v["f1"]; best_state = {k: vv.detach().cpu().clone() for k, vv in model.state_dict().items()}; bad = 0
        else:
            bad += 1
            if bad >= patience: print(f"[{label} LP] early stop"); break

    # ---- Stage 2: full fine-tune (unfreeze backbone, low LR) ----
    if best_state is not None: model.load_state_dict(best_state)
    for p in backbone_group["params"]: p.requires_grad = True
    opt = torch.optim.AdamW([
        {"params": head_group["params"],     "lr": 1e-4},
        {"params": backbone_group["params"], "lr": 1e-5},
    ], weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=ft_epochs)
    bad = 0
    for ep in range(ft_epochs):
        model.train(); t0 = time.time(); loss_sum, n = 0.0, 0
        for x, y in dl_tr:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            if np.random.rand() < 0.3:
                x, ysoft = mixup(x, y, 0.2, num_classes); soft = True
            else: ysoft = y; soft = False
            opt.zero_grad(set_to_none=True)
            with autocast():
                out = model(x)
                loss = -(ysoft * F.log_softmax(out, 1)).sum(1).mean() if soft else F.cross_entropy(out, ysoft)
            scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
            loss_sum += loss.item()*x.size(0); n += x.size(0)
        sched.step()
        v = evaluate(model, dl_va, device, num_classes)
        history.append({"phase": "ft", "epoch": ep, "loss": loss_sum/n, "val_acc": v["acc"], "val_f1": v["f1"]})
        print(f"[{label} FT] ep {ep+1}/{ft_epochs}  loss {loss_sum/n:.4f}  val_acc {v['acc']*100:5.2f}  val_f1 {v['f1']*100:5.2f}  ({time.time()-t0:.0f}s)", flush=True)
        if v["f1"] > best_f1 + 1e-4:
            best_f1 = v["f1"]; best_state = {k: vv.detach().cpu().clone() for k, vv in model.state_dict().items()}; bad = 0
        else:
            bad += 1
            if bad >= patience: print(f"[{label} FT] early stop"); break

    if best_state is not None: model.load_state_dict(best_state)
    return model, history, best_f1, image_size, mean, std


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--weights-dir", required=True)
    ap.add_argument("--retfound-weights", default="weights_retfound.pth")
    ap.add_argument("--models", nargs="+", default=["dinov2_l", "swin_b", "retfound"])
    ap.add_argument("--batch-size", type=int, default=24)
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--lp-epochs", type=int, default=20)
    ap.add_argument("--ft-epochs", type=int, default=15)
    ap.add_argument("--patience", type=int, default=8)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    torch.manual_seed(args.seed); np.random.seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu"); print(f"device: {device}")
    out_dir = Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)
    w_dir = Path(args.weights_dir); w_dir.mkdir(parents=True, exist_ok=True)

    M = json.load(open(args.manifest))
    num_classes = len(M["classes"])
    samples_tr = [tuple(x) for x in M["splits"]["train"]]
    samples_va = [tuple(x) for x in M["splits"]["val"]]
    samples_te = [tuple(x) for x in M["splits"]["test"]]
    print(f"train {len(samples_tr)} | val {len(samples_va)} | test {len(samples_te)} | {num_classes} classes")

    builders = {
        "dinov2_l": lambda: build_dinov2_large(num_classes),
        "swin_b":   lambda: build_swin_base(num_classes),
        "retfound": lambda: build_retfound(num_classes, args.retfound_weights),
    }

    summary = {}
    for name in args.models:
        # Skip RETFound if weights file missing or empty (HF gated)
        if name == "retfound":
            wp = args.retfound_weights
            if not (wp and os.path.exists(wp) and os.path.getsize(wp) > 1_000_000):
                print(f"\n[retfound] SKIPPING — weights file '{wp}' missing or empty (HF gated). Use DINOv2/Swin instead.")
                continue
        print(f"\n======== {name} ========")
        try:
            model, hist, best_f1, image_size, mean, std = train_foundation(
                name, builders[name], samples_tr + samples_va, samples_va,
                num_classes, device, args.batch_size, args.workers,
                args.lp_epochs, args.ft_epochs, args.patience, label=name)
        except Exception as e:
            print(f"[{name}] FAILED: {e}"); continue

        tf_te = make_transforms(image_size, train=False, mean=mean, std=std)
        dl_te = DataLoader(ImageListDataset(samples_te, tf_te), batch_size=args.batch_size,
                           shuffle=False, num_workers=args.workers, pin_memory=True)
        print(f"[{name}] evaluating on test with TTA ...")
        res = evaluate(model, dl_te, device, num_classes, use_tta=True)
        labels = np.array(res["labels"]); preds = np.array(res["preds"])
        acc_lo, acc_hi = bootstrap_ci(labels, preds, accuracy_score)
        f1_lo, f1_hi = bootstrap_ci(labels, preds,
            lambda l, p: precision_recall_fscore_support(l, p, average="macro", zero_division=0)[2])
        summary[name] = {
            "test_acc": res["acc"], "test_acc_ci": [acc_lo, acc_hi],
            "test_f1": res["f1"], "test_f1_ci": [f1_lo, f1_hi],
            "test_precision": res["precision"], "test_recall": res["recall"],
            "roc_auc": res["roc_auc"], "pr_auc": res["pr_auc"], "ece": res["ece"],
        }
        with open(out_dir / f"{name}_test.json", "w") as f: json.dump(summary[name], f, indent=2)
        with open(out_dir / f"{name}_test_preds.json", "w") as f:
            json.dump({"labels": res["labels"], "preds": res["preds"], "probs": res["probs"]}, f)
        torch.save(model.state_dict(), w_dir / f"{name}_v2.pth")
        print(f"[{name}] test acc {res['acc']*100:.2f} [{acc_lo*100:.1f},{acc_hi*100:.1f}]  f1 {res['f1']*100:.2f}  roc {res['roc_auc']:.4f}")
        del model; torch.cuda.empty_cache()

    with open(out_dir / "summary_foundation.json", "w") as f: json.dump(summary, f, indent=2)
    print("\nDone (Phase 2).")


if __name__ == "__main__":
    main()