from __future__ import annotations

import os
import sys
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

os.environ["PATH"] = os.path.dirname(sys.executable) + os.pathsep + os.environ["PATH"]


_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cooperative_groups.h>

#define HINDEX(row, column) ((column) ^ (((row) & 7) * 4))

using torch::Tensor;

__device__ __forceinline__ unsigned __half2_as_uint(__half2 value) {
    return reinterpret_cast<unsigned&>(value);
}

__device__ __forceinline__ __half2 __uint_as_half2(unsigned value) {
    return reinterpret_cast<__half2&>(value);
}

__device__ __forceinline__ unsigned pack_half(float first, float second) {
    return __half2_as_uint(__floats2half2_rn(first, second));
}

__global__ void pack_weights(const float* weights, uint4* packed) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= 3 * 16 * 3 * 32 * 32) return;
    int lane = index % 32;
    int tile = index / 32 % 32;
    int gate = index / (32 * 32) % 3;
    int inner = index / (32 * 32 * 3) % 16;
    int layer = index / (32 * 32 * 3 * 16);
    if (layer == 0) return;
    int neuron = tile * 8 + lane / 4;
    int column = inner * 16 + lane % 4 * 2;
    const float* source = weights + (layer * 768 + gate * 256 + neuron) * 256 + column;
    float first = source[0], second = source[1];
    float third = source[8], fourth = source[9];
    unsigned high0 = pack_half(first, second);
    unsigned high1 = pack_half(third, fourth);
    float2 recon0 = __half22float2(__uint_as_half2(high0));
    float2 recon1 = __half22float2(__uint_as_half2(high1));
    packed[index] = make_uint4(high0, high1, pack_half(first-recon0.x, second-recon0.y),
        pack_half(third-recon1.x, fourth-recon1.y));
}

__global__ void compose_encoder(const float* weights, const float* enc, const float* enc_bias, float* composed) {
    int lane = threadIdx.x % 32;
    int neuron = blockIdx.x * 8 + threadIdx.x / 32;
    float sums[5] = {};
    #pragma unroll
    for (int inner = lane; inner < 256; inner += 32) {
        float weight = weights[neuron * 256 + inner];
        #pragma unroll
        for (int channel = 0; channel < 4; ++channel)
            sums[channel] = fmaf(weight, enc[inner * 4 + channel], sums[channel]);
        sums[4] = fmaf(weight, enc_bias[inner], sums[4]);
    }
    #pragma unroll
    for (int channel = 0; channel < 5; ++channel) {
        #pragma unroll
        for (int offset = 16; offset > 0; offset /= 2)
            sums[channel] += __shfl_down_sync(0xffffffff, sums[channel], offset);
        if (lane == 0) composed[neuron * 5 + channel] = sums[channel];
    }
}

__device__ __forceinline__ void mma(float (&acc)[4], const unsigned (&left)[4], unsigned right0, unsigned right1) {
    asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
        : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3])
        : "r"(left[0]), "r"(left[1]), "r"(left[2]), "r"(left[3]), "r"(right0), "r"(right1));
}

__device__ __forceinline__ void load_matrix(unsigned (&fragment)[4], const unsigned (*source)[128], int rowtile, int inner, int lane) {
    int row = rowtile * 16 + lane / 4;
    int column = inner * 8 + lane % 4;
    fragment[0] = source[row][HINDEX(row,column)];
    fragment[1] = source[row+8][HINDEX(row+8,column)];
    fragment[2] = source[row][HINDEX(row,column+4)];
    fragment[3] = source[row+8][HINDEX(row+8,column+4)];
}

__device__ __forceinline__ unsigned long long next_rng(unsigned long long rng) {
    return (rng * 6364136223846793005ULL + 1ULL) & 0x7fffffffffffffffULL;
}

