from __future__ import annotations

import os
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load

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

BOARD = 11
OBS_DIM = 4
HIDDEN = 256
GRU_LAYERS = 3
NUM_ACTIONS = 4
GRU_OUT = 3 * HIDDEN
ENV_BLOCK = 128
CHUNK_ENVS = 16384
RNG_SEED_STRIDE = 10007

_ext = None


def _module():
    global _ext
    if _ext is None:
        os.environ["TORCH_CUDA_ARCH_LIST"] = "12.0"
        build_dir = HERE / "build"
        build_dir.mkdir(exist_ok=True)
        _ext = load(
            name="grid_mingru_ext",
            sources=[str(HERE / "grid_mingru.cu"), str(HERE / "grid_mingru_bind.cpp")],
            extra_cuda_cflags=["-O3"],
            build_directory=str(build_dir),
            verbose=False,
        )
    return _ext


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 policy_forward(model: Model, obs: torch.Tensor, state: torch.Tensor):
    logits, new_state, value = _module().policy(
        obs.contiguous(),
        state.contiguous(),
        model.w_enc,
        model.b_enc,
        model.w_gru,
        model.w_a,
        model.b_a,
        model.w_v,
        model.b_v,
    )
    return logits, new_state, value


def env_step(agent, food, actions, rng_state):
    n = agent.shape[0]
    agent_out = torch.empty_like(agent)
    food_out = torch.empty_like(food)
    reward = torch.empty(n, dtype=torch.float32, device=agent.device)
    rng_out = torch.empty_like(rng_state)
    any_hit = torch.zeros(1, dtype=torch.int32, device=agent.device)
    _module().env_step(
        agent.contiguous(),
        food.contiguous(),
        actions.to(torch.int64).contiguous(),
        rng_state.contiguous(),
        agent_out,
        food_out,
        reward,
        rng_out,
        any_hit,
    )
    return agent_out, food_out, reward, rng_out


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()

    n = int(num_envs)
    npad = -(-n // ENV_BLOCK) * ENV_BLOCK
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    agent = torch.zeros(npad, 2, dtype=torch.int32, device=device)
    food = torch.zeros(npad, 2, dtype=torch.int32, device=device)
    agent[:n] = torch.randint(0, BOARD, (n, 2), generator=g).to(device=device, dtype=torch.int32)
    food[:n] = torch.randint(0, BOARD, (n, 2), generator=g).to(device=device, dtype=torch.int32)
    rewards = torch.zeros(npad, dtype=torch.float32, device=device)
    state = torch.empty(npad, GRU_LAYERS, HIDDEN, dtype=torch.float32, device=device)
    logits = torch.zeros(npad, NUM_ACTIONS, dtype=torch.float32, device=device)
    ctr = torch.zeros(2, dtype=torch.int32, device=device)

    _module().rollout(
        model.w_enc,
        model.b_enc,
        model.w_gru,
        model.w_a,
        model.b_a,
        agent,
        food,
        rewards,
        state,
        logits,
        ctr,
        seed * RNG_SEED_STRIDE,
        n,
        horizon,
        CHUNK_ENVS,
    )
    return {
        "rewards": rewards[:n],
        "positions": agent[:n].to(torch.int64),
        "last_logits": logits[:n],
    }


# ==================================================================
# ===== sidecar: grid_mingru.cu (21547 bytes, loaded by solution.py) =====
# ==================================================================

#include <cuda_runtime.h>
#include <cstddef>
#include <cstdint>

namespace {

constexpr int NACT = 4;
constexpr int HID = 256;
constexpr int GATES = 3 * HID;
constexpr int BOARD = 11;
constexpr unsigned long long LCG_A = 6364136223846793005ULL;
constexpr unsigned long long LCG_MASK = 0x7FFFFFFFFFFFFFFFULL;

constexpr int L0_ENVS = 32;
constexpr int L0_THREADS = 256;

constexpr int GM = 128;
constexpr int GU = 64;
constexpr int GT = 256;
constexpr int GKC = 16;
constexpr int NKC = HID / GKC;
constexpr int HLD = GM;
constexpr int WLD = 196;

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

__device__ __forceinline__ unsigned long long lcg_step(unsigned long long x) {
    return (x * LCG_A + 1ULL) & LCG_MASK;
}

__device__ __forceinline__ void cp_async16(void* smem, const void* gmem) {
    const unsigned s = static_cast<unsigned>(__cvta_generic_to_shared(smem));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(s), "l"(gmem));
}
__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;\n" ::); }
template <int N>
__device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }

