#!/usr/bin/env python3
"""Feature 2: protein + RNA → binding probability (RnaRealismClassifierV3)."""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import h5py
import torch
from transformers import AutoTokenizer

from common import PKG_ROOT, device, load_env, resolve
from encode_protein import encode_protein_to_h5

sys.path.insert(0, str(PKG_ROOT / "generator"))
sys.path.insert(0, str(PKG_ROOT / "classifier"))

from labeled_dataset import make_collate  # noqa: E402
from model import RnaRealismClassifierV3  # noqa: E402
from src.model import EsmConfig, EsmForMaskedLM  # noqa: E402
from src.utils import _remap_legacy_checkpoint_state_dict, base_config  # noqa: E402


def _load_mlm(ckpt_path: Path, tokenizer, torch_device: torch.device, model_cfg: dict) -> EsmForMaskedLM:
    mcfg = model_cfg["model"]
    cfg = EsmConfig(**base_config)
    cfg.vocab_size = len(tokenizer)
    cfg.pad_token_id = tokenizer.pad_token_id
    cfg.mask_token_id = tokenizer.mask_token_id
    cfg.use_FiLM = mcfg["use_FiLM"]
    cfg.use_AdaLN = mcfg["use_AdaLN"]
    cfg.use_gated_bias = mcfg["use_gated_bias"]
    cfg.use_protein_conditioning_attention = mcfg["use_protein_conditioning_attention"]
    cfg.protein_dim = mcfg["protein_dim"]
    mlm = EsmForMaskedLM(cfg)
    ckpt = torch.load(ckpt_path, map_location="cpu")
    sd = _remap_legacy_checkpoint_state_dict(ckpt["model_state_dict"])
    mlm.load_state_dict(sd, strict=True)
    mlm.to(torch_device)
    return mlm


@torch.no_grad()
def score_binding(
    protein_seq: str,
    rna_seq: str,
    *,
    p_id: str = "QUERY",
    protein_h5: Path | None = None,
) -> dict:
    load_env()
    torch_device = torch.device(device() if torch.cuda.is_available() else "cpu")

    if protein_h5 is None:
        work = resolve(os.environ.get("WORKSPACE", "workspace"))
        work.mkdir(parents=True, exist_ok=True)
        protein_h5 = encode_protein_to_h5(
            protein_seq, p_id=p_id, output_h5=work / f"{p_id}_vesm3b.h5"
        )

    with h5py.File(protein_h5, "r") as f:
        p_ids = [x.decode() if isinstance(x, bytes) else str(x) for x in f["p_ids"][:]]
        if p_id not in p_ids:
            p_id = p_ids[0]
        idx = p_ids.index(p_id)
        start = int(f["starts"][idx])
        length = int(f["lengths"][idx])
        emb = torch.from_numpy(f["embeddings"][start : start + length]).float()

    rna = rna_seq.upper().replace("T", "U")
    rna = "".join(ch for ch in rna if ch in "ACGU")
    if not rna:
        raise ValueError("RNA sequence is empty or invalid")

    cls_ckpt = resolve(os.environ.get("CLASSIFIER_CKPT", "checkpoints/classifier.pt"))
    backbone = resolve(
        os.environ.get("BACKBONE_CKPT")
        or os.environ.get("D3LM_BACKBONE_CKPT")  # legacy env name
        or "checkpoints/backbone.pt"
    )
    tok_dir = resolve(os.environ.get("TOKENIZER_DIR", "generator/tokenizer"))

    bundle = torch.load(cls_ckpt, map_location="cpu")
    # Checkpoints may still store config under the legacy key "d3lm_config".
    model_cfg = bundle.get("model_config") or bundle["d3lm_config"]
    train_args = bundle.get("train_args") or {}
    data_cfg = model_cfg["data"]

    tokenizer = AutoTokenizer.from_pretrained(str(tok_dir), trust_remote_code=True)
    mlm = _load_mlm(backbone, tokenizer, torch_device, model_cfg)

    model = RnaRealismClassifierV3(
        mlm,
        mlm.config.hidden_size,
        int(model_cfg["model"]["protein_dim"]),
        freeze_mode=train_args.get("freeze_mode", "conditioning_only"),
        head_dropout=float(train_args.get("head_dropout", 0.1)),
        cross_attn_heads=int(train_args.get("cross_attn_heads", 8)),
    ).to(torch_device)

    state_key = "classifier_state_dict" if "classifier_state_dict" in bundle else "model_state_dict"
    model.load_state_dict(bundle[state_key], strict=True)
    model.eval()

    sample = {
        "r_id": "query",
        "rna": rna,
        "p_id": p_id,
        "protein_emb": emb,
        "rna_len": len(rna),
        "protein_len": int(length),
        "label": 0,
    }
    append_eos = bool(data_cfg.get("append_eos_token", False))
    max_rna_nt = data_cfg.get("max_generated_rna_bp")
    max_rna_nt = int(max_rna_nt) if max_rna_nt is not None else None
    collate = make_collate(str(tok_dir), append_eos, max_rna_nt)
    batch = collate([sample])

    use_bf16 = torch_device.type == "cuda"
    with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16, enabled=use_bf16):
        logits = model(
            protein_cond=batch["protein_cond"].to(torch_device),
            protein_attention_mask=batch["protein_attention_mask"].to(torch_device),
            input_ids=batch["rna_input_ids_clean"].to(torch_device),
            attention_mask=batch["rna_attention_mask"].to(torch_device),
        )
    if isinstance(logits, (tuple, list)):
        logits = logits[0]
    logit_t = logits.reshape(-1)[0].float()
    prob_t = torch.sigmoid(logit_t)
    return {
        "p_id": p_id,
        "rna": rna,
        "rna_len": len(rna),
        "logit": float(logit_t.cpu()),
        "binding_prob": float(prob_t.cpu()),
        "protein_h5": str(protein_h5),
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Protein + RNA → binding probability")
    ap.add_argument("--protein", required=True, help="AA sequence, or @path/to.fasta")
    ap.add_argument("--rna", required=True, help="RNA sequence (ACGU/T)")
    ap.add_argument("--p-id", default="QUERY")
    ap.add_argument("--protein-h5", type=Path, default=None, help="Skip encoding if H5 exists")
    args = ap.parse_args()

    load_env()
    protein = args.protein
    if protein.startswith("@"):
        text = Path(protein[1:]).read_text(encoding="utf-8")
        lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
        if lines and lines[0].startswith(">"):
            protein = "".join(lines[1:])
            if args.p_id == "QUERY":
                args.p_id = lines[0][1:].split()[0]
        else:
            protein = "".join(lines)

    out = score_binding(protein, args.rna, p_id=args.p_id, protein_h5=args.protein_h5)
    print(json.dumps(out, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()