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 quantile_normalize_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int batch_size,
    int dim,
    float min_val,
    float max_val) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * dim;

    if (idx < total) {
        int b = idx / dim;
        int i = idx % dim;

        const float* x_ptr = input + b * dim;
        float val = x_ptr[i];

        int rank = 0;
        for (int j = 0; j < dim; j++) {
            if (x_ptr[j] < val || (x_ptr[j] == val && j < i)) {
                rank++;
            }
        }

        float norm = (float)rank / (dim - 1);
        float range = max_val - min_val;
        output[idx] = norm * range + min_val;
    }
}

torch::Tensor quantile_normalize_cuda(
    torch::Tensor x,
    float min_val,
    float max_val) {

    int batch_size = x.size(0);
    int dim = x.size(1);

    auto output = torch::empty_like(x);

    int total = batch_size * dim;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    quantile_normalize_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        dim,
        min_val,
        max_val
    );

    return output;
}
"""

cpp_source = """
torch::Tensor quantile_normalize_cuda(
    torch::Tensor x,
    float min_val,
    float max_val);
"""

cuda_module = load_inline(
    name="quantile_normalize_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["quantile_normalize_cuda"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, min_val, max_val):
        super(ModelNew, self).__init__()
        self.min_val = min_val
        self.max_val = max_val
        self.range = max_val - min_val

    def forward(self, x):
        return cuda_module.quantile_normalize_cuda(x, self.min_val, self.max_val)