"""Frozen Arc Institute STATE embedding (SE-600M) spot embeddings of the 5 RNA sets (each spot = one pseudo-cell).

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

Model     : arcinstitute/SE-600M (HF), revision 5a9a80f44f7ce32ce57059933ef0d735d7c10ce5, files model.safetensors
            (metadata: converted from se600m_epoch16.ckpt, i.e. the final 16-epoch checkpoint), config.yaml and
            protein_embeddings.pt (19,790 gene symbols -> ESM2-15B mean protein embedding, 5120-d; identical, in
            the same order, to the checkpoint's pe_embedding.weight -- asserted). Licence: Arc Research Institute
            State Model Non-Commercial License + Acceptable Use Policy (weights; non-commercial research use only),
            code CC BY-NC-SA 4.0. Official code: github.com/ArcInstitute/state (commit 9bbfe78a, cloned to
            ext/state/state, read but NOT imported / installed): the forward pass of src/state/emb/nn/model.py
            (StateEmbeddingModel) + nn/flash_transformer.py is re-implemented below and loaded with a strict
            state-dict load of all 238 tensors (the encoder / gene_embedding_layer tie is asserted); the unused
            heads (binary_decoder, dataset_encoder, gene_embedding_layer) are dropped after loading.
            16-layer post-LN transformer, d=2048, 16 heads, FFN 2048, GELU, no positional encoding, no attention mask.
Inputs    : rna_sets.load(set) raw(-like) counts + HGNC symbols. Genes = the set's genes whose symbol is a key of
            protein_embeddings.pt (official FilteredGenesCounts valid mask, set gene order kept); other genes dropped.
Preprocess: official `state emb transform` path = Inference.encode_adata -> create_dataloader ->
            VCIDatasetSentenceCollator.sample_cell_sentences (is_train=False):
              * counts are raw -> log1p (the official per-cell raw-vs-log1p heuristic, max>35 or sum(expm1)>5e6, is
                bypassed because all inputs are known raw counts; how many spots it would misclassify is logged);
              * expression weight w = log1p / sum(log1p) over the valid genes (uniform if the spot is all-zero);
              * cell sentence (pad_length 2048) = [CLS] + 2047 genes ranked by descending expression with random
                tie-breaking (every set has > 2047 valid genes, so: all detected genes in rank order, top 2047 if more,
                then randomly drawn undetected genes without replacement); seeded numpy RNG (seed 0) instead of
                torch.randperm;
              * token count value = 100 * w[gene]; position 0 carries 100 * w[local gene index 3] exactly as the
                official collator (cls_token_idx = 3 is written as a local gene index before the count lookup);
              * token embedding = L2-normalised ESM2 embedding, position 0 replaced by the learned cls_token, the
                learned dataset_token appended (count embedding 0) -> 2049 tokens.
Forward   : encoder(Linear 5120->2048, LayerNorm, SiLU) * sqrt(2048) + soft-binned count embedding
            (softmax(count_encoder(c)) @ bin_encoder) -> 16 layers -> decoder (SkipBlock + Linear).
            Whole model in bf16 (official get_precision_config('cuda') casts the model to bf16 + bf16 autocast).
Readout   : official embedding = L2-normalised decoder output at the CLS position (StateEmbeddingModel.forward
            `embedding`), D = 2048. `state emb transform` additionally concatenates a 10-d dataset-token vector
            (dataset_embedder(decoder output at the dataset token)) -> 2058-d X_state; that batch-identity part is NOT
            saved here (only the 2048-d cell CLS embedding).
Output    : rna_sets.save("se600m", set, E) -> cache/rnafm/se600m_<set>.npy (n_spot, 2048) float16.
Sanity    : finite, mean pairwise cosine of 200 random spots, (hest) organ agreement of cosine 10-NN vs base rate,
            and (hest, 300 spots) cosine between two runs with different tie-breaking seeds (stochasticity of the
            undetected-gene filler tokens).
"""
import math
import os
import sys
import time

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import yaml
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

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

