from __future__ import annotations

from typing import Any, Optional

import numpy as np

try:
    from .two_roi_3d_utils import (
        CUT_STATE_IN,
        CUT_STATE_NORMAL,
        CUT_STATE_OUT,
        FACE_CORNERS,
        FACE_OFFSETS_41,
        FACE_VISIBILITY_SCORE_THRESH,
        _decode_yaw_from_prediction,
        _decoded_edge_to_points_3d,
        _edge_batches_to_list,
        _point_inside_image,
        _resolve_pred_cut_state_for_decode,
        _stack_edge_batches,
        back_project_2d_to_3d,
        compute_3d_box_corners,
        decode_visible_face_edge_from_prediction,
        get_pred_cut_primary_face,
        get_pred_cut_state,
        reconstruct_3d_box_from_face,
        reconstruct_3d_box_from_whole,
        select_pred_visible_faces,
        select_pred_visible_faces_for_decode,
    )
    from .two_roi_types import DecodedVisibleEdge, ResizedCalib
except ImportError:
    from two_roi_3d_utils import (
        CUT_STATE_IN,
        CUT_STATE_NORMAL,
        CUT_STATE_OUT,
        FACE_CORNERS,
        FACE_OFFSETS_41,
        FACE_VISIBILITY_SCORE_THRESH,
        _decode_yaw_from_prediction,
        _decoded_edge_to_points_3d,
        _edge_batches_to_list,
        _point_inside_image,
        _resolve_pred_cut_state_for_decode,
        _stack_edge_batches,
        back_project_2d_to_3d,
        compute_3d_box_corners,
        decode_visible_face_edge_from_prediction,
        get_pred_cut_primary_face,
        get_pred_cut_state,
        reconstruct_3d_box_from_face,
        reconstruct_3d_box_from_whole,
        select_pred_visible_faces,
        select_pred_visible_faces_for_decode,
    )
    from two_roi_types import DecodedVisibleEdge, ResizedCalib


EDGE_YAW_VALID_VISIBILITY_SCORE_THRESH = 0.1
EDGE_YAW_CUT_SIDE_MIN_VISIBLE_LENGTH_RATIO = 0.5
EDGE_YAW_MAX_LATERAL_DIST_M = 30.0


def edge_points_to_yaw(points_3d: Any, face_type: int) -> float:
    if points_3d is None or len(points_3d) < 2 or face_type not in range(4):
        return float("nan")

    pts = np.asarray(points_3d, dtype=np.float64)
    valid = np.isfinite(pts).all(axis=1)
    pts = pts[valid]
    if len(pts) < 2:
        return float("nan")

    tangent = np.array([pts[-1, 0] - pts[0, 0], pts[-1, 2] - pts[0, 2]], dtype=np.float64)
    tangent_norm = float(np.linalg.norm(tangent))
    if tangent_norm < 1e-8:
        return float("nan")
    tangent /= tangent_norm

    midpoint = np.mean(pts[:, [0, 2]], axis=0)

    def _rot_cw(v: np.ndarray) -> np.ndarray:
        return np.array([v[1], -v[0]], dtype=np.float64)

    def _rot_ccw(v: np.ndarray) -> np.ndarray:
        return np.array([-v[1], v[0]], dtype=np.float64)

    if face_type in (0, 1):
        forward_candidates = (_rot_cw(tangent), -_rot_cw(tangent))
    else:
        forward_candidates = (tangent, -tangent)

    def _face_normal(forward: np.ndarray) -> np.ndarray:
        if face_type == 0:
            return forward
        if face_type == 1:
            return -forward
        if face_type == 2:
            return _rot_ccw(forward)
        return -_rot_ccw(forward)

    best_forward = min(forward_candidates, key=lambda forward: float(np.dot(_face_normal(forward), midpoint)))
    yaw = np.arctan2(-best_forward[1], best_forward[0])
    return float((yaw + np.pi) % (2 * np.pi) - np.pi)


def _decoded_edge_points_are_drawable(
    points_2d: Optional[np.ndarray],
    img_w: Optional[float] = None,
    img_h: Optional[float] = None,
    min_endpoint_dist_px: float = 2.0,
) -> bool:
    if points_2d is None:
        return False

    pts = np.asarray(points_2d, dtype=np.float32)
    if pts.ndim != 2 or pts.shape[0] < 2 or pts.shape[1] != 2 or not np.isfinite(pts).all():
        return False
    if img_w is not None and img_h is not None:
        if not np.all([_point_inside_image(point_2d, img_w, img_h) for point_2d in pts]):
            return False
    return float(np.linalg.norm(pts[-1] - pts[0])) >= float(min_endpoint_dist_px)


def _edge_segment_length_3d(points_3d: Optional[np.ndarray]) -> Optional[float]:
    if points_3d is None:
        return None
    pts = np.asarray(points_3d, dtype=np.float32)
    if pts.ndim != 2 or pts.shape[0] < 2 or pts.shape[1] != 3 or not np.isfinite(pts).all():
        return None
    return float(np.linalg.norm(pts[-1, [0, 2]] - pts[0, [0, 2]]))


def _prediction_lateral_distance_m_from_center(center: Optional[np.ndarray]) -> Optional[float]:
    if center is None:
        return None
    center_arr = np.asarray(center, dtype=np.float32).reshape(-1)
    if center_arr.shape[0] < 1 or not np.isfinite(center_arr[0]):
        return None
    return float(abs(center_arr[0]))


