"""CUDA grid-foraging + 3x MinGRU(h=256) policy. Fast rollout on RTX PRO 6000.

Exposes Model / policy_forward / env_step / run matching reference.py exactly.
The heavy lifting is a real CUDA C++ kernel (kernels.cu) compiled via
torch.utils.cpp_extension.load.
"""
from __future__ import annotations

import os
from pathlib import Path

import torch
import torch.nn as nn

BOARD = 11
OBS_DIM = 4
HIDDEN = 256
GRU_LAYERS = 3
NUM_ACTIONS = 4
GRU_OUT = 3 * HIDDEN

_HERE = Path(__file__).resolve().parent


def _setup_env() -> None:
    # Prefer a real nvcc over any shim; ensure ninja is findable.
    cur = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
    if not (cur and (Path(cur) / "bin/nvcc").exists()):
        for cand in ("/usr/local/cuda-12.8", "/usr/local/cuda", "/usr/local/cuda-13.0"):
            if (Path(cand) / "bin/nvcc").exists():
                os.environ["CUDA_HOME"] = cand
                break
    # Locate ninja (torch's cpp_extension requires it).
    cand_bins = []
    if os.environ.get("VIRTUAL_ENV"):
        cand_bins.append(str(Path(os.environ["VIRTUAL_ENV"]) / "bin"))
    try:
        import ninja  # type: ignore

        cand_bins.append(ninja.BIN_DIR)
    except Exception:
        pass
    for cand_bin in cand_bins:
        if os.path.isdir(cand_bin) and os.path.exists(os.path.join(cand_bin, "ninja")):
            os.environ["PATH"] = cand_bin + os.pathsep + os.environ.get("PATH", "")
            break


_setup_env()

from torch.utils.cpp_extension import load as _ext_load  # noqa: E402

_ext = _ext_load(
    name="grid_mingru_ext",
    sources=[str(_HERE / "kernels.cu")],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.w_enc = nn.Parameter(torch.empty(HIDDEN, OBS_DIM))
        self.b_enc = nn.Parameter(torch.zeros(HIDDEN))
        self.w_gru = nn.Parameter(torch.empty(GRU_LAYERS, GRU_OUT, HIDDEN))
        self.w_a = nn.Parameter(torch.empty(NUM_ACTIONS, HIDDEN))
        self.b_a = nn.Parameter(torch.zeros(NUM_ACTIONS))
        self.w_v = nn.Parameter(torch.empty(1, HIDDEN))
        self.b_v = nn.Parameter(torch.zeros(1))
        self.reset_parameters(0)

    def reset_parameters(self, seed: int = 0) -> None:
        g = torch.Generator(device="cpu")
        g.manual_seed(seed)
        for p in self.parameters():
            tmp = torch.empty(p.shape, dtype=p.dtype, device="cpu")
            tmp.normal_(0.0, 0.02, generator=g)
            p.data.copy_(tmp)

    def forward(self, obs: torch.Tensor, state: torch.Tensor):
        return policy_forward(self, obs, state)


def _weights(model: nn.Module, device=None):
    dev = device or next(model.parameters()).device
    return tuple(
        p.detach().to(device=dev, dtype=torch.float32).contiguous()
        for p in (model.w_enc, model.b_enc, model.w_gru, model.w_a,
                  model.b_a, model.w_v, model.b_v)
    )


def policy_forward(model: nn.Module, obs: torch.Tensor, state: torch.Tensor):
    w = _weights(model, obs.device)
    logits, state_out, value = _ext.policy_forward(*w, obs.contiguous(),
                                                   state.contiguous())
    return logits, state_out, value


def env_step(agent: torch.Tensor, food: torch.Tensor,
             actions: torch.Tensor, rng_state: torch.Tensor):
    agent = agent.contiguous()
    food = food.contiguous()
    actions = actions.contiguous()
    rng_state = rng_state.contiguous()
    agent_o, food_o, reward, rng_o = _ext.env_step(agent, food, actions, rng_state)
    return agent_o, food_o, reward, rng_o


def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
    device = torch.device("cuda:0")
    if model is None:
        model = Model()
    model = model.to(device).eval()
    w = _weights(model, device)

    # Initial conditions must match reference.run exactly (CPU MT19937).
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    agent = torch.randint(0, BOARD, (num_envs, 2), generator=g)
    food = torch.randint(0, BOARD, (num_envs, 2), generator=g)
    rng_state = torch.arange(num_envs, device=device, dtype=torch.int64) + (seed * 10007)

    rewards, positions, last_logits = _ext.rollout(
        *w, agent, food, rng_state, int(num_envs), int(horizon)
    )

    return {
        "rewards": rewards,
        "positions": positions,
        "last_logits": last_logits,
    }


# ==================================================================
# ===== sidecar: kernels.cu (18907 bytes, loaded by solution.py) =====
# ==================================================================

// Grid-foraging + 3x MinGRU(h=256) CUDA kernels.
// Policy is a block-tiled SGEMM: each block handles BN envs, keeping the
// recurrent activation h in shared memory between layers.
#include <cuda_runtime.h>
#include <cuda_pipeline.h>
#include <cstdint>

#include <torch/extension.h>

#define BOARD 11
#define HIDDEN 256
#define GRU_LAYERS 3
#define GRU_OUT 768
#define NUM_ACTIONS 4
#define LCG_A 6364136223846793005ULL
#define LCG_MASK 0x7FFFFFFFFFFFFFFFULL

// ---- tiling config ----
#define BN 32            // envs per block
#define BM 192           // gate rows per M-pass
#define THREADS 256
#define TM 12            // M-rows per thread (multiple of 3 -> TM/3 hidden units)
#define TN 2             // envs per thread
#define KC 8             // K-chunk
#define N_MPASS (GRU_OUT / BM)          // 4
#define N_KC (HIDDEN / KC)              // 32
#define N_HU (TM / 3)                   // hidden units per thread per pass

#define SHMEM_WS_OFF (HIDDEN * BN)                      // floats (single h buffer)
#define SHMEM_OBS_OFF (SHMEM_WS_OFF + 2 * KC * BM)      // floats (double-buffered Ws)
#define SHMEM_TOTAL_FLOATS (SHMEM_OBS_OFF + BN * 4)


__device__ __forceinline__ float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }

