ai/environments/vector_env_gpu.py

"""
GPU-Native Vectorized Game Environment.

This module provides VectorEnvGPU - a GPU-resident implementation using CuPy
and Numba CUDA for maximum throughput. All game state arrays live in GPU VRAM,
eliminating PCI-E transfer overhead during RL training.

Usage:
    Set USE_GPU_ENV=1 to enable GPU environment in training.
"""

import json
import os
import time

import numpy as np

# CUDA detection
HAS_CUDA = False
try:
    import cupy as cp
    from numba import cuda

    if cuda.is_available():
        HAS_CUDA = True
        from numba.cuda.random import create_xoroshiro128p_states
except ImportError:
    pass

# Mock objects for CPU fallback
if not HAS_CUDA:

    class MockCP:
        int32 = np.int32
        int8 = np.int8
        float32 = np.float32
        bool_ = np.bool_

        def full(self, shape, val, dtype=None):
            return np.full(shape, val, dtype=dtype)

        def zeros(self, shape, dtype=None):
            return np.zeros(shape, dtype=dtype)

        def ones(self, shape, dtype=None):
            return np.ones(shape, dtype=dtype)

        def asnumpy(self, arr):
            return np.array(arr)

        def array(self, arr, dtype=None):
            return np.array(arr, dtype=dtype)

        def asarray(self, arr, dtype=None):
            return np.asarray(arr, dtype=dtype)

        def arange(self, n, dtype=None):
            return np.arange(n, dtype=dtype)

        def get_default_memory_pool(self):
            class MockPool:
                def used_bytes(self):
                    return 0

            return MockPool()

    cp = MockCP()

    class MockCudaMod:
        def to_device(self, arr):
            return arr

        def device_array(self, shape, dtype=None):
            return np.zeros(shape, dtype=dtype)

        def synchronize(self):
            pass

        def jit(self, *args, **kwargs):
            return lambda x: x

        def grid(self, x):
            return 0

    cuda = MockCudaMod()

    def create_xoroshiro128p_states(n, seed):
        return None


