"""Clean-room VisDrone2019-DET -> YOLO builder.

Written from the public VisDrone annotation spec (no third-party converter
code copied). Each annotation line is:
    bbox_left,bbox_top,bbox_width,bbox_height,score,category,truncation,occlusion
Categories: 0=ignored-regions 1=pedestrian 2=people 3=bicycle 4=car 5=van
6=truck 7=tricycle 8=awning-tricycle 9=bus 10=motor 11=others. We keep 1..10
(remapped to 0..9) and drop 0/11 and score==0 (ignored) boxes.

Builds:  <datasets>/visdrone/{images,labels}/{train,val}/  + train.txt/val.txt
         + visdrone.yaml
from the extracted <datasets>/visdrone_raw/VisDrone2019-DET-{train,val}/.

Usage:
    python build_visdrone.py --datasets /path/to/datasets
"""

import argparse
import shutil
from pathlib import Path

from PIL import Image

SPLITS = {"train": "VisDrone2019-DET-train", "val": "VisDrone2019-DET-val"}
NAMES = ["pedestrian", "people", "bicycle", "car", "van", "truck",
         "tricycle", "awning-tricycle", "bus", "motor"]  # VisDrone 1..10 -> 0..9


def convert_split(raw: Path, out: Path, split: str, raw_name: str) -> int:
    src = raw / raw_name
    img_src, ann_src = src / "images", src / "annotations"
    img_dst = out / "images" / split
    lbl_dst = out / "labels" / split
    img_dst.mkdir(parents=True, exist_ok=True)
    lbl_dst.mkdir(parents=True, exist_ok=True)

    listing = []
    n_imgs = n_boxes = n_dropped = 0
    for img_path in sorted(img_src.glob("*.jpg")):
        with Image.open(img_path) as im:
            W, H = im.size
        ann = ann_src / (img_path.stem + ".txt")
        lines = []
        if ann.exists():
            for raw in ann.read_text().splitlines():
                raw = raw.strip().rstrip(",")
                if not raw:
                    continue
                p = raw.split(",")
                if len(p) < 6:
                    continue
                left, top, w, h, score, cat = (int(float(x)) for x in p[:6])
                if score == 0 or cat < 1 or cat > 10 or w <= 0 or h <= 0:
                    n_dropped += 1
                    continue
                cx = (left + w / 2) / W
                cy = (top + h / 2) / H
                nw, nh = w / W, h / H
                # clip to [0,1] (a few VisDrone boxes bleed past the edge)
                cx, cy = min(max(cx, 0), 1), min(max(cy, 0), 1)
                nw, nh = min(nw, 1), min(nh, 1)
                lines.append(f"{cat - 1} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}")
                n_boxes += 1
        # copy image into the YOLO tree (move would empty the raw dir; copy is safe)
        shutil.copy2(img_path, img_dst / img_path.name)
        (lbl_dst / (img_path.stem + ".txt")).write_text("\n".join(lines))
        listing.append(f"images/{split}/{img_path.name}")
        n_imgs += 1

    (out / f"{split}.txt").write_text("\n".join(listing))
    print(f"{split}: {n_imgs} imgs, {n_boxes} boxes kept, {n_dropped} dropped")
    return n_imgs


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--datasets",
        required=True,
        help="Directory containing visdrone_raw/ with the extracted official zips",
    )
    args = parser.parse_args()
    datasets = Path(args.datasets)
    raw = datasets / "visdrone_raw"
    out = datasets / "visdrone"

    for split, raw_name in SPLITS.items():
        convert_split(raw, out, split, raw_name)
    names_block = "\n".join(f"  {i}: {n}" for i, n in enumerate(NAMES))
    (out / "visdrone.yaml").write_text(
        f"# VisDrone2019-DET (research/non-commercial license)\n"
        f"path: visdrone\ntrain: train.txt\nval: val.txt\n"
        f"nc: {len(NAMES)}\nnames:\n{names_block}\n"
    )
    print("WROTE", out / "visdrone.yaml")


if __name__ == "__main__":
    main()