"""
Phase 1 — v2 training pipeline for fundus classification.

Differences from v1 (run_final_experiments.py):
  * Reads holdout_split_augmented.json (group-aware split over the full
    Original+Augmented union; no filename-level leakage).
  * Adds CLAHE preprocessing (luminance channel) before all transforms.
  * Adds RandAugment(n=2, m=9) on the training transforms.
  * Adds WeightedRandomSampler (inverse class frequency).
  * Adds MixUp/CutMix (α=0.2, alternating per batch with p=0.5).
  * 100 epochs, EarlyStop patience 12, warmup (3 ep) + cosine.
  * 6-view TTA at inference (original + hflip + 4 corner crops).
"""

import argparse, json, math, os, random, time
from pathlib import Path
from collections import defaultdict

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, models
import cv2
from PIL import Image
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support,
    roc_auc_score, average_precision_score,
)
from scipy.stats import binom
from tqdm import tqdm


# ---------------------------- repro ----------------------------

def set_seed(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)


# ------------------------- CLAHE preprocessing -------------------------

class CLAHEPreprocess:
    """Apply CLAHE on the L channel of LAB color space. PIL in, PIL out."""
    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])
        rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        return Image.fromarray(rgb)


# ------------------------- dataset -------------------------

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


# ------------------------- transforms -------------------------

IMAGENET_MEAN = [0.485, 0.456, 0.406]; IMAGENET_STD = [0.229, 0.224, 0.225]
CLIP_MEAN = [0.4815, 0.4578, 0.4082]; CLIP_STD = [0.2686, 0.2613, 0.2758]

def build_transforms(image_size, use_clip_norm=False, train=True, use_clahe=True):
    mean = CLIP_MEAN if use_clip_norm else IMAGENET_MEAN
    std = CLIP_STD if use_clip_norm else IMAGENET_STD
    pre = [CLAHEPreprocess()] if use_clahe else []
    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.RandomVerticalFlip(p=0.2),
            transforms.RandomRotation(20),
            transforms.RandAugment(num_ops=2, magnitude=9),
            transforms.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.1),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
            transforms.RandomErasing(p=0.25, scale=(0.02, 0.15)),
        ])
    return transforms.Compose(pre + [
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])


# ------------------------- models -------------------------

def build_model(name, num_classes):
    name = name.lower()
    if name == "vgg19":
        m = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1)
        m.classifier[6] = nn.Linear(m.classifier[6].in_features, num_classes)
        return m, 224, False
    if name == "resnet50":
        m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
        m.fc = nn.Linear(m.fc.in_features, num_classes); return m, 224, False
    if name == "resnet101":
        m = models.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V2)
        m.fc = nn.Linear(m.fc.in_features, num_classes); return m, 224, False
    if name == "densenet121":
        m = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1)
        m.classifier = nn.Linear(m.classifier.in_features, num_classes); return m, 224, False
    if name == "inception_v3":
        m = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1, aux_logits=True)
        m.fc = nn.Linear(m.fc.in_features, num_classes)
        m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, num_classes)
        return m, 299, False
    if name == "clip_openai":
        import open_clip
        model, _, _ = open_clip.create_model_and_transforms("ViT-B-16", pretrained="openai")
        class CLIPClf(nn.Module):
            def __init__(self, backbone, nc):
                super().__init__(); self.backbone = backbone.visual
                d = self.backbone.output_dim if hasattr(self.backbone, "output_dim") else 512
                self.head = nn.Linear(d, nc)
            def forward(self, x):
                f = self.backbone(x); return self.head(f)
        return CLIPClf(model, num_classes), 224, True
    raise ValueError(name)


# ------------------------- MixUp / CutMix -------------------------

def mixup(x, y, alpha=0.2, num_classes=10):
    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1.0
    idx = torch.randperm(x.size(0), device=x.device)
    x = lam * x + (1 - lam) * x[idx]
    y_oh = F.one_hot(y, num_classes).float()
    y_mix = lam * y_oh + (1 - lam) * y_oh[idx]
    return x, y_mix

