"""The pinhole camera model a depth map is interpreted through.

The node does not need intrinsics to draw anything -- it displays them, and
holds them for whatever reads them next. Back-projecting a depth map into a
point cloud is the obvious such reader: every pixel needs fx, fy, cx and cy to
become an XYZ point, and nothing else in a dataflow carries them.

Two senders are handled because dora-zed publishes the same numbers twice: once
per camera session on `camera_info`, and again in the metadata of every frame.
The per-frame copy is what makes a viewer that started after the camera did
still end up with intrinsics.
"""

from dataclasses import dataclass

from dora_image_view.payload import as_float32

# [fx, fy, cx, cy], the four numbers dora-zed publishes. Also accepted is a
# nine-element row-major K, which is how ROS spells the same model in
# sensor_msgs/CameraInfo.
COMPACT = 4
K_MATRIX = 9

# Indices of fx, fy, cx, cy in a row-major K:
#   [fx,  0, cx,
#     0, fy, cy,
#     0,  0,  1]
K_INDICES = (0, 4, 2, 5)


@dataclass(frozen=True)
class Intrinsics:
    """A pinhole model: focal lengths and principal point, both in pixels."""

    fx: float
    fy: float
    cx: float
    cy: float
    distortion: tuple[float, ...] = ()
    model: str = ""

    def __str__(self) -> str:
        """Render for the overlay, short enough to fit across a panel."""
        return f"fx {self.fx:.1f} fy {self.fy:.1f} cx {self.cx:.1f} cy {self.cy:.1f}"


def parse(storage, metadata: dict) -> Intrinsics:
    """Read one `camera_info` payload.

    Accepts the four-element [fx, fy, cx, cy] dora-zed sends and the
    nine-element row-major K that ROS does.
    """
    values = as_float32(storage)

    if values.size >= K_MATRIX:
        fx, fy, cx, cy = (float(values[index]) for index in K_INDICES)
    elif values.size >= COMPACT:
        fx, fy, cx, cy = (float(value) for value in values[:COMPACT])
    else:
        raise RuntimeError(
            f"camera_info holds {values.size} values, "
            f"expected {COMPACT} as [fx, fy, cx, cy] or {K_MATRIX} as K",
        )

    return Intrinsics(fx, fy, cx, cy, _distortion(metadata), _model(metadata))


def from_metadata(metadata: dict) -> Intrinsics | None:
    """Read intrinsics a sender attached to a frame, or None if it attached none."""
    try:
        fx, fy, cx, cy = (
            float(metadata["fx"]),
            float(metadata["fy"]),
            float(metadata["cx"]),
            float(metadata["cy"]),
        )
    except (KeyError, TypeError, ValueError):
        return None

    return Intrinsics(fx, fy, cx, cy, _distortion(metadata), _model(metadata))


def _distortion(metadata: dict) -> tuple[float, ...]:
    values = metadata.get("distortion") or ()
    return tuple(float(value) for value in values)


def _model(metadata: dict) -> str:
    return str(metadata.get("distortion_model", ""))