def _get_camera_facing_side_face_from_corners(corners_3d: Optional[np.ndarray]) -> Optional[int]:
    if corners_3d is None:
        return None
    corners = np.asarray(corners_3d, dtype=np.float64)
    if corners.shape != (8, 3) or not np.isfinite(corners).all():
        return None

    box_center = corners.mean(axis=0)
    best_face_type, best_score = None, -np.inf
    for face_type in (2, 3):
        face_points = corners[list(FACE_CORNERS[face_type])]
        face_center = face_points.mean(axis=0)
        view_dir = -face_center
        view_norm = float(np.linalg.norm(view_dir))
        if view_norm < 1e-8:
            continue

        edge_a = face_points[1] - face_points[0]
        edge_b = face_points[2] - face_points[1]
        normal = np.cross(edge_a, edge_b)
        normal_norm = float(np.linalg.norm(normal))
        if normal_norm < 1e-8:
            continue
        if np.dot(normal, face_center - box_center) < 0:
            normal = -normal
        score = float(np.dot(normal / normal_norm, view_dir / view_norm))
        if score > best_score:
            best_face_type, best_score = face_type, score
    return best_face_type


def _get_cut_object_side_face(face_type_or_state: int, cut_side: Optional[str] = None, corners_3d: Optional[np.ndarray] = None) -> Optional[int]:
    if face_type_or_state not in {CUT_STATE_IN, CUT_STATE_OUT}:
        return None
    side_face_type = _get_camera_facing_side_face_from_corners(corners_3d)
    if side_face_type in (2, 3):
        return side_face_type
    if cut_side not in {"left", "right"}:
        return None
    return 3 if cut_side == "left" else 2


def _reconstruct_pred_corners_for_cut_edge(
    pred_41: np.ndarray,
    anchor_xy: np.ndarray,
    stride: float,
    calib: ResizedCalib,
    cut_state: Optional[int] = None,
) -> Optional[np.ndarray]:
    if calib is None:
        return None

    resolved_cut_state = get_pred_cut_state(pred_41) if cut_state is None else int(cut_state)
    dims = np.asarray(pred_41[27:30], dtype=np.float32)
    rot_y = _decode_yaw_from_prediction(pred_41)
    if np.any(np.isnan(dims)) or not np.isfinite(rot_y):
        return None

    primary_face = get_pred_cut_primary_face(resolved_cut_state)
    if primary_face is not None:
        off = FACE_OFFSETS_41[primary_face]
        z_face = float(pred_41[off])
        uv_face_offset = np.asarray(pred_41[off + 1 : off + 3], dtype=np.float32)
        if np.isfinite(z_face) and z_face > 0 and np.isfinite(uv_face_offset).all():
            u_face = float((anchor_xy[0] + uv_face_offset[0]) * stride)
            v_face = float((anchor_xy[1] + uv_face_offset[1]) * stride)
            corners = reconstruct_3d_box_from_face((u_face, v_face), z_face, dims, rot_y, primary_face, calib)
            if corners is not None:
                return corners

    z_whole = float(pred_41[24])
    uv_whole_offset = np.asarray(pred_41[25:27], dtype=np.float32)
    if not np.isfinite(z_whole) or z_whole <= 0 or not np.isfinite(uv_whole_offset).all():
        return None
    u_whole = float((anchor_xy[0] + uv_whole_offset[0]) * stride)
    v_whole = float((anchor_xy[1] + uv_whole_offset[1]) * stride)
    return reconstruct_3d_box_from_whole((u_whole, v_whole), z_whole, dims, rot_y, calib)


def _decode_cut_partial_side_edge_from_prediction(
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    img_w: float,
    cut_side: Optional[str] = None,
    calib: Optional[ResizedCalib] = None,
    corners_3d: Optional[np.ndarray] = None,
) -> Optional[DecodedVisibleEdge]:
    if pred_edge_60 is None:
        return None

    cut_state = get_pred_cut_state(pred_41)
    if cut_state == CUT_STATE_NORMAL:
        return None

    if corners_3d is None and calib is not None:
        corners_3d = _reconstruct_pred_corners_for_cut_edge(pred_41, anchor_xy, stride, calib, cut_state=cut_state)

    side_face_type = _get_cut_object_side_face(cut_state, cut_side, corners_3d=corners_3d)
    if side_face_type is None:
        return None

    return decode_visible_face_edge_from_prediction(pred_edge_60, side_face_type, anchor_xy, stride)


def _bev_edge_points(points_3d: np.ndarray) -> Optional[np.ndarray]:
    pts = np.asarray(points_3d, dtype=np.float64)
    if pts.ndim != 2 or pts.shape[0] < 2 or pts.shape[1] != 3:
        return None
    valid = np.isfinite(pts).all(axis=1)
    pts = pts[valid]
    if len(pts) < 2:
        return None
    return pts[:, [0, 2]]


def _fit_bev_edge_axis(points_3d: np.ndarray) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
    bev_points = _bev_edge_points(points_3d)
    if bev_points is None:
        return None, None
    midpoint = np.mean(bev_points, axis=0)
    centered = bev_points - midpoint
    try:
        _, _, vh = np.linalg.svd(centered, full_matrices=False)
    except np.linalg.LinAlgError:
        return None, None
    axis = np.asarray(vh[0], dtype=np.float64)
    norm = float(np.linalg.norm(axis))
    if norm < 1e-8:
        return None, None
    return axis / norm, midpoint


