"""CUDA grid foraging and MinGRU rollout."""
from __future__ import annotations

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

BOARD = 11
HIDDEN = 256
GRU_LAYERS = 3
_CUDA = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cmath>

__device__ __forceinline__ int clip11(int x) { return max(0, min(10, x)); }
__device__ __forceinline__ int64_t lcg(int64_t x) {
    return (int64_t)(((uint64_t)x * 6364136223846793005ULL + 1ULL) & 0x7fffffffffffffffULL);
}
__global__ void encode_all_half_kernel(const float* agent, const float* food,
        const float* w, const float* b, __half* h, int n) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= n * 256) return;
    int e = k >> 8, j = k & 255;
    float x = agent[2*e], y = agent[2*e+1];
    float dx = (food[2*e] - x) / 11.0f;
    float dy = (food[2*e+1] - y) / 11.0f;
    const float* ww = w + 4*j;
    float v = b[j] + dx*ww[0] + dy*ww[1] + (x/10.0f)*ww[2] + (y/10.0f)*ww[3];
    h[k] = __float2half_rn(v);
}

__global__ void gate_all_half_kernel(const __half* gates, __half* state, __half* h, int layer, int n) {
    int e = blockIdx.x*16 + (threadIdx.x >> 5);
    int lane = threadIdx.x & 31;
    if (e>=n) return;
    const __half* g = gates + e*768;
    __half* stp = state + e*768 + layer*256;
    __half* hp = h + e*256;
    #pragma unroll
    for (int i=0;i<8;++i) {
        int j=lane+32*i;
        float st=__half2float(stp[j]);
        float zh=__half2float(g[j]);
        float zg=__half2float(g[256+j]);
        float zp=__half2float(g[512+j]);
        float z=1.0f/(1.0f+__expf(-zg));
        float out=st+z*(__tanhf(zh)-st);
        float p=1.0f/(1.0f+__expf(-zp));
        float old=__half2float(hp[j]);
        stp[j]=__float2half_rn(out);
        hp[j]=__float2half_rn(p*out+(1.0f-p)*old);
    }
}

__global__ void advance_kernel(const float* logits, const float* bias, const int64_t* actions,
        float* agent, const float* food, float* rewards, unsigned char* hit,
        int* anyhit, float* last, int n) {
    int e = blockIdx.x * blockDim.x + threadIdx.x;
    if (e >= n) return;
    int action;
    if (actions) {
        action = (int)actions[e];
    } else {
        float best = logits[e*4] + bias[0];
        action = 0;
        if (last) last[e*4] = best;
        #pragma unroll
        for (int a = 1; a < 4; ++a) {
            float v = logits[e*4+a] + bias[a];
            if (last) last[e*4+a] = v;
            if (v > best) {best = v; action = a;}
        }
    }
    int x = (int)agent[e*2], y = (int)agent[e*2+1];
    x = clip11(x + (action == 3) - (action == 2));
    y = clip11(y + (action == 1) - (action == 0));
    agent[e*2] = (float)x;
    agent[e*2+1] = (float)y;
    bool ate = x == (int)food[e*2] && y == (int)food[e*2+1];
    hit[e] = (unsigned char)ate;
    if (ate) { rewards[e] += 1.0f; atomicOr(anyhit, 1); }
}

__global__ void respawn_kernel(const float* agent, float* food, int64_t* rng,
        const unsigned char* hit, const int* anyhit, const float* w, const float* b,
        float* h, int n) {
    int e = blockIdx.x * blockDim.x + threadIdx.x;
    if (e >= n) return;
    if (anyhit[0]) {
        int64_t r1 = lcg(rng[e]);
        int64_t r2 = lcg(r1);
        rng[e] = r2;
        if (hit[e]) {
            food[2*e] = (float)(r1 % 11);
            food[2*e+1] = (float)(r2 % 11);
        }
    }
    if (h) {
        float x = agent[2*e], y = agent[2*e+1];
        float dx = (food[2*e] - x) / 11.0f;
        float dy = (food[2*e+1] - y) / 11.0f;
        for (int j = 0; j < 256; ++j) {
            const float* ww = w + 4*j;
            h[e*256+j] = b[j] + dx*ww[0] + dy*ww[1] + (x/10.0f)*ww[2] + (y/10.0f)*ww[3];
        }
    }
}