template<int ROWS, int WARPS, bool ROLLOUT, bool KEEP_STATE = false>
__device__ __forceinline__ void fused_policy_body(const float* obs, const float* old_state, float* new_state,
    const float* __restrict__ enc, const float* __restrict__ enc_bias, const uint4* __restrict__ weights,
    const float* __restrict__ action_weight, const float* __restrict__ action_bias, const float* __restrict__ value_weight,
    const float* __restrict__ value_bias, float* logits, float* values, long long* agent,
    long long* food, unsigned long long* rng, float* rewards, int* flags,
    int count, int step, int last_step, unsigned long long seed, float4* register_state = nullptr) {
    constexpr int TILES = 32 / WARPS;
    constexpr int MTILES = ROWS / 16;
    extern __shared__ __align__(16) unsigned shared_storage[];
    auto high = reinterpret_cast<unsigned (*)[128]>(shared_storage);
    auto low = reinterpret_cast<unsigned (*)[128]>(shared_storage + ROWS * 128);
    auto observations = reinterpret_cast<float (*)[4]>(shared_storage + ROWS * 256);
    auto heads = reinterpret_cast<float (*)[5]>(shared_storage + ROWS * 260);
    auto shared_state = reinterpret_cast<float4*>(shared_storage + ROWS * 265);
    int thread = threadIdx.x;
    int lane = thread % 32;
    int warp = thread / 32;
    int base = blockIdx.x * ROWS;

    if (thread < ROWS) {
        int env = base + thread;
        if constexpr (ROLLOUT) {
            if (env < count) {
                if (step == 0) {
                    rng[env] = static_cast<unsigned long long>(env) + seed * 10007ULL;
                    rewards[env] = 0.0f;
                } else if (flags[step - 1]) {
                    unsigned long long first = next_rng(rng[env]);
                    unsigned long long second = next_rng(first);
                    rng[env] = second;
                    if (agent[env * 2] == food[env * 2] && agent[env * 2 + 1] == food[env * 2 + 1]) {
                        food[env * 2] = first % 11ULL;
                        food[env * 2 + 1] = second % 11ULL;
                    }
                }
                observations[thread][0] = float(food[env * 2] - agent[env * 2]) * (1.0f / 11.0f);
                observations[thread][1] = float(food[env * 2 + 1] - agent[env * 2 + 1]) * (1.0f / 11.0f);
                observations[thread][2] = float(agent[env * 2]) * 0.1f;
                observations[thread][3] = float(agent[env * 2 + 1]) * 0.1f;
            } else {
                for (int channel = 0; channel < 4; ++channel) observations[thread][channel] = 0;
            }
        } else {
            for (int channel = 0; channel < 4; ++channel)
                observations[thread][channel] = env < count ? obs[env * 4 + channel] : 0;
        }
    }
    __syncthreads();
    for (int index = thread; index < ROWS * 128; index += WARPS * 32) {
        int row = index / 128;
        int column = index % 128 * 2;
        float first = 0, second = 0;
        #pragma unroll
        for (int channel = 0; channel < 4; ++channel) {
            first = fmaf(observations[row][channel], enc[column * 4 + channel], first);
            second = fmaf(observations[row][channel], enc[(column + 1) * 4 + channel], second);
        }
        first += enc_bias[column];
        second += enc_bias[column+1];
        unsigned rounded = pack_half(first, second);
        float2 restored = __half22float2(__uint_as_half2(rounded));
        high[row][HINDEX(row,column/2)] = rounded;
        low[row][HINDEX(row,column/2)] = pack_half(first-restored.x, second-restored.y);
    }
    __syncthreads();

    #pragma unroll
    for (int layer = 0; layer < 3; ++layer) {
        float accum[MTILES][TILES][3][4] = {};
        if (layer == 0) {
            #pragma unroll
            for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
                #pragma unroll
                for (int tile = 0; tile < TILES; ++tile) {
                    #pragma unroll
                    for (int gate = 0; gate < 3; ++gate) {
                        #pragma unroll
                        for (int part = 0; part < 4; ++part) {
                            int row = rowtile * 16 + lane / 4 + (part / 2) * 8;
                            int column = gate * 256 + (warp * TILES + tile) * 8 + lane % 4 * 2 + part % 2;
                            const float* composed = reinterpret_cast<const float*>(weights) + column * 5;
                            float sum = 0;
                            #pragma unroll
                            for (int channel = 0; channel < 4; ++channel)
                                sum = fmaf(observations[row][channel], composed[channel], sum);
                            accum[rowtile][tile][gate][part] = sum + composed[4];
                        }
                    }
                }
            }
        } else {
            #pragma unroll 1
            for (int inner = 0; inner < 16; ++inner) {
                unsigned left_high[MTILES][4], left_low[MTILES][4];
                #pragma unroll
                for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
                    load_matrix(left_high[rowtile], high, rowtile, inner, lane);
                    load_matrix(left_low[rowtile], low, rowtile, inner, lane);
                }
                #pragma unroll
                for (int tile = 0; tile < TILES; ++tile) {
                    if constexpr (MTILES == 1) {
                        #pragma unroll
                        for (int gate = 0; gate < 3; ++gate) {
                            uint4 right = weights[(((layer * 16 + inner) * 3 + gate) * 32 + warp * TILES + tile) * 32 + lane];
                            mma(accum[0][tile][gate], left_high[0], right.z, right.w);
                            mma(accum[0][tile][gate], left_low[0], right.x, right.y);
                            mma(accum[0][tile][gate], left_high[0], right.x, right.y);
                        }
                    } else {
                        uint4 right[3];
                        #pragma unroll
                        for (int gate = 0; gate < 3; ++gate) {
                            right[gate] = weights[(((layer * 16 + inner) * 3 + gate) * 32 + warp * TILES + tile) * 32 + lane];
                        }
                        #pragma unroll
                        for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
                            #pragma unroll
                            for (int gate = 0; gate < 3; ++gate)
                                mma(accum[rowtile][tile][gate], left_high[rowtile], right[gate].z, right[gate].w);
                        }
                        #pragma unroll
                        for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
                            #pragma unroll
                            for (int gate = 0; gate < 3; ++gate)
                                mma(accum[rowtile][tile][gate], left_low[rowtile], right[gate].x, right[gate].y);
                        }
                        #pragma unroll
                        for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
                            #pragma unroll
                            for (int gate = 0; gate < 3; ++gate)
                                mma(accum[rowtile][tile][gate], left_high[rowtile], right[gate].x, right[gate].y);
                        }
                    }
                }
            }
        }
        __syncthreads();
        #pragma unroll
        for (int rowtile = 0; rowtile < MTILES; ++rowtile) {
            #pragma unroll
            for (int tile = 0; tile < TILES; ++tile) {
                int packed_state_index = ((((base/16 + rowtile) * 3 + layer) * 32 + warp * TILES + tile) * 32 + lane);
                float4 history = make_float4(0,0,0,0);
                float updated[4];
                if constexpr (ROLLOUT) {
                    if (step && base + rowtile * 16 < count) {
                        if constexpr (KEEP_STATE) {
                            if (layer == 2) history = register_state[rowtile * TILES + tile];
                            else history = shared_state[((rowtile * 2 + layer) * 32 + warp * TILES + tile) * 32 + lane];
                        } else history = reinterpret_cast<const float4*>(old_state)[packed_state_index];
                    }
                }
                #pragma unroll
                for (int rowpart = 0; rowpart < 2; ++rowpart) {
                    int row = rowtile * 16 + lane / 4 + rowpart * 8;
                    int column = (warp * TILES + tile) * 8 + lane % 4 * 2;
                    int env = base + row;
                    float2 old_high = __half22float2(__uint_as_half2(high[row][HINDEX(row,column/2)]));
                    float2 old_low = __half22float2(__uint_as_half2(low[row][HINDEX(row,column/2)]));
                    float output[2];
                    float2 old_values = make_float2(0,0);
                    if constexpr (!ROLLOUT) {
                        if (env < count) old_values = *reinterpret_cast<const float2*>(old_state + (env * 3 + layer) * 256 + column);
                    }
                    #pragma unroll
                    for (int element = 0; element < 2; ++element) {
                        int part = rowpart * 2 + element;
                        float previous;
                        if constexpr (ROLLOUT) previous = reinterpret_cast<float*>(&history)[part];
                        else previous = element == 0 ? old_values.x : old_values.y;
                        float candidate = tanhf(accum[rowtile][tile][0][part]);
                        float gate = fmaf(0.5f, tanhf(0.5f * accum[rowtile][tile][1][part]), 0.5f);
                        float out = previous + gate * (candidate - previous);
                        float highway = fmaf(0.5f, tanhf(0.5f * accum[rowtile][tile][2][part]), 0.5f);
                        float old_hidden = element == 0 ? old_high.x + old_low.x : old_high.y + old_low.y;
                        output[element] = highway * out + (1.0f - highway) * old_hidden;
                        updated[part] = out;
                    }
                    if constexpr (!ROLLOUT) {
                        if (env < count) *reinterpret_cast<float2*>(new_state + (env * 3 + layer) * 256 + column) = make_float2(updated[rowpart*2], updated[rowpart*2+1]);
                    }
                    if (layer == 2) {
                        high[row][HINDEX(row,column/2)] = __float_as_uint(output[0]);
                        low[row][HINDEX(row,column/2)] = __float_as_uint(output[1]);
                    } else {
                        unsigned rounded = pack_half(output[0], output[1]);
                        float2 restored = __half22float2(__uint_as_half2(rounded));
                        high[row][HINDEX(row,column/2)] = rounded;
                        low[row][HINDEX(row,column/2)] = pack_half(output[0]-restored.x, output[1]-restored.y);
                    }
                }
                if constexpr (ROLLOUT) {
                    if (base + rowtile * 16 < count && step < last_step) {
                        float4 next_state = make_float4(updated[0],updated[1],updated[2],updated[3]);
                        if constexpr (KEEP_STATE) {
                            if (layer == 2) register_state[rowtile * TILES + tile] = next_state;
                            else shared_state[((rowtile * 2 + layer) * 32 + warp * TILES + tile) * 32 + lane] = next_state;
                        } else reinterpret_cast<float4*>(new_state)[packed_state_index] = next_state;
                    }
                }
            }
        }
        __syncthreads();
    }

    for (int headrow = warp; headrow < ROWS * (ROLLOUT ? 4 : 5); headrow += WARPS) {
        int row = headrow / (ROLLOUT ? 4 : 5);
        int head = headrow % (ROLLOUT ? 4 : 5);
        float sum = 0;
        const float* head_weight = head < 4 ? action_weight + head * 256 : value_weight;
        #pragma unroll
        for (int index = 0; index < 4; ++index) {
            int column = (lane + index * 32) * 2;
            float first = __uint_as_float(high[row][HINDEX(row,column/2)]);
            float second = __uint_as_float(low[row][HINDEX(row,column/2)]);
            sum = fmaf(first, head_weight[column], sum);
            sum = fmaf(second, head_weight[column+1], sum);
        }
        #pragma unroll
        for (int offset = 16; offset > 0; offset /= 2) sum += __shfl_down_sync(0xffffffff, sum, offset);
        if (lane == 0) {
            sum += head < 4 ? action_bias[head] : value_bias[0];
            heads[row][head] = sum;
            if (base + row < count) {
                if (head < 4) {
                    if (!ROLLOUT || step == last_step) logits[(base + row) * 4 + head] = sum;
                }
                else values[base + row] = sum;
            }
        }
    }
    if constexpr (ROLLOUT) {
        __syncthreads();
        if (thread < ROWS && base + thread < count) {
            int env = base + thread;
            int action = 0;
            #pragma unroll
            for (int head = 1; head < 4; ++head)
                if (heads[thread][head] > heads[thread][action]) action = head;
            int pos_x = int(agent[env*2]) + (action == 3) - (action == 2);
            int pos_y = int(agent[env*2+1]) + (action == 1) - (action == 0);
            pos_x = max(0, min(10, pos_x));
            pos_y = max(0, min(10, pos_y));
            agent[env*2] = pos_x;
            agent[env*2+1] = pos_y;
            if (pos_x == food[env*2] && pos_y == food[env*2+1]) {
                rewards[env] += 1.0f;
                atomicExch(flags + step, 1);
            }
        }
    }
}

