from __future__ import annotations

import math
from typing import Any, Optional

import cv2
import numpy as np

try:
    from .two_roi_types import (
        Decoded3DPrediction,
        DecodedVisibleEdge,
        Prediction3DAttrs,
        ResizedCalib,
    )
except ImportError:
    from two_roi_types import (
        Decoded3DPrediction,
        DecodedVisibleEdge,
        Prediction3DAttrs,
        ResizedCalib,
    )


PROJECTION_Z_MIN = 0.1
YAW_BIN_OFFSETS = (0.0, np.pi / 2, -np.pi / 2, np.pi)
FACE_OFFSETS_41 = (0, 6, 12, 18)
FACE_EDGE_OFFSETS_60 = (0, 15, 30, 45)
FACE_CORNERS = {0: (4, 5, 6, 7), 1: (0, 1, 2, 3), 2: (1, 2, 5, 6), 3: (0, 3, 4, 7)}
FACE_BOTTOM_EDGE_CORNERS = {0: (6, 7), 1: (2, 3), 2: (2, 6), 3: (3, 7)}
FACE_CENTER_OFFSETS = {0: [1, 0.5, 0.5], 1: [0, 0.5, 0.5], 2: [0.5, 0.5, 1], 3: [0.5, 0.5, 0]}
FACE_VISIBILITY_SCORE_THRESH = 0.05
CUT_STATE_NORMAL = 0
CUT_STATE_IN = 1
CUT_STATE_OUT = 2
FACE_COLORS = ((0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255))


def rotation_3d_in_axis(points, angles, axis=1):
    rot_sin = np.sin(angles)
    rot_cos = np.cos(angles)
    ones = np.ones_like(rot_cos)
    zeros = np.zeros_like(rot_cos)
    if axis == 1:
        rot_mat = np.stack(
            [
                np.stack([rot_cos, zeros, -rot_sin]),
                np.stack([zeros, ones, zeros]),
                np.stack([rot_sin, zeros, rot_cos]),
            ]
        )
    elif axis == 2:
        rot_mat = np.stack(
            [
                np.stack([rot_cos, rot_sin, zeros]),
                np.stack([-rot_sin, rot_cos, zeros]),
                np.stack([zeros, zeros, ones]),
            ]
        )
    elif axis == 0:
        rot_mat = np.stack(
            [
                np.stack([ones, zeros, zeros]),
                np.stack([zeros, rot_cos, rot_sin]),
                np.stack([zeros, -rot_sin, rot_cos]),
            ]
        )
    else:
        raise ValueError(f"axis should be in [0, 1, 2], got {axis}")
    return np.dot(points, rot_mat)


def compute_3d_box_corners(center_3d, dimensions, rotation, face_type=-1):
    l, h, w = dimensions
    corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1).astype(np.float64)
    corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]]
    corners_norm -= FACE_CENTER_OFFSETS.get(face_type, [0.5, 0.5, 0.5])
    corners = np.array([l, h, w]).reshape(1, 3) * corners_norm.reshape(8, 3)
    corners = rotation_3d_in_axis(corners, rotation, axis=1)
    corners += np.array(center_3d).reshape(1, 3)
    return corners


def apply_fisheye_distortion(x, y, distort_coeffs):
    if distort_coeffs is None or len(distort_coeffs) < 4:
        return x, y
    k1, k2, k3, k4 = distort_coeffs[:4]
    r = np.sqrt(x * x + y * y)
    if r < 1e-8:
        return x, y
    theta = np.arctan(r)
    theta2 = theta * theta
    theta4 = theta2 * theta2
    theta6 = theta4 * theta2
    theta8 = theta4 * theta4
    theta_d = theta * (1 + k1 * theta2 + k2 * theta4 + k3 * theta6 + k4 * theta8)
    scale = theta_d / r
    return x * scale, y * scale


def remove_fisheye_distortion(xd, yd, distort_coeffs, max_iter=20):
    if distort_coeffs is None or len(distort_coeffs) < 4:
        return xd, yd
    k1, k2, k3, k4 = distort_coeffs[:4]
    r_d = np.sqrt(xd * xd + yd * yd)
    if r_d < 1e-8:
        return xd, yd
    theta_d = r_d
    theta_d2 = theta_d * theta_d
    theta = theta_d / (1 + k1 * theta_d2)
    for _ in range(max_iter):
        theta2 = theta * theta
        theta4 = theta2 * theta2
        theta6 = theta4 * theta2
        theta8 = theta4 * theta4
        f = theta * (1 + k1 * theta2 + k2 * theta4 + k3 * theta6 + k4 * theta8) - theta_d
        f_prime = 1 + 3 * k1 * theta2 + 5 * k2 * theta4 + 7 * k3 * theta6 + 9 * k4 * theta8
        theta_new = theta - f / f_prime
        if abs(theta_new - theta) < 1e-8:
            theta = theta_new
            break
        theta = theta_new
    r = np.tan(theta)
    scale = r / r_d
    return xd * scale, yd * scale


