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-8

assert (C * H * W) % 4 == 0


class ModelNew(nn.Module):
    def __init__(self, unbiased=True):
        super().__init__()
        self.block_size = 512
        self.unbiased = unbiased
        self.register_buffer('temp_buffer', torch.zeros((N, 5), dtype=torch.float32))
        self._compile_cuda_kernel()

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

        void pearson_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            torch::Tensor temp_buffer,
            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 = WARP_SIZE / 2; offset > 0; offset /= 2) {{
                val += __shfl_down_sync(0xffffffff, val, offset);
            }}
            return val;
        }}

        __global__ __launch_bounds__(BLOCK_SIZE)
        void pearson_split_kernel(
            const float* __restrict__ x,
            const float* __restrict__ y,
            float* __restrict__ temp_buffer,
            int D_vec_total
        ) {{
            const int n_idx = blockIdx.y;
            const int split_idx = blockIdx.x;
            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_x = reinterpret_cast<const float4*>(x + batch_offset) + start_idx + threadIdx.x;
            const float4* curr_y = reinterpret_cast<const float4*>(y + batch_offset) + start_idx + threadIdx.x;
            const float4* end_ptr = reinterpret_cast<const float4*>(x + batch_offset) + end_idx;

            float sum_x[ILP];
            float sum_y[ILP];
            float sum_xx[ILP];
            float sum_yy[ILP];
            float sum_xy[ILP];

            #pragma unroll
            for (int k=0; k<ILP; ++k) {{
                sum_x[k] = 0.0f;
                sum_y[k] = 0.0f;
                sum_xx[k] = 0.0f;
                sum_yy[k] = 0.0f;
                sum_xy[k] = 0.0f;
            }}

            const int stride = BLOCK_SIZE * ILP;

            while (curr_x + (ILP - 1) * BLOCK_SIZE < end_ptr) {{
                float4 r_x[ILP];
                float4 r_y[ILP];

                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    r_x[k] = __ldg(curr_x + k * BLOCK_SIZE);
                    r_y[k] = __ldg(curr_y + k * BLOCK_SIZE);
                }}

                #pragma unroll
                for (int k = 0; k < ILP; ++k) {{
                    float vx = r_x[k].x; float vy = r_y[k].x;
                    sum_x[k] += vx; sum_y[k] += vy;
                    sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;

                    vx = r_x[k].y; vy = r_y[k].y;
                    sum_x[k] += vx; sum_y[k] += vy;
                    sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;

                    vx = r_x[k].z; vy = r_y[k].z;
                    sum_x[k] += vx; sum_y[k] += vy;
                    sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;

                    vx = r_x[k].w; vy = r_y[k].w;
                    sum_x[k] += vx; sum_y[k] += vy;
                    sum_xx[k] += vx * vx; sum_yy[k] += vy * vy; sum_xy[k] += vx * vy;
                }}

                curr_x += stride;
                curr_y += stride;
            }}

            while (curr_x < end_ptr) {{
                float4 vx = __ldg(curr_x);
                float4 vy = __ldg(curr_y);

                float v1 = vx.x; float v2 = vy.x;
                sum_x[0] += v1; sum_y[0] += v2;
                sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;

                v1 = vx.y; v2 = vy.y;
                sum_x[0] += v1; sum_y[0] += v2;
                sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;

                v1 = vx.z; v2 = vy.z;
                sum_x[0] += v1; sum_y[0] += v2;
                sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;

                v1 = vx.w; v2 = vy.w;
                sum_x[0] += v1; sum_y[0] += v2;
                sum_xx[0] += v1*v1; sum_yy[0] += v2*v2; sum_xy[0] += v1*v2;

                curr_x += BLOCK_SIZE;
                curr_y += BLOCK_SIZE;
            }}

            float l_x = 0.0f, l_y = 0.0f, l_xx = 0.0f, l_yy = 0.0f, l_xy = 0.0f;
            #pragma unroll
            for (int k=0; k<ILP; ++k) {{
                l_x += sum_x[k];
                l_y += sum_y[k];
                l_xx += sum_xx[k];
                l_yy += sum_yy[k];
                l_xy += sum_xy[k];
            }}

            l_x = warp_reduce_sum(l_x);
            l_y = warp_reduce_sum(l_y);
            l_xx = warp_reduce_sum(l_xx);
            l_yy = warp_reduce_sum(l_yy);
            l_xy = warp_reduce_sum(l_xy);

            __shared__ float s_x[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_y[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_xx[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_yy[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_xy[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_x[warp_id] = l_x;
                s_y[warp_id] = l_y;
                s_xx[warp_id] = l_xx;
                s_yy[warp_id] = l_yy;
                s_xy[warp_id] = l_xy;
            }}
            __syncthreads();

            if (warp_id == 0) {{
                float b_x = 0.0f, b_y = 0.0f, b_xx = 0.0f, b_yy = 0.0f, b_xy = 0.0f;
                if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
                    b_x = s_x[lane_id];
                    b_y = s_y[lane_id];
                    b_xx = s_xx[lane_id];
                    b_yy = s_yy[lane_id];
                    b_xy = s_xy[lane_id];
                }}

                b_x = warp_reduce_sum(b_x);
                b_y = warp_reduce_sum(b_y);
                b_xx = warp_reduce_sum(b_xx);
                b_yy = warp_reduce_sum(b_yy);
                b_xy = warp_reduce_sum(b_xy);

                if (lane_id == 0) {{
                    float* dst = temp_buffer + n_idx * 5;
                    atomicAdd(&dst[0], b_x);
                    atomicAdd(&dst[1], b_y);
                    atomicAdd(&dst[2], b_xx);
                    atomicAdd(&dst[3], b_yy);
                    atomicAdd(&dst[4], b_xy);
                }}
            }}
        }}

        void pearson_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            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 + 1024 - 1) / 1024;
            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);

            pearson_split_kernel<<<blocks, threads>>>(
                x.data_ptr<float>(),
                y.data_ptr<float>(),
                temp_buffer.data_ptr<float>(),
                D_vec
            );
        }}
        """

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

    def forward(self, x, y):
        if not x.is_contiguous(): x = x.contiguous()
        if not y.is_contiguous(): y = y.contiguous()

        if not x.is_cuda: x = x.cuda()
        if not y.is_cuda: y = y.cuda()

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

        self.temp_buffer.zero_()

        self.op.pearson_sum_cuda(x, y, self.temp_buffer, N, D)

        sum_x = self.temp_buffer[:, 0]
        sum_y = self.temp_buffer[:, 1]
        sum_xy = self.temp_buffer[:, 4]

        mean_x = sum_x / D
        mean_y = sum_y / D

        cov_sum = sum_xy - D * mean_x * mean_y

        if self.unbiased:
            divisor = D - 1
        else:
            divisor = D

        if divisor <= 0:
            return cov_sum * 0.0

        return cov_sum / divisor