import torch
from torch.utils.cpp_extension import load_inline

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

#define BLOCK_SIZE 256

__global__ void canberra_focal_loss_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ target,
    float* __restrict__ losses,
    int feature_dim,
    float alpha,
    float gamma
) {
    extern __shared__ float sdata[];
    
    int tid = threadIdx.x;
    int sample_idx = blockIdx.x;
    
    float partial_sum = 0.0f;
    int base_addr = sample_idx * feature_dim;
    
    for (int i = tid; i < feature_dim; i += blockDim.x) {
        float x_val = x[base_addr + i];
        float y_val = y[base_addr + i];
        
        float abs_diff = fabsf(x_val - y_val);
        float abs_sum = fabsf(x_val) + fabsf(y_val);
        
        float canberra_term = 0.0f;
        if (abs_sum > 0.0f) {
            canberra_term = abs_diff / abs_sum;
        }
        
        // Apply focal modulation element-wise
        float focal_weight = alpha * powf(1.0f + canberra_term, -gamma);
        partial_sum += focal_weight * canberra_term;
    }
    
    sdata[tid] = partial_sum;
    __syncthreads();
    
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        losses[sample_idx] = sdata[0];
    }
}

torch::Tensor canberra_focal_loss_cuda(
    torch::Tensor x, 
    torch::Tensor y, 
    torch::Tensor target,
    float alpha,
    float gamma
) {
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
    TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
    TORCH_CHECK(x.sizes() == y.sizes(), "X and Y must have the same shape");

    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    auto target_contig = target.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    auto losses = torch::zeros({batch_size}, x.options());
    
    const int block_size = BLOCK_SIZE;
    size_t shared_mem = block_size * sizeof(float);
    
    canberra_focal_loss_kernel<<<batch_size, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        target_contig.data_ptr<float>(),
        losses.data_ptr<float>(),
        feature_dim,
        alpha,
        gamma
    );
    
    return losses;
}
"""

canberra_focal_cpp_source = """
torch::Tensor canberra_focal_loss_cuda(
    torch::Tensor x, 
    torch::Tensor y, 
    torch::Tensor target,
    float alpha,
    float gamma
);
"""

canberra_focal = load_inline(
    name="canberra_focal_loss_fixed",
    cpp_sources=canberra_focal_cpp_source,
    cuda_sources=canberra_focal_source,
    functions=["canberra_focal_loss_cuda"],
    extra_cuda_cflags=["-O3"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, alpha=1.0, gamma=2.0):
        super(ModelNew, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.canberra_focal = canberra_focal

    def forward(self, x, y, target):
        return self.canberra_focal.canberra_focal_loss_cuda(x, y, target, self.alpha, self.gamma)
