import gradio as gr
from pipeline.predict import predict_fake_news, load_model_and_predict
from collections import Counter

# Model label mapping
model_labels = {
    0: "BiLSTM Neural Network",
    1: "Logistic Regression",
    2: "Naive Bayes",
    3: "XGBoost",
    4: "LightGBM",
    5: "Random Forest"
}

model_choices = [(name, idx) for idx, name in model_labels.items()]


def run_single_model(text, model_index):
    if model_index == 0:
        result = predict_fake_news(text)
    else:
        result = load_model_and_predict(model_index, [text])[0]

    return f"Model: {model_labels[model_index]}\n" \
           f"Prediction: {result['verdict']}\n" \
           f"Confidence: {result['confidence']}"


def run_all_models(text=None, file=None):
    if not text and not file:
        return "Please enter text or upload a file."

    if not text and file:
        with open(file.name, 'r', encoding='utf-8') as f:
            text = f.read()

    results = {}
    predictions = []

    for idx in range(6):
        if idx == 0:
            res = predict_fake_news(text)
        else:
            res = load_model_and_predict(idx, [text])[0]
        results[model_labels[idx]] = res
        predictions.append(res["label"])

    majority_label = Counter(predictions).most_common(1)[0][0]
    majority_verdict = "Real News" if majority_label == 1 else "Fake News"

    result_table = "\n".join(
        f"{model}: {res['verdict']} (Confidence: {res['confidence']})"
        for model, res in results.items()
    )

    return f"{result_table}\n\n🗳️ Overall Majority Verdict: **{majority_verdict}**"


# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 📰 Fake News Classifier - Multi-Model Evaluation")

    with gr.Tab("🔍 Single Model Prediction"):
        with gr.Row():
            text_input = gr.Textbox(label="Enter News Text", lines=10, placeholder="Paste news text here...")
            model_dropdown = gr.Dropdown(choices=model_choices, label="Select Model", value=0)
        single_output = gr.Textbox(label="Prediction", lines=5)
        predict_btn = gr.Button("Predict")
        predict_btn.click(fn=run_single_model, inputs=[text_input, model_dropdown], outputs=single_output)

    with gr.Tab("📊 All Models Comparison"):
        gr.Markdown("Paste news text below **or** upload a `.txt` file")
        text_input_all = gr.Textbox(label="Paste News Text (optional)", lines=10, placeholder="Paste here...")
        file_input_all = gr.File(label="Upload .txt File (optional)", file_types=[".txt"])
        compare_btn = gr.Button("Run All Models")
        multi_output = gr.Textbox(label="Model-wise Prediction & Majority Verdict", lines=12)
        compare_btn.click(fn=run_all_models, inputs=[text_input_all, file_input_all], outputs=multi_output)


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