"""
Batch integration downstream example:
- Extract embeddings with frozen GeneMamba
- Evaluate simple batch mixing score proxy (silhouette by batch)

Expected h5ad columns:
- obs['batch']
"""

import argparse
import numpy as np
import scanpy as sc
import torch
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import LabelEncoder
from transformers import AutoModel


def build_ranked_input_ids(adata, symbol2id, seq_len=2048, pad_id=1):
    gene_names = np.array(adata.var_names)
    X = adata.X
    out = np.full((adata.n_obs, seq_len), pad_id, dtype=np.int64)

    for i in range(adata.n_obs):
        row = X[i]
        if hasattr(row, "toarray"):
            expr = row.toarray().ravel()
        else:
            expr = np.asarray(row).ravel()

        nz = np.where(expr > 0)[0]
        if len(nz) == 0:
            continue

        genes = gene_names[nz]
        vals = expr[nz]
        order = np.argsort(-vals)
        ranked_genes = genes[order]
        ids = [symbol2id[g] for g in ranked_genes if g in symbol2id][:seq_len]
        out[i, : len(ids)] = ids

    return out


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", required=True)
    parser.add_argument("--h5ad", required=True)
    parser.add_argument("--symbol2id_npy", required=True)
    parser.add_argument("--seq_len", type=int, default=2048)
    parser.add_argument("--batch_size", type=int, default=64)
    args = parser.parse_args()

    adata = sc.read_h5ad(args.h5ad)
    assert "batch" in adata.obs, "h5ad must include obs['batch']"

    symbol2id = np.load(args.symbol2id_npy, allow_pickle=True).item()
    input_ids = build_ranked_input_ids(adata, symbol2id, seq_len=args.seq_len)

    model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True)
    model.eval().cuda()

    embeds = []
    with torch.no_grad():
        for s in range(0, input_ids.shape[0], args.batch_size):
            batch = torch.tensor(input_ids[s : s + args.batch_size], dtype=torch.long, device="cuda")
            out = model(batch)
            embeds.append(out.pooled_embedding.detach().cpu().numpy())
    embeds = np.concatenate(embeds, axis=0)

    batch_labels = LabelEncoder().fit_transform(adata.obs["batch"].values)
    score = silhouette_score(embeds, batch_labels, metric="euclidean")

    print("silhouette_by_batch:", score)
    print("(Closer to 0 typically indicates better batch mixing than very high positive values.)")


if __name__ == "__main__":
    main()