from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional, TypedDict

import numpy as np


ColorBGR = tuple[int, int, int]
ImagePoint2D = tuple[float, float]


class RawCameraCalib(TypedDict, total=False):
    focal_u: float
    focal_v: float
    cu: float
    cv: float
    roll: float
    pitch: float
    yaw: float
    pos: list[float]
    distort_coeffs: list[float]
    image_width: Any
    image_height: Any
    source_format: str
    angle_unit: str


class ROICropCalib(TypedDict):
    focal_u: float
    focal_v: float
    cu: float
    cv: float
    src_w: int
    src_h: int
    distort_coeffs: list[float]


class ResizedCalib(TypedDict):
    fx: float
    fy: float
    cx: float
    cy: float
    distort_coeffs: list[float]
    depth_scale: float


@dataclass
class DecodedVisibleEdge:
    face_type: int
    points_2d: np.ndarray
    depths: np.ndarray
    points_3d: Optional[np.ndarray] = None


@dataclass
class Decoded3DPrediction:
    corners_3d: np.ndarray
    face_center_2d: Optional[ImagePoint2D]
    face_color: Optional[ColorBGR]
    visible_face_type: Optional[int]
    visible_face_types: tuple[int, ...]
    edge_points_2d: Optional[np.ndarray]
    edge_points_3d: Optional[np.ndarray]
    cls_id: int


@dataclass
class Prediction3DAttrs:
    center: np.ndarray
    depth: float
    dims: np.ndarray
    yaw: float
    edge_yaw: float
    uv: np.ndarray
    visible_face_type: Optional[int]
    face_center: Optional[np.ndarray]


class SerializedPredictionRecord(TypedDict, total=False):
    bbox_xyxy: Any
    confidence: float
    cls_id: int
    cls_name: str
    difficulty_logit: Optional[float]
    difficulty_prob: Optional[float]
    difficulty_label: Optional[int]
    difficulty_name: Optional[str]
    edge_head_available: bool
    xyzlhwyaw: Any
    xyzlhwyaw_ego: Any
    box_center_xyz: Any
    box_center_xyz_ego: Any
    depth_m: Optional[float]
    box_depth_m: Optional[float]
    lateral_distance_m: Optional[float]
    euclidean_distance_m: Optional[float]
    xz_distance_m: Optional[float]
    attribute: Any
    yaw_rad: Optional[float]
    edge_yaw_rad: Optional[float]
    edge_yaw_confident: bool
    edge_yaw_lateral_distance_m: Optional[float]
    edge_yaw_lateral_ok: bool
    edge_yaw_two_face_eligible: bool
    edge_yaw_selected_face_types: Any
    edge_yaw_selected_face_is_partial: Any
    edge_vs_reg_yaw_rad: Optional[float]
    selected_edge_direct_box_fit_available: bool
    selected_edge_direct_box_fit_mean_px: Optional[float]
    selected_edge_direct_box_fit_max_px: Optional[float]
    selected_edge_direct_box_fit_per_face_mean_px: Any
    selected_edge_edgeyaw_box_fit_available: bool
    selected_edge_edgeyaw_box_fit_mean_px: Optional[float]
    selected_edge_edgeyaw_box_fit_max_px: Optional[float]
    selected_edge_edgeyaw_box_fit_per_face_mean_px: Any
    selected_edge_fit_gain_px: Optional[float]
    edge_box_center_3d: Any
    edge_box_dims: Any
    edge_box_length_m: Optional[float]
    edge_box_width_m: Optional[float]
    edge_box_mode: Optional[str]
    edge_box_length_source: Optional[str]
    edge_box_width_source: Optional[str]
    all_edge_predictions: Any
    edge_selection: Any
    center_uv: Any
    center_3d: Any
    dims: Any
    cut_cls: Optional[int]
    roi_id: Optional[int]
    visible_face_type: Any
    visible_face_count: int
    visible_face_types: Any
    crop_bounds: list[int]
    original_bbox_xyxy: Any


class SerializedROIPayload(TypedDict):
    crop_bounds: list[int]
    vp_x: float
    vp_y: float
    crop_center_x: float
    crop_center_y: float
    edge_head_available: bool
    edge_yaw_max_lateral_dist_m: float
    calib: dict[str, Any]
    predictions: list[SerializedPredictionRecord]


class SerializedMergedPayload(TypedDict, total=False):
    method: str
    edge_head_available: bool
    roi_bounds: dict[str, list[int]]
    predictions: list[SerializedPredictionRecord]
    visualization: str


class SerializedFramePayload(TypedDict, total=False):
    frame_index: int
    frame_name: str
    rois: dict[str, SerializedROIPayload]
    merged: SerializedMergedPayload
    merged_vru: SerializedMergedPayload
    visualization: str


class SerializedPredictionsPayload(TypedDict):
    case_name: str
    images_dir: str
    calib_file: str
    exported_model_path: str
    edge_head_available: bool
    edge_yaw_max_lateral_dist_m: float
    frames: list[SerializedFramePayload]