def _estimate_two_edge_yaw_from_candidates(
    longitudinal_candidate: Optional[dict[str, Any]],
    side_candidate: Optional[dict[str, Any]],
    reference_yaw: Optional[float] = None,
) -> float:
    if longitudinal_candidate is None or side_candidate is None:
        return float("nan")
    if int(longitudinal_candidate["face_type"]) not in (0, 1) or int(side_candidate["face_type"]) not in (2, 3):
        return float("nan")

    side_axis, side_midpoint = _fit_bev_edge_axis(side_candidate["points_3d"])
    _, long_midpoint = _fit_bev_edge_axis(longitudinal_candidate["points_3d"])
    if side_axis is None or long_midpoint is None or side_midpoint is None:
        return float("nan")

    long_face_type = int(longitudinal_candidate["face_type"])
    side_face_type = int(side_candidate["face_type"])

    def _rot_ccw(v: np.ndarray) -> np.ndarray:
        return np.array([-v[1], v[0]], dtype=np.float64)

    def _face_normal(forward: np.ndarray, face_type: int) -> np.ndarray:
        if face_type == 0:
            return forward
        if face_type == 1:
            return -forward
        if face_type == 2:
            return _rot_ccw(forward)
        return -_rot_ccw(forward)

    forward_candidates = (side_axis, -side_axis)
    best_forward = min(
        forward_candidates,
        key=lambda forward: float(np.dot(_face_normal(forward, long_face_type), long_midpoint))
        + float(np.dot(_face_normal(forward, side_face_type), side_midpoint)),
    )
    if reference_yaw is not None and np.isfinite(reference_yaw):
        ref_forward = np.array([np.cos(float(reference_yaw)), -np.sin(float(reference_yaw))], dtype=np.float64)
        if float(np.dot(best_forward, ref_forward)) < 0.0:
            best_forward = -best_forward
    yaw = np.arctan2(-best_forward[1], best_forward[0])
    return float((yaw + np.pi) % (2 * np.pi) - np.pi)


def visible_face_edges_to_yaw(
    face_edges_3d: Any,
    face_scores: Optional[Any] = None,
) -> float:
    if face_edges_3d is None:
        return float("nan")

    items = list(face_edges_3d.items() if hasattr(face_edges_3d, "items") else face_edges_3d)
    weighted_candidates = []
    for face_type, points_3d in items:
        weight = 1.0
        if face_scores is not None:
            if hasattr(face_scores, "get"):
                weight = face_scores.get(face_type, 1.0)
            else:
                weight = face_scores[face_type]
        if not np.isfinite(weight) or weight <= 0:
            weight = 1.0
        weighted_candidates.append(
            {
                "face_type": int(face_type),
                "points_3d": np.asarray(points_3d, dtype=np.float32),
                "score": float(weight),
            }
        )

    longitudinal_candidates = [candidate for candidate in weighted_candidates if candidate["face_type"] in (0, 1)]
    side_candidates = [candidate for candidate in weighted_candidates if candidate["face_type"] in (2, 3)]
    if longitudinal_candidates and side_candidates:
        longitudinal_candidate = max(longitudinal_candidates, key=lambda item: item["score"])
        side_candidate = max(side_candidates, key=lambda item: item["score"])
        yaw = _estimate_two_edge_yaw_from_candidates(longitudinal_candidate, side_candidate)
        if np.isfinite(yaw):
            return yaw

    yaws, weights = [], []
    for face_type, points_3d in items:
        yaw = edge_points_to_yaw(points_3d, face_type)
        if not np.isfinite(yaw):
            continue

        weight = 1.0
        if face_scores is not None:
            if hasattr(face_scores, "get"):
                weight = face_scores.get(face_type, 1.0)
            else:
                weight = face_scores[face_type]
        if not np.isfinite(weight) or weight <= 0:
            weight = 1.0

        yaws.append(float(yaw))
        weights.append(float(weight))

    if not yaws:
        return float("nan")
    if len(yaws) == 1:
        return float(yaws[0])

    forward = np.stack([np.cos(yaws), -np.sin(yaws)], axis=1)
    mean_forward = np.sum(forward * np.asarray(weights, dtype=np.float64)[:, None], axis=0)
    norm = float(np.linalg.norm(mean_forward))
    if norm < 1e-8:
        return float(yaws[int(np.argmax(weights))])

    mean_forward /= norm
    yaw = np.arctan2(-mean_forward[1], mean_forward[0])
    return float((yaw + np.pi) % (2 * np.pi) - np.pi)


def _resolve_two_face_candidate_roles(candidates: list[dict[str, Any]], yaw: float) -> Optional[dict[str, Any]]:
    if candidates is None or len(candidates) < 2 or not np.isfinite(float(yaw)):
        return None

    forward_bev = np.array([np.cos(float(yaw)), -np.sin(float(yaw))], dtype=np.float64)
    right_bev = np.array([np.sin(float(yaw)), np.cos(float(yaw))], dtype=np.float64)

    role_candidates = []
    for index, candidate in enumerate(candidates[:2]):
        axis, midpoint = _fit_bev_edge_axis(candidate["points_3d"])
        if axis is None or midpoint is None:
            return None
        role_candidates.append(
            {
                "index": int(index),
                "candidate": candidate,
                "axis": axis,
                "midpoint": midpoint,
                "forward_align": abs(float(np.dot(axis, forward_bev))),
                "right_align": abs(float(np.dot(axis, right_bev))),
            }
        )

    def _role_label_penalty(info: dict[str, Any], role: str) -> int:
        face_type = int(info["candidate"].get("face_type", -1))
        if role == "longitudinal":
            return 0 if face_type in (0, 1) else 1
        return 0 if face_type in (2, 3) else 1

    assignments = ((0, 1), (1, 0))
    best_assignment = min(
        assignments,
        key=lambda assignment: (
            (1.0 - role_candidates[assignment[0]]["right_align"]) + (1.0 - role_candidates[assignment[1]]["forward_align"]),
            _role_label_penalty(role_candidates[assignment[0]], "longitudinal")
            + _role_label_penalty(role_candidates[assignment[1]], "side"),
            -(role_candidates[assignment[0]]["right_align"] + role_candidates[assignment[1]]["forward_align"]),
        ),
    )
    longitudinal_info = role_candidates[best_assignment[0]]
    side_info = role_candidates[best_assignment[1]]
    return {
        "forward_bev": forward_bev,
        "right_bev": right_bev,
        "longitudinal": longitudinal_info,
        "side": side_info,
    }


