from transformers import PreTrainedModel, PretrainedConfig
import torch
import torch.nn as nn
from tape import ProteinBertForSequenceClassification, TAPETokenizer


class KinaseSubstrateConfig(PretrainedConfig):
    model_type = "kinase_substrate_bert"

    def __init__(
        self,
        tape_model_name="bert-base",
        num_labels=2,
        threshold=0.5,
        with_sep=True,
        max_len=1024,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.tape_model_name = tape_model_name
        self.num_labels = num_labels
        self.threshold = threshold
        self.with_sep = with_sep
        self.max_len = max_len


class KinaseSubstrateModel(PreTrainedModel):
    config_class = KinaseSubstrateConfig

    def __init__(self, config: KinaseSubstrateConfig):
        super().__init__(config)
        self.backbone = ProteinBertForSequenceClassification.from_pretrained(
            config.tape_model_name, num_labels=config.num_labels
        )

    def forward(self, input_ids, input_mask=None, targets=None):
        return self.backbone(
            input_ids=input_ids,
            input_mask=input_mask,
            targets=targets,
        )

    def predict_proba(self, input_ids, input_mask):
        self.eval()
        with torch.no_grad():
            (_, _), logits = self.forward(input_ids=input_ids, input_mask=input_mask)
            return torch.softmax(logits, dim=-1)[:, 1]