"""zero-shot Aldrich (GSE220977, ~25k WT + En1-cKO volar snRNA-seq) through frozen PANDA."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, pickle
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import anndata as ad
import scanpy as sc
import scipy.sparse as sp
import torch
import torch.nn.functional as F

from pathlib import Path as _P_root
ROOT = _P_root(__file__).resolve().parents[2]
ROOT_STR = str(ROOT)
sys.path.insert(0, ROOT_STR)
from panda import PANDAEncoder

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CKPT_DIR = Path(f"{ROOT_STR}/checkpoints/pan_skin")
HARM     = Path(f"{ROOT_STR}/data/corpus/pan_skin/harmonized")
OUT_DIR  = Path(f"{ROOT_STR}/discovery/pan_skin/marker")
OUT_DIR.mkdir(parents=True, exist_ok=True)
TARGET   = Path(f"{ROOT_STR}/data/processed/skin/adata_processed.h5ad")


def project_to_corpus(a: ad.AnnData, shared_hvgs, mu, sig):
    """z-scored dense matrix aligned to shared_hvgs, zero-impute missing."""
    n = a.n_obs
    G = len(shared_hvgs)
    hvg_to_ix = {g: i for i, g in enumerate(shared_hvgs)}
    var_names = list(a.var_names.astype(str))
    common = [g for g in var_names if g in hvg_to_ix]
    present_frac = len(common) / G
    print(f"[proj] {len(common)}/{G} shared HVGs present ({present_frac:.1%})", flush=True)

    a_c = a[:, common].copy()
    sc.pp.normalize_total(a_c, target_sum=1e4)
    sc.pp.log1p(a_c)
    X = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32)
    Xf = np.zeros((n, G), dtype=np.float32)
    cols = [hvg_to_ix[g] for g in common]
    Xf[:, cols] = X
    Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10)
    return Xz, present_frac


def main():
    ck = torch.load(CKPT_DIR / "panda_final.pt", map_location=DEVICE, weights_only=False)
    classes = ck["classes"]; datasets = ck["datasets"]
    with (CKPT_DIR / "label_encoding.json").open() as f:
        enc = json.load(f)
    model = PANDAEncoder(n_pca=50, n_classes=len(classes),
                         n_datasets=len(datasets)).to(DEVICE).eval()
    model.load_state_dict(ck["model"])
    prototypes = torch.tensor(ck["prototypes"]).to(DEVICE)
    print(f"[model] classes: {classes}", flush=True)

    stats = np.load(HARM / "corpus_stats.npz", allow_pickle=True)
    shared_hvgs = [str(g) for g in stats["shared_hvgs"]]
    mu, sig = stats["mean"], stats["std"]
    with open(HARM / "pca_basis.pkl", "rb") as f:
        pca_basis = pickle.load(f)

    a = ad.read_h5ad(TARGET)
    if a.raw is not None:
        raw = a.raw.to_adata(); raw.obs = a.obs.copy(); a = raw
    print(f"[target] Aldrich shape: {a.shape}, "
          f"genotype counts: {a.obs['genotype'].value_counts().to_dict()}", flush=True)

    Xz, present_frac = project_to_corpus(a, shared_hvgs, mu, sig)
    counts = np.asarray(a.X.sum(axis=1)).ravel()
    logc = np.log10(counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6)
    Xpca = pca_basis.transform(Xz).astype(np.float32)
    mhf = np.full(a.n_obs, 1.0 - present_frac, dtype=np.float32)

    batch = 4096
    N = a.n_obs
    all_z, all_logits = [], []
    with torch.no_grad():
        for i in range(0, N, batch):
            xb = torch.from_numpy(Xpca[i:i+batch]).to(DEVICE)
            aux = torch.from_numpy(np.stack([mhf[i:i+batch], logc[i:i+batch]], 1).astype(np.float32)).to(DEVICE)
            out = model(xb, aux, lam_dann=0.0)
            all_z.append(out["z"].cpu().numpy())
            all_logits.append(out["logits"].cpu().numpy())
    Z = np.concatenate(all_z, axis=0)
    logits = np.concatenate(all_logits, axis=0)

    protos_np = prototypes.cpu().numpy()
    protos_np = protos_np / (np.linalg.norm(protos_np, axis=1, keepdims=True) + 1e-8)
    cos = Z @ protos_np.T  # (N, K)

    tau = 0.07
    proto_probs = np.exp((cos - cos.max(axis=1, keepdims=True)) / tau)
    proto_probs = proto_probs / proto_probs.sum(axis=1, keepdims=True)

    pred_ix = cos.argmax(axis=1)
    pred_conf = cos.max(axis=1)
    entropy = -(proto_probs * np.log(proto_probs + 1e-12)).sum(axis=1)

    abstain = pred_conf < 0.3
    pred_label = np.array([classes[i] for i in pred_ix], dtype=object)
    pred_label[abstain] = "UNK/abstain"

    # BBSE label-shift correction
    prior_train = np.bincount(np.arange(len(classes)), minlength=len(classes)).astype(float) + 1.0
    prior_train /= prior_train.sum()
    prior_test = proto_probs.mean(axis=0)
    ratio = np.log(prior_test / prior_train + 1e-8)
    bbse_logits = np.log(proto_probs + 1e-12) + ratio[None, :]
    bbse_probs = np.exp(bbse_logits - bbse_logits.max(axis=1, keepdims=True))
    bbse_probs = bbse_probs / bbse_probs.sum(axis=1, keepdims=True)
    bbse_pred = bbse_probs.argmax(axis=1)

    a.obs["pred_label"] = pred_label
    a.obs["pred_conf"]  = pred_conf.astype(np.float32)
    a.obs["pred_entropy"] = entropy.astype(np.float32)
    a.obs["abstain"] = abstain
    a.obs["pred_bbse_label"] = [classes[i] for i in bbse_pred]

    tbl = a.obs[[c for c in ["genotype","sample","cell_type","fate_label",
                             "pred_label","pred_conf","pred_entropy",
                             "abstain","pred_bbse_label"] if c in a.obs.columns]].copy()
    tbl.to_csv(OUT_DIR / "50_aldrich_predictions.csv")
    print(f"[out] wrote {OUT_DIR/'50_aldrich_predictions.csv'} ({len(tbl)} rows)", flush=True)

    print("\n[summary] naive pred_label breakdown:")
    print(a.obs["pred_label"].value_counts().head(20))
    print(f"[summary] abstain rate: {abstain.mean():.1%}")
    if "genotype" in a.obs.columns:
        print("\n[summary] pred_label x genotype:")
        print(pd.crosstab(a.obs["pred_label"], a.obs["genotype"]))
    print("\n[summary] BBSE-corrected pred_label breakdown:")
    print(a.obs["pred_bbse_label"].value_counts().head(20))

    # slim h5ad: obs + obsm only
    slim = ad.AnnData(X=sp.csr_matrix((a.n_obs, 1), dtype=np.float32), obs=a.obs)
    slim.obsm["Z_projection"] = Z.astype(np.float32)
    slim.obsm["proto_cos"]    = cos.astype(np.float32)
    slim.obsm["proto_probs"]  = proto_probs.astype(np.float32)
    slim.obsm["bbse_probs"]   = bbse_probs.astype(np.float32)
    slim.uns["classes"] = classes
    slim.uns["prototypes"] = protos_np
    slim.write_h5ad(OUT_DIR / "50_aldrich_projections.h5ad", compression="gzip")
    print(f"[out] wrote {OUT_DIR/'50_aldrich_projections.h5ad'}", flush=True)


if __name__ == "__main__":
    main()