def _resolve_two_face_center_from_geometry(
    longitudinal_info: dict[str, Any],
    side_info: dict[str, Any],
    length_m: float,
    width_m: float,
) -> Optional[dict[str, Any]]:
    if longitudinal_info is None or side_info is None:
        return None

    forward_bev = np.asarray(longitudinal_info["forward_bev"], dtype=np.float64)
    right_bev = np.asarray(longitudinal_info["right_bev"], dtype=np.float64)
    long_mid = np.asarray(longitudinal_info["midpoint"], dtype=np.float64)
    side_mid = np.asarray(side_info["midpoint"], dtype=np.float64)
    if not np.isfinite(long_mid).all() or not np.isfinite(side_mid).all():
        return None

    raw_longitudinal_face_type = int(longitudinal_info["candidate"].get("face_type", -1))
    if raw_longitudinal_face_type == 0:
        longitudinal_options = ((1.0, 0),)
    elif raw_longitudinal_face_type == 1:
        longitudinal_options = ((-1.0, 1),)
    else:
        longitudinal_options = ((1.0, 0), (-1.0, 1))

    best = None
    for longitudinal_sign, longitudinal_face_type in longitudinal_options:
        center_from_longitudinal = long_mid - longitudinal_sign * forward_bev * (float(length_m) * 0.5)
        for side_sign, side_face_type in ((1.0, 2), (-1.0, 3)):
            center_from_side = side_mid - side_sign * right_bev * (float(width_m) * 0.5)
            disagreement = float(np.linalg.norm(center_from_longitudinal - center_from_side))
            if best is None or disagreement < best["disagreement"]:
                best = {
                    "center_from_longitudinal": center_from_longitudinal,
                    "center_from_side": center_from_side,
                    "longitudinal_face_type": int(longitudinal_face_type),
                    "side_face_type": int(side_face_type),
                    "disagreement": disagreement,
                }

    if best is None:
        return None

    longitudinal_coord = float(np.dot(best["center_from_longitudinal"], forward_bev))
    lateral_coord = float(np.dot(best["center_from_side"], right_bev))
    center_bev = longitudinal_coord * forward_bev + lateral_coord * right_bev
    return {
        "center_bev": center_bev,
        "longitudinal_face_type": int(best["longitudinal_face_type"]),
        "side_face_type": int(best["side_face_type"]),
        "center_from_longitudinal": best["center_from_longitudinal"],
        "center_from_side": best["center_from_side"],
    }


def _align_yaw_to_reference_pi_periodic(yaw: float, reference_yaw: float) -> float:
    if not np.isfinite(yaw) or not np.isfinite(reference_yaw):
        return float(yaw)
    base = float((float(yaw) + np.pi) % (2 * np.pi) - np.pi)
    alt = float((float(yaw) + 2 * np.pi) % (2 * np.pi) - np.pi)
    return min(
        (base, alt),
        key=lambda candidate: abs(float((candidate - float(reference_yaw) + np.pi) % (2 * np.pi) - np.pi)),
    )


def _estimate_single_edge_yaw_with_cut_primary_face(
    candidate: Optional[dict[str, Any]],
    cut_state: int,
    reference_yaw: Optional[float] = None,
) -> float:
    if candidate is None or cut_state not in (CUT_STATE_IN, CUT_STATE_OUT):
        return float("nan")

    face_type = int(candidate["face_type"])
    if face_type in (0, 1):
        yaw = edge_points_to_yaw(candidate["points_3d"], face_type)
        if reference_yaw is not None and np.isfinite(reference_yaw):
            return _align_yaw_to_reference_pi_periodic(yaw, reference_yaw)
        primary_face = get_pred_cut_primary_face(cut_state)
        if primary_face in (0, 1) and int(primary_face) != face_type:
            return float((float(yaw) + 2 * np.pi) % (2 * np.pi) - np.pi)
        return float(yaw)
    if face_type not in (2, 3):
        return float("nan")

    axis, midpoint = _fit_bev_edge_axis(candidate["points_3d"])
    midpoint_points = _bev_edge_points(candidate["points_3d"])
    if axis is None or midpoint is None or midpoint_points is None:
        return float("nan")
    midpoint = np.mean(midpoint_points, axis=0)

    yaw_candidates = [float((np.arctan2(-forward[1], forward[0]) + np.pi) % (2 * np.pi) - np.pi) for forward in (axis, -axis)]
    primary_face = get_pred_cut_primary_face(cut_state)
    if primary_face in (0, 1):
        matched = []
        for yaw in yaw_candidates:
            forward = np.array([np.cos(float(yaw)), -np.sin(float(yaw))], dtype=np.float64)
            longitudinal_score = float(np.dot(forward, midpoint))
            if (int(primary_face) == 0 and longitudinal_score > 0.0) or (int(primary_face) == 1 and longitudinal_score < 0.0):
                matched.append(float(yaw))
        candidates = matched or yaw_candidates
    else:
        candidates = yaw_candidates
    yaw = float(candidates[0])
    if reference_yaw is not None and np.isfinite(reference_yaw):
        return _align_yaw_to_reference_pi_periodic(yaw, reference_yaw)
    return yaw


def extract_face_regressed_size_priors_from_prediction(pred_41: np.ndarray) -> dict[int, dict[str, float]]:
    p = np.asarray(pred_41, dtype=np.float32).reshape(-1)
    priors: dict[int, dict[str, float]] = {}
    for face_type, off in enumerate(FACE_OFFSETS_41):
        size_pair = np.asarray(p[off + 3 : off + 5], dtype=np.float32).reshape(-1)
        if size_pair.shape != (2,) or not np.isfinite(size_pair).all():
            continue
        if face_type in (0, 1):
            priors[int(face_type)] = {"height": float(abs(size_pair[0])), "width": float(abs(size_pair[1]))}
        else:
            priors[int(face_type)] = {"length": float(abs(size_pair[0])), "height": float(abs(size_pair[1]))}
    return priors


