"""Concatenate the per-source manifests, filter, split by scene, assert sanity.

This is where the thresholds from schema.py are applied and where the
cross-dataset agreement checks live. The checks are real asserts, not things to
eyeball, because a coordinate-convention mistake between nuScenes and AV2
produces a manifest that trains YOLO perfectly well and makes the depth numbers
nonsense a day later.

Also writes class_priors.yaml -- per-class mean physical dimensions derived from
gt_dims_hwl. Phase 2's known-size estimator consumes it. Data-driven beats
looked-up, and it is one groupby.

Usage:
    python -m src.data.merge
    python -m src.data.merge --val-fraction 0.2
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
import yaml

from src.common import paths, schema

# The US MUTCD R1-1 stop sign is 0.75 m across the flats. Stop signs have no
# cuboid dimensions in COCO and are pole-mounted, so phase 2 uses width, not
# height, for them. This is the one prior we cannot derive from the data.
STOP_SIGN_WIDTH_M = 0.75


# ---------------------------------------------------------------------------
# Load and filter
# ---------------------------------------------------------------------------


def load_parts(root: Path) -> pd.DataFrame:
    part_paths = sorted(paths.parts_dir(root).glob("*.parquet"))
    if not part_paths:
        raise SystemExit(f"no per-source manifests under {paths.parts_dir(root)}; "
                         f"run an extractor first")

    frames = []
    for path in part_paths:
        frame = schema.read_manifest(path)
        print(f"  {path.name:28s} {len(frame):7d} rows")
        frames.append(frame)
    return pd.concat(frames, ignore_index=True)


def apply_filters(frame: pd.DataFrame) -> pd.DataFrame:
    """Drop objects we do not intend to detect. Thresholds live in schema.py."""
    objects = frame["class"].notna()
    height = frame["y2"] - frame["y1"]

    too_small = objects & (height < schema.MIN_BOX_HEIGHT_PX)
    too_occluded = objects & frame["visibility"].notna() & (
        frame["visibility"] < schema.MIN_VISIBILITY
    )
    too_truncated = objects & frame["truncation"].notna() & (
        frame["truncation"] > schema.MAX_TRUNCATION
    )

    print(f"\nfilters (thresholds in src/common/schema.py)")
    print(f"  box height < {schema.MIN_BOX_HEIGHT_PX:g} px : {int(too_small.sum())}")
    print(f"  visibility < {schema.MIN_VISIBILITY:g}       : {int(too_occluded.sum())}")
    print(f"  truncation > {schema.MAX_TRUNCATION:g}       : {int(too_truncated.sum())}")

    return frame[~(too_small | too_occluded | too_truncated)].copy()


def restore_frames_as_negatives(filtered: pd.DataFrame,
                                original: pd.DataFrame) -> pd.DataFrame:
    """Keep images that lost every object to the filters, as negatives.

    The alternative is dropping those images entirely. We keep them: everything
    the filters removed is below the size we intend to detect at all, so
    teaching the model not to fire there matches the deployment target. Use
    --drop-filtered-frames to choose the other behaviour.
    """
    surviving = set(filtered["image_path"])
    lost = original[~original["image_path"].isin(surviving)]
    if lost.empty:
        return filtered

    negatives = lost.drop_duplicates("image_path").apply(
        lambda row: schema.negative_row(
            image_path=row["image_path"],
            sensor_id=row["sensor_id"],
            source=row["source"],
            scene_id=row["scene_id"],
        ),
        axis=1,
    )
    print(f"  {len(negatives)} images lost every object -> kept as negatives")
    return pd.concat(
        [filtered, schema.rows_to_frame(list(negatives))], ignore_index=True
    )


# ---------------------------------------------------------------------------
# Split
# ---------------------------------------------------------------------------


def assign_split(frame: pd.DataFrame, val_fraction: float, seed: int) -> pd.DataFrame:
    """Split by scene, never by frame, and do it independently per source.

    Adjacent video frames are near-duplicates, so a frame-level random split
    leaks the val set into training and inflates val mAP a long way. Splitting
    per source also guarantees every source is represented in val -- otherwise
    a random scene draw can leave val with no nuScenes cones, which is the one
    thing phase 2 needs.

    Within a source, val is seeded with one scene per class before being filled
    to the frame target. With few scenes this can overshoot --val-fraction; that
    is the better failure, since a class with no val instances cannot be scored
    at all.
    """
    frame = frame.copy()
    frame["split"] = None
    rng = np.random.default_rng(seed)

    for source in sorted(frame["source"].unique()):
        mask = frame["source"] == source
        source_frame = frame[mask]
        frames_per_scene = source_frame.groupby("scene_id")["image_path"].nunique()
        scene_ids = frames_per_scene.index.to_numpy()
        rng.shuffle(scene_ids)

        # Which scenes contain each class. Instances cluster hard by scene -- in
        # AV2 nearly all cones live in a few logs -- so a purely random draw over
        # a handful of scenes can leave a class with zero val instances and no
        # way to score it. Seed val with one scene per class before filling.
        objects = source_frame[source_frame["class"].notna()]
        scenes_with = {name: set(group["scene_id"])
                       for name, group in objects.groupby("class")}

        target = val_fraction * frames_per_scene.sum()
        val_scenes: list[str] = []
        accumulated = 0

        for class_name in sorted(scenes_with):
            if any(scene in scenes_with[class_name] for scene in val_scenes):
                continue
            # scene_ids is already shuffled, so the first match is a random one.
            for scene_id in scene_ids:
                if scene_id in scenes_with[class_name]:
                    val_scenes.append(scene_id)
                    accumulated += int(frames_per_scene[scene_id])
                    break

        for scene_id in scene_ids:
            if accumulated >= target:
                break
            if scene_id not in val_scenes:
                val_scenes.append(scene_id)
                accumulated += int(frames_per_scene[scene_id])

        val_scenes = set(val_scenes)
        frame.loc[mask, "split"] = np.where(
            frame.loc[mask, "scene_id"].isin(val_scenes), "val", "train"
        )
        print(f"  {source:10s} {len(val_scenes):5d}/{len(scene_ids):5d} scenes -> val "
              f"({accumulated} frames)")
    return frame


# ---------------------------------------------------------------------------
# Sanity checks -- asserts, not eyeballing
# ---------------------------------------------------------------------------


def check(frame: pd.DataFrame) -> None:
    objects = schema.objects_only(frame)

    # 1. Boxes are non-degenerate.
    assert (objects["x2"] > objects["x1"]).all(), "boxes with x2 <= x1 survived merge"
    assert (objects["y2"] > objects["y1"]).all(), "boxes with y2 <= y1 survived merge"

    # 2. Distance is always positive where it exists. A negative distance means
    #    an object behind the camera was projected as if it were in front.
    with_distance = objects[objects["gt_distance_m"].notna()]
    assert (with_distance["gt_distance_m"] > 0).all(), (
        f"{int((with_distance['gt_distance_m'] <= 0).sum())} objects with "
        f"non-positive gt_distance_m"
    )

    # 3. Distance grows toward the horizon. Within one camera, an object whose
    #    box bottom is lower in the frame is nearer. A sign error or a swapped
    #    axis in the camera-frame transform flips or destroys this.
    ground = with_distance[with_distance["class"].isin(schema.GROUND_CONTACT_CLASSES)]
    for sensor_id, group in ground.groupby("sensor_id"):
        if len(group) < 50:
            continue
        correlation = float(np.corrcoef(group["y2"], group["gt_distance_m"])[0, 1])
        assert correlation < 0.0, (
            f"{sensor_id}: box bottom row and distance correlate at {correlation:+.2f}. "
            f"Objects lower in the frame should be NEARER. The camera-frame "
            f"transform is wrong."
        )
        if correlation > -0.3:
            print(f"  !! {sensor_id}: weak horizon correlation {correlation:+.2f}")
        else:
            print(f"  {sensor_id}: horizon correlation {correlation:+.2f}")

    # 4. Median cone distance is plausible. Getting this wrong by a factor of
    #    100 (millimetres, or a distance that is really an index) is the classic
    #    unit mistake and it survives every check above.
    cones = with_distance[with_distance["class"] == "cone"]
    if len(cones):
        median = float(cones["gt_distance_m"].median())
        low, high = schema.PLAUSIBLE_MEDIAN_CONE_DISTANCE_M
        assert low <= median <= high, (
            f"median cone distance is {median:.2f} m, expected {low}-{high} m. "
            f"Check the units and the camera-frame transform."
        )
        print(f"  median cone distance {median:.1f} m")

    # 5. The thing phase 2 actually runs on. Report every class, because a
    #    scene-level split on a small number of scenes can easily land zero
    #    instances of a class in val, and a silent zero here means phase 2 has
    #    nothing to score that class against.
    print("  val instances with gt_distance_m, per class:")
    val = with_distance[with_distance["split"] == "val"]
    train_classes = set(objects[objects["split"] == "train"]["class"])
    for class_name in sorted(train_classes):
        count = int((val["class"] == class_name).sum())
        flag = "  <-- ZERO, change --seed or --val-fraction" if count == 0 else ""
        print(f"    {class_name:10s} {count:6d}{flag}")
    if len(val[val["class"] == "cone"]) < 200:
        print("  !! fewer than 200 val cones -- phase 2 wants at least a few hundred")


# ---------------------------------------------------------------------------
# Class priors
# ---------------------------------------------------------------------------


def write_class_priors(frame: pd.DataFrame, root: Path) -> Path:
    """Per-class mean physical size, for phase 2's known-size estimator."""
    objects = schema.objects_only(frame)
    with_dims = objects[objects["gt_dims_hwl"].notna()]

    priors: dict[str, dict] = {}
    for class_name, group in with_dims.groupby("class"):
        dims = np.array(group["gt_dims_hwl"].tolist(), dtype=float)
        priors[class_name] = {
            "height_m": float(np.mean(dims[:, 0])),
            "width_m": float(np.mean(dims[:, 1])),
            "length_m": float(np.mean(dims[:, 2])),
            "height_std_m": float(np.std(dims[:, 0])),
            "n_samples": int(len(dims)),
            "source": "derived from manifest gt_dims_hwl",
        }

    if "stop_sign" not in priors:
        priors["stop_sign"] = {
            "width_m": STOP_SIGN_WIDTH_M,
            "source": "US MUTCD R1-1, across the flats (no cuboid dims available)",
        }

    path = Path(root) / "class_priors.yaml"
    with open(path, "w") as handle:
        yaml.safe_dump(priors, handle, sort_keys=True)
    print(f"\nwrote {path}")
    for class_name, prior in sorted(priors.items()):
        if "height_m" in prior:
            print(f"  {class_name:10s} height {prior['height_m']:.2f} m "
                  f"(n={prior['n_samples']})")
        else:
            print(f"  {class_name:10s} width  {prior['width_m']:.2f} m ({prior['source']})")
    return path


