import argparse
from collections import defaultdict
from pathlib import Path
import random

import numpy as np

from tools.pdcl_inference.analyze_val_two_roi_badcases import (
    FACE_VISIBILITY_BUCKET_ORDER,
    add_interval_visual_record,
    build_grouped_metric_bucket_interval_rows,
    build_badcase_draw_kwargs,
    flatten_grouped_metric_bucket_interval_rows,
    is_signed_lateral_yaw_compare_longitudinal_eligible,
    interval_label,
    make_metric_bucket,
    make_interval_visual_reservoirs,
    make_record,
    select_interval_visual_records,
)
from tools.pdcl_inference.two_roi_inference import (
    DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    _project_selected_face_edges_from_box,
    _selected_edge_residual_stats,
    _build_edge_heading_decoded,
    _restore_depth_scale,
    add_two_roi_inference_args,
)
from ultralytics.utils.plotting_3d import (
    compute_3d_box_corners,
    decode_edge_yaw_selection_from_prediction,
    decode_3d_prediction,
    extract_3d_attrs_from_prediction,
    face_center_from_corners,
    project_3d_to_2d,
    project_face_bottom_edge,
    project_partial_face_bottom_edge,
    reconstruct_edge_based_box_from_selection,
    reconstruct_two_face_box_from_edge_selection,
)


def _encode_face_branch(pred, corners, face_type, score, anchor, stride, calib):
    face_center = face_center_from_corners(corners, face_type)
    face_uv = project_3d_to_2d(face_center[None, :], calib)[0]
    off = face_type * 6
    pred[off] = face_center[2]
    pred[off + 1 : off + 3] = [face_uv[0] / stride - anchor[0], face_uv[1] / stride - anchor[1]]
    pred[off + 5] = score


def _encode_edge_branch(pred_edge, face_type, points_3d, points_2d, anchor, stride):
    face_block = pred_edge[face_type * 15 : (face_type + 1) * 15].reshape(5, 3)
    face_block[:, 0] = points_2d[:, 0] / stride - anchor[0]
    face_block[:, 1] = points_2d[:, 1] / stride - anchor[1]
    face_block[:, 2] = points_3d[:, 2]


def test_signed_lateral_yaw_compare_longitudinal_eligibility_uses_lt_50m():
    assert is_signed_lateral_yaw_compare_longitudinal_eligible({"gt_z_m": 49.99}) is True
    assert is_signed_lateral_yaw_compare_longitudinal_eligible({"gt_z_m": 50.0}) is False
    assert is_signed_lateral_yaw_compare_longitudinal_eligible({"gt_z_m": 65.0}) is False
    assert is_signed_lateral_yaw_compare_longitudinal_eligible({"gt_z_m": -1.0}) is False
    assert is_signed_lateral_yaw_compare_longitudinal_eligible({"gt_z_m": None}) is False


def test_build_grouped_metric_bucket_interval_rows_splits_face_visibility_buckets():
    stores = {bucket_name: defaultdict(make_metric_bucket) for bucket_name in FACE_VISIBILITY_BUCKET_ORDER}

    front_bucket = stores["front_rear_only"][-5.0]
    front_bucket["matched_3d"] = 4
    front_bucket["yaw_compare_eligible_count"] = 4
    front_bucket["yaw_compare_pair_count"] = 2
    front_bucket["direct_visible_yaw_abs_deg"] = [4.0, 6.0]
    front_bucket["edge_visible_yaw_abs_deg"] = [5.0, 7.0]
    front_bucket["direct_minus_edge_visible_yaw_abs_deg"] = [-1.0, -1.0]

    two_face_bucket = stores["two-face"][0.0]
    two_face_bucket["matched_3d"] = 5
    two_face_bucket["yaw_compare_eligible_count"] = 5
    two_face_bucket["yaw_compare_pair_count"] = 3
    two_face_bucket["direct_visible_yaw_abs_deg"] = [2.0, 3.0, 4.0]
    two_face_bucket["edge_visible_yaw_abs_deg"] = [1.0, 2.0, 3.0]
    two_face_bucket["direct_minus_edge_visible_yaw_abs_deg"] = [1.0, 1.0, 1.0]

    rows_by_group = build_grouped_metric_bucket_interval_rows(
        stores,
        bin_width_m=5.0,
        prefix="lateral_bin",
        group_order=FACE_VISIBILITY_BUCKET_ORDER,
    )

    assert list(rows_by_group) == list(FACE_VISIBILITY_BUCKET_ORDER)
    assert [row["lateral_bin_label"] for row in rows_by_group["front_rear_only"]] == ["[-5,0)m"]
    assert rows_by_group["side only"] == []
    assert [row["lateral_bin_label"] for row in rows_by_group["two-face"]] == ["[0,5)m"]
    assert np.isclose(rows_by_group["two-face"][0]["direct_regression_yaw_mae_deg"], 3.0)
    assert np.isclose(rows_by_group["two-face"][0]["edge_based_yaw_mae_deg"], 2.0)

    flat_rows = flatten_grouped_metric_bucket_interval_rows(rows_by_group, group_key="face_visibility_bucket")
    assert [row["face_visibility_bucket"] for row in flat_rows] == ["front_rear_only", "two-face"]
    assert [row["lateral_bin_label"] for row in flat_rows] == ["[-5,0)m", "[0,5)m"]


