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 mahalanobis_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mean,
    const float* __restrict__ inv_cov,
    float* __restrict__ output,
    int batch_size,
    int dim
) {
    extern __shared__ float s_diff[]; // Size: dim

    int b = blockIdx.x;
    if (b >= batch_size) return;

    int tid = threadIdx.x;


    for (int i = tid; i < dim; i += blockDim.x) {
        s_diff[i] = x[b * dim + i] - mean[i];
    }
    __syncthreads();



    float local_sum = 0.0f;
    int num_elements = dim * dim;

    for (int idx = tid; idx < num_elements; idx += blockDim.x) {
        int row = idx / dim;
        int col = idx % dim;
        float val = inv_cov[idx];
        local_sum += s_diff[row] * val * s_diff[col];
    }



    for (int offset = 16; offset > 0; offset /= 2) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }



    if (tid % 32 == 0) {
        s_diff[tid / 32] = local_sum;
    }
    __syncthreads();

    if (tid < 32) {
        float warp_sum = (tid < (blockDim.x + 31) / 32) ? s_diff[tid] : 0.0f;
        for (int offset = 16; offset > 0; offset /= 2) {
            warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
        }
        if (tid == 0) {
            output[b] = sqrtf(warp_sum + 1e-6f);
        }
    }
}

torch::Tensor mahalanobis_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor inv_cov) {
    int batch_size = x.size(0);
    int dim = x.size(1);

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

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

    mahalanobis_kernel<<<blocks, threads, shared_mem>>>(
        x.data_ptr<float>(),
        mean.data_ptr<float>(),
        inv_cov.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        dim
    );

    return output.mean();
}
"""

cpp_source = """
torch::Tensor mahalanobis_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor inv_cov);
"""

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


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

    def forward(self, x, mean, inv_cov):
        return mahalanobis_loss.mahalanobis_cuda(x, mean, inv_cov)