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

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

__global__ void gamma_divergence_kernel(
    const float* __restrict__ p,
    const float* __restrict__ q,
    float* __restrict__ output,
    int num_classes,
    float gamma
) {
    extern __shared__ float sdata[];

    int tid = threadIdx.x;
    int bid = blockIdx.x; 

    int row_offset = bid * num_classes;

    float* s_p = sdata;
    float* s_pq = sdata + blockDim.x;
    float* s_q = sdata + 2 * blockDim.x;

    float local_p = 0.0f;
    float local_pq = 0.0f;
    float local_q = 0.0f;

    for (int i = tid; i < num_classes; i += blockDim.x) {
        float p_val = p[row_offset + i];
        float q_val = q[row_offset + i];

        float q_pow_g = powf(q_val, gamma);

        local_p += powf(p_val, 1.0f + gamma);
        local_pq += p_val * q_pow_g;
        local_q += powf(q_val, 1.0f + gamma);
    }

    s_p[tid] = local_p;
    s_pq[tid] = local_pq;
    s_q[tid] = local_q;
    __syncthreads();

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

    if (tid == 0) {
        float sum_p = s_p[0];
        float sum_pq = s_pq[0];
        float sum_q = s_q[0];

        float term1 = logf(sum_p) / (gamma * (1.0f + gamma));
        float term2 = logf(sum_pq) / gamma;
        float term3 = logf(sum_q) / (1.0f + gamma);

        output[bid] = term1 - term2 + term3;
    }
}

torch::Tensor gamma_divergence_cuda(torch::Tensor p, torch::Tensor q, float gamma) {
    int batch_size = p.size(0);
    int num_classes = p.size(1);

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

    int threads = 256;
    int blocks = batch_size;
    int shared_mem = 3 * threads * sizeof(float);

    gamma_divergence_kernel<<<blocks, threads, shared_mem>>>(
        p.data_ptr<float>(),
        q.data_ptr<float>(),
        output.data_ptr<float>(),
        num_classes,
        gamma
    );

    return output.mean();
}
"""

cpp_source = """
torch::Tensor gamma_divergence_cuda(torch::Tensor p, torch::Tensor q, float gamma);
"""

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


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

    def forward(self, p, q):
        p_prob = F.softmax(p, dim=1)
        q_prob = F.softmax(q, dim=1)
        return gamma_divergence_loss.gamma_divergence_cuda(p_prob, q_prob, self.gamma)