def test_restore_depth_scale_scales_edge_depth_channels_too():
    preds_3d = np.ones((2, 41), dtype=np.float32)
    preds_edge = np.ones((2, 60), dtype=np.float32)

    restored_3d, restored_edge = _restore_depth_scale(preds_3d, preds_edge, 2.5)

    for channel in (0, 6, 12, 18, 24):
        assert np.allclose(restored_3d[:, channel], 2.5)
    untouched_3d = np.setdiff1d(np.arange(preds_3d.shape[1]), np.array([0, 6, 12, 18, 24]))
    assert np.allclose(restored_3d[:, untouched_3d], 1.0)

    assert restored_edge is not None
    assert np.allclose(restored_edge[:, 2::3], 2.5)
    assert np.allclose(restored_edge[:, 0::3], 1.0)
    assert np.allclose(restored_edge[:, 1::3], 1.0)

    assert np.allclose(preds_3d, 1.0)
    assert np.allclose(preds_edge, 1.0)


def test_build_edge_heading_decoded_anchors_cut_objects_on_primary_face():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([0.0, 120.0, 80.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_partial_face_bottom_edge(corners, 3, calib, 640, 480, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None
    _encode_edge_branch(pred_edge, 0, front_points_3d, front_points_2d, anchor, stride)
    _encode_edge_branch(pred_edge, 3, side_points_3d, side_points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    assert edge_confident
    assert np.isfinite(edge_yaw)
    assert heading_decoded is not None
    assert heading_decoded["visible_face_type"] == 0
    assert heading_decoded["visible_face_types"] == (0, 3)
    diff = (edge_yaw - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-3)
    assert np.allclose(
        face_center_from_corners(heading_decoded["corners_3d"], 0),
        face_center_from_corners(decoded["corners_3d"], 0),
        atol=1e-4,
    )


def test_build_edge_heading_decoded_falls_back_to_single_front_face_when_cut_side_edge_is_short():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([0.0, 120.0, 80.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, _ = project_partial_face_bottom_edge(corners, 3, calib, 640, 480, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None
    short_len = float(dims[0] * 0.4)
    side_dir = side_points_3d[-1] - side_points_3d[0]
    side_dir = side_dir / np.linalg.norm(side_dir)
    short_side_points_3d = side_points_3d[0] + np.linspace(0.0, short_len, 5, dtype=np.float32)[:, None] * side_dir[None, :]
    short_side_points_2d = project_3d_to_2d(short_side_points_3d, calib)
    assert np.all(np.isfinite(short_side_points_2d))
    _encode_edge_branch(pred_edge, 0, front_points_3d, front_points_2d, anchor, stride)
    _encode_edge_branch(pred_edge, 3, short_side_points_3d, short_side_points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert edge_selection["cut_side_visible_length_ratio"] is not None
    assert edge_selection["cut_side_visible_length_ratio"] < 0.5
    assert edge_selection["cut_side_visible_ratio_ok"] is False
    assert edge_selection["is_valid"] is False
    assert edge_selection["face_types"] == (0,)
    assert heading_decoded is not None
    assert edge_confident
    assert np.isfinite(edge_yaw)


def test_reconstruct_two_face_box_from_edge_selection_uses_selected_edge_lengths():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    edge_box = reconstruct_two_face_box_from_edge_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "edge_points_3d": np.stack([front_points_3d, side_points_3d], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_height_m=float(dims[1]),
    )

    assert edge_box is not None
    assert np.allclose(edge_box["center"], center, atol=1e-4)
    assert np.allclose(edge_box["dims"], dims, atol=1e-4)
    diff = (edge_box["yaw"] - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-5)


def test_reconstruct_edge_based_box_from_selection_trusts_longitudinal_label_and_ignores_side_label():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 2),
            "face_is_partial": (False, False),
            "edge_points_3d": np.stack([front_points_3d, side_points_3d], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
    )

    assert edge_box is not None
    assert edge_box["mode"] == "two-face"
    assert np.allclose(edge_box["center"], center, atol=1e-4)
    assert np.allclose(edge_box["dims"], dims, atol=1e-4)
    assert edge_box["face_types"] == (0, 3)


def test_reconstruct_edge_based_box_from_selection_uses_edge_y_as_vertical_anchor():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.3, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None
    noisy_front = front_points_3d.copy()
    noisy_side = side_points_3d.copy()
    noisy_front[:, 1] -= 0.9
    noisy_side[:, 1] -= 0.7

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "edge_points_3d": np.stack([noisy_front, noisy_side], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
    )

    assert edge_box is not None
    expected_center_y = float(np.mean(np.concatenate([noisy_front[:, 1], noisy_side[:, 1]], axis=0)) - float(dims[1]) * 0.5)
    assert np.isclose(float(edge_box["center"][1]), expected_center_y, atol=1e-5)
    assert np.allclose(edge_box["dims"], dims, atol=1e-4)


def test_reconstruct_edge_based_box_from_selection_ignores_vertical_noise_in_edge_length():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    noisy_side = side_points_3d.copy()
    noisy_side[:, 1] += np.linspace(0.0, 8.0, noisy_side.shape[0], dtype=np.float32)
    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "edge_points_3d": np.stack([front_points_3d, noisy_side], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
    )

    assert edge_box is not None
    assert np.isclose(float(edge_box["dims"][0]), float(dims[0]), atol=1e-4)
    assert np.isclose(float(edge_box["dims"][2]), float(dims[2]), atol=1e-4)


def test_reconstruct_edge_based_box_from_front_face_uses_regressed_length():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.0, 0.1, 20.0], dtype=np.float32)
    true_dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    regressed_dims = np.array([5.2, 1.7, 1.8], dtype=np.float32)
    rot_y = 0.45

    corners = compute_3d_box_corners(center, true_dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0,),
            "edge_points_3d": front_points_3d,
            "edge_points_2d": front_points_2d,
        },
        box_center_y_m=float(center[1]),
        regressed_dims=regressed_dims,
    )

    assert edge_box is not None
    assert edge_box["mode"] == "front-rear"
    assert np.isclose(float(edge_box["dims"][0]), float(regressed_dims[0]), atol=1e-5)
    assert np.isclose(float(edge_box["dims"][1]), float(regressed_dims[1]), atol=1e-5)
    assert np.isclose(float(edge_box["dims"][2]), float(true_dims[2]), atol=1e-4)
    expected_face_center = face_center_from_corners(corners, 0)
    assert expected_face_center is not None
    assert np.isclose(float(edge_box["center"][0]), float(expected_face_center[0]), atol=1e-4)
    assert np.isclose(float(edge_box["center"][2]), float(expected_face_center[2]), atol=1e-4)
    rebuilt_face_center = face_center_from_corners(edge_box["corners_3d"], 0)
    assert rebuilt_face_center is not None and expected_face_center is not None
    assert np.allclose(rebuilt_face_center, expected_face_center, atol=1e-4)


def test_reconstruct_edge_based_box_from_side_face_uses_regressed_width():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-1.2, -0.2, 18.0], dtype=np.float32)
    true_dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    regressed_dims = np.array([4.0, 1.6, 2.4], dtype=np.float32)
    rot_y = -0.55

    corners = compute_3d_box_corners(center, true_dims, rot_y, face_type=-1)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert side_points_3d is not None and side_points_2d is not None

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (3,),
            "edge_points_3d": side_points_3d,
            "edge_points_2d": side_points_2d,
        },
        box_center_y_m=float(center[1]),
        regressed_dims=regressed_dims,
    )

    assert edge_box is not None
    assert edge_box["mode"] == "side"
    assert np.isclose(float(edge_box["dims"][0]), float(true_dims[0]), atol=1e-4)
    assert np.isclose(float(edge_box["dims"][1]), float(regressed_dims[1]), atol=1e-5)
    assert np.isclose(float(edge_box["dims"][2]), float(regressed_dims[2]), atol=1e-5)
    expected_face_center = face_center_from_corners(corners, 3)
    assert expected_face_center is not None
    assert np.isclose(float(edge_box["center"][0]), float(expected_face_center[0]), atol=1e-4)
    assert np.isclose(float(edge_box["center"][2]), float(expected_face_center[2]), atol=1e-4)
    rebuilt_face_center = face_center_from_corners(edge_box["corners_3d"], 3)
    assert rebuilt_face_center is not None and expected_face_center is not None
    assert np.allclose(rebuilt_face_center, expected_face_center, atol=1e-4)


