"""
STEP 1 ─ prepare_dataset.py
===========================
Run this FIRST. It creates the required folder structure and
a dataset.yaml that YOLOv8 expects.

Usage:
    python prepare_dataset.py

After running. your folder tree will look like this:
    dataset/
      images/
        train/   <- put your labeled images here.
        val/     <- validation split (auto-filled by augment.py)

      labels/
        train/   <- put your .txt label files here.
        val/     <- auto-generated.
    dataset.yaml <- auto-generated.

Label format (one .txt per image, same name):
    <class_id> <cx> <cy> <w> <h> (all values 0.0-1.0, normalised)

Recommended free labeling tools:
    • Roboflow   https://roboflow.com        (easiest, web-based)
    • LabelImg   pip install labelImg        (desktop, saves YOLO format directly)
"""
import os
import yaml

# ── Edit this list to match your product classes ────────
CLASS_NAMES = [
        "1001",
        "Aroma",
        "Aroma Gold",
        "Bonal",
        "Facto",
        "Famico",
        "Famico Exclusive",
        "Fegalo",
        "Gosto",
        "Molino",
        "Oscar",
        "Primo",
        "Siglo"
]

DATASET_ROOT = "dataset"

DIRS = [
    f"{DATASET_ROOT}/images/train",
    f"{DATASET_ROOT}/images/val",
    f"{DATASET_ROOT}/labels/train",
    f"{DATASET_ROOT}/labels/val",
]

def Structure_Create():
    for d in DIRS:
        os.makedirs(d, exist_ok=True)
    print("[1/3] Folder structure created:")

    for d in DIRS:
        print(f"         {d}/")

def YAML_Write():
    cfg = {
            "path" : os.path.abspath(DATASET_ROOT),
            "train": "images/train",
            "val"  : "images/val",
            "nc"   : len(CLASS_NAMES),
            "names": CLASS_NAMES,
        }

    yaml_path = "dataset.yml"
    with open(yaml_path, "w") as f:
        yaml.dump(cfg, f, default_flow_style=False, sort_keys=False)

    print("[2/3] dataset.yaml written -> {yaml_path}")
    return yaml_path

def Print_NextSetups():
    print("""
[3/3] NEXT STEPS
─────────────────────────────────────────────────────────────
1. Put your RAW images (one per product view) into:
       dataset/images/train/

2. Label them using Roboflow or LabelImg (YOLO format).
   Each image  →  a matching .txt in dataset/labels/train/
   Example label file  (famico_veloute_01.txt):
       0 0.512 0.483 0.320 0.541

3. Run the augmentation script:
       python 2_augment.py

4. Train:
       python 3_train.py

5. Detect on new scenes:
       python 4_detect.py --source scene.jpg
─────────────────────────────────────────────────────────────
    """)

if __name__ == "__main__":
    Structure_Create()
    YAML_Write()
    Print_NextSetups()