data/depth_encoding/power_transform.py

"""
PyTorch implementation of Barron's power transform for metric depth encoding:

    f(d; lam, c) = 1 - (1 + d / (|lam| * c)) ** (lam + 1)
                 = -expm1( (lam + 1) * log1p( d / (|lam| * c) ) )

The second form is what we actually evaluate: it is numerically stable for
every `lam < -1` and avoids `0 ** 0` / `0 ** neg` pathologies that hit the
naive closed form.

The paper's depth encoding uses `(lam, c) = (-3, 10/3)` so that
`f(d) = 1 - 1 / (1 + d/10)^2`; keep these as defaults.

Valid range for `lam`:  `lam < -1` (strictly).
  * At `lam = -1` the family degenerates (all outputs are 0).
  * For `lam > -1` the image is no longer contained in `[0, 1)`, so the
    depth encoding is ill-defined.

Practical range for `lam`:
  * `lam ~= -3` (the paper's default) gives roughly uniform resolution per
    octave of depth over ~[0.3 m, 100 m].
  * Useful range is roughly `lam in [-10, -1.5]`.  Past `lam ~ -15` the
    transform saturates so aggressively that `y = 2f - 1` cannot resolve
    large depths even in float64 (every `d > c` maps to within one ulp of 1);
    this is a property of the mapping, not of the implementation.

Functions
---------
power_transform(d, lam, c, mask)            : [0, inf)   -> [0, 1)
inverse_power_transform(f, lam, c, ...)     : [0, 1)     -> [0, inf)
depth_to_signed_unit(d, lam, c, mask, ...)  : [0, inf)   -> [-1, 1)
signed_unit_to_depth(y, lam, c, ...)        : [-1, 1)    -> [0, inf)

`lam` and `c` may each be either a Python float, a 0-d tensor, or a tensor
broadcastable against `d` (e.g. a per-sample `lam` of shape `[B, 1, 1, 1]`).

Class
-----
DepthPowerTransform : plain container that bundles `(lam, c, inf_cap,
invalid_fill)` and exposes `encode` / `decode` (also callable via `()`).
It is NOT an `nn.Module` -- this is purely a normalization utility, so it has
no buffers / parameters and doesn't need `.to(device)`.  Device/dtype follow
the input tensor automatically.
"""

from __future__ import annotations

from typing import Optional, Tuple, Union

import torch
from torch import Tensor

# A scalar "shape parameter" may be passed as either a Python float or a
# Tensor (so it can be a learnable `nn.Parameter`, a per-sample value, ...).
ScalarLike = Union[float, Tensor]


# ---------------------------------------------------------------------------
# Small utility: bring a ScalarLike onto the same dtype / device as a ref
# tensor while preserving its shape (so broadcasting Just Works).
# ---------------------------------------------------------------------------

def _as_like(x: ScalarLike, ref: Tensor) -> Tensor:
    if isinstance(x, Tensor):
        return x.to(dtype=ref.dtype, device=ref.device)
    return torch.as_tensor(x, dtype=ref.dtype, device=ref.device)


# ---------------------------------------------------------------------------
# Core forward / inverse
# ---------------------------------------------------------------------------

def power_transform(
    d: Tensor,
    lam: ScalarLike = -3.0,
    c: ScalarLike = 10.0 / 3.0,
    mask: Optional[Tensor] = None,
    invalid_fill: float = 0.0,
) -> Tensor:
    """Map metric depth `d` in `[0, inf)` to normalized `f` in `[0, 1)`.

    Computed as
        f = -expm1( (lam + 1) * log1p( d / (|lam| * c) ) )
    which is stable for any `lam < -1` and well-defined for any `d >= 0`.

    Args:
        d            : depth tensor, any shape (e.g. (1, H, W), (B, 1, H, W)).
        lam, c       : power-transform parameters; float, 0-d tensor, or tensor
                       broadcastable to `d` (lets you learn them or set them
                       per-sample).
        mask         : optional bool tensor broadcastable to `d`.  If given,
                       pixels with `mask == False` or `d <= 0` are written as
                       `invalid_fill` in the output.  If `None`, we just clamp
                       `d` to be non-negative.
        invalid_fill : scalar value written at invalid pixels (default 0.0).

    Returns:
        f : same shape as `d`, values in `[0, 1)`.
    """
    lam_t = _as_like(lam, d)
    c_t   = _as_like(c,   d)

    d_safe = d.clamp(min=0.0)
    # |lam| * c is always > 0 in the regime we care about (lam < -1, c > 0),
    # but we still guard against accidental zero.
    denom = lam_t.abs() * c_t
    denom = denom.clamp(min=torch.finfo(d.dtype).tiny)

    f = -torch.expm1((lam_t + 1.0) * torch.log1p(d_safe / denom))

    if mask is not None:
        valid = mask.to(torch.bool) & (d > 0)
        f = torch.where(valid, f, torch.full_like(f, invalid_fill))
    return f