__device__ __forceinline__ uint64_t lcg_step(uint64_t r) {
    return (r * LCG_A + 1ULL) & LCG_MASK;
}

__device__ __forceinline__ void stage_ws(float* Ws, int buf, const float* wlT,
                                         int m_base, int k0) {
    // copy KC*BM floats (16B chunks) from global wlT into shared Ws[buf]
    constexpr int n4 = (KC * BM) >> 2;
    for (int i = threadIdx.x; i < n4; i += THREADS) {
        int kk = i / (BM >> 2);
        int mm4 = i % (BM >> 2);
        const float* src = wlT + (k0 + kk) * GRU_OUT + m_base + (mm4 * 4);
        float4* dst = reinterpret_cast<float4*>(Ws + buf * (KC * BM) + kk * BM + (mm4 * 4));
        __pipeline_memcpy_async(dst, src, 16);
    }
    __pipeline_commit();
}

// Block-tiled MLP. smem layout:
//   [0, HIDDEN*BN)          : h buffer (k-major: [k][n])
//   [SHMEM_WS_OFF, +2*KC*BM): double-buffered W staging chunk (k-major: [kk][m])
//   [SHMEM_OBS_OFF, +BN*4)  : obs buffer [n][i]
// mode: 0 = obs direct, 1 = obs from agent/food
__global__ void __launch_bounds__(THREADS) policy_tiled_kernel(
    const float* __restrict__ w_enc, const float* __restrict__ b_enc,
    const float* __restrict__ w_permT,   // (3, 256, 768) permuted-transposed GRU weights
    const float* __restrict__ w_a, const float* __restrict__ b_a,
    const float* __restrict__ w_v, const float* __restrict__ b_v,
    const float* __restrict__ obs_in,    // (N,4) for mode 0
    const int* __restrict__ agent, const int* __restrict__ food,  // for mode 1
    const float* __restrict__ state_in, float* __restrict__ state_out,
    float* __restrict__ logits_out, float* __restrict__ value_out,
    int* __restrict__ actions_out, int N, int mode)
{
    extern __shared__ float smem[];
    float* buf0 = smem;              // [HIDDEN*BN]  (h, k-major [k][n])
    float* Ws = smem + SHMEM_WS_OFF;    // [2*KC*BM]
    float* obs_s = smem + SHMEM_OBS_OFF; // [BN*4]

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;           // 0..7
    const int warp_m = warp >> 1;        // 0..3
    const int warp_n = warp & 1;         // 0..1
    const int thread_m = lane >> 3;      // 0..3
    const int thread_n = lane & 7;       // 0..7
    const int env_base = blockIdx.x * BN;

    // m offset (relative to M-pass base) and env offset (relative to block envs)
    const int m_off = warp_m * (BM / 4) + thread_m * TM;   // 0..BM
    const int env_off = warp_n * (BN / 2) + thread_n * TN; // 0..BN

    // ---- load obs into shared ----
    #pragma unroll
    for (int idx = tid; idx < BN * 4; idx += THREADS) {
        int n = idx >> 2, i = idx & 3;
        int n_abs = env_base + n;
        if (n_abs < N) {
            if (mode == 0) {
                obs_s[n * 4 + i] = obs_in[n_abs * 4 + i];
            } else {
                int ax = agent[n_abs * 2], ay = agent[n_abs * 2 + 1];
                int fx = food[n_abs * 2], fy = food[n_abs * 2 + 1];
                if (i == 0) obs_s[n * 4 + 0] = (float)(fx - ax) / (float)BOARD;
                else if (i == 1) obs_s[n * 4 + 1] = (float)(fy - ay) / (float)BOARD;
                else if (i == 2) obs_s[n * 4 + 2] = (float)ax / 10.0f;
                else obs_s[n * 4 + 3] = (float)ay / 10.0f;
            }
        }
    }
    __syncthreads();

    // ---- enc: hT[u][n] = b_enc[u] + sum_i obs_s[n*4+i]*w_enc[u*4+i] ----
    #pragma unroll
    for (int idx = tid; idx < HIDDEN * BN; idx += THREADS) {
        int u = idx / BN, n = idx % BN;
        int n_abs = env_base + n;
        float s = b_enc[u];
        if (n_abs < N) {
            const float* o = obs_s + n * 4;
            const float* we = w_enc + u * 4;
            s += we[0] * o[0] + we[1] * o[1] + we[2] * o[2] + we[3] * o[3];
        }
        buf0[u * BN + n] = s;
    }
    __syncthreads();

    // ---- GRU layers ----
    #pragma unroll
    for (int l = 0; l < GRU_LAYERS; ++l) {
        const float* wlT = w_permT + l * HIDDEN * GRU_OUT;  // (256, 768)
        float acc[TM][TN];
        float nh[N_MPASS * N_HU * TN];  // new hidden values, written back after all passes

        #pragma unroll
        for (int mp = 0; mp < N_MPASS; ++mp) {
            const int m_base = mp * BM;
            // zero accumulators
            #pragma unroll
            for (int tm = 0; tm < TM; ++tm)
                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) acc[tm][tn] = 0.0f;

            // pipelined K-loop: double-buffered Ws via cp.async
            stage_ws(Ws, 0, wlT, m_base, 0);
            for (int kc = 0; kc < N_KC; ++kc) {
                const int k0 = kc * KC;
                if (kc + 1 < N_KC) stage_ws(Ws, (kc + 1) & 1, wlT, m_base, k0 + KC);
                __pipeline_wait_prior(kc + 1 < N_KC ? 1 : 0);
                __syncthreads();
                const float* Wp = Ws + (kc & 1) * (KC * BM);
                #pragma unroll
                for (int kk = 0; kk < KC; ++kk) {
                    float h0 = buf0[(k0 + kk) * BN + env_off + 0];
                    float h1 = buf0[(k0 + kk) * BN + env_off + 1];
                    const float* wp = Wp + kk * BM + m_off;
                    #pragma unroll
                    for (int tm = 0; tm < TM; ++tm) {
                        float w = wp[tm];
                        acc[tm][0] += w * h0;
                        acc[tm][1] += w * h1;
                    }
                }
            }

            // ---- GRU update for this M-pass: keep new h in registers ----
            const int u_base = (m_base + m_off) / 3;
            #pragma unroll
            for (int ul = 0; ul < N_HU; ++ul) {
                const int u = u_base + ul;
                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) {
                    int n_abs = env_base + env_off + tn;
                    if (n_abs >= N) continue;
                    float zh = acc[ul * 3 + 0][tn];
                    float zg = acc[ul * 3 + 1][tn];
                    float zp = acc[ul * 3 + 2][tn];
                    float st = state_in[n_abs * GRU_LAYERS * HIDDEN + l * HIDDEN + u];
                    float out = st + sigmoidf(zg) * (tanhf(zh) - st);
                    float p = sigmoidf(zp);
                    float hn = p * out + (1.0f - p) * buf0[u * BN + env_off + tn];
                    nh[(mp * N_HU + ul) * TN + tn] = hn;
                    state_out[n_abs * GRU_LAYERS * HIDDEN + l * HIDDEN + u] = out;
                }
            }
        }
        // write new hidden back to buf0 (all passes done; buf0 reads complete)
        #pragma unroll
        for (int mp = 0; mp < N_MPASS; ++mp) {
            const int u_base = (mp * BM + m_off) / 3;
            #pragma unroll
            for (int ul = 0; ul < N_HU; ++ul) {
                const int u = u_base + ul;
                #pragma unroll
                for (int tn = 0; tn < TN; ++tn) {
                    buf0[u * BN + env_off + tn] = nh[(mp * N_HU + ul) * TN + tn];
                }
            }
        }
        __syncthreads();  // new-h writes visible before next layer reads
    }

    // ---- output: logits, value, argmax ----
    #pragma unroll
    for (int tn = 0; tn < TN; ++tn) {
        int n_abs = env_base + env_off + tn;
        if (n_abs >= N) continue;
        float lp[4];
        #pragma unroll
        for (int a = 0; a < NUM_ACTIONS; ++a) {
            float s = b_a[a];
            const float* wa = w_a + a * HIDDEN;
            #pragma unroll 4
            for (int k = 0; k < HIDDEN; ++k) s += wa[k] * buf0[k * BN + env_off + tn];
            lp[a] = s;
        }
        float v = b_v[0];
        #pragma unroll 4
        for (int k = 0; k < HIDDEN; ++k) v += w_v[k] * buf0[k * BN + env_off + tn];
        int best = 0;
        #pragma unroll
        for (int a = 1; a < NUM_ACTIONS; ++a) if (lp[a] > lp[best]) best = a;
        logits_out[n_abs * NUM_ACTIONS + 0] = lp[0];
        logits_out[n_abs * NUM_ACTIONS + 1] = lp[1];
        logits_out[n_abs * NUM_ACTIONS + 2] = lp[2];
        logits_out[n_abs * NUM_ACTIONS + 3] = lp[3];
        value_out[n_abs] = v;
        actions_out[n_abs] = best;
    }
}

