import gradio as gr
import json
from clinical_ner import ClinicalNERProcessor

# Initialize the NER processor
ner_processor = ClinicalNERProcessor(use_pos=True, use_anatomy=True)

# Example text
EXAMPLE_TEXT = "Patient presents with pain in the left ventricle and elevated cardiac enzymes. The heart shows signs of inflammation."

def format_entities(entities):
    """Format entities for display"""
    if not entities:
        return "No entities found."
    
    result = []
    for i, entity in enumerate(entities, 1):
        result.append(f"{i}. **{entity['word']}** - Type: {entity['entity_group']} (Score: {entity['score']:.4f})")
    return "\n".join(result)

def format_pos_tags(pos_tags):
    """Format POS tags for display"""
    if not pos_tags:
        return "No POS tags found."
    
    result = []
    for i, tag in enumerate(pos_tags, 1):
        result.append(f"{i}. **{tag['token']}** - POS: {tag['pos']}, Tag: {tag['tag']}, Lemma: {tag['lemma']}")
    return "\n".join(result)

def clinical_ner_basic(text):
    """Clinical NER only"""
    if not text.strip():
        return "Please enter some text."
    try:
        entities = ner_processor.basic_ner(text)
        return format_entities(entities)
    except Exception as e:
        return f"Error: {str(e)}"

def clinical_ner_prolog(text):
    """Clinical NER as Prolog facts"""
    if not text.strip():
        return "Please enter some text."
    try:
        prolog_facts = ner_processor.prolog_ner(text)
        return prolog_facts if prolog_facts else "No entities found."
    except Exception as e:
        return f"Error: {str(e)}"

def anatomy_ner_basic(text):
    """Anatomy NER only"""
    if not text.strip():
        return "Please enter some text."
    try:
        entities = ner_processor.anatomy_ner(text)
        return format_entities(entities)
    except Exception as e:
        return f"Error: {str(e)}"

def anatomy_ner_prolog(text):
    """Anatomy NER as Prolog facts"""
    if not text.strip():
        return "Please enter some text."
    try:
        prolog_facts = ner_processor.prolog_anatomy(text)
        return prolog_facts if prolog_facts else "No entities found."
    except Exception as e:
        return f"Error: {str(e)}"

def pos_tagging_basic(text):
    """POS tagging only"""
    if not text.strip():
        return "Please enter some text."
    try:
        pos_tags = ner_processor.pos_tagging(text)
        return format_pos_tags(pos_tags)
    except Exception as e:
        return f"Error: {str(e)}"

def pos_tagging_prolog(text):
    """POS tagging as Prolog facts"""
    if not text.strip():
        return "Please enter some text."
    try:
        prolog_facts = ner_processor.prolog_pos(text)
        return prolog_facts if prolog_facts else "No POS tags found."
    except Exception as e:
        return f"Error: {str(e)}"

def combined_analysis(text):
    """Combined analysis"""
    if not text.strip():
        return "Please enter some text.", "Please enter some text.", "Please enter some text."
    try:
        result = ner_processor.combined_analysis(text)
        clinical = format_entities(result['clinical_entities'])
        anatomy = format_entities(result['anatomy_entities'])
        pos = format_pos_tags(result['pos_tags'])
        return clinical, anatomy, pos
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return error_msg, error_msg, error_msg

