from __future__ import annotations

import os

import torch
import torch.nn as nn
import torch.nn.functional as F

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


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)


_CUDA_CPP = r"""
#include <torch/extension.h>
#include <vector>

std::vector<torch::Tensor> env_step_cuda(torch::Tensor agent,
                                         torch::Tensor food,
                                         torch::Tensor actions,
                                         torch::Tensor rng_state);
std::vector<torch::Tensor> env_step_logits_cuda(torch::Tensor agent,
                                                torch::Tensor food,
                                                torch::Tensor logits,
                                                torch::Tensor rng_state);
torch::Tensor observe_cuda(torch::Tensor agent, torch::Tensor food);
torch::Tensor mingru_cuda(torch::Tensor gates, torch::Tensor h,
                          torch::Tensor state, torch::Tensor new_state,
                          int64_t layer);
torch::Tensor mingru_half_cuda(torch::Tensor gates, torch::Tensor h,
                               torch::Tensor state, torch::Tensor new_state,
                               int64_t layer);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("env_step", &env_step_cuda, "Grid-foraging environment step (CUDA)");
  m.def("env_step_logits", &env_step_logits_cuda,
        "Greedy policy action and environment step (CUDA)");
  m.def("observe", &observe_cuda, "Grid-foraging observation (CUDA)");
  m.def("mingru", &mingru_cuda, "MinGRU activation (CUDA)");
  m.def("mingru_half", &mingru_half_cuda,
        "MinGRU activation with FP16 gates and hidden state (CUDA)");
}
"""

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <vector>

__global__ void move_agents(const float* agent, const float* food,
                            const int64_t* actions, const float* logits,
                            bool use_logits, float* agent_out,
                            float* reward, int* any_hit, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  int x = __float2int_rn(agent[2 * i]);
  int y = __float2int_rn(agent[2 * i + 1]);
  int a;
  if (use_logits) {
    int action = 0;
    float best = logits[4 * i];
    for (int k = 1; k < 4; ++k) {
      float v = logits[4 * i + k];
      if (v > best) {
        best = v;
        action = k;
      }
    }
    a = action;
  } else {
    a = static_cast<int>(actions[i]);
  }
  if (a == 0) --y;
  else if (a == 1) ++y;
  else if (a == 2) --x;
  else if (a == 3) ++x;
  x = max(0, min(10, x));
  y = max(0, min(10, y));
  agent_out[2 * i] = static_cast<float>(x);
  agent_out[2 * i + 1] = static_cast<float>(y);
  bool hit = (static_cast<float>(x) == food[2 * i]) &&
             (static_cast<float>(y) == food[2 * i + 1]);
  reward[i] = hit ? 1.0f : 0.0f;
  if (hit) atomicExch(any_hit, 1);
}

__device__ __forceinline__ uint64_t lcg(uint64_t x) {
  return (x * 6364136223846793005ULL + 1ULL) & 0x7fffffffffffffffULL;
}

__device__ __forceinline__ void store_obs(const float* agent,
                                           const float* food,
                                           float* obs, int i) {
  float ax = agent[2 * i];
  float ay = agent[2 * i + 1];
  obs[4 * i] = __fdiv_rn(food[2 * i] - ax, 11.0f);
  obs[4 * i + 1] = __fdiv_rn(food[2 * i + 1] - ay, 11.0f);
  obs[4 * i + 2] = __fdiv_rn(ax, 10.0f);
  obs[4 * i + 3] = __fdiv_rn(ay, 10.0f);
}

__global__ void observe_kernel(const float* agent, const float* food,
                               float* obs, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) store_obs(agent, food, obs, i);
}

