import gradio as gr
from transformers import pipeline
import pandas as pd
import io

# Load once
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

LIKERT_OPTIONS = ["Strongly disagree", "Disagree", "Neutral", "Agree", "Strongly agree"]

response_table = []

def classify_likert(questions_text):
    questions = [q.strip() for q in questions_text.strip().split("\n") if q.strip()]
    
    global response_table
    response_table = []
    output_lines = []

    for i, question in enumerate(questions, 1):
        result = classifier(question, LIKERT_OPTIONS, multi_label=False)
        probs = dict(zip(result['labels'], result['scores']))
        output_lines.append(f"{i}. {question}")
        for label in LIKERT_OPTIONS:
            prob = round(probs.get(label, 0.0), 3)
            output_lines.append(f"→ {label}: {prob}")
        output_lines.append("")

        row = {"Item #": i, "Item": question}
        row.update({label: round(probs.get(label, 0.0), 3) for label in LIKERT_OPTIONS})
        response_table.append(row)

    return "\n".join(output_lines)

def download_csv():
    global response_table
    if not response_table:
        return None
    df = pd.DataFrame(response_table)
    csv_buffer = io.StringIO()
    df.to_csv(csv_buffer, index=False)
    return csv_buffer.getvalue()

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Likert-Style Zero-Shot Classifier")
    gr.Markdown("Paste questionnaire items. Each will be classified into: Strongly disagree → Strongly agree, with probabilities.")

    questions_input = gr.Textbox(label="Enter multiple items (one per line)", lines=10, placeholder="e.g.\nI feel in control of my life.\nI enjoy being around others...")
    output_box = gr.Textbox(label="Classification Output", lines=20)
    submit_btn = gr.Button("Classify Items")
    csv_btn = gr.Button("📥 Download CSV")
    file_output = gr.File(label="Download CSV", visible=False)

    submit_btn.click(fn=classify_likert, inputs=questions_input, outputs=output_box)
    csv_btn.click(fn=download_csv, inputs=[], outputs=file_output)

demo.launch()