"""Frozen single-cell foundation-model embeddings of HEST Visium spot expression (each spot = one pseudo-cell).

    python 22_extract_gene_fm.py scgpt | geneformer

scgpt      : wanglab/scGPT-human (same re-implementation / strict load as 15_extract_marker_fm.py); input = expressed
             genes among the 2000 HVGs that are in the scGPT vocab, raw counts -> log1p CP10k -> per-spot quantile
             binning (51 bins); at most 1200 genes (random subset, as scgpt.tasks.embed_data max_length=1200);
             <cls> output, L2-normalised.
geneformer : ctheodoris/Geneformer V2-104M; input = rank-value encoding of all common genes (count / total * 1e4 /
             gc104M gene median, descending, nonzero), truncated to 4096 tokens incl. <cls>/<eos>;
             penultimate layer, mean over gene tokens.
Output cache/hest/gene_<name>.npy (n_spot, 1, D) float16
"""
import importlib
import os
import pickle
import sys
import time

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch

from common import CACHE

M15 = importlib.import_module("15_extract_marker_fm")
OUT = os.path.join(CACHE, "hest")
name = sys.argv[1]
dev = "cuda"
X = sp.load_npz(os.path.join(OUT, "counts.npz")).tocsr()
genes = pd.read_csv(os.path.join(OUT, "genes.csv")).gene.values
tot = np.asarray(X.sum(1)).ravel()
t0 = time.time()


def batches(seqs, bs):
    for i in range(0, len(seqs), bs):
        yield seqs[i:i + bs]


if name == "scgpt":
    import json
    from huggingface_hub import hf_hub_download
    args = json.load(open(hf_hub_download(M15.SCGPT_REPO, "args.json")))
    vocab = json.load(open(hf_hub_download(M15.SCGPT_REPO, "vocab.json")))
    sd = torch.load(hf_hub_download(M15.SCGPT_REPO, "best_model.pt"), map_location="cpu", weights_only=True)
    sd = {k.replace("self_attn.Wqkv.", "self_attn.in_proj_"): v for k, v in sd.items()}
    PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"]
    model = M15.ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], PAD)
    model.load_state_dict(sd, strict=True)
    model.eval().to(dev)
    hvg = pd.read_csv(os.path.join(OUT, "hvg_genes.csv")).gene.values
    gi = np.array([i for i, g in enumerate(genes) if g in set(hvg) and g in vocab])
    gid = np.array([vocab[g] for g in genes[gi]])
    L = sp.diags(1e4 / tot) @ X[:, gi]
    L.data = np.log1p(L.data)
    L = L.tocsr()
    rng = np.random.default_rng(0)
    seqs = []
    for r in range(L.shape[0]):
        row = L.getrow(r)
        nz, v = row.indices, row.data
        if len(nz) > 1199:
            keep = np.sort(rng.choice(len(nz), 1199, replace=False))
            nz, v = nz[keep], v[keep]
        vals = M15.scgpt_binning(v, args["n_bins"], rng).astype(np.float32)
        seqs.append((np.r_[CLS, gid[nz]], np.r_[np.float32(args["pad_value"]), vals]))
    print(f"scGPT: {len(gi)} HVGs in vocab; genes/spot median {np.median([len(s[0]) - 1 for s in seqs])}", flush=True)
    embs = []
    with torch.no_grad():
        for b in batches(seqs, 32):
            Lm = max(len(g) for g, _ in b)
            g = torch.full((len(b), Lm), PAD, dtype=torch.long)
            v = torch.full((len(b), Lm), float(args["pad_value"]))
            for j, (gg, vv) in enumerate(b):
                g[j, :len(gg)], v[j, :len(vv)] = torch.from_numpy(gg), torch.from_numpy(vv)
            g, v = g.to(dev), v.to(dev)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                h = model.encode(g, v, g.eq(PAD))
            embs.append(h[:, 0].float().cpu())
    E = torch.cat(embs).numpy()
    E /= np.linalg.norm(E, axis=1, keepdims=True)
else:
    from transformers import BertModel
    tokd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/token_dictionary_gc104M.pkl"), "rb"))
    medd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_median_dictionary_gc104M.pkl"), "rb"))
    n2id = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_name_id_dict_gc104M.pkl"), "rb"))
    mdir = os.path.dirname(M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/model.safetensors"))
    M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/config.json")
    gi = np.array([i for i, g in enumerate(genes) if g in n2id and n2id[g] in tokd and n2id[g] in medd])
    tok = np.array([tokd[n2id[g]] for g in genes[gi]])
    med = np.array([medd[n2id[g]] for g in genes[gi]])
    CLS, EOS, PAD = tokd["<cls>"], tokd["<eos>"], tokd["<pad>"]
    model = BertModel.from_pretrained(mdir, add_pooling_layer=False, attn_implementation="sdpa").eval().to(dev)
    Xg = (sp.diags(1e4 / tot) @ X[:, gi]).tocsr()
    seqs = []
    for r in range(Xg.shape[0]):
        row = Xg.getrow(r)
        nz, v = row.indices, row.data / med[row.indices]
        order = nz[np.argsort(-v, kind="stable")][:4094]
        seqs.append([CLS] + tok[order].tolist() + [EOS])
    print(f"Geneformer: {len(gi)} genes in dictionary; tokens/spot median {np.median([len(s) for s in seqs])}", flush=True)
    order_len = np.argsort([len(s) for s in seqs])        # length-sorted batches (less padding), restored below
    embs = np.zeros((len(seqs), model.config.hidden_size), np.float32)
    with torch.no_grad():
        for k in range(0, len(seqs), 8):
            idx = order_len[k:k + 8]
            b = [seqs[i] for i in idx]
            Lm = max(map(len, b))
            ids = torch.full((len(b), Lm), PAD, dtype=torch.long)
            for j, s in enumerate(b):
                ids[j, :len(s)] = torch.tensor(s)
            ids = ids.to(dev)
            att = (ids != PAD).long()
            with torch.autocast("cuda", dtype=torch.bfloat16):
                h = model(input_ids=ids, attention_mask=att, output_hidden_states=True).hidden_states[-2].float()
            gm = att.clone().float()
            gm[:, 0] = 0
            gm[torch.arange(len(b)), att.sum(1) - 1] = 0
            embs[idx] = ((h * gm[..., None]).sum(1) / gm.sum(1, keepdim=True).clamp(min=1)).cpu().numpy()
            if k % 800 == 0:
                print(f"  {k}/{len(seqs)} {time.time() - t0:.0f}s", flush=True)
    E = embs
np.save(os.path.join(OUT, f"gene_{name}.npy"), E.astype(np.float16)[:, None])
print(f"[{name}] saved {E.shape}  finite={np.isfinite(E).all()}  {time.time() - t0:.0f}s")