// ---- env movement + hit detection + reward accumulation ----
__global__ void env_move_kernel(
    const int* __restrict__ agent, const int* __restrict__ food,
    const int* __restrict__ actions, int* __restrict__ agent_out,
    unsigned char* __restrict__ hit, int* __restrict__ hit_any,
    float* __restrict__ rewards, int N)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;
    int a = actions[idx];
    int ax = agent[idx * 2], ay = agent[idx * 2 + 1];
    int nx = ax, ny = ay;
    if (a == 0) { if (ny > 0) --ny; }
    else if (a == 1) { if (ny < BOARD - 1) ++ny; }
    else if (a == 2) { if (nx > 0) --nx; }
    else { if (nx < BOARD - 1) ++nx; }
    agent_out[idx * 2] = nx;
    agent_out[idx * 2 + 1] = ny;
    int h = (nx == food[idx * 2] && ny == food[idx * 2 + 1]) ? 1 : 0;
    hit[idx] = (unsigned char)h;
    if (h) { atomicOr(hit_any, 1); rewards[idx] += 1.0f; }
}

// ---- respawn: multi-block; if any hit, advance all rng and respawn hit foods ----
__global__ void respawn_kernel(
    const unsigned char* __restrict__ hit, const int* __restrict__ hit_any,
    uint64_t* __restrict__ rng, int* __restrict__ food, int N)
{
    int any = *hit_any;
    if (!any) return;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N) return;
    uint64_t r = lcg_step(rng[idx]);
    uint64_t r2 = lcg_step(r);
    rng[idx] = r2;
    if (hit[idx]) {
        food[idx * 2] = (int)(r % BOARD);
        food[idx * 2 + 1] = (int)(r2 % BOARD);
    }
}