def _select_edge_or_regressed_size(
    measured_size_m: Optional[float],
    regressed_size_m: float,
    min_fraction: float = 0.85,
    max_fraction: float = 1.35,
) -> tuple[Optional[float], Optional[str]]:
    regressed = float(abs(regressed_size_m))
    if not np.isfinite(regressed) or regressed <= 1e-6:
        return None, None
    measured = None if measured_size_m is None else float(abs(measured_size_m))
    if measured is None or not np.isfinite(measured) or measured <= 1e-6:
        return regressed, "regressed"
    fraction = measured / regressed
    if fraction < float(min_fraction) or fraction > float(max_fraction):
        return regressed, "regressed"
    return measured, "edge"


def reconstruct_edge_based_box_from_selection(
    edge_selection: Optional[dict[str, Any]],
    box_center_y_m: Optional[float],
    regressed_dims: np.ndarray,
    face_regressed_dims_by_type: Optional[dict[int, dict[str, float]]] = None,
) -> Optional[dict[str, Any]]:
    if edge_selection is None:
        return None
    yaw = float(edge_selection.get("yaw", float("nan")))
    if not np.isfinite(yaw):
        return None

    dims_reg = np.asarray(regressed_dims, dtype=np.float32).reshape(-1)
    if dims_reg.shape != (3,) or not np.isfinite(dims_reg).all():
        return None
    reg_length = float(abs(dims_reg[0]))
    box_height = float(abs(dims_reg[1]))
    reg_width = float(abs(dims_reg[2]))
    if reg_length <= 1e-6 or box_height <= 1e-6 or reg_width <= 1e-6:
        return None

    face_types = tuple(int(face_type) for face_type in (edge_selection.get("face_types") or ()))
    edge_batches = _edge_batches_to_list(edge_selection.get("edge_points_3d"))
    if len(face_types) != len(edge_batches):
        return None

    face_is_partial = tuple(bool(flag) for flag in (edge_selection.get("face_is_partial") or ()))
    if len(face_is_partial) < len(face_types):
        face_is_partial = face_is_partial + (False,) * (len(face_types) - len(face_is_partial))

    candidates = []
    for face_type, points_3d, is_partial in zip(face_types, edge_batches, face_is_partial):
        pts = np.asarray(points_3d, dtype=np.float32)
        if pts.ndim != 2 or pts.shape[0] < 2 or pts.shape[1] != 3 or not np.isfinite(pts).all():
            return None
        candidates.append({"face_type": int(face_type), "points_3d": pts, "is_partial": bool(is_partial)})

    forward = np.array([np.cos(yaw), 0.0, -np.sin(yaw)], dtype=np.float64)
    right = np.array([np.sin(yaw), 0.0, np.cos(yaw)], dtype=np.float64)
    center_x = None
    center_z = None
    length_m = None
    width_m = None
    length_source = None
    width_source = None
    mode = None
    resolved_face_types = list(face_types)
    resolved_longitudinal_face_type = None
    resolved_side_face_type = None
    face_regressed_dims_by_type = face_regressed_dims_by_type or {}

    def _face_size_prior(candidate, key, fallback, max_ratio=1.25):
        if candidate is None or bool(candidate.get("is_partial")):
            return float(fallback)
        prior = face_regressed_dims_by_type.get(int(candidate["face_type"]), {})
        value = prior.get(key)
        if value is None or not np.isfinite(float(value)) or float(value) <= 1e-6:
            return float(fallback)
        prior_value = float(value)
        fallback_value = float(abs(fallback))
        if fallback_value <= 1e-6:
            return prior_value
        ratio = max(prior_value / fallback_value, fallback_value / prior_value)
        if ratio > float(max_ratio):
            return fallback_value
        return prior_value

    role_resolution = _resolve_two_face_candidate_roles(candidates, yaw) if len(candidates) >= 2 else None
    if role_resolution is not None:
        longitudinal_info = {**role_resolution["longitudinal"], "forward_bev": role_resolution["forward_bev"], "right_bev": role_resolution["right_bev"]}
        side_info = {**role_resolution["side"], "forward_bev": role_resolution["forward_bev"], "right_bev": role_resolution["right_bev"]}
        longitudinal_candidate = longitudinal_info["candidate"]
        side_candidate = side_info["candidate"]
        side_length_m = None if bool(side_candidate.get("is_partial")) else _edge_segment_length_3d(side_candidate["points_3d"])
        width_from_long_m = None if bool(longitudinal_candidate.get("is_partial")) else _edge_segment_length_3d(longitudinal_candidate["points_3d"])

        length_m, length_source = _select_edge_or_regressed_size(side_length_m, _face_size_prior(side_candidate, "length", reg_length))
        width_m, width_source = _select_edge_or_regressed_size(width_from_long_m, _face_size_prior(longitudinal_candidate, "width", reg_width))
        if length_m is None or width_m is None:
            return None

        center_resolution = _resolve_two_face_center_from_geometry(longitudinal_info, side_info, length_m, width_m)
        if center_resolution is None:
            return None
        center_bev = np.asarray(center_resolution["center_bev"], dtype=np.float64)
        if center_bev.shape != (2,) or not np.isfinite(center_bev).all():
            return None
        center_x = float(center_bev[0])
        center_z = float(center_bev[1])
        resolved_longitudinal_face_type = int(center_resolution["longitudinal_face_type"])
        resolved_side_face_type = int(center_resolution["side_face_type"])
        resolved_face_types[int(longitudinal_info["index"])] = resolved_longitudinal_face_type
        resolved_face_types[int(side_info["index"])] = resolved_side_face_type
        mode = "two-face"
    else:
        longitudinal_candidate = next((candidate for candidate in candidates if candidate["face_type"] in (0, 1)), None)
        side_candidate = next((candidate for candidate in candidates if candidate["face_type"] in (2, 3)), None)

    if mode == "two-face":
        pass
    elif longitudinal_candidate is not None:
        long_mid = np.mean(np.asarray(longitudinal_candidate["points_3d"], dtype=np.float64), axis=0)
        width_from_long_m = None if bool(longitudinal_candidate.get("is_partial")) else _edge_segment_length_3d(longitudinal_candidate["points_3d"])
        if not np.isfinite(long_mid).all():
            return None
        width_m, width_source = _select_edge_or_regressed_size(width_from_long_m, _face_size_prior(longitudinal_candidate, "width", reg_width))
        if width_m is None:
            return None
        longitudinal_sign = 1.0 if int(longitudinal_candidate["face_type"]) == 0 else -1.0
        center_from_longitudinal = long_mid - longitudinal_sign * forward * (float(reg_length) * 0.5)
        center_x = float(center_from_longitudinal[0])
        center_z = float(center_from_longitudinal[2])
        length_m = float(reg_length)
        width_source = width_source or "regressed"
        length_source = "regressed"
        resolved_longitudinal_face_type = int(longitudinal_candidate["face_type"])
        mode = "front-rear"
    elif side_candidate is not None:
        side_mid = np.mean(np.asarray(side_candidate["points_3d"], dtype=np.float64), axis=0)
        side_length_m = None if bool(side_candidate.get("is_partial")) else _edge_segment_length_3d(side_candidate["points_3d"])
        if not np.isfinite(side_mid).all():
            return None
        length_m, length_source = _select_edge_or_regressed_size(side_length_m, _face_size_prior(side_candidate, "length", reg_length))
        if length_m is None:
            return None
        side_sign = 1.0 if int(side_candidate["face_type"]) == 2 else -1.0
        center_from_side = side_mid - side_sign * right * (float(reg_width) * 0.5)
        center_x = float(center_from_side[0])
        center_z = float(center_from_side[2])
        width_m = float(reg_width)
        width_source = "regressed"
        resolved_side_face_type = int(side_candidate["face_type"])
        mode = "side"
    else:
        return None

    all_y = np.concatenate([candidate["points_3d"][:, 1] for candidate in candidates], axis=0)
    if all_y.size == 0 or not np.isfinite(all_y).all():
        if box_center_y_m is None or not np.isfinite(float(box_center_y_m)):
            return None
        center_y = float(box_center_y_m)
    else:
        center_y = float(np.mean(all_y) - box_height * 0.5)

    center = np.array([float(center_x), float(center_y), float(center_z)], dtype=np.float32)
    if not np.isfinite(center).all():
        return None

    dims = np.array([float(length_m), float(box_height), float(width_m)], dtype=np.float32)
    corners_3d = compute_3d_box_corners(center, dims, float(yaw), face_type=-1)
    return {
        "center": center,
        "dims": dims,
        "yaw": float(yaw),
        "corners_3d": corners_3d.astype(np.float32),
        "mode": mode,
        "side_length_m": float(length_m),
        "width_m": float(width_m),
        "length_source": length_source,
        "width_source": width_source,
        "face_types": tuple(int(face_type) for face_type in resolved_face_types),
        "longitudinal_face_type": resolved_longitudinal_face_type,
        "side_face_type": resolved_side_face_type,
    }


