"""Frozen Cell2Sentence (C2S-Scale-Gemma-2-2B) spot embeddings for the RNA sets of rna_sets.py (spot = pseudo-cell).

    python 43_rnafm_c2s.py [set ...]          (default: all of rna_sets.SETS)

Model     : HF vandijklab/C2S-Scale-Gemma-2-2B @ 5ddf28b8f1c81b7ab7a9be192924da82b6c5d512 (cc-by-4.0), a Gemma-2 2B
            decoder (26 layers, d=2304) instruction-tuned on cell sentences from >800 CellxGene/HCA datasets. Loaded as
            a plain transformers Gemma2Model (no remote code, strict: no missing / unexpected keys), bf16, eager
            attention (keeps Gemma-2's attention-logit soft-capping exactly as in training).
Tokenizer : the repo ships only the sentencepiece tokenizer.model; sentencepiece is not in the shared env, so it was
            converted once to ext/cell2sentence/tokenizer/fast/tokenizer.json with transformers' own GemmaConverter
            inside an isolated venv (ext/cell2sentence/venv) and checked token-identical to the sentencepiece
            GemmaTokenizer on sampled prompts of every set.
Cell sentence (exactly as cell2sentence src/cell2sentence/utils.py generate_sentences + prompt_formatter.py, repo
            vandijklab/cell2sentence @ 6599c9d, read-only clone in ext/cell2sentence/src):
            raw counts -> sc.pp.normalize_total (target = median total counts of the set) -> log1p base 10
            (tutorial 0; rank-preserving within a spot) -> gene symbols upper-cased (generate_vocabulary) ->
            nonzero genes of the spot, shuffled with np.random.seed(42) / sklearn.utils.shuffle and stable-sorted by
            descending expression (= their random tie-breaking) -> top K=200 symbols (default n_genes=200 of
            tasks.embed_cells / predict_cell_types_of_data and the tutorials), space-separated.
            Not applied: tutorial-0 QC filters (filter_cells min_genes=200, filter_genes min_cells=3), because they
            would drop spots / make a spot's sentence depend on the other spots of the set.
Prompt    : the cell-type-prediction template used by tasks.embed_cells (C2SPromptFormatter task
            "cell_type_prediction"); that formatter picks one of 15 paraphrases at random, we fix the first one, which
            is also the template on the C2S-Scale model card:
            "The following is a list of {K} gene names ordered by descending expression level in a Homo sapiens cell.
             Your task is to give the cell type which this cell belongs to based on its gene expression.\\n
             Cell sentence: {sentence}.\\nThe cell type corresponding to these genes is:"   (+ <bos>, no chat template)
            {K} is always 200, as in the formatter (spots with < 200 expressed genes keep "200").
Readout   : last hidden layer (post final RMSNorm = hidden_states[-1], as CSModel.embed_cells_batched).
            "c2s"      : mean over the cell-sentence tokens only (tokens whose char span overlaps the gene list;
                         official code averages over the whole padded prompt, which mixes in the constant template
                         and padding).
            "c2s-last" : hidden state of the last prompt token (":" after "is"), i.e. the state that would generate
                         the cell type.
Output    : cache/rnafm/c2s_<set>.npy and c2s-last_<set>.npy, (n_spot, 2304) float16, set order. Log logs/rnafm/c2s.log.
"""
import os
import sys
import time

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
from sklearn.utils import shuffle

from common import CACHE, ROOT
from rna_sets import SETS, load, save

REPO = "vandijklab/C2S-Scale-Gemma-2-2B"
REV = "5ddf28b8f1c81b7ab7a9be192924da82b6c5d512"
TOK_DIR = os.path.join(ROOT, "ext", "cell2sentence", "tokenizer", "fast")
K = 200
ORGANISM = "Homo sapiens"
PREFIX = (f"The following is a list of {K} gene names ordered by descending expression level in a {ORGANISM} cell. "
          "Your task is to give the cell type which this cell belongs to based on its gene expression.\nCell sentence: ")
SUFFIX = ".\nThe cell type corresponding to these genes is:"
GPU_GB = 8.0
TOK_BUDGET = 8 * 800           # tokens per forward pass (halved on OOM; ~7 GB peak at L~750)
dev = "cuda"


def log(*a):
    print(*a, flush=True)


