[tab] data/jasmine_data/atari/visualize_array_record.py (python)
import os
import math
import argparse
import pickle
import random
from typing import Optional, List

import numpy as np

from PIL import Image, ImageDraw

from array_record.python.array_record_module import ArrayRecordReader


def infer_hw_from_bytes(total_elements: int, seq_len: int, channels: int) -> tuple[int, int]:
    assert seq_len > 0 and channels > 0, "sequence_length and channels must be positive"
    base = total_elements // (seq_len * channels)
    side = int(math.isqrt(base))
    if side * side != base:
        raise ValueError(
            f"Could not infer square HxW from buffer. elements={total_elements}, seq_len={seq_len}, channels={channels}"
        )
    return side, side


def load_random_record(input_path: str, seed: Optional[int] = None) -> dict:
    assert ArrayRecordReader is not None, "array_record is not installed. pip install array-record"
    reader = ArrayRecordReader(str(input_path))

    # Reservoir sample one record uniformly at random without loading all into memory
    rng = random.Random(seed)
    chosen: Optional[bytes] = None
    count = 0
    for raw in reader:
        count += 1
        if rng.random() < 1.0 / count:
            chosen = raw

    if chosen is None:
        raise RuntimeError(f"No records found in {input_path}")

    return pickle.loads(chosen)


def get_action_meanings(env_id: Optional[str]) -> Optional[List[str]]:
    if env_id is None:
        return None
    try:
        import gymnasium as gym
        import ale_py  # noqa: F401

        gym.register_envs(ale_py)  # ensure Atari envs are registered
        env = gym.make(env_id)
        try:
            meanings = list(env.unwrapped.get_action_meanings())  # type: ignore[attr-defined]
        finally:
            env.close()
        return meanings
    except Exception:
        return None


def save_frames_with_actions(
    record: dict,
    output_dir: str,
    channels: int,
    height: Optional[int],
    width: Optional[int],
    action_meanings: Optional[List[str]],
    fps: int,
    gif_path: Optional[str] = None,
) -> None:
    os.makedirs(output_dir, exist_ok=True)

    seq_len: int = int(record["sequence_length"])
    raw_video: bytes = record["raw_video"]
    actions: Optional[np.ndarray] = record.get("actions")
    breakpoint()
    if actions is not None:
        actions = np.asarray(actions).reshape(-1)

    # Decode frames
    arr = np.frombuffer(raw_video, dtype=np.uint8)
    total_elements = arr.size

    if height is None or width is None:
        h, w = infer_hw_from_bytes(total_elements, seq_len, channels)
    else:
        h, w = int(height), int(width)
        expected = seq_len * h * w * channels
        if expected != total_elements:
            # Fall back to inference if mismatch
            h, w = infer_hw_from_bytes(total_elements, seq_len, channels)

    frames_cl = arr.reshape(seq_len, h, w, channels)

    # Basic alignment checks
    if actions is not None and actions.shape[0] != seq_len:
        print(
            f"Warning: actions length ({actions.shape[0]}) != sequence_length ({seq_len}). Will use min length."
        )
    eff_len = seq_len if actions is None else min(seq_len, actions.shape[0])

    # Save an actions.txt index for quick inspection
    if actions is not None:
        with open(os.path.join(output_dir, "actions.txt"), "w") as f:
            for t in range(eff_len):
                a = int(actions[t])
                name = (
                    action_meanings[a]
                    if action_meanings is not None and 0 <= a < len(action_meanings)
                    else None
                )
                f.write(f"{t}	{a}" + (f"	{name}" if name is not None else "") + "
")

    # Build frames with overlays and save as GIF
    duration_ms = max(1, int(1000 / max(1, fps)))
    imgs: List[Image.Image] = []
    for t in range(eff_len):
        # frames_cl is (T, H, W, C) so we can use it directly for RGB
        img_np_rgb = frames_cl[t]
        img = Image.fromarray(img_np_rgb, mode="RGB")
        draw = ImageDraw.Draw(img)
        if actions is not None:
            a = int(actions[t])
            name = (
                action_meanings[a]
                if action_meanings is not None and 0 <= a < len(action_meanings)
                else None
            )
            text = f"{a}" if name is None else f"{a} {name}"
        else:
            text = "?"
        draw.text((2, 2), text, fill=(255, 255, 255))
        imgs.append(img)

    if not imgs:
        raise RuntimeError("No frames to render.")

    if gif_path is None:
        os.makedirs(output_dir, exist_ok=True)
        gif_path = os.path.join(output_dir, "sequence.gif")

    imgs[0].save(
        gif_path,
        save_all=True,
        append_images=imgs[1:],
        duration=duration_ms,
        loop=0,
    )

    print(f"Saved GIF to {gif_path} with {eff_len} frames (H={h}, W={w}, C={channels}, fps={fps}).")


def main():
    parser = argparse.ArgumentParser(
        description="Visualize a random sequence from an ArrayRecord by saving frames with actions."
    )
    parser.add_argument("--input", required=True, help="Path to .array_record file")
    parser.add_argument(
        "--output_dir",
        required=True,
        help="Directory to save output images and actions.txt",
    )
    parser.add_argument("--seed", type=int, default=None, help="Random seed for selection")
    parser.add_argument("--channels", type=int, default=3, help="Number of channels in frames (default 3 for RGB)")
    parser.add_argument("--height", type=int, default=84, help="Frame height (if mismatch, will be inferred)")
    parser.add_argument("--width", type=int, default=84, help="Frame width (if mismatch, will be inferred)")
    parser.add_argument("--env_id", type=str, default=None, help="Gymnasium env id to map action indices to names")
    parser.add_argument("--fps", type=int, default=10, help="GIF frames per second")
    parser.add_argument("--gif_path", type=str, default=None, help="Output path for the GIF file")

    args = parser.parse_args()

    assert args.channels == 3, "Only 3 channels are currently supported"

    record = load_random_record(args.input, seed=args.seed)

    # Print quick summary
    seq_len = int(record.get("sequence_length", -1))
    has_actions = "actions" in record and record["actions"] is not None
    print(
        f"Loaded record with sequence_length={seq_len}, actions_present={has_actions}"
    )

    action_meanings = get_action_meanings(args.env_id)

    save_frames_with_actions(
        record=record,
        output_dir=args.output_dir,
        channels=int(args.channels),
        height=int(args.height) if args.height else None,
        width=int(args.width) if args.width else None,
        action_meanings=action_meanings,
        fps=int(args.fps),
        gif_path=args.gif_path,
    )


if __name__ == "__main__":
    main()