def cutmix(x, y, alpha=1.0, num_classes=10):
    lam = np.random.beta(alpha, alpha) if alpha > 0 else 1.0
    idx = torch.randperm(x.size(0), device=x.device)
    H, W = x.size(2), x.size(3)
    cut_rat = math.sqrt(1.0 - lam)
    cw, ch = int(W * cut_rat), int(H * cut_rat)
    cx, cy = np.random.randint(W), np.random.randint(H)
    x1 = np.clip(cx - cw // 2, 0, W); x2 = np.clip(cx + cw // 2, 0, W)
    y1 = np.clip(cy - ch // 2, 0, H); y2 = np.clip(cy + ch // 2, 0, H)
    x[:, :, y1:y2, x1:x2] = x[idx, :, y1:y2, x1:x2]
    lam = 1 - ((x2 - x1) * (y2 - y1) / (W * H))
    y_oh = F.one_hot(y, num_classes).float()
    y_mix = lam * y_oh + (1 - lam) * y_oh[idx]
    return x, y_mix


# ------------------------- metrics -------------------------

def expected_calibration_error(probs, labels, n_bins=15):
    conf = probs.max(axis=1); pred = probs.argmax(axis=1); correct = (pred == labels).astype(float)
    bins = np.linspace(0, 1, n_bins + 1); ece = 0.0
    for i in range(n_bins):
        mask = (conf > bins[i]) & (conf <= bins[i+1])
        if mask.sum() > 0:
            ece += (mask.mean()) * abs(correct[mask].mean() - conf[mask].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(vals.mean()), float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))


# ------------------------- TTA inference -------------------------

@torch.no_grad()
def tta_predict(model, images, device):
    """6 views: original + hflip + 4 corner crops of 90% size resized back."""
    model.eval()
    out_probs = None; n_views = 0
    B, C, H, W = images.shape
    crop = int(H * 0.9)
    views = [images, torch.flip(images, dims=[3])]
    for (y, x) in [(0, 0), (0, W - crop), (H - crop, 0), (H - crop, W - crop)]:
        c = images[:, :, y:y+crop, x:x+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_probs = p if out_probs is None else out_probs + p
        n_views += 1
    return (out_probs / n_views).cpu().numpy()


@torch.no_grad()
def evaluate(model, loader, device, num_classes, use_tta=False):
    model.eval(); all_probs, all_labels = [], []
    for x, y in loader:
        if use_tta:
            p = tta_predict(model, x, device)
        else:
            x = x.to(device); out = model(x)
            if isinstance(out, tuple): out = out[0]
            p = F.softmax(out, dim=1).cpu().numpy()
        all_probs.append(p); all_labels.append(y.numpy())
    probs = np.concatenate(all_probs); labels = np.concatenate(all_labels)
    preds = probs.argmax(axis=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")
    ece = expected_calibration_error(probs, labels)
    return {
        "acc": acc, "precision": p, "recall": r, "f1": f1,
        "roc_auc": roc, "pr_auc": pr_auc, "ece": ece,
        "labels": labels.tolist(), "preds": preds.tolist(), "probs": probs.tolist(),
    }


# ------------------------- train one model -------------------------

def train_model(name, samples_train, samples_val, num_classes, device,
                epochs, batch_size, workers, patience, label, use_clahe):
    model, image_size, use_clip = build_model(name, num_classes)
    model = model.to(device)
    tf_train = build_transforms(image_size, use_clip_norm=use_clip, train=True,  use_clahe=use_clahe)
    tf_val   = build_transforms(image_size, use_clip_norm=use_clip, train=False, use_clahe=use_clahe)
    ds_train = ImageListDataset(samples_train, tf_train)
    ds_val   = ImageListDataset(samples_val,   tf_val)

    # Weighted sampler
    labels_arr = np.array([s[1] for s in samples_train])
    class_counts = np.bincount(labels_arr, minlength=num_classes)
    class_weights = 1.0 / np.maximum(class_counts, 1)
    sample_weights = class_weights[labels_arr]
    sampler = WeightedRandomSampler(sample_weights.tolist(), num_samples=len(sample_weights), replacement=True)

    dl_train = DataLoader(ds_train, batch_size=batch_size, sampler=sampler,
                          num_workers=workers, pin_memory=True, drop_last=True)
    dl_val   = DataLoader(ds_val, batch_size=batch_size, shuffle=False,
                          num_workers=workers, pin_memory=True)

    opt = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-4)
    warmup_epochs = 3
    def lr_lambda(epoch):
        if epoch < warmup_epochs: return (epoch + 1) / warmup_epochs
        prog = (epoch - warmup_epochs) / max(1, epochs - warmup_epochs)
        return 0.5 * (1 + math.cos(math.pi * prog))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
    scaler = GradScaler()

    best_f1 = -1; best_state = None; bad = 0
    history = []
    for ep in range(epochs):
        model.train()
        t0 = time.time(); n = 0; loss_sum = 0.0
        for x, y in dl_train:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            r = np.random.rand()
            if r < 0.4:
                x_m, y_soft = mixup(x, y, alpha=0.2, num_classes=num_classes); use_soft = True
            elif r < 0.7:
                x_m, y_soft = cutmix(x, y, alpha=1.0, num_classes=num_classes); use_soft = True
            else:
                x_m, y_soft = x, y; use_soft = False
            opt.zero_grad(set_to_none=True)
            with autocast():
                out = model(x_m)
                if isinstance(out, tuple):
                    main_out, aux_out = out
                    if use_soft:
                        loss = -(y_soft * F.log_softmax(main_out, 1)).sum(1).mean()
                        loss += 0.4 * (-(y_soft * F.log_softmax(aux_out, 1)).sum(1).mean())
                    else:
                        loss = F.cross_entropy(main_out, y_soft) + 0.4 * F.cross_entropy(aux_out, y_soft)
                else:
                    if use_soft:
                        loss = -(y_soft * F.log_softmax(out, 1)).sum(1).mean()
                    else:
                        loss = F.cross_entropy(out, y_soft)
            scaler.scale(loss).backward(); scaler.step(opt); scaler.update()
            loss_sum += loss.item() * x.size(0); n += x.size(0)
        sched.step()
        val = evaluate(model, dl_val, device, num_classes, use_tta=False)
        dt = time.time() - t0
        history.append({"epoch": ep, "loss": loss_sum/n, "val_acc": val["acc"], "val_f1": val["f1"], "dt": dt})
        print(f"[{label}] ep {ep+1:3d}/{epochs}  loss {loss_sum/n:.4f}  val_acc {val['acc']*100:5.2f}  val_f1 {val['f1']*100:5.2f}  ({dt:.0f}s)", flush=True)
        if val["f1"] > best_f1 + 1e-4:
            best_f1 = val["f1"]; best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}; bad = 0
        else:
            bad += 1
            if bad >= patience:
                print(f"[{label}] early stop at epoch {ep+1}", flush=True); break
    if best_state is not None:
        model.load_state_dict(best_state)
    return model, history, best_f1


# ------------------------- main -------------------------

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("--models", nargs="+",
                    default=["vgg19", "resnet50", "resnet101", "densenet121", "inception_v3", "clip_openai"])
    ap.add_argument("--epochs", type=int, default=100)
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--patience", type=int, default=12)
    ap.add_argument("--use-clahe", action="store_true", default=True)
    ap.add_argument("--skip-cv", action="store_true",
                    help="Only do final-train + indep test (skip k-fold CV)")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    set_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))
    classes = M["classes"]; num_classes = len(classes)
    print(f"classes ({num_classes}): {classes}")
    samples_train = [tuple(x) for x in M["splits"]["train"]]
    samples_val   = [tuple(x) for x in M["splits"]["val"]]
    samples_test  = [tuple(x) for x in M["splits"]["test"]]
    print(f"train {len(samples_train)} | val {len(samples_val)} | test {len(samples_test)}")

    pool_paths = M["pool_paths"]; pool_labels = M["pool_labels"]
    folds = M["folds"]

    summary = {}
    test_preds_all = {}
    for name in args.models:
        print(f"\n========================  {name}  ========================")
        per_fold = []
        if not args.skip_cv:
            cv_epochs = max(20, args.epochs // 2)  # CV uses half-budget; final uses full
            for fi, fold in enumerate(folds[:args.folds]):
                tr = [(pool_paths[i], pool_labels[i]) for i in fold["train_idx"]]
                va = [(pool_paths[i], pool_labels[i]) for i in fold["val_idx"]]
                print(f"\n--- fold {fi+1}/{args.folds}  train {len(tr)} val {len(va)} ---")
                fmodel, hist, best_f1 = train_model(
                    name, tr, va, num_classes, device,
                    cv_epochs, args.batch_size, args.workers, args.patience,
                    label=f"{name}-f{fi+1}", use_clahe=args.use_clahe)
                # Eval (no TTA) for fold metrics
                _, image_size_f, use_clip_f = build_model(name, num_classes)
                tf_vf = build_transforms(image_size_f, use_clip_norm=use_clip_f, train=False, use_clahe=args.use_clahe)
                dl_vf = DataLoader(ImageListDataset(va, tf_vf),
                                   batch_size=args.batch_size, shuffle=False,
                                   num_workers=args.workers, pin_memory=True)
                fres = evaluate(fmodel, dl_vf, device, num_classes, use_tta=False)
                per_fold.append({
                    "fold": fi, "best_val_f1": best_f1,
                    "val_acc": fres["acc"], "val_f1": fres["f1"],
                    "val_roc_auc": fres["roc_auc"], "val_ece": fres["ece"],
                    "history": hist,
                })
                del fmodel; torch.cuda.empty_cache()

        # Final train: combine train+val for stronger final model, evaluate on test
        print(f"\n--- {name} FINAL train on train+val ({len(samples_train)+len(samples_val)} samples) ---")
        final_model, hist, _ = train_model(
            name, samples_train + samples_val, samples_val, num_classes, device,
            args.epochs, args.batch_size, args.workers, args.patience,
            label=f"{name}-final", use_clahe=args.use_clahe)
        # Test eval with TTA
        _, image_size, use_clip = build_model(name, num_classes)
        tf_test = build_transforms(image_size, use_clip_norm=use_clip, train=False, use_clahe=args.use_clahe)
        dl_test = DataLoader(ImageListDataset(samples_test, tf_test),
                             batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True)
        print(f"[{name}] evaluating on test with TTA ...")
        test_res = evaluate(final_model, dl_test, device, num_classes, use_tta=True)
        labels = np.array(test_res["labels"]); preds = np.array(test_res["preds"])
        acc_mean, acc_lo, acc_hi = bootstrap_ci(labels, preds, accuracy_score)
        f1_mean, f1_lo, f1_hi = bootstrap_ci(labels, preds,
            lambda l, p: precision_recall_fscore_support(l, p, average="macro", zero_division=0)[2])
        test_res["acc_ci"] = [acc_lo, acc_hi]; test_res["f1_ci"] = [f1_lo, f1_hi]
        summary[name] = {
            "test_acc": test_res["acc"], "test_acc_ci": test_res["acc_ci"],
            "test_f1": test_res["f1"], "test_f1_ci": test_res["f1_ci"],
            "test_precision": test_res["precision"], "test_recall": test_res["recall"],
            "roc_auc": test_res["roc_auc"], "pr_auc": test_res["pr_auc"], "ece": test_res["ece"],
            "n_folds_run": len(per_fold),
        }
        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": test_res["labels"], "preds": test_res["preds"], "probs": test_res["probs"]}, f)
        if per_fold:
            with open(out_dir / f"{name}_kfold.json", "w") as f: json.dump(per_fold, f, indent=2)
        torch.save(final_model.state_dict(), w_dir / f"{name}_v2_final.pth")
        test_preds_all[name] = test_res

        print(f"[{name}] test acc {test_res['acc']*100:.2f}  f1 {test_res['f1']*100:.2f}  roc {test_res['roc_auc']:.4f}  ece {test_res['ece']:.4f}")

    # McNemar
    print("\n=== McNemar pairwise ===")
    mcnemar = {}
    keys = list(test_preds_all.keys())
    labels = np.array(test_preds_all[keys[0]]["labels"])
    for i in range(len(keys)):
        for j in range(i+1, len(keys)):
            p1 = np.array(test_preds_all[keys[i]]["preds"]); p2 = np.array(test_preds_all[keys[j]]["preds"])
            c1 = p1 == labels; c2 = p2 == labels
            b = int(((c1) & (~c2)).sum()); c = int(((~c1) & (c2)).sum())
            n = b + c
            if n == 0: pval = 1.0
            else:
                k = min(b, c); pval = float(2 * binom.cdf(k, n, 0.5))
                if pval > 1: pval = 1.0
            mcnemar[f"{keys[i]}_vs_{keys[j]}"] = {"b": b, "c": c, "p": pval}
            print(f"  {keys[i]} vs {keys[j]}: b={b} c={c} p={pval:.4g}")
    with open(out_dir / "mcnemar.json", "w") as f: json.dump(mcnemar, f, indent=2)
    with open(out_dir / "summary.json", "w") as f: json.dump(summary, f, indent=2)
    print("\nDone.")


if __name__ == "__main__":
    main()