def test_decode_edge_yaw_selection_single_side_face_uses_cut_cls_for_yaw_disambiguation():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-0.5, 0.0, 18.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.55
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([80.0, 120.0, 220.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=3, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert side_points_3d is not None and side_points_2d is not None
    _encode_edge_branch(pred_edge, 3, side_points_3d, side_points_2d, anchor, stride)

    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        max_faces=1,
    )

    assert edge_selection["face_types"] == (3,)
    diff = (float(edge_selection["yaw"]) - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-4)


def test_decode_edge_yaw_selection_single_side_face_left_cut_in_keeps_negative_yaw():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([0.0, 120.0, 80.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=3, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    side_points_3d, side_points_2d = project_partial_face_bottom_edge(corners, 3, calib, 640, 480, num_samples=5)
    assert side_points_3d is not None and side_points_2d is not None
    _encode_edge_branch(pred_edge, 3, side_points_3d, side_points_2d, anchor, stride)

    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        max_faces=1,
    )

    assert edge_selection["face_types"] == (3,)
    diff = (float(edge_selection["yaw"]) - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-4)


def test_decode_edge_yaw_selection_single_rear_face_cut_in_uses_reference_yaw_to_avoid_pi_flip():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-0.5, 0.0, 18.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([0.0, 120.0, 80.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=1, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    rear_points_3d, rear_points_2d = project_face_bottom_edge(corners, 1, calib, num_samples=5)
    assert rear_points_3d is not None and rear_points_2d is not None
    _encode_edge_branch(pred_edge, 1, rear_points_3d, rear_points_2d, anchor, stride)

    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        max_faces=1,
    )

    assert edge_selection["face_types"] == (1,)
    diff = (float(edge_selection["yaw"]) - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-4)


def test_reconstruct_edge_based_box_from_selection_falls_back_when_edge_metric_size_is_unstable():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.5, 0.0, 30.0], dtype=np.float32)
    dims = np.array([4.5, 1.5, 2.0], dtype=np.float32)
    rot_y = 0.6

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    inflated_front = front_points_3d.copy()
    inflated_side = side_points_3d.copy()
    front_mid = np.mean(inflated_front[:, [0, 2]], axis=0, keepdims=True)
    side_mid = np.mean(inflated_side[:, [0, 2]], axis=0, keepdims=True)
    inflated_front[:, [0, 2]] = front_mid + (inflated_front[:, [0, 2]] - front_mid) * 2.8
    inflated_side[:, [0, 2]] = side_mid + (inflated_side[:, [0, 2]] - side_mid) * 2.8

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "face_is_partial": (False, False),
            "edge_points_3d": np.stack([inflated_front, inflated_side], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
        face_regressed_dims_by_type={
            0: {"height": float(dims[1]), "width": float(dims[2])},
            3: {"length": float(dims[0]), "height": float(dims[1])},
        },
    )

    assert edge_box is not None
    assert np.allclose(edge_box["dims"], dims, atol=1e-4)
    assert edge_box["length_source"] == "regressed"
    assert edge_box["width_source"] == "regressed"


def test_reconstruct_edge_based_box_from_selection_does_not_use_partial_edge_length_as_full_box_size():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    partial_side = side_points_3d.copy()
    full_len = np.linalg.norm(partial_side[-1, [0, 2]] - partial_side[0, [0, 2]])
    partial_dir = partial_side[-1] - partial_side[0]
    partial_dir = partial_dir / np.linalg.norm(partial_dir)
    partial_side = partial_side[0] + np.linspace(0.0, float(full_len) * 0.875, 5, dtype=np.float32)[:, None] * partial_dir[None, :]

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "face_is_partial": (False, True),
            "edge_points_3d": np.stack([front_points_3d, partial_side], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
    )

    assert edge_box is not None
    assert np.isclose(float(edge_box["dims"][0]), float(dims[0]), atol=1e-4)
    assert np.isclose(float(edge_box["dims"][2]), float(dims[2]), atol=1e-4)
    assert edge_box["length_source"] == "regressed"
    assert edge_box["width_source"] == "edge"


def test_reconstruct_edge_based_box_from_selection_ignores_face_size_priors_that_disagree_with_whole_box():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.5, 0.0, 30.0], dtype=np.float32)
    dims = np.array([4.5, 1.5, 2.0], dtype=np.float32)
    rot_y = 0.6

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    shrunk_front = front_points_3d.copy()
    shrunk_side = side_points_3d.copy()
    front_mid = np.mean(shrunk_front[:, [0, 2]], axis=0, keepdims=True)
    side_mid = np.mean(shrunk_side[:, [0, 2]], axis=0, keepdims=True)
    shrunk_front[:, [0, 2]] = front_mid + (shrunk_front[:, [0, 2]] - front_mid) * 0.7
    shrunk_side[:, [0, 2]] = side_mid + (shrunk_side[:, [0, 2]] - side_mid) * 0.7

    edge_box = reconstruct_edge_based_box_from_selection(
        {
            "yaw": float(rot_y),
            "face_types": (0, 3),
            "face_is_partial": (False, False),
            "edge_points_3d": np.stack([shrunk_front, shrunk_side], axis=0),
            "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        },
        box_center_y_m=float(center[1]),
        regressed_dims=dims,
        face_regressed_dims_by_type={
            0: {"height": float(dims[1]), "width": float(dims[2] * 0.7)},
            3: {"length": float(dims[0] * 0.7), "height": float(dims[1])},
        },
    )

    assert edge_box is not None
    assert np.allclose(edge_box["dims"], dims, atol=1e-4)
    assert edge_box["length_source"] == "regressed"
    assert edge_box["width_source"] == "regressed"


def test_build_edge_heading_decoded_uses_edge_geometry_for_box_size():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.0, 20.0], dtype=np.float32)
    true_dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    regressed_dims = np.array([6.5, 1.5, 2.6], dtype=np.float32)
    rot_y = 0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([120.0, 120.0, 240.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = regressed_dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)

    corners = compute_3d_box_corners(center, true_dims, rot_y, face_type=-1)
    for face_type, score in ((0, 0.95), (3, 0.9)):
        _encode_face_branch(pred, corners, face_type=face_type, score=score, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    for face_type in (0, 3):
        points_3d, points_2d = project_face_bottom_edge(corners, face_type, calib, num_samples=5)
        assert points_3d is not None and points_2d is not None
        _encode_edge_branch(pred_edge, face_type, points_3d, points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert edge_confident
    assert heading_decoded is not None
    diff = (edge_yaw - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-3)
    front_edge_3d, _ = project_face_bottom_edge(heading_decoded["corners_3d"], 0, calib, num_samples=5)
    side_edge_3d, _ = project_face_bottom_edge(heading_decoded["corners_3d"], 3, calib, num_samples=5)
    assert front_edge_3d is not None and side_edge_3d is not None
    rebuilt_width = float(np.linalg.norm(front_edge_3d[-1] - front_edge_3d[0]))
    rebuilt_length = float(np.linalg.norm(side_edge_3d[-1] - side_edge_3d[0]))
    assert np.isclose(rebuilt_length, float(true_dims[0]), atol=1e-4)
    assert np.isclose(rebuilt_width, float(true_dims[2]), atol=1e-4)
    assert not np.isclose(rebuilt_length, float(regressed_dims[0]), atol=1e-3)
    assert not np.isclose(rebuilt_width, float(regressed_dims[2]), atol=1e-3)


def test_build_edge_heading_decoded_keeps_same_best_score_face_anchor_as_direct_box():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([120.0, 120.0, 240.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    for face_type, score in ((0, 0.95), (1, 0.99), (3, 0.80)):
        _encode_face_branch(pred, corners, face_type=face_type, score=score, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    for face_type in (0, 1, 3):
        points_3d, points_2d = project_face_bottom_edge(corners, face_type, calib, num_samples=5)
        assert points_3d is not None and points_2d is not None
        _encode_edge_branch(pred_edge, face_type, points_3d, points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert decoded is not None
    assert heading_decoded is not None
    assert edge_confident
    assert edge_selection["face_types"] == (1, 3)
    # Direct-yaw and edge-yaw reconstruction should share the same score-chosen face anchor; only yaw changes.
    assert heading_decoded["visible_face_type"] == decoded["visible_face_type"] == 1
    diff = (edge_yaw - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-5), edge_yaw


def test_build_edge_heading_decoded_keeps_primary_face_when_edge_points_leak_outside_image():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-12.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.9
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert float(np.min(front_points_2d[:, 0])) < 0.0
    _encode_edge_branch(pred_edge, 0, front_points_3d, front_points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=None,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert decoded is not None
    assert edge_selection["face_types"] == (0,)
    assert heading_decoded is not None
    assert edge_confident
    diff = (edge_yaw - rot_y + np.pi) % (2 * np.pi) - np.pi
    assert np.isclose(diff, 0.0, atol=1e-4)


def test_selected_edge_residual_stats_are_zero_for_matching_reprojected_edges():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    yaw = 0.75
    corners = compute_3d_box_corners(center, dims, yaw, face_type=-1)

    edge_points_3d, edge_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    assert edge_points_3d is not None and edge_points_2d is not None

    projected = _project_selected_face_edges_from_box(
        corners,
        face_types=(0,),
        face_is_partial=(False,),
        calib=calib,
        img_w=640,
        img_h=480,
    )
    stats = _selected_edge_residual_stats(edge_points_2d, projected)

    assert stats["available"] is True
    assert np.isclose(stats["mean_px"], 0.0, atol=1e-6)
    assert np.isclose(stats["max_px"], 0.0, atol=1e-6)


def test_build_edge_heading_decoded_rejects_far_lateral_edge_yaw():
    threshold_m = 5.0
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([6.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([200.0, 120.0, 320.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.9, anchor=anchor, stride=stride, calib=calib)
    _encode_face_branch(pred, corners, face_type=3, score=0.85, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    for face_type in (0, 3):
        points_3d, points_2d = project_face_bottom_edge(corners, face_type, calib, num_samples=5)
        assert points_3d is not None and points_2d is not None
        _encode_edge_branch(pred_edge, face_type, points_3d, points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=threshold_m,
    )

    assert whole_attrs is not None
    assert float(abs(whole_attrs["center"][0])) > threshold_m
    assert np.isfinite(edge_yaw)
    assert heading_decoded is None
    assert not edge_confident


def test_build_edge_heading_decoded_reconstructs_single_front_face_from_edge_and_regressed_length():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-2.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.35
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([160.0, 120.0, 260.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.9, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    points_3d, points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    assert points_3d is not None and points_2d is not None
    _encode_edge_branch(pred_edge, 0, points_3d, points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    assert decoded is not None
    assert decoded["visible_face_types"] == (0,)
    assert heading_decoded is not None
    assert np.isfinite(edge_yaw)
    assert edge_confident
    rebuilt_front_edge_3d, _ = project_face_bottom_edge(heading_decoded["corners_3d"], 0, calib, num_samples=5)
    assert rebuilt_front_edge_3d is not None
    rebuilt_width = float(np.linalg.norm(rebuilt_front_edge_3d[-1] - rebuilt_front_edge_3d[0]))
    assert np.isclose(rebuilt_width, float(dims[2]), atol=1e-4)


def test_build_edge_heading_decoded_reconstructs_cut_object_from_single_front_face_when_side_edge_invalid():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([0.0, 120.0, 80.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.95, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, _ = project_partial_face_bottom_edge(corners, 3, calib, 640, 480, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None
    short_len = float(dims[0] * 0.4)
    side_dir = side_points_3d[-1] - side_points_3d[0]
    side_dir = side_dir / np.linalg.norm(side_dir)
    short_side_points_3d = side_points_3d[0] + np.linspace(0.0, short_len, 5, dtype=np.float32)[:, None] * side_dir[None, :]
    short_side_points_2d = project_3d_to_2d(short_side_points_3d, calib)
    _encode_edge_branch(pred_edge, 0, front_points_3d, front_points_2d, anchor, stride)
    _encode_edge_branch(pred_edge, 3, short_side_points_3d, short_side_points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)
    edge_selection = decode_edge_yaw_selection_from_prediction(
        pred,
        pred_edge,
        anchor,
        stride,
        calib,
        bbox_xyxy=bbox_xyxy,
        img_w=640,
        img_h=480,
        max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert edge_selection["is_valid"] is False
    assert edge_selection["face_types"] == (0,)
    assert heading_decoded is not None
    assert np.isfinite(edge_yaw)
    assert edge_confident


def test_build_edge_heading_decoded_limits_near_cut_boxes_to_two_edge_yaw_faces():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([-4.0, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = -0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([140.0, 120.0, 260.0, 240.0], dtype=np.float32)

    pred = np.zeros(41, dtype=np.float32)
    pred[24] = center[2]
    pred[25:27] = [center[0] / center[2] * calib["fx"] / stride + (calib["cx"] / stride - anchor[0]), 0.0]
    pred[27:30] = dims
    pred[30] = 10.0
    pred[34] = np.sin(rot_y)
    pred[39] = 10.0  # cut_in-like, but not actually clipped at the border

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    _encode_face_branch(pred, corners, face_type=0, score=0.82, anchor=anchor, stride=stride, calib=calib)
    _encode_face_branch(pred, corners, face_type=1, score=0.95, anchor=anchor, stride=stride, calib=calib)
    _encode_face_branch(pred, corners, face_type=3, score=0.9, anchor=anchor, stride=stride, calib=calib)

    pred_edge = np.zeros(60, dtype=np.float32)
    for face_type in (0, 1, 3):
        points_3d, points_2d = project_face_bottom_edge(corners, face_type, calib, num_samples=5)
        assert points_3d is not None and points_2d is not None
        _encode_edge_branch(pred_edge, face_type, points_3d, points_2d, anchor, stride)

    decoded = decode_3d_prediction(
        pred,
        anchor,
        stride,
        calib,
        640,
        480,
        {0},
        set(),
        0,
        pred_edge_60=pred_edge,
        bbox_xyxy=bbox_xyxy,
    )
    whole_attrs = extract_3d_attrs_from_prediction(pred, anchor, stride, calib, pred_edge_60=pred_edge)

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded=decoded,
        pred_41=pred,
        pred_edge_60=pred_edge,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        whole_attrs=whole_attrs,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
    )

    assert decoded is not None
    assert decoded["visible_face_types"] == (0, 1, 3)
    assert edge_confident
    assert np.isfinite(edge_yaw)
    assert heading_decoded is not None
    assert heading_decoded["visible_face_types"] == (0, 3)
    assert np.asarray(heading_decoded["edge_points_2d"], dtype=np.float32).shape == (2, 5, 2)


def test_build_edge_heading_decoded_uses_resolved_side_face_type_for_edge_fit():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([1.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    rot_y = 0.75
    anchor = np.array([10.0, 20.0], dtype=np.float32)
    stride = 8.0
    bbox_xyxy = np.array([120.0, 120.0, 240.0, 240.0], dtype=np.float32)

    corners = compute_3d_box_corners(center, dims, rot_y, face_type=-1)
    front_points_3d, front_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    side_points_3d, side_points_2d = project_face_bottom_edge(corners, 3, calib, num_samples=5)
    assert front_points_3d is not None and front_points_2d is not None
    assert side_points_3d is not None and side_points_2d is not None

    pred = np.zeros(41, dtype=np.float32)
    pred[27:30] = dims
    edge_selection = {
        "yaw": float(rot_y),
        "face_types": (0, 2),
        "face_is_partial": (False, False),
        "edge_points_3d": np.stack([front_points_3d, side_points_3d], axis=0),
        "edge_points_2d": np.stack([front_points_2d, side_points_2d], axis=0),
        "lateral_ok": True,
    }

    heading_decoded, edge_yaw, edge_confident = _build_edge_heading_decoded(
        base_decoded={},
        pred_41=pred,
        pred_edge_60=None,
        anchor_xy=anchor,
        stride=stride,
        bbox_xyxy=bbox_xyxy,
        calib=calib,
        img_w=640,
        img_h=480,
        whole_attrs=None,
        edge_yaw_max_lateral_dist_m=DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M,
        edge_selection=edge_selection,
    )

    assert edge_confident
    assert np.isclose(edge_yaw, rot_y, atol=1e-5)
    assert heading_decoded is not None
    assert heading_decoded["visible_face_types"] == (0, 3)
    edge_box_selected_edges_2d = _project_selected_face_edges_from_box(
        heading_decoded["corners_3d"],
        heading_decoded["visible_face_types"],
        edge_selection["face_is_partial"],
        calib,
        640,
        480,
    )
    edge_fit = _selected_edge_residual_stats(edge_selection["edge_points_2d"], edge_box_selected_edges_2d)
    assert edge_fit["available"] is True
    assert edge_fit["max_px"] is not None and edge_fit["max_px"] < 1e-3


def test_add_two_roi_inference_args_includes_edge_yaw_lateral_threshold():
    parser = argparse.ArgumentParser()
    add_two_roi_inference_args(parser, include_output_dir=False)
    args = parser.parse_args([])
    assert args.edge_yaw_max_lateral_dist == DEFAULT_EDGE_YAW_MAX_LATERAL_DIST_M


def test_make_record_preserves_badcase_geometry_shapes_for_visualization():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.5, 0.0, 20.0], dtype=np.float32)
    pred_center = center + np.array([0.3, 0.1, 0.2], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    yaw = 0.2
    pred_yaw = 0.35
    corners = compute_3d_box_corners(center, dims, yaw, face_type=-1).astype(np.float32)
    pred_corners = compute_3d_box_corners(pred_center, dims, pred_yaw, face_type=-1).astype(np.float32)
    face_center = face_center_from_corners(corners, 0).astype(np.float32)
    pred_face_center = face_center_from_corners(pred_corners, 0).astype(np.float32)
    edge_points_3d, edge_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)

    assert edge_points_3d is not None
    assert edge_points_2d is not None

    gt_attrs = {
        "center": center,
        "depth": float(center[2]),
        "dims": dims,
        "yaw": yaw,
        "visible_face_type": 0,
        "face_center": face_center,
        "corners_3d": corners,
    }
    pred_attrs = {
        "center": pred_center,
        "depth": float(pred_center[2]),
        "dims": dims,
        "yaw": pred_yaw,
        "visible_face_type": 0,
        "face_center": pred_face_center,
        "corners_3d": pred_corners,
    }
    gt_decoded = {
        "face_center_2d": project_3d_to_2d(face_center[None, :], calib)[0],
        "face_color": (0, 0, 255),
        "visible_face_types": (0,),
        "edge_points_2d": edge_points_2d,
    }
    pred_decoded = {
        "face_center_2d": project_3d_to_2d(pred_face_center[None, :], calib)[0],
        "face_color": (0, 255, 0),
        "visible_face_types": (0,),
        "edge_points_2d": edge_points_2d,
    }

    record = make_record(
        sample_index=0,
        roi_name="roi0",
        image_path=Path("frame.png"),
        label_path=Path("frame.txt"),
        cls_name="car",
        gt_index=0,
        pred_index=0,
        gt_box=np.array([100.0, 120.0, 220.0, 260.0], dtype=np.float32),
        pred_box=np.array([102.0, 121.0, 222.0, 261.0], dtype=np.float32),
        match_iou=0.9,
        prediction={"cls_id": 0, "confidence": 0.95},
        gt_attrs=gt_attrs,
        pred_attrs=pred_attrs,
        gt_decoded=gt_decoded,
        pred_decoded=pred_decoded,
        gt_visible_faces=[(0, np.zeros(8, dtype=np.float32))],
        gt_visible_yaw=yaw,
        pred_edge_visible_yaw=pred_yaw,
        is_cut_object_flag=False,
        position_eligible=True,
        yaw_compare_max_lateral_dist_m=5.0,
    )

    assert np.asarray(record["gt_corners"], dtype=np.float32).shape == (8, 3)
    assert np.asarray(record["pred_corners"], dtype=np.float32).shape == (8, 3)
    assert np.asarray(record["gt_edge_points_2d"], dtype=np.float32).shape == (5, 2)
    assert np.asarray(record["pred_edge_points_2d"], dtype=np.float32).shape == (5, 2)

    draw_kwargs = build_badcase_draw_kwargs(record["gt_corners"], calib, visible_face_type=0, visible_face_types=[0])
    assert "edge_points_2d" in draw_kwargs


def test_make_record_invalidates_side_edge_length_for_front_rear_only_edge_box():
    center = np.array([0.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    attrs = {
        "center": center,
        "depth": float(center[2]),
        "dims": dims,
        "yaw": 0.0,
        "visible_face_type": 0,
        "face_center": center,
        "corners_3d": np.zeros((8, 3), dtype=np.float32),
    }
    decoded = {"visible_face_types": (0,), "face_center_2d": None, "face_color": None, "edge_points_2d": None}

    record = make_record(
        sample_index=0,
        roi_name="roi1",
        image_path=Path("frame.png"),
        label_path=Path("frame.txt"),
        cls_name="car",
        gt_index=0,
        pred_index=0,
        gt_box=np.array([0.0, 0.0, 10.0, 10.0], dtype=np.float32),
        pred_box=np.array([0.0, 0.0, 10.0, 10.0], dtype=np.float32),
        match_iou=1.0,
        prediction={"cls_id": 0, "confidence": 0.95},
        gt_attrs=attrs,
        pred_attrs=attrs,
        gt_decoded=decoded,
        pred_decoded=decoded,
        gt_visible_faces=[(0, np.zeros(8, dtype=np.float32))],
        gt_visible_yaw=0.0,
        pred_edge_visible_yaw=0.0,
        pred_edge_bucket_visible_yaw=0.0,
        pred_edge_box={"mode": "front-rear", "dims": np.array([9.9, 1.5, 1.8], dtype=np.float32)},
        is_cut_object_flag=False,
        position_eligible=True,
        yaw_compare_max_lateral_dist_m=5.0,
        pred_yaw_compare_face_types=(0,),
        pred_yaw_compare_valid=False,
    )

    assert record["pred_edge_box_mode"] == "front-rear"
    assert record["pred_edge_length_m"] is None
    assert record["edge_length_abs_err_m"] is None


def test_make_record_tracks_strict_yaw_compare_face_bucket_metadata_separately():
    calib = {"fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 240.0}
    center = np.array([0.5, 0.0, 20.0], dtype=np.float32)
    dims = np.array([4.0, 1.5, 1.8], dtype=np.float32)
    yaw = 0.2
    corners = compute_3d_box_corners(center, dims, yaw, face_type=-1).astype(np.float32)
    face_center = face_center_from_corners(corners, 0).astype(np.float32)
    side_center = face_center_from_corners(corners, 3).astype(np.float32)
    edge_points_3d, edge_points_2d = project_face_bottom_edge(corners, 0, calib, num_samples=5)
    assert edge_points_3d is not None and edge_points_2d is not None

    attrs = {
        "center": center,
        "depth": float(center[2]),
        "dims": dims,
        "yaw": yaw,
        "visible_face_type": 0,
        "face_center": face_center,
        "corners_3d": corners,
    }
    decoded = {
        "face_center_2d": project_3d_to_2d(face_center[None, :], calib)[0],
        "face_color": (0, 0, 255),
        "visible_face_types": (0, 3),
        "edge_points_2d": edge_points_2d,
    }

    record = make_record(
        sample_index=0,
        roi_name="roi0",
        image_path=Path("frame.png"),
        label_path=Path("frame.txt"),
        cls_name="car",
        gt_index=0,
        pred_index=0,
        gt_box=np.array([100.0, 120.0, 220.0, 260.0], dtype=np.float32),
        pred_box=np.array([100.0, 120.0, 220.0, 260.0], dtype=np.float32),
        match_iou=1.0,
        prediction={"cls_id": 0, "confidence": 0.95},
        gt_attrs=attrs,
        pred_attrs=attrs,
        gt_decoded=decoded,
        pred_decoded=decoded,
        gt_visible_faces=[(0, np.zeros(8, dtype=np.float32)), (3, np.zeros(8, dtype=np.float32))],
        gt_visible_yaw=yaw,
        pred_edge_visible_yaw=yaw,
        is_cut_object_flag=False,
        position_eligible=True,
        yaw_compare_max_lateral_dist_m=5.0,
        pred_yaw_compare_face_types=(0,),
        pred_yaw_compare_valid=False,
    )

    assert record["visible_face_count"] == 2
    assert record["has_side_face_visible"] is True
    assert record["yaw_compare_visible_face_count"] == 1
    assert record["yaw_compare_has_side_face_visible"] is False
    assert record["yaw_compare_visible_faces"] == "front"
    assert record["yaw_compare_face_bucket"] == "front_rear_only"
    assert record["yaw_compare_pred_valid"] is False


def test_interval_label_supports_metric_units():
    assert interval_label(5.0, 1.0, unit="deg") == "[5,6)deg"
    assert interval_label(0.5, 0.5, unit="m") == "[0.5,1.0)m"


def test_select_interval_visual_records_round_robins_high_error_bins_first():
    stores = make_interval_visual_reservoirs()
    add_rng = random.Random(0)
    for value in (5.2, 5.4, 6.2, 6.4, 7.2, 7.4):
        add_interval_visual_record(
            stores,
            value=value,
            record={"value": value},
            bin_width=1.0,
            per_bin_capacity=10,
            rng=add_rng,
        )

    selected = select_interval_visual_records(stores, max_total=5, rng=random.Random(1), reverse=True)

    assert list(sorted(selected.keys(), reverse=True)) == [7.0, 6.0, 5.0]
    assert sum(len(records) for records in selected.values()) == 5
    assert len(selected[7.0]) == 2
    assert len(selected[6.0]) == 2
    assert len(selected[5.0]) == 1