import gradio as gr
from transformers import pipeline

# ===============================
# Model dictionary (lazy loaded)
# ===============================
model_names = {
    "DistilGPT-2": "distilgpt2",
    "Bloom-560M": "bigscience/bloom-560m",
    "OPT-350M": "facebook/opt-350m",
    "Flan-T5-Base": "google/flan-t5-base",
    "Phi-2": "microsoft/phi-2"
}

loaded_models = {}
summarizer = None  # Flan-T5 for cleanup

# ===============================
# Lazy-load helper
# ===============================
def get_model(name):
    if name not in loaded_models:
        mdl = model_names[name]
        if "flan" in mdl.lower():
            loaded_models[name] = pipeline("text2text-generation", model=mdl)
        else:
            loaded_models[name] = pipeline("text-generation", model=mdl)
    return loaded_models[name]

def get_summarizer():
    global summarizer
    if summarizer is None:
        summarizer = pipeline("text2text-generation", model="google/flan-t5-base")
    return summarizer

# ===============================
# Compare function
# ===============================
def compare_models(user_input, max_new_tokens=100, temperature=0.7):
    raw_outputs, clean_outputs = {}, {}
    for name in model_names.keys():
        try:
            generator = get_model(name)
            if generator.task == "text-generation":
                output = generator(
                    user_input,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature
                )[0]["generated_text"]
            else:  # text2text-generation (Flan)
                output = generator(user_input, max_new_tokens=max_new_tokens)[0]["generated_text"]

            raw_outputs[name] = output

            # Summarize
            summary = get_summarizer()("Summarize this: " + output, max_new_tokens=60)[0]["generated_text"]
            clean_outputs[name] = summary

        except Exception as e:
            raw_outputs[name] = f"⚠️ Error: {str(e)}"
            clean_outputs[name] = "N/A"

    return [raw_outputs[m] for m in model_names.keys()], [clean_outputs[m] for m in model_names.keys()]

# ===============================
# Gradio UI
# ===============================
with gr.Blocks(css="style.css") as demo:
    gr.Markdown("## 🤖 Open-Source Model Comparator\n"
                "Compare outputs from open-source LLMs side by side.\n"
                "Raw output + a cleaned summary from Flan-T5.")

    with gr.Row():
        user_input = gr.Textbox(label="Your prompt", placeholder="Try: 'Explain quantum computing in simple terms'", lines=2)
        generate_btn = gr.Button("Generate", variant="primary")

    with gr.Row():
        max_tokens = gr.Slider(20, 200, value=100, step=10, label="Max new tokens")
        temp = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Creativity (temperature)")

    gr.Markdown("### 🔎 Raw Outputs")
    with gr.Row():
        raw_boxes = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in model_names.keys()]

    gr.Markdown("### ✨ Cleaned Summaries (Flan-T5)")
    with gr.Row():
        clean_boxes = [gr.Textbox(label=f"{name} (Summary)", elem_classes="output-box", interactive=False) for name in model_names.keys()]

    examples = [
        ["Explain quantum computing in simple terms."],
        ["Write a haiku about autumn leaves."],
        ["What are the pros and cons of nuclear energy?"],
        ["Describe a futuristic city in the year 2200."],
        ["Write a funny short story about a robot learning to cook."]
    ]
    gr.Examples(examples=examples, inputs=[user_input])

    generate_btn.click(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes)
    user_input.submit(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes)

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