template<int ROWS, int WARPS, bool ROLLOUT>
__global__ __launch_bounds__(WARPS * 32) void fused_policy(const float* obs, const float* old_state, float* new_state,
    const float* __restrict__ enc, const float* __restrict__ enc_bias, const uint4* __restrict__ weights,
    const float* __restrict__ action_weight, const float* __restrict__ action_bias, const float* __restrict__ value_weight,
    const float* __restrict__ value_bias, float* logits, float* values, long long* agent,
    long long* food, unsigned long long* rng, float* rewards, int* flags,
    int count, int step, int last_step, unsigned long long seed) {
    fused_policy_body<ROWS,WARPS,ROLLOUT>(obs, old_state, new_state, enc, enc_bias, weights,
        action_weight, action_bias, value_weight, value_bias, logits, values, agent, food, rng,
        rewards, flags, count, step, last_step, seed);
}

template<int ROWS, int WARPS>
__global__ __launch_bounds__(WARPS * 32) void persistent_policy(const float* obs, const float* old_state, float* new_state,
    const float* __restrict__ enc, const float* __restrict__ enc_bias, const uint4* __restrict__ weights,
    const float* __restrict__ action_weight, const float* __restrict__ action_bias, const float* __restrict__ value_weight,
    const float* __restrict__ value_bias, float* logits, float* values, long long* agent,
    long long* food, unsigned long long* rng, float* rewards, int* flags,
    int count, int first_step, int last_step, unsigned long long seed) {
    auto grid = cooperative_groups::this_grid();
    float4 register_state[ROWS / 16 * (32 / WARPS)];
    for (int step = 0; step <= last_step; ++step) {
        fused_policy_body<ROWS,WARPS,true,true>(obs, old_state, new_state, enc, enc_bias, weights,
            action_weight, action_bias, value_weight, value_bias, logits, values, agent, food, rng,
            rewards, flags, count, step, last_step, seed, register_state);
        if (step != last_step) grid.sync();
    }
}

