"""test if veres alpha-pool polyhormonal (Ins+/Gcg+/Sst+) cells form a distinct sub-cluster vs graded."""
from __future__ import annotations

import json
import warnings
from pathlib import Path

import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc

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)

VERES = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad"
PRED  = ROOT / "discovery/pancreas/marker/veres_predictions.csv"


def col(sub, g):
    if g not in sub.var_names:
        return np.zeros(sub.n_obs)
    j = sub.var_names.get_loc(g)
    x = sub.X[:, j]
    if hasattr(x, "toarray"):
        x = x.toarray()
    return np.asarray(x).ravel()


def main():
    print("[load]", flush=True)
    a = ad.read_h5ad(VERES)
    pred = pd.read_csv(PRED)
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    a.obs["pred_label"] = pd.Categorical(
        [pred_map.get(c, "unknown") for c in a.obs_names])

    sub = a[a.obs["pred_label"].astype(str).isin(
        ["alpha_progenitor", "alpha"])].copy()
    print(f"[filter] {sub.n_obs} alpha-pool cells", flush=True)

    ins1 = col(sub, "Ins1"); ins2 = col(sub, "Ins2")
    gcg  = col(sub, "Gcg");  sst  = col(sub, "Sst"); iapp = col(sub, "Iapp")

    ins_level = ins1 + ins2
    q_ins = np.quantile(ins_level, 0.75)
    q_gcg = np.quantile(gcg,       0.75)
    q_sst = np.quantile(sst,       0.75)

    n_pos = (
        (ins_level >= q_ins).astype(int)
        + (gcg      >= q_gcg).astype(int)
        + (sst      >= q_sst).astype(int)
    )
    sub.obs["INS_level"] = ins_level
    sub.obs["GCG_level"] = gcg
    sub.obs["SST_level"] = sst
    sub.obs["n_hormones_positive"] = n_pos

    print("[cluster] PCA + Leiden", flush=True)
    sc.pp.highly_variable_genes(sub, n_top_genes=2000, flavor="seurat_v3",
                                subset=False, batch_key=None)
    sc.pp.pca(sub, n_comps=30)
    sc.pp.neighbors(sub, n_neighbors=15, n_pcs=30)
    sc.tl.leiden(sub, resolution=0.5, random_state=0, key_added="leiden_alpha")

    df = sub.obs[[
        "pred_label", "paper_label", "leiden_alpha",
        "INS_level", "GCG_level", "SST_level", "n_hormones_positive"
    ]].copy()
    df["Ins1"] = ins1; df["Ins2"] = ins2; df["Gcg"] = gcg
    df["Sst"] = sst;   df["Iapp"] = iapp
    df.reset_index().rename(columns={"index": "cell_id"}).to_csv(
        OUT / "110_veres_polyhormonal_alpha_scores.csv", index=False)

    baseline_polyhormonal = float((df["n_hormones_positive"] >= 2).mean())
    per_clus = df.groupby("leiden_alpha", observed=True).agg(
        n_cells=("n_hormones_positive", "size"),
        frac_polyhormonal=("n_hormones_positive",
                           lambda s: float((s >= 2).mean())),
        frac_gcg_hi=("GCG_level",
                     lambda s: float((s >= q_gcg).mean())),
        frac_ins_hi=("INS_level",
                     lambda s: float((s >= q_ins).mean())),
        frac_sst_hi=("SST_level",
                     lambda s: float((s >= q_sst).mean())),
        mean_gcg=("GCG_level", "mean"),
        mean_ins=("INS_level", "mean"),
        mean_sst=("SST_level", "mean"),
    ).sort_values("frac_polyhormonal", ascending=False).reset_index()
    per_clus["enrichment_vs_baseline"] = per_clus["frac_polyhormonal"] \
                                          / max(baseline_polyhormonal, 1e-6)
    per_clus.to_csv(OUT / "110_veres_polyhormonal_alpha_per_cluster.csv",
                    index=False)

    n_2x_clusters = int((per_clus["enrichment_vs_baseline"] >= 2.0).sum())
    verdict = ("distinct_polyhormonal_subcluster" if n_2x_clusters in (1, 2)
               else "graded_phenotype" if n_2x_clusters == 0
               else "diffuse_enrichment")

    summary = {
        "n_alpha_pool":               int(sub.n_obs),
        "baseline_polyhormonal_frac": round(baseline_polyhormonal, 4),
        "q75_thresholds":             {"Ins": float(q_ins),
                                       "Gcg": float(q_gcg),
                                       "Sst": float(q_sst)},
        "leiden_resolution":          0.5,
        "n_clusters":                 int(per_clus["leiden_alpha"].nunique()),
        "n_clusters_enriched_2x":     n_2x_clusters,
        "verdict":                    verdict,
        "per_cluster":                per_clus.to_dict("records"),
    }
    with open(OUT / "110_veres_polyhormonal_alpha_summary.json", "w") as f:
        json.dump(summary, f, indent=2)

    print("\n[done] verdict:", verdict, flush=True)
    print(per_clus.round(3).to_string(index=False))


if __name__ == "__main__":
    main()