def project_3d_to_2d_with_distortion(points_3d, calib: ResizedCalib):
    fx, fy = calib["fx"], calib["fy"]
    cx, cy = calib["cx"], calib["cy"]
    distort_coeffs = calib.get("distort_coeffs", [])
    points_2d = np.full((len(points_3d), 2), np.nan)
    for index, (x, y, z) in enumerate(points_3d):
        if z > PROJECTION_Z_MIN:
            xn, yn = x / z, y / z
            xd, yd = apply_fisheye_distortion(xn, yn, distort_coeffs)
            points_2d[index] = [fx * xd + cx, fy * yd + cy]
    return points_2d


def project_3d_to_2d_with_calib(points_3d, calib: ResizedCalib):
    fx, fy = calib["fx"], calib["fy"]
    cx, cy = calib["cx"], calib["cy"]
    points_2d = np.full((len(points_3d), 2), np.nan)
    for index, (x, y, z) in enumerate(points_3d):
        if z > PROJECTION_Z_MIN:
            points_2d[index] = [fx * x / z + cx, fy * y / z + cy]
    return points_2d


def project_3d_to_2d(points_3d, calib: ResizedCalib):
    if calib is None:
        return np.full((len(points_3d), 2), np.nan)
    distort_coeffs = calib.get("distort_coeffs", [])
    if distort_coeffs is not None and len(distort_coeffs) >= 4:
        return project_3d_to_2d_with_distortion(points_3d, calib)
    return project_3d_to_2d_with_calib(points_3d, calib)


def sample_3d_edge(p1, p2, num_samples=10):
    t = np.linspace(0.0, 1.0, num_samples, dtype=np.float64).reshape(-1, 1)
    return p1 + t * (p2 - p1)


def _point_inside_image(point_2d, img_w, img_h):
    x, y = float(point_2d[0]), float(point_2d[1])
    return np.isfinite(x) and np.isfinite(y) and 0.0 <= x <= img_w - 1 and 0.0 <= y <= img_h - 1


def _solve_edge_image_boundary_t(p0_2d, p1_2d, img_w, img_h):
    p0 = np.asarray(p0_2d, dtype=np.float64)
    p1 = np.asarray(p1_2d, dtype=np.float64)
    if not np.isfinite(p0).all() or not np.isfinite(p1).all():
        return None
    dx, dy = p1 - p0
    t_min, t_max = 0.0, 1.0
    for p, q in ((-dx, p0[0]), (dx, (img_w - 1) - p0[0]), (-dy, p0[1]), (dy, (img_h - 1) - p0[1])):
        if abs(p) < 1e-12:
            if q < 0:
                return None
            continue
        t = q / p
        if p < 0:
            t_min = max(t_min, t)
        else:
            t_max = min(t_max, t)
        if t_min > t_max:
            return None
    return t_min, t_max


def _project_edge_point_at_t(p1, p2, t, calib: ResizedCalib):
    point_3d = np.asarray(p1, dtype=np.float64) + float(t) * (np.asarray(p2, dtype=np.float64) - np.asarray(p1, dtype=np.float64))
    point_2d = project_3d_to_2d(point_3d[None, :], calib)[0]
    return point_3d, point_2d


def _refine_visible_edge_boundary(p1, p2, calib: ResizedCalib, img_w, img_h, t_out, t_in, steps=12):
    lo, hi = (float(t_out), float(t_in)) if t_out < t_in else (float(t_in), float(t_out))
    for _ in range(steps):
        mid = 0.5 * (lo + hi)
        _, point_2d = _project_edge_point_at_t(p1, p2, mid, calib)
        if _point_inside_image(point_2d, img_w, img_h):
            hi = mid
        else:
            lo = mid
    return hi if t_out < t_in else lo