template<typename Position>
__global__ void move_agents(const Position* agent, const Position* food, const long long* actions,
    Position* new_agent, float* reward, int* hit_any, int count) {
    int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= count) return;
    long long action = actions[env];
    Position pos_x = min(Position(10), max(Position(0), Position(agent[env*2] + (action == 3) - (action == 2))));
    Position pos_y = min(Position(10), max(Position(0), Position(agent[env*2+1] + (action == 1) - (action == 0))));
    new_agent[env*2] = pos_x;
    new_agent[env*2+1] = pos_y;
    bool hit = pos_x == food[env*2] && pos_y == food[env*2+1];
    reward[env] = float(hit);
    if (hit) atomicExch(hit_any, 1);
}

template<typename Position>
__global__ void respawn_food(const Position* food, const unsigned long long* rng, const float* reward,
    Position* new_food, unsigned long long* new_rng, const int* hit_any, int count) {
    int env = blockIdx.x * blockDim.x + threadIdx.x;
    if (env >= count) return;
    unsigned long long original = rng[env];
    unsigned long long first = hit_any[0] ? next_rng(original) : original;
    unsigned long long second = hit_any[0] ? next_rng(first) : original;
    new_rng[env] = second;
    new_food[env*2] = reward[env] ? Position(first % 11ULL) : food[env*2];
    new_food[env*2+1] = reward[env] ? Position(second % 11ULL) : food[env*2+1];
}