// Advances the env rng by 2K global steps in closed form, then draws the respawn food.
__device__ __forceinline__ void respawn_food(unsigned long long rng_off, int env, int K, int& fx, int& fy) {
    unsigned long long ra = 1ULL, rb = 0ULL, pa = LCG_A, pb = 1ULL;
    for (unsigned n = 2u * static_cast<unsigned>(K); n; n >>= 1) {
        if (n & 1u) {
            rb = ra * pb + rb;
            ra = ra * pa;
        }
        pb = pa * pb + pb;
        pa = pa * pa;
    }
    const unsigned long long x0 = static_cast<unsigned long long>(env) + rng_off;
    const unsigned long long x = (ra * x0 + rb) & LCG_MASK;
    const unsigned long long r1 = lcg_step(x);
    const unsigned long long r2 = lcg_step(r1);
    fx = static_cast<int>(r1 % BOARD);
    fy = static_cast<int>(r2 % BOARD);
}

__global__ void precompute_m0(const float* __restrict__ w0, const float* __restrict__ w_enc,
                              const float* __restrict__ b_enc, float* __restrict__ m0,
                              float* __restrict__ c0) {
    __shared__ float red[5][HID];
    const int j = blockIdx.x;
    const int k = threadIdx.x;
    const float w = w0[j * HID + k];
    red[0][k] = w * w_enc[k * 4 + 0];
    red[1][k] = w * w_enc[k * 4 + 1];
    red[2][k] = w * w_enc[k * 4 + 2];
    red[3][k] = w * w_enc[k * 4 + 3];
    red[4][k] = w * b_enc[k];
    __syncthreads();
    for (int s = HID / 2; s > 0; s >>= 1) {
        if (k < s) {
#pragma unroll
            for (int c = 0; c < 5; ++c) red[c][k] += red[c][k + s];
        }
        __syncthreads();
    }
    if (k == 0) {
        for (int c = 0; c < 4; ++c) m0[j * 4 + c] = red[c][0];
        c0[j] = red[4][0];
    }
}

