"""PoseConv3D Preprocessing: (T, 17, 2) keypoints -> (1, 32, 64, 64) heatmaps."""
import numpy as np

def resample_keypoints(kp_seq: np.ndarray, target_t: int = 32) -> np.ndarray:
    """(L, J, 2) in [-1, 1] range -> (target_t, J, 2) linearly interpolated over time."""
    L, J, D = kp_seq.shape
    if L == target_t:
        return kp_seq.copy()
    src = np.arange(L)
    dst = np.linspace(0, L - 1, target_t)
    out = np.zeros((target_t, J, D), dtype=np.float32)
    for j in range(J):
        for d in range(D):
            col = kp_seq[:, j, d]
            bad = np.isnan(col)
            if bad.all():
                out[:, j, d] = np.nan
            elif bad.any():
                col_clean = np.interp(src, src[~bad], col[~bad])
                out[:, j, d] = np.interp(dst, src, col_clean)
            else:
                out[:, j, d] = np.interp(dst, src, col)
    return out

def rasterize_heatmaps(kp: np.ndarray, h: int = 64, w: int = 64,
                       sigma: float = 1.8, format: str = "pytorch") -> np.ndarray:
    """(target_t, J, 2) keypoints -> (1, target_t, h, w) for PyTorch or (target_t, h, w, 1) for MLX."""
    T, J, _ = kp.shape
    y_grid, x_grid = np.ogrid[:h, :w]
    inv_2s2 = 1.0 / (2.0 * sigma * sigma)

    cx = (kp[..., 0] + 1.0) * 0.5 * (w - 1)
    cy = (kp[..., 1] + 1.0) * 0.5 * (h - 1)

    dx = x_grid[None, None, :, :] - cx[:, :, None, None]
    dy = y_grid[None, None, :, :] - cy[:, :, None, None]
    dist_sq = dx * dx + dy * dy
    g = np.exp(-dist_sq * inv_2s2)
    g = np.nan_to_num(g, nan=0.0)
    vol = np.max(g, axis=1, keepdims=True)  # (T, 1, H, W)
    if format.lower() == "mlx":
        return np.transpose(vol, (0, 2, 3, 1)).astype(np.float32)  # (T, H, W, 1)
    return np.transpose(vol, (1, 0, 2, 3)).astype(np.float32)      # (1, T, H, W)