"""adult-beta canonical panel enrichment on Veres: mean log1p in adult-beta vs beta vs overall. reports vacuous if n_adult_beta=0."""
from pathlib import Path
import warnings, json, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
warnings.filterwarnings("ignore"); sc.settings.verbosity = 0

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))
OUT = ROOT / "discovery/pancreas/marker"
OUT.mkdir(parents=True, exist_ok=True)

PANEL = ["MAFA", "UCN3", "IAPP", "INS", "SIX3", "MAFB", "MNX1", "NEUROD1"]


def load_veres():
    SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract"
    parts = []
    for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")):
        counts_file = str(meta_file).replace("cell_metadata", "processed_counts")
        if not Path(counts_file).exists(): continue
        meta = pd.read_csv(meta_file, sep="\t", compression="gzip")
        counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0)
        obs = meta.set_index("library.barcode")
        obs = obs.loc[obs.index.intersection(counts.index)]
        counts_al = counts.loc[obs.index]
        X = sp.csr_matrix(counts_al.values.astype(np.float32))
        a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns))
        a.var_names_make_unique()
        parts.append(a)
    return ad.concat(parts, join="outer")


print("[load] Veres + predictions", flush=True)
raw = load_veres()
pred_df = pd.read_csv(ROOT / "discovery/pancreas/marker/veres_predictions.csv")
pred_df["cell_id"] = pred_df["cell_id"].astype(str).str.replace(r"^veres_", "", regex=True)
common = raw.obs_names.intersection(pd.Index(pred_df["cell_id"].astype(str)))
raw = raw[list(common)].copy()
pred_map = dict(zip(pred_df["cell_id"].astype(str), pred_df["pred_label"]))
raw.obs["pred_label"] = pd.Categorical([pred_map.get(c, "unknown") for c in raw.obs_names])
print(f"[align] {raw.n_obs} cells", flush=True)

sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw)

n_adult_beta = int((raw.obs["pred_label"] == "adult-beta").sum())
n_beta       = int((raw.obs["pred_label"] == "beta").sum())
print(f"[counts] adult-beta={n_adult_beta}  beta={n_beta}  overall={raw.n_obs}", flush=True)

def mean_expr(mask, gene):
    if gene not in raw.var_names or mask.sum() == 0:
        return float("nan")
    col = raw[mask, gene].X
    if sp.issparse(col): col = col.toarray()
    return float(col.mean())

mask_ab = (raw.obs["pred_label"] == "adult-beta").values
mask_b  = (raw.obs["pred_label"] == "beta").values

result = {
    "cluster": "adult-beta",
    "n_adult_beta": n_adult_beta,
    "n_beta": n_beta,
    "n_overall": int(raw.n_obs),
    "vacuous": n_adult_beta == 0,
    "marker": {},
}

for g in PANEL:
    ab = mean_expr(mask_ab, g)
    b  = mean_expr(mask_b, g)
    ov = mean_expr(np.ones(raw.n_obs, dtype=bool), g)
    enr = (ab / b) if (b and not np.isnan(b) and b > 0) else float("nan")
    result["marker"][g] = {
        "adult_beta_mean_log1p": None if np.isnan(ab) else round(ab, 4),
        "beta_mean_log1p":       None if np.isnan(b)  else round(b, 4),
        "overall_mean_log1p":    None if np.isnan(ov) else round(ov, 4),
        "enrichment_adult_beta_vs_beta": None if np.isnan(enr) else round(enr, 3),
    }

if n_adult_beta == 0:
    result["interpretation"] = (
        "VACUOUS: current PANDA-Marker (Jul-23 checkpoint) predicts 0 adult-beta cells on Veres. "
        "The 'beta' cluster (n={}) captures INS/IAPP/MAFB/ADCYAP1 signal instead; adult-vs-juvenile "
        "distinction is not resolved on this dataset. Enrichment ratios below use 0/beta and are NaN."
    ).format(n_beta)
else:
    result["interpretation"] = (
        "adult-beta cluster (n={}) canonical panel enrichment vs beta cluster (n={})."
    ).format(n_adult_beta, n_beta)

with open(OUT / "95_adult_beta_validation.json", "w") as f:
    json.dump(result, f, indent=2)
print(f"[write] {OUT}/95_adult_beta_validation.json", flush=True)
print(json.dumps(result, indent=2))