import os
import shutil
from pycocotools.coco import COCO
from tqdm import tqdm

def prepare_coco_three_classes(coco_json_path, images_dir, output_dir, split):
    # Cibles :bagages
    target_classes = ['suitcase', "backpack", "handbag"]

    # Initialisation COCO
    coco = COCO(coco_json_path)

    # Récupérer les IDs de catégories pour les classes ciblées
    class_id_map = {}
    for idx, class_name in enumerate(target_classes):
        cat_ids = coco.getCatIds(catNms=[class_name])
        if cat_ids:
            class_id_map[cat_ids[0]] = idx

    print(f"Classes sélectionnées (COCO id → YOLO id) : {class_id_map}")

    # Récupérer tous les IDs d’images contenant au moins une des classes cibles
    img_ids = set()
    for cat_id in class_id_map.keys():
        img_ids.update(coco.getImgIds(catIds=[cat_id]))

    imgs = coco.loadImgs(list(img_ids))

    # Créer les dossiers de sortie
    output_images = os.path.join(output_dir, "images", split)
    output_labels = os.path.join(output_dir, "labels", split)
    os.makedirs(output_images, exist_ok=True)
    os.makedirs(output_labels, exist_ok=True)

    for img in tqdm(imgs, desc=f"Processing {split}"):
        img_path = os.path.join(images_dir, img["file_name"])
        if not os.path.exists(img_path):
            continue  # Éviter les erreurs si l’image n’est pas là

        shutil.copy(img_path, os.path.join(output_images, img["file_name"]))

        # Récupérer les annotations pour les catégories cibles
        ann_ids = coco.getAnnIds(imgIds=img['id'], catIds=list(class_id_map.keys()))
        anns = coco.loadAnns(ann_ids)

        label_path = os.path.join(output_labels, img["file_name"].replace(".jpg", ".txt"))
        with open(label_path, "w") as f:
            for ann in anns:
                if ann.get("iscrowd", 0):
                    continue  # On évite les annotations crowd

                cat_id = ann["category_id"]
                yolo_class_id = class_id_map[cat_id]

                x, y, w, h = ann["bbox"]
                cx = x + w / 2
                cy = y + h / 2
                cx /= img["width"]
                cy /= img["height"]
                w /= img["width"]
                h /= img["height"]

                f.write(f"{yolo_class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")