"""Decoding and colouring of the `depth` input.

A depth map is a single-channel distance image, not a picture: it is turned into
something displayable by mapping a distance range onto a colormap. Decoding
normalises to float32 millimetres with 0 for missing, so the colouring only ever
sees one representation.
"""

import os

import cv2
import numpy as np

from dora_image_view.decode import is_image_struct
from dora_image_view.payload import as_float32

# Perceptually uniform and readable in both directions, unlike JET.
COLORMAP = cv2.COLORMAP_TURBO

MISSING = 0.0

# Single-channel depth encodings and the pixel type each names. The ROS spellings
# are the ones senders use; a depth map is an image whose pixels are distances,
# which is why it arrives as std/media/v1/Image rather than a type of its own --
# dora's registry has no depth type.
DEPTH_DTYPES = {
    "32fc1": np.float32,
    "float32": np.float32,
    "16uc1": np.uint16,
    "mono16": np.uint16,
    "uint16": np.uint16,
}


def to_millimetres(storage, metadata: dict) -> np.ndarray:
    """Decode one `depth` payload into a float32 array of millimetres.

    Missing depth is 0. Non-finite values are normalised to 0 as well, so
    everything downstream only has to test for it.
    """
    depth, height, width = _values(storage, metadata)

    if depth.size < height * width:
        raise RuntimeError(
            f"depth payload holds {depth.size} values, "
            f"expected at least {height * width} for the declared shape",
        )

    depth = depth[: height * width].reshape(height, width).copy()
    depth[~np.isfinite(depth)] = MISSING

    if units(metadata) == "m":
        depth *= 1000.0

    return depth


def colorize(depth_mm: np.ndarray) -> np.ndarray:
    """Map millimetres onto the colormap, leaving missing depth black."""
    low, high = resolve_range(depth_mm)

    normalised = np.clip((depth_mm - low) / (high - low), 0.0, 1.0)
    frame = cv2.applyColorMap((normalised * 255).astype(np.uint8), COLORMAP)
    frame[depth_mm <= MISSING] = 0
    return frame


def resolve_range(depth_mm: np.ndarray) -> tuple[float, float]:
    """Pick the millimetre range to map onto the colormap.

    Each end is taken from the environment if it is set there and auto-ranged
    over the valid pixels otherwise, so pinning one end does not force the
    other. Auto-ranging re-reads the frame every time, which is what makes an
    unconfigured sensor legible and also what makes the colours shift between
    frames.

    The auto range is the full extent of the valid pixels, not a percentile of
    them. Trimming the tails reads as the safer choice but breaks the scene
    this viewer exists for: a small object in front of a large flat wall is a
    few percent of the pixels, so both tails land on the wall, the range
    collapses, and the object -- the thing being looked at -- clips to one
    colour with everything else.
    """
    low = _env("VIEW_DEPTH_MIN_MM")
    high = _env("VIEW_DEPTH_MAX_MM")

    if low is None or high is None:
        valid = depth_mm[depth_mm > MISSING]
        auto_low, auto_high = (
            (float(valid.min()), float(valid.max())) if valid.size else (0.0, 1.0)
        )
        low = auto_low if low is None else low
        high = auto_high if high is None else high

    # An empty or single-valued frame would otherwise divide by zero.
    if high <= low:
        high = low + 1.0

    return float(low), float(high)


def _values(storage, metadata: dict) -> tuple[np.ndarray, int, int]:
    """Read the payload of either shape into float32 values plus its shape."""
    if is_image_struct(storage):
        row = storage[0].as_py()
        encoding = str(row["encoding"]).lower()
        dtype = DEPTH_DTYPES.get(encoding)

        if dtype is None:
            raise RuntimeError(
                f"unsupported depth encoding: {encoding}; "
                f"expected one of {', '.join(sorted(DEPTH_DTYPES))}",
            )

        values = np.frombuffer(row["data"], dtype=dtype).astype(np.float32)
        return values, int(row["height"]), int(row["width"])

    values = as_float32(storage)
    height, width = _hw(metadata, values.size)
    return values, height, width


def _hw(metadata: dict, size: int) -> tuple[int, int]:
    shape = metadata.get("shape")

    if shape is not None and len(shape) >= 2:
        return int(shape[0]), int(shape[1])

    height, width = metadata.get("height"), metadata.get("width")
    if height and width:
        return int(height), int(width)

    raise RuntimeError(
        "depth needs a `shape` [height, width] or `height`/`width` metadata; "
        f"got a payload of {size} values and neither",
    )


def units(metadata: dict) -> str:
    """Read the distance unit the sender declared, defaulting to millimetres."""
    return str(metadata.get("units", "mm")).lower()


def _env(name: str) -> float | None:
    value = os.getenv(name)
    return float(value) if value else None