// Layer 0 has a 4-wide input, so its gates are M0 * obs + c0 and its input h0 = W_enc * obs + b_enc.
__global__ void __launch_bounds__(L0_THREADS) layer0_kernel(
    const int* __restrict__ agent, const int* __restrict__ food, const float* __restrict__ m0,
    const float* __restrict__ c0, const float* __restrict__ w_enc, const float* __restrict__ b_enc,
    float* __restrict__ state, float* __restrict__ h0, int* ctr, int npad, int env_base, bool bump, bool first) {
    __shared__ __align__(16) float sM0[GATES * 4];
    __shared__ __align__(16) float sC0[GATES];
    __shared__ float tile[HID][L0_ENVS];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int e0 = env_base + blockIdx.x * L0_ENVS;

    if (bump && blockIdx.x == 0 && threadIdx.x == 0) {
        ctr[0] += ctr[1];
        ctr[1] = 0;
    }

    float stv[HID / 32][4];
#pragma unroll
    for (int p = 0; p < HID / 32; ++p)
#pragma unroll
        for (int q = 0; q < 4; ++q)
            stv[p][q] = first ? 0.f : state[static_cast<size_t>(e0 + warp + 8 * q) * GATES + lane + 32 * p];
    for (int i = threadIdx.x; i < GATES; i += L0_THREADS)
        reinterpret_cast<float4*>(sM0)[i] = reinterpret_cast<const float4*>(m0)[i];
    for (int i = threadIdx.x; i < GATES / 4; i += L0_THREADS)
        reinterpret_cast<float4*>(sC0)[i] = reinterpret_cast<const float4*>(c0)[i];
    __syncthreads();

    float ob[4][4];
#pragma unroll
    for (int q = 0; q < 4; ++q) {
        const int env = e0 + warp + 8 * q;
        const int ax = agent[2 * env];
        const int ay = agent[2 * env + 1];
        const int fx = food[2 * env];
        const int fy = food[2 * env + 1];
        ob[q][0] = (float)(fx - ax) / (float)BOARD;
        ob[q][1] = (float)(fy - ay) / (float)BOARD;
        ob[q][2] = (float)ax / (float)(BOARD - 1);
        ob[q][3] = (float)ay / (float)(BOARD - 1);
    }

#pragma unroll
    for (int p = 0; p < HID / 32; ++p) {
        const int u = lane + 32 * p;
        const float4 mz = *reinterpret_cast<const float4*>(&sM0[u * 4]);
        const float4 mg = *reinterpret_cast<const float4*>(&sM0[(HID + u) * 4]);
        const float4 mp = *reinterpret_cast<const float4*>(&sM0[(2 * HID + u) * 4]);
        const float cz0 = sC0[u];
        const float cz1 = sC0[HID + u];
        const float cz2 = sC0[2 * HID + u];
        const float4 we = *reinterpret_cast<const float4*>(&w_enc[u * 4]);
        const float be = b_enc[u];
#pragma unroll
        for (int q = 0; q < 4; ++q) {
            const int le = warp + 8 * q;
            const int env = e0 + le;
            const float o0 = ob[q][0], o1 = ob[q][1], o2 = ob[q][2], o3 = ob[q][3];
            const float zh = fmaf(mz.w, o3, fmaf(mz.z, o2, fmaf(mz.y, o1, fmaf(mz.x, o0, cz0))));
            const float zg = fmaf(mg.w, o3, fmaf(mg.z, o2, fmaf(mg.y, o1, fmaf(mg.x, o0, cz1))));
            const float zp = fmaf(mp.w, o3, fmaf(mp.z, o2, fmaf(mp.y, o1, fmaf(mp.x, o0, cz2))));
            const float hin = fmaf(we.w, o3, fmaf(we.z, o2, fmaf(we.y, o1, fmaf(we.x, o0, be))));
            float* sp = state + static_cast<size_t>(env) * GATES + u;
            const float st = stv[p][q];
            const float out = st + sigm(zg) * (tanhf(zh) - st);
            const float pz = sigm(zp);
            tile[u][le ^ lane] = fmaf(pz, out, (1.f - pz) * hin);
            *sp = out;
        }
    }
    __syncthreads();
    for (int idx = threadIdx.x; idx < HID * L0_ENVS; idx += L0_THREADS) {
        const int u = idx / L0_ENVS;
        const int c = idx % L0_ENVS;
        h0[static_cast<size_t>(u) * npad + e0 + c] = tile[u][c ^ (u & 31)];
    }
}

