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


class ModelNew(nn.Module):
    def __init__(self, gamma):
        super().__init__()
        self.gamma = gamma.item() if isinstance(gamma, torch.Tensor) else gamma
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor dqn_loss_cuda(
            torch::Tensor q_values,
            torch::Tensor actions,
            torch::Tensor rewards,
            torch::Tensor next_q_values,
            torch::Tensor dones,
            float gamma);
        """

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

        __global__ void dqn_loss_kernel(
            const float* __restrict__ q_values,
            const long* __restrict__ actions,
            const float* __restrict__ rewards,
            const float* __restrict__ next_q_values,
            const float* __restrict__ dones,
            float* __restrict__ output,
            float gamma,
            int batch_size,
            int action_dim)
        {
            extern __shared__ float sdata[];
            unsigned int tid = threadIdx.x;
            unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
            unsigned int gridSize = blockDim.x * gridDim.x;

            float local_sum = 0.0f;

            while (i < batch_size) {
                int row_offset = i * action_dim;
                long action = actions[i];
                float q_pred = q_values[row_offset + action];

                float max_next_q = -1e20f;
                for (int a = 0; a < action_dim; ++a) {
                    float val = next_q_values[row_offset + a];
                    if (val > max_next_q) {
                        max_next_q = val;
                    }
                }

                float q_target = rewards[i] + gamma * max_next_q * (1.0f - dones[i]);
                float diff = q_pred - q_target;
                local_sum += diff * diff;

                i += gridSize;
            }

            sdata[tid] = local_sum;
            __syncthreads();

            for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
                if (tid < s) {
                    sdata[tid] += sdata[tid + s];
                }
                __syncthreads();
            }

            if (tid == 0) {
                atomicAdd(output, sdata[0] / batch_size);
            }
        }

        torch::Tensor dqn_loss_cuda(
            torch::Tensor q_values,
            torch::Tensor actions,
            torch::Tensor rewards,
            torch::Tensor next_q_values,
            torch::Tensor dones,
            float gamma)
        {
            auto q_values_c = q_values.contiguous();
            auto actions_c = actions.contiguous();
            auto rewards_c = rewards.contiguous();
            auto next_q_values_c = next_q_values.contiguous();
            auto dones_c = dones.contiguous();

            int batch_size = q_values_c.size(0);
            int action_dim = q_values_c.size(1);

            auto output = torch::zeros({1}, q_values.options());

            const int threads = 256;
            const int blocks = min((batch_size + threads - 1) / threads, 1024);
            const int shared_mem = threads * sizeof(float);

            dqn_loss_kernel<<<blocks, threads, shared_mem>>>(
                q_values_c.data_ptr<float>(),
                actions_c.data_ptr<long>(),
                rewards_c.data_ptr<float>(),
                next_q_values_c.data_ptr<float>(),
                dones_c.data_ptr<float>(),
                output.data_ptr<float>(),
                gamma,
                batch_size,
                action_dim
            );

            return output[0];
        }
        """

        self.op = load_inline(
            name="dqn_loss_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["dqn_loss_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, q_values, actions, rewards, next_q_values, dones):
        return self.op.dqn_loss_cuda(q_values, actions, rewards, next_q_values, dones, self.gamma)