"""per-class 20x20 hessian of prototype cosine over top-attributed genes."""
from __future__ import annotations
from pathlib import Path
import warnings, json, pickle, sys, numpy as np, pandas as pd, torch
warnings.filterwarnings("ignore")
import anndata as ad, scanpy as sc, scipy.sparse as sp
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 = Path(f"{ROOT_STR}/checkpoints")
CORP = Path(f"{ROOT_STR}/data/corpus")
OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True)
TOP_K_GENES = 20


def compute_hessian_for_class(model, protos, class_idx, x_pca_c, pca_V, mu, sig, gene_idx_subset):
    # parametrise perturbation as delta in gene-z space of the D selected genes
    V_sub = torch.from_numpy(pca_V[:, gene_idx_subset]).float().to(DEVICE)   # (50, D)

    def f_of_delta(delta):
        x = x_pca_c + V_sub @ delta
        x = x.clamp(-10, 10).unsqueeze(0)
        aux = torch.zeros(1, 2, device=DEVICE)
        h = model.trunk(x)
        z_raw = model.projection(h)
        z = torch.nn.functional.normalize(z_raw, dim=1)
        return (z * protos[class_idx].unsqueeze(0)).sum()

    D = len(gene_idx_subset)
    delta0 = torch.zeros(D, device=DEVICE)
    H = torch.autograd.functional.hessian(f_of_delta, delta0)
    return H.detach().cpu().numpy()


def main(sys):
    print(f"\n===== {sys} =====", flush=True)
    stats = np.load(CORP / sys / "harmonized/corpus_stats.npz", allow_pickle=True)
    hvgs = [str(g) for g in stats["shared_hvgs"]]
    mu = np.asarray(stats["mean"], dtype=np.float32)
    sig = np.asarray(stats["std"], dtype=np.float32)
    pca = pickle.load(open(CORP / sys / "harmonized/pca_basis.pkl", "rb"))

    ck = torch.load(CKPT / sys / "marker" / "panda_final.pt", map_location=DEVICE, weights_only=False)
    model = PANDAEncoder(variant="marker", n_pca=50, n_markers=len(ck.get("marker_genes",[])), n_classes=len(ck["classes"]), n_sub=3,
                         n_datasets=len(ck["datasets"])).to(DEVICE).eval()
    model.load_state_dict(ck["model"])
    classes = ck["classes"]
    protos = torch.from_numpy(ck["prototypes"]).to(DEVICE)
    protos = protos / (protos.norm(dim=1, keepdim=True) + 1e-8)

    corp = ad.read_h5ad(CORP / sys / "harmonized/corpus.h5ad")
    label_key = "canonical_label" if "canonical_label" in corp.obs else "cell_type"
    y = corp.obs[label_key].astype(str).values
    hvg2i = {g: i for i, g in enumerate(hvgs)}
    common = [g for g in corp.var_names.astype(str) if g in hvg2i]
    a = corp[:, common].copy()
    sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a)
    X = a.X.toarray().astype(np.float32) if sp.issparse(a.X) else a.X.astype(np.float32)
    Xf = np.zeros((corp.n_obs, len(hvgs)), dtype=np.float32)
    cols = np.array([hvg2i[g] for g in common])
    Xf[:, cols] = X

    A = np.load(OUT / f"80_{sys}_gene_attribution_full.npy")

    K = len(classes)
    H_all = np.zeros((K, TOP_K_GENES, TOP_K_GENES), dtype=np.float32)
    top_gene_names_per_class = []
    pair_rows = []

    for ci, cls_name in enumerate(classes):
        mask = y == cls_name
        if mask.sum() == 0:
            print(f"[!] {cls_name}: no cells in corpus", flush=True)
            continue
        x_gene_log_mean = Xf[mask].mean(axis=0)
        x_z = np.clip((x_gene_log_mean - mu) / sig, -10, 10)
        x_pca_c = torch.from_numpy(pca.transform(x_z.reshape(1, -1))[0]).float().to(DEVICE)

        att_c = A[ci]
        gene_idx = np.argsort(-np.abs(att_c))[:TOP_K_GENES]
        top_names = [hvgs[i] for i in gene_idx]
        top_gene_names_per_class.append(top_names)

        H = compute_hessian_for_class(model, protos, ci, x_pca_c,
                                       pca.components_.astype(np.float32), mu, sig,
                                       gene_idx.tolist())
        H_all[ci] = H

        off = H.copy()
        np.fill_diagonal(off, 0)
        rows, cols_ = np.triu_indices(TOP_K_GENES, k=1)
        vals = off[rows, cols_]
        order = np.argsort(-np.abs(vals))[:30]
        for r_i in order:
            gi, gj = int(rows[r_i]), int(cols_[r_i])
            pair_rows.append({
                "class": cls_name,
                "gene_a": top_names[gi],
                "gene_b": top_names[gj],
                "hessian_off_diag": float(vals[r_i]),
                "abs_h": float(abs(vals[r_i])),
                "attribution_a": float(att_c[gene_idx[gi]]),
                "attribution_b": float(att_c[gene_idx[gj]]),
            })

        print(f"[{cls_name}] Hessian |diag|_max={float(np.abs(np.diag(H)).max()):.4f}  "
              f"|offdiag|_max={float(np.abs(off).max()):.4f}  "
              f"top-3 pairs: {', '.join([f'{top_names[int(rows[r_i])]}·{top_names[int(cols_[r_i])]}({vals[r_i]:+.4f})' for r_i in order[:3]])}",
              flush=True)

    np.save(OUT / f"85_{sys}_hessian_top20.npy", H_all)
    with open(OUT / f"85_{sys}_hessian_top20_genes.json", "w") as f:
        json.dump({classes[i]: top_gene_names_per_class[i] for i in range(K)}, f, indent=2)
    pd.DataFrame(pair_rows).to_csv(OUT / f"85_{sys}_hessian_pairs.csv", index=False)
    print(f"[wrote] {OUT}/85_{sys}_hessian_*", flush=True)


for sys in ["pan_skin", "hematopoiesis", "pancreas"]:
    try:
        main(sys)
    except Exception as e:
        import traceback; traceback.print_exc()
        print(f"[!] {sys}: {e}", flush=True)

print("\n=== DONE ===", flush=True)