"""Compose a rendered model image over a live frame using a mask.

Two styles:
- :func:`compose_render_over_live` — alpha-blends the full mujoco RGB
  render where mask=True. Misalignment shows as ghosting.
- :func:`compose_mask_over_live` — tints the mask region with a solid
  color (default red in BGR). Misalignment shows as a flat shape
  drifting from the real arm.
"""
from __future__ import annotations

import numpy as np


def compose_render_over_live(
    live_rgb: np.ndarray,
    render_rgb: np.ndarray,
    mask: np.ndarray,
    alpha: float = 0.5,
) -> np.ndarray:
    """Return composite uint8 RGB. ``mask`` is (H, W) 0/1; render shape (H, W, 3)."""
    out = live_rgb.copy()
    m = mask.astype(bool)
    out[m] = (alpha * render_rgb[m] + (1 - alpha) * live_rgb[m]).astype(np.uint8)
    return out


def compose_mask_over_live(
    live_rgb: np.ndarray,
    mask: np.ndarray,
    color=(0, 0, 255),    # BGR-order tuple; pass (255, 0, 0) for RGB-input images
    alpha: float = 0.5,
) -> np.ndarray:
    """Tint the mask region with a solid color over the live image.

    Use this for a clean "silhouette" view — easier to read calibration
    alignment than the textured render in some lighting.
    """
    out = live_rgb.copy()
    m = mask.astype(bool)
    c = np.asarray(color, dtype=np.float32)
    out[m] = (alpha * c + (1 - alpha) * out[m]).astype(np.uint8)
    return out