__global__ void mingru_kernel(const float* gates, const float* h,
                              const float* state, float* new_state,
                              float* h_out, int64_t state_stride0,
                              int64_t state_stride1, int layer, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  int total = n * 256;
  if (i >= total) return;
  int row = i >> 8;
  int col = i & 255;
  float zh = gates[row * 768 + col];
  float zg = gates[row * 768 + 256 + col];
  float zp = gates[row * 768 + 512 + col];
  float st = state[row * state_stride0 + layer * state_stride1 + col];
  float sh = 1.0f / (1.0f + expf(-zg));
  float candidate = tanhf(zh);
  float out = __fadd_rn(st, __fmul_rn(sh, __fsub_rn(candidate, st)));
  float p = 1.0f / (1.0f + expf(-zp));
  float mixed = __fadd_rn(__fmul_rn(p, out),
                          __fmul_rn(__fsub_rn(1.0f, p), h[i]));
  new_state[row * 768 + layer * 256 + col] = out;
  h_out[i] = mixed;
}

__global__ void mingru_half_kernel(const __half* gates, const __half* h,
                                   const __half* state, __half* new_state,
                                   __half* h_out, int64_t state_stride0,
                                   int64_t state_stride1, int layer, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  int total = n * 256;
  if (i >= total) return;
  int row = i >> 8;
  int col = i & 255;
  float zh = __half2float(gates[row * 768 + col]);
  float zg = __half2float(gates[row * 768 + 256 + col]);
  float zp = __half2float(gates[row * 768 + 512 + col]);
  float st = __half2float(state[row * state_stride0 + layer * state_stride1 + col]);
  float sh = 1.0f / (1.0f + expf(-zg));
  float candidate = tanhf(zh);
  float out = __fadd_rn(st, __fmul_rn(sh, __fsub_rn(candidate, st)));
  float p = 1.0f / (1.0f + expf(-zp));
  float mixed = __fadd_rn(__fmul_rn(p, out),
                          __fmul_rn(__fsub_rn(1.0f, p),
                                    __half2float(h[i])));
  new_state[row * 768 + layer * 256 + col] = __float2half_rn(out);
  h_out[i] = __float2half_rn(mixed);
}

__global__ void respawn_food(const float* food, const float* reward,
                             const float* agent, const int64_t* rng_in,
                             const int* any_hit,
                             float* food_out, float* obs_out,
                             int64_t* rng_out, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  float fx = food[2 * i];
  float fy = food[2 * i + 1];
  uint64_t r = static_cast<uint64_t>(rng_in[i]);
  if (*any_hit) {
    r = lcg(r);
    if (reward[i] != 0.0f) {
      fx = static_cast<float>(r % 11ULL);
    }
    r = lcg(r);
    if (reward[i] != 0.0f) {
      fy = static_cast<float>(r % 11ULL);
    }
  }
  food_out[2 * i] = fx;
  food_out[2 * i + 1] = fy;
  rng_out[i] = static_cast<int64_t>(r);
  store_obs(agent, food_out, obs_out, i);
}