def sample_partial_3d_edge(p1, p2, calib: ResizedCalib, img_w, img_h, num_samples=5, dense_samples=129):
    endpoints_3d = np.asarray([p1, p2], dtype=np.float64)
    dense_t = np.linspace(0.0, 1.0, dense_samples, dtype=np.float64)
    dense_points_3d = endpoints_3d[0:1] + dense_t[:, None] * (endpoints_3d[1:2] - endpoints_3d[0:1])
    dense_points_2d = project_3d_to_2d(dense_points_3d, calib)
    visible = np.array([_point_inside_image(point_2d, img_w, img_h) for point_2d in dense_points_2d], dtype=bool)
    if not visible.any():
        return None, None

    visible_idx = np.flatnonzero(visible)
    split_idx = np.where(np.diff(visible_idx) > 1)[0] + 1
    visible_runs = np.split(visible_idx, split_idx)
    visible_run = max(visible_runs, key=len)
    first_idx, last_idx = int(visible_run[0]), int(visible_run[-1])

    t_start = dense_t[first_idx]
    if first_idx > 0:
        t_start = _refine_visible_edge_boundary(
            endpoints_3d[0],
            endpoints_3d[1],
            calib,
            img_w,
            img_h,
            dense_t[first_idx - 1],
            dense_t[first_idx],
        )

    t_end = dense_t[last_idx]
    if last_idx < len(dense_t) - 1:
        t_end = _refine_visible_edge_boundary(
            endpoints_3d[0],
            endpoints_3d[1],
            calib,
            img_w,
            img_h,
            dense_t[last_idx + 1],
            dense_t[last_idx],
        )

    if t_end - t_start < 1e-6:
        return None, None

    sample_t = np.linspace(t_start, t_end, num_samples, dtype=np.float64)
    sample_points_3d = endpoints_3d[0:1] + sample_t[:, None] * (endpoints_3d[1:2] - endpoints_3d[0:1])
    sample_points_2d = project_3d_to_2d(sample_points_3d, calib)
    if np.any(np.isnan(sample_points_2d)):
        return None, None
    if not np.all([_point_inside_image(point_2d, img_w, img_h) for point_2d in sample_points_2d]):
        return None, None

    order = np.argsort(sample_points_2d[:, 0], kind="stable")
    return sample_points_3d[order], sample_points_2d[order]


def project_3d_box_edges_with_distortion(corners_3d, calib: ResizedCalib, samples_per_edge=10):
    edges = {
        "back_0": (4, 5),
        "back_1": (5, 6),
        "back_2": (6, 7),
        "back_3": (7, 4),
        "connect_0": (0, 4),
        "connect_1": (1, 5),
        "connect_2": (2, 6),
        "connect_3": (3, 7),
        "front_0": (0, 1),
        "front_1": (1, 2),
        "front_2": (2, 3),
        "front_3": (3, 0),
        "front_x1": (0, 2),
        "front_x2": (1, 3),
    }
    edge_points_2d = {}
    for edge_name, (i, j) in edges.items():
        sampled_3d = sample_3d_edge(corners_3d[i], corners_3d[j], samples_per_edge)
        edge_points_2d[edge_name] = project_3d_to_2d_with_distortion(sampled_3d, calib)
    return edge_points_2d


def plot_box3d_on_img_with_distortion(
    img,
    edge_points_2d,
    color_front=(0, 0, 255),
    color_back=(255, 0, 0),
    color_side=(255, 255, 0),
    thickness=1,
):
    front_edges = {"front_0", "front_1", "front_2", "front_3", "front_x1", "front_x2"}
    back_edges = {"back_0", "back_1", "back_2", "back_3", "back_x1", "back_x2"}
    for edge_name, points in edge_points_2d.items():
        if np.any(np.isnan(points)):
            continue
        pts = points.astype(np.int32)
        color = color_front if edge_name in front_edges else color_back if edge_name in back_edges else color_side
        cv2.polylines(img, [pts], isClosed=False, color=color, thickness=thickness, lineType=cv2.LINE_AA)
    return img


def plot_box3d_on_img(img, corners_2d, color_front=(0, 0, 255), color_back=(255, 0, 0), color_side=(255, 255, 0), thickness=1):
    line_indices = (
        (4, 5),
        (5, 6),
        (6, 7),
        (7, 4),
        (0, 4),
        (1, 5),
        (2, 6),
        (3, 7),
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 0),
        (0, 2),
        (1, 3),
    )
    front_edges = {(0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3)}
    back_edges = {(4, 5), (5, 6), (6, 7), (7, 4)}
    pts = corners_2d.astype(np.int32)
    for i, j in line_indices:
        color = color_front if (i, j) in front_edges else color_back if (i, j) in back_edges else color_side
        cv2.line(img, tuple(pts[i]), tuple(pts[j]), color, thickness, cv2.LINE_AA)
    return img


def back_project_2d_to_3d(uv, depth, calib: ResizedCalib):
    if calib is None or depth <= 0:
        return None
    fx, fy = calib["fx"], calib["fy"]
    cx, cy = calib["cx"], calib["cy"]
    u, v = uv
    xd = (u - cx) / fx
    yd = (v - cy) / fy
    distort_coeffs = calib.get("distort_coeffs", [])
    if distort_coeffs is not None and len(distort_coeffs) >= 4:
        xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs)
    else:
        xn, yn = xd, yd
    return np.array([xn * depth, yn * depth, depth], dtype=np.float64)


