"""Gene-expression modality: HEST-1k human Visium spots (H&E patch + whole-transcriptome spot expression).

Input  ../HEST/visium_humandisease_adata.pickle            list of 421 AnnData (1% spot subsample per sample)
       ../HEST/UNI/visium_hest_univ2_imagefeature.pickle    list of 421 (n_spot, 1536) UNI2-h spot features (aligned)
       ../HEST/prov-gigapath-conchinside/visium_hest_gigapath_imagefeature.pickle   (GigaPath, aligned)
       ../HEST/HEST_v1_1_0.csv                              sample metadata; list order = Visium/Homo sapiens rows
Kept   samples whose panel covers the common whole-transcriptome genes (targeted ~1k-gene panels dropped).
Groups for CV: study + patient (patient missing -> whole study), i.e. no patient leaks across folds.
Outputs cache/hest/
  meta.csv            one row per spot (sample id, organ, oncotree, group, total counts)
  counts.npz          raw counts (n_spot, n_common_genes) sparse + gene names
  hvg.npy             log1p(CP10k) of the top-2000 HVGs (n_spot, 2000) float32      -> MLP gene encoder input
  he_uni2.npy, he_gigapath.npy   (n_spot, 1, 1536) float16                            -> frozen H&E anchor
"""
import os
import pickle
import weakref

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
from anndata._core import file_backing

from common import CACHE, ROOT

HEST = os.path.join(ROOT, "..", "HEST")
OUT = os.path.join(CACHE, "hest")
os.makedirs(OUT, exist_ok=True)


def _ss(self, state):  # old anndata pickles store '_adata' instead of '_adata_ref'
    self.__dict__ = state.copy()
    ref = state.get("_adata_ref", state.get("_adata"))
    self.__dict__["_adata_ref"] = weakref.ref(ref) if ref is not None else (lambda: None)


file_backing.AnnDataFileManager.__setstate__ = _ss
A = pickle.load(open(os.path.join(HEST, "visium_humandisease_adata.pickle"), "rb"))
FU = pickle.load(open(os.path.join(HEST, "UNI", "visium_hest_univ2_imagefeature.pickle"), "rb"))
FG = pickle.load(open(os.path.join(HEST, "prov-gigapath-conchinside", "visium_hest_gigapath_imagefeature.pickle"), "rb"))
H = pd.read_csv(os.path.join(HEST, "HEST_v1_1_0.csv"))
H = H[(H.st_technology == "Visium") & (H.species == "Homo sapiens")].reset_index(drop=True)
assert len(H) == len(A) == len(FU) == len(FG)

# --- check the list order against the metadata (full-res size from thumbnail / scalefactor, oncotree code)
ok = []
for a, (_, h) in zip(A, H.iterrows()):
    st = a.uns["spatial"]["ST"]
    sh = np.array(st["images"]["downscaled_fullres"].shape[:2]) / st["scalefactors"]["tissue_downscaled_fullres_scalef"]
    size_ok = np.allclose(sh, [h.fullres_px_height, h.fullres_px_width], rtol=0.02)
    onc = a.obs.oncotree_code.iloc[0]
    onc_ok = (pd.isna(onc) and pd.isna(h.oncotree_code)) or onc == h.oncotree_code
    ok.append(size_ok and onc_ok)
print(f"metadata order check: {sum(ok)}/{len(ok)} samples match (image size + oncotree)")
assert np.mean(ok) > 0.95

# --- harmonise gene names to HGNC symbols (samples use symbols, Ensembl ids + SYMBOL column, or 'GRCh38__' prefixes)
e2s = {}
for a in A:
    if "gene_ids" in a.var and not str(a.var_names[0]).startswith(("ENSG", "GRCh")):
        e2s.update(dict(zip(a.var.gene_ids.astype(str), a.var_names)))


def symbols(a):
    if "SYMBOL" in a.var:
        return a.var.SYMBOL.astype(str).tolist()
    out = []
    for x in map(str, a.var_names):
        x = x.split("__")[-1] if x.startswith("GRCh") else x
        out.append(e2s.get(x.split(".")[0], x) if x.startswith("ENSG") else x)
    return out


for a in A:
    a.var_names = symbols(a)

# --- common gene set over whole-transcriptome samples
wt = [i for i, a in enumerate(A) if a.n_vars >= 17000]
common = set(A[wt[0]].var_names)
for i in wt[1:]:
    common &= set(A[i].var_names)
common = sorted(g for g in common if not g.startswith(("MT-", "RPL", "RPS")))
print(f"whole-transcriptome samples: {len(wt)}/{len(A)}; common genes (no MT/ribo): {len(common)}")

rows, Xs, fu, fg = [], [], [], []
for i in wt:
    a, h = A[i], H.iloc[i]
    a = a[:, ~a.var_names.duplicated()]
    X = a[:, common].X
    Xs.append(sp.csr_matrix(X, dtype=np.float32))
    fu.append(FU[i].numpy()); fg.append(FG[i].numpy())
    study = str(h.dataset_title)
    grp = f"{study}|{h.patient}" if pd.notna(h.patient) else study
    for b in a.obs_names:
        rows.append({"sample": h.id, "spot": b, "organ": h.organ, "oncotree": h.oncotree_code,
                     "disease": h.disease_state, "study": study, "patient": h.patient, "group": grp})
M = pd.DataFrame(rows)
X = sp.vstack(Xs).tocsr()
M["total_counts"] = np.asarray(X.sum(1)).ravel()
keep = (M.total_counts >= 200).values
M, X = M[keep].reset_index(drop=True), X[keep]
fu, fg = np.concatenate(fu)[keep], np.concatenate(fg)[keep]
M["n_genes"] = np.asarray((X > 0).sum(1)).ravel()

# log1p CP10k, HVG = top-2000 by variance of the log-normalised values (genes expressed in >=1% spots)
L = sp.diags(1e4 / M.total_counts.values) @ X
L.data = np.log1p(L.data)
frac = np.asarray((X > 0).mean(0)).ravel()
mu = np.asarray(L.mean(0)).ravel()
var = np.asarray(L.multiply(L).mean(0)).ravel() - mu ** 2
var[frac < 0.01] = -1
hvg = np.sort(np.argsort(-var)[:2000])
np.save(os.path.join(OUT, "hvg.npy"), L[:, hvg].toarray().astype(np.float32))
sp.save_npz(os.path.join(OUT, "counts.npz"), X)
pd.Series(common).to_csv(os.path.join(OUT, "genes.csv"), index=False, header=["gene"])
pd.Series(np.array(common)[hvg]).to_csv(os.path.join(OUT, "hvg_genes.csv"), index=False, header=["gene"])
np.save(os.path.join(OUT, "he_uni2.npy"), fu.astype(np.float16)[:, None])
np.save(os.path.join(OUT, "he_gigapath.npy"), fg.astype(np.float16)[:, None])
M.to_csv(os.path.join(OUT, "meta.csv"), index_label="idx")
print(f"spots kept: {len(M)} (dropped {int((~keep).sum())} with <200 counts); samples {M['sample'].nunique()}, "
      f"groups {M.group.nunique()}, studies {M.study.nunique()}")
print("organs:", M.organ.value_counts().to_dict())
print("median counts / genes per spot:", M.total_counts.median(), M.n_genes.median())