import gradio as gr
from huggingface_hub import InferenceClient
from model import load_model, load_tokenizer
from utils import clean_output, get_shap_values
import torch
import shap
import matplotlib.pyplot as plt

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

model = load_model()
tokenizer = load_tokenizer()

def gradio_generate(context, num_questions, max_length):
    input_prompt = f"generate question: {context.strip()}"
    inputs = tokenizer(input_prompt, return_tensors="pt", truncation=True, padding="longest").to(model.device)
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=max_length,
        num_return_sequences=num_questions,
        do_sample=True,
        top_p=0.95,
        temperature=1.0
    )
    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    questions = clean_output(decoded)
    return "\n".join([f"{i+1}. {q}" for i, q in enumerate(questions)])

def gradio_shap(context):
    input_prompt = f"generate question: {context.strip()}"
    try:
        shap_values, tokens = get_shap_values(tokenizer, model, input_prompt)
        fig, ax = plt.subplots(figsize=(10, 2))
        shap.plots.text(shap.Explanation(values=shap_values, data=tokens), display=False)
        plt.tight_layout()
        return fig
    except Exception as e:
        fig, ax = plt.subplots(figsize=(8, 2))
        ax.text(0.5, 0.5, f"SHAP explanation failed:\n{e}", ha='center', va='center', wrap=True, fontsize=12)
        ax.axis('off')
        return fig

with gr.Blocks() as demo:
    gr.Markdown("# 🧠 C3QG – Context-Controlled Question Generation with FLAN-T5")
    context = gr.Textbox(label="📄 Paste your context paragraph here:", lines=6)
    num_questions = gr.Slider(1, 5, value=3, label="Number of Questions")
    max_length = gr.Slider(32, 128, value=64, label="Max Output Tokens")
    generate_btn = gr.Button("🔄 Generate Questions")
    questions_output = gr.Textbox(label="🎯 Generated Questions")
    shap_btn = gr.Button("Show SHAP Explanation")
    shap_output = gr.Plot(label="SHAP Token Importance")

    generate_btn.click(
        gradio_generate,
        inputs=[context, num_questions, max_length],
        outputs=questions_output
    )
if __name__ == "__main__":
    demo.launch()