import os
import gradio as gr
from run_llm import main
import argparse

theme = gr.themes.Soft()

def run_llm_interface(model, task, sentence):
    args = argparse.Namespace(
        model_path=model,
        prompt=task,
        start=0,
        end=1  # Set to 1 to process a single sentence
    )
    main(args)
    
    # Create directories if they don't exist
    os.makedirs(f'result/prompt1_qa/{model}/ptb/per_ent/NOUN', exist_ok=True)
    os.makedirs(f'result/prompt2_instruction/chunking/{model}/ptb', exist_ok=True)
    os.makedirs(f'result/prompt3_structured_prompt/chunking/{model}/ptb', exist_ok=True)
    
    # Read the outputs from the result files
    with open(f'result/prompt1_qa/{model}/ptb/per_ent/NOUN/0.txt', 'r') as f:
        output_1 = f.read()
    
    with open(f'result/prompt2_instruction/chunking/{model}/ptb/0.txt', 'r') as f:
        output_2 = f.read()
    
    with open(f'result/prompt3_structured_prompt/chunking/{model}/ptb/0.txt', 'r') as f:
        output_3 = f.read()

    return {"output_1": output_1, "output_2": output_2, "output_3": output_3}

# Define example instructions for testing
instruction_examples = [
    ["gpt3.5", "POS Tagging", "Describe the origin of the universe"],
    ["vicuna-7b", "Chunking", "Explain the concept of artificial intelligence"],
    ["fastchat-t5", "Parsing", "Describe the most common types of cancer"],
]

with gr.Interface(
    fn=run_llm_interface,
    inputs=[
        gr.Dropdown(['gpt3.5', 'vicuna-7b', 'vicuna-13b', 'fastchat-t5', 'llama-7b', 'llama-13b', 'llama-30b', 'alpaca'], label="Select Model", default='gpt3.5', key="model"),
        gr.Dropdown(['POS Tagging', 'Chunking', 'Parsing'], label="Select Task", default='POS Tagging', key="task"),
        gr.Textbox("", label="Enter Sentence", key="sentence", placeholder="Enter a sentence..."),
    ],
    outputs=[
        gr.Textbox("", label="Strategy 1 Output", key="output_1", readonly=True),
        gr.Textbox("", label="Strategy 2 Output", key="output_2", readonly=True),
        gr.Textbox("", label="Strategy 3 Output", key="output_3", readonly=True),
    ],
    examples=instruction_examples, 
    live=False,
    title="LLM Evaluator with Linguistic Scrutiny",
    theme=theme
) as iface:
    iface.launch()