// One CTA owns GM envs and GU hidden units (3*GU gate rows). Each thread owns 8 envs x 4 units x 3 gates.
// Layers 1 and 2 use this kernel. Layer 2 writes per-slice partial logits instead of h.
template <bool FINAL>
__global__ void __launch_bounds__(GT) gru_gemm(const float* __restrict__ h_in, float* __restrict__ h_out,
                                              const float* __restrict__ w, float* __restrict__ state, int layer,
                                              bool first, const float* __restrict__ w_a,
                                              float* __restrict__ partial, int npad, int env_base) {
    __shared__ __align__(16) float Hs[2][GKC][HLD];
    __shared__ __align__(16) float Ws[2][GKC][WLD];

    const int tid = threadIdx.x;
    const int ty = tid >> 4;
    const int tx = tid & 15;
    const int env0 = env_base + blockIdx.x * GM;
    const int ub = blockIdx.y * GU;

    int wrow[3], wkq[3], wcol[3];
#pragma unroll
    for (int q = 0; q < 3; ++q) {
        const int p = tid + GT * q;
        const int c = p >> 2;
        wcol[q] = c;
        wkq[q] = p & 3;
        wrow[q] = (c % 3) * HID + ub + c / 3;
    }

    float acc[8][4][3];
#pragma unroll
    for (int j = 0; j < 8; ++j)
#pragma unroll
        for (int i = 0; i < 4; ++i)
#pragma unroll
            for (int g = 0; g < 3; ++g) acc[j][i][g] = 0.f;

    float4 wr[3];

    auto load_w = [&](int kc) {
#pragma unroll
        for (int q = 0; q < 3; ++q)
            wr[q] = *reinterpret_cast<const float4*>(w + static_cast<size_t>(wrow[q]) * HID + kc * GKC + 4 * wkq[q]);
    };
    auto store_w = [&](int buf) {
#pragma unroll
        for (int q = 0; q < 3; ++q) {
            const float* v = reinterpret_cast<const float*>(&wr[q]);
#pragma unroll
            for (int j = 0; j < 4; ++j) Ws[buf][4 * wkq[q] + j][wcol[q]] = v[j];
        }
    };
    auto issue_h = [&](int kc, int buf) {
#pragma unroll
        for (int q = 0; q < 2; ++q) {
            const int p = tid + GT * q;
            const int kr = p >> 5;
            const int c4 = p & 31;
            cp_async16(&Hs[buf][kr][4 * c4], h_in + static_cast<size_t>(kc * GKC + kr) * npad + env0 + 4 * c4);
        }
    };
    auto compute = [&](int buf) {
#pragma unroll
        for (int kk = 0; kk < GKC; ++kk) {
            const float4 h0 = *reinterpret_cast<const float4*>(&Hs[buf][kk][ty * 8]);
            const float4 h1 = *reinterpret_cast<const float4*>(&Hs[buf][kk][ty * 8 + 4]);
            const float4 w0 = *reinterpret_cast<const float4*>(&Ws[buf][kk][tx * 12]);
            const float4 w1 = *reinterpret_cast<const float4*>(&Ws[buf][kk][tx * 12 + 4]);
            const float4 w2 = *reinterpret_cast<const float4*>(&Ws[buf][kk][tx * 12 + 8]);
            const float hv[8] = {h0.x, h0.y, h0.z, h0.w, h1.x, h1.y, h1.z, h1.w};
            const float wv[12] = {w0.x, w0.y, w0.z, w0.w, w1.x, w1.y, w1.z, w1.w, w2.x, w2.y, w2.z, w2.w};
#pragma unroll
            for (int j = 0; j < 8; ++j)
#pragma unroll
                for (int i = 0; i < 4; ++i)
#pragma unroll
                    for (int g = 0; g < 3; ++g) acc[j][i][g] = fmaf(hv[j], wv[3 * i + g], acc[j][i][g]);
        }
    };

    issue_h(0, 0);
    cp_async_commit();
    load_w(0);
    store_w(0);
#pragma unroll 1
    for (int kc = 0; kc < NKC; ++kc) {
        const int buf = kc & 1;
        cp_async_wait<0>();
        __syncthreads();
        if (kc + 1 < NKC) {
            issue_h(kc + 1, buf ^ 1);
            cp_async_commit();
            load_w(kc + 1);
        }
        compute(buf);
        if (kc + 1 < NKC) store_w(buf ^ 1);
    }

    float hin[4][8];
#pragma unroll
    for (int i = 0; i < 4; ++i) {
        const float* hp = h_in + static_cast<size_t>(ub + 4 * tx + i) * npad + env0 + ty * 8;
        const float4 a = *reinterpret_cast<const float4*>(hp);
        const float4 b = *reinterpret_cast<const float4*>(hp + 4);
        hin[i][0] = a.x; hin[i][1] = a.y; hin[i][2] = a.z; hin[i][3] = a.w;
        hin[i][4] = b.x; hin[i][5] = b.y; hin[i][6] = b.z; hin[i][7] = b.w;
    }

    float hout[4][8];
#pragma unroll
    for (int j = 0; j < 8; ++j) {
        const int env = env0 + ty * 8 + j;
        float* sp = state + static_cast<size_t>(env) * GATES + layer * HID + ub + 4 * tx;
        float st[4] = {0.f, 0.f, 0.f, 0.f};
        if (!first) {
            const float4 s4 = *reinterpret_cast<const float4*>(sp);
            st[0] = s4.x; st[1] = s4.y; st[2] = s4.z; st[3] = s4.w;
        }
        float outv[4];
#pragma unroll
        for (int i = 0; i < 4; ++i) {
            const float zh = acc[j][i][0];
            const float zg = acc[j][i][1];
            const float zp = acc[j][i][2];
            outv[i] = st[i] + sigm(zg) * (tanhf(zh) - st[i]);
            const float p = sigm(zp);
            hout[i][j] = fmaf(p, outv[i], (1.f - p) * hin[i][j]);
        }
        *reinterpret_cast<float4*>(sp) = make_float4(outv[0], outv[1], outv[2], outv[3]);
    }

    if (!FINAL) {
#pragma unroll
        for (int i = 0; i < 4; ++i) {
            float* op = h_out + static_cast<size_t>(ub + 4 * tx + i) * npad + env0 + ty * 8;
            *reinterpret_cast<float4*>(op) = make_float4(hout[i][0], hout[i][1], hout[i][2], hout[i][3]);
            *reinterpret_cast<float4*>(op + 4) = make_float4(hout[i][4], hout[i][5], hout[i][6], hout[i][7]);
        }
    } else {
        float pa[8][NACT];
#pragma unroll
        for (int j = 0; j < 8; ++j)
#pragma unroll
            for (int a = 0; a < NACT; ++a) pa[j][a] = 0.f;
#pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int u = ub + 4 * tx + i;
#pragma unroll
            for (int a = 0; a < NACT; ++a) {
                const float wa = w_a[a * HID + u];
#pragma unroll
                for (int j = 0; j < 8; ++j) pa[j][a] = fmaf(wa, hout[i][j], pa[j][a]);
            }
        }
#pragma unroll
        for (int j = 0; j < 8; ++j)
#pragma unroll
            for (int a = 0; a < NACT; ++a) {
                float v = pa[j][a];
                v += __shfl_xor_sync(0xffffffffu, v, 8);
                v += __shfl_xor_sync(0xffffffffu, v, 4);
                v += __shfl_xor_sync(0xffffffffu, v, 2);
                v += __shfl_xor_sync(0xffffffffu, v, 1);
                pa[j][a] = v;
            }
        if (tx == 0) {
#pragma unroll
            for (int j = 0; j < 8; ++j)
#pragma unroll
                for (int a = 0; a < NACT; ++a)
                    partial[(static_cast<size_t>(blockIdx.y) * NACT + a) * npad + env0 + ty * 8 + j] = pa[j][a];
        }
    }
}

