#!/usr/bin/env python3
"""
Script para corregir la configuración YAML del dataset
"""

import yaml
import os
from pathlib import Path

def fix_dataset_config():
    """Corrige el archivo de configuración YAML"""
    
    config_path = "/home/leonel/sistema_polinizador/Dataset/Classification/dataset_config.yaml"
    
    print("🔧 CORRIGIENDO CONFIGURACIÓN YAML")
    print("=" * 50)
    
    # Verificar que existe el archivo
    if not os.path.exists(config_path):
        print(f"❌ No se encontró el archivo: {config_path}")
        return False
    
    # Verificar estructura del dataset
    dataset_base_path = "/home/leonel/sistema_polinizador/Dataset/Classification"
    train_path = os.path.join(dataset_base_path, "Train")
    val_path = os.path.join(dataset_base_path, "Validation")
    test_path = os.path.join(dataset_base_path, "Test")
    
    print(f"📁 Verificando estructura del dataset:")
    print(f"   Base: {dataset_base_path} - {'✅' if os.path.exists(dataset_base_path) else '❌'}")
    print(f"   Train: {train_path} - {'✅' if os.path.exists(train_path) else '❌'}")
    print(f"   Validation: {val_path} - {'✅' if os.path.exists(val_path) else '❌'}")
    print(f"   Test: {test_path} - {'✅' if os.path.exists(test_path) else '❌'}")
    
    if not all(os.path.exists(p) for p in [dataset_base_path, train_path, val_path, test_path]):
        print("❌ Faltan carpetas del dataset")
        return False
    
    # Obtener clases del directorio Train
    class_names = []
    if os.path.exists(train_path):
        class_names = sorted([d for d in os.listdir(train_path) 
                            if os.path.isdir(os.path.join(train_path, d))])
    
    print(f"🏷️  Clases encontradas: {len(class_names)}")
    for i, name in enumerate(class_names):
        print(f"   {i}: {name}")
    
    # Crear configuración corregida
    corrected_config = {
        'path': dataset_base_path,  # Ruta absoluta al directorio base
        'train': 'Train',           # Ruta relativa desde 'path'
        'val': 'Validation',        # Ruta relativa desde 'path'
        'test': 'Test',             # Ruta relativa desde 'path'
        'nc': len(class_names),     # Número de clases
        'names': {i: name for i, name in enumerate(class_names)}  # Mapeo índice -> nombre
    }
    
    # Guardar configuración corregida
    with open(config_path, 'w') as f:
        yaml.dump(corrected_config, f, default_flow_style=False, sort_keys=False)
    
    print(f"\n✅ Configuración corregida guardada en: {config_path}")
    
    # Mostrar contenido del archivo corregido
    print(f"\n📄 Contenido del archivo YAML:")
    with open(config_path, 'r') as f:
        content = f.read()
        print(content)
    
    return True

if __name__ == "__main__":
    success = fix_dataset_config()
    if success:
        print("\n🎯 Configuración corregida. Ahora puedes ejecutar el entrenamiento.")
    else:
        print("\n❌ Error corrigiendo la configuración.")