import torch
import torch.nn as nn
import numpy as np
from huggingface_hub import hf_hub_download 

# ==============================================================================
# 1. Configuration et Architecture (DOIT correspondre au modèle entraîné)
# ==============================================================================

# --- Constantes du Modèle ---
REPO_ID = "Clemylia/BeeAI-Bena"
FILENAME = "pytorch_model.bin"

FLOWER_NAMES = [
    'lavande', 'coquelicot', 'muguet', 'lilas', 'jasmin', 
    'marguerite', 'rose', 'tournesol', 'acacia', 'tulipe',
    'pissenlit', 'trèfle', 'bruyère', 'romarin', 'thym',
    'sauge', 'bourrache', 'bleuet', 'primevère', 'camomille'
]
NUM_FLOWERS = len(FLOWER_NAMES) # 20
NUM_DANCE_REPORTS = 10          # 0 à 9
INPUT_SIZE = NUM_FLOWERS + 2    # 22 (20 fleurs + 2 positions)

# --- Rapports de Danse (Pour décoder la sortie 0-9) ---
RAPPORT_DESCRIPTIONS = [
    "0: La fleur n'est pas butinable, c'est hautement probable qu'elle soit vide de nectar.",
    "1: La fleur n'est pas butinable, c'est très probable qu'elle soit vide de nectar, ou qu'elle en est pas assez.",
    "2: La fleur n'est pas trop butinable, emplacement peu favorable (risque de toiles d'araignée).",
    "3: La fleur est butinable, mais dans la moindre mesure (pour une petite ruche).",
    "4: La fleur est butinable, mais ses ressources sont limitées.",
    "5: La fleur se situe trop loin de la ruche.",
    "6: La fleur se situe à un endroit susceptible d'attirer des frelons, c'est trop dangereux.",
    "7: La fleur est butinable, et possède beaucoup de nectar, mais pas assez pour toute la ruche et les larves.",
    "8: La fleur est butinable, et se situe dans une prairie remplis de fleurs hautement butinable, mais il y a un danger (grenouilles).",
    "9: La fleur est parfaitement butinable."
]

# --- Définition de la classe du modèle ---
class BeeAI(nn.Module):
    def __init__(self, input_size, num_classes):
        super(BeeAI, self).__init__()
        # L'architecture doit être identique à celle utilisée pour la sauvegarde
        self.fc1 = nn.Linear(input_size, 128) 
        self.relu = nn.ReLU() 
        self.fc2 = nn.Linear(128, 64)
        self.fc_out = nn.Linear(64, num_classes)
        
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc_out(out)
        return out

# ==============================================================================
# 2. FONCTIONS UTILITAIRES POUR LE CHARGEMENT ET L'INFÉRENCE
# ==============================================================================

def load_bee_ai_model(repo_id, filename, input_size, num_classes):
    """Télécharge les poids et charge le modèle PyTorch."""
    print(f"🔄 Tentative de téléchargement des poids depuis {repo_id}...")
    
    # 1. Télécharge le fichier de poids
    try:
        weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
    except Exception as e:
        print(f"❌ Erreur de téléchargement : {e}")
        return None
        
    # 2. Crée et charge le modèle
    model = BeeAI(input_size, num_classes)
    model.load_state_dict(torch.load(weights_path))
    model.eval() # Passe en mode évaluation pour l'inférence
    
    print("✅ Modèle BeeAI chargé avec succès.")
    return model

def flower_to_input(flower_name, x_pos, y_pos):
    """Convertit les entrées utilisateur en tenseur d'entrée (22 dimensions)."""
    try:
        idx = FLOWER_NAMES.index(flower_name.lower())
    except ValueError:
        print(f"⚠️ Avertissement: Fleur '{flower_name}' inconnue. Utilisation d'un vecteur nul.")
        idx = -1
        
    one_hot = np.zeros(NUM_FLOWERS, dtype=np.float32)
    if idx != -1:
        one_hot[idx] = 1.0
        
    position = np.array([x_pos, y_pos], dtype=np.float32)
    input_vector = np.concatenate((one_hot, position))
    # Ajout d'une dimension pour le batch
    return torch.tensor(input_vector, dtype=torch.float32).unsqueeze(0) 

def generate_dance_report(model, flower_name, x_pos, y_pos):
    """Effectue l'inférence et retourne le code de danse décodé."""
    
    # Prépare l'entrée
    X_test = flower_to_input(flower_name, x_pos, y_pos)

    # Inférence
    with torch.no_grad(): 
        output = model(X_test)
        
    # Décodage (obtient l'indice de la probabilité maximale)
    _, predicted_index = torch.max(output.data, 1) 
    report_code = predicted_index.item()
    
    return report_code, RAPPORT_DESCRIPTIONS[report_code]

# ==============================================================================
# 3. EXÉCUTION (L'expérience utilisateur)
# ==============================================================================

# 1. Charger le modèle
bee_ai_model = load_bee_ai_model(REPO_ID, FILENAME, INPUT_SIZE, NUM_DANCE_REPORTS)

if bee_ai_model:
    
    # 2. Définir la situation à tester
    test_flower = 'pissenlit' # Une fleur connue et généralement bonne
    test_x = 4.5
    test_y = 5.5
    
    # 3. Générer le rapport !
    code, description = generate_dance_report(bee_ai_model, test_flower, test_x, test_y)

    # 4. Afficher le résultat
    print(f"\n==============================================")
    print(f"      Rapport de Danse de Bena (Bee AI)       ")
    print(f"==============================================")
    print(f"Fleur évaluée : {test_flower.upper()} à la position ({test_x}, {test_y})")
    print(f"CODE GÉNÉRÉ : {code}")
    print(f"DESCRIPTION : {description.split(':')[1].strip()}")
    print(f"==============================================")