import torch
import torch.nn.functional as F
from dotenv import load_dotenv
import os
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

#Mapping entre les ID des classes et les labels textuels
id2label = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}


def zero_shot_inference(text):
    """
    Effectue une classification zero-shot à l'aide du modèle BART MNLI.
    
    Args:
        text (str): Texte à classifier.
    
    Returns:
        tuple:
            - str: Label prédit.
            - dict: Dictionnaire {label: score} pour chaque classe.
    """
    classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
    candidate_labels = list(id2label.values())
    result = classifier(text, candidate_labels)
    prediction = result["labels"][0]
    # Formatage des scores avec 4 décimales
    scores = {label: float(f"{score:.4f}") for label, score in zip(result["labels"], result["scores"])}
    return prediction, scores


def few_shot_inference(text):
    """
    Classification few-shot avec FLAN-T5 : génère uniquement le label (World, Sports, etc.).
    
    Args:
        text (str): Texte à classifier.
    
    Returns:
        tuple:
            - str: Label prédit (nettoyé et validé).
            - dict: Détails du texte généré brut.
    """
    model_name = "google/flan-t5-small"
    classifier = pipeline("text2text-generation", model=model_name, max_new_tokens=10)

    examples = [
        ("The president met the UN delegation to discuss global peace.", "World"),
        ("The football team won their match last night.", "Sports"),
        ("The company reported a big profit this quarter.", "Business"),
        ("New research in AI shows promising results.", "Sci/Tech")
    ]

    # Prompt few-shot
    prompt = "Classify the following text into one of the following categories: World, Sports, Business, Sci/Tech.\n\n"
    for ex_text, ex_label in examples:
        prompt += f"Text: {ex_text}\nCategory: {ex_label}\n\n"
    prompt += f"Text: {text}\nCategory:"

    # Génération
    output = classifier(prompt)[0]["generated_text"].strip()

    # Nettoyage du label
    output_clean = output.split()[0].rstrip(".").capitalize()  # ex : "sci/tech." → "Sci/tech"

    # Mapping pour être sûr que ça correspond à une catégorie connue
    candidate_labels = ["World", "Sports", "Business", "Sci/Tech"]
    prediction = next((label for label in candidate_labels if label.lower() in output_clean.lower()), "Unknown")

    # Fausse distribution (1.0 pour la classe prédite, 0.0 pour les autres)
    scores = {label: 1.0 if label == prediction else 0.0 for label in candidate_labels}

    return prediction, scores


def base_model_inference(text):
    """
    Utilise un modèle BERT préentraîné sur AG News (sans fine-tuning personnalisé).
    
    Args:
        text (str): Texte à classifier.
    
    Returns:
        tuple:
            - str: Label prédit.
            - dict: Scores softmax par classe.
    """
    model_name = "textattack/bert-base-uncased-ag-news"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)

    #Encodage du texte
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)

    #Prédiction sans calcul de gradients
    with torch.no_grad():
        outputs = model(**inputs)

    #Calcul des probabilités avec softmax
    probs = F.softmax(outputs.logits, dim=-1)[0].cpu().numpy()

    pred_id = probs.argmax()
    prediction = id2label[pred_id]
    scores = {id2label[i]: float(f"{p:.4f}") for i, p in enumerate(probs)}
    return prediction, scores


def fine_tuned_inference(text, model_path="Merwan611/agnews-finetuned-bert"):
    """
    Utilise un modèle BERT fine-tuné personnalisé sur AG News, avec authentification Hugging Face si nécessaire.

    Args:
        text (str): Texte à classifier.
        model_path (str): Nom du modèle Hugging Face ou chemin local.

    Returns:
        tuple:
            - str: Label prédit.
            - dict: Scores softmax par classe.
    """

    #Récupération du token d'auth depuis les variables d'environnement
    token = os.getenv("CLE")

    tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=token)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)

    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)

    probs = F.softmax(outputs.logits, dim=-1)[0].cpu().numpy()
    pred_id = probs.argmax()
    prediction = id2label[pred_id]
    scores = {id2label[i]: float(f"{p:.4f}") for i, p in enumerate(probs)}
    return prediction, scores