"""Argoverse 2 -> manifest rows + calibration YAML.

Second calibration source, and the only source of stop signs that come with
ground-truth distance. Droppable if time runs short -- nuScenes alone is a
complete submission.

This reads the AV2 feather files directly with pandas rather than going through
the devkit dataloader. The layout is stable and documented, and it means the
projection maths is literally the same code path as the nuScenes extractor
instead of a second implementation that might disagree.

Expected layout under --dataroot:
    <log_id>/annotations.feather
    <log_id>/calibration/intrinsics.feather
    <log_id>/calibration/egovehicle_SE3_sensor.feather
    <log_id>/sensors/cameras/<camera>/<timestamp_ns>.jpg

Usage:
    python -m src.data.extract_av2 --dataroot /data/av2/sensor/train
"""

from __future__ import annotations

import argparse
from collections import Counter
from pathlib import Path

import numpy as np
import pandas as pd

from src.common import calib, paths, schema
from src.common.geometry import (
    box_from_points,
    clip_box,
    cuboid_corners,
    point_in_box,
    quaternion_to_rotation_matrix,
)

SOURCE = "av2"

# Barrel is deliberately NOT mapped to barrier. An AV2 CONSTRUCTION_BARREL is an
# orange drum; a nuScenes barrier is a jersey/fence barrier. Merging them tanks
# the class and makes it look like a training problem.
# CONSTRUCTION_BARREL is deliberately not mapped. It is not an evaluation class,
# and an orange drum is a different object from a jersey barrier -- mapping it to
# barrier would merge two visually unrelated things into one label.
AV2_CLASS_MAP = {
    "CONSTRUCTION_CONE": "cone",
    "STOP_SIGN": "stop_sign",
}

MIN_CORNER_DEPTH_M = 0.1

# Annotations are at 10 Hz, ring cameras at 20 Hz. Pair each annotation sweep
# with the single nearest image rather than duplicating labels onto both.
MAX_TIME_OFFSET_NS = 25_000_000  # 25 ms


# Ground-contact categories, used to locate the road surface in the ego frame.
# Used only to locate the road surface, not to produce labels. Barrels rest on
# the road and are plentiful, so they sharpen the estimate even though the class
# itself is not kept.
GROUND_CATEGORIES = ("CONSTRUCTION_CONE", "CONSTRUCTION_BARREL")

# Objects used for that estimate. Near enough that road grade has not diverged
# from the plane under the vehicle, far enough to be outside the bumper.
GROUND_FIT_RANGE_M = (8.0, 40.0)
MIN_GROUND_SAMPLES = 50


def sensor_id_for(camera: str, log_id: str) -> str:
    """One sensor per log, not one per camera.

    Unlike nuScenes, AV2 intrinsics and extrinsics vary between logs -- across
    the fifteen logs here, focal length moves by 1% and camera height by 6%.
    Collapsing them into a single "av2_ring_front_center" would hand phase 2 one
    arbitrary log's calibration for every image.
    """
    return f"{SOURCE}_{log_id[:8]}_{camera.lower()}"


def ground_height_in_ego(annotations: pd.DataFrame) -> float | None:
    """Where the road surface sits in the ego frame, from objects resting on it.

    nuScenes puts its ego origin on the ground, so its camera height is simply
    the sensor's z translation. AV2 does not: its origin sits at the rear axle
    centre, roughly 0.26 m up, and taking z translation as camera height makes
    every ground-plane distance read about 18% short. Rather than hardcode a
    vehicle constant, measure it -- cones and barrels rest on the road, so the
    median underside of their cuboids is the road surface.
    """
    ground = annotations[annotations["category"].isin(GROUND_CATEGORIES)]
    if ground.empty:
        return None
    distance = np.hypot(ground["tx_m"], ground["ty_m"])
    low, high = GROUND_FIT_RANGE_M
    near = ground[(distance > low) & (distance < high) & (ground["tx_m"] > 0)]
    if len(near) < MIN_GROUND_SAMPLES:
        return None
    return float((near["tz_m"] - near["height_m"] / 2).median())


