import torch
import torch.nn as nn
import torch.nn.functional as F
import esm
import joblib  # ✅ added

class BiTransformerAttentionModel(nn.Module):
    def __init__(self, input_dim, seq_len=8, embed_dim=512, n_heads=8, ff_hidden=256, dropout=0.2):
        super().__init__()
        assert input_dim % seq_len == 0
        self.seq_len = seq_len
        self.token_dim = input_dim // seq_len
        self.fc_in = nn.Linear(self.token_dim, embed_dim)
        enc = lambda: nn.TransformerEncoderLayer(
            d_model=embed_dim, nhead=n_heads,
            dim_feedforward=ff_hidden, dropout=dropout,
            batch_first=True
        )
        self.fwd = nn.TransformerEncoder(enc(), num_layers=2)
        self.bwd = nn.TransformerEncoder(enc(), num_layers=2)
        self.attn = nn.Linear(embed_dim*2,1)
        self.ln = nn.LayerNorm(embed_dim*2)
        self.fc = nn.Sequential(
            nn.Linear(embed_dim*2, ff_hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ff_hidden, ff_hidden//2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ff_hidden//2,1)
        )
        self.pos = nn.Parameter(torch.randn(1, seq_len, embed_dim)*0.01)

    def forward(self,x):
        x = x.view(x.size(0), self.seq_len, self.token_dim)
        x = self.fc_in(x)
        f = self.fwd(x+self.pos)
        b = self.bwd(torch.flip(x+self.pos,[1]))
        x = torch.cat([f,b],dim=-1)
        x = self.ln(x)
        w = torch.softmax(self.attn(x),dim=1)
        x = (w*x).sum(dim=1)
        return self.fc(x).squeeze()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ✅ Load checkpoint
checkpoint = torch.load("BiTransformer_BEST.pt", map_location=device)

model = BiTransformerAttentionModel(input_dim=checkpoint["input_dim"])
model.load_state_dict(checkpoint["model_state"])
model.to(device)
model.eval()

# ✅ Load scaler
scaler = joblib.load("scaler.pkl")

# -----------------------------
# ESM embedding
# -----------------------------
esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
esm_model = esm_model.eval().to(device)
batch_converter = alphabet.get_batch_converter()

def sequence_to_embedding(sequences):
    batch_labels, batch_strs, batch_tokens = batch_converter(sequences)
    batch_tokens = batch_tokens.to(device)
    with torch.no_grad():
        results = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)
    token_reps = results["representations"][33]
    mask = batch_tokens != alphabet.padding_idx
    embeddings = (token_reps * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True)
    return embeddings

# -----------------------------
# Prediction
# -----------------------------
def predict(sequence: str):
    seq_tuple = [("seq", sequence)]
    X_embed = sequence_to_embedding(seq_tuple)

    # match input size
    X_embed = X_embed[:, :checkpoint["input_dim"]]

    # ✅ APPLY SCALER (CRITICAL FIX)
    X_embed = scaler.transform(X_embed.cpu().numpy())
    X_embed = torch.tensor(X_embed, dtype=torch.float32).to(device)

    model.eval()
    with torch.no_grad():
        logits = model(X_embed)
        prob = torch.sigmoid(logits)

    return float(prob.cpu().item())