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

N, C, H, W = 32, 64, 56, 56
EPS = 1e-6

assert (C * H * W) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"


class ModelNew(nn.Module):

    def __init__(self):
        super().__init__()
        self.block_size = 512
        self.eps = EPS

        self.register_buffer('temp_buffer', torch.zeros((N, 3), dtype=torch.float32))

        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        #include <cuda_runtime.h>

        void jaccard_fused_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            torch::Tensor temp_buffer, // [N, 3]
            int N,
            int D);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE {self.block_size}
        #define WARP_SIZE 32
        #define ILP 4  


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


        __global__ __launch_bounds__(BLOCK_SIZE)
        void jaccard_split_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            float* __restrict__ temp_buffer, // [N, 3]
            int D_vec_total  // D / 4
        ) {{

            const int n_idx = blockIdx.y;     // Batch Index
            const int split_idx = blockIdx.x; // Split Index
            const int num_splits = gridDim.x;


            const int chunk_size = (D_vec_total + num_splits - 1) / num_splits;
            const int start_idx = split_idx * chunk_size;
            const int end_idx = min(start_idx + chunk_size, D_vec_total);

            if (start_idx >= D_vec_total) return;


            const int64_t batch_offset = (int64_t)n_idx * D_vec_total * 4;

            const float4* curr_x1 = reinterpret_cast<const float4*>(x1 + batch_offset) + start_idx + threadIdx.x;
            const float4* curr_x2 = reinterpret_cast<const float4*>(x2 + batch_offset) + start_idx + threadIdx.x;
            const float4* end_ptr = reinterpret_cast<const float4*>(x1 + batch_offset) + end_idx;


            float acc_dot[ILP];
            float acc_sq1[ILP];
            float acc_sq2[ILP];

            #pragma unroll
            for (int k=0; k<ILP; ++k) {{
                acc_dot[k] = 0.0f;
                acc_sq1[k] = 0.0f;
                acc_sq2[k] = 0.0f;
            }}

            const int stride = BLOCK_SIZE * ILP;


            while (curr_x1 + (ILP - 1) * BLOCK_SIZE < end_ptr) {{
                float4 r1[ILP];
                float4 r2[ILP];


                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    r1[k] = __ldg(curr_x1 + k * BLOCK_SIZE);
                    r2[k] = __ldg(curr_x2 + k * BLOCK_SIZE);
                }}

                // Compute (FMA 优化)
                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    // Dot: x1 * x2
                    acc_dot[k] += r1[k].x * r2[k].x;
                    acc_dot[k] += r1[k].y * r2[k].y;
                    acc_dot[k] += r1[k].z * r2[k].z;
                    acc_dot[k] += r1[k].w * r2[k].w;

                    // Sq1: x1 * x1
                    acc_sq1[k] += r1[k].x * r1[k].x;
                    acc_sq1[k] += r1[k].y * r1[k].y;
                    acc_sq1[k] += r1[k].z * r1[k].z;
                    acc_sq1[k] += r1[k].w * r1[k].w;

                    // Sq2: x2 * x2
                    acc_sq2[k] += r2[k].x * r2[k].x;
                    acc_sq2[k] += r2[k].y * r2[k].y;
                    acc_sq2[k] += r2[k].z * r2[k].z;
                    acc_sq2[k] += r2[k].w * r2[k].w;
                }}

                curr_x1 += stride;
                curr_x2 += stride;
            }}


            while (curr_x1 < end_ptr) {{
                float4 v1 = __ldg(curr_x1);
                float4 v2 = __ldg(curr_x2);

                acc_dot[0] += v1.x * v2.x + v1.y * v2.y + v1.z * v2.z + v1.w * v2.w;
                acc_sq1[0] += v1.x * v1.x + v1.y * v1.y + v1.z * v1.z + v1.w * v1.w;
                acc_sq2[0] += v2.x * v2.x + v2.y * v2.y + v2.z * v2.z + v2.w * v2.w;

                curr_x1 += BLOCK_SIZE;
                curr_x2 += BLOCK_SIZE;
            }}


            float sum_dot = 0.0f;
            float sum_sq1 = 0.0f;
            float sum_sq2 = 0.0f;

            #pragma unroll
            for (int k=0; k<ILP; ++k) {{
                sum_dot += acc_dot[k];
                sum_sq1 += acc_sq1[k];
                sum_sq2 += acc_sq2[k];
            }}

            // 8. Warp 归约
            sum_dot = warp_reduce_sum(sum_dot);
            sum_sq1 = warp_reduce_sum(sum_sq1);
            sum_sq2 = warp_reduce_sum(sum_sq2);

            // 9. Block 归约 (使用 Shared Memory)
            __shared__ float s_dot[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_sq1[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_sq2[BLOCK_SIZE / WARP_SIZE];

            const int lane_id = threadIdx.x % WARP_SIZE;
            const int warp_id = threadIdx.x / WARP_SIZE;

            if (lane_id == 0) {{
                s_dot[warp_id] = sum_dot;
                s_sq1[warp_id] = sum_sq1;
                s_sq2[warp_id] = sum_sq2;
            }}
            __syncthreads();

            if (warp_id == 0) {{
                float block_dot = 0.0f;
                float block_sq1 = 0.0f;
                float block_sq2 = 0.0f;

                if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
                    block_dot = s_dot[lane_id];
                    block_sq1 = s_sq1[lane_id];
                    block_sq2 = s_sq2[lane_id];
                }}

                block_dot = warp_reduce_sum(block_dot);
                block_sq1 = warp_reduce_sum(block_sq1);
                block_sq2 = warp_reduce_sum(block_sq2);

                // 10. 原子累加到全局 Temp Buffer
                // Buffer Layout: [N, 3] -> Row n: [dot, sq1, sq2]
                if (lane_id == 0) {{
                    float* dst = temp_buffer + n_idx * 3;
                    atomicAdd(&dst[0], block_dot);
                    atomicAdd(&dst[1], block_sq1);
                    atomicAdd(&dst[2], block_sq2);
                }}
            }}
        }}

        void jaccard_fused_cuda(
            torch::Tensor x1,
            torch::Tensor x2,
            torch::Tensor temp_buffer,
            int N,
            int D) 
        {{
            int D_vec = D / 4;

            int device_id;
            cudaGetDevice(&device_id);
            int sm_count;
            cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id);


            int target_blocks = sm_count * 4; 
            int splits = (target_blocks + N - 1) / N;


            int max_splits = (D_vec + 512 - 1) / 512;
            if (splits > max_splits) splits = max_splits;
            if (splits < 1) splits = 1;

            if (splits > 512) splits = 512;

            dim3 blocks(splits, N);
            dim3 threads(BLOCK_SIZE);

            jaccard_split_kernel<<<blocks, threads>>>(
                x1.data_ptr<float>(),
                x2.data_ptr<float>(),
                temp_buffer.data_ptr<float>(),
                D_vec
            );
        }}
        """

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

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        if not x1.is_contiguous(): x1 = x1.contiguous()
        if not x2.is_contiguous(): x2 = x2.contiguous()

        N, C, H, W = x1.size()
        D = C * H * W

        self.temp_buffer.zero_()

        self.op.jaccard_fused_cuda(
            x1,
            x2,
            self.temp_buffer,
            N,
            D
        )

        dot = self.temp_buffer[:, 0]
        sq1 = self.temp_buffer[:, 1]
        sq2 = self.temp_buffer[:, 2]

        union = sq1 + sq2 - dot
        return dot / (union + self.eps)