from typing import TypedDict, TypeAlias, Any
from typing_extensions import NotRequired

from jaxtyping import Num, jaxtyped
from beartype import beartype
from jax import numpy as jnp, Array

CameraID: TypeAlias = str


@jaxtyped(typechecker=beartype)
class CameraParams(TypedDict):
    """
    Camera parameters: intrinsic matrix, extrinsic matrix, and distortion coefficients
    """

    K: Num[Array, "3 3"]
    """
    intrinsic matrix
    """
    Rt: Num[Array, "4 4"]
    """
    [R|t] extrinsic matrix

    R and t are the rotation and translation that describe the change of
    coordinates from world to camera coordinate systems (or camera frame)
    """
    dist_coeffs: Num[Array, "N"]
    """
    An array of distortion coefficients of the form
    [k1, k2, [p1, p2, [k3]]], where ki is the ith
    radial distortion coefficient and pi is the ith
    tangential distortion coeff.
    """


@jaxtyped(typechecker=beartype)
class Camera(TypedDict):
    """
    a description of a camera
    """

    id: CameraID
    """
    Camera ID
    """
    params: CameraParams
    """
    Camera parameters
    """
    size: tuple[int, int]
    """
    Image size
    """