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


class ModelNew(nn.Module):
    def __init__(self, feature_dim):
        super().__init__()
        self.register_buffer("mean", torch.zeros(feature_dim))
        self.register_buffer("inv_cov", torch.eye(feature_dim))
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor mahalanobis_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor inv_cov);
        """

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

        static cublasHandle_t handle = nullptr;

        void init_cublas() {
            if (!handle) {
                cublasCreate(&handle);
                // 尝试使用 Tensor Cores 加速（如果硬件支持）
                cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);
            }
        }

        __global__ void sub_mean_kernel(
            const float* __restrict__ x,
            const float* __restrict__ mean,
            float* __restrict__ diff,
            int batch_size,
            int feature_dim) 
        {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= batch_size * feature_dim) return;

            int col = idx % feature_dim;
            diff[idx] = x[idx] - mean[col];
        }

        __device__ __forceinline__ double warp_sum(double val) {
            for (int offset = 16; offset > 0; offset /= 2) {
                val += __shfl_down_sync(0xffffffff, val, offset);
            }
            return val;
        }

        __device__ __forceinline__ double block_sum(double val) {
            static __shared__ double shared[32];
            int lane = threadIdx.x % 32;
            int wid = threadIdx.x / 32;

            val = warp_sum(val);
            if (lane == 0) shared[wid] = val;
            __syncthreads();

            val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0;
            if (wid == 0) val = warp_sum(val);
            return val;
        }

        __global__ void reduction_kernel(
            const float* __restrict__ diff,
            const float* __restrict__ temp,
            float* __restrict__ output,
            int batch_size,
            int feature_dim)
        {
            int bid = blockIdx.x;
            if (bid >= batch_size) return;

            double sum = 0.0;
            int offset = bid * feature_dim;

            for (int i = threadIdx.x; i < feature_dim; i += blockDim.x) {
                sum += (double)diff[offset + i] * (double)temp[offset + i];
            }

            sum = block_sum(sum);

            if (threadIdx.x == 0) {
                output[bid] = sqrtf((float)max(sum, 0.0));
            }
        }

        torch::Tensor mahalanobis_cuda(torch::Tensor x, torch::Tensor mean, torch::Tensor inv_cov) {
            init_cublas();

            auto x_c = x.contiguous();
            auto m_c = mean.contiguous();
            auto iv_c = inv_cov.contiguous();

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

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

            int total_elements = batch_size * feature_dim;
            int threads = 256;
            int blocks = (total_elements + threads - 1) / threads;

            sub_mean_kernel<<<blocks, threads>>>(
                x_c.data_ptr<float>(),
                m_c.data_ptr<float>(),
                diff.data_ptr<float>(),
                batch_size,
                feature_dim
            );

            float alpha = 1.0f;
            float beta = 0.0f;

            cublasStatus_t status = cublasGemmEx(
                handle,
                CUBLAS_OP_N, CUBLAS_OP_N,
                feature_dim, batch_size, feature_dim,
                &alpha,
                iv_c.data_ptr<float>(), CUDA_R_32F, feature_dim,
                diff.data_ptr<float>(), CUDA_R_32F, feature_dim,
                &beta,
                temp.data_ptr<float>(), CUDA_R_32F, feature_dim,
                CUBLAS_COMPUTE_32F_FAST_TF32,  // 使用 TF32 加速（Ampere+）
                CUBLAS_GEMM_DEFAULT
            );

            if (status != CUBLAS_STATUS_SUCCESS) {
                cublasSgemm(
                    handle,
                    CUBLAS_OP_N, CUBLAS_OP_N,
                    feature_dim, batch_size, feature_dim,
                    &alpha,
                    iv_c.data_ptr<float>(), feature_dim,
                    diff.data_ptr<float>(), feature_dim,
                    &beta,
                    temp.data_ptr<float>(), feature_dim
                );
            }

            reduction_kernel<<<batch_size, 256>>>(
                diff.data_ptr<float>(),
                temp.data_ptr<float>(),
                output.data_ptr<float>(),
                batch_size,
                feature_dim
            );

            return output;
        }
        """

        self.op = load_inline(
            name="mahalanobis_fixed_v3",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["mahalanobis_cuda"],
            extra_cuda_cflags=["-O3"],
            extra_ldflags=["-lcublas"],
            verbose=False
        )

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