def cell_sentences(X, genes):
    """C2S normalisation + rank -> list of top-K space-separated (upper-case) gene symbols, plus #expressed genes."""
    X = X.tocsr().astype(np.float64)
    tot = np.asarray(X.sum(1)).ravel()
    target = np.median(tot[tot > 0])                       # sc.pp.normalize_total default target_sum=None
    Xn = sp.diags(target / np.maximum(tot, 1e-12)) @ X
    Xn = Xn.tocsr()
    Xn.data = np.log1p(Xn.data) / np.log(10)               # sc.pp.log1p(base=10)
    enc = np.char.upper(np.asarray(genes, dtype=str))
    np.random.seed(42)                                     # generate_sentences(random_state=42)
    sents, n_expr = [], np.zeros(Xn.shape[0], int)
    for i in range(Xn.shape[0]):
        cols = Xn.indices[Xn.indptr[i]:Xn.indptr[i + 1]]
        vals = Xn.data[Xn.indptr[i]:Xn.indptr[i + 1]]
        cols, vals = shuffle(cols, vals)
        order = cols[np.argsort(-vals, kind="stable")]
        n_expr[i] = len(order)
        sents.append(" ".join(enc[order[:K]]))
    return sents, n_expr


def tokenize(tok, sents):
    """Prompt ids + mask of cell-sentence tokens (by char offsets)."""
    prompts = [PREFIX + s + SUFFIX for s in sents]
    enc = tok(prompts, return_offsets_mapping=True, add_special_tokens=True)
    out = []
    for p, s, ids, off in zip(prompts, sents, enc["input_ids"], enc["offset_mapping"]):
        a, b = len(PREFIX), len(PREFIX) + len(s)
        off = np.asarray(off)
        m = (off[:, 0] < b) & (off[:, 1] > a) & (off[:, 1] > off[:, 0])
        assert ids[0] == tok.bos_token_id and m.sum() > 0
        out.append((np.asarray(ids), m))
    return prompts, out


