#!/usr/bin/env python3
"""
Script simplificado de entrenamiento YOLOv8 clasificación
"""

import torch
from ultralytics import YOLO
import os

def main():
    print("🚀 ENTRENAMIENTO YOLO CLASIFICACIÓN")
    print("=" * 50)
    
    # Verificar CUDA
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"💻 Dispositivo: {device}")
    print(f"🎯 Objetivo: >90% precisión")
    
    # Dataset path
    dataset_path = "/home/leonel/sistema_polinizador/Dataset/Classification_YOLO"
    
    if not os.path.exists(dataset_path):
        print(f"❌ Dataset no encontrado: {dataset_path}")
        print("💡 Ejecuta primero: python fix_structure.py")
        return
    
    # Configuraciones de entrenamiento
    configs = [
        {
            "name": "nano_quick",
            "model": "yolov8n-cls.pt",
            "epochs": 30,
            "imgsz": 224,
            "batch": 32
        },
        {
            "name": "small_balanced", 
            "model": "yolov8s-cls.pt",
            "epochs": 60,
            "imgsz": 256,
            "batch": 16
        },
        {
            "name": "medium_accurate",
            "model": "yolov8m-cls.pt", 
            "epochs": 100,
            "imgsz": 320,
            "batch": 8
        }
    ]
    
    best_accuracy = 0
    best_model = None
    
    for i, config in enumerate(configs, 1):
        print(f"\n{i}️⃣ MODELO: {config['name']}")
        print("=" * 40)
        
        try:
            # Cargar modelo
            model = YOLO(config["model"])
            print(f"📥 Modelo cargado: {config['model']}")
            
            # Entrenar
            print(f"⏰ Iniciando entrenamiento...")
            results = model.train(
                data=dataset_path,
                epochs=config["epochs"],
                imgsz=config["imgsz"],
                batch=config["batch"],
                device=device,
                project="pollinator_final",
                name=config["name"],
                patience=20,
                save=True,
                verbose=False,
                plots=True
            )
            
            # Evaluar
            print(f"📊 Evaluando en test set...")
            test_results = model.val(split='test')
            accuracy = float(test_results.top1) * 100
            
            print(f"✅ Entrenamiento completado")
            print(f"🎯 Precisión: {accuracy:.2f}%")
            
            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_model = f"pollinator_final/{config['name']}/weights/best.pt"
            
            # Verificar objetivo
            if accuracy >= 90:
                print(f"🎉 ¡OBJETIVO ALCANZADO! {accuracy:.2f}% ≥ 90%")
                break
            else:
                print(f"⚠️  Faltan {90-accuracy:.2f}% para objetivo")
                
        except Exception as e:
            print(f"❌ Error: {e}")
            continue
    
    # Resultados finales
    print(f"\n" + "=" * 50)
    print("📊 RESULTADOS FINALES")
    print("=" * 50)
    print(f"🏆 Mejor precisión: {best_accuracy:.2f}%")
    
    if best_accuracy >= 90:
        print(f"✅ OBJETIVO ALCANZADO!")
    else:
        print(f"❌ Objetivo no alcanzado")
        print(f"💡 Recomendación: Entrenar modelo YOLOv8l o YOLOv8x")
    
    if best_model:
        print(f"📁 Mejor modelo: {best_model}")
        
        # Crear script de predicción simple
        pred_script = f'''#!/usr/bin/env python3
from ultralytics import YOLO

# Cargar modelo entrenado
model = YOLO('{best_model}')

# Función para clasificar
def classify_insect(image_path):
    results = model(image_path, verbose=False)
    probs = results[0].probs
    
    classes = [
        'Acmaeodera Flavomarginata', 'Acromyrmex Octospinosus', 
        'Adelpha Basiloides', 'Adelpha Iphicleola', 'Aedes Aegypti',
        'Agrius Cingulata', 'Anaea Aidea', 'Anartia fatima',
        'Anartia jatrophae', 'Anoplolepis Gracilipes'
    ]
    
    top_class = classes[probs.top1]
    confidence = probs.top1conf.item() * 100
    
    print(f"🔍 Predicción: {{top_class}}")
    print(f"📊 Confianza: {{confidence:.1f}}%")
    
    return top_class, confidence

# Ejemplo de uso
if __name__ == "__main__":
    image_path = input("Ruta de imagen: ")
    if image_path:
        classify_insect(image_path)
'''
        
        with open('predict_final.py', 'w') as f:
            f.write(pred_script)
        
        print(f"✅ Script de predicción: predict_final.py")
    
    return best_accuracy

if __name__ == "__main__":
    final_accuracy = main()
    print(f"\n🎯 Entrenamiento completado. Precisión final: {final_accuracy:.2f}%")