"""variant B — fully supervised panda on replicated dingwall Derm0..Derm11 labels from script 103."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, time
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import anndata as ad
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from scipy.stats import fisher_exact

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.model import (
    PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce
)

ROOT = Path(str(PANDA_ROOT))
REPLICA_H5 = ROOT / "data/processed/dingwall_replica/dingwall_replica.h5ad"
OUT_DIR = ROOT / "discovery/pan_skin/marker"
CK_DIR = ROOT / "checkpoints/pan_skin_dingwall_derm"

TRAIN_FRAC = 0.7
SEED = 0
N_PCA = 40
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Training config (mirrors 20_train_panda.py)
GUARANTEED_PER_CLASS = 6
NATURAL_SLOTS = 96
STAGE_EPOCHS = [15, 25, 40, 40]
BALANCE_MIX = 0.5


# ---------- split ----------

def genotype_stratified_split(labels: np.ndarray, genotypes: np.ndarray,
                              frac_train: float = TRAIN_FRAC, seed: int = SEED
                              ) -> tuple[np.ndarray, np.ndarray]:
    """stratified 70/30 within each (label, genotype) group; preserves cKO/WT ratio per class."""
    rng = np.random.default_rng(seed)
    n = len(labels); train = np.zeros(n, dtype=bool); test = np.zeros(n, dtype=bool)
    for lab in np.unique(labels):
        for g in np.unique(genotypes):
            idx = np.where((labels == lab) & (genotypes == g))[0]
            if len(idx) == 0: continue
            rng.shuffle(idx)
            k = max(1, int(len(idx) * frac_train)) if len(idx) > 1 else len(idx)
            train[idx[:k]] = True
            if len(idx) > 1:
                test[idx[k:]] = True
    return train, test


# ---------- PANDA training (identical to variant A) ----------

class CorpusDataset(Dataset):
    def __init__(self, X, y, d, aux):
        self.X = X.astype(np.float32); self.y = y.astype(np.int64)
        self.d = d.astype(np.int64); self.aux = aux.astype(np.float32)
    def __len__(self): return self.X.shape[0]
    def __getitem__(self, i):
        return (torch.from_numpy(self.X[i]), torch.tensor(self.y[i]),
                torch.tensor(self.d[i]), torch.from_numpy(self.aux[i]))


class HybridSampler:
    def __init__(self, y, n_batches=100, seed=0):
        self.y = np.asarray(y); self.n_batches = n_batches
        self.rng = np.random.default_rng(seed)
        self.classes = np.unique(self.y)
        self.by_cls = {int(c): np.where(self.y == c)[0] for c in self.classes}
        counts = np.bincount(self.y, minlength=int(self.classes.max()) + 1).astype(float)
        self.natural_p = counts / counts.sum()
    def __iter__(self):
        for _ in range(self.n_batches):
            batch = []
            for c in self.classes:
                idx = self.by_cls[int(c)]
                take = min(GUARANTEED_PER_CLASS, len(idx))
                if take > 0:
                    batch.extend(self.rng.choice(idx, size=take, replace=(len(idx) < take)).tolist())
            for _ in range(NATURAL_SLOTS):
                c = self.rng.choice(len(self.natural_p), p=self.natural_p)
                idx = self.by_cls.get(int(c), self.by_cls[int(self.classes[0])])
                batch.append(int(self.rng.choice(idx)))
            yield batch
    def __len__(self): return self.n_batches


def train_panda(X_tr, y_tr, d_tr, aux_tr, n_classes, n_datasets, ck_out: Path):
    ck_out.mkdir(parents=True, exist_ok=True)
    counts = np.bincount(y_tr, minlength=n_classes)
    inv_sqrt = 1.0 / np.sqrt(counts + 1); inv_sqrt = inv_sqrt / inv_sqrt.mean()
    class_w = BALANCE_MIX * inv_sqrt + (1 - BALANCE_MIX) * np.ones_like(inv_sqrt)
    class_w = torch.tensor(class_w, dtype=torch.float32, device=DEVICE)

    ds = CorpusDataset(X_tr, y_tr, d_tr, aux_tr)
    loader = DataLoader(ds, batch_sampler=HybridSampler(y_tr, n_batches=100), num_workers=0)

    model = PANDAEncoder(variant="pca", n_pca=X_tr.shape[1], n_classes=n_classes,
                         n_datasets=n_datasets).to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)

    for stage, n_ep in enumerate(STAGE_EPOCHS):
        print(f"[panda-B] stage {stage} ({n_ep} epochs)", flush=True)
        for e in range(n_ep):
            t0 = time.time(); losses = []
            for X_b, y_b, d_b, aux_b in loader:
                X_b = X_b.to(DEVICE); y_b = y_b.to(DEVICE); d_b = d_b.to(DEVICE); aux_b = aux_b.to(DEVICE)
                lam = 1.0 if stage >= 2 else 0.0
                out = model(X_b, aux_b, lam_dann=lam)
                L_supcon = supcon_loss(out["z"], y_b)
                L_vic = vicreg_loss(out["z"])
                L_ce = F.cross_entropy(out["logits"], y_b, weight=class_w, label_smoothing=0.05)
                total = L_supcon + 1.0 * L_vic + 0.4 * L_ce
                if stage >= 1:
                    proto_ref = model.prototypes.detach().clone()
                    total = total + 0.6 * subcenter_angular_infonce(out["z"], y_b, proto_ref)
                if stage >= 2:
                    total = total + F.cross_entropy(out["dom"], d_b)
                    total = total + 0.3 * F.mse_loss(out["depth"].squeeze(1), aux_b[:, 1])
                    total = total + 0.05 * hsic_biased(out["repr"], aux_b[:, 1:2])
                opt.zero_grad(); total.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                opt.step()
                if stage >= 1:
                    model.update_prototypes(out["z"].detach(), y_b)
                losses.append(float(total.item()))
            if (e + 1) % 5 == 0:
                print(f"  ep {e+1}/{n_ep} loss={np.mean(losses):.3f} dt={time.time()-t0:.1f}s", flush=True)
        torch.save({"model": model.state_dict()}, ck_out / f"panda_stage{stage}.pt")
    torch.save({"model": model.state_dict(),
                "prototypes": model.prototypes.detach().cpu().numpy()},
               ck_out / "panda_final.pt")
    return model


@torch.no_grad()
def infer(model, X, aux):
    model.eval()
    Xt = torch.from_numpy(X.astype(np.float32)).to(DEVICE)
    at = torch.from_numpy(aux.astype(np.float32)).to(DEVICE)
    B = 4096; preds = []; confs = []
    for i in range(0, len(Xt), B):
        out = model(Xt[i:i+B], at[i:i+B])
        p = F.softmax(out["logits"], dim=1)
        preds.append(p.argmax(dim=1).cpu().numpy())
        confs.append(p.max(dim=1).values.cpu().numpy())
    return np.concatenate(preds), np.concatenate(confs)


# ---------- reporting ----------

def depletion_table(true_or_pred: np.ndarray, genotype: np.ndarray, class_names: list[str]
                    ) -> pd.DataFrame:
    n_wt = int((genotype == "WT").sum()); n_cko = int((genotype == "En1-cKO").sum())
    base = n_cko / max(n_wt + n_cko, 1)
    rows = []
    for i, cn in enumerate(class_names):
        m = true_or_pred == i
        w = int(((genotype == "WT") & m).sum()); k = int(((genotype == "En1-cKO") & m).sum())
        if w + k == 0: continue
        try:
            odds, p = fisher_exact([[w, n_wt - w], [k, n_cko - k]], alternative="two-sided")
        except ValueError:
            odds, p = 1.0, 1.0
        rows.append({"derm_label": cn, "n": w + k, "n_WT": w, "n_cKO": k,
                     "cko_frac": k / (w + k), "baseline_cko": base,
                     "odds_ratio": float(odds), "fisher_p": float(p)})
    return pd.DataFrame(rows).sort_values("cko_frac")


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True); CK_DIR.mkdir(parents=True, exist_ok=True)

    print("[B] load replica", flush=True)
    if not REPLICA_H5.exists():
        raise FileNotFoundError(f"Run 103 first — {REPLICA_H5} missing")
    a = ad.read_h5ad(REPLICA_H5)

    dermal = a[a.obs["derm_label"].astype(str) != "non_dermal"].copy()
    print(f"[B] dermal n={dermal.n_obs}", flush=True)
    labels_str = dermal.obs["derm_label"].astype(str).values
    classes = sorted(set(labels_str))
    cls_ix = {c: i for i, c in enumerate(classes)}
    y_all = np.array([cls_ix[c] for c in labels_str])
    genotype = dermal.obs["genotype"].astype(str).values

    # get embedding from replica (harmony-corrected PCA)
    rep_key = dermal.uns.get("_replica_rep", "X_pca_harmony")
    if rep_key not in dermal.obsm:
        rep_key = "X_pca_harmony" if "X_pca_harmony" in dermal.obsm else "X_pca"
    X_all = np.asarray(dermal.obsm[rep_key])
    print(f"[B] using {rep_key} (d={X_all.shape[1]})", flush=True)

    sample_ix = {s: i for i, s in enumerate(sorted(dermal.obs["sample"].astype(str).unique()))}
    d_all = np.array([sample_ix[s] for s in dermal.obs["sample"].astype(str)])
    total_counts = np.asarray(dermal.X.sum(axis=1)).ravel()
    logc = np.log10(total_counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6)
    aux_all = np.stack([np.zeros(dermal.n_obs, dtype=np.float32), logc.astype(np.float32)], axis=1)

    print("[B] genotype-stratified 70/30 split", flush=True)
    tr, te = genotype_stratified_split(labels_str, genotype, frac_train=TRAIN_FRAC, seed=SEED)
    print(f"[B] train={tr.sum()} test={te.sum()}", flush=True)

    manifest = pd.DataFrame({
        "cell_id": dermal.obs_names.astype(str).values,
        "derm_label": labels_str,
        "genotype": genotype,
        "split": np.where(tr, "train", np.where(te, "test", "unassigned")),
    })
    manifest.to_csv(OUT_DIR / "104_dingwall_derm_split_manifest.csv", index=False)

    print("[B] train PANDA", flush=True)
    model = train_panda(X_all[tr], y_all[tr], d_all[tr], aux_all[tr],
                        n_classes=len(classes), n_datasets=len(sample_ix), ck_out=CK_DIR)

    print("[B] infer on held-out", flush=True)
    pred_ix, conf = infer(model, X_all[te], aux_all[te])
    pred = pd.DataFrame({
        "cell_id": dermal.obs_names.astype(str).values[te],
        "derm_true": labels_str[te],
        "derm_pred": [classes[p] for p in pred_ix],
        "confidence": conf,
        "genotype": genotype[te],
    })
    pred.to_csv(OUT_DIR / "104_dingwall_derm_predictions.csv", index=False)

    # depletion — reported for TEST set only, using PANDA predictions
    pred_ix_full = np.array([cls_ix[c] for c in pred["derm_pred"].values])
    dep_pred = depletion_table(pred_ix_full, genotype[te], classes)
    dep_true = depletion_table(y_all[te], genotype[te], classes)
    dep_pred.to_csv(OUT_DIR / "104_dingwall_derm_depletion_pred.csv", index=False)
    dep_true.to_csv(OUT_DIR / "104_dingwall_derm_depletion_true.csv", index=False)

    d10_true = dep_true[dep_true["derm_label"] == "Derm10"].to_dict("records")
    d10_pred = dep_pred[dep_pred["derm_label"] == "Derm10"].to_dict("records")
    acc = float((pred_ix == y_all[te]).mean())

    summary = {
        "variant": "B_fully_supervised_replica_labels",
        "n_dermal_total": int(dermal.n_obs),
        "n_train": int(tr.sum()), "n_test": int(te.sum()),
        "classes": classes,
        "test_accuracy": acc,
        "expected_paper_derm10": {"wt_pct": 1.99, "cko_pct": 0.08,
                                   "or_approx": 24.5, "wt_n_approx": 346, "cko_n_approx": 7},
        "test_derm10_true": d10_true,
        "test_derm10_pred": d10_pred,
        "test_depletion_true": dep_true.to_dict("records"),
        "test_depletion_pred": dep_pred.to_dict("records"),
    }
    (OUT_DIR / "104_dingwall_derm_summary.json").write_text(json.dumps(summary, indent=2, default=str))
    print(f"[B] done -> {OUT_DIR}/104_dingwall_derm_*", flush=True)


if __name__ == "__main__":
    main()