Tensor prepare_weights(Tensor weights, Tensor enc, Tensor enc_bias) {
    auto packed = torch::empty({3 * 16 * 3 * 32 * 32 * 4}, weights.options().dtype(torch::kInt32));
    pack_weights<<<(3 * 16 * 3 * 32 * 32 + 255)/256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(weights.data_ptr<float>(), reinterpret_cast<uint4*>(packed.data_ptr<int>()));
    compose_encoder<<<96, 256, 0, at::cuda::getCurrentCUDAStream()>>>(weights.data_ptr<float>(), enc.data_ptr<float>(), enc_bias.data_ptr<float>(), reinterpret_cast<float*>(packed.data_ptr<int>()));
    return packed;
}

std::vector<Tensor> forward_cuda(Tensor obs, Tensor state, Tensor enc, Tensor enc_bias,
    Tensor weights, Tensor aw, Tensor ab, Tensor vw, Tensor vb) {
    c10::cuda::CUDAGuard guard(obs.device());
    int count = obs.size(0);
    auto packed = prepare_weights(weights, enc, enc_bias);
    auto logits = torch::empty({count, 4}, obs.options());
    auto new_state = torch::empty_like(state);
    auto values = torch::empty({count}, obs.options());
    if (count) fused_policy<16,8,false><<<(count+15)/16,256,16*1060,at::cuda::getCurrentCUDAStream()>>>(
        obs.data_ptr<float>(), state.data_ptr<float>(), new_state.data_ptr<float>(),
        enc.data_ptr<float>(), enc_bias.data_ptr<float>(), reinterpret_cast<uint4*>(packed.data_ptr<int>()),
        aw.data_ptr<float>(), ab.data_ptr<float>(), vw.data_ptr<float>(), vb.data_ptr<float>(),
        logits.data_ptr<float>(), values.data_ptr<float>(), nullptr, nullptr, nullptr, nullptr, nullptr, count, 0, 0, 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {logits, new_state, values};
}

template<int ROWS, int WARPS, typename... Args>
bool launch_persistent(int count, Args... args) {
    auto kernel = persistent_policy<ROWS,WARPS>;
    constexpr int SHARED_SIZE = ROWS * 3108;
    if constexpr (SHARED_SIZE > 48000)
        C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SHARED_SIZE));
    int active_blocks;
    C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active_blocks, kernel, WARPS*32, SHARED_SIZE));
    int blocks = (count + ROWS - 1) / ROWS;
    if (blocks > active_blocks * at::cuda::getCurrentDeviceProperties()->multiProcessorCount) return false;
    void* arguments[] = {static_cast<void*>(&args)...};
    C10_CUDA_CHECK(cudaLaunchCooperativeKernel(reinterpret_cast<const void*>(kernel), dim3(blocks), dim3(WARPS*32), arguments, SHARED_SIZE, at::cuda::getCurrentCUDAStream()));
    return true;
}