class VectorEnvGPU:
    """
    GPU-Resident Vectorized Game Environment.

    All state arrays are CuPy arrays in GPU VRAM.
    Observations and actions are passed as GPU tensors with zero-copy.

    Args:
        num_envs: Number of parallel environments
        opp_mode: Opponent mode (0=Heuristic, 1=Random)
        force_start_order: -1=Random, 0=P1, 1=P2
    """

    def __init__(self, num_envs: int = 4096, opp_mode: int = 0, force_start_order: int = -1, seed: int = 42):
        self.num_envs = num_envs
        self.opp_mode = opp_mode  # 0=Heuristic, 1=Random, 2=Solitaire
        self.force_start_order = force_start_order
        self.seed = seed

        print(f" [VectorEnvGPU] Initializing {num_envs} environments. CUDA: {HAS_CUDA}")

        # =========================================================
        # AGENT STATE (GPU-Resident)
        # =========================================================
        self.batch_stage = cp.full((num_envs, 3), -1, dtype=cp.int32)
        self.batch_energy_vec = cp.zeros((num_envs, 3, 32), dtype=cp.int32)
        self.batch_energy_count = cp.zeros((num_envs, 3), dtype=cp.int32)
        self.batch_continuous_vec = cp.zeros((num_envs, 32, 10), dtype=cp.int32)
        self.batch_continuous_ptr = cp.zeros(num_envs, dtype=cp.int32)
        self.batch_tapped = cp.zeros((num_envs, 16), dtype=cp.int32)
        self.batch_live = cp.zeros((num_envs, 50), dtype=cp.int32)
        self.batch_opp_tapped = cp.zeros((num_envs, 16), dtype=cp.int32)
        self.batch_scores = cp.zeros(num_envs, dtype=cp.int32)

        self.batch_flat_ctx = cp.zeros((num_envs, 64), dtype=cp.int32)
        self.batch_global_ctx = cp.zeros((num_envs, 128), dtype=cp.int32)

        self.batch_hand = cp.zeros((num_envs, 60), dtype=cp.int32)
        self.batch_deck = cp.zeros((num_envs, 60), dtype=cp.int32)
        self.batch_trash = cp.zeros((num_envs, 60), dtype=cp.int32)
        self.batch_opp_history = cp.zeros((num_envs, 6), dtype=cp.int32)

        # =========================================================
        # OPPONENT STATE (GPU-Resident)
        # =========================================================
        self.opp_stage = cp.full((num_envs, 3), -1, dtype=cp.int32)
        self.opp_energy_vec = cp.zeros((num_envs, 3, 32), dtype=cp.int32)
        self.opp_energy_count = cp.zeros((num_envs, 3), dtype=cp.int32)
        self.opp_tapped = cp.zeros((num_envs, 16), dtype=cp.int8)
        self.opp_live = cp.zeros((num_envs, 50), dtype=cp.int32)
        self.opp_scores = cp.zeros(num_envs, dtype=cp.int32)
        self.opp_global_ctx = cp.zeros((num_envs, 128), dtype=cp.int32)
        self.opp_hand = cp.zeros((num_envs, 60), dtype=cp.int32)
        self.opp_deck = cp.zeros((num_envs, 60), dtype=cp.int32)
        self.opp_trash = cp.zeros((num_envs, 60), dtype=cp.int32)

        # =========================================================
        # TRACKING STATE
        # =========================================================
        self.prev_scores = cp.zeros(num_envs, dtype=cp.int32)
        self.prev_opp_scores = cp.zeros(num_envs, dtype=cp.int32)
        self.prev_phases = cp.zeros(num_envs, dtype=cp.int32)
        self.episode_returns = cp.zeros(num_envs, dtype=cp.float32)
        self.episode_lengths = cp.zeros(num_envs, dtype=cp.int32)

        # =========================================================
        # OBSERVATION MODE
        # =========================================================
        self.obs_mode = os.getenv("OBS_MODE", "STANDARD")
        if self.obs_mode == "COMPRESSED":
            self.obs_dim = 512
        elif self.obs_mode == "IMAX":
            self.obs_dim = 8192
        elif self.obs_mode == "ATTENTION":
            self.obs_dim = 2240
        else:
            self.obs_dim = 2304
        print(f" [VectorEnvGPU] Observation Mode: {self.obs_mode} ({self.obs_dim}-dim)")

        self.batch_obs = cp.zeros((num_envs, self.obs_dim), dtype=cp.float32)
        self.terminal_obs_buffer = cp.zeros((num_envs, self.obs_dim), dtype=cp.float32)

        # Rewards and Dones
        self.rewards = cp.zeros(num_envs, dtype=cp.float32)
        self.dones = cp.zeros(num_envs, dtype=cp.bool_)
        self.term_scores_agent = cp.zeros(num_envs, dtype=cp.int32)
        self.term_scores_opp = cp.zeros(num_envs, dtype=cp.int32)

        # =========================================================
        # GAME CONFIG
        # =========================================================
        self.scenario_reward_scale = float(os.getenv("SCENARIO_REWARD_SCALE", "1.0"))
        if os.getenv("USE_SCENARIOS", "0") == "1" and self.scenario_reward_scale != 1.0:
            print(f" [VectorEnvGPU] Scenario Reward Scale: {self.scenario_reward_scale}")

        self.game_config = cp.zeros(10, dtype=cp.float32)
        self.game_config[0] = float(os.getenv("GAME_TURN_LIMIT", "100"))
        self.game_config[1] = float(os.getenv("GAME_STEP_LIMIT", "1000"))
        self.game_config[2] = float(os.getenv("GAME_REWARD_WIN", "100.0"))
        self.game_config[3] = float(os.getenv("GAME_REWARD_LOSE", "-100.0"))
        self.game_config[4] = float(os.getenv("GAME_REWARD_SCORE_SCALE", "50.0"))
        self.game_config[5] = float(os.getenv("GAME_REWARD_TURN_PENALTY", "-0.05"))

        # =========================================================
        # GPU RNG
        # =========================================================
        if HAS_CUDA:
            self.rng_states = create_xoroshiro128p_states(num_envs, seed=seed)
        else:
            self.rng_states = None

        # =========================================================
        # KERNEL CONFIGURATION
        # =========================================================
        self.threads_per_block = 128
        self.blocks_per_grid = (num_envs + self.threads_per_block - 1) // self.threads_per_block

        # =========================================================
        # LOAD DATA
        # =========================================================
        self._load_bytecode()
        self._load_card_stats()
        self._load_deck_pool()

        # Memory stats
        if HAS_CUDA:
            mempool = cp.get_default_memory_pool()
            used_mb = mempool.used_bytes() / 1024 / 1024
            print(f" [VectorEnvGPU] GPU VRAM used: {used_mb:.2f} MB")

    def _load_bytecode(self):
        """Load compiled bytecode to GPU."""
        host_map = np.zeros((100, 128, 4), dtype=np.int32)
        host_idx = np.zeros((2000, 8), dtype=np.int32)

        try:
            with open("data/cards_numba.json", "r") as f:
                raw_map = json.load(f)

            max_cards = 2000
            max_abilities = 8
            max_len = 128

            unique_entries = len(raw_map)
            host_map = np.zeros((unique_entries + 1, max_len, 4), dtype=np.int32)
            host_idx = np.full((max_cards, max_abilities), 0, dtype=np.int32)

            idx_counter = 1
            for key, bc_list in raw_map.items():
                cid, aid = map(int, key.split("_"))
                if cid < max_cards and aid < max_abilities:
                    bc_arr = np.array(bc_list, dtype=np.int32).reshape(-1, 4)
                    length = min(bc_arr.shape[0], max_len)
                    host_map[idx_counter, :length] = bc_arr[:length]
                    host_idx[cid, aid] = idx_counter
                    idx_counter += 1

            print(f" [VectorEnvGPU] Loaded {unique_entries} compiled abilities.")
        except FileNotFoundError:
            print(" [VectorEnvGPU] Warning: cards_numba.json not found.")
        except Exception as e:
            print(f" [VectorEnvGPU] Warning: Failed to load bytecode: {e}")

        self.bytecode_map = cp.asarray(host_map)
        self.bytecode_index = cp.asarray(host_idx)

    def _load_card_stats(self):
        """Load card statistics to GPU."""
        host_stats = np.zeros((2000, 80), dtype=np.int32)

        try:
            with open("data/cards_compiled.json", "r", encoding="utf-8") as f:
                db = json.load(f)

            count = 0
            if "member_db" in db:
                for cid_str, card in db["member_db"].items():
                    cid = int(cid_str)
                    if cid < 2000:
                        host_stats[cid, 0] = card.get("cost", 0)
                        host_stats[cid, 1] = card.get("blades", 0)
                        host_stats[cid, 2] = sum(card.get("hearts", []))
                        host_stats[cid, 10] = 1  # Type: Member

                        # Hearts breakdown
                        h_arr = card.get("hearts", [])
                        for r_idx in range(min(len(h_arr), 7)):
                            host_stats[cid, 12 + r_idx] = h_arr[r_idx]

                        # Traits
                        mask = 0
                        for g in card.get("groups", []):
                            try:
                                mask |= 1 << (int(g) % 20)
                            except:
                                pass
                        host_stats[cid, 11] = mask
                        count += 1

            if "live_db" in db:
                for cid_str, card in db["live_db"].items():
                    cid = int(cid_str)
                    if cid < 2000:
                        host_stats[cid, 10] = 2  # Type: Live
                        reqs = card.get("required_hearts", [])
                        for r_idx in range(min(len(reqs), 7)):
                            host_stats[cid, 12 + r_idx] = reqs[r_idx]
                        host_stats[cid, 38] = card.get("score", 0)
                        count += 1

            print(f" [VectorEnvGPU] Loaded stats for {count} cards.")
        except Exception as e:
            print(f" [VectorEnvGPU] Warning: Failed to load card stats: {e}")

        self.card_stats = cp.asarray(host_stats)

    def _load_deck_pool(self):
        """Load verified card pool for deck generation."""
        ability_member_ids = []
        ability_live_ids = []

        try:
            with open("data/verified_card_pool.json", "r", encoding="utf-8") as f:
                verified_data = json.load(f)

            with open("data/cards_compiled.json", "r", encoding="utf-8") as f:
                db_data = json.load(f)

            member_no_map = {}
            live_no_map = {}
            for cid, cdata in db_data.get("member_db", {}).items():
                member_no_map[cdata["card_no"]] = int(cid)
            for cid, cdata in db_data.get("live_db", {}).items():
                live_no_map[cdata["card_no"]] = int(cid)

            if isinstance(verified_data, list):
                for v_no in verified_data:
                    if v_no in member_no_map:
                        ability_member_ids.append(member_no_map[v_no])
                    elif v_no in live_no_map:
                        ability_live_ids.append(live_no_map[v_no])
            else:
                source_members = verified_data.get("verified_abilities", []) + verified_data.get("members", [])
                for v_no in source_members:
                    if v_no in member_no_map:
                        ability_member_ids.append(member_no_map[v_no])

                source_lives = verified_data.get("verified_lives", []) + verified_data.get("lives", [])
                for v_no in source_lives:
                    if v_no in live_no_map:
                        ability_live_ids.append(live_no_map[v_no])

                if not ability_member_ids:
                    for v_no in verified_data.get("vanilla_members", []):
                        if v_no in member_no_map:
                            ability_member_ids.append(member_no_map[v_no])
                if not ability_live_ids:
                    for v_no in verified_data.get("vanilla_lives", []):
                        if v_no in live_no_map:
                            ability_live_ids.append(live_no_map[v_no])

            if not ability_member_ids:
                ability_member_ids = [1]
            if not ability_live_ids:
                ability_live_ids = [999]

            print(f" [VectorEnvGPU] Deck Pool: {len(ability_member_ids)} members, {len(ability_live_ids)} lives")
        except Exception as e:
            print(f" [VectorEnvGPU] Deck Load Error: {e}")
            ability_member_ids = [1]
            ability_live_ids = [999]

        self.ability_member_ids = cp.array(ability_member_ids, dtype=cp.int32)
        self.ability_live_ids = cp.array(ability_live_ids, dtype=cp.int32)

    # =========================================================
    # PYTORCH INTERFACE
    # =========================================================

    def get_observations_tensor(self):
        """Return observations as PyTorch CUDA tensor (zero-copy)."""
        import torch

        return torch.as_tensor(self.batch_obs, device="cuda")

    def get_action_masks_tensor(self):
        """Return action masks as PyTorch CUDA tensor."""
        import torch

        masks = self.get_action_masks()
        return torch.as_tensor(masks, device="cuda")

    def get_rewards_tensor(self):
        """Return rewards as PyTorch CUDA tensor."""
        import torch

        return torch.as_tensor(self.rewards, device="cuda")

    def get_dones_tensor(self):
        """Return dones as PyTorch CUDA tensor."""
        import torch

        return torch.as_tensor(self.dones, device="cuda")

    # =========================================================
    # ENVIRONMENT INTERFACE
    # =========================================================

    def reset(self, indices=None):
        """Reset environments."""
        if not HAS_CUDA:
            # CPU fallback
            self.batch_stage.fill(-1)
            self.batch_scores.fill(0)
            self.batch_global_ctx.fill(0)
            self.batch_hand.fill(0)
            self.batch_deck.fill(0)
            return self.batch_obs

        from ai.cuda_kernels import encode_observations_attention_kernel, encode_observations_kernel, reset_kernel

        if indices is None:
            indices_gpu = cp.arange(self.num_envs, dtype=cp.int32)
        else:
            indices_gpu = cp.array(indices, dtype=cp.int32)

        blocks = (len(indices_gpu) + self.threads_per_block - 1) // self.threads_per_block

        reset_kernel[blocks, self.threads_per_block](
            indices_gpu,
            self.batch_stage,
            self.batch_energy_vec,
            self.batch_energy_count,
            self.batch_continuous_vec,
            self.batch_continuous_ptr,
            self.batch_tapped,
            self.batch_live,
            self.batch_scores,
            self.batch_flat_ctx,
            self.batch_global_ctx,
            self.batch_hand,
            self.batch_deck,
            self.batch_trash,
            self.batch_opp_history,
            self.opp_stage,
            self.opp_energy_vec,
            self.opp_energy_count,
            self.opp_tapped,
            self.opp_live,
            self.opp_scores,
            self.opp_global_ctx,
            self.opp_hand,
            self.opp_deck,
            self.opp_trash,
            self.ability_member_ids,
            self.ability_live_ids,
            self.rng_states,
            self.force_start_order,
            self.batch_obs,
            self.card_stats,
        )

        # Encode initial observations
        if self.obs_mode == "ATTENTION":
            encode_observations_attention_kernel[self.blocks_per_grid, self.threads_per_block](
                self.num_envs,
                self.batch_hand,
                self.batch_stage,
                self.batch_energy_count,
                self.batch_tapped,
                self.batch_scores,
                self.opp_scores,
                self.opp_stage,
                self.opp_tapped,
                self.card_stats,
                self.batch_global_ctx,
                self.batch_live,
                self.batch_opp_history,
                self.opp_global_ctx,
                1,
                self.batch_obs,
            )
        else:
            encode_observations_kernel[self.blocks_per_grid, self.threads_per_block](
                self.num_envs,
                self.batch_hand,
                self.batch_stage,
                self.batch_energy_count,
                self.batch_tapped,
                self.batch_scores,
                self.opp_scores,
                self.opp_stage,
                self.opp_tapped,
                self.card_stats,
                self.batch_global_ctx,
                self.batch_live,
                1,
                self.batch_obs,
            )

        # Reset tracking
        if indices is None:
            self.prev_scores.fill(0)
            self.prev_opp_scores.fill(0)
            self.episode_returns.fill(0)
            self.episode_lengths.fill(0)
        else:
            self.prev_scores[indices_gpu] = 0
            self.prev_opp_scores[indices_gpu] = 0
            self.episode_returns[indices_gpu] = 0
            self.episode_lengths[indices_gpu] = 0

        return self.batch_obs

    def step(self, actions):
        """
        Step all environments.

        Args:
            actions: CuPy array or PyTorch tensor of actions

        Returns:
            obs, rewards, dones, infos
        """
        if not HAS_CUDA:
            # Fallback
            return self.batch_obs, self.rewards, self.dones, [{}] * self.num_envs

        import torch
        from ai.cuda_kernels import (
            encode_observations_attention_kernel,
            encode_observations_kernel,
            reset_kernel,
            step_kernel,
        )

        # Convert to CuPy if needed
        if isinstance(actions, torch.Tensor):
            actions_gpu = cp.asarray(actions.cpu().numpy(), dtype=cp.int32)
        elif isinstance(actions, np.ndarray):
            actions_gpu = cp.asarray(actions, dtype=cp.int32)
        else:
            actions_gpu = actions

        # 1. Step kernel
        step_kernel[self.blocks_per_grid, self.threads_per_block](
            self.num_envs,
            actions_gpu,
            self.batch_hand,
            self.batch_deck,
            self.batch_stage,
            self.batch_energy_vec,
            self.batch_energy_count,
            self.batch_continuous_vec,
            self.batch_continuous_ptr,
            self.batch_tapped,
            self.batch_live,
            self.batch_scores,
            self.batch_flat_ctx,
            self.batch_global_ctx,
            self.opp_hand,
            self.opp_deck,
            self.opp_stage,
            self.opp_energy_vec,
            self.opp_energy_count,
            self.opp_tapped,
            self.opp_live,
            self.opp_scores,
            self.opp_global_ctx,
            self.card_stats,
            self.bytecode_map,
            self.bytecode_index,
            self.batch_obs,
            self.rewards,
            self.dones,
            self.prev_scores,
            self.prev_opp_scores,
            self.prev_phases,
            self.terminal_obs_buffer,
            self.batch_trash,
            self.opp_trash,
            self.batch_opp_history,
            self.term_scores_agent,
            self.term_scores_opp,
            self.ability_member_ids,
            self.ability_live_ids,
            self.rng_states,
            self.game_config,
            self.opp_mode,
            self.force_start_order,
        )

        # Apply Scenario Reward Scaling
        if self.scenario_reward_scale != 1.0 and os.getenv("USE_SCENARIOS", "0") == "1":
            self.rewards *= self.scenario_reward_scale

        # 2. Update Episodic Returns/Lengths (Vectorized GPU)
        self.episode_returns += self.rewards
        self.episode_lengths += 1

        # 3. Handle Auto-Reset (High Performance)
        dones_cpu = cp.asnumpy(self.dones)

        # Pre-allocate infos list (reused or created)
        infos = [{} for _ in range(self.num_envs)]

        if np.any(dones_cpu):
            done_indices = np.where(dones_cpu)[0]
            done_indices_gpu = cp.array(done_indices, dtype=cp.int32)

            # A. Capture Terminal Observations (from UNRESET state)
            # Efficient Device-to-Device copy
            # NOTE: step_kernel leaves env in finished state, so batch_obs has terminal state.
            # We must encode it?
            # Actually, step_kernel calls encode at end? No, step_kernel does NOT encode obs in my implementation.
            # I removed the Python-side encode calls from previous impl?
            # Wait, step_kernel logic in my head vs file.
            # In ai/cuda_kernels.py, step_kernel does NOT call encode.
            # So batch_obs is STALE (from previous step)!
            # We MUST encode the terminal state first.

            # Encode CURRENT state (Terminal) for ALL envs? Or just done?
            # Usually we encode all envs at end of step.
            # BUT we need to reset done envs and encode AGAIN.

            # OPTIMIZATION:
            # 1. Encode ALL envs (Next state for running, Terminal for done).
            turn_num = 1  # Dummy, kernels use ctx
            if self.obs_mode == "ATTENTION":
                encode_observations_attention_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    self.batch_opp_history,
                    self.opp_global_ctx,
                    turn_num,
                    self.batch_obs,
                )
            else:
                encode_observations_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    turn_num,
                    self.batch_obs,
                )

            # 2. For Done Envs: Copy encoded terminal state to buffer
            # We can use fancy indexing copy on GPU
            self.terminal_obs_buffer[done_indices_gpu] = self.batch_obs[done_indices_gpu]

            # 3. Fetch Terminal Info Metrics (Bulk D2H)
            final_returns = cp.asnumpy(self.episode_returns[done_indices_gpu])
            final_lengths = cp.asnumpy(self.episode_lengths[done_indices_gpu])
            term_obs_cpu = cp.asnumpy(self.terminal_obs_buffer[done_indices_gpu])
            term_scores_ag = cp.asnumpy(self.term_scores_agent[done_indices_gpu])
            term_scores_op = cp.asnumpy(self.term_scores_opp[done_indices_gpu])

            # 4. Populate Infos (CPU Loop over SMALL subset)
            for k, idx in enumerate(done_indices):
                infos[idx] = {
                    "terminal_observation": term_obs_cpu[k],
                    "episode": {"r": float(final_returns[k]), "l": int(final_lengths[k])},
                    "terminal_score_agent": int(term_scores_ag[k]),
                    "terminal_score_opp": int(term_scores_op[k]),
                }

            # 5. Reset Done Envs
            # Reset accumulators
            self.episode_returns[done_indices_gpu] = 0
            self.episode_lengths[done_indices_gpu] = 0

            # Launch Reset Kernel
            blocks_reset = (len(done_indices) + self.threads_per_block - 1) // self.threads_per_block
            reset_kernel[blocks_reset, self.threads_per_block](
                done_indices_gpu,
                self.batch_stage,
                self.batch_energy_vec,
                self.batch_energy_count,
                self.batch_continuous_vec,
                self.batch_continuous_ptr,
                self.batch_tapped,
                self.batch_live,
                self.batch_scores,
                self.batch_flat_ctx,
                self.batch_global_ctx,
                self.batch_hand,
                self.batch_deck,
                self.batch_trash,
                self.batch_opp_history,
                self.opp_stage,
                self.opp_energy_vec,
                self.opp_energy_count,
                self.opp_tapped,
                self.opp_live,
                self.opp_scores,
                self.opp_global_ctx,
                self.opp_hand,
                self.opp_deck,
                self.opp_trash,
                self.ability_member_ids,
                self.ability_live_ids,
                self.rng_states,
                self.force_start_order,
                self.batch_obs,
                self.card_stats,
            )

            # 6. Re-Encode Reset Envs (to get initial state)
            # We assume reset_kernel updates state but NOT obs.
            # We need to re-run encode kernel ONLY for done indices?
            # Or run global encode again? Global is waste.
            # We need an encode kernel that takes indices.
            # The current kernel takes `num_envs` and assumes `0..N`.
            # We can reuse the global kernel if we are clever or modify it.
            # Modifying kernel to accept indices is best.
            # However, for now, to save complexity, we can re-run global encode.
            # It's redundant for non-done envs but correct.
            # Better: Reset modifies batch_obs directly? No, reset_kernel doesn't encode.

            # Let's re-run global encode. It's fast (GPU) compared to CPU loop.
            if self.obs_mode == "ATTENTION":
                encode_observations_attention_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    self.batch_opp_history,
                    self.opp_global_ctx,
                    turn_num,
                    self.batch_obs,
                )
            else:
                encode_observations_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    turn_num,
                    self.batch_obs,
                )

        else:
            # No resets needed. Just encode once to get next states.
            # Encode observations
            turn_num = 1
            if self.obs_mode == "ATTENTION":
                encode_observations_attention_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    self.batch_opp_history,
                    self.opp_global_ctx,
                    turn_num,
                    self.batch_obs,
                )
            else:
                encode_observations_kernel[self.blocks_per_grid, self.threads_per_block](
                    self.num_envs,
                    self.batch_hand,
                    self.batch_stage,
                    self.batch_energy_count,
                    self.batch_tapped,
                    self.batch_scores,
                    self.opp_scores,
                    self.opp_stage,
                    self.opp_tapped,
                    self.card_stats,
                    self.batch_global_ctx,
                    self.batch_live,
                    turn_num,
                    self.batch_obs,
                )

        return self.batch_obs, self.rewards, self.dones, infos

    def get_observations(self):
        """Return observation buffer (CuPy array)."""
        return self.batch_obs

    def get_action_masks(self):
        """Compute and return action masks (CuPy array)."""
        if not HAS_CUDA:
            return cp.ones((self.num_envs, 2000), dtype=cp.bool_)

        from ai.cuda_kernels import compute_action_masks_kernel

        masks = cp.zeros((self.num_envs, 2000), dtype=cp.bool_)

        compute_action_masks_kernel[self.blocks_per_grid, self.threads_per_block](
            self.num_envs,
            self.batch_hand,
            self.batch_stage,
            self.batch_tapped,
            self.batch_global_ctx,
            self.batch_live,
            self.card_stats,
            masks,
        )

        return masks


