#!/usr/bin/env python
"""Parity check: standalone hmoe_da vs the study's cached Kamath predictions.

The cached file Kamath_DA_predicted.h5ad holds `pred_subtype` produced by the
original torch/awatramani_lab path with Gad2 EXCLUDED. So the standalone package
run with include_gad2=False must reproduce those labels exactly, and the stored
obsm['X_hmoe_P_sub'] to <1e-5.
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import scanpy as sc

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from hmoe_da import predict  # noqa: E402

REPO = Path(__file__).resolve().parents[2]
RAW = REPO / "Data" / "Human" / "Kamath" / "Human_Kamath_Dopamine.h5ad"
CACHED = REPO / "Data" / "Human" / "Kamath" / "Kamath_DA_predicted.h5ad"


def main() -> int:
    raw = sc.read_h5ad(RAW)
    cached = sc.read_h5ad(CACHED)
    print(f"raw {raw.shape}  cached {cached.shape}")

    # Align cached rows to raw order by barcode.
    idx = cached.obs_names.get_indexer(raw.obs_names)
    assert (idx >= 0).all(), "some raw barcodes missing from cached file"
    ref_sub = cached.obs["pred_subtype"].values[idx]
    ref_P = np.asarray(cached.obsm["X_hmoe_P_sub"])[idx]

    ad = predict(raw.copy(), include_gad2=False, verbose=True)
    got_sub = ad.obs["hmoe_subtype"].values
    got_P = np.asarray(ad.obsm["X_hmoe_P_sub"])

    max_dP = float(np.max(np.abs(got_P - ref_P)))
    label_match = float(np.mean(got_sub == ref_sub))
    n_mismatch = int(np.sum(got_sub != ref_sub))
    print(f"\nP_sub max abs diff : {max_dP:.2e}")
    print(f"label agreement    : {label_match*100:.4f}%  ({n_mismatch} mismatched)")

    ok = (max_dP < 1e-5) and (n_mismatch == 0)
    if not ok and n_mismatch:
        # show a few mismatches for diagnosis
        mm = np.where(got_sub != ref_sub)[0][:10]
        for i in mm:
            print(f"  cell {i}: got {got_sub[i]!r} vs ref {ref_sub[i]!r} "
                  f"(dP={np.max(np.abs(got_P[i]-ref_P[i])):.2e})")

    # sanity: include_gad2=True should surface some Gad2 calls
    ad2 = predict(raw.copy(), include_gad2=True, verbose=False)
    n_gad = int((ad2.obs["hmoe_family"] == "Gad2").sum())
    print(f"include_gad2=True  : {n_gad} cells called Gad2")

    print("\nPARITY", "PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())