def reconstruct_3d_box_from_face(face_uv, face_z, dims, rot_y, face_type, calib: ResizedCalib):
    if calib is None or face_z <= 0:
        return None
    center_3d = back_project_2d_to_3d(face_uv, face_z, calib)
    if center_3d is None:
        return None
    if np.any(np.isnan(np.asarray(dims, dtype=np.float64))) or not np.isfinite(float(rot_y)):
        return None
    return compute_3d_box_corners(center_3d, dims, rot_y, face_type)


def reconstruct_3d_box_from_whole(uv, z3d, dims, rot_y, calib: ResizedCalib):
    if calib is None or z3d <= 0:
        return None
    center_3d = back_project_2d_to_3d(uv, z3d, calib)
    if center_3d is None:
        return None
    if np.any(np.isnan(np.asarray(dims, dtype=np.float64))) or not np.isfinite(float(rot_y)):
        return None
    return compute_3d_box_corners(center_3d, dims, rot_y, face_type=-1)


def get_face_bottom_edge_points(corners_3d, face_type, num_samples=5):
    if corners_3d is None or face_type not in FACE_BOTTOM_EDGE_CORNERS:
        return None
    start_idx, end_idx = FACE_BOTTOM_EDGE_CORNERS[face_type]
    return sample_3d_edge(corners_3d[start_idx], corners_3d[end_idx], num_samples=num_samples)


def project_face_bottom_edge(corners_3d, face_type, calib: ResizedCalib, num_samples=5):
    points_3d = get_face_bottom_edge_points(corners_3d, face_type, num_samples=num_samples)
    if points_3d is None:
        return None, None
    points_2d = project_3d_to_2d(points_3d, calib)
    if np.any(np.isnan(points_2d)):
        return points_3d, None
    order = np.argsort(points_2d[:, 0], kind="stable")
    return points_3d[order], points_2d[order]


def project_partial_face_bottom_edge(corners_3d, face_type, calib: ResizedCalib, img_w, img_h, num_samples=5):
    if corners_3d is None or face_type not in FACE_BOTTOM_EDGE_CORNERS:
        return None, None
    start_idx, end_idx = FACE_BOTTOM_EDGE_CORNERS[face_type]
    return sample_partial_3d_edge(corners_3d[start_idx], corners_3d[end_idx], calib, img_w, img_h, num_samples=num_samples)


def collect_face_bottom_edges(corners_3d, face_types, calib: ResizedCalib, num_samples=5):
    if corners_3d is None:
        return None, None
    edge_points_3d, edge_points_2d = [], []
    for face_type in face_types:
        points_3d, points_2d = project_face_bottom_edge(corners_3d, face_type, calib, num_samples=num_samples)
        if points_3d is None or points_2d is None:
            continue
        edge_points_3d.append(points_3d.astype(np.float32, copy=False))
        edge_points_2d.append(points_2d.astype(np.float32, copy=False))
    if not edge_points_2d:
        return None, None
    if len(edge_points_2d) == 1:
        return edge_points_3d[0], edge_points_2d[0]
    return np.stack(edge_points_3d, axis=0), np.stack(edge_points_2d, axis=0)


def _edge_batches_to_list(edge_points):
    if edge_points is None:
        return []
    arr = np.asarray(edge_points, dtype=np.float32)
    if arr.ndim == 2:
        return [arr]
    return [arr[i] for i in range(arr.shape[0])]


def _stack_edge_batches(edge_batches):
    if not edge_batches:
        return None
    if len(edge_batches) == 1:
        return edge_batches[0]
    return np.stack(edge_batches, axis=0)


def _append_edge_batch(edge_points_3d, edge_points_2d, decoded_edge: DecodedVisibleEdge):
    if decoded_edge.points_3d is None:
        return edge_points_3d, edge_points_2d
    edge3d_list = _edge_batches_to_list(edge_points_3d)
    edge2d_list = _edge_batches_to_list(edge_points_2d)
    edge3d_list.append(np.asarray(decoded_edge.points_3d, dtype=np.float32))
    edge2d_list.append(np.asarray(decoded_edge.points_2d, dtype=np.float32))
    return _stack_edge_batches(edge3d_list), _stack_edge_batches(edge2d_list)


def decode_visible_face_edge_from_prediction(pred_edge_60, face_type, anchor_xy, stride) -> Optional[DecodedVisibleEdge]:
    if pred_edge_60 is None or face_type not in range(4):
        return None
    off = FACE_EDGE_OFFSETS_60[face_type]
    face = np.asarray(pred_edge_60[off : off + 15], dtype=np.float32).reshape(5, 3)
    points_2d = np.empty((5, 2), dtype=np.float32)
    points_2d[:, 0] = (anchor_xy[0] + face[:, 0]) * stride
    points_2d[:, 1] = (anchor_xy[1] + face[:, 1]) * stride
    order = np.argsort(points_2d[:, 0], kind="stable")
    return DecodedVisibleEdge(
        face_type=int(face_type),
        points_2d=points_2d[order],
        depths=face[order, 2].astype(np.float32),
    )