def combined_prolog(text):
    """Combined analysis as Prolog facts"""
    if not text.strip():
        return "Please enter some text."
    try:
        prolog_facts = ner_processor.prolog_combined(text)
        return prolog_facts if prolog_facts else "No results found."
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface with tabs
with gr.Blocks(title="Clinical NER & Anatomy Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Clinical NER, Anatomy Detection, and POS Tagging
        
        This application provides Named Entity Recognition (NER) for clinical text, 
        anatomy detection, and Part-of-Speech (POS) tagging using state-of-the-art models:
        - **Clinical NER**: Bio_ClinicalBERT
        - **Anatomy NER**: OpenMed AnatomyDetect
        - **POS Tagging**: spaCy en_core_web_sm
        """
    )
    
    with gr.Tabs():
        # Tab 1: Clinical NER
        with gr.Tab("Clinical NER"):
            with gr.Row():
                with gr.Column():
                    clinical_input = gr.Textbox(
                        label="Enter Clinical Text",
                        placeholder="Enter medical text here...",
                        lines=5,
                        value=EXAMPLE_TEXT
                    )
                    clinical_format = gr.Radio(
                        choices=["Basic", "Prolog"],
                        value="Basic",
                        label="Output Format"
                    )
                    clinical_btn = gr.Button("Extract Clinical Entities", variant="primary")
                
                with gr.Column():
                    clinical_output = gr.Textbox(
                        label="Clinical Entities",
                        lines=15,
                        show_copy_button=True
                    )
            
            def clinical_ner_process(text, format_type):
                if format_type == "Basic":
                    return clinical_ner_basic(text)
                else:
                    return clinical_ner_prolog(text)
            
            clinical_btn.click(
                fn=clinical_ner_process,
                inputs=[clinical_input, clinical_format],
                outputs=clinical_output
            )
        
        # Tab 2: Anatomy NER
        with gr.Tab("Anatomy Detection"):
            with gr.Row():
                with gr.Column():
                    anatomy_input = gr.Textbox(
                        label="Enter Clinical Text",
                        placeholder="Enter medical text here...",
                        lines=5,
                        value=EXAMPLE_TEXT
                    )
                    anatomy_format = gr.Radio(
                        choices=["Basic", "Prolog"],
                        value="Basic",
                        label="Output Format"
                    )
                    anatomy_btn = gr.Button("Detect Anatomy", variant="primary")
                
                with gr.Column():
                    anatomy_output = gr.Textbox(
                        label="Anatomy Entities",
                        lines=15,
                        show_copy_button=True
                    )
            
            def anatomy_ner_process(text, format_type):
                if format_type == "Basic":
                    return anatomy_ner_basic(text)
                else:
                    return anatomy_ner_prolog(text)
            
            anatomy_btn.click(
                fn=anatomy_ner_process,
                inputs=[anatomy_input, anatomy_format],
                outputs=anatomy_output
            )
        
        # Tab 3: POS Tagging
        with gr.Tab("POS Tagging"):
            with gr.Row():
                with gr.Column():
                    pos_input = gr.Textbox(
                        label="Enter Text",
                        placeholder="Enter text here...",
                        lines=5,
                        value=EXAMPLE_TEXT
                    )
                    pos_format = gr.Radio(
                        choices=["Basic", "Prolog"],
                        value="Basic",
                        label="Output Format"
                    )
                    pos_btn = gr.Button("Tag POS", variant="primary")
                
                with gr.Column():
                    pos_output = gr.Textbox(
                        label="POS Tags",
                        lines=15,
                        show_copy_button=True
                    )
            
            def pos_process(text, format_type):
                if format_type == "Basic":
                    return pos_tagging_basic(text)
                else:
                    return pos_tagging_prolog(text)
            
            pos_btn.click(
                fn=pos_process,
                inputs=[pos_input, pos_format],
                outputs=pos_output
            )
        
        # Tab 4: Combined Analysis
        with gr.Tab("Combined Analysis"):
            with gr.Row():
                with gr.Column():
                    combined_input = gr.Textbox(
                        label="Enter Clinical Text",
                        placeholder="Enter medical text here...",
                        lines=5,
                        value=EXAMPLE_TEXT
                    )
                    combined_format = gr.Radio(
                        choices=["Basic (Separated)", "Prolog (Combined)"],
                        value="Basic (Separated)",
                        label="Output Format"
                    )
                    combined_btn = gr.Button("Analyze All", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    combined_clinical = gr.Textbox(
                        label="Clinical Entities",
                        lines=10,
                        show_copy_button=True,
                        visible=True
                    )
                
                with gr.Column():
                    combined_anatomy = gr.Textbox(
                        label="Anatomy Entities",
                        lines=10,
                        show_copy_button=True,
                        visible=True
                    )
                
                with gr.Column():
                    combined_pos = gr.Textbox(
                        label="POS Tags",
                        lines=10,
                        show_copy_button=True,
                        visible=True
                    )
            
            combined_prolog_output = gr.Textbox(
                label="Combined Prolog Output",
                lines=20,
                show_copy_button=True,
                visible=False
            )
            
            def combined_process(text, format_type):
                if format_type == "Basic (Separated)":
                    clinical, anatomy, pos = combined_analysis(text)
                    return {
                        combined_clinical: gr.update(value=clinical, visible=True),
                        combined_anatomy: gr.update(value=anatomy, visible=True),
                        combined_pos: gr.update(value=pos, visible=True),
                        combined_prolog_output: gr.update(visible=False)
                    }
                else:
                    prolog = combined_prolog(text)
                    return {
                        combined_clinical: gr.update(visible=False),
                        combined_anatomy: gr.update(visible=False),
                        combined_pos: gr.update(visible=False),
                        combined_prolog_output: gr.update(value=prolog, visible=True)
                    }
            
            combined_btn.click(
                fn=combined_process,
                inputs=[combined_input, combined_format],
                outputs=[combined_clinical, combined_anatomy, combined_pos, combined_prolog_output]
            )
    
    gr.Markdown(
        """
        ---
        ### Models Used:
        - Clinical NER: `samrawal/bert-base-uncased_clinical-ner`
        - Anatomy Detection: `OpenMed/OpenMed-NER-AnatomyDetect-BioPatient-108M`
        - POS Tagging: spaCy `en_core_web_sm`
        """
    )

if __name__ == "__main__":
    demo.launch()