"""CUDA grid-foraging rollout with a globally synchronized MinGRU policy."""
import os
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.cpp_extension import load_inline

_ROOT = Path(__file__).resolve().parent
os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '12.0')
os.environ.setdefault('MAX_JOBS', '4')
_CPP = r'''
#include <torch/extension.h>
std::vector<at::Tensor> gru_cuda(at::Tensor,at::Tensor,at::Tensor,int);
std::vector<at::Tensor> env_cuda(at::Tensor,at::Tensor,at::Tensor,at::Tensor);
void rollout_cuda(at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,at::Tensor,int);
PYBIND11_MODULE(TORCH_EXTENSION_NAME,m) {
 m.def("gru", &gru_cuda); m.def("env", &env_cuda); m.def("rollout", &rollout_cuda);
}
'''
_CUDA = r'''
#include <ATen/ATen.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cooperative_groups.h>
#include <vector>
namespace cg=cooperative_groups;
using at::Tensor;
__device__ __forceinline__ float sig(float x) { return 1.f/(1.f+expf(-x)); }
__global__ void gru_kernel(const float* h,const float* gates,const float* state,float* hn,float* sn,int n,int layer) {
 int i=blockIdx.x*blockDim.x+threadIdx.x;
 if(i>=n*256) return;
 int e=i/256,c=i%256;
 float zh=gates[e*768+c],zg=gates[e*768+256+c],zp=gates[e*768+512+c];
 float st=state[e*768+layer*256+c];
 float o=st+sig(zg)*(tanhf(zh)-st),p=sig(zp);
 sn[e*256+c]=o;
 hn[i]=p*o+(1.f-p)*h[i];
}
std::vector<Tensor> gru_cuda(Tensor h,Tensor gates,Tensor state,int layer) {
 c10::cuda::CUDAGuard guard(h.device());
 auto hn=at::empty_like(h),sn=at::empty_like(h);
 gru_kernel<<<(h.numel()+255)/256,256,0,c10::cuda::getCurrentCUDAStream()>>>(h.data_ptr<float>(),gates.data_ptr<float>(),state.data_ptr<float>(),hn.data_ptr<float>(),sn.data_ptr<float>(),h.size(0),layer);
 C10_CUDA_KERNEL_LAUNCH_CHECK(); return {hn,sn};
}
__global__ void move_kernel(const float* a,const float* f,const int64_t* act,float* ao,unsigned char* hits,float* rew,int* any,int n) {
 int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=n)return;
 int ac=act[i]; float x=fminf(10.f,fmaxf(0.f,a[i*2]+(ac==3)-(ac==2))),y=fminf(10.f,fmaxf(0.f,a[i*2+1]+(ac==1)-(ac==0)));
 ao[i*2]=x;ao[i*2+1]=y;
 bool hit=(x==f[i*2] && y==f[i*2+1]);hits[i]=hit;rew[i]=hit;
 if(hit)atomicExch(any,1);
}
__device__ __forceinline__ unsigned long long lcg(unsigned long long x) { return (x*6364136223846793005ULL+1ULL)&0x7fffffffffffffffULL; }
__global__ void respawn_kernel(const float* f,const unsigned long long* rng,const unsigned char* hits,const int* any,float* fo,unsigned long long* ro,int n) {
 int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=n)return;
 unsigned long long r=rng[i];float x=f[2*i],y=f[2*i+1];
 if(*any){r=lcg(r);if(hits[i])x=r%11ULL;r=lcg(r);if(hits[i])y=r%11ULL;}
 fo[2*i]=x;fo[2*i+1]=y;ro[i]=r;
}
std::vector<Tensor> env_cuda(Tensor a,Tensor f,Tensor act,Tensor rng) {
 c10::cuda::CUDAGuard guard(a.device());
 auto ao=at::empty_like(a),fo=at::empty_like(f),ro=at::empty_like(rng),rew=at::empty({a.size(0)},a.options());
 auto hits=at::empty({a.size(0)},a.options().dtype(at::kByte)),any=at::zeros({1},a.options().dtype(at::kInt));
 auto stream=c10::cuda::getCurrentCUDAStream();int n=a.size(0);
 move_kernel<<<(n+255)/256,256,0,stream>>>(a.data_ptr<float>(),f.data_ptr<float>(),act.data_ptr<int64_t>(),ao.data_ptr<float>(),hits.data_ptr<unsigned char>(),rew.data_ptr<float>(),any.data_ptr<int>(),n);
 respawn_kernel<<<(n+255)/256,256,0,stream>>>(f.data_ptr<float>(),(const unsigned long long*)rng.data_ptr<int64_t>(),hits.data_ptr<unsigned char>(),any.data_ptr<int>(),fo.data_ptr<float>(),(unsigned long long*)ro.data_ptr<int64_t>(),n);
 C10_CUDA_KERNEL_LAUNCH_CHECK();return {ao,fo,rew,ro};
}
__device__ __forceinline__ void mma(float* c,const unsigned* a,const unsigned* b) {
 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"(c[0]),"+f"(c[1]),"+f"(c[2]),"+f"(c[3])
 : "r"(a[0]),"r"(a[1]),"r"(a[2]),"r"(a[3]),"r"(b[0]),"r"(b[1]));
}
// Each CTA evaluates 16 environments. The recurrent state is FP32, and the
// matrix products use three FP16 tensor-core products with FP32 accumulation.
// Splitting both operands into a high part and a rounding residual preserves
// FP32 accuracy without rounding recurrent state or highway activations.
__global__ __launch_bounds__(256) void rollout_kernel(
 const float* we,const float* be,const unsigned short* wh,const unsigned short* wl,
 const float* wa,const float* ba,const float* wv,const float* bv,
 int64_t* agent,int64_t* food,unsigned long long* rng,float* state,float* rewards,float* logits,
 unsigned char* hits,int* flags,int n,int horizon) {
 __shared__ float h[16*264];
 __shared__ unsigned short ah[16*264],al[16*264];
 int lane=threadIdx.x%32,warp=threadIdx.x/32;
 int row=lane/4,col=(lane%4)*2;
 auto grid=cg::this_grid();
 for(int t=0;t<horizon;t++) {
  for(int tile=blockIdx.x;tile*16<n;tile+=gridDim.x) {
   int base=tile*16;
   for(int i=threadIdx.x;i<4096;i+=256) {
    int e=base+i/256,c=i%256;
    float z=0.f;
    if(e<n){
     float x=agent[2*e],y=agent[2*e+1];
     float ox=(float)(food[2*e]-agent[2*e])/11.f,oy=(float)(food[2*e+1]-agent[2*e+1])/11.f;
     z=be[c]+we[c*4]*ox+we[c*4+1]*oy+we[c*4+2]*(x/10.f)+we[c*4+3]*(y/10.f);
    }
    int idx=(i/256)*264+c;h[idx]=z;__half hi=__float2half_rn(z);ah[idx]=__half_as_ushort(hi);al[idx]=__half_as_ushort(__float2half_rn(z-__half2float(hi)));
   }
   __syncthreads();
   #pragma unroll
   for(int layer=0;layer<3;layer++) {
    float acc[3][4][4]={};
    #pragma unroll 1
    for(int k=0;k<256;k+=16) {
     unsigned a[4],ar[4];
     const unsigned* ap=(const unsigned*)ah;const unsigned* arp=(const unsigned*)al;
     a[0]=ap[(row*264+k+col)/2];a[1]=ap[((row+8)*264+k+col)/2];
     a[2]=ap[(row*264+k+col+8)/2];a[3]=ap[((row+8)*264+k+col+8)/2];
     ar[0]=arp[(row*264+k+col)/2];ar[1]=arp[((row+8)*264+k+col)/2];
     ar[2]=arp[(row*264+k+col+8)/2];ar[3]=arp[((row+8)*264+k+col+8)/2];
     #pragma unroll
     for(int g=0;g<3;g++) {
      #pragma unroll
      for(int j=0;j<4;j++) {
       int wi=((layer*96+g*32+warp*4+j)*16+k/16)*32+lane;
       uint2 bb=((const uint2*)wh)[wi],bbr=((const uint2*)wl)[wi];
       unsigned b[2]={bb.x,bb.y},br[2]={bbr.x,bbr.y};
       mma(acc[g][j],a,br);mma(acc[g][j],ar,b);mma(acc[g][j],a,b);
      }
     }
    }
    __syncthreads();
    #pragma unroll
    for(int j=0;j<4;j++) {
     #pragma unroll
     for(int q=0;q<4;q++) {
      int r=row+(q/2)*8,c=warp*32+j*8+col+(q%2),idx=r*264+c,e=base+r;
      float st=e<n?state[e*768+layer*256+c]:0.f;
      float o=st+sig(acc[1][j][q])*(tanhf(acc[0][j][q])-st),p=sig(acc[2][j][q]);
      float z=p*o+(1.f-p)*h[idx];
      if(e<n)state[e*768+layer*256+c]=o;
      h[idx]=z;__half hi=__float2half_rn(z);ah[idx]=__half_as_ushort(hi);al[idx]=__half_as_ushort(__float2half_rn(z-__half2float(hi)));
     }
    }
    __syncthreads();
   }
   // Two environment rows per warp; action and value head use FP32 SIMT.
   #pragma unroll
   for(int rr=0;rr<2;rr++) {
    int r=warp+rr*8,e=base+r;
    float sums[4]={};
    #pragma unroll
    for(int k=0;k<8;k++) {
     int c=lane+k*32;float x=h[r*264+c];
     #pragma unroll
     for(int a=0;a<4;a++)sums[a]+=x*wa[a*256+c];
    }
    #pragma unroll
    for(int d=16;d>0;d/=2) {
     #pragma unroll
     for(int a=0;a<4;a++)sums[a]+=__shfl_down_sync(0xffffffff,sums[a],d);
    }
    if(lane==0 && e<n) {
     int action=0;float best=sums[0]+ba[0];
     #pragma unroll
     for(int a=0;a<4;a++){float z=sums[a]+ba[a];logits[e*4+a]=z;if(z>best){best=z;action=a;}}
     int64_t x=agent[e*2]+(action==3)-(action==2),y=agent[e*2+1]+(action==1)-(action==0);
     x=x<0?0:(x>10?10:x);y=y<0?0:(y>10?10:y);
     agent[e*2]=x;agent[e*2+1]=y;
     bool hit=x==food[e*2] && y==food[e*2+1];hits[e]=hit;rewards[e]+=hit;
     if(hit)atomicExch(flags+t,1);
    }
   }
   __syncthreads();
  }
  grid.sync();
  if(flags[t]) {
   for(int e=blockIdx.x*256+threadIdx.x;e<n;e+=gridDim.x*256) {
    unsigned long long r=lcg(rng[e]);if(hits[e])food[e*2]=r%11ULL;
    r=lcg(r);if(hits[e])food[e*2+1]=r%11ULL;rng[e]=r;
   }
  }
  grid.sync();
 }
}
void rollout_cuda(Tensor we,Tensor be,Tensor wh,Tensor wl,Tensor wa,Tensor ba,Tensor wv,Tensor bv,
 Tensor agent,Tensor food,Tensor rng,Tensor state,Tensor rewards,Tensor logits,Tensor flags,int horizon) {
 c10::cuda::CUDAGuard guard(agent.device());
 auto hits=at::empty({agent.size(0)},agent.options().dtype(at::kByte));
 int n=agent.size(0),blocks_per_sm=0,dev=agent.get_device();cudaDeviceProp prop;C10_CUDA_CHECK(cudaGetDeviceProperties(&prop,dev));
 C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks_per_sm,rollout_kernel,256,0));
 int blocks=std::min((n+15)/16,prop.multiProcessorCount*blocks_per_sm);
 const float *pwe=we.data_ptr<float>(),*pbe=be.data_ptr<float>(),*pwa=wa.data_ptr<float>(),*pba=ba.data_ptr<float>(),*pwv=wv.data_ptr<float>(),*pbv=bv.data_ptr<float>();
 const unsigned short *pwh=(const unsigned short*)wh.data_ptr<at::Half>(),*pwl=(const unsigned short*)wl.data_ptr<at::Half>();
 int64_t *pa=agent.data_ptr<int64_t>(),*pf=food.data_ptr<int64_t>();
 auto pr=(unsigned long long*)rng.data_ptr<int64_t>();float *ps=state.data_ptr<float>(),*prew=rewards.data_ptr<float>(),*pl=logits.data_ptr<float>();
 auto ph=hits.data_ptr<unsigned char>();auto pflags=flags.data_ptr<int>();
 void* args[]={&pwe,&pbe,&pwh,&pwl,&pwa,&pba,&pwv,&pbv,&pa,&pf,&pr,&ps,&prew,&pl,&ph,&pflags,&n,&horizon};
 C10_CUDA_CHECK(cudaLaunchCooperativeKernel((void*)rollout_kernel,blocks,256,args,0,c10::cuda::getCurrentCUDAStream()));
}
'''
_EXT = None