def decode_pred_face_anchor(
    pred_41: np.ndarray,
    anchor_xy: np.ndarray,
    stride: float,
    calib: ResizedCalib,
    face_type: int,
) -> Optional[dict[str, Any]]:
    if face_type not in range(4):
        return None
    off = FACE_OFFSETS_41[face_type]
    z_face = float(pred_41[off])
    uv_face_offset = np.asarray(pred_41[off + 1 : off + 3], dtype=np.float32)
    if not np.isfinite(z_face) or z_face <= 0 or not np.isfinite(uv_face_offset).all():
        return None
    u_face = float((anchor_xy[0] + uv_face_offset[0]) * stride)
    v_face = float((anchor_xy[1] + uv_face_offset[1]) * stride)
    center_3d = back_project_2d_to_3d((u_face, v_face), z_face, calib)
    if center_3d is None:
        return None
    center_arr = np.asarray(center_3d, dtype=np.float32)
    if center_arr.shape != (3,) or not np.isfinite(center_arr).all():
        return None
    return {"face_type": int(face_type), "center_3d": center_arr, "center_2d": np.array([u_face, v_face], dtype=np.float32)}


def select_best_score_pred_face_anchor(
    pred_41: np.ndarray,
    anchor_xy: np.ndarray,
    stride: float,
    calib: ResizedCalib,
    visible_faces: list[tuple[int, float]],
) -> Optional[dict[str, Any]]:
    if not visible_faces:
        return None
    best_face_type, _ = max(
        ((int(face_type), float(score)) for face_type, score in visible_faces if int(face_type) in range(4)),
        key=lambda item: item[1],
        default=(-1, float("-inf")),
    )
    if best_face_type not in range(4):
        return None
    return decode_pred_face_anchor(pred_41, anchor_xy, stride, calib, best_face_type)


