"""held-out CV for panda. GroupKFold by dataset when possible, else StratifiedKFold."""
from __future__ import annotations
import argparse, sys, json, pickle, warnings, numpy as np, pandas as pd, torch, torch.nn.functional as F
from pathlib import Path
import anndata as ad, scanpy as sc, scipy.sparse as sp, yaml
from sklearn.model_selection import GroupKFold, StratifiedKFold
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, classification_report
warnings.filterwarnings("ignore")

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
from panda import (
    PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased,
    subcenter_angular_infonce, prototype_repulsion,
)
from scripts.common.train_panda import prepare_batches, load_corpus, get_marker_gene_list

ROOT = Path(str(PANDA_ROOT))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def train_one_fold(Xpca, Xmark, y, y_dset, log10cz, classes, variant,
                   train_idx, epochs=6, batch=256, lr=1e-3, seed=0):
    n_classes = len(classes)
    n_datasets = int(max(y_dset[train_idx].max() + 1, 1))
    n_markers = Xmark.shape[1] if Xmark is not None else 0

    torch.manual_seed(seed); np.random.seed(seed)
    model = PANDAEncoder(
        variant=variant, n_pca=50, n_markers=n_markers,
        n_classes=n_classes, n_sub=3, n_datasets=n_datasets, dropout=0.2,
    ).to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)

    train_n = len(train_idx)
    rng = np.random.default_rng(seed)
    Xt = Xpca[train_idx]; yt = y[train_idx]; ydt = y_dset[train_idx]; dt = log10cz[train_idx]
    Xmt = Xmark[train_idx] if Xmark is not None else None

    for epoch in range(epochs):
        stage = 0 if epoch < 1 else 1 if epoch < 3 else 2 if epoch < 6 else 3
        for g in opt.param_groups: g["lr"] = lr * (0.5 if epoch >= epochs - 1 else 1.0)
        perm = rng.permutation(train_n)
        for bstart in range(0, train_n, batch):
            idx = perm[bstart:bstart+batch]
            x = torch.from_numpy(Xt[idx]).to(DEVICE)
            xm = torch.from_numpy(Xmt[idx]).to(DEVICE) if Xmt is not None else None
            yy = torch.from_numpy(yt[idx]).to(DEVICE)
            yd = torch.from_numpy(ydt[idx]).to(DEVICE)
            dd = torch.from_numpy(dt[idx]).float().to(DEVICE).unsqueeze(1)
            aux = torch.zeros(len(idx), 2, device=DEVICE)
            lam = 0.1 if stage >= 2 else 0.0
            out = model(x, aux, x_markers=xm, lam_dann=lam)
            z = out["z"]
            L = supcon_loss(z, yy, 0.1) + 1.0 * vicreg_loss(z) + 0.4 * F.cross_entropy(out["logits"], yy)
            if stage >= 1:
                L = L + 0.6 * subcenter_angular_infonce(z, yy, model.prototypes.detach().clone(),
                                                        margin=0.15, temperature=0.07)
            if stage >= 2:
                L = L + F.cross_entropy(out["dom"], yd) + 0.3 * F.mse_loss(out["depth"], dd) \
                      + 0.05 * hsic_biased(out["repr"], dd)
            if stage >= 3:
                L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone())
            opt.zero_grad(); L.backward(); opt.step()
            if stage >= 1:
                with torch.no_grad(): model.update_prototypes(z.detach(), yy)
    return model


def evaluate(model, Xpca, Xmark, y, val_idx, classes):
    model.eval()
    preds, probs = [], []
    Xv = Xpca[val_idx]; Xmv = Xmark[val_idx] if Xmark is not None else None
    with torch.no_grad():
        for i in range(0, len(val_idx), 2048):
            xb = torch.from_numpy(Xv[i:i+2048]).to(DEVICE)
            xmb = torch.from_numpy(Xmv[i:i+2048]).to(DEVICE) if Xmv is not None else None
            aux = torch.zeros(len(xb), 2, device=DEVICE)
            out = model(xb, aux, x_markers=xmb, lam_dann=0.0)
            z = out["z"]
            mc = model.max_sub_cos(z)   # (B, K)
            preds.append(mc.argmax(dim=1).cpu().numpy())
            probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy())
    preds = np.concatenate(preds); probs = np.concatenate(probs)
    yv = y[val_idx]
    acc = accuracy_score(yv, preds)
    f1 = f1_score(yv, preds, average="macro", zero_division=0)
    # macro AUROC — nan if val has <2 classes
    n_classes = len(classes)
    try:
        y_onehot = np.eye(n_classes)[yv]
        auc = roc_auc_score(y_onehot, probs, average="macro", multi_class="ovr")
    except Exception:
        auc = float("nan")
    rep = classification_report(yv, preds, labels=list(range(n_classes)),
                                target_names=classes, output_dict=True, zero_division=0)
    return acc, f1, auc, rep