# ============================================================================
# BENCHMARK
# ============================================================================


def benchmark_gpu_env(num_envs=4096, steps=1000):
    """Benchmark GPU environment throughput."""
    print("\n=== GPU Environment Benchmark ===")
    print(f"Environments: {num_envs}")
    print(f"Steps: {steps}")

    env = VectorEnvGPU(num_envs=num_envs)
    env.reset()

    # Warmup
    for _ in range(10):
        actions = cp.zeros(num_envs, dtype=cp.int32)
        env.step(actions)

    if HAS_CUDA:
        cuda.synchronize()

    # Benchmark
    start = time.time()
    for _ in range(steps):
        actions = cp.zeros(num_envs, dtype=cp.int32)  # Pass action
        env.step(actions)

    if HAS_CUDA:
        cuda.synchronize()

    elapsed = time.time() - start
    total_steps = num_envs * steps
    sps = total_steps / elapsed

    print("\nResults:")
    print(f"  Total Steps: {total_steps:,}")
    print(f"  Time: {elapsed:.2f}s")
    print(f"  Throughput: {sps:,.0f} steps/sec")

    return sps


if __name__ == "__main__":
    # Quick test
    env = VectorEnvGPU(num_envs=128)
    obs = env.reset()
    print(f"Observation shape: {obs.shape}")

    actions = cp.zeros(128, dtype=cp.int32)
    obs, rewards, dones, infos = env.step(actions)
    print(f"Step completed. Rewards shape: {rewards.shape}")

    # Benchmark
    benchmark_gpu_env(num_envs=1024, steps=100)