"""COCO -> manifest rows for stop signs.

Free stop-sign data that the pretrained weights have already seen. No depth
ground truth, so gt_distance_m is null -- these rows train the detector and are
invisible to the phase 2 scoring.

Domain caveat, worth remembering when stop-sign mAP looks strange: COCO stop
signs are photographs from arbitrary viewpoints, including close-ups and
non-road contexts. nuScenes and AV2 frames are dashcam views. A model trained
on both can do well on COCO-style stop signs and poorly on dashcam ones.

Usage:
    python -m src.data.extract_coco --dataroot /data/coco --split train2017
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path

from src.common import paths, schema
from src.common.geometry import clip_box

SOURCE = "coco"
COCO_CATEGORY = "stop sign"
UNIFIED_CLASS = "stop_sign"

# COCO images have no shared calibration. Phase 2 must not try to estimate
# distance from them, and this sensor_id has no calib/*.yaml on purpose.
SENSOR_ID = "coco_unknown"


def extract(dataroot: Path, split: str, root: Path, max_images: int | None) -> None:
    annotation_path = dataroot / "annotations" / f"instances_{split}.json"
    print(f"reading {annotation_path} (this takes a moment, it is a large file)")
    with open(annotation_path) as handle:
        coco = json.load(handle)

    category_ids = [c["id"] for c in coco["categories"] if c["name"] == COCO_CATEGORY]
    if not category_ids:
        raise SystemExit(f"no category named {COCO_CATEGORY!r} in {annotation_path}")
    category_id = category_ids[0]

    annotations_by_image: dict[int, list[dict]] = defaultdict(list)
    for annotation in coco["annotations"]:
        if annotation["category_id"] == category_id and not annotation.get("iscrowd", 0):
            annotations_by_image[annotation["image_id"]].append(annotation)

    images = {image["id"]: image for image in coco["images"]}
    image_ids = sorted(annotations_by_image)
    if max_images is not None:
        image_ids = image_ids[:max_images]

    paths.link_source(root, SOURCE, dataroot)

    rows: list[dict] = []
    skipped = 0
    for image_id in image_ids:
        image = images[image_id]
        image_path = f"{SOURCE}/{split}/{image['file_name']}"

        for annotation in annotations_by_image[image_id]:
            x, y, w, h = annotation["bbox"]
            clipped, truncation = clip_box(
                (x, y, x + w, y + h), image["width"], image["height"]
            )
            if clipped[2] - clipped[0] <= 1.0 or clipped[3] - clipped[1] <= 1.0:
                skipped += 1
                continue

            rows.append(
                schema.object_row(
                    image_path=image_path,
                    sensor_id=SENSOR_ID,
                    source=SOURCE,
                    # Each COCO photo is its own scene; there is no video to leak.
                    scene_id=f"{SOURCE}_{image_id}",
                    class_name=UNIFIED_CLASS,
                    box=clipped,
                    gt_distance_m=None,
                    gt_dims_hwl=None,
                    visibility=None,
                    truncation=truncation,
                )
            )

    frame = schema.rows_to_frame(rows)
    part = paths.part_path(root, SOURCE)
    part.parent.mkdir(parents=True, exist_ok=True)
    schema.write_manifest(frame, part)

    print(f"\nwrote {part}")
    print(f"  frames  : {frame['image_path'].nunique()}")
    print(f"  objects : {len(frame)}")
    if skipped:
        print(f"  dropped : {skipped} degenerate boxes")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--dataroot", required=True, type=Path,
                        help="COCO root, containing annotations/ and train2017/")
    parser.add_argument("--split", default="train2017")
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--max-images", type=int, default=None)
    args = parser.parse_args()

    extract(
        dataroot=args.dataroot,
        split=args.split,
        root=paths.unified_root(args.unified_root),
        max_images=args.max_images,
    )


if __name__ == "__main__":
    main()