static std::vector<torch::Tensor> env_step_impl(
    torch::Tensor agent, torch::Tensor food, const int64_t* actions,
    const float* logits, bool use_logits, torch::Tensor rng_state) {
  c10::cuda::CUDAGuard guard(agent.device());
  const int n = static_cast<int>(agent.size(0));
  auto agent_out = torch::empty_like(agent);
  auto food_out = torch::empty_like(food);
  auto obs_out = torch::empty({n, 4}, agent.options());
  auto reward = torch::empty({n}, agent.options());
  auto rng_out = torch::empty_like(rng_state);
  auto flag = torch::empty({1}, agent.options().dtype(torch::kInt32));
  auto stream = at::cuda::getCurrentCUDAStream();
  cudaMemsetAsync(flag.data_ptr<int>(), 0, sizeof(int), stream);
  constexpr int threads = 256;
  const int blocks = (n + threads - 1) / threads;
  move_agents<<<blocks, threads, 0, stream>>>(
      agent.data_ptr<float>(), food.data_ptr<float>(),
      actions, logits, use_logits, agent_out.data_ptr<float>(),
      reward.data_ptr<float>(), flag.data_ptr<int>(), n);
  respawn_food<<<blocks, threads, 0, stream>>>(
      food.data_ptr<float>(), reward.data_ptr<float>(), agent_out.data_ptr<float>(),
      rng_state.data_ptr<int64_t>(), flag.data_ptr<int>(),
      food_out.data_ptr<float>(), obs_out.data_ptr<float>(),
      rng_out.data_ptr<int64_t>(), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return {agent_out, food_out, reward, rng_out, obs_out};
}

std::vector<torch::Tensor> env_step_cuda(torch::Tensor agent,
                                         torch::Tensor food,
                                         torch::Tensor actions,
                                         torch::Tensor rng_state) {
  return env_step_impl(agent, food, actions.data_ptr<int64_t>(), nullptr,
                       false, rng_state);
}

std::vector<torch::Tensor> env_step_logits_cuda(torch::Tensor agent,
                                                torch::Tensor food,
                                                torch::Tensor logits,
                                                torch::Tensor rng_state) {
  return env_step_impl(agent, food, nullptr, logits.data_ptr<float>(),
                       true, rng_state);
}

torch::Tensor observe_cuda(torch::Tensor agent, torch::Tensor food) {
  c10::cuda::CUDAGuard guard(agent.device());
  const int n = static_cast<int>(agent.size(0));
  auto obs = torch::empty({n, 4}, agent.options());
  constexpr int threads = 256;
  const int blocks = (n + threads - 1) / threads;
  auto stream = at::cuda::getCurrentCUDAStream();
  observe_kernel<<<blocks, threads, 0, stream>>>(
      agent.data_ptr<float>(), food.data_ptr<float>(), obs.data_ptr<float>(), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return obs;
}

torch::Tensor mingru_cuda(torch::Tensor gates, torch::Tensor h,
                          torch::Tensor state, torch::Tensor new_state,
                          int64_t layer) {
  c10::cuda::CUDAGuard guard(h.device());
  auto h_out = torch::empty_like(h);
  const int n = static_cast<int>(h.size(0));
  constexpr int threads = 256;
  const int blocks = (n * 256 + threads - 1) / threads;
  auto stream = at::cuda::getCurrentCUDAStream();
  mingru_kernel<<<blocks, threads, 0, stream>>>(
      gates.data_ptr<float>(), h.data_ptr<float>(), state.data_ptr<float>(),
      new_state.data_ptr<float>(), h_out.data_ptr<float>(),
      state.stride(0), state.stride(1), static_cast<int>(layer), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return h_out;
}

torch::Tensor mingru_half_cuda(torch::Tensor gates, torch::Tensor h,
                               torch::Tensor state, torch::Tensor new_state,
                               int64_t layer) {
  c10::cuda::CUDAGuard guard(h.device());
  auto h_out = torch::empty_like(h);
  const int n = static_cast<int>(h.size(0));
  constexpr int threads = 256;
  const int blocks = (n * 256 + threads - 1) / threads;
  auto stream = at::cuda::getCurrentCUDAStream();
  mingru_half_kernel<<<blocks, threads, 0, stream>>>(
      reinterpret_cast<const __half*>(gates.data_ptr<at::Half>()),
      reinterpret_cast<const __half*>(h.data_ptr<at::Half>()),
      reinterpret_cast<const __half*>(state.data_ptr<at::Half>()),
      reinterpret_cast<__half*>(new_state.data_ptr<at::Half>()),
      reinterpret_cast<__half*>(h_out.data_ptr<at::Half>()),
      state.stride(0), state.stride(1), static_cast<int>(layer), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
  return h_out;
}
"""

_EXT = None
_GRAPH_CACHE: dict[tuple, tuple] = {}


def _cuda_ext():
    global _EXT
    if _EXT is None:
        from torch.utils.cpp_extension import load_inline

        build_directory = os.path.join(
            os.path.dirname(__file__), ".torch_extensions", "grid_mingru_env_sm120_v6"
        )
        os.makedirs(build_directory, exist_ok=True)
        _EXT = load_inline(
            name="grid_mingru_env_sm120_v6",
            build_directory=build_directory,
            cpp_sources=_CUDA_CPP,
            cuda_sources=_CUDA_SRC,
            functions=None,
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3", "-arch=sm_120"],
            with_cuda=True,
            verbose=False,
        )
    return _EXT


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


def _policy_forward_rollout(model, obs, state):
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        h = F.linear(obs, model.w_enc, model.b_enc)
        new_state = torch.empty((obs.shape[0], GRU_LAYERS, HIDDEN),
                                device=obs.device, dtype=torch.float16)
        for layer in range(GRU_LAYERS):
            gates = F.linear(h, model.w_gru[layer])
            h = _cuda_ext().mingru_half(gates, h, state, new_state, layer)
        logits = F.linear(h, model.w_a, model.b_a)
    return logits.float(), new_state


def env_step(agent, food, actions, rng_state):
    return tuple(_cuda_ext().env_step(agent, food, actions, rng_state)[:4])


def _rollout(model, agent, food, rng_state, state, horizon):
    rewards = torch.zeros(agent.shape[0], device=agent.device, dtype=torch.float32)
    last_logits = torch.empty((agent.shape[0], NUM_ACTIONS),
                              device=agent.device, dtype=torch.float32)
    obs = _cuda_ext().observe(agent, food)
    for _ in range(horizon):
        logits, state = _policy_forward_rollout(model, obs, state)
        last_logits = logits
        agent, food, r, rng_state, obs = _cuda_ext().env_step_logits(
            agent, food, logits, rng_state
        )
        rewards = rewards + r
    return rewards, agent, food, rng_state, state, last_logits


def _make_graph(model, num_envs, horizon, device):
    _cuda_ext()
    static_agent = torch.zeros((num_envs, 2), device=device, dtype=torch.float32)
    static_food = torch.ones((num_envs, 2), device=device, dtype=torch.float32)
    static_rng = torch.arange(num_envs, device=device, dtype=torch.int64)
    static_state = torch.zeros((num_envs, GRU_LAYERS, HIDDEN),
                               device=device, dtype=torch.float16)
    stream = torch.cuda.Stream(device=device)
    stream.wait_stream(torch.cuda.current_stream(device))
    old_tf32 = torch.backends.cuda.matmul.allow_tf32
    torch.backends.cuda.matmul.allow_tf32 = True
    try:
        with torch.cuda.stream(stream), torch.inference_mode():
            for _ in range(2):
                _rollout(model, static_agent, static_food, static_rng,
                         static_state, min(horizon, 2))
        torch.cuda.current_stream(device).wait_stream(stream)
        graph = torch.cuda.CUDAGraph()
        with torch.inference_mode(), torch.cuda.graph(graph):
            captured = _rollout(model, static_agent, static_food, static_rng,
                                 static_state, horizon)
    finally:
        torch.backends.cuda.matmul.allow_tf32 = old_tf32
    return (graph, static_agent, static_food, static_rng,
            captured[0], captured[1], static_state, model, captured[5])


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

    params_key = tuple(p.data_ptr() for p in model.parameters())
    key = (id(model), params_key, num_envs, horizon, device.index)
    entry = _GRAPH_CACHE.get(key)
    if entry is None:
        entry = _make_graph(model, num_envs, horizon, device)
        _GRAPH_CACHE[key] = entry

    graph, static_agent, static_food, static_rng, out_rewards, out_agent, _static_state, _keep_model, out_logits = entry
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    agent = torch.randint(0, BOARD, (num_envs, 2), generator=g).float()
    food = torch.randint(0, BOARD, (num_envs, 2), generator=g).float()
    static_agent.copy_(agent, non_blocking=True)
    static_food.copy_(food, non_blocking=True)
    static_rng.copy_(torch.arange(num_envs, device=device, dtype=torch.int64) + seed * 10007)
    graph.replay()
    return {
        "rewards": out_rewards.clone(),
        "positions": out_agent.long(),
        "last_logits": out_logits.clone(),
    }