def get_cut_side_from_bbox_xyxy(bbox_xyxy, img_w, tol=1.0):
    if bbox_xyxy is None:
        return None
    x1, _, x2, _ = np.asarray(bbox_xyxy, dtype=np.float64)
    touch_left = x1 <= tol and x2 > tol
    touch_right = x2 >= img_w - 1 - tol and x1 < img_w - 1 - tol
    if touch_left == touch_right:
        return None
    return "left" if touch_left else "right"


def get_cut_object_side_face(face_type_or_state, cut_side=None):
    if cut_side not in {"left", "right"}:
        return None
    if face_type_or_state not in {CUT_STATE_IN, CUT_STATE_OUT}:
        return None
    return 3 if cut_side == "left" else 2


def get_pred_cut_state(pred_41):
    cut_logits = np.asarray(pred_41[38:41], dtype=np.float32)
    return int(np.argmax(cut_logits))


def get_pred_cut_primary_face(cut_state):
    if cut_state == CUT_STATE_IN:
        return 0
    if cut_state == CUT_STATE_OUT:
        return 1
    return None


def _resolve_pred_cut_state_for_decode(pred_41, bbox_xyxy=None, img_w=None):
    cut_state = get_pred_cut_state(pred_41)
    if cut_state == CUT_STATE_NORMAL:
        return cut_state, None
    cut_side = None
    if bbox_xyxy is not None and img_w is not None:
        cut_side = get_cut_side_from_bbox_xyxy(bbox_xyxy, img_w)
    if cut_side not in {"left", "right"}:
        return CUT_STATE_NORMAL, None
    return cut_state, cut_side


def select_pred_visible_faces(pred_41, score_thr=FACE_VISIBILITY_SCORE_THRESH):
    selected = []
    for face_type, off in enumerate(FACE_OFFSETS_41):
        score = float(pred_41[off + 5])
        if np.isnan(score) or score < score_thr:
            continue
        selected.append((face_type, score))
    return selected


def _select_best_pred_face_score(pred_41):
    best_face_type, best_score = None, float("-inf")
    for face_type, off in enumerate(FACE_OFFSETS_41):
        score = float(pred_41[off + 5])
        if not np.isfinite(score):
            continue
        if score > best_score:
            best_face_type = int(face_type)
            best_score = float(score)
    if best_face_type is None:
        return None
    return best_face_type, best_score


def select_pred_visible_faces_for_decode(pred_41, score_thr=FACE_VISIBILITY_SCORE_THRESH, bbox_xyxy=None, img_w=None):
    cut_state, _ = _resolve_pred_cut_state_for_decode(pred_41, bbox_xyxy=bbox_xyxy, img_w=img_w)
    primary_face = get_pred_cut_primary_face(cut_state)
    if primary_face is not None:
        off = FACE_OFFSETS_41[primary_face]
        return [(primary_face, float(pred_41[off + 5]))]
    visible_faces = list(select_pred_visible_faces(pred_41, score_thr=score_thr))
    best_face = _select_best_pred_face_score(pred_41)
    if best_face is None:
        return visible_faces
    best_face_type, best_score = best_face
    if all(int(face_type) != int(best_face_type) for face_type, _ in visible_faces):
        visible_faces.append((int(best_face_type), float(best_score)))
    return visible_faces