def inverse_power_transform(
    f: Tensor,
    lam: ScalarLike = -3.0,
    c: ScalarLike = 10.0 / 3.0,
    inf_cap: Optional[float] = None,
    _one_minus_f: Optional[Tensor] = None,
) -> Tensor:
    """Map normalized `f` in `[0, 1)` back to metric depth `d` in `[0, inf)`.

    Computed as
        d = |lam| * c * expm1( log(1 - f) / (lam + 1) )

    `_one_minus_f` is an optional precomputed `(1 - f)` supplied by upstream
    callers that already have it in higher precision.  This matters because
    when `f` is close to 1, computing `1 - f` directly in float32 throws away
    much of the precision in `f`; `signed_unit_to_depth` passes `(1 - y)/2`
    through this path.

    At `f = 1` the depth is `+inf`; pass an `inf_cap` (e.g. 1e4) to keep
    gradients finite during VAE training.
    """
    lam_t = _as_like(lam, f)
    c_t   = _as_like(c,   f)

    if _one_minus_f is None:
        one_minus_f = (1.0 - f).clamp(min=0.0, max=1.0)
    else:
        one_minus_f = _one_minus_f.clamp(min=0.0, max=1.0)

    # Avoid exact zero so `log(0) = -inf -> exp(-inf) = 0` is a finite path.
    eps = torch.finfo(one_minus_f.dtype).tiny
    one_minus_f = one_minus_f.clamp(min=eps)

    denom_lam = (lam_t + 1.0)
    # Guard against lam_t -> -1 producing a 0-division at inference time.
    denom_lam = torch.where(
        denom_lam.abs() < eps,
        torch.full_like(denom_lam, -eps),
        denom_lam,
    )

    u = torch.expm1(torch.log(one_minus_f) / denom_lam)
    d = lam_t.abs() * c_t * u

    if inf_cap is not None:
        d = d.clamp(max=inf_cap)
    return d


# ---------------------------------------------------------------------------
# VAE-friendly: [0, inf) <-> [-1, 1) with optional mask
# ---------------------------------------------------------------------------

def depth_to_signed_unit(
    d: Tensor,
    lam: ScalarLike = -3.0,
    c: ScalarLike = 10.0 / 3.0,
    mask: Optional[Tensor] = None,
    invalid_fill: float = 0.0,
    return_mask: bool = False,
) -> Union[Tensor, Tuple[Tensor, Tensor]]:
    """Encode metric depth in `[0, inf)` to `y` in `[-1, 1)` via

        y = 2 * power_transform(d, lam, c) - 1

    Invalid pixels (either `mask == False` or `d <= 0`) are filled with
    `invalid_fill`.  Using `invalid_fill = 0` (the neutral centre of
    `[-1, 1]`) and passing the mask to the loss separately keeps real
    `d = 0` pixels (which encode to `y = -1`) distinct from "unknown"
    pixels.

    Args:
        d            : depth tensor, any shape.
        lam, c       : scalars or tensors broadcastable to `d`.
        mask         : optional bool tensor, same broadcastable shape as `d`.
                       True = valid, False = invalid.  If None, every `d > 0`
                       is considered valid.
        invalid_fill : value written at invalid pixels.
        return_mask  : if True, also return the effective valid-mask.

    Returns:
        y            : same shape as `d`, in `[-1, 1)`.
        valid (opt.) : bool tensor of same shape as `d`.
    """
    f = power_transform(d, lam=lam, c=c)                # raw [0,1)
    y = 2.0 * f - 1.0

    d_valid = d > 0
    valid = d_valid if mask is None else (mask.to(torch.bool) & d_valid)
    y = torch.where(valid, y, torch.full_like(y, invalid_fill))

    if return_mask:
        return y, valid
    return y


