import os
import sys
from ultralytics import YOLO
import yaml
import shutil
from pathlib import Path
import cv2
import numpy as np

def verify_images(dataset_path):
    """Überprüfe alle Bilder auf Lesbarkeit und korrekte Dimensionen."""
    print("\nÜberprüfe Bilder...")
    corrupted = []
    for split in ['train', 'val']:
        img_dir = dataset_path / split / 'images'
        if not img_dir.exists():
            continue
        
        for img_path in img_dir.glob('*'):
            try:
                img = cv2.imread(str(img_path))
                if img is None:
                    raise ValueError("Bild konnte nicht gelesen werden")
                if img.shape[0] < 10 or img.shape[1] < 10:
                    raise ValueError("Bild ist zu klein")
            except Exception as e:
                print(f"Warnung: Problem mit {img_path}: {str(e)}")
                corrupted.append(img_path)
    
    if corrupted:
        print(f"\nGefundene problematische Bilder: {len(corrupted)}")
        return False
    return True

def create_yolo_dataset():
    """Erstelle die YOLO-Dateistruktur und konvertiere die Annotationen."""
    try:
        # Erstelle Verzeichnisse
        dataset_path = Path('yolo_dataset')
        if not dataset_path.exists():
            print("Fehler: YOLO Dataset nicht gefunden! Bitte führen Sie zuerst convert_annotations.py aus.")
            sys.exit(1)

        # Überprüfe Bilder
        if not verify_images(dataset_path):
            print("Warnung: Einige Bilder könnten problematisch sein!")

        # Erstelle dataset.yaml
        dataset_yaml = {
            'path': str(dataset_path.absolute()),
            'train': 'train/images',
            'val': 'val/images',
            'names': {
                0: 'NAO-Roboter'
            },
            'nc': 1  # Anzahl der Klassen
        }

        yaml_path = dataset_path / 'dataset.yaml'
        with open(yaml_path, 'w') as f:
            yaml.dump(dataset_yaml, f)

        return str(yaml_path)
    except Exception as e:
        print(f"Fehler beim Erstellen des Datasets: {str(e)}")
        sys.exit(1)

def train_yolo():
    """Trainiere das YOLOv8 Modell."""
    # Erstelle und konfiguriere das Dataset
    dataset_yaml = create_yolo_dataset()

    # Initialisiere das Modell
    model = YOLO('yolov8n.pt')  # Verwende das kleine YOLOv8 Modell

    # Trainiere das Modell
    results = model.train(
        data=dataset_yaml,
        epochs=50,
        imgsz=640,  # Bildgröße
        batch=16,    # Batch-Größe
        device='cpu',  # Verwende CPU (oder 'cuda' für GPU)
        patience=5,   # Early Stopping
        save=True,    # Speichere die besten Gewichte
        project='yolo_training',  # Projektname
        name='NAO_detector',  # Experimentname
    )

    # Exportiere das Modell für den NAO
    print("\nExportiere Modell in verschiedene Formate...")
    
    # ONNX Format (gut für eingebettete Systeme)
    model.export(format='onnx', imgsz=640)
    
    # OpenVINO Format (optimiert für Intel Hardware)
    model.export(format='openvino', imgsz=640)
    
    print("\nTraining abgeschlossen. Die Modelle wurden im 'yolo_training/NAO_detector' Ordner gespeichert.")

if __name__ == '__main__':
    train_yolo()