__global__ void head_all_half_kernel(const __half* h, const float* w, const float* bias,
        float* agent, const float* food, float* rewards, unsigned char* hit,
        int* anyhit, float* last, int n) {
    int e = blockIdx.x * 8 + (threadIdx.x >> 5);
    int lane = threadIdx.x & 31;
    if (e >= n) return;
    float s0=0.0f, s1=0.0f, s2=0.0f, s3=0.0f;
    #pragma unroll
    for (int i=0; i<8; ++i) {
        int j=lane+32*i;
        float v=__half2float(h[e*256+j]);
        s0 += v*w[j];
        s1 += v*w[256+j];
        s2 += v*w[512+j];
        s3 += v*w[768+j];
    }
    for (int offset=16; offset>0; offset>>=1) {
        s0 += __shfl_down_sync(0xffffffff,s0,offset);
        s1 += __shfl_down_sync(0xffffffff,s1,offset);
        s2 += __shfl_down_sync(0xffffffff,s2,offset);
        s3 += __shfl_down_sync(0xffffffff,s3,offset);
    }
    if (lane==0) {
        s0+=bias[0]; s1+=bias[1]; s2+=bias[2]; s3+=bias[3];
        last[e*4]=s0; last[e*4+1]=s1; last[e*4+2]=s2; last[e*4+3]=s3;
        float best=s0; int action=0;
        if (s1>best) {best=s1;action=1;}
        if (s2>best) {best=s2;action=2;}
        if (s3>best) {best=s3;action=3;}
        int x=(int)agent[e*2], y=(int)agent[e*2+1];
        x=clip11(x+(action==3)-(action==2));
        y=clip11(y+(action==1)-(action==0));
        agent[e*2]=(float)x; agent[e*2+1]=(float)y;
        bool ate=x==(int)food[e*2] && y==(int)food[e*2+1];
        hit[e]=(unsigned char)ate;
        if (ate) {rewards[e]+=1.0f; atomicOr(anyhit,1);}
    }
}

__global__ void respawn_encode_all_half_kernel(const float* agent, float* food, int64_t* rng,
        const unsigned char* hit, const int* anyhit, const float* w, const float* b,
        __half* h, int n) {
    int e = blockIdx.x * 8 + (threadIdx.x >> 5);
    int lane = threadIdx.x & 31;
    if (e >= n) return;
    float x = 0.0f, y = 0.0f, fx = 0.0f, fy = 0.0f;
    if (lane == 0) {
        x = agent[2*e]; y = agent[2*e+1];
        fx = food[2*e]; fy = food[2*e+1];
        if (anyhit[0]) {
            int64_t r1 = lcg(rng[e]);
            int64_t r2 = lcg(r1);
            rng[e] = r2;
            if (hit[e]) {
                fx = (float)(r1 % 11); fy = (float)(r2 % 11);
                food[2*e] = fx; food[2*e+1] = fy;
            }
        }
    }
    x = __shfl_sync(0xffffffff, x, 0);
    y = __shfl_sync(0xffffffff, y, 0);
    fx = __shfl_sync(0xffffffff, fx, 0);
    fy = __shfl_sync(0xffffffff, fy, 0);
    float dx = (fx - x) / 11.0f;
    float dy = (fy - y) / 11.0f;
    float xx = x / 10.0f, yy = y / 10.0f;
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        int j = lane + 32*i;
        const float* ww = w + 4*j;
        h[e*256+j] = __float2half_rn(b[j] + dx*ww[0] + dy*ww[1] + xx*ww[2] + yy*ww[3]);
    }
}

__global__ void positions_kernel(const float* agent, int64_t* positions, int n) {
    int k=blockIdx.x*blockDim.x+threadIdx.x;
    if (k<n*2) positions[k]=(int64_t)agent[k];
}

void encode_all_half(torch::Tensor agent, torch::Tensor food, torch::Tensor w, torch::Tensor b, torch::Tensor h) {
    int n=agent.size(0);
    encode_all_half_kernel<<<(n*256+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(
        agent.data_ptr<float>(),food.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),
        reinterpret_cast<__half*>(h.data_ptr<at::Half>()),n);
}