def decode_cut_partial_side_edge_from_prediction(pred_41, pred_edge_60, anchor_xy, stride, img_w, cut_side=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
    side_face_type = get_cut_object_side_face(cut_state, cut_side)
    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 _decoded_edge_to_points_3d(decoded_edge: Optional[DecodedVisibleEdge], calib: ResizedCalib) -> Optional[np.ndarray]:
    if decoded_edge is None:
        return None
    points = [back_project_2d_to_3d(tuple(pt), depth, calib) for pt, depth in zip(decoded_edge.points_2d, decoded_edge.depths)]
    if any(point is None for point in points):
        return None
    return np.asarray(points, dtype=np.float32)


def edge_points_to_yaw(points_3d, face_type):
    points = np.asarray(points_3d, dtype=np.float32)
    if points.shape[0] < 2:
        return float("nan")
    direction = points[-1] - points[0]
    yaw = math.atan2(float(direction[0]), float(direction[2]))
    if face_type == 0:
        return yaw
    if face_type == 1:
        return yaw + np.pi
    if face_type == 2:
        return yaw + np.pi / 2
    if face_type == 3:
        return yaw - np.pi / 2
    return float("nan")


def visible_face_edges_to_yaw(face_edges_3d, face_scores=None):
    if not face_edges_3d:
        return float("nan")
    if face_scores:
        face_type = max(face_edges_3d.keys(), key=lambda ft: face_scores.get(ft, 0.0))
        return edge_points_to_yaw(face_edges_3d[face_type], face_type)
    face_type = next(iter(face_edges_3d))
    return edge_points_to_yaw(face_edges_3d[face_type], face_type)


def _draw_edge_points(img, edge_points_2d=None, edge_color=(0, 255, 0), thickness=1):
    if edge_points_2d is None:
        return
    points = np.asarray(edge_points_2d, dtype=np.float32)
    if points.ndim == 2:
        points = points[None, ...]
    for batch in points:
        for point in batch:
            cv2.circle(img, tuple(np.round(point).astype(np.int32)), max(thickness + 1, 2), edge_color, -1, cv2.LINE_AA)


def _decode_yaw_from_prediction(pred_41):
    yaw_cls_logits = pred_41[30:34]
    yaw_residual_sin = np.clip(pred_41[34:38], -1.0, 1.0)
    best_bin = int(np.argmax(yaw_cls_logits))
    return np.arcsin(yaw_residual_sin[best_bin]) + YAW_BIN_OFFSETS[best_bin]


def decode_3d_prediction(
    pred_41,
    anchor_xy,
    stride,
    calib,
    img_w,
    img_h,
    face_3d_classes,
    complete_3d_classes,
    cls_id,
    pred_edge_60=None,
    score_thr=FACE_VISIBILITY_SCORE_THRESH,
    bbox_xyxy=None,
) -> Optional[Decoded3DPrediction]:
    pred = pred_41
    rot_y = _decode_yaw_from_prediction(pred)
    z_whole = pred[24]
    uv_whole_offset = pred[25:27]
    dims_whole = pred[27:30]
    u_whole = (anchor_xy[0] + uv_whole_offset[0]) * stride
    v_whole = (anchor_xy[1] + uv_whole_offset[1]) * stride

    if cls_id in face_3d_classes:
        _, cut_side = _resolve_pred_cut_state_for_decode(pred, bbox_xyxy=bbox_xyxy, img_w=img_w)
        visible_faces = select_pred_visible_faces_for_decode(pred, score_thr=score_thr, bbox_xyxy=bbox_xyxy, img_w=img_w)
        best_type, _ = (-1, -1.0) if not visible_faces else max(visible_faces, key=lambda item: item[1])
        if best_type < 0:
            return None

        off = best_type * 6
        z_face = pred[off]
        uv_face_offset = pred[off + 1 : off + 3]
        u_face = (anchor_xy[0] + uv_face_offset[0]) * stride
        v_face = (anchor_xy[1] + uv_face_offset[1]) * stride
        corners = reconstruct_3d_box_from_face((u_face, v_face), z_face, dims_whole, rot_y, best_type, calib)
        if corners is None:
            return None

        edge_points_3d, edge_points_2d = collect_face_bottom_edges(
            corners,
            [face_type for face_type, _ in visible_faces],
            calib,
            num_samples=5,
        )
        if pred_edge_60 is not None:
            pred_edge_points_2d, pred_edge_points_3d = [], []
            for face_type, _ in visible_faces:
                pred_edge = decode_visible_face_edge_from_prediction(pred_edge_60, face_type, anchor_xy, stride)
                if pred_edge is None:
                    continue
                points_3d = [back_project_2d_to_3d(tuple(pt), depth, calib) for pt, depth in zip(pred_edge.points_2d, pred_edge.depths)]
                if any(point is None for point in points_3d):
                    continue
                pred_edge_points_2d.append(pred_edge.points_2d.astype(np.float32, copy=False))
                pred_edge_points_3d.append(np.asarray(points_3d, dtype=np.float32))
            if pred_edge_points_2d:
                edge_points_2d = _stack_edge_batches(pred_edge_points_2d)
                edge_points_3d = _stack_edge_batches(pred_edge_points_3d)

            partial_edge = decode_cut_partial_side_edge_from_prediction(
                pred,
                pred_edge_60,
                anchor_xy,
                stride,
                img_w,
                cut_side=cut_side,
            )
            if partial_edge is not None:
                partial_points_3d = [back_project_2d_to_3d(tuple(pt), depth, calib) for pt, depth in zip(partial_edge.points_2d, partial_edge.depths)]
                if all(point is not None for point in partial_points_3d):
                    partial_edge.points_3d = np.asarray(partial_points_3d, dtype=np.float32)
                    visible_face_types = {face_type for face_type, _ in visible_faces}
                    if partial_edge.face_type not in visible_face_types:
                        edge_points_3d, edge_points_2d = _append_edge_batch(edge_points_3d, edge_points_2d, partial_edge)
                        visible_faces = [*visible_faces, (partial_edge.face_type, 1.0)]

        return Decoded3DPrediction(
            corners_3d=np.asarray(corners, dtype=np.float32),
            face_center_2d=(u_face, v_face),
            face_color=FACE_COLORS[best_type],
            visible_face_type=best_type,
            visible_face_types=tuple(face_type for face_type, _ in visible_faces),
            edge_points_2d=None if edge_points_2d is None else np.asarray(edge_points_2d, dtype=np.float32),
            edge_points_3d=None if edge_points_3d is None else np.asarray(edge_points_3d, dtype=np.float32),
            cls_id=cls_id,
        )

    if cls_id in complete_3d_classes:
        corners = reconstruct_3d_box_from_whole((u_whole, v_whole), z_whole, dims_whole, rot_y, calib)
        if corners is None:
            return None
        return Decoded3DPrediction(
            corners_3d=np.asarray(corners, dtype=np.float32),
            face_center_2d=None,
            face_color=None,
            visible_face_type=None,
            visible_face_types=(),
            edge_points_2d=None,
            edge_points_3d=None,
            cls_id=cls_id,
        )

    return None


def draw_3d_box(img, corners_3d, calib: ResizedCalib, face_center_2d=None, face_color=None, edge_points_2d=None, edge_color=(0, 255, 0), thickness=1):
    corners_3d = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]]
    color_front = (0, 0, 255)
    color_back = (255, 0, 0)
    color_side = (255, 255, 0)

    distort_coeffs = calib.get("distort_coeffs", []) if calib is not None else []
    if distort_coeffs is not None and len(distort_coeffs) >= 4:
        edge_points_2d_box = project_3d_box_edges_with_distortion(corners_3d, calib, samples_per_edge=15)
        plot_box3d_on_img_with_distortion(
            img,
            edge_points_2d_box,
            color_front=color_front,
            color_back=color_back,
            color_side=color_side,
            thickness=thickness,
        )
    else:
        corners_2d = project_3d_to_2d(corners_3d, calib)
        if np.any(np.isnan(corners_2d)):
            return img
        plot_box3d_on_img(
            img,
            corners_2d,
            color_front=color_front,
            color_back=color_back,
            color_side=color_side,
            thickness=thickness,
        )

    if face_center_2d is not None and face_color is not None:
        cv2.circle(img, (int(face_center_2d[0]), int(face_center_2d[1])), 2, face_color, -1, cv2.LINE_AA)

    _draw_edge_points(img, edge_points_2d=edge_points_2d, edge_color=edge_color, thickness=thickness)
    return img