def _ext():
    global _EXT
    if _EXT is None:
        build = _ROOT / '.cuda_build'
        build.mkdir(exist_ok=True)
        _EXT = load_inline(name='grid_rollout_cuda', cpp_sources=_CPP, cuda_sources=_CUDA,
                           build_directory=str(build), extra_cflags=['-O3'],
                           extra_cuda_cflags=['-O3', '--use_fast_math', '-lineinfo'], 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').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)
        self._packed = None

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

    def _weights(self):
        key = (self.w_gru.data_ptr(), self.w_gru._version)
        cached = getattr(self, '_packed', None)
        if cached is None or cached[0] != key:
            wh = self.w_gru.detach().to(torch.float16)
            wl = (self.w_gru.detach() - wh.float()).to(torch.float16)
            wh = wh.reshape(3, 96, 8, 16, 2, 4, 2).permute(0, 1, 3, 2, 5, 4, 6).contiguous()
            wl = wl.reshape(3, 96, 8, 16, 2, 4, 2).permute(0, 1, 3, 2, 5, 4, 6).contiguous()
            self._packed = (key, wh, wl)
        return self._packed[1:]

@torch.no_grad()
def policy_forward(model, obs, state):
    obs, state = obs.contiguous(), state.contiguous()
    h = F.linear(obs, model.w_enc, model.b_enc)
    states = []
    for layer in range(3):
        gates = F.linear(h, model.w_gru[layer])
        h, st = _ext().gru(h, gates, state, layer)
        states.append(st)
    return F.linear(h, model.w_a, model.b_a), torch.stack(states, dim=1), F.linear(h, model.w_v, model.b_v).squeeze(-1)

@torch.no_grad()
def env_step(agent, food, actions, rng_state):
    if agent.shape[0] == 0:
        return agent.clone(), food.clone(), torch.empty(0, device=agent.device, dtype=torch.float32), rng_state.clone()
    a, f, reward, rng = _ext().env(agent.float().contiguous(), food.float().contiguous(),
                                  actions.long().contiguous(), rng_state.contiguous())
    return a.to(agent.dtype), f.to(food.dtype), reward, rng

@torch.no_grad()
def run(num_envs: int, horizon: int, seed: int, model=None) -> dict:
    if model is None:
        model = Model()
    model = model.to('cuda:0').eval()
    ext = _ext()
    g = torch.Generator(device='cpu').manual_seed(seed)
    agent = torch.randint(0, 11, (num_envs, 2), generator=g).to('cuda:0')
    food = torch.randint(0, 11, (num_envs, 2), generator=g).to('cuda:0')
    rng = torch.arange(num_envs, device='cuda:0', dtype=torch.int64) + seed * 10007
    state = torch.zeros(num_envs, 3, 256, device='cuda:0')
    rewards = torch.zeros(num_envs, device='cuda:0')
    logits = torch.zeros(num_envs, 4, device='cuda:0')
    flags = torch.zeros(horizon, device='cuda:0', dtype=torch.int32)
    wh, wl = model._weights()
    if num_envs and horizon:
        ext.rollout(model.w_enc, model.b_enc, wh, wl, model.w_a, model.b_a, model.w_v, model.b_v,
                    agent, food, rng, state, rewards, logits, flags, horizon)
    return {'rewards': rewards, 'positions': agent, 'last_logits': logits, 'state': state}