import gradio as gr
from transformers import pipeline

# --------------------------------------------------------------------------
#  1. Load your NER model from the Hub using a pipeline
# --------------------------------------------------------------------------
# Replace this with your actual model name on the Hugging Face Hub
model_name = "swardiantara/ADFLER-xlnet-base-cased" 
ner_pipeline = pipeline("ner", model=model_name, aggregation_strategy="simple")

# --------------------------------------------------------------------------
#  2. Define the prediction function
# --------------------------------------------------------------------------
# This function takes raw text and returns a format that Gradio's HighlightedText component understands.
def recognize_log_events(text):
    """
    Performs NER and robustly formats the output for Gradio's HighlightedText.
    """
    if not text:
        return []

    ner_results = ner_pipeline(text)
    
    # Sort entities by their start index to process them in order
    ner_results.sort(key=lambda x: x['start'])
    
    highlighted_output = []
    last_end = 0

    for entity in ner_results:
        # Add the text between the last entity and this one (un-highlighted)
        if entity['start'] > last_end:
            highlighted_output.append((text[last_end:entity['start']], None))
        
        # Add the highlighted entity text
        highlighted_output.append((entity['word'], entity['entity_group']))
        
        last_end = entity['end']
        
    # Add any remaining text after the last entity (un-highlighted)
    if last_end < len(text):
        highlighted_output.append((text[last_end:], None))

    return highlighted_output

# --------------------------------------------------------------------------
#  3. Create the Gradio Interface
# --------------------------------------------------------------------------
# A brief description of your project to display in the app
description = """
This demo showcases an NER model for recognizing key events in drone flight logs, 
a part of my PhD research in digital forensics at ITS. 
Enter a line from a drone log to see the model identify events like 'Takeoff', 'Landing', 'GPS Lock', etc.
"""

# An article providing more context and links
article = """
<div style='text-align: center;'>
    <p>For more details, check out the project on GitHub:</p>
    <a href='https://dronenlp.github.io/documentation' target='_blank'>DroneNLP Project</a> | 
    <a href='https://huggingface.co/swardiantara/ADFLER-xlnet-base-cased' target='_blank'>Model Card</a> |
</div>
"""
# <a href='https' target='_blank'>Research Paper (if applicable)</a>

# Example drone log entries for users to try
examples = [
    ["No image transmission. RTH.; Press Brake button to cancel RTH.; No image transmission. Aircraft returning to home.; Image transmission signal weak. Adjust antennas and make sure they are perpendicular to flight direction of aircraft.; Flight mode changed to Go Home."],
    ["Battery temperature is below 15 degrees Celsius. Warm up the battery temperature to above 25 degree Celsius to ensure a safe flight."],
    ["Strong wireless interference. Please fly with caution. Obstacle Avoidance Disabled. Landing gear lowered. Obstacle Avoidance Disabled."],
    ["Cannot switch flight mode. Turn on 'Multiple Flight Modes' to enable Atti and Sport Modes."]
]

# The main interface
iface = gr.Interface(
    fn=recognize_log_events,
    inputs=gr.Textbox(
        lines=5, 
        label="Drone Flight Log Entry", 
        placeholder="Paste a log entry here..."
    ),
    outputs=gr.HighlightedText(
        label="Recognized Events",
        color_map={"Event": "green"} # Customize these labels and colors!
    ),
    title="🚁 Drone Flight Log Event Recognizer",
    description=description,
    article=article,
    examples=examples
)

# Launch the app!
iface.launch()