@torch.no_grad()
def embed(model, seqs, pad_id):
    n, D = len(seqs), model.config.hidden_size
    E_mean, E_last = np.zeros((n, D), np.float32), np.zeros((n, D), np.float32)
    order = np.argsort([len(s[0]) for s in seqs])[::-1]   # longest first: OOM shows up immediately
    budget, k, t0, done = TOK_BUDGET, 0, time.time(), 0
    while k < n:
        Lm = len(seqs[order[k]][0])
        bs = max(1, budget // Lm)
        idx = order[k:k + bs]
        ids = torch.full((len(idx), Lm), pad_id, dtype=torch.long)
        msk = torch.zeros((len(idx), Lm), dtype=torch.bool)
        att = torch.zeros((len(idx), Lm), dtype=torch.long)
        for j, i in enumerate(idx):
            s, m = seqs[i]
            ids[j, :len(s)], msk[j, :len(s)], att[j, :len(s)] = torch.from_numpy(s), torch.from_numpy(m), 1
        try:
            h = model(input_ids=ids.to(dev), attention_mask=att.to(dev)).last_hidden_state.float()
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            if budget <= Lm:
                raise
            budget //= 2
            log(f"  OOM -> token budget {budget}")
            continue
        mk = msk.to(dev).float()
        E_mean[idx] = ((h * mk[..., None]).sum(1) / mk.sum(1, keepdim=True)).cpu().numpy()
        last = att.sum(1).to(dev) - 1
        E_last[idx] = h[torch.arange(len(idx), device=dev), last].cpu().numpy()
        k += len(idx)
        if k // 1000 > done:
            done = k // 1000
            log(f"  {k}/{n}  {time.time() - t0:.0f}s  bs={len(idx)} L={Lm}  peak {torch.cuda.max_memory_allocated() / 2**30:.2f} GB")
    return E_mean, E_last


def sanity(E, tag, organ=None):
    rng = np.random.default_rng(0)
    Z = E / np.linalg.norm(E, axis=1, keepdims=True)
    s = Z[rng.choice(len(Z), min(200, len(Z)), replace=False)]
    C = s @ s.T
    mc = C[~np.eye(len(s), dtype=bool)].mean()
    msg = f"[{tag}] finite={np.isfinite(E).all()}  mean pairwise cos(200)={mc:.4f}  std/dim={E.std(0).mean():.4f}"
    if organ is not None:
        Zc = Z - Z.mean(0)
        Zc /= np.linalg.norm(Zc, axis=1, keepdims=True)
        o = np.asarray(organ)
        same, raw = [], []
        for ZZ, acc in ((Zc, same), (Z, raw)):
            Zt = torch.from_numpy(ZZ).to(dev)
            for i in range(0, len(Zt), 2048):
                S = Zt[i:i + 2048] @ Zt.T
                S[torch.arange(S.shape[0]), torch.arange(i, i + S.shape[0])] = -2
                nn = S.topk(10, dim=1).indices.cpu().numpy()
                acc.append((o[nn] == o[i:i + S.shape[0], None]).mean(1))
        cnt = pd.Series(o).value_counts()
        base = ((cnt * (cnt - 1)).sum()) / (len(o) * (len(o) - 1))
        msg += (f"  10-NN same-organ: {np.concatenate(raw).mean():.3f} (cos), {np.concatenate(same).mean():.3f}"
                f" (centred cos)  base rate {base:.3f}")
    log(msg)


def main():
    from transformers import Gemma2Model, PreTrainedTokenizerFast
    from huggingface_hub import snapshot_download
    sets = sys.argv[1:] or SETS
    torch.cuda.set_per_process_memory_fraction(min(1.0, GPU_GB * 2**30 / torch.cuda.get_device_properties(0).total_memory))
    tok = PreTrainedTokenizerFast.from_pretrained(TOK_DIR)
    tok.padding_side = "right"
    mdir = snapshot_download(REPO, revision=REV)
    model, info = Gemma2Model.from_pretrained(mdir, dtype=torch.bfloat16, attn_implementation="eager",
                                              output_loading_info=True)
    log(f"loaded {REPO}@{REV}: missing={info['missing_keys']} unexpected={info['unexpected_keys']} "
        f"mismatched={info['mismatched_keys']}")
    assert not info["missing_keys"] and not info["mismatched_keys"] and not info["unexpected_keys"]
    model.eval().to(dev)
    log(f"weights {torch.cuda.memory_allocated() / 2**30:.2f} GB; prompt prefix tokens {len(tok(PREFIX)['input_ids'])}")
    meta = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv"))
    tot_seq, tot_t = 0, 0.0
    for name in sets:
        t0 = time.time()
        X, genes = load(name)
        sents, n_expr = cell_sentences(X, genes)
        prompts, seqs = tokenize(tok, sents)
        ntok = np.array([len(s[0]) for s in seqs])
        nsent = np.array([m.sum() for _, m in seqs])
        upg = np.char.upper(genes.astype(str))
        nonhgnc = np.mean([("." in g) or g.startswith("ENSG") for g in upg])
        log(f"== {name}: {X.shape[0]} spots x {X.shape[1]} genes (all symbols usable as text; {nonhgnc:.1%} non-HGNC-"
            f"style ids e.g. ENSG / RP11-..), {len(set(upg))} unique upper-case symbols; expressed genes/spot median "
            f"{np.median(n_expr):.0f}, spots with < {K} expressed genes {np.mean(n_expr < K):.1%}; sentence genes/spot "
            f"median {np.median(np.minimum(n_expr, K)):.0f}; tokens/prompt median {np.median(ntok):.0f} "
            f"(sentence tokens {np.median(nsent):.0f}, {np.median(nsent) / K:.2f} tok/gene)")
        log("  example: " + prompts[0][:400].replace("\n", " | ") + " ...")
        with open(os.path.join(ROOT, "ext", "cell2sentence", f"prompts_sample_{name}.txt"), "w") as f:  # tokenizer cross-check
            for i in np.random.default_rng(1).choice(len(prompts), min(200, len(prompts)), replace=False):
                f.write(repr(prompts[i]) + "\n")
        t1 = time.time()
        E_mean, E_last = embed(model, seqs, tok.pad_token_id)
        dt = time.time() - t1
        tot_seq, tot_t = tot_seq + len(seqs), tot_t + dt
        log(f"  forward {dt:.0f}s = {len(seqs) / dt:.2f} seq/s, {ntok.sum() / dt:.0f} tok/s; total {time.time() - t0:.0f}s;"
            f" peak GPU {torch.cuda.max_memory_allocated() / 2**30:.2f} GB")
        save("c2s", name, E_mean)
        save("c2s-last", name, E_last)
        organ = meta.organ.values if name in ("hest", "hest_v") else None
        if organ is not None:
            assert len(organ) == len(E_mean)
        sanity(E_mean, f"c2s {name}", organ)
        sanity(E_last, f"c2s-last {name}", organ)
    log(f"overall throughput {tot_seq / max(tot_t, 1e-9):.2f} seq/s over {tot_seq} sequences ({tot_t / 60:.1f} min)")


if __name__ == "__main__":
    main()