"""Per-(T, bin) heatmap grids for rotation / gripper / height heads.

Each grid is shape ``(T, n_bins)`` — one row per predicted timestep, one
column per bin. Cell color is the softmax prob (jet). Easy to spot at a
glance: GT bin should be a clear bright cell per row.

If ``target`` is provided, the GT cell gets a small black outline.

Used by training (wandb) and inference (rerun).
"""
from __future__ import annotations

from typing import Optional

import cv2
import numpy as np


def bin_grid_panel(
    logits_tk: np.ndarray,                  # (T, K) raw logits
    target: Optional[np.ndarray] = None,    # (T,) int — GT bin per timestep
    cell_h: int = 16,
    cell_w: int = 8,
) -> np.ndarray:
    """Return (T*cell_h, K*cell_w, 3) BGR uint8 heatmap with GT outlined."""
    T, K = logits_tk.shape
    x = logits_tk.astype(np.float32)
    x -= x.max(axis=1, keepdims=True)
    e = np.exp(x); e /= e.sum(axis=1, keepdims=True)
    norm = (e - e.min(axis=1, keepdims=True))
    norm /= (norm.max(axis=1, keepdims=True) - norm.min(axis=1, keepdims=True) + 1e-8) * 1.0
    img_small = (norm * 255).astype(np.uint8)               # (T, K)
    color = cv2.applyColorMap(img_small, cv2.COLORMAP_JET)  # (T, K, 3)
    grid = cv2.resize(color, (K * cell_w, T * cell_h),
                       interpolation=cv2.INTER_NEAREST)
    if target is not None:
        for t in range(T):
            k = int(target[t])
            tl = (k * cell_w, t * cell_h)
            br = ((k + 1) * cell_w - 1, (t + 1) * cell_h - 1)
            cv2.rectangle(grid, tl, br, (0, 0, 0), 1)
    return grid