__global__ void head_kernel(const float* __restrict__ partial, const float* __restrict__ b_a,
                            int* __restrict__ agent, int* __restrict__ food, float* __restrict__ rewards,
                            float* __restrict__ logits, int* ctr, unsigned long long rng_off, int npad,
                            int n_valid, int env_base) {
    const int e = env_base + blockIdx.x * blockDim.x + threadIdx.x;
    if (e >= npad) return;
    float lg[NACT];
#pragma unroll
    for (int a = 0; a < NACT; ++a) {
        float s = b_a[a];
#pragma unroll
        for (int sl = 0; sl < 4; ++sl) s += partial[(static_cast<size_t>(sl) * NACT + a) * npad + e];
        lg[a] = s;
    }
    int act = 0;
    float best = lg[0];
    if (lg[1] > best) { best = lg[1]; act = 1; }
    if (lg[2] > best) { best = lg[2]; act = 2; }
    if (lg[3] > best) { best = lg[3]; act = 3; }
    if (logits != nullptr) {
#pragma unroll
        for (int a = 0; a < NACT; ++a) logits[e * NACT + a] = lg[a];
    }

    int ax = agent[2 * e];
    int ay = agent[2 * e + 1];
    const int fx0 = food[2 * e];
    const int fy0 = food[2 * e + 1];
    if (act == 0) ay -= 1;
    else if (act == 1) ay += 1;
    else if (act == 2) ax -= 1;
    else ax += 1;
    ax = min(max(ax, 0), BOARD - 1);
    ay = min(max(ay, 0), BOARD - 1);
    agent[2 * e] = ax;
    agent[2 * e + 1] = ay;

    int fx = fx0;
    int fy = fy0;
    if (e < n_valid && ax == fx0 && ay == fy0) {
        atomicOr(&ctr[1], 1);
        respawn_food(rng_off, e, ctr[0], fx, fy);
        rewards[e] += 1.f;
    }
    food[2 * e] = fx;
    food[2 * e + 1] = fy;
}

