import gradio as gr

import torch,re
from transformers import EsmTokenizer,EsmForMaskedLM

model_name = "facebook/esm2_t6_8M_UR50D"
tokenizer = EsmTokenizer.from_pretrained(model_name)

mask_token = tokenizer.mask_token

device = torch.device("cpu")

def predict_mt(sequence, mutation):
  protein_sequence = sequence
  mutatation = mutation
  variant_position=match = re.search(r'\d+', mutation)
  variant_position_0_based =  int(variant_position.group()) -1# Index for R
  wild_type_aa = mutation[0]
  mutant_aa = mutation[-1]
# 2. Prepare the masked sequence (replace WT AA with '<mask>')
  masked_sequence_list = list(protein_sequence)
  masked_sequence_list[variant_position_0_based] = mask_token # Replace with the model's mask token
  masked_sequence = "".join(masked_sequence_list)

# 3. Encode the masked sequence using the tokenizer
# The tokenizer automatically adds CLS and EOS tokens
  encoded_inputs = tokenizer(masked_sequence, return_tensors="pt", add_special_tokens=True)
  batch_tokens = encoded_inputs['input_ids'].to(device)
  model_mlm = EsmForMaskedLM.from_pretrained(model_name).to(device)
  results = model_mlm(batch_tokens)
  logits = results.logits
  masked_token_index_in_batch = variant_position_0_based + 1
  logits_at_masked_position = logits[0, masked_token_index_in_batch]
  wt_token_id = tokenizer.encode(wild_type_aa, add_special_tokens=False)[0]
  mut_token_id = tokenizer.encode(mutant_aa, add_special_tokens=False)[0]
  log_prob_mutant = logits_at_masked_position[mut_token_id]
  log_prob_wild_type = logits_at_masked_position[wt_token_id]
  llr_score = log_prob_mutant - log_prob_wild_type
  return f"LLR SCore: {llr_score:.2f}"

demo = gr.Interface(
    fn=predict_mt,
    inputs=[
        gr.Textbox(label="Enter Protein Amino Acid Sequence (1-letter code)",
        placeholder="ACDEFGHIKLMNPQRSTVWY"),
        gr.Textbox(label="Enter Missense Mutation",
        placeholder="R5G")
    ],
    outputs="text",
    title="Nano Protein Language Model for Missense Mutation Prediction",
    description="Enter an amino acid sequence (using the 1-letter code) and Missense Mutation (Eg. R5G) to predict its effect.",
    examples=[
        ["MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG","R5G"], # Example sequence
        #["MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKA"], # Example sequence 2
    ]
)
demo.launch()