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>

__inline__ __device__ float warp_reduce(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

__global__ void wasserstein_energy_gelu_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ target, 
    float* __restrict__ y, 
    int batch_size, 
    int width) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    if (row >= batch_size) return;

    // Shared memory for Sorting
    // width * 2 floats
    extern __shared__ float s_mem[];
    float* s_x = s_mem;
    float* s_t = s_mem + width;

    // 1. Load Data
    for (int i = tid; i < width; i += blockDim.x) {
        s_x[i] = x[row * width + i];
        s_t[i] = target[i];
    }
    __syncthreads();

    // 2. Bitonic Sort X
    // Assumes width is power of 2
    for (int k = 2; k <= width; k <<= 1) {
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int i = tid; i < width; i += blockDim.x) {
                int ixj = i ^ j;
                if (ixj > i) {
                    bool up = ((i & k) == 0);
                    float v1 = s_x[i];
                    float v2 = s_x[ixj];
                    if (up ? (v1 > v2) : (v1 < v2)) {
                        s_x[i] = v2;
                        s_x[ixj] = v1;
                    }
                }
            }
            __syncthreads();
        }
    }

    // 3. Bitonic Sort Target
    for (int k = 2; k <= width; k <<= 1) {
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int i = tid; i < width; i += blockDim.x) {
                int ixj = i ^ j;
                if (ixj > i) {
                    bool up = ((i & k) == 0);
                    float v1 = s_t[i];
                    float v2 = s_t[ixj];
                    if (up ? (v1 > v2) : (v1 < v2)) {
                        s_t[i] = v2;
                        s_t[ixj] = v1;
                    }
                }
            }
            __syncthreads();
        }
    }

    // 4. Compute Squared Difference Sum (Energy Distance term)
    float sum_sq_diff = 0.0f;
    for (int i = tid; i < width; i += blockDim.x) {
        float diff = s_x[i] - s_t[i];
        sum_sq_diff += diff * diff;
    }

    sum_sq_diff = warp_reduce(sum_sq_diff);

    // Block Reduction
    static __shared__ float shared_sum[32];
    int lane = tid % 32;
    int wid = tid / 32;

    if (lane == 0) shared_sum[wid] = sum_sq_diff;
    __syncthreads();

    sum_sq_diff = (tid < blockDim.x / 32) ? shared_sum[lane] : 0.0f;
    if (wid == 0) sum_sq_diff = warp_reduce(sum_sq_diff);

    // 5. Final Calculation and GELU
    if (tid == 0) {
        float energy = sum_sq_diff / (float)width;

        float val = energy;
        float gelu = val * 0.5f * (1.0f + erff(val * 0.70710678f));
        y[row] = gelu;
    }
}

torch::Tensor launch_wasserstein_energy_gelu(torch::Tensor x, torch::Tensor target) {
    auto batch_size = x.size(0);
    auto width = x.size(1);
    auto y = torch::empty({batch_size}, x.options());

    const int threads = 256; 
    const int blocks = batch_size;
    int shared_mem = 2 * width * sizeof(float);

    wasserstein_energy_gelu_kernel<<<blocks, threads, shared_mem>>>(
        x.data_ptr<float>(), 
        target.data_ptr<float>(), 
        y.data_ptr<float>(), 
        batch_size, 
        width
    );
    return y;
}
"""

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

wasserstein_energy_gelu_module = load_inline(
    name='wasserstein_energy_gelu_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_wasserstein_energy_gelu'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, target):
        super(ModelNew, self).__init__()
        self.target = nn.Parameter(target)
        self.op = wasserstein_energy_gelu_module

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