def decode_visible_face_yaw_from_prediction(pred_41, pred_edge_60, anchor_xy, stride, face_type, calib: ResizedCalib):
    if pred_edge_60 is None or face_type not in range(4):
        return float("nan")
    decoded = decode_visible_face_edge_from_prediction(pred_edge_60, face_type, anchor_xy, stride)
    points_3d = _decoded_edge_to_points_3d(decoded, calib)
    if points_3d is None:
        return float("nan")
    return edge_points_to_yaw(points_3d, face_type)


def decode_multi_visible_face_yaw_from_prediction(
    pred_41,
    pred_edge_60,
    anchor_xy,
    stride,
    calib,
    fallback_face_type=None,
    score_thr=FACE_VISIBILITY_SCORE_THRESH,
    bbox_xyxy=None,
    img_w=None,
):
    if pred_edge_60 is None:
        if fallback_face_type in range(4):
            return decode_visible_face_yaw_from_prediction(pred_41, pred_edge_60, anchor_xy, stride, fallback_face_type, calib)
        return float("nan")

    inferred_img_w = float(img_w) if img_w 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)

    cut_state, cut_side = _resolve_pred_cut_state_for_decode(pred_41, bbox_xyxy=bbox_xyxy, img_w=inferred_img_w)
    face_edges_3d, face_scores = {}, {}
    for face_type, score in select_pred_visible_faces_for_decode(pred_41, score_thr=score_thr, bbox_xyxy=bbox_xyxy, img_w=inferred_img_w):
        decoded = decode_visible_face_edge_from_prediction(pred_edge_60, face_type, anchor_xy, stride)
        points_3d = _decoded_edge_to_points_3d(decoded, calib)
        if points_3d is None:
            continue
        face_edges_3d[face_type] = points_3d
        face_scores[face_type] = float(score)

    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,
    )
    partial_points_3d = _decoded_edge_to_points_3d(partial_edge, calib)
    if cut_state != CUT_STATE_NORMAL:
        if partial_edge is not None and partial_points_3d is not None:
            return edge_points_to_yaw(partial_points_3d, int(partial_edge.face_type))
        return float("nan")

    if any(face_type in (2, 3) for face_type in face_edges_3d):
        side_face_type = max(
            (face_type for face_type in face_edges_3d if face_type in (2, 3)),
            key=lambda face_type: face_scores.get(face_type, 0.0),
        )
        return edge_points_to_yaw(face_edges_3d[side_face_type], side_face_type)

    if partial_points_3d is not None:
        face_edges_3d[partial_edge.face_type] = partial_points_3d
        face_scores[partial_edge.face_type] = max(face_scores.get(partial_edge.face_type, 0.0), 1.0)

    if len(face_edges_3d) >= 2:
        yaw = visible_face_edges_to_yaw(face_edges_3d, face_scores=face_scores)
        if np.isfinite(yaw):
            return yaw

    if fallback_face_type in range(4):
        return decode_visible_face_yaw_from_prediction(pred_41, pred_edge_60, anchor_xy, stride, fallback_face_type, calib)
    return visible_face_edges_to_yaw(face_edges_3d, face_scores=face_scores)


