import gradio as gr
import torch
import json
import re
from huggingface_hub import hf_hub_download
from tape import ProteinBertForSequenceClassification, TAPETokenizer

WHITESPACE_RE = re.compile(r"\s+")

def sanitize(s):
    return WHITESPACE_RE.sub("", str(s).upper())

# Load config
config_path = hf_hub_download("steveyu323/kinbert_v2_long", "config.json")
with open(config_path) as f:
    config_dict = json.load(f)

THRESHOLD = float(config_dict.get("threshold", 0.5))
MAX_LEN = int(config_dict.get("max_len", 1024))

# Load model
model_path = hf_hub_download("steveyu323/kinbert_v2_long", "pytorch_model.bin")
model = ProteinBertForSequenceClassification.from_pretrained("bert-base", num_labels=2)
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

tokenizer = TAPETokenizer(vocab="iupac")

def predict(kinase_seq, substrate_seq):
    kinase_seq = sanitize(kinase_seq)
    substrate_seq = sanitize(substrate_seq)

    if not kinase_seq or not substrate_seq:
        return "Please enter both sequences.", None

    kin_toks = tokenizer.tokenize(kinase_seq)
    sub_toks = tokenizer.tokenize(substrate_seq)
    toks = kin_toks + ["<sep>"] + sub_toks
    toks = tokenizer.add_special_tokens(toks)
    ids = tokenizer.convert_tokens_to_ids(toks)[:MAX_LEN]

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

    with torch.no_grad():
        outputs = model(input_ids=input_ids, input_mask=input_mask)
        # TAPE returns (loss_tuple, logits) only when targets provided
        # without targets it returns just logits
        logits = outputs[0] if isinstance(outputs, tuple) else outputs
        prob = float(torch.softmax(logits, dim=-1)[0, 1])

    label = "✅ Interaction" if prob >= THRESHOLD else "❌ No Interaction"
    return label, round(prob, 4)

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(
            lines=3,
            label="Kinase Domain Sequence",
            value="LVLGKTLGEGEFGKVVKATAFHLKGRAGYTTVAVKMLKENASPSELRDLLSEFNVLKQVNHPHVIKLYGACSQDGPLLLIVEYAKYGSLRGFLRESRKVGPGYLGSGGSRNSSSLDHPDERALTMGDLISFAWQISQGMQYLAEMKLVHRDLAARNILVAEGRKMKISDFGLSRDVYEEDSYVKRSQGRIPVKWMAIESLFDHIYTTQSDVWSFGVLLWEIVTLGGNPYPGIPPERLFNLLKTGHRMERPDNCSEEMYRLMLQCWKQEPDKRPVFADISKDLEKMMVKRRDYL",  # default kinase
        ),
        gr.Textbox(
            lines=3,
            label="Substrate Sequence",
            value="TWIENKLYGMSDPNW",  # default substrate
        ),
    ],
    outputs=[
        gr.Text(label="Prediction"),
        gr.Number(label="Interaction Probability"),
    ],
    title="KinBERT — Kinase–Substrate Interaction Classifier",
    description="Predicts whether a kinase will phosphorylate a given substrate sequence.",
)

demo.launch(ssr_mode=False)