__global__ void policy_naive(const float* __restrict__ obs, const float* __restrict__ st_in,
                             float* __restrict__ st_out, float* __restrict__ logits, float* __restrict__ value,
                             const float* __restrict__ w_enc, const float* __restrict__ b_enc,
                             const float* __restrict__ w_gru, const float* __restrict__ w_a,
                             const float* __restrict__ b_a, const float* __restrict__ w_v,
                             const float* __restrict__ b_v, int n) {
    const int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= n) return;
    float h[HID];
    float hn[HID];
    const float o0 = obs[4 * env], o1 = obs[4 * env + 1], o2 = obs[4 * env + 2], o3 = obs[4 * env + 3];
    for (int u = 0; u < HID; ++u) {
        const float* we = w_enc + 4 * u;
        h[u] = b_enc[u] + we[0] * o0 + we[1] * o1 + we[2] * o2 + we[3] * o3;
    }
    for (int l = 0; l < 3; ++l) {
        const float* W = w_gru + (size_t)l * GATES * HID;
        for (int u = 0; u < HID; ++u) {
            const float* rh = W + (size_t)u * HID;
            const float* rg = W + (size_t)(HID + u) * HID;
            const float* rp = W + (size_t)(2 * HID + u) * HID;
            float zh = 0.f, zg = 0.f, zp = 0.f;
            for (int k = 0; k < HID; ++k) {
                zh = fmaf(rh[k], h[k], zh);
                zg = fmaf(rg[k], h[k], zg);
                zp = fmaf(rp[k], h[k], zp);
            }
            const size_t si = ((size_t)env * 3 + l) * HID + u;
            const float st = st_in[si];
            const float out = st + sigm(zg) * (tanhf(zh) - st);
            const float p = sigm(zp);
            st_out[si] = out;
            hn[u] = fmaf(p, out, (1.f - p) * h[u]);
        }
        for (int u = 0; u < HID; ++u) h[u] = hn[u];
    }
    for (int ac = 0; ac < NACT; ++ac) {
        float s = b_a[ac];
        for (int k = 0; k < HID; ++k) s = fmaf(w_a[ac * HID + k], h[k], s);
        logits[env * NACT + ac] = s;
    }
    float v = b_v[0];
    for (int k = 0; k < HID; ++k) v = fmaf(w_v[k], h[k], v);
    value[env] = v;
}

__global__ void env_move(const float* __restrict__ agent, const float* __restrict__ food,
                         const int64_t* __restrict__ act, float* __restrict__ agent_out,
                         float* __restrict__ reward, int* __restrict__ any, int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float ax = agent[2 * i];
    float ay = agent[2 * i + 1];
    const long long ac = act[i];
    if (ac == 0) ay -= 1.f;
    else if (ac == 1) ay += 1.f;
    else if (ac == 2) ax -= 1.f;
    else if (ac == 3) ax += 1.f;
    ax = fminf(fmaxf(ax, 0.f), (float)(BOARD - 1));
    ay = fminf(fmaxf(ay, 0.f), (float)(BOARD - 1));
    agent_out[2 * i] = ax;
    agent_out[2 * i + 1] = ay;
    const bool hit = (ax == food[2 * i]) && (ay == food[2 * i + 1]);
    reward[i] = hit ? 1.f : 0.f;
    if (hit) atomicOr(any, 1);
}