def signed_unit_to_depth(
    y: Tensor,
    lam: ScalarLike = -3.0,
    c: ScalarLike = 10.0 / 3.0,
    inf_cap: Optional[float] = 1e4,
) -> Tensor:
    """Decode `y` in `[-1, 1)` back to metric depth `d` in `[0, inf)`.

    Internally uses `(1 - f) = (1 - y) / 2` so that no precision is lost when
    `y` is close to 1 (large depth).
    """
    y_c = y.clamp(min=-1.0, max=1.0)
    one_minus_f = 0.5 * (1.0 - y_c)
    f = 1.0 - one_minus_f                               # for API consistency
    return inverse_power_transform(
        f, lam=lam, c=c, inf_cap=inf_cap, _one_minus_f=one_minus_f
    )


# ---------------------------------------------------------------------------
# Class wrapper: define once, use everywhere
# ---------------------------------------------------------------------------

class DepthPowerTransform:
    """Plain container bundling the power-transform depth codec's settings.

    Pure normalization utility -- no buffers, no parameters, not an `nn.Module`.
    The stored `lam` / `c` are Python floats; the actual forward / inverse
    computations happen entirely inside `depth_to_signed_unit` /
    `signed_unit_to_depth`, and device & dtype are picked up automatically from
    whatever tensor you pass in.

    Example
    -------
    >>> codec = DepthPowerTransform()                     # lam=-3, c=10/3
    >>> y          = codec.encode(depth, mask=mask)       # [0, inf) -> [-1, 1)
    >>> y          = codec(depth, mask=mask)              # __call__ == encode
    >>> y, valid   = codec(depth, mask=mask, return_mask=True)
    >>> d_hat      = codec.decode(y_pred)                 # [-1, 1) -> [0, inf)
    """

    def __init__(
        self,
        lam: float = -3.0,
        c: float = 10.0 / 3.0,
        inf_cap: Optional[float] = 1e4,
        invalid_fill: float = 0.0,
    ):
        if lam >= -1.0:
            raise ValueError(f"lam must be < -1, got {lam}")
        if c <= 0.0:
            raise ValueError(f"c must be > 0, got {c}")

        self.lam = float(lam)
        self.c = float(c)
        self.inf_cap = inf_cap
        self.invalid_fill = float(invalid_fill)

    def encode(
        self,
        d: Tensor,
        mask: Optional[Tensor] = None,
        return_mask: bool = False,
    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
        """Depth in `[0, inf)` -> y in `[-1, 1)`.  See `depth_to_signed_unit`."""
        return depth_to_signed_unit(
            d,
            lam=self.lam,
            c=self.c,
            mask=mask,
            invalid_fill=self.invalid_fill,
            return_mask=return_mask,
        )

    def decode(self, y: Tensor) -> Tensor:
        """y in `[-1, 1)` -> depth in `[0, inf)`.  See `signed_unit_to_depth`."""
        return signed_unit_to_depth(
            y,
            lam=self.lam,
            c=self.c,
            inf_cap=self.inf_cap,
        )

    def __call__(
        self,
        d: Tensor,
        valid_mask: Optional[Tensor] = None,
        return_mask: bool = False,
    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
        return self.encode(d, mask=valid_mask, return_mask=return_mask)

    def __repr__(self) -> str:
        return (f"DepthPowerTransform(lam={self.lam:.4f}, c={self.c:.4f}, "
                f"inf_cap={self.inf_cap}, invalid_fill={self.invalid_fill})")


# ---------------------------------------------------------------------------
# Self-test: round-trip, gradient w.r.t. lam, learnable-lam example
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    torch.manual_seed(0)

    # 1) round-trip on a random (1, H, W) depth map with a mask
    d = torch.rand(1, 32, 32) * 100.0
    mask = d > 0.5
    y, valid = depth_to_signed_unit(d, mask=mask, return_mask=True)
    d_hat = signed_unit_to_depth(y)
    print(f"shape              : {tuple(d.shape)}")
    print(f"y range            : [{y.min().item():+.4f}, {y.max().item():+.4f}]")
    print(f"y at invalid pixels: {y[~valid].unique().tolist()}   (expect [0.0])")
    print(f"round-trip max err : {(d - d_hat)[valid].abs().max():.2e} m")

    # 2) match the legacy closed form at (lam, c) = (-3, 10/3) for sanity
    d2 = torch.linspace(0.0, 200.0, 500, dtype=torch.float64)
    f_legacy = 1.0 - 1.0 / (1.0 + d2 / 10.0) ** 2
    f_new    = power_transform(d2, lam=-3.0, c=10.0 / 3.0)
    print(f"\nmax|f_new - f_legacy| = {(f_new - f_legacy).abs().max():.2e}")

    # 3) vary `lam` freely: round-trip for several values (float64 for headroom)
    print("\nround-trip for varying lam (float64, should all be ~0):")
    for lam in [-1.5, -2.0, -3.0, -5.0, -10.0]:
        d3  = torch.rand(1, 64, 64, dtype=torch.float64) * 200.0
        y3  = depth_to_signed_unit(d3, lam=lam, c=10.0 / 3.0)
        d3h = signed_unit_to_depth(y3, lam=lam, c=10.0 / 3.0)
        err = (d3 - d3h).abs().max().item()
        print(f"   lam = {lam:>6.2f}   max|d - d_hat| = {err:.2e} m")

    # 4) per-sample `lam`: different shape parameter for each image in batch
    d4 = torch.rand(4, 1, 32, 32, dtype=torch.float64) * 100.0
    lam4 = torch.tensor([-1.5, -3.0, -5.0, -10.0]).view(4, 1, 1, 1)
    y4 = depth_to_signed_unit(d4, lam=lam4)
    d4h = signed_unit_to_depth(y4, lam=lam4)
    print(f"\nper-sample lam round-trip max err: "
          f"{(d4 - d4h).abs().max():.2e} m")

    # 5) `lam` as a learnable nn.Parameter: gradient flows through it
    lam_p = torch.nn.Parameter(torch.tensor(-3.0))
    c_p   = torch.nn.Parameter(torch.tensor(10.0 / 3.0))
    d5    = (torch.rand(2, 1, 16, 16) * 50.0).requires_grad_(True)
    y5    = depth_to_signed_unit(d5, lam=lam_p, c=c_p)
    loss  = (y5 - 0.1).pow(2).mean()
    loss.backward()
    print(f"\nlearnable lam  |  loss = {loss.item():.4f}"
          f"   d/dlam = {lam_p.grad.item():+.4f}"
          f"   d/dc = {c_p.grad.item():+.4f}"
          f"   max|d/dd| = {d5.grad.abs().max():.4f}")

    # 6) DepthPowerTransform class: define once, use everywhere
    print("\n--- DepthPowerTransform class ---")
    codec = DepthPowerTransform()
    print(codec)
    d6 = torch.rand(2, 1, 48, 48) * 80.0
    m6 = d6 > 0.2
    y6, v6 = codec(d6, mask=m6, return_mask=True)         # __call__ == encode
    d6h = codec.decode(y6)
    print(f"  round-trip err   : {(d6 - d6h)[v6].abs().max():.2e} m")

    # Non-default parameters.
    codec_b = DepthPowerTransform(lam=-4.0, c=5.0, inf_cap=1e3, invalid_fill=-1.0)
    print(codec_b)

    # Device/dtype follow the input tensor; no `.to()` / `.double()` needed.
    d8 = torch.rand(1, 64, 64, dtype=torch.float64) * 200.0
    y8 = codec(d8)
    print(f"  float64 in/out   : y.dtype={y8.dtype}, "
          f"round-trip = {(d8 - codec.decode(y8)).abs().max():.2e} m")