import argparse
import glob
import os
import torch
from ultralytics import YOLO

# uv run inference.py path_to_model.pt path_to_images ./test


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Prédiction YOLO sur un dossier d'images avec sauvegarde des résultats annotés et des labels .txt"
    )
    parser.add_argument(
        "model_path",
        type=str,
        help="Chemin vers le modèle YOLO (.pt)",
    )
    parser.add_argument(
        "images_path",
        type=str,
        help="Dossier contenant les images à traiter",
    )
    parser.add_argument(
        "output_path",
        type=str,
        help="Dossier de sortie pour les images annotées et les labels",
    )
    parser.add_argument(
        "--imgsz",
        type=int,
        default=640,
        help="Taille d'image pour l'inférence (défaut: 640)",
    )
    parser.add_argument(
        "--half",
        action="store_true",
        help="Utiliser la précision FP16 pour réduire la consommation mémoire",
    )

    return parser.parse_args()


def save_yolo_txt(result, txt_path):
    """
    Sauvegarde les prédictions d'un résultat Ultralytics au format YOLO :
    <class_id> <x_center> <y_center> <width> <height> <confidence>
    (coordonnées normalisées entre 0 et 1)
    """
    boxes = result.boxes
    with open(txt_path, "w") as f:
        for box in boxes:
            class_id = int(box.cls[0])
            conf = float(box.conf[0])
            x_center, y_center, width, height = box.xywhn[0].tolist()
            f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {conf:.6f}\n")


def predict(model_path: str, images_path: str, output_path: str, imgsz: int, half: bool):
    model = YOLO(model_path)

    images = glob.glob(os.path.join(images_path, "*.jpg"))
    selection = images
    print(f"{len(selection)} images sélectionnées pour la prédiction")

    # Nom du modèle (sans extension) pour organiser les résultats par modèle testé
    model_name = model_path.split('/')
    model_name = model_name[0]

    images_dir = os.path.join(output_path, model_name, "images")
    labels_dir = os.path.join(output_path, model_name, "labels")
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    total = len(selection)
    for i, image_path in enumerate(selection, start=1):
        image_name = os.path.splitext(os.path.basename(image_path))[0]

        # Inférence sur une seule image à la fois
        with torch.no_grad():
            results = model(image_path, imgsz=imgsz, verbose=False)
        result = results[0]

        # Sauvegarde de l'image annotée
        result.save(filename=os.path.join(images_dir, f"{image_name}.jpg"))

        # Sauvegarde des prédictions au format YOLO .txt
        save_yolo_txt(result, os.path.join(labels_dir, f"{image_name}.txt"))

        # Libération explicite de la mémoire GPU après chaque image
        del results, result
        torch.cuda.empty_cache()

        if i % 10 == 0 or i == total:
            print(f"[{i}/{total}] traitées")

    print(f"Modèle testé                      : {model_name}")
    print(f"Images annotées sauvegardées dans : {images_dir}")
    print(f"Labels YOLO sauvegardés dans      : {labels_dir}")


def main() -> None:
    args = parse_args()
    predict(args.model_path, args.images_path, args.output_path, args.imgsz, args.half)


if __name__ == "__main__":
    main()