def _back_project_metric_point(u, v, z, calib: ResizedCalib) -> np.ndarray:
    if calib is not None and z > 0:
        center_3d = back_project_2d_to_3d((u, v), z, calib)
        if center_3d is None:
            x3d, y3d = float("nan"), float("nan")
        else:
            x3d, y3d = center_3d[0], center_3d[1]
    else:
        x3d, y3d = float("nan"), float("nan")
    return np.array([x3d, y3d, z], dtype=np.float32)


def extract_3d_attrs_from_prediction(
    pred_41,
    anchor_xy,
    stride,
    calib: ResizedCalib,
    face_type=None,
    pred_edge_60=None,
) -> Optional[Prediction3DAttrs]:
    pred = pred_41
    rot_y = _decode_yaw_from_prediction(pred)
    dims = pred[27:30].astype(np.float32)

    if face_type is None:
        z = float(pred[24])
        uv_offset = pred[25:27]
        edge_yaw = float("nan")
    else:
        off = FACE_OFFSETS_41[face_type]
        z = float(pred[off])
        uv_offset = pred[off + 1 : off + 3]
        edge_yaw = decode_multi_visible_face_yaw_from_prediction(
            pred,
            pred_edge_60,
            anchor_xy,
            stride,
            calib,
            fallback_face_type=face_type,
        )

    u = float((anchor_xy[0] + uv_offset[0]) * stride)
    v = float((anchor_xy[1] + uv_offset[1]) * stride)
    center = _back_project_metric_point(u, v, z, calib)
    return Prediction3DAttrs(
        center=center,
        depth=z,
        dims=dims,
        yaw=float(rot_y),
        edge_yaw=float(edge_yaw),
        uv=np.array([u, v], dtype=np.float32),
        visible_face_type=None if face_type is None else int(face_type),
        face_center=None if face_type is None else center,
    )


def face_center_from_corners(corners_3d, face_type):
    if corners_3d is None or face_type not in FACE_CORNERS:
        return None
    corners = np.asarray(corners_3d, dtype=np.float32)
    if corners.shape != (8, 3) or not np.isfinite(corners).all():
        return None
    return corners[list(FACE_CORNERS[face_type])].mean(axis=0)


def rebuild_box_corners_for_visualization(
    corners_3d,
    dims,
    yaw,
    visible_face_type=None,
    face_center_3d=None,
    box_center_3d=None,
):
    dims_arr = np.asarray(dims, dtype=np.float32)
    if dims_arr.shape != (3,) or not np.isfinite(dims_arr).all() or not np.isfinite(float(yaw)):
        return None

    if visible_face_type is not None:
        if face_center_3d is None:
            face_center_3d = face_center_from_corners(corners_3d, int(visible_face_type))
        else:
            face_center_3d = np.asarray(face_center_3d, dtype=np.float32)
        if face_center_3d is None or face_center_3d.shape != (3,) or not np.isfinite(face_center_3d).all():
            return None
        return compute_3d_box_corners(face_center_3d, dims_arr, float(yaw), face_type=int(visible_face_type))

    if box_center_3d is not None:
        box_center_3d = np.asarray(box_center_3d, dtype=np.float32)
        if box_center_3d.shape != (3,) or not np.isfinite(box_center_3d).all():
            return None
        return compute_3d_box_corners(box_center_3d, dims_arr, float(yaw), face_type=-1)

    corners = np.asarray(corners_3d, dtype=np.float32)
    if corners.shape != (8, 3) or not np.isfinite(corners).all():
        return None
    return compute_3d_box_corners(corners.mean(axis=0), dims_arr, float(yaw), face_type=-1)