def cv(system, variant, folds=5, epochs=6, split_mode="auto", seed=0):
    a, hvgs, mu, sig, pca = load_corpus(system)
    marker_genes = get_marker_gene_list(system) if variant == "marker" else []
    Xpca, Xmark, y, classes, y_dset, dset_classes, log10cz = prepare_batches(
        a, hvgs, mu, sig, pca, marker_genes, variant
    )
    print(f"[cv] {system}/{variant} n={a.n_obs} K={len(classes)} datasets={len(dset_classes)}", flush=True)

    use_group = (split_mode == "group") or (split_mode == "auto" and len(dset_classes) >= folds)
    if use_group:
        splitter = GroupKFold(n_splits=folds)
        splits = list(splitter.split(np.zeros(len(y)), y, y_dset))
        print(f"[cv] GroupKFold by dataset ({len(dset_classes)} groups)", flush=True)
    else:
        splitter = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
        splits = list(splitter.split(np.zeros(len(y)), y))
        print(f"[cv] StratifiedKFold on labels ({split_mode}) seed={seed}", flush=True)

    per_fold_acc, per_fold_f1, per_fold_auc = [], [], []
    last_report = None
    for fold, (tr, va) in enumerate(splits):
        print(f"[fold {fold+1}/{folds}] train={len(tr)} val={len(va)}", flush=True)
        model = train_one_fold(Xpca, Xmark, y, y_dset, log10cz, classes, variant, tr,
                               epochs=epochs, seed=seed * 100 + fold)
        acc, f1, auc, rep = evaluate(model, Xpca, Xmark, y, va, classes)
        per_fold_acc.append(acc); per_fold_f1.append(f1); per_fold_auc.append(auc)
        last_report = rep
        print(f"[fold {fold+1}] acc={acc:.4f} f1={f1:.4f} auc={auc:.4f}", flush=True)

    result = {
        "system": system, "variant": variant, "folds": folds, "epochs": epochs, "seed": seed,
        "n_cells": int(a.n_obs), "n_classes": len(classes),
        "per_fold_acc": per_fold_acc, "per_fold_f1": per_fold_f1, "per_fold_auc": per_fold_auc,
        "mean_acc": float(np.mean(per_fold_acc)), "std_acc": float(np.std(per_fold_acc)),
        "mean_f1":  float(np.mean(per_fold_f1)),  "std_f1":  float(np.std(per_fold_f1)),
        "mean_auc": float(np.nanmean(per_fold_auc)),
        "std_auc":  float(np.nanstd(per_fold_auc)),
        "per_class_report": last_report,
    }
    out_dir = ROOT / f"discovery/{system}/{variant}"
    out_dir.mkdir(parents=True, exist_ok=True)
    suffix = f"_seed{seed}" if seed != 0 else ""
    (out_dir / f"cv_{folds}fold{suffix}.json").write_text(json.dumps(result, indent=2, default=str))
    print(f"\n[cv] mean acc={result['mean_acc']:.4f}±{result['std_acc']:.4f}  "
          f"f1={result['mean_f1']:.4f}  auc={result['mean_auc']:.4f}", flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("system", choices=["pan_skin", "hematopoiesis", "pancreas"])
    ap.add_argument("--variant", choices=["pca", "marker"], required=True)
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--epochs", type=int, default=6)
    ap.add_argument("--split", choices=["auto", "group", "stratified"], default="auto")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    cv(args.system, args.variant, args.folds, args.epochs, args.split, args.seed)