std::vector<Tensor> rollout_cuda(Tensor initial, int horizon, long long seed,
    Tensor enc, Tensor enc_bias, Tensor weights, Tensor aw, Tensor ab, Tensor vw, Tensor vb) {
    c10::cuda::CUDAGuard guard(initial.device());
    auto agent = initial.select(0, 0);
    auto food = initial.select(0, 1);
    int count = agent.size(0);
    auto options = enc.options();
    auto packed = prepare_weights(weights, enc, enc_bias);
    auto logits = horizon ? torch::empty({count,4}, options) : torch::zeros({count,4}, options);
    auto state = torch::empty({((count+15)/16)*16,3,256}, options);
    auto rewards = horizon ? torch::empty({count}, options) : torch::zeros({count}, options);
    auto rng = torch::empty({count}, agent.options());
    auto flags = torch::zeros({horizon}, agent.options().dtype(torch::kInt32));
    #define KERNEL_ARGS(STEP) static_cast<const float*>(nullptr), state.data_ptr<float>(), state.data_ptr<float>(), enc.data_ptr<float>(), enc_bias.data_ptr<float>(), \
        reinterpret_cast<uint4*>(packed.data_ptr<int>()), aw.data_ptr<float>(), ab.data_ptr<float>(), \
        vw.data_ptr<float>(), vb.data_ptr<float>(), logits.data_ptr<float>(), static_cast<float*>(nullptr), \
        reinterpret_cast<long long*>(agent.data_ptr<int64_t>()), reinterpret_cast<long long*>(food.data_ptr<int64_t>()), \
        reinterpret_cast<unsigned long long*>(rng.data_ptr<int64_t>()), rewards.data_ptr<float>(), flags.data_ptr<int>(), count, STEP, horizon - 1, seed
    bool persistent = false;
    if (count && count <= 4096 && horizon)
        persistent = launch_persistent<32,16>(count, KERNEL_ARGS(0));
    if (!persistent && count && horizon) {
        C10_CUDA_CHECK(cudaFuncSetAttribute(fused_policy<48,16,true>, cudaFuncAttributeMaxDynamicSharedMemorySize, 48*1060));
        for (int step = 0; step < horizon; ++step)
            fused_policy<48,16,true><<<(count+47)/48,512,48*1060,at::cuda::getCurrentCUDAStream()>>>(KERNEL_ARGS(step));
    }
    #undef KERNEL_ARGS
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {rewards, agent, logits};
}