def read_calibration(log_dir: Path, camera: str) -> calib.Calibration:
    """Read K, lens distortion and the ego<-cam transform for one log's camera."""
    intrinsics = pd.read_feather(log_dir / "calibration" / "intrinsics.feather")
    extrinsics = pd.read_feather(log_dir / "calibration" / "egovehicle_SE3_sensor.feather")

    row = intrinsics[intrinsics["sensor_name"] == camera]
    if row.empty:
        raise KeyError(f"camera {camera!r} not in {log_dir/'calibration'/'intrinsics.feather'}")
    row = row.iloc[0]
    K = np.array([
        [row["fx_px"], 0.0, row["cx_px"]],
        [0.0, row["fy_px"], row["cy_px"]],
        [0.0, 0.0, 1.0],
    ], dtype=float)

    pose = extrinsics[extrinsics["sensor_name"] == camera].iloc[0]
    rotation = quaternion_to_rotation_matrix(pose["qw"], pose["qx"], pose["qy"], pose["qz"])

    # The ego origin is NOT on the ground here, so measure where the ground is.
    ground_z = ground_height_in_ego(pd.read_feather(log_dir / "annotations.feather"))
    if ground_z is None:
        ground_z = 0.0
        print(f"  !! {log_dir.name[:8]}: too few ground objects to locate the road "
              f"surface; falling back to the ego origin, distances will read short")

    return calib.Calibration(
        sensor_id=sensor_id_for(camera, log_dir.name),
        width=int(row["width_px"]),
        height=int(row["height_px"]),
        K=K,
        cam_height_m=float(pose["tz_m"]) - ground_z,
        R_ego_from_cam=rotation,
        distortion=(float(row["k1"]), float(row["k2"]), float(row["k3"])),
    )


def camera_timestamps(log_dir: Path, camera: str) -> tuple[np.ndarray, list[Path]]:
    """Sorted image timestamps (ns) and their paths."""
    image_dir = log_dir / "sensors" / "cameras" / camera
    image_paths = sorted(image_dir.glob("*.jpg"))
    timestamps = np.array([int(path.stem) for path in image_paths], dtype=np.int64)
    return timestamps, image_paths


def nearest_image(timestamp_ns: int, timestamps: np.ndarray,
                  image_paths: list[Path]) -> Path | None:
    if len(timestamps) == 0:
        return None
    index = int(np.argmin(np.abs(timestamps - timestamp_ns)))
    if abs(int(timestamps[index]) - timestamp_ns) > MAX_TIME_OFFSET_NS:
        return None
    return image_paths[index]


def sign_face_only(centre_ego: np.ndarray, rotation_ego: np.ndarray,
                   width_m: float, height_m: float) -> tuple[np.ndarray, float]:
    """Shrink a STOP_SIGN cuboid from the whole assembly down to the sign face.

    AV2 annotates a stop sign as one cuboid spanning the ground to the top of the
    sign -- median 3.22 m tall by 0.86 m wide -- so projecting it whole gives a
    2D box with an aspect ratio near 4:1 that is mostly pole. COCO stop-sign
    boxes are tight on the octagon at roughly 1:1. Training on both definitions
    at once teaches the detector two incompatible ideas of the same class, and
    the pole-inclusive box is the wrong thing to feed a width-based distance
    estimator.

    A stop sign is a regular octagon, so its face is as tall as it is wide. Take
    the cuboid's own width as the face height and keep the top slice of the
    cuboid -- data-driven, no assumed sign standard, and it degrades sensibly for
    the oversized signs used on multilane roads.
    """
    face_height = min(width_m, height_m)
    # Object frame is +Z up, so lift the centre to the middle of the top slice.
    offset = rotation_ego @ np.array([0.0, 0.0, (height_m - face_height) / 2.0])
    return centre_ego + offset, face_height