void gate_all_half(torch::Tensor gates, torch::Tensor state, torch::Tensor h, int layer) {
    int n=h.size(0);
    gate_all_half_kernel<<<(n+15)/16,512,0,at::cuda::getCurrentCUDAStream()>>>(
        reinterpret_cast<const __half*>(gates.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(state.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(h.data_ptr<at::Half>()),layer,n);
}

void head_all_half(torch::Tensor h, torch::Tensor w, torch::Tensor bias,
        torch::Tensor agent, torch::Tensor food, torch::Tensor rewards,
        torch::Tensor hit, torch::Tensor anyhit, torch::Tensor last) {
    int n=agent.size(0);
    cudaMemsetAsync(anyhit.data_ptr<int>(),0,sizeof(int),at::cuda::getCurrentCUDAStream());
    head_all_half_kernel<<<(n+7)/8,256,0,at::cuda::getCurrentCUDAStream()>>>(
        reinterpret_cast<const __half*>(h.data_ptr<at::Half>()),w.data_ptr<float>(),bias.data_ptr<float>(),
        agent.data_ptr<float>(),food.data_ptr<float>(),rewards.data_ptr<float>(),
        hit.data_ptr<unsigned char>(),anyhit.data_ptr<int>(),last.data_ptr<float>(),n);
}

void respawn_all_half(torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
        torch::Tensor hit, torch::Tensor anyhit, torch::Tensor w, torch::Tensor b, torch::Tensor h) {
    int n=agent.size(0);
    respawn_encode_all_half_kernel<<<(n+7)/8,256,0,at::cuda::getCurrentCUDAStream()>>>(
        agent.data_ptr<float>(),food.data_ptr<float>(),rng.data_ptr<int64_t>(),
        hit.data_ptr<unsigned char>(),anyhit.data_ptr<int>(),w.data_ptr<float>(),b.data_ptr<float>(),
        reinterpret_cast<__half*>(h.data_ptr<at::Half>()),n);
}

void positions(torch::Tensor agent, torch::Tensor out) {
    int n=agent.size(0);
    positions_kernel<<<(n*2+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(
        agent.data_ptr<float>(),out.data_ptr<int64_t>(),n);
}

void advance_actions(torch::Tensor actions, torch::Tensor agent,
        torch::Tensor food, torch::Tensor rewards, torch::Tensor hit, torch::Tensor anyhit) {
    int n = agent.size(0);
    cudaMemsetAsync(anyhit.data_ptr<int>(), 0, sizeof(int), at::cuda::getCurrentCUDAStream());
    advance_kernel<<<(n+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(
        nullptr,nullptr,actions.data_ptr<int64_t>(),agent.data_ptr<float>(),
        food.data_ptr<float>(),rewards.data_ptr<float>(),hit.data_ptr<unsigned char>(),
        anyhit.data_ptr<int>(),nullptr,n);
}

void respawn(torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
        torch::Tensor hit, torch::Tensor anyhit, torch::Tensor w, torch::Tensor b, torch::Tensor h) {
    int n=agent.size(0);
    respawn_kernel<<<(n+255)/256,256,0,at::cuda::getCurrentCUDAStream()>>>(
        agent.data_ptr<float>(),food.data_ptr<float>(),rng.data_ptr<int64_t>(),
        hit.data_ptr<unsigned char>(),anyhit.data_ptr<int>(),nullptr,nullptr,nullptr,n);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("encode_all_half", &encode_all_half);
    m.def("gate_all_half", &gate_all_half);
    m.def("head_all_half", &head_all_half);
    m.def("respawn_all_half", &respawn_all_half);
    m.def("positions", &positions);
    m.def("advance_actions", &advance_actions);
    m.def("respawn", &respawn);
}
'''
_ext = None

def _get_ext():
    global _ext
    if _ext is None:
        os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '12.0')
        _ext = load_inline(
            name='grid_mingru_cuda_v11', cpp_sources='', cuda_sources=_CUDA,
            functions=None, extra_cuda_cflags=['-O3', '-arch=sm_120', '--use_fast_math'],
            verbose=False,
        )
    return _ext

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=0):
        g = torch.Generator(device='cpu')
        g.manual_seed(seed)
        for p in self.parameters():
            t = torch.empty(p.shape, dtype=p.dtype, device='cpu')
            t.normal_(0.0, 0.02, generator=g)
            p.data.copy_(t)
        self._clear_rollout_cache()
    def _clear_rollout_cache(self):
        cache = globals().get('_GRAPH_CACHE')
        if cache is not None:
            for key in [key for key in cache if key[2] == id(self)]:
                del cache[key]
    def load_state_dict(self, *args, **kwargs):
        result = super().load_state_dict(*args, **kwargs)
        self._clear_rollout_cache()
        return result
    def forward(self, obs, state):
        return policy_forward(self, obs, state)

def policy_forward(model, obs, state):
    h = F.linear(obs, model.w_enc, model.b_enc)
    new_states = []
    for layer in range(3):
        gates = F.linear(h, model.w_gru[layer])
        zh, zg, zp = gates.split(256, dim=-1)
        out = state[:,layer,:] + torch.sigmoid(zg)*(torch.tanh(zh)-state[:,layer,:])
        p = torch.sigmoid(zp)
        h = p*out + (1.0-p)*h
        new_states.append(out)
    return F.linear(h,model.w_a,model.b_a), torch.stack(new_states,dim=1), F.linear(h,model.w_v,model.b_v).squeeze(-1)

def env_step(agent, food, actions, rng_state):
    ext = _get_ext()
    a, f, rng = agent.clone(), food.clone(), rng_state.clone()
    rewards = torch.zeros(agent.shape[0], device=agent.device, dtype=torch.float32)
    hit = torch.empty(agent.shape[0], device=agent.device, dtype=torch.uint8)
    anyhit = torch.empty(1,device=agent.device,dtype=torch.int32)
    empty = torch.empty(0,device=agent.device,dtype=torch.float32)
    ext.advance_actions(actions.contiguous(), a, f, rewards, hit, anyhit)
    ext.respawn(a,f,rng,hit,anyhit,empty,empty,empty)
    return a,f,rewards,rng

_GRAPH_CACHE = {}

def _graph_run(num_envs, horizon, seed, model):
    ext = _get_ext()
    device = model.w_enc.device
    key = (num_envs,horizon,id(model),device,tuple(p._version for p in model.parameters()))
    ctx = _GRAPH_CACHE.get(key)
    if ctx is None:
        n = num_envs
        agent = torch.empty((n,2),device=device)
        food = torch.empty((n,2),device=device)
        rng = torch.empty((n,),device=device,dtype=torch.int64)
        base_rng = torch.arange(n,device=device,dtype=torch.int64)
        cpu_agent = torch.empty((n,2),dtype=torch.float32,pin_memory=True)
        cpu_food = torch.empty((n,2),dtype=torch.float32,pin_memory=True)
        generator = torch.Generator(device='cpu')
        copy_done = torch.cuda.Event(blocking=False)
        rewards = torch.empty(n,device=device)
        state = torch.empty((n,3,256),device=device,dtype=torch.float16)
        h = torch.empty((n,256),device=device,dtype=torch.float16)
        gates = torch.empty((n,768),device=device,dtype=torch.float16)
        last = torch.empty((n,4),device=device)
        positions_out = torch.empty((n,2),device=device,dtype=torch.int64)
        hit = torch.empty(n,device=device,dtype=torch.uint8)
        anyhit = torch.empty(1,device=device,dtype=torch.int32)
        weights = [model.w_gru[i].t().half() for i in range(3)]
        def body():
            rewards.zero_()
            state.zero_()
            last.zero_()
            ext.encode_all_half(agent,food,model.w_enc,model.b_enc,h)
            for _ in range(horizon):
                for layer in range(3):
                    torch.mm(h,weights[layer],out=gates)
                    ext.gate_all_half(gates,state,h,layer)
                ext.head_all_half(h,model.w_a,model.b_a,agent,food,rewards,hit,anyhit,last)
                ext.respawn_all_half(agent,food,rng,hit,anyhit,model.w_enc,model.b_enc,h)
            ext.positions(agent,positions_out)
        agent.zero_()
        food.zero_()
        rng.zero_()
        with torch.no_grad():
            # Set up cuBLAS on a side stream before capture.
            stream = torch.cuda.Stream(device=device)
            stream.wait_stream(torch.cuda.current_stream(device))
            with torch.cuda.stream(stream):
                body()
            torch.cuda.current_stream(device).wait_stream(stream)
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                body()
        ctx = (base_rng,cpu_agent,cpu_food,generator,copy_done,agent,food,rng,rewards,state,h,gates,last,positions_out,hit,anyhit,graph,weights,model)
        _GRAPH_CACHE[key] = ctx
    base_rng,cpu_agent,cpu_food,generator,copy_done,agent,food,rng,rewards,state,h,gates,last,positions_out,hit,anyhit,graph,weights,model = ctx
    copy_done.synchronize()
    generator.manual_seed(seed)
    torch.randint(0,11,(num_envs,2),generator=generator,out=cpu_agent)
    torch.randint(0,11,(num_envs,2),generator=generator,out=cpu_food)
    agent.copy_(cpu_agent,non_blocking=True)
    food.copy_(cpu_food,non_blocking=True)
    copy_done.record(torch.cuda.current_stream(device))
    torch.add(base_rng,seed*10007,out=rng)
    graph.replay()
    return {'rewards':rewards.clone(), 'positions':positions_out.clone(), 'last_logits':last.clone()}

_DEFAULT_MODEL = None
def run(num_envs:int,horizon:int,seed:int,model=None)->dict:
    global _DEFAULT_MODEL
    if model is None:
        if _DEFAULT_MODEL is None:
            _DEFAULT_MODEL = Model()
        model = _DEFAULT_MODEL
    model = model.to('cuda:0').eval()
    return _graph_run(num_envs,horizon,seed,model)