__global__ void env_rng(const float* __restrict__ agent_new, const float* __restrict__ food,
                        float* __restrict__ food_out, const int64_t* __restrict__ rng,
                        int64_t* __restrict__ rng_out, const int* __restrict__ any, int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    const unsigned long long r = (unsigned long long)rng[i];
    if (*any) {
        const unsigned long long r1 = lcg_step(r);
        const unsigned long long r2 = lcg_step(r1);
        rng_out[i] = (int64_t)r2;
        const bool hit = (agent_new[2 * i] == food[2 * i]) && (agent_new[2 * i + 1] == food[2 * i + 1]);
        if (hit) {
            food_out[2 * i] = (float)(r1 % BOARD);
            food_out[2 * i + 1] = (float)(r2 % BOARD);
        } else {
            food_out[2 * i] = food[2 * i];
            food_out[2 * i + 1] = food[2 * i + 1];
        }
    } else {
        rng_out[i] = rng[i];
        food_out[2 * i] = food[2 * i];
        food_out[2 * i + 1] = food[2 * i + 1];
    }
}

}  // namespace

void rollout_launch(const float* w_enc, const float* b_enc, const float* w_gru, const float* w_a, const float* b_a,
                    int* agent, int* food, float* rewards, float* state, float* logits, float* m0, float* c0,
                    float* h_a, float* h_b, float* partial, int* ctr, unsigned long long rng_off, int n_valid,
                    int npad, int horizon, int chunk, void* stream_v) {
    cudaStream_t stream = static_cast<cudaStream_t>(stream_v);
    precompute_m0<<<GATES, HID, 0, stream>>>(w_gru, w_enc, b_enc, m0, c0);
    for (int t = 0; t < horizon; ++t) {
        const bool first = (t == 0);
        const bool last = (t == horizon - 1);
        for (int eb = 0; eb < npad; eb += chunk) {
            const int cur = (npad - eb < chunk) ? (npad - eb) : chunk;
            layer0_kernel<<<cur / L0_ENVS, L0_THREADS, 0, stream>>>(agent, food, m0, c0, w_enc, b_enc, state, h_a,
                                                                    ctr, npad, eb, eb == 0, first);
            gru_gemm<false><<<dim3(cur / GM, HID / GU), GT, 0, stream>>>(h_a, h_b, w_gru + GATES * HID, state, 1,
                                                                         first, nullptr, nullptr, npad, eb);
            gru_gemm<true><<<dim3(cur / GM, HID / GU), GT, 0, stream>>>(h_b, nullptr, w_gru + 2 * GATES * HID, state,
                                                                        2, first, w_a, partial, npad, eb);
            head_kernel<<<cur / 32, 32, 0, stream>>>(partial, b_a, agent, food, rewards, last ? logits : nullptr, ctr,
                                                     rng_off, npad, n_valid, eb);
        }
    }
}

void policy_launch(const float* obs, const float* st_in, float* st_out, float* logits, float* value,
                   const float* w_enc, const float* b_enc, const float* w_gru, const float* w_a, const float* b_a,
                   const float* w_v, const float* b_v, int n, void* stream_v) {
    cudaStream_t stream = static_cast<cudaStream_t>(stream_v);
    const int block = 128;
    policy_naive<<<(n + block - 1) / block, block, 0, stream>>>(obs, st_in, st_out, logits, value, w_enc, b_enc,
                                                                w_gru, w_a, b_a, w_v, b_v, n);
}

void env_step_launch(const float* agent, const float* food, const int64_t* actions, float* agent_out,
                     float* food_out, float* reward, const int64_t* rng, int64_t* rng_out, int* any, int n,
                     void* stream_v) {
    cudaStream_t stream = static_cast<cudaStream_t>(stream_v);
    const int block = 128;
    const int grid = (n + block - 1) / block;
    env_move<<<grid, block, 0, stream>>>(agent, food, actions, agent_out, reward, any, n);
    env_rng<<<grid, block, 0, stream>>>(agent_out, food, food_out, rng, rng_out, any, n);
}


# ==================================================================
# ===== sidecar: grid_mingru_bind.cpp (4452 bytes, loaded by solution.py) =====
# ==================================================================

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <vector>

