"""Deprecated MitoInteract v1 architecture retained for artifact forensics.

The v1 checkpoint failed its evaluation gates and uses unsafe pickle
serialization. New work lives under ``recovery/``. Loading and inference are
disabled by default so this module cannot silently present an undocumented
legacy score as pKd.
"""

import torch
import torch.nn as nn
from transformers import EsmModel, AutoModel, AutoTokenizer


class MitoInteract(nn.Module):
    def __init__(
        self,
        esm_model_name="facebook/esm2_t33_650M_UR50D",
        mol_model_name="seyonec/ChemBERTa-zinc-base-v1",
        protein_dim=1280,
        mol_dim=768,
        proj_dim=256,
        n_heads=8,
        dropout=0.1,
        freeze_encoders=True,
    ):
        super().__init__()
        self.freeze_encoders = freeze_encoders
        self.esm = EsmModel.from_pretrained(esm_model_name)
        self.protein_dim = protein_dim
        self.mol_encoder = AutoModel.from_pretrained(mol_model_name)
        self.mol_dim = mol_dim
        if freeze_encoders:
            for p in self.esm.parameters():
                p.requires_grad = False
            for p in self.mol_encoder.parameters():
                p.requires_grad = False
        self.prot_proj = nn.Sequential(
            nn.Linear(protein_dim, proj_dim),
            nn.LayerNorm(proj_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.mol_proj = nn.Sequential(
            nn.Linear(mol_dim, proj_dim),
            nn.LayerNorm(proj_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.cross_attn_mol2prot = nn.MultiheadAttention(
            proj_dim, n_heads, dropout=dropout, batch_first=True
        )
        self.cross_attn_prot2mol = nn.MultiheadAttention(
            proj_dim, n_heads, dropout=dropout, batch_first=True
        )
        self.ln_mol2prot = nn.LayerNorm(proj_dim)
        self.ln_prot2mol = nn.LayerNorm(proj_dim)
        fused_dim = proj_dim * 2
        self.mlp = nn.Sequential(
            nn.Linear(fused_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 1),
        )

    def encode_protein(self, input_ids, attention_mask):
        ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
        with ctx:
            out = self.esm(input_ids=input_ids, attention_mask=attention_mask)
        mask = attention_mask.unsqueeze(-1).float()
        pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
        return pooled, out.last_hidden_state

    def encode_molecule(self, input_ids, attention_mask):
        ctx = torch.no_grad() if self.freeze_encoders else torch.enable_grad()
        with ctx:
            out = self.mol_encoder(input_ids=input_ids, attention_mask=attention_mask)
        return out.pooler_output, out.last_hidden_state

    def forward(
        self, prot_input_ids, prot_attention_mask, mol_input_ids, mol_attention_mask
    ):
        prot_pooled, prot_seq = self.encode_protein(prot_input_ids, prot_attention_mask)
        mol_pooled, mol_seq = self.encode_molecule(mol_input_ids, mol_attention_mask)
        prot_seq_proj = self.prot_proj(prot_seq)
        mol_seq_proj = self.mol_proj(mol_seq)
        prot_q = self.prot_proj(prot_pooled).unsqueeze(1)
        mol_q = self.mol_proj(mol_pooled).unsqueeze(1)
        prot_pad_mask = prot_attention_mask == 0
        mol_pad_mask = mol_attention_mask == 0
        h_prot2mol, _ = self.cross_attn_prot2mol(
            prot_q, mol_seq_proj, mol_seq_proj, key_padding_mask=mol_pad_mask
        )
        h_mol2prot, _ = self.cross_attn_mol2prot(
            mol_q, prot_seq_proj, prot_seq_proj, key_padding_mask=prot_pad_mask
        )
        h_prot2mol = self.ln_prot2mol(h_prot2mol.squeeze(1))
        h_mol2prot = self.ln_mol2prot(h_mol2prot.squeeze(1))
        fused = torch.cat([h_prot2mol, h_mol2prot], dim=-1)
        return self.mlp(fused).squeeze(-1)


def load_model(checkpoint_path, device="cpu", *, allow_unsafe_legacy=False):
    """Load the failed v1 pickle only after explicit risk acknowledgement."""
    if not allow_unsafe_legacy:
        raise RuntimeError(
            "MitoInteract v1 is deprecated and full_model.pt requires unsafe pickle "
            "deserialization. See AUDIT.md. Pass allow_unsafe_legacy=True only for "
            "controlled artifact forensics, never for untrusted files."
        )
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    config = checkpoint["config"]
    model = MitoInteract(
        esm_model_name=config["esm_model"],
        mol_model_name=config["mol_model"],
        protein_dim=config["protein_dim"],
        mol_dim=config["mol_dim"],
        proj_dim=config["proj_dim"],
        n_heads=config["n_heads"],
        dropout=config["dropout"],
        freeze_encoders=True,
    )
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model, config


def predict_binding(
    model, protein_seq, smiles, device="cpu", *, allow_invalid_legacy_output=False
):
    """Return the undocumented v1 score only after explicit acknowledgement."""
    if not allow_invalid_legacy_output:
        raise RuntimeError(
            "v1 output has undocumented target semantics and failed validation. "
            "It cannot be interpreted as pKd or converted to Kd. See AUDIT.md."
        )
    prot_tokenizer = AutoTokenizer.from_pretrained(model.esm.name_or_path)
    mol_tokenizer = AutoTokenizer.from_pretrained(model.mol_encoder.name_or_path)

    prot_enc = prot_tokenizer(
        protein_seq, return_tensors="pt", padding=True, truncation=True, max_length=512
    )
    mol_enc = mol_tokenizer(
        smiles, return_tensors="pt", padding=True, truncation=True, max_length=200
    )

    model = model.to(device)
    with torch.no_grad():
        legacy_score = model(
            prot_enc["input_ids"].to(device),
            prot_enc["attention_mask"].to(device),
            mol_enc["input_ids"].to(device),
            mol_enc["attention_mask"].to(device),
        )

    return {
        "legacy_score": legacy_score.item(),
        "warning": "Do not interpret this value as pKd, Kd, or validated affinity.",
    }