from transformers import pipeline
import gradio as gr
from flair.data import Sentence
from flair.models import SequenceTagger
import torch



#constants

GR_TXT1 = """
    Recent studies have identified several key mutations associated with familial breast cancer.
    The BRCA1 c.5266dupC mutation (p.Gln1756Profs74) was detected in 12% of patients, while the BRCA2 p.Ser1982fs variant showed reduced penetrance.
    Analysis revealed that the TP53 R175H substitution co-occurred with EGFR L858R in 8 cases.
    Additionally, the HER2 amplification correlated with response to trastuzumab therapy.
    Novel variants including PIK3CA E545K and AKT1 E17K were identified in endocrine-resistant tumors.
    The KRAS G12D mutation remained a negative predictive marker for anti-EGFR therapy.
"""

GR_TXT2 = """
    During the 12-week treatment period, 23 patients experienced adverse events.
    The most common was grade 2 nausea occurring in 15 participants, followed by fatigue in 11 cases.
    Three patients developed severe neutropenia requiring dose reduction, and two experienced thrombocytopenia.
    One case of drug-induced hepatotoxicity was documented with elevated ALT levels.
    Mild headache and dizziness were reported by 8 patients during the first week.
    Serious adverse events included one instance of anaphylaxis and two cases of Stevens-Johnson syndrome, leading to treatment discontinuation.
    No cardiovascular events or nephrotoxicity were observed.
  """


GR_MARKDOWN_TXT = """ 
# 🧬 OTAR3088 Work-in-progress NER-models demo space 
This space allows for the **visualisation of outputs / review of our WIP NER-models.** Presently these models include... 👀:  
- **Adverse Events**  
- **Variants**   

👉 Try it out:  
1. Select a model from the dropdown menu.  
2. Enter/paste your text into the input box, or use our provided example biomedical paragraphs.  
3. View the extracted entities highlighted directly in **"Tagged Entities"** box.  
**Note📢:** Models in this demo are continuously updated and improved as part of our ongoing research.  
"""


GR_THEME = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="rose",
    neutral_hue="gray"
)




MODEL_REGISTRY = {
    "Variants-V1": "OTAR3088/Variants-V1",
    "Variants-Reinit-LLRD": "OTAR3088/Variant_reinit-llrd_PuBMedBert_V1",
    "PDBe-Variants-V2.1": "PDBEurope/BiomedNLP-PubMedBERT-ProteinStructure-NER-v2.1",
    "AdverseEvent-PHEE-V1": "OTAR3088/hf-phee-V1"
   }


hf_pipes = {}
flair_pipes = {}



def load_model(model_name):
  if model_name.lower().startswith("flair"):
    if not model_name in flair_pipes:
      # Force default tensor type
      torch.set_default_dtype(torch.float32)

      flair_pipes[model_name] = SequenceTagger.load(MODEL_REGISTRY[model_name])
      flair_pipes[model_name].to(torch.device("cpu"))
      flair_pipes[model_name].float()
    return flair_pipes[model_name], "flair"
  else:
    if not model_name in hf_pipes:
      hf_pipes[model_name] = pipeline("ner", model=MODEL_REGISTRY[model_name], aggregation_strategy='simple')
    return hf_pipes[model_name], "hf"



def tagger(text, model_name):
  model, model_type = load_model(model_name)
  if model_type == "flair":
    sentence = Sentence(text)
    model.predict(sentence)
    entities = [{"start": ent.start_position, 
                 "end": ent.end_position, 
                 "score": ent.score, 
                 "entity": ent.tag} for ent in sentence.get_spans('ner')]

  elif model_type == "hf":
    entities = model(text)

  return {"text": text, "entities": entities}


def gradio_ui():
  with gr.Blocks(theme=GR_THEME) as demo:
    with gr.Row():
        with gr.Column(scale=7):
            gr.Markdown(GR_MARKDOWN_TXT,elem_classes="full-width")
            input_text = gr.Textbox(label="Enter your text here", type='text', placeholder="Biomedical Input text", lines=8)
        

            gr.Examples(
                    examples=[GR_TXT1, GR_TXT2],
                inputs=[input_text],
                label= "Example Biomedical texts to try (fabricated texts, not from literature)"
            )
            run_btn = gr.Button("Submit Text", variant="primary")

        with gr.Column(scale=7):
            model_choice = gr.Dropdown(choices=list(MODEL_REGISTRY.keys()), label="Select a model for Inference")
            output_highlight = gr.HighlightedText(label="Tagged Entities")
            

    # with gr.Row():
    #   with gr.Column(scale=7):
    #     input_text = gr.Textbox(label="Enter your text here", type='text', placeholder="Biomedical Input text", lines=8)
        

    #     gr.Examples(
    #             examples=[GR_TXT1, GR_TXT2],
    #         inputs=[input_text],
    #         label= "Example Biomedical texts to try"
    #     )
    #     run_btn = gr.Button("Submit Text", variant="primary")


    run_btn.click(
          fn=tagger,
          inputs=[input_text, model_choice],
          outputs=[output_highlight]
      )

        

    return demo




if __name__ == "__main__":
  app = gradio_ui()
  app.launch()