import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
import pandas as pd
import util
import folium

# Model names (replace with your actual Hugging Face repo names)
MODEL_NAMES = {
    "BioBERT": "nattkorat/biobert-base-uncased-ner",
    "SciBERT": "nattkorat/scibert-base-uncased-ner",
    "BERT": "nattkorat/bert-base-uncased-ner"
}

# Cache loaded models to avoid reloading every time
loaded_models = {}

def load_model(model_key):
    model_name = MODEL_NAMES[model_key]
    if model_name not in loaded_models:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForTokenClassification.from_pretrained(model_name)
        ner = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="average")
        loaded_models[model_name] = ner
    return loaded_models[model_name]

def extract_entities(model_choice, text):
    ner = load_model(model_choice)
    entities = ner(text)

    if not entities:
        return pd.DataFrame(columns=["outbreak", "cases", "deaths", "date", "location", "latitude", "longtitude"])
    
    data = {}
    for ent in entities:
        if ent['entity_group'] == 'VIRUS':
            data['outbreak'] = ent['word']
        elif ent['entity_group'] == 'CASES':
            data['cases'] = ent['word']
        elif ent['entity_group'] == 'DEATHS':
            data['deaths'] = ent['word']
        elif ent['entity_group'] == 'DATE':
            data['date'] = util.parse_date(ent['word'])
        elif ent['entity_group'] == 'COUNTRY':
            location_info = util.get_location(ent['word'])
            data['location'] = location_info['name']
            data['latitude'] = location_info['latitude']
            data['longtitude'] = location_info['longitude']
    
    m = None
    if 'latitude' in data:
        m = folium.Map(location=[data['latitude'], data['longtitude']])
        info = f"""
        <h3>Outbreak Information</h3>
        <strong>Location:</strong> {data.get('location', 'N/A')}<br>
        <strong>Outbreak:</strong> {data.get('outbreak', 'N/A')}<br>
        <strong>Cases:</strong> {data.get('cases', 'N/A')}<br>
        <strong>Deaths:</strong> {data.get('deaths', 'N/A')}<br>
        <strong>Date:</strong> {data.get('date', 'N/A')}<br>
        <strong>Location:</strong> {data.get('location', 'N/A')}
        """
        folium.Marker(
            location=[data['latitude'], data['longtitude']],
            popup=info,
            icon=folium.Icon(icon="warning", color="red"),
        ).add_to(m)

    return pd.DataFrame([data]), m._repr_html_()

# Sample example texts for testing
examples = [
    ["Outbreak of Ebola occurred in Congo on 2007-09-11, with confirmed 372 cases and 166 deaths."],
    ["Outbreak of Avian influenza occurred in The Netherlands on 2010-12-09, with confirmed 1 cases and 141 deaths."],
    ["Cholera occurred in Mozambique on 1998-08-14, with confirmed 26783 cases and 619 deaths."]
]

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Outbreak Extraction DEMO")
    gr.Markdown("Select a model, enter text, and see extracted entities below:")

    with gr.Row():
        model_dropdown = gr.Dropdown(choices=list(MODEL_NAMES.keys()), label="Select Model", value="BioBERT")
        text_input = gr.Textbox(label="Input Text", placeholder="Enter your sentence here...", lines=4)
        
    gr.Examples(
        examples=examples,
        inputs=[text_input],
        label="Try Examples",
    )
    
    run_button = gr.Button("Extract Entities")
    output_table = gr.Dataframe(headers=["outbreak", "cases", "deaths", "date", "location", "latitude", "longtitude"], interactive=False)
    output_map = gr.HTML(label="Map")

    run_button.click(fn=extract_entities, inputs=[model_dropdown, text_input], outputs=[output_table, output_map])

demo.launch()