std::vector<Tensor> environment_cuda(Tensor agent, Tensor food, Tensor actions, Tensor rng) {
    c10::cuda::CUDAGuard guard(agent.device());
    int count = agent.size(0);
    auto new_agent = torch::empty_like(agent);
    auto new_food = torch::empty_like(food);
    auto new_rng = torch::empty_like(rng);
    auto rewards = torch::empty({count}, agent.options().dtype(torch::kFloat32));
    auto hit = torch::zeros({1}, agent.options().dtype(torch::kInt32));
    if (count) {
        AT_DISPATCH_ALL_TYPES(agent.scalar_type(), "environment", [&] {
            move_agents<scalar_t><<<(count+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(agent.data_ptr<scalar_t>(), food.data_ptr<scalar_t>(), reinterpret_cast<long long*>(actions.data_ptr<int64_t>()), new_agent.data_ptr<scalar_t>(), rewards.data_ptr<float>(), hit.data_ptr<int>(), count);
            respawn_food<scalar_t><<<(count+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(food.data_ptr<scalar_t>(), reinterpret_cast<unsigned long long*>(rng.data_ptr<int64_t>()), rewards.data_ptr<float>(), new_food.data_ptr<scalar_t>(), reinterpret_cast<unsigned long long*>(new_rng.data_ptr<int64_t>()), hit.data_ptr<int>(), count);
        });
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {new_agent, new_food, rewards, new_rng};
}
"""

_CPP = r"""
#include <torch/extension.h>
std::vector<torch::Tensor> forward_cuda(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> rollout_cuda(torch::Tensor, int, long long, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> environment_cuda(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor);
"""

_extension = load_inline(
    name="grid_mingru_cuda",
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    functions=["forward_cuda", "rollout_cuda", "environment_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo", "-Xptxas=-v"],
)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.w_enc = nn.Parameter(torch.empty(256, 4))
        self.b_enc = nn.Parameter(torch.zeros(256))
        self.w_gru = nn.Parameter(torch.empty(3, 768, 256))
        self.w_a = nn.Parameter(torch.empty(4, 256))
        self.b_a = nn.Parameter(torch.zeros(4))
        self.w_v = nn.Parameter(torch.empty(1, 256))
        self.b_v = nn.Parameter(torch.zeros(1))
        self.reset_parameters(0)

    def reset_parameters(self, seed: int = 0):
        generator = torch.Generator(device="cpu")
        generator.manual_seed(seed)
        for parameter in self.parameters():
            temporary = torch.empty(parameter.shape, dtype=parameter.dtype, device="cpu")
            temporary.normal_(0.0, 0.02, generator=generator)
            parameter.data.copy_(temporary)

    def forward(self, obs, state):
        return policy_forward(self, obs, state)


def _parameters(model):
    return model.w_enc, model.b_enc, model.w_gru, model.w_a, model.b_a, model.w_v, model.b_v


@torch.no_grad()
def policy_forward(model, obs, state):
    return tuple(_extension.forward_cuda(obs.contiguous(), state.contiguous(), *_parameters(model)))


@torch.no_grad()
def env_step(agent, food, actions, rng_state):
    return tuple(_extension.environment_cuda(agent.contiguous(), food.contiguous(), actions.contiguous(), rng_state.contiguous()))


@torch.no_grad()
def run(num_envs: int, horizon: int, seed: int, model=None) -> dict:
    if model is None:
        model = Model().cuda()
    elif not model.w_enc.is_cuda:
        model = model.cuda()
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    initial = torch.randint(0, 11, (2, num_envs, 2), generator=generator).to(model.w_enc.device)
    rewards, positions, logits = _extension.rollout_cuda(initial, horizon, seed, *_parameters(model))
    return {"rewards": rewards, "positions": positions, "last_logits": logits}