"""Per-timestep marginal heatmap panels.

Two functions:
- :func:`marginal_heatmap_panel` — one timestep, blended over the live
  image at the live resolution.
- :func:`marginal_heatmap_grid` — a row of T cells at the heatmap's
  native resolution, each cell = colormap blended over a downsampled
  background. Lets the viewer see the predicted spatial distribution
  across the whole window at a glance.

Shared between training (wandb) and inference (rerun).
"""
from __future__ import annotations

import cv2
import numpy as np


def marginal_heatmap_panel(
    img_bgr: np.ndarray,                  # (H, W, 3) uint8 background
    logits_tzhw: np.ndarray,              # (T, Z, P, P) OR (T, P, P)
    timestep: int = 0,                    # which T to render
    alpha: float = 0.5,
) -> np.ndarray:
    """Return uint8 BGR panel with heatmap blended over ``img_bgr``."""
    H, W = img_bgr.shape[:2]
    x = logits_tzhw.astype(np.float32)
    if x.ndim == 4:
        T, Z, P, _ = x.shape
        # Joint softmax over (Z, P, P), then marginalize Z.
        flat = x.reshape(T, -1)
        flat = flat - flat.max(axis=1, keepdims=True)
        e = np.exp(flat)
        e /= e.sum(axis=1, keepdims=True)
        probs = e.reshape(T, Z, P, P).sum(axis=1)         # (T, P, P)
    else:
        T, P, _ = x.shape
        flat = x.reshape(T, -1)
        flat = flat - flat.max(axis=1, keepdims=True)
        e = np.exp(flat)
        probs = (e / e.sum(axis=1, keepdims=True)).reshape(T, P, P)

    p = probs[timestep]                                    # (P, P)
    p = (p - p.min()) / (p.max() - p.min() + 1e-8)
    p_resized = cv2.resize((p * 255).astype(np.uint8), (W, H),
                           interpolation=cv2.INTER_LINEAR)
    color = cv2.applyColorMap(p_resized, cv2.COLORMAP_JET)
    blend = cv2.addWeighted(color, alpha, img_bgr, 1 - alpha, 0)
    return blend


def marginal_heatmap_grid(
    img_bgr: np.ndarray,                  # (H, W, 3) uint8 background
    logits_tzhw: np.ndarray,              # (T, Z, P, P) OR (T, P, P)
    cell_size: int = 128,
    alpha: float = 0.5,
    pad_px: int = 4,
) -> np.ndarray:
    """Horizontal strip of T heatmap cells at the heatmap's native res.

    Each cell is ``cell_size × cell_size`` (default 128, the model's
    pred grid) with the live RGB downsampled underneath. White gaps
    between cells. Returns ``(cell_size, T * cell_size + (T - 1) * pad, 3)``
    uint8 BGR.
    """
    bg = cv2.resize(img_bgr, (cell_size, cell_size), interpolation=cv2.INTER_AREA)

    x = logits_tzhw.astype(np.float32)
    if x.ndim == 4:
        T, Z, P, _ = x.shape
        flat = x.reshape(T, -1)
        flat = flat - flat.max(axis=1, keepdims=True)
        e = np.exp(flat); e /= e.sum(axis=1, keepdims=True)
        probs = e.reshape(T, Z, P, P).sum(axis=1)
    else:
        T, P, _ = x.shape
        flat = x.reshape(T, -1)
        flat = flat - flat.max(axis=1, keepdims=True)
        e = np.exp(flat); e /= e.sum(axis=1, keepdims=True)
        probs = e.reshape(T, P, P)

    pad = np.full((cell_size, pad_px, 3), 255, dtype=np.uint8)
    cells = []
    for t in range(T):
        p = probs[t]
        p = (p - p.min()) / (p.max() - p.min() + 1e-8)
        p_resized = cv2.resize((p * 255).astype(np.uint8),
                                (cell_size, cell_size),
                                interpolation=cv2.INTER_LINEAR)
        color = cv2.applyColorMap(p_resized, cv2.COLORMAP_JET)
        cell = cv2.addWeighted(color, alpha, bg, 1 - alpha, 0)
        cells.append(cell)
        if t < T - 1:
            cells.append(pad)
    return np.concatenate(cells, axis=1)