stateplay/utils/state.py

"""State loading and denormalization for StatePlay inference."""
from pathlib import Path
from typing import Sequence

import math

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F

STATE_COLUMNS = ("timer", "hp1", "hp2", "meter1", "meter2")
STATE_NORM_MAX = torch.tensor([99.0, 160.0, 160.0, 104.0, 96.0], dtype=torch.float32)
_ALIASES = {
    "timer": ("timer",),
    "hp1": ("hp1", "p1_hp"),
    "hp2": ("hp2", "p2_hp"),
    "meter1": ("meter1", "p1_meter"),
    "meter2": ("meter2", "p2_meter"),
}


def _resolve_columns(df: pd.DataFrame, columns: Sequence[str]) -> list[str]:
    available = set(df.columns)
    resolved = []
    missing = []
    for col in columns:
        aliases = _ALIASES.get(col, (col,))
        hit = next((name for name in aliases if name in available), None)
        if hit is None:
            missing.append((col, aliases))
        else:
            resolved.append(hit)
    if missing:
        details = "; ".join(
            f"{col} (accepted: {', '.join(aliases)})" for col, aliases in missing
        )
        raise KeyError(
            "Required StatePlay state columns are missing: "
            f"{details}. Available columns: {sorted(df.columns.tolist())}"
        )
    return resolved


def normalize_state(raw_state: torch.Tensor, norm_max: torch.Tensor = STATE_NORM_MAX) -> torch.Tensor:
    norm = norm_max.to(device=raw_state.device, dtype=raw_state.dtype)
    state = torch.minimum(torch.clamp(raw_state, min=0.0), norm) / norm
    return state * 2.0 - 1.0


def denormalize_state(norm_state: torch.Tensor, norm_max: torch.Tensor = STATE_NORM_MAX) -> torch.Tensor:
    norm = norm_max.to(device=norm_state.device, dtype=norm_state.dtype)
    return (norm_state + 1.0) * 0.5 * norm


def load_initial_state(
    parquet_path: str,
    num_frames: int,
    latent_frames: int,
    sampling: str = "interpolation",
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    columns: Sequence[str] = STATE_COLUMNS,
) -> torch.Tensor:
    """Return the first normalized latent-state token as [1, 1, 5]."""
    df = pd.read_parquet(parquet_path)
    cols = _resolve_columns(df, columns)
    arr = df[cols].values[:num_frames].astype(np.float32)
    if arr.shape[0] == 0:
        raise ValueError(f"{parquet_path} contains no state rows.")
    state = normalize_state(torch.from_numpy(arr))
    if sampling == "interpolation":
        initial = state[:1]
    elif sampling == "end":
        end_idx = math.ceil(state.shape[0] / latent_frames) - 1
        initial = state[end_idx:end_idx + 1]
    elif sampling == "avg":
        end_idx = math.ceil(state.shape[0] / latent_frames)
        initial = state[:end_idx].mean(dim=0, keepdim=True)
    else:
        raise ValueError(f"unsupported state sampling: {sampling!r}")
    return initial.unsqueeze(0).to(device=device, dtype=dtype)


def downsample_state_to_latent(
    state: torch.Tensor,
    latent_frames: int,
    sampling: str = "interpolation",
) -> torch.Tensor:
    """Downsample normalized frame-level state [B, T, 5] to latent frames [B, F, 5]."""
    if state.ndim == 4 and state.shape[1] == 1:
        state = state[:, 0]
    if state.ndim != 3 or state.shape[-1] != len(STATE_COLUMNS):
        raise ValueError(
            f"expected state shape [B, T, {len(STATE_COLUMNS)}] or [B, 1, T, {len(STATE_COLUMNS)}], "
            f"got {tuple(state.shape)}"
        )
    if state.shape[1] == latent_frames:
        return state
    if sampling == "interpolation":
        return F.interpolate(
            state.transpose(1, 2).float(), size=latent_frames, mode="linear", align_corners=True,
        ).transpose(1, 2).to(dtype=state.dtype)
    if sampling == "avg":
        return F.adaptive_avg_pool1d(
            state.transpose(1, 2).float(), output_size=latent_frames,
        ).transpose(1, 2).to(dtype=state.dtype)
    if sampling == "end":
        total_frames = state.shape[1]
        end_idx = torch.ceil(
            torch.arange(1, latent_frames + 1, device=state.device, dtype=torch.float32)
            * total_frames
            / latent_frames
        ).long() - 1
        end_idx = end_idx.clamp(min=0, max=total_frames - 1)
        return state.index_select(1, end_idx)
    raise ValueError(f"unsupported state sampling: {sampling!r}")


def load_true_state(
    parquet_path: str,
    num_frames: int,
    latent_frames: int,
    sampling: str = "interpolation",
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    columns: Sequence[str] = STATE_COLUMNS,
) -> torch.Tensor:
    """Return normalized latent-frame ground-truth state as [1, F, 5]."""
    df = pd.read_parquet(parquet_path)
    cols = _resolve_columns(df, columns)
    arr = df[cols].values[:num_frames].astype(np.float32)
    if arr.shape[0] == 0:
        raise ValueError(f"{parquet_path} contains no state rows.")
    state = normalize_state(torch.from_numpy(arr)).unsqueeze(0)
    state = downsample_state_to_latent(state, latent_frames=latent_frames, sampling=sampling)
    return state.to(device=device, dtype=dtype)


def save_state_txt(
    state: torch.Tensor,
    path: str,
    columns: Sequence[str] = STATE_COLUMNS,
    true_state: torch.Tensor | None = None,
) -> None:
    """Save normalized state after inverse scaling, optionally with true/error columns."""
    pred = denormalize_state(state.detach().float().cpu()).squeeze(0)
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    if true_state is not None:
        true = denormalize_state(true_state.detach().float().cpu()).squeeze(0)
        if true.shape != pred.shape:
            raise ValueError(f"true_state shape {tuple(true.shape)} does not match pred shape {tuple(pred.shape)}")
        err = (pred - true).abs()
        with out.open("w", encoding="utf-8") as f:
            f.write(f"{'frame':>5}")
            for col in columns:
                f.write(f" {col + '_pred':>12} {col + '_true':>12} {col + '_err':>12}")
            f.write("\n")
            for i in range(pred.shape[0]):
                f.write(f"{i:5d}")
                for j in range(len(columns)):
                    f.write(
                        f" {pred[i, j].item():12.3f}"
                        f" {true[i, j].item():12.3f}"
                        f" {err[i, j].item():12.3f}"
                    )
                f.write("\n")
        return

    raw = pred.numpy()
    header = "{:>12} ".format("latent_frame") + " ".join("{:>12}".format(col) for col in columns)
    with out.open("w") as f:
        f.write(header + "\n")
        for i, row in enumerate(raw):
            values = " ".join("{:12.6f}".format(float(x)) for x in row)
            f.write("{:12d} {}\n".format(i, values))