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 trpo_kernel(const float* log_probs, const float* old_log_probs, const float* advantages, const float* old_probs, const float* new_probs, float* surrogate_out, float* kl_out, int batch_size, int action_dim) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < batch_size) {
        float ratio = expf(log_probs[idx] - old_log_probs[idx]);
        surrogate_out[idx] = -ratio * advantages[idx];

        float kl_sum = 0.0f;
        for (int i = 0; i < action_dim; i++) {
            float old_p = old_probs[idx * action_dim + i];
            float new_p = new_probs[idx * action_dim + i];
            if (old_p > 1e-8f && new_p > 1e-8f) {
                kl_sum += old_p * (logf(old_p) - logf(new_p));
            }
        }
        kl_out[idx] = kl_sum;
    }
}

torch::Tensor trpo_cuda(torch::Tensor log_probs, torch::Tensor old_log_probs, torch::Tensor advantages, torch::Tensor old_probs, torch::Tensor new_probs, float max_kl) {
    auto batch_size = log_probs.size(0);
    auto action_dim = old_probs.size(1);
    auto surrogate_out = torch::empty({batch_size}, log_probs.options());
    auto kl_out = torch::empty({batch_size}, log_probs.options());
    const int block_size = 256;
    int num_blocks = (batch_size + block_size - 1) / block_size;
    trpo_kernel<<<num_blocks, block_size>>>(log_probs.data_ptr<float>(), old_log_probs.data_ptr<float>(), advantages.data_ptr<float>(), old_probs.data_ptr<float>(), new_probs.data_ptr<float>(), surrogate_out.data_ptr<float>(), kl_out.data_ptr<float>(), batch_size, action_dim);
    return surrogate_out.mean() + max_kl * kl_out.mean();
}
"""

cpp_source = """
torch::Tensor trpo_cuda(torch::Tensor log_probs, torch::Tensor old_log_probs, torch::Tensor advantages, torch::Tensor old_probs, torch::Tensor new_probs, float max_kl);
"""

trpo_module = load_inline(
    name="trpo_loss",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["trpo_cuda"],
    verbose=True
)


class ModelNew(torch.nn.Module):
    def __init__(self, max_kl):
        super(ModelNew, self).__init__()
        self.max_kl = max_kl
        self.trpo_module = trpo_module

    def forward(self, log_probs, old_log_probs, advantages, old_probs, new_probs):
        return self.trpo_module.trpo_cuda(log_probs, old_log_probs, advantages, old_probs, new_probs, self.max_kl)