"""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

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

MISSING = 0.0


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 = as_float32(storage)
    height, width = _hw(metadata, depth.size)

    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 as_float32(storage) -> np.ndarray:
    """Read the payload as float32, whether it arrived as floats or as bytes.

    A node that sends the map as a raw byte slice -- dora-zed does -- produces a
    UInt8 arrow array even though its manifest declares Float32, so the buffer
    has to be reinterpreted rather than converted.
    """
    buffer = storage.to_numpy()

    if buffer.dtype == np.uint8:
        return buffer.view(np.float32)

    return buffer.astype(np.float32, copy=False)


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