"""
SERAPH model architecture definition and weight-loading helper.
"""
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import EsmModel, EsmTokenizer

from config import HF_REPO_ID, WEIGHTS_FILE, ESM_MODEL_ID


class SERAPH(nn.Module):
    """ESM-2 backbone + Conv1D + BiLSTM head for 3-state secondary
    structure prediction (Helix / Sheet / Coil)."""

    def __init__(self, esm_model, conv_channels=256, kernel_size=7,
                 lstm_hidden=256, num_classes=3, dropout=0.3, freeze_esm=True):
        super().__init__()
        self.esm = esm_model
        if freeze_esm:
            for param in self.esm.encoder.layer[:-2].parameters():
                param.requires_grad = False
        esm_embed_dim = self.esm.config.hidden_size
        self.conv = nn.Conv1d(esm_embed_dim, conv_channels, kernel_size=kernel_size, padding=kernel_size // 2)
        self.bn = nn.BatchNorm1d(conv_channels)
        self.dropout = nn.Dropout(dropout)
        self.bilstm = nn.LSTM(conv_channels, lstm_hidden, num_layers=2, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(lstm_hidden * 2, num_classes)

    def forward(self, input_ids, attention_mask=None):
        x = self.esm(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        x = x.transpose(1, 2)
        x = torch.relu(self.bn(self.conv(x)))
        x = self.dropout(x)
        x = x.transpose(1, 2)
        x, _ = self.bilstm(x)
        x = self.dropout(x)
        return self.fc(x)


def load_model():
    """Downloads the ESM-2 backbone + tokenizer and the fine-tuned SERAPH
    checkpoint, and returns a ready-to-use (model, tokenizer) pair."""
    print("Loading ESM2 backbone...")
    esm = EsmModel.from_pretrained(ESM_MODEL_ID)
    tokenizer = EsmTokenizer.from_pretrained(ESM_MODEL_ID)

    print("Downloading SERAPH weights...")
    weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=WEIGHTS_FILE)
    checkpoint = torch.load(weights_path, map_location="cpu")

    model = SERAPH(esm_model=esm)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    print("SERAPH model ready.")

    return model, tokenizer