# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Configuration for the camera sensor."""


from dataclasses import MISSING
from typing import List, Tuple

from omni.isaac.orbit.utils import configclass

from ..sensor_base_cfg import SensorBaseCfg

__all__ = ["PinholeCameraCfg", "FisheyeCameraCfg"]


@configclass
class PinholeCameraCfg(SensorBaseCfg):
    """Configuration for a pinhole camera sensor."""

    @configclass
    class UsdCameraCfg:
        """USD related configuration of the sensor.

        The parameter is kept default from USD if it is set to :obj:`None`. This includes the default
        parameters (in case the sensor is created) or the ones set by the user (in case the sensor is
        loaded from existing USD stage).

        Reference:
            * https://docs.omniverse.nvidia.com/prod_extensions/prod_extensions/ext_replicator/annotators_details.html
            * https://graphics.pixar.com/usd/docs/api/class_usd_geom_camera.html
        """

        clipping_range: Tuple[float, float] = None
        """Near and far clipping distances (in stage units)."""
        focal_length: float = None
        """Perspective focal length (in mm). Longer lens lengths narrower FOV, shorter lens lengths wider FOV."""
        focus_distance: float = None
        """Distance from the camera to the focus plane (in stage units).

        The distance at which perfect sharpness is achieved.
        """
        f_stop: float = None
        """Lens aperture. Defaults to 0.0, which turns off focusing.

        Controls Distance Blurring. Lower Numbers decrease focus range, larger numbers increase it.
        """
        horizontal_aperture: float = None
        """Horizontal aperture (in mm). Emulates sensor/film width on a camera."""
        horizontal_aperture_offset: float = None
        """Offsets Resolution/Film gate horizontally."""
        vertical_aperture_offset: float = None
        """Offsets Resolution/Film gate vertically."""

    data_types: List[str] = ["rgb"]
    """List of sensor names/types to enable for the camera. Defaults to ["rgb"]."""
    width: int = MISSING
    """Width of the image in pixels."""
    height: int = MISSING
    """Height of the image in pixels."""
    semantic_types: List[str] = ["class"]
    """List of allowed semantic types the types. Defaults to ["class"].

    For example, if semantic types is [“class”], only the bounding boxes for prims with semantics of
    type “class” will be retrieved.

    More information available at:
        https://docs.omniverse.nvidia.com/app_code/prod_extensions/ext_replicator/semantic_schema_editor.html
    """
    colorize: bool = False
    """whether to output colorized semantic information or non-colorized one. Defaults to False.

    If True, the semantic images will be a 2D array of RGBA values, where each pixel is colored according to
    the semantic type. Accordingly, the information output will contain mapping from color to semantic labels.

    If False, the semantic images will be a 2D array of integers, where each pixel is an integer representing
    the semantic ID. Accordingly, the information output will contain mapping from semantic ID to semantic labels.
    """
    projection_type: str = "pinhole"
    """Type of projection to use for the camera. Defaults to "pinhole"."""
    usd_params: UsdCameraCfg = UsdCameraCfg()
    """Parameters for setting USD camera settings."""


@configclass
class FisheyeCameraCfg(PinholeCameraCfg):
    """Configuration for a fisheye camera sensor."""

    @configclass
    class UsdCameraCfg(PinholeCameraCfg.UsdCameraCfg):
        """USD related configuration of the sensor for the fisheye model."""

        fisheye_nominal_width: float = None
        """Nominal width of fisheye lens model."""
        fisheye_nominal_height: float = None
        """Nominal height of fisheye lens model."""
        fisheye_optical_centre_x: float = None
        """Horizontal optical centre position of fisheye lens model."""
        fisheye_optical_centre_y: float = None
        """Vertical optical centre position of fisheye lens model."""
        fisheye_max_fov: float = None
        """Maximum field of view of fisheye lens model."""
        fisheye_polynomial_a: float = None
        """First component of fisheye polynomial."""
        fisheye_polynomial_b: float = None
        """Second component of fisheye polynomial."""
        fisheye_polynomial_c: float = None
        """Third component of fisheye polynomial."""
        fisheye_polynomial_d: float = None
        """Fourth component of fisheye polynomial."""
        fisheye_polynomial_e: float = None
        """Fifth component of fisheye polynomial."""

    projection_type: str = "fisheye_polynomial"
    """Type of projection to use for the camera. Defaults to "fisheye_polynomial"."""
    usd_params: UsdCameraCfg = UsdCameraCfg()
    """Parameters for setting USD camera settings."""