// ---- reset the any-hit flag (1 thread) ----
__global__ void reset_flag_kernel(int* __restrict__ hit_any) {
    if (threadIdx.x == 0) *hit_any = 0;
}

// ===================== host helpers =====================

static void check_cuda(const char* tag) {
    cudaError_t e = cudaGetLastError();
    TORCH_CHECK(e == cudaSuccess, tag, ": ", cudaGetErrorString(e));
}

static const float* wptr(torch::Tensor& t) { return t.data_ptr<float>(); }
static const int* iptr(torch::Tensor& t) { return t.data_ptr<int>(); }
static float* wptr_nc(torch::Tensor& t) { return t.data_ptr<float>(); }
static int* iptr_nc(torch::Tensor& t) { return t.data_ptr<int>(); }

// w_gru (3,768,256) -> w_permT (3,256,768): w_permT[l][k][m] = w_gru[l][(m%3)*256 + m/3][k]
static torch::Tensor permute_gru(torch::Tensor w_gru) {
    w_gru = w_gru.contiguous();
    auto opts = w_gru.options();
    auto out = torch::empty({GRU_LAYERS, HIDDEN, GRU_OUT}, opts);
    auto idx = torch::arange(GRU_OUT, w_gru.options().dtype(torch::kInt64));
    auto orig = (idx % 3) * HIDDEN + (idx / 3);
    auto permuted = w_gru.index_select(1, orig.to(opts.dtype(torch::kLong)));
    out.copy_(permuted.permute({0, 2, 1}));
    return out;
}

