import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


class ModelNew(nn.Module):
    def __init__(self, eps=1e-12):
        super().__init__()
        self.eps = eps
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor gower_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor ranges, float eps);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cuda_fp16.h>

        __device__ __forceinline__ float warp_sum_fast(float val) {
            #pragma unroll
            for (int mask = 16; mask > 0; mask >>= 1) {
                val += __shfl_xor_sync(0xffffffff, val, mask);
            }
            return val;
        }


        __global__ void gower_kernel_ultra(
            const float* __restrict__ x,
            const float* __restrict__ y,
            const float* __restrict__ ranges,
            float* __restrict__ output,
            const int feature_dim,
            const int batch_size,
            const float eps) 
        {
            extern __shared__ float s_ranges[];

            const int bid = blockIdx.x;
            const int tid = threadIdx.x;
            const int warp_id = tid >> 5;  // tid / 32
            const int lane_id = tid & 31;  // tid % 32

            for (int i = tid; i < feature_dim; i += blockDim.x) {
                s_ranges[i] = ranges[i];
            }
            __syncthreads();

            if (bid >= batch_size) return;

            const float* x_row = x + bid * feature_dim;
            const float* y_row = y + bid * feature_dim;

            float sum = 0.0f;  

            const int vec_loops = feature_dim >> 2;  // / 4
            const int vec_remainder = feature_dim & 3;  // % 4

            const float4* x_vec = reinterpret_cast<const float4*>(x_row);
            const float4* y_vec = reinterpret_cast<const float4*>(y_row);
            const float4* r_vec = reinterpret_cast<const float4*>(s_ranges);

            // ILP (Instruction-Level Parallelism): 展开 2 次
            #pragma unroll 2
            for (int i = tid; i < vec_loops; i += blockDim.x) {
                float4 vx = __ldg(&x_vec[i]);  // 使用只读缓存
                float4 vy = __ldg(&y_vec[i]);
                float4 vr = r_vec[i];

                float rcp_eps = __frcp_rn(eps);  
                sum += fabsf(vx.x - vy.x) * __frcp_rn(vr.x + eps);
                sum += fabsf(vx.y - vy.y) * __frcp_rn(vr.y + eps);
                sum += fabsf(vx.z - vy.z) * __frcp_rn(vr.z + eps);
                sum += fabsf(vx.w - vy.w) * __frcp_rn(vr.w + eps);
            }

            const int tail_start = vec_loops << 2;
            for (int i = tail_start + tid; i < feature_dim; i += blockDim.x) {
                float diff = fabsf(x_row[i] - y_row[i]);
                sum += diff * __frcp_rn(s_ranges[i] + eps);
            }

            sum = warp_sum_fast(sum);

            __shared__ float warp_sums[32];  // 最多支持 1024 线程 (32 warps)

            if (lane_id == 0) {
                warp_sums[warp_id] = sum;
            }
            __syncthreads();

            if (warp_id == 0) {
                sum = (lane_id < (blockDim.x >> 5)) ? warp_sums[lane_id] : 0.0f;
                sum = warp_sum_fast(sum);

                if (lane_id == 0) {
                    output[bid] = sum * __frcp_rn((float)feature_dim);  
                }
            }
        }

        __global__ void gower_kernel_large_batch(
            const float* __restrict__ x,
            const float* __restrict__ y,
            const float* __restrict__ ranges,
            float* __restrict__ output,
            const int feature_dim,
            const int batch_size,
            const float eps) 
        {
            extern __shared__ float s_ranges[];

            const int tid = threadIdx.x;

            for (int i = tid; i < feature_dim; i += blockDim.x) {
                s_ranges[i] = ranges[i];
            }
            __syncthreads();

            for (int bid = blockIdx.x; bid < batch_size; bid += gridDim.x) {
                const float* x_row = x + bid * feature_dim;
                const float* y_row = y + bid * feature_dim;

                float sum = 0.0f;

                const int vec_loops = feature_dim >> 2;
                const float4* x_vec = reinterpret_cast<const float4*>(x_row);
                const float4* y_vec = reinterpret_cast<const float4*>(y_row);
                const float4* r_vec = reinterpret_cast<const float4*>(s_ranges);

                #pragma unroll 2
                for (int i = tid; i < vec_loops; i += blockDim.x) {
                    float4 vx = __ldg(&x_vec[i]);
                    float4 vy = __ldg(&y_vec[i]);
                    float4 vr = r_vec[i];

                    sum += fabsf(vx.x - vy.x) * __frcp_rn(vr.x + eps);
                    sum += fabsf(vx.y - vy.y) * __frcp_rn(vr.y + eps);
                    sum += fabsf(vx.z - vy.z) * __frcp_rn(vr.z + eps);
                    sum += fabsf(vx.w - vy.w) * __frcp_rn(vr.w + eps);
                }

                const int tail_start = vec_loops << 2;
                for (int i = tail_start + tid; i < feature_dim; i += blockDim.x) {
                    sum += fabsf(x_row[i] - y_row[i]) * __frcp_rn(s_ranges[i] + eps);
                }

                sum = warp_sum_fast(sum);

                __shared__ float warp_sums[32];
                const int lane_id = tid & 31;
                const int warp_id = tid >> 5;

                if (lane_id == 0) warp_sums[warp_id] = sum;
                __syncthreads();

                if (warp_id == 0) {
                    sum = (lane_id < (blockDim.x >> 5)) ? warp_sums[lane_id] : 0.0f;
                    sum = warp_sum_fast(sum);
                    if (lane_id == 0) {
                        output[bid] = sum * __frcp_rn((float)feature_dim);
                    }
                }
                __syncthreads();  // 确保所有线程完成当前样本
            }
        }

        torch::Tensor gower_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor ranges, float eps) {
            auto x_c = x.contiguous();
            auto y_c = y.contiguous();
            auto r_c = ranges.contiguous();

            const int batch_size = x_c.size(0);
            const int feature_dim = x_c.size(1);

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

            const int shared_mem_size = feature_dim * sizeof(float);


            cudaDeviceProp prop;
            cudaGetDeviceProperties(&prop, 0);
            const int max_blocks = prop.multiProcessorCount * 8;  // 每个 SM 8 个 blocks

            if (batch_size <= max_blocks) {
                int threads = 256;
                gower_kernel_ultra<<<batch_size, threads, shared_mem_size>>>(
                    x_c.data_ptr<float>(),
                    y_c.data_ptr<float>(),
                    r_c.data_ptr<float>(),
                    output.data_ptr<float>(),
                    feature_dim,
                    batch_size,
                    eps
                );
            } else {
                // 大批量：Grid-Stride Loop
                int threads = 256;
                int blocks = min(max_blocks, batch_size);
                gower_kernel_large_batch<<<blocks, threads, shared_mem_size>>>(
                    x_c.data_ptr<float>(),
                    y_c.data_ptr<float>(),
                    r_c.data_ptr<float>(),
                    output.data_ptr<float>(),
                    feature_dim,
                    batch_size,
                    eps
                );
            }

            return output;
        }
        """

        self.op = load_inline(
            name="gower_ultra_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["gower_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math", "-Xptxas=-v"],
            verbose=False
        )

    def forward(self, x, y):

        if hasattr(torch, 'aminmax'):
            combined = torch.cat([x, y], dim=0)
            min_vals, max_vals = torch.aminmax(combined, dim=0)
            ranges = max_vals - min_vals
        else:

            combined = torch.cat([x, y], dim=0)
            ranges = combined.max(dim=0).values - combined.min(dim=0).values

        ranges = torch.clamp(ranges, min=self.eps)

        return self.op.gower_cuda(x, y, ranges, self.eps)