# ---------------------------------------------------------------------------


def merge(root: Path, val_fraction: float, seed: int, drop_filtered_frames: bool) -> None:
    print("loading per-source manifests")
    frame = load_parts(root)
    print(f"  {'total':28s} {len(frame):7d} rows")

    filtered = apply_filters(frame)
    if not drop_filtered_frames:
        filtered = restore_frames_as_negatives(filtered, frame)

    print("\nsplitting by scene")
    filtered = assign_split(filtered, val_fraction, seed)

    print("\nsanity checks")
    check(filtered)

    manifest_path = Path(root) / "manifest.parquet"
    schema.write_manifest(filtered, manifest_path)

    objects = schema.objects_only(filtered)
    print(f"\nwrote {manifest_path}")
    print(f"  frames : {filtered['image_path'].nunique()}")
    print(f"  objects: {len(objects)}")
    print("\ninstances per source x class")
    print(pd.crosstab(objects["source"], objects["class"]).to_string())
    print("\nframes per split")
    print(filtered.groupby("split")["image_path"].nunique().to_string())

    write_class_priors(filtered, root)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--val-fraction", type=float, default=0.15)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--drop-filtered-frames", action="store_true",
                        help="drop images that lost every object to the filters, "
                             "instead of keeping them as negatives")
    args = parser.parse_args()

    merge(
        root=paths.unified_root(args.unified_root),
        val_fraction=args.val_fraction,
        seed=args.seed,
        drop_filtered_frames=args.drop_filtered_frames,
    )


if __name__ == "__main__":
    main()