static const int SHMEM_BYTES = SHMEM_TOTAL_FLOATS * 4;

static void ensure_smem(int shmem_bytes) {
    static bool configured = false;
    if (!configured) {
        cudaFuncAttributes attr;
        cudaFuncGetAttributes(&attr, policy_tiled_kernel);  // init attrs
        cudaGetLastError();
        cudaError_t e = cudaFuncSetAttribute(
            policy_tiled_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_BYTES);
        TORCH_CHECK(e == cudaSuccess, "cudaFuncSetAttribute: ", cudaGetErrorString(e));
        configured = true;
    }
}

static void launch_policy(
    torch::Tensor& w_enc, torch::Tensor& b_enc, torch::Tensor& w_permT,
    torch::Tensor& w_a, torch::Tensor& b_a, torch::Tensor& w_v, torch::Tensor& b_v,
    torch::Tensor& obs, torch::Tensor& agent, torch::Tensor& food,
    torch::Tensor& state_in, torch::Tensor& state_out,
    torch::Tensor& logits, torch::Tensor& value, torch::Tensor& actions,
    int mode)
{
    int N = (int)state_in.size(0);
    ensure_smem(SHMEM_BYTES);
    int blocks = (N + BN - 1) / BN;
    policy_tiled_kernel<<<blocks, THREADS, SHMEM_BYTES>>>(
        wptr(w_enc), wptr(b_enc), wptr(w_permT),
        wptr(w_a), wptr(b_a), wptr(w_v), wptr(b_v),
        mode == 0 ? wptr(obs) : nullptr,
        mode == 1 ? iptr(agent) : nullptr, mode == 1 ? iptr(food) : nullptr,
        wptr(state_in), wptr_nc(state_out),
        wptr_nc(logits), wptr_nc(value), iptr_nc(actions), N, mode);
    check_cuda("policy_tiled");
}

std::vector<torch::Tensor> policy_forward(
    torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_gru,
    torch::Tensor w_a, torch::Tensor b_a, torch::Tensor w_v, torch::Tensor b_v,
    torch::Tensor obs, torch::Tensor state)
{
    auto N = obs.size(0);
    auto w_permT = permute_gru(w_gru);
    auto logits = torch::empty({N, NUM_ACTIONS}, obs.options());
    auto state_out = torch::empty({N, GRU_LAYERS, HIDDEN}, obs.options());
    auto value = torch::empty({N}, obs.options());
    auto actions = torch::empty({N}, obs.options().dtype(torch::kInt32));
    auto obs_c = obs.contiguous();
    auto state_c = state.contiguous();
    auto dummy = torch::empty({0}, obs.options());
    launch_policy(w_enc, b_enc, w_permT, w_a, b_a, w_v, b_v,
                  obs_c, dummy, dummy, state_c, state_out,
                  logits, value, actions, 0);
    return {logits, state_out, value};
}

