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>
#include <math.h>

__global__ void compute_cost_matrix_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ target, 
    float* __restrict__ M, 
    int batch_size, 
    int dim, 
    float epsilon) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elements = batch_size * batch_size;

    if (idx < total_elements) {
        int row = idx / batch_size;
        int col = idx % batch_size;

        float dist_sq = 0.0f;
        const float* x_row = x + row * dim;
        const float* t_row = target + col * dim;

        for (int k = 0; k < dim; ++k) {
            float diff = x_row[k] - t_row[k];
            dist_sq += diff * diff;
        }

        M[idx] = -dist_sq / epsilon;
    }
}

__global__ void sinkhorn_solver_kernel(
    const float* __restrict__ M, 
    float* __restrict__ out_loss, 
    int batch_size, 
    float epsilon,
    int iters) 
{
    extern __shared__ float s_mem[];
    float* s_f = s_mem;
    float* s_g = s_f + batch_size;
    float* s_red = s_g + batch_size; 

    int tid = threadIdx.x;
    int b = batch_size;

    if (tid < b) {
        s_f[tid] = 0.0f;
        s_g[tid] = 0.0f;
    }
    __syncthreads();

    for (int iter = 0; iter < iters; ++iter) {
        if (tid < b) {
            float max_val = -3.402823466e+38F;
            for (int j = 0; j < b; ++j) {
                float val = M[tid * b + j] + s_g[j];
                if (val > max_val) max_val = val;
            }

            float sum_exp = 0.0f;
            for (int j = 0; j < b; ++j) {
                sum_exp += expf(M[tid * b + j] + s_g[j] - max_val);
            }
            s_f[tid] = -(logf(sum_exp) + max_val);
        }
        __syncthreads();

        if (tid < b) {
            float max_val = -3.402823466e+38F;
            for (int i = 0; i < b; ++i) {
                float val = M[i * b + tid] + s_f[i];
                if (val > max_val) max_val = val;
            }

            float sum_exp = 0.0f;
            for (int i = 0; i < b; ++i) {
                sum_exp += expf(M[i * b + tid] + s_f[i] - max_val);
            }
            s_g[tid] = -(logf(sum_exp) + max_val);
        }
        __syncthreads();
    }

    float row_sum = 0.0f;
    if (tid < b) {
        float fi = s_f[tid];
        for (int j = 0; j < b; ++j) {
            float Mij = M[tid * b + j];
            float gj = s_g[j];

            float log_P = Mij + fi + gj;
            float P = expf(log_P);
            float C = -Mij * epsilon;

            row_sum += P * C;
        }
    }

    float val = row_sum;
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);

    if ((tid % 32) == 0) {
        s_red[tid / 32] = val;
    }
    __syncthreads();

    if (tid == 0) {
        float total_loss = 0.0f;
        int warps = (blockDim.x + 31) / 32;
        for (int i = 0; i < warps; ++i) {
            total_loss += s_red[i];
        }
        out_loss[0] = total_loss / (float)b;
    }
}

torch::Tensor launch_sinkhorn_loss(torch::Tensor x, torch::Tensor target) {
    auto batch_size = x.size(0);
    auto dim = x.size(1);
    auto M = torch::empty({batch_size, batch_size}, x.options());
    auto out = torch::empty({1}, x.options());

   
    int total_elements = batch_size * batch_size;
    int threads_m = 256;
    int blocks_m = (total_elements + threads_m - 1) / threads_m;

    compute_cost_matrix_kernel<<<blocks_m, threads_m>>>(
        x.data_ptr<float>(),
        target.data_ptr<float>(),
        M.data_ptr<float>(),
        batch_size,
        dim,
        0.1f
    );


    int threads_s = batch_size; 
    int blocks_s = 1;
    int shared_mem = (2 * batch_size + 32) * sizeof(float);

    sinkhorn_solver_kernel<<<blocks_s, threads_s, shared_mem>>>(
        M.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        0.1f,
        5
    );

    return out;
}
"""

cpp_source = """
torch::Tensor launch_sinkhorn_loss(torch::Tensor x, torch::Tensor target);
"""

sinkhorn_loss_module = load_inline(
    name='hungarian_loss_fast',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_sinkhorn_loss'],
    verbose=False
)


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

    def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.op.launch_sinkhorn_loss(x.contiguous(), target.contiguous())