def rows_for_log(log_dir: Path, camera: str, calibration: calib.Calibration,
                 dataroot: Path, dropped: Counter) -> list[dict]:
    annotations = pd.read_feather(log_dir / "annotations.feather")
    annotations = annotations[annotations["category"].isin(AV2_CLASS_MAP)]

    timestamps, image_paths = camera_timestamps(log_dir, camera)
    scene_id = f"{SOURCE}_{log_dir.name}"

    # Camera frame from ego frame: p_cam = R^T (p_ego - t).
    extrinsics = pd.read_feather(log_dir / "calibration" / "egovehicle_SE3_sensor.feather")
    pose = extrinsics[extrinsics["sensor_name"] == camera].iloc[0]
    R = calibration.R_ego_from_cam
    t = np.array([pose["tx_m"], pose["ty_m"], pose["tz_m"]], dtype=float)

    rows: list[dict] = []
    seen_images: set[str] = set()

    for timestamp_ns, sweep in annotations.groupby("timestamp_ns"):
        image_path = nearest_image(int(timestamp_ns), timestamps, image_paths)
        if image_path is None:
            dropped["no_image_in_window"] += len(sweep)
            continue

        relative = f"{SOURCE}/{image_path.relative_to(dataroot)}"
        seen_images.add(relative)

        for _, annotation in sweep.iterrows():
            class_name = AV2_CLASS_MAP[annotation["category"]]

            centre_ego = np.array(
                [annotation["tx_m"], annotation["ty_m"], annotation["tz_m"]], dtype=float
            )
            rotation_ego = quaternion_to_rotation_matrix(
                annotation["qw"], annotation["qx"], annotation["qy"], annotation["qz"]
            )
            length_m = float(annotation["length_m"])
            width_m = float(annotation["width_m"])
            height_m = float(annotation["height_m"])

            if class_name == "stop_sign":
                centre_ego, height_m = sign_face_only(
                    centre_ego, rotation_ego, width_m, height_m
                )

            corners_ego = cuboid_corners(
                centre_ego, rotation_ego, length_m, width_m, height_m
            )

            corners_cam = R.T @ (corners_ego - t.reshape(3, 1))
            if corners_cam[2].min() < MIN_CORNER_DEPTH_M:
                dropped["behind_camera"] += 1
                continue

            uv = calibration.project(corners_cam)
            raw_box = box_from_points(uv)
            clipped, truncation = clip_box(raw_box, calibration.width, calibration.height)
            if clipped[2] - clipped[0] <= 1.0 or clipped[3] - clipped[1] <= 1.0:
                dropped["degenerate"] += 1
                continue

            centre_cam = R.T @ (centre_ego - t)
            distance = float(centre_cam[2])
            if distance <= 0.0:
                dropped["non_positive_distance"] += 1
                continue

            centre_uv = calibration.project(centre_cam.reshape(3, 1))
            assert point_in_box((float(centre_uv[0, 0]), float(centre_uv[1, 0])), raw_box), (
                f"projected centre outside its own box in {log_dir.name} "
                f"-- the ego->camera transform is wrong"
            )

            rows.append(
                schema.object_row(
                    image_path=relative,
                    sensor_id=calibration.sensor_id,
                    source=SOURCE,
                    scene_id=scene_id,
                    class_name=class_name,
                    box=clipped,
                    gt_distance_m=distance,
                    gt_dims_hwl=[height_m, width_m, length_m],
                    # AV2 has no visibility field; num_interior_pts is a LiDAR
                    # count, not an image-space occlusion measure, so leave null.
                    visibility=None,
                    truncation=truncation,
                )
            )

    # Frames in this log with none of our classes become negatives.
    annotated = {row["image_path"] for row in rows}
    for path in image_paths:
        relative = f"{SOURCE}/{path.relative_to(dataroot)}"
        if relative not in annotated:
            rows.append(
                schema.negative_row(
                    image_path=relative,
                    sensor_id=calibration.sensor_id,
                    source=SOURCE,
                    scene_id=scene_id,
                )
            )
    return rows


def extract(dataroot: Path, cameras: list[str], root: Path, max_logs: int | None) -> None:
    # A log needs both annotations and images. The fetcher keeps annotations for
    # every log it scanned but downloads images only for the ones worth having,
    # so skip the metadata-only leftovers instead of walking them to find nothing.
    log_dirs = sorted(
        p for p in dataroot.iterdir()
        if (p / "annotations.feather").exists()
        and any((p / "sensors" / "cameras").glob(f"{cameras[0]}/*.jpg"))
    )
    if max_logs is not None:
        log_dirs = log_dirs[:max_logs]
    if not log_dirs:
        raise SystemExit(f"no AV2 logs with annotations.feather under {dataroot}")

    paths.link_source(root, SOURCE, dataroot)

    rows: list[dict] = []
    dropped: Counter = Counter()
    written_calibrations: set[str] = set()

    for index, log_dir in enumerate(log_dirs):
        for camera in cameras:
            calibration = read_calibration(log_dir, camera)
            calib.save(calibration, root)
            written_calibrations.add(calibration.sensor_id)
            rows.extend(rows_for_log(log_dir, camera, calibration, dataroot, dropped))
        print(f"  {index + 1}/{len(log_dirs)} logs  "
              f"cam_height_m={calibration.cam_height_m:.3f}")

    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)

    objects = schema.objects_only(frame)
    print(f"\nwrote {part}")
    print(f"  frames  : {frame['image_path'].nunique()}")
    print(f"  objects : {len(objects)}")
    print(f"  by class: {dict(objects['class'].value_counts())}")
    if dropped:
        print(f"  dropped : {dict(dropped)}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--dataroot", required=True, type=Path,
                        help="directory containing AV2 log folders")
    parser.add_argument("--cameras", nargs="+", default=["ring_front_center"])
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--max-logs", type=int, default=None)
    args = parser.parse_args()

    extract(
        dataroot=args.dataroot,
        cameras=args.cameras,
        root=paths.unified_root(args.unified_root),
        max_logs=args.max_logs,
    )


if __name__ == "__main__":
    main()