def decode_edge_yaw_selection_from_prediction(
    pred_41: np.ndarray,
    pred_edge_60: Optional[np.ndarray],
    anchor_xy: np.ndarray,
    stride: float,
    calib: ResizedCalib,
    score_thr: float = EDGE_YAW_VALID_VISIBILITY_SCORE_THRESH,
    bbox_xyxy: Optional[np.ndarray] = None,
    img_w: Optional[float] = None,
    img_h: Optional[float] = None,
    max_lateral_dist_m: Optional[float] = None,
    cut_side_min_visible_length_ratio: float = EDGE_YAW_CUT_SIDE_MIN_VISIBLE_LENGTH_RATIO,
    max_faces: int = 2,
) -> dict[str, Any]:
    empty = {
        "yaw": float("nan"),
        "face_types": (),
        "face_is_partial": (),
        "edge_points_2d": None,
        "edge_points_3d": None,
        "two_face_eligible": False,
        "lateral_distance_m": None,
        "lateral_ok": False if max_lateral_dist_m is not None else True,
        "cut_side_visible_length_m": None,
        "cut_side_visible_length_ratio": None,
        "cut_side_visible_ratio_ok": None,
        "is_valid": False,
    }
    if pred_edge_60 is None:
        return empty

    inferred_img_w = float(img_w) if img_w is not None else None
    inferred_img_h = float(img_h) if img_h is not None else None
    if inferred_img_w is None:
        if bbox_xyxy is not None:
            inferred_img_w = max(float(np.asarray(bbox_xyxy, dtype=np.float64)[2]), 1.0)
        else:
            inferred_img_w = max(float((anchor_xy[0] + pred_41[25]) * stride) * 2.0, 1.0)

    decode_visible_faces = list(
        select_pred_visible_faces_for_decode(pred_41, score_thr=FACE_VISIBILITY_SCORE_THRESH, bbox_xyxy=bbox_xyxy, img_w=inferred_img_w)
    )
    anchor_face = select_best_score_pred_face_anchor(pred_41, anchor_xy, stride, calib, decode_visible_faces)
    lateral_distance_m = None if anchor_face is None else _prediction_lateral_distance_m_from_center(anchor_face.get("center_3d"))
    lateral_ok = bool(max_lateral_dist_m is None or (lateral_distance_m is not None and lateral_distance_m < float(max_lateral_dist_m)))
    primary_candidate_face_type = max(
        ((int(face_type), float(score)) for face_type, score in decode_visible_faces if int(face_type) in range(4)),
        key=lambda item: item[1],
        default=(-1, float("-inf")),
    )[0]

    raw_cut_state = get_pred_cut_state(pred_41)
    primary_face = get_pred_cut_primary_face(raw_cut_state)
    visible_faces = list(select_pred_visible_faces(pred_41, score_thr=score_thr))
    longitudinal_faces = {face_type for face_type, _ in visible_faces if face_type in (0, 1)}
    if primary_face in longitudinal_faces and len(longitudinal_faces) > 1:
        visible_faces = [(face_type, score) for face_type, score in visible_faces if face_type not in (0, 1) or face_type == primary_face]

    def _decode_face_candidate(face_type: int, score: float, require_in_image: bool = True) -> Optional[dict[str, Any]]:
        if face_type not in range(4):
            return None
        decoded = decode_visible_face_edge_from_prediction(pred_edge_60, face_type, anchor_xy, stride)
        if decoded is None:
            return None
        drawable = _decoded_edge_points_are_drawable(
            decoded.points_2d,
            inferred_img_w if require_in_image else None,
            inferred_img_h if require_in_image else None,
        )
        if not drawable:
            return None
        points_3d = _decoded_edge_to_points_3d(decoded, calib)
        if points_3d is None:
            return None
        return {
            "face_type": int(face_type),
            "score": float(score),
            "is_partial": False,
            "points_2d": np.asarray(decoded.points_2d, dtype=np.float32),
            "points_3d": np.asarray(points_3d, dtype=np.float32),
        }

    face_candidates = {}
    for face_type, score in visible_faces:
        candidate = _decode_face_candidate(face_type, score)
        if candidate is not None:
            face_candidates[int(face_type)] = candidate

    primary_candidate = None
    if primary_candidate_face_type in range(4):
        primary_score = next((float(score) for face_type, score in decode_visible_faces if int(face_type) == int(primary_candidate_face_type)), float("-inf"))
        primary_candidate = _decode_face_candidate(int(primary_candidate_face_type), primary_score, require_in_image=False)
        if primary_candidate is not None:
            face_candidates.pop(int(primary_candidate_face_type), None)

    resolved_cut_state, cut_side = _resolve_pred_cut_state_for_decode(pred_41, bbox_xyxy=bbox_xyxy, img_w=inferred_img_w)
    partial_candidate = None
    cut_side_visible_length_m = None
    cut_side_visible_length_ratio = None
    cut_side_visible_ratio_ok = None
    cut_corners = None
    if resolved_cut_state != CUT_STATE_NORMAL:
        cut_corners = _reconstruct_pred_corners_for_cut_edge(pred_41, anchor_xy, stride, calib, cut_state=resolved_cut_state)
        partial_edge = _decode_cut_partial_side_edge_from_prediction(
            pred_41,
            pred_edge_60,
            anchor_xy,
            stride,
            img_w=inferred_img_w,
            cut_side=cut_side,
            calib=calib,
            corners_3d=cut_corners,
        )
        if partial_edge is not None and not _decoded_edge_points_are_drawable(partial_edge.points_2d, inferred_img_w, inferred_img_h):
            partial_edge = None
        partial_points_3d = _decoded_edge_to_points_3d(partial_edge, calib)
        cut_side_visible_length_m = _edge_segment_length_3d(partial_points_3d)
        box_length_m = float(abs(pred_41[27])) if np.isfinite(pred_41[27]) else None
        if cut_side_visible_length_m is not None and box_length_m is not None and box_length_m > 1e-6:
            cut_side_visible_length_ratio = float(cut_side_visible_length_m / box_length_m)
            cut_side_visible_ratio_ok = bool(cut_side_visible_length_ratio > float(cut_side_min_visible_length_ratio))
        else:
            cut_side_visible_ratio_ok = False
        if partial_edge is not None and partial_points_3d is not None:
            partial_face_type = int(partial_edge.face_type)
            partial_score = face_candidates.get(partial_face_type, {}).get("score", 1.0)
            partial_candidate = {
                "face_type": partial_face_type,
                "score": float(partial_score),
                "is_partial": True,
                "points_2d": np.asarray(partial_edge.points_2d, dtype=np.float32),
                "points_3d": np.asarray(partial_points_3d, dtype=np.float32),
            }
            face_candidates.pop(partial_face_type, None)

    if resolved_cut_state != CUT_STATE_NORMAL and not cut_side_visible_ratio_ok:
        partial_candidate = None

    selected_candidates: list[dict[str, Any]] = []

    def _best_candidate(candidates: list[dict[str, Any]]) -> Optional[dict[str, Any]]:
        if not candidates:
            return None
        return max(candidates, key=lambda item: (float(item["score"]), -int(item["face_type"])))

    cut_expected_side_face = None
    if raw_cut_state != CUT_STATE_NORMAL:
        cut_corners_for_side = cut_corners if resolved_cut_state != CUT_STATE_NORMAL and cut_corners is not None else _reconstruct_pred_corners_for_cut_edge(pred_41, anchor_xy, stride, calib, cut_state=raw_cut_state)
        cut_expected_side_face = _get_cut_object_side_face(raw_cut_state, corners_3d=cut_corners_for_side)

    if primary_candidate is not None:
        selected_candidates.append(primary_candidate)

    if len(selected_candidates) < int(max_faces):
        secondary_candidate = None
        if primary_candidate is not None and int(primary_candidate["face_type"]) in (0, 1):
            secondary_candidate = partial_candidate
            if secondary_candidate is None and cut_expected_side_face in (2, 3) and (resolved_cut_state == CUT_STATE_NORMAL or cut_side_visible_ratio_ok):
                secondary_candidate = face_candidates.pop(int(cut_expected_side_face), None)
            if secondary_candidate is None and (resolved_cut_state == CUT_STATE_NORMAL or cut_side_visible_ratio_ok):
                secondary_candidate = _best_candidate([candidate for candidate in face_candidates.values() if candidate["face_type"] in (2, 3)])
                if secondary_candidate is not None:
                    face_candidates.pop(int(secondary_candidate["face_type"]), None)
        elif primary_candidate is not None and int(primary_candidate["face_type"]) in (2, 3):
            longitudinal_candidate = None
            if primary_face is not None:
                longitudinal_candidate = face_candidates.pop(int(primary_face), None)
            if longitudinal_candidate is None:
                longitudinal_candidate = _best_candidate([candidate for candidate in face_candidates.values() if candidate["face_type"] in (0, 1)])
                if longitudinal_candidate is not None:
                    face_candidates.pop(int(longitudinal_candidate["face_type"]), None)
            secondary_candidate = longitudinal_candidate
        if secondary_candidate is not None:
            selected_candidates.append(secondary_candidate)

    if not selected_candidates:
        return {
            **empty,
            "cut_side_visible_length_m": cut_side_visible_length_m,
            "cut_side_visible_length_ratio": cut_side_visible_length_ratio,
            "cut_side_visible_ratio_ok": cut_side_visible_ratio_ok,
        }

    edge_points_3d = _stack_edge_batches([candidate["points_3d"] for candidate in selected_candidates])
    edge_points_2d = _stack_edge_batches([candidate["points_2d"] for candidate in selected_candidates])
    face_types = tuple(int(candidate["face_type"]) for candidate in selected_candidates)
    face_is_partial = tuple(bool(candidate.get("is_partial", False)) for candidate in selected_candidates)

    if len(selected_candidates) >= 2:
        longitudinal_selected = next((candidate for candidate in selected_candidates if candidate["face_type"] in (0, 1)), None)
        side_selected = next((candidate for candidate in selected_candidates if candidate["face_type"] in (2, 3)), None)
        yaw = _estimate_two_edge_yaw_from_candidates(longitudinal_selected, side_selected, reference_yaw=_decode_yaw_from_prediction(pred_41))
        if not np.isfinite(yaw):
            yaw = visible_face_edges_to_yaw(
                {candidate["face_type"]: candidate["points_3d"] for candidate in selected_candidates},
                face_scores={candidate["face_type"]: candidate["score"] for candidate in selected_candidates},
            )
    else:
        only_candidate = selected_candidates[0]
        if raw_cut_state in (CUT_STATE_IN, CUT_STATE_OUT):
            yaw = _estimate_single_edge_yaw_with_cut_primary_face(only_candidate, cut_state=raw_cut_state, reference_yaw=_decode_yaw_from_prediction(pred_41))
            if not np.isfinite(yaw):
                yaw = edge_points_to_yaw(only_candidate["points_3d"], only_candidate["face_type"])
        else:
            yaw = edge_points_to_yaw(only_candidate["points_3d"], only_candidate["face_type"])

    has_longitudinal = any(candidate["face_type"] in (0, 1) for candidate in selected_candidates)
    has_side = any(candidate["face_type"] in (2, 3) for candidate in selected_candidates)
    two_face_eligible = len(selected_candidates) >= 2 and has_longitudinal and has_side
    is_valid = bool(two_face_eligible and np.isfinite(yaw) and lateral_ok)
    return {
        "yaw": float(yaw),
        "face_types": face_types,
        "face_is_partial": face_is_partial,
        "edge_points_2d": edge_points_2d,
        "edge_points_3d": edge_points_3d,
        "two_face_eligible": bool(two_face_eligible),
        "lateral_distance_m": lateral_distance_m,
        "lateral_ok": lateral_ok,
        "cut_side_visible_length_m": cut_side_visible_length_m,
        "cut_side_visible_length_ratio": cut_side_visible_length_ratio,
        "cut_side_visible_ratio_ok": cut_side_visible_ratio_ok,
        "is_valid": bool(is_valid),
    }
