"""
Script 1 — Crop Dataset Generator
===================================
Runs your fine-tuned YOLOv8n model on a folder of shelf images,
crops every detection, and dumps everything into a single flat folder.

You classify ONCE, then run script 1b_split_dataset.py to create
the train / val / test split automatically.

Output structure
----------------
data/
    unknown/
        img001_crop0000.jpg
        img001_crop0001.jpg
        img002_crop0000.jpg
        ...

After generation
----------------
1. Rename / move crops into class subfolders inside data/:
       data/
           cola_can/    img001_crop0000.jpg  …
           pepsi_can/   img003_crop0002.jpg  …
           …

   Use Label Studio, a file manager, or any annotation tool.
   The `unknown/` folder holds anything you haven't labelled yet.

2. Then run:
       python 1b_split_dataset.py --data_dir data --output_dir crops_dataset

Usage
-----
python 1_generate_crop_dataset.py \
    --source     /path/to/shelf/images \
    --weights    /path/to/best.pt \
    --output_dir data \
    --conf       0.25 \
    --iou        0.45 \
    --img_size   640 \
    --padding    10 \
    --min_size   32
"""

import argparse
from pathlib import Path

import cv2
import numpy as np
from ultralytics import YOLO
from tqdm import tqdm


# ─────────────────────────── helpers ─────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="Generate flat crop dataset from YOLOv8 detections")
    p.add_argument("--source",      required=True             ,  help="Directory with shelf images (jpg/png/…)")
    p.add_argument("--weights",     default="detector/best.pt",  help="Path to your fine-tuned YOLOv8n .pt file")
    p.add_argument("--output_dir",  default="data"            , help="Root output dir — crops go into <output_dir>/unknown/")
    p.add_argument("--conf",        type=float, default=0.25  , help="Detection confidence threshold")
    p.add_argument("--iou",         type=float, default=0.45  , help="NMS IoU threshold")
    p.add_argument("--img_size",    type=int,   default=640   ,  help="YOLO inference image size")
    p.add_argument("--padding",     type=int,   default=10    ,   help="Extra pixels around each crop (px)")
    p.add_argument("--min_size",    type=int,   default=32    ,   help="Discard crops smaller than this (px)")
    p.add_argument("--save_labels", action="store_true",
                   help="Save a sidecar .txt with YOLO cls_id / confidence next to each crop")
    return p.parse_args()


def get_image_paths(source: str):
    exts = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
    p = Path(source)

    if p.is_file():
        # Single image passed directly  e.g. --source 01.jpg
        if p.suffix.lower() not in exts:
            raise ValueError(f"File is not a supported image type: {p}")
        return [p]

    if p.is_dir():
        # Folder — recurse for all images
        paths = sorted(f for f in p.rglob("*") if f.suffix.lower() in exts)
        if not paths:
            raise FileNotFoundError(f"No images found in directory: {p}")
        return paths

    # Glob pattern  e.g. --source "frames/*.jpg"
    paths = sorted(f for f in Path(".").glob(source) if f.suffix.lower() in exts)
    if not paths:
        raise FileNotFoundError(f"No images matched: {source}")
    return paths


def crop_detection(image: np.ndarray, box, padding: int, min_size: int):
    """
    Extract one padded crop from *image* using *box* (xyxy).
    Returns the crop array, or None if it is too small.
    """
    x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
    h, w = image.shape[:2]

    x1 = max(0, x1 - padding)
    y1 = max(0, y1 - padding)
    x2 = min(w, x2 + padding)
    y2 = min(h, y2 + padding)

    if (x2 - x1) < min_size or (y2 - y1) < min_size:
        return None

    return image[y1:y2, x1:x2]


# ─────────────────────────── main ────────────────────────────────────────────

def main():
    args    = parse_args()
    out_dir = Path(args.output_dir) / "unknown"
    out_dir.mkdir(parents=True, exist_ok=True)

    # ── Load model ────────────────────────────────────────────────────────────
    print(f"[1/2] Loading model: {args.weights}")
    model = YOLO(args.weights)

    # ── Detect & crop ─────────────────────────────────────────────────────────
    image_paths = get_image_paths(args.source)
    print(f"[2/2] Found {len(image_paths)} images — running detection …\n")

    total_saved = 0
    total_skip  = 0

    for img_path in tqdm(image_paths, desc="Detecting & cropping"):
        image = cv2.imread(str(img_path))
        if image is None:
            print(f"  ⚠  Cannot read {img_path.name} — skipping")
            continue

        results = model(
            image,
            conf=args.conf,
            iou=args.iou,
            imgsz=args.img_size,
            verbose=False,
        )[0]

        stem = img_path.stem
        for i, box in enumerate(results.boxes):
            crop = crop_detection(image, box, args.padding, args.min_size)

            if crop is None:
                total_skip += 1
                continue

            crop_name = f"{stem}_crop{i:04d}.jpg"
            crop_path = out_dir / crop_name
            cv2.imwrite(str(crop_path), crop)
            total_saved += 1

            if args.save_labels:
                cls_id   = int(box.cls[0])
                conf     = float(box.conf[0])
                cls_name = model.names.get(cls_id, str(cls_id))
                crop_path.with_suffix(".txt").write_text(
                    f"{cls_id} {cls_name} {conf:.4f}\n"
                )

    # ── Summary ───────────────────────────────────────────────────────────────
    print("\n" + "=" * 55)
    print("  Crop generation complete!")
    print(f"  Saved   : {total_saved:,} crops  ->  {out_dir.resolve()}")
    print(f"  Skipped : {total_skip:,}  (below --min_size={args.min_size}px)")
    print("=" * 55)
    print("""
  Next steps:
    1. Organise crops into class subfolders inside data/:
           data/
               cola_can/    *.jpg
               pepsi_can/   *.jpg
               ...
       Leave unlabelled crops in data/unknown/ — the split
       script ignores that folder automatically.

    2. Run the splitter:
           python 1b_split_dataset.py \\
               --data_dir   data \\
               --output_dir crops_dataset \\
               --val_split  0.15 \\
               --test_split 0.05
    """)


if __name__ == "__main__":
    main()