"""split veres beta-predicted cells into mature (MAFA/UCN3 hi) vs SC-beta by INS x MAT quadrants."""
from __future__ import annotations

import json
import warnings
from pathlib import Path

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

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])))
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 z(x):
    x = np.asarray(x, dtype=float)
    s = x.std()
    return (x - x.mean()) / (s if s > 0 else 1.0)


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])

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

    # veres X is already log-normalised (range 0..9)
    def col(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()

    ins1 = col("Ins1"); ins2 = col("Ins2"); iapp = col("Iapp")
    mafa = col("Mafa"); ucn3 = col("Ucn3")
    insulin = ins1 + ins2
    mature = z(mafa) + z(ucn3)

    # split thresholds are the median within the veres beta-predicted set
    ins_thr = float(np.median(insulin))
    mat_thr = float(np.median(mature))
    ins_hi  = insulin > ins_thr
    mat_hi  = mature  > mat_thr

    quadrant = np.array(
        ["INS+/MAT+" if (ih and mh) else
         "INS+/MAT-" if (ih and not mh) else
         "INS-/MAT+" if (not ih and mh) else
         "INS-/MAT-"
         for ih, mh in zip(ins_hi, mat_hi)])
    sub.obs["INS_level"]     = insulin
    sub.obs["MAT_score"]     = mature
    sub.obs["quadrant"]      = pd.Categorical(quadrant)

    df = sub.obs[["quadrant", "INS_level", "MAT_score"]].copy()
    df["Ins1"] = ins1; df["Ins2"] = ins2; df["Iapp"] = iapp
    df["Mafa"] = mafa; df["Ucn3"] = ucn3
    df["paper_label"] = sub.obs["paper_label"].astype(str).values
    df.reset_index().rename(columns={"index": "cell_id"}).to_csv(
        OUT / "109_veres_mature_beta_scores.csv", index=False)

    counts = df["quadrant"].value_counts().to_dict()
    paper_by_quad = df.groupby(["quadrant", "paper_label"], observed=True) \
                      .size().unstack(fill_value=0)
    paper_by_quad.to_csv(OUT / "109_veres_mature_beta_paper_x_quadrant.csv")

    n_total = int(sub.n_obs)
    n_mature = int((df["quadrant"] == "INS+/MAT+").sum())
    n_scbeta = int((df["quadrant"] == "INS+/MAT-").sum())

    fraction_mature = n_mature / n_total if n_total else 0.0

    means = df.groupby("quadrant", observed=True)[
        ["Mafa", "Ucn3", "Ins1", "Ins2", "Iapp"]].mean().round(3).to_dict()

    summary = {
        "n_beta_predicted": n_total,
        "ins_threshold":    ins_thr,
        "mat_threshold":    mat_thr,
        "quadrant_counts":  counts,
        "n_mature_INS+MAT+":  n_mature,
        "n_SCbeta_INS+MAT-":  n_scbeta,
        "fraction_mature":    round(fraction_mature, 4),
        "mean_expression_per_quadrant": means,
        "paper_label_x_quadrant": {
            q: paper_by_quad.loc[q].to_dict()
            for q in paper_by_quad.index
        } if len(paper_by_quad) else {},
    }
    with open(OUT / "109_veres_mature_beta_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    print("\n[done]")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()