MODEL = "se600m"
REPO, REV = "arcinstitute/SE-600M", "5a9a80f44f7ce32ce57059933ef0d735d7c10ce5"
dev = "cuda"
DT = torch.bfloat16


# ----------------------------------------------------------------------------- model (re-implemented from ext/state)
class SkipBlock(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.intermediate_dense = nn.Linear(d, 2 * d)
        self.dense = nn.Linear(2 * d, d)
        self.layer_norm = nn.LayerNorm(d)

    def forward(self, x):
        return self.layer_norm(x + self.dense(F.relu(self.intermediate_dense(x))))


class Layer(nn.Module):                    # FlashTransformerEncoderLayer (post-LN, dropout 0 at inference)
    def __init__(self, d, nhead, d_hid):
        super().__init__()
        self.nhead = nhead
        self.qkv_proj = nn.Linear(d, 3 * d)
        self.out_proj = nn.Linear(d, d)
        self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.linear1, self.linear2 = nn.Linear(d, d_hid), nn.Linear(d_hid, d)

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
        q, k, v = (t.view(B, T, self.nhead, D // self.nhead).transpose(1, 2) for t in (q, k, v))
        a = F.scaled_dot_product_attention(q, k, v, is_causal=False).transpose(1, 2).reshape(B, T, D)
        x = self.norm1(x + self.out_proj(a))
        return self.norm2(x + self.linear2(F.gelu(self.linear1(x))))


class TransformerEncoder(nn.Module):
    def __init__(self, layers):
        super().__init__()
        self.layers = nn.ModuleList(layers)


class SE(nn.Module):
    def __init__(self, token_dim, d, nhead, d_hid, nlayers, out_dim, n_pe, z_dim, n_ds):
        super().__init__()
        self.d = d
        self.cls_token = nn.Parameter(torch.zeros(1, token_dim))
        self.dataset_token = nn.Parameter(torch.zeros(1, token_dim))
        self.encoder = nn.Sequential(nn.Linear(token_dim, d), nn.LayerNorm(d), nn.SiLU())
        self.gene_embedding_layer = self.encoder                       # tied, as in the official model
        self.transformer_encoder = TransformerEncoder([Layer(d, nhead, d_hid) for _ in range(nlayers)])
        self.decoder = nn.Sequential(SkipBlock(d), nn.Linear(d, out_dim))
        self.binary_decoder = nn.Sequential(SkipBlock(out_dim + d + z_dim), SkipBlock(out_dim + d + z_dim),
                                            nn.Linear(out_dim + d + z_dim, 1))
        self.bin_encoder = nn.Embedding(10, d)
        self.count_encoder = nn.Sequential(nn.Linear(1, 512), nn.LeakyReLU(), nn.Linear(512, 10))
        self.dataset_embedder = nn.Linear(out_dim, 10)
        self.dataset_encoder = nn.Sequential(nn.Linear(out_dim, d), nn.SiLU(), nn.LayerNorm(d), nn.Dropout(0.1),
                                             nn.Linear(d, n_ds))
        self.pe_embedding = nn.Embedding(n_pe, token_dim)

    def forward(self, tok, cnt):
        """tok: (B, L) global gene ids (position 0 = placeholder for CLS); cnt: (B, L) count values."""
        x = F.normalize(self.pe_embedding(tok), dim=2)
        x[:, 0, :] = self.cls_token.expand(x.size(0), -1)
        x = torch.cat([x, self.dataset_token.expand(x.size(0), -1).unsqueeze(1)], 1)
        x = self.encoder(x) * math.sqrt(self.d)
        w = F.softmax(self.count_encoder(cnt.unsqueeze(-1)), dim=-1)
        ce = w @ self.bin_encoder.weight
        x = x + torch.cat([ce, torch.zeros_like(ce[:, :1])], 1)
        for lyr in self.transformer_encoder.layers:
            x = lyr(x)
        return F.normalize(self.decoder(x)[:, 0, :], dim=1)


def build():
    cfg = yaml.safe_load(open(hf_hub_download(REPO, "config.yaml", revision=REV)))
    m, ds = cfg["model"], cfg["dataset"]
    assert m["counts"] and m["rda"] and m["dataset_correction"] and ds["pad_length"] == 2048
    assert ds["cls_token_idx"] == 3 and cfg["loss"]["name"] == "tabular"
    pe = torch.load(hf_hub_download(REPO, "protein_embeddings.pt", revision=REV), map_location="cpu",
                    weights_only=True)
    sd = load_file(hf_hub_download(REPO, "model.safetensors", revision=REV))
    n_ds = cfg["dataset"][cfg["dataset"]["current"]]["num_datasets"]
    model = SE(cfg["tokenizer"]["token_dim"], m["emsize"], m["nhead"], m["d_hid"], m["nlayers"], m["output_dim"],
               len(pe), 1 + 10, n_ds)
    for k in ["0.weight", "0.bias", "1.weight", "1.bias"]:
        assert torch.equal(sd[f"encoder.{k}"], sd[f"gene_embedding_layer.{k}"])
    model.load_state_dict(sd, strict=True)
    assert torch.equal(model.pe_embedding.weight, torch.vstack(list(pe.values())))   # CLI's protein_embeddings.pt
    del model.binary_decoder, model.dataset_encoder, model.gene_embedding_layer
    model = model.eval().to(dev, DT)
    print(f"[{MODEL}] loaded {len(sd)} tensors strictly; config pad_length={ds['pad_length']} "
          f"{m['nlayers']}x{m['emsize']} heads={m['nhead']}", flush=True)
    return model, cfg, {g: i for i, g in enumerate(pe)}


# ----------------------------------------------------------------------------- tokenisation (official collator)
def sentences(X, genes, g2i, pad_len, cls_idx, seed):
    valid = np.array([g in g2i for g in genes])
    Xv = X[:, valid].tocsr()
    gid = np.array([g2i[g] for g in genes[valid]], dtype=np.int64)
    n = Xv.shape[1]
    assert n >= pad_len - 1                                         # -> official "top pad_len-1 genes" branch
    rng = np.random.default_rng(seed)
    T = np.zeros((Xv.shape[0], pad_len), np.int64)
    C = np.zeros((Xv.shape[0], pad_len), np.float32)
    with np.errstate(over="ignore"):                                # official is_raw_integer_counts heuristic
        n_heur = int(sum((r.data.max(initial=0) <= 35) and (np.expm1(r.data.astype(np.float64)).sum() <= 5e6)
                         for r in (X.getrow(i) for i in range(X.shape[0]))))   # (on the full row, as official)
    n_trunc = 0
    for r in range(Xv.shape[0]):
        a, b = Xv.indptr[r], Xv.indptr[r + 1]
        nz, raw = Xv.indices[a:b], Xv.data[a:b].astype(np.float64)
        lv = np.log1p(raw)
        s = lv.sum()
        w = np.zeros(n)
        if s > 0:
            w[nz] = lv / s
        else:
            w[:] = 1.0 / n                                          # softmax of an all-zero row
        p = rng.permutation(len(nz))                                # random tie-break among expressed genes
        order = nz[p][np.argsort(-lv[p], kind="stable")][:pad_len - 1]
        n_trunc += len(nz) > pad_len - 1
        k = pad_len - 1 - len(order)
        if k > 0:                                                   # random undetected genes, no replacement
            zero = np.setdiff1d(np.arange(n), nz, assume_unique=True)
            order = np.r_[order, rng.choice(zero, k, replace=False)]
        loc = np.r_[cls_idx, order]
        T[r] = gid[loc]
        C[r] = 100.0 * w[loc]
    stats = dict(valid=int(valid.sum()), n_gene=len(genes), heur_log=n_heur, trunc=n_trunc,
                 det_med=float(np.median(np.diff(Xv.indptr))))
    return T, C, stats


def embed(model, T, C, bs=16):
    E = np.zeros((len(T), model.decoder[1].out_features), np.float32)
    i = 0
    with torch.no_grad(), torch.autocast("cuda", dtype=DT):
        while i < len(T):
            try:
                t = torch.from_numpy(T[i:i + bs]).to(dev)
                c = torch.from_numpy(C[i:i + bs]).to(dev, DT)
                E[i:i + bs] = model(t, c).float().cpu().numpy()
                i += bs
            except torch.OutOfMemoryError:
                torch.cuda.empty_cache()
                assert bs > 1, "OOM at batch size 1"
                bs //= 2
                print(f"  OOM -> batch size {bs}", flush=True)
    return E


def sanity(E, name):
    rng = np.random.default_rng(0)
    S = E[rng.choice(len(E), min(200, len(E)), replace=False)]
    S = S / np.linalg.norm(S, axis=1, keepdims=True)
    mc = float((S @ S.T)[np.triu_indices(len(S), 1)].mean())
    msg = f"[{MODEL}] {name}: finite={np.isfinite(E).all()} std={E.std(0).mean():.4f} mean-pair-cos(200)={mc:.4f}"
    if name == "hest":
        org = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")).organ.values
        assert len(org) == len(E)
        Z = torch.from_numpy(E / np.linalg.norm(E, axis=1, keepdims=True)).to(dev)
        nn_idx = []
        for k in range(0, len(Z), 2048):
            s = Z[k:k + 2048] @ Z.T
            s[torch.arange(s.shape[0]), torch.arange(k, k + s.shape[0])] = -2
            nn_idx.append(s.topk(10, dim=1).indices.cpu().numpy())
        nn_idx = np.concatenate(nn_idx)
        knn = float((org[nn_idx] == org[:, None]).mean())
        p = pd.Series(org).value_counts(normalize=True).values
        msg += f"  10-NN same-organ={knn:.3f} (base rate {float((p ** 2).sum()):.3f})"
    print(msg, flush=True)
    return mc


if __name__ == "__main__":
    torch.cuda.set_per_process_memory_fraction(min(1.0, 8 * 2 ** 30 / torch.cuda.get_device_properties(0).total_memory))
    t0 = time.time()
    model, cfg, g2i = build()
    PAD, CLS = cfg["dataset"]["pad_length"], cfg["dataset"]["cls_token_idx"]
    for name in (sys.argv[1:] or SETS):
        t1 = time.time()
        X, genes = load(name)
        T, C, st = sentences(X, genes, g2i, PAD, CLS, seed=0)
        print(f"[{MODEL}] {name}: {X.shape[0]} spots; gene coverage {st['valid']}/{st['n_gene']} "
              f"({st['valid'] / st['n_gene']:.1%}) in protein_embeddings; detected valid genes/spot median "
              f"{st['det_med']:.0f}; spots with >2047 detected (truncated) {st['trunc']}; spots the official "
              f"raw-vs-log1p heuristic would treat as log1p {st['heur_log']}", flush=True)
        E = embed(model, T, C)
        if name == "hest":                                          # tie-breaking stochasticity check
            idx = np.random.default_rng(1).choice(len(T), 300, replace=False)
            T2, C2, _ = sentences(X[idx], genes, g2i, PAD, CLS, seed=1)
            c = (embed(model, T2, C2) * E[idx]).sum(1)
            print(f"  seed-0 vs seed-1 tie-breaking: per-spot cosine median {np.median(c):.4f} min {c.min():.4f}",
                  flush=True)
        sanity(E, name)
        save(MODEL, name, E)
        print(f"  {name}: {time.time() - t1:.0f}s  peak GPU {torch.cuda.max_memory_allocated() / 2 ** 30:.2f} GB",
              flush=True)
    print(f"[{MODEL}] done in {time.time() - t0:.0f}s", flush=True)