// env_step on float tensors -> tuple (agent, food, reward, rng)
std::vector<torch::Tensor> env_step(
    torch::Tensor agent_f, torch::Tensor food_f, torch::Tensor actions,
    torch::Tensor rng)
{
    auto N = agent_f.size(0);
    auto agent_i = agent_f.to(torch::kInt32).contiguous();
    auto food_i = food_f.to(torch::kInt32).contiguous();
    auto acts_i = actions.to(torch::kInt32).contiguous();
    auto agent_o = torch::empty_like(agent_i);
    auto hit = torch::zeros({N}, agent_f.options().dtype(torch::kUInt8));
    auto hit_any = torch::zeros({1}, agent_f.options().dtype(torch::kInt32));
    auto reward = torch::zeros({N}, agent_f.options());
    auto rng_c = rng.to(torch::kInt64).contiguous();
    int n = (int)N;
    int threads = 256;
    int blocks = (n + threads - 1) / threads;
    env_move_kernel<<<blocks, threads>>>(
        agent_i.data_ptr<int>(), food_i.data_ptr<int>(), acts_i.data_ptr<int>(),
        agent_o.data_ptr<int>(), hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
        reward.data_ptr<float>(), n);
    check_cuda("env_move");
    uint64_t* rng_ptr = reinterpret_cast<uint64_t*>(rng_c.data_ptr<int64_t>());
    respawn_kernel<<<blocks, threads>>>(hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
                                        rng_ptr, food_i.data_ptr<int>(), n);
    check_cuda("respawn");
    reset_flag_kernel<<<1, 1>>>(hit_any.data_ptr<int>());
    check_cuda("reset");
    auto agent_f_out = agent_o.to(torch::kFloat32);
    auto food_f_out = food_i.to(torch::kFloat32);
    return {agent_f_out, food_f_out, reward, rng_c};
}

// Full rollout. Returns (rewards, positions, last_logits).
std::vector<torch::Tensor> rollout(
    torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_gru,
    torch::Tensor w_a, torch::Tensor b_a, torch::Tensor w_v, torch::Tensor b_v,
    torch::Tensor agent_cpu, torch::Tensor food_cpu, torch::Tensor rng,
    int64_t num_envs, int64_t horizon)
{
    int N = (int)num_envs;
    int T = (int)horizon;
    auto opts = torch::TensorOptions().device(torch::kCUDA);
    auto w_permT = permute_gru(w_gru);
    auto agent = agent_cpu.to(torch::kCUDA).to(torch::kInt32).contiguous();
    auto food = food_cpu.to(torch::kCUDA).to(torch::kInt32).contiguous();
    auto rng_i = rng.to(torch::kCUDA).contiguous();
    uint64_t* rng_ptr = reinterpret_cast<uint64_t*>(rng_i.data_ptr<int64_t>());

    auto state = torch::zeros({N, GRU_LAYERS * HIDDEN}, opts.dtype(torch::kFloat32));
    auto state2 = torch::zeros({N, GRU_LAYERS * HIDDEN}, opts.dtype(torch::kFloat32));
    auto logits = torch::empty({N, NUM_ACTIONS}, opts.dtype(torch::kFloat32));
    auto value = torch::empty({N}, opts.dtype(torch::kFloat32));
    auto actions = torch::empty({N}, opts.dtype(torch::kInt32));
    auto hit = torch::zeros({N}, opts.dtype(torch::kUInt8));
    auto hit_any = torch::zeros({1}, opts.dtype(torch::kInt32));
    auto rewards = torch::zeros({N}, opts.dtype(torch::kFloat32));
    auto agent_tmp = torch::empty_like(agent);

    ensure_smem(SHMEM_BYTES);
    int blocks = (N + BN - 1) / BN;

    for (int t = 0; t < T; ++t) {
        policy_tiled_kernel<<<blocks, THREADS, SHMEM_BYTES>>>(
            wptr(w_enc), wptr(b_enc), wptr(w_permT),
            wptr(w_a), wptr(b_a), wptr(w_v), wptr(b_v),
            nullptr, agent.data_ptr<int>(), food.data_ptr<int>(),
            state.data_ptr<float>(), state2.data_ptr<float>(),
            logits.data_ptr<float>(), value.data_ptr<float>(),
            actions.data_ptr<int>(), N, 1);
        check_cuda("policy");

        env_move_kernel<<<blocks, 256>>>(
            agent.data_ptr<int>(), food.data_ptr<int>(), actions.data_ptr<int>(),
            agent_tmp.data_ptr<int>(), hit.data_ptr<unsigned char>(),
            hit_any.data_ptr<int>(), rewards.data_ptr<float>(), N);
        check_cuda("move");

        respawn_kernel<<<blocks, 256>>>(
            hit.data_ptr<unsigned char>(), hit_any.data_ptr<int>(),
            rng_ptr, food.data_ptr<int>(), N);
        check_cuda("respawn");
        reset_flag_kernel<<<1, 1>>>(hit_any.data_ptr<int>());
        check_cuda("reset");

        std::swap(state, state2);
        std::swap(agent, agent_tmp);
    }

    auto positions = agent.to(torch::kInt64);
    return {rewards, positions, logits};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("policy_forward", &policy_forward, "policy forward");
    m.def("env_step", &env_step, "env step");
    m.def("rollout", &rollout, "rollout");
}