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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void bellman_loss_kernel(
    const float* __restrict__ q_values,
    const int64_t* __restrict__ actions,
    const float* __restrict__ rewards,
    const float* __restrict__ next_q_values,
    const float* __restrict__ dones,
    float* __restrict__ output,
    int batch_size,
    int num_actions,
    float gamma
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < batch_size) {
        int64_t action = actions[idx];

        float curr_q = q_values[idx * num_actions + action];

        float max_next_q = -1e20f; // -inf
        int next_q_offset = idx * num_actions;
        for (int a = 0; a < num_actions; ++a) {
            float val = next_q_values[next_q_offset + a];
            if (val > max_next_q) {
                max_next_q = val;
            }
        }

        float target = rewards[idx] + gamma * max_next_q * (1.0f - dones[idx]);

        float diff = curr_q - target;
        output[idx] = diff * diff;
    }
}

torch::Tensor bellman_loss_cuda(
    torch::Tensor q_values,
    torch::Tensor actions,
    torch::Tensor rewards,
    torch::Tensor next_q_values,
    torch::Tensor dones,
    float gamma
) {
    int batch_size = q_values.size(0);
    int num_actions = q_values.size(1);

    auto output = at::empty({batch_size}, q_values.options());

    int threads = 256;
    int blocks = (batch_size + threads - 1) / threads;

    bellman_loss_kernel<<<blocks, threads>>>(
        q_values.data_ptr<float>(),
        actions.data_ptr<int64_t>(),
        rewards.data_ptr<float>(),
        next_q_values.data_ptr<float>(),
        dones.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        num_actions,
        gamma
    );

    return output.mean();
}
"""

cpp_source = """
torch::Tensor bellman_loss_cuda(
    torch::Tensor q_values,
    torch::Tensor actions,
    torch::Tensor rewards,
    torch::Tensor next_q_values,
    torch::Tensor dones,
    float gamma
);
"""

bellman_loss = load_inline(
    name="bellman_loss",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["bellman_loss_cuda"],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, gamma=0.99):
        super(ModelNew, self).__init__()
        self.gamma = gamma

    def forward(self, q_values, actions, rewards, next_q_values, dones):
        return bellman_loss.bellman_loss_cuda(
            q_values, actions, rewards, next_q_values, dones, self.gamma
        )