void rollout_launch(const float* w_enc, const float* b_enc, const float* w_gru, const float* w_a, const float* b_a,
                    int* agent, int* food, float* rewards, float* state, float* logits, float* m0, float* c0,
                    float* h_a, float* h_b, float* partial, int* ctr, unsigned long long rng_off, int n_valid,
                    int npad, int horizon, int chunk, void* stream);
void policy_launch(const float* obs, const float* st_in, float* st_out, float* logits, float* value,
                   const float* w_enc, const float* b_enc, const float* w_gru, const float* w_a, const float* b_a,
                   const float* w_v, const float* b_v, int n, void* stream);
void env_step_launch(const float* agent, const float* food, const int64_t* actions, float* agent_out,
                     float* food_out, float* reward, const int64_t* rng, int64_t* rng_out, int* any, int n,
                     void* stream);

namespace {
constexpr int64_t HID = 256;
constexpr int64_t GATES = 768;
constexpr int64_t NACT = 4;
}  // namespace

void rollout(torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_gru, torch::Tensor w_a, torch::Tensor b_a,
             torch::Tensor agent, torch::Tensor food, torch::Tensor rewards, torch::Tensor state,
             torch::Tensor logits, torch::Tensor ctr, int64_t rng_off, int64_t n_valid, int64_t horizon,
             int64_t chunk) {
    const int64_t npad = agent.size(0);
    const auto opts = w_enc.options();
    torch::Tensor m0 = torch::empty({GATES, 4}, opts);
    torch::Tensor c0 = torch::empty({GATES}, opts);
    torch::Tensor h_a = torch::empty({HID, npad}, opts);
    torch::Tensor h_b = torch::empty({HID, npad}, opts);
    torch::Tensor partial = torch::empty({4, NACT, npad}, opts);
    rollout_launch(w_enc.data_ptr<float>(), b_enc.data_ptr<float>(), w_gru.data_ptr<float>(), w_a.data_ptr<float>(),
                   b_a.data_ptr<float>(), agent.data_ptr<int>(), food.data_ptr<int>(), rewards.data_ptr<float>(),
                   state.data_ptr<float>(), logits.data_ptr<float>(), m0.data_ptr<float>(), c0.data_ptr<float>(),
                   h_a.data_ptr<float>(), h_b.data_ptr<float>(), partial.data_ptr<float>(), ctr.data_ptr<int>(),
                   static_cast<unsigned long long>(rng_off), static_cast<int>(n_valid), static_cast<int>(npad),
                   static_cast<int>(horizon), static_cast<int>(chunk), at::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

std::vector<torch::Tensor> policy(torch::Tensor obs, torch::Tensor st_in, 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) {
    const int64_t n = obs.size(0);
    const auto opts = obs.options();
    torch::Tensor logits = torch::empty({n, NACT}, opts);
    torch::Tensor st_out = torch::empty_like(st_in);
    torch::Tensor value = torch::empty({n}, opts);
    policy_launch(obs.data_ptr<float>(), st_in.data_ptr<float>(), st_out.data_ptr<float>(), logits.data_ptr<float>(),
                  value.data_ptr<float>(), w_enc.data_ptr<float>(), b_enc.data_ptr<float>(), w_gru.data_ptr<float>(),
                  w_a.data_ptr<float>(), b_a.data_ptr<float>(), w_v.data_ptr<float>(), b_v.data_ptr<float>(),
                  static_cast<int>(n), at::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {logits, st_out, value};
}

void env_step(torch::Tensor agent, torch::Tensor food, torch::Tensor actions, torch::Tensor rng,
              torch::Tensor agent_out, torch::Tensor food_out, torch::Tensor reward, torch::Tensor rng_out,
              torch::Tensor any) {
    env_step_launch(agent.data_ptr<float>(), food.data_ptr<float>(), actions.data_ptr<int64_t>(),
                    agent_out.data_ptr<float>(), food_out.data_ptr<float>(), reward.data_ptr<float>(),
                    rng.data_ptr<int64_t>(), rng_out.data_ptr<int64_t>(), any.data_ptr<int>(),
                    static_cast<int>(agent.size(0)), at::cuda::getCurrentCUDAStream().stream());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

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