import gradio as gr
import pandas as pd
from scripts.inference import (
    zero_shot_inference,
    few_shot_inference,
    base_model_inference,
    fine_tuned_inference
)

def predict_with_model(text, model_type):
    """
    Applique la stratégie de classification sélectionnée sur un texte donné
    et retourne la catégorie prédite avec les scores de confiance.

    Args:
        text (str): Le texte à analyser (actualité).
        model_type (str): Le type de modèle sélectionné ("Zero-shot", "Few-shot", etc.).

    Returns:
        tuple:
            - str: Catégorie prédite.
            - pandas.DataFrame: Tableau des scores de confiance par classe.
    """

    #Sélection du modèle d'inférence en fonction du choix utilisateur
    if model_type == "Zero-shot":
        prediction, scores = zero_shot_inference(text)
    elif model_type == "Few-shot":
        prediction, scores = few_shot_inference(text)
    elif model_type == "Fine-tuned":
        prediction, scores = fine_tuned_inference(text)
    elif model_type == "Base model":
        prediction, scores = base_model_inference(text)
    else:
        return "Modèle inconnu", pd.DataFrame()

    #Convertit les scores (dict) en DataFrame pour affichage dans Gradio
    scores_df = pd.DataFrame([
        {"Classe": label, "Score": score} for label, score in scores.items()
    ])

    return prediction, scores_df

#Définition de l'interface utilisateur avec Gradio
iface = gr.Interface(
    fn=predict_with_model,  #Fonction appelée au clic de l'utilisateur
    inputs=[
        gr.Textbox(
            lines=4,
            placeholder="Entrez une phrase d'actualité ici...",
            label="Texte à classifier"
        ),
        gr.Radio(
            choices=["Base model", "Zero-shot", "Few-shot", "Fine-tuned"],
            label="Choisir le modèle",
            value="Base model"                  #Valeur par défaut
        )
    ],
    outputs=[
        gr.Label(label="Catégorie prédite"),    #Affiche la prédiction principale
        gr.BarPlot(                             #Affiche les scores de confiance
            label="Scores de confiance",
            x="Classe",
            y="Score",
            color="Classe"
        )
    ],
    title="Classification AG News (4 stratégies)",
    description="Comparer un modèle préentraîné, Zero-shot, Few-shot et Fine-tuned sur AG News"
)

#Lancement de l'application
if __name__ == "__main__":
    iface.launch()