from transformers import Pipeline
from tape import TAPETokenizer
import torch
import numpy as np


class KinaseSubstratePipeline(Pipeline):
    """
    Usage:
        pipe = pipeline(
            "kinase-substrate",
            model="your-username/kinase-substrate-classifier",
            trust_remote_code=True,
        )
        result = pipe({"kinase_seq": "MGSSHHH...", "substrate_seq": "ARTKQTAR..."})
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tape_tokenizer = TAPETokenizer(vocab="iupac")

    def _sanitize_parameters(self, **kwargs):
        return {}, {}, {}

    def preprocess(self, inputs):
        kinase_seq = inputs["kinase_seq"].upper().replace(" ", "")
        substrate_seq = inputs["substrate_seq"].upper().replace(" ", "")

        cfg = self.model.config
        kin_toks = self.tape_tokenizer.tokenize(kinase_seq)
        sub_toks = self.tape_tokenizer.tokenize(substrate_seq)
        toks = kin_toks + (["<sep>"] if cfg.with_sep else []) + sub_toks
        toks = self.tape_tokenizer.add_special_tokens(toks)
        ids = self.tape_tokenizer.convert_tokens_to_ids(toks)[: cfg.max_len]

        input_ids = torch.tensor([ids], dtype=torch.long)
        input_mask = torch.ones_like(input_ids)

        return {"input_ids": input_ids, "input_mask": input_mask}

    def _forward(self, model_inputs):
        prob = self.model.predict_proba(
            input_ids=model_inputs["input_ids"].to(self.device),
            input_mask=model_inputs["input_mask"].to(self.device),
        )
        return {"prob": prob.cpu().numpy()}

    def postprocess(self, model_outputs):
        prob = float(model_outputs["prob"][0])
        threshold = self.model.config.threshold
        return {
            "probability": round(prob, 4),
            "label": "interaction" if prob >= threshold else "no_interaction",
            "threshold": threshold,
        }