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._compile_cuda_kernel()

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

        void bhattacharyya_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            torch::Tensor out,
            float eps,
            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 bhattacharyya_split_kernel(
            const float* __restrict__ x,
            const float* __restrict__ y,
            float* __restrict__ out,
            float eps,
            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_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[ILP];
            #pragma unroll
            for (int k=0; k<ILP; ++k) sum[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) {{
                    // x
                    float p_x = r_x[k].x * r_y[k].x; sum[k] += sqrtf(p_x + eps);
                    // y
                    float p_y = r_x[k].y * r_y[k].y; sum[k] += sqrtf(p_y + eps);
                    // z
                    float p_z = r_x[k].z * r_y[k].z; sum[k] += sqrtf(p_z + eps);
                    // w
                    float p_w = r_x[k].w * r_y[k].w; sum[k] += sqrtf(p_w + eps);
                }}

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

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

                sum[0] += sqrtf(vx.x * vy.x + eps);
                sum[0] += sqrtf(vx.y * vy.y + eps);
                sum[0] += sqrtf(vx.z * vy.z + eps);
                sum[0] += sqrtf(vx.w * vy.w + eps);

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

            float local_sum = 0.0f;
            #pragma unroll
            for (int k=0; k<ILP; ++k) local_sum += sum[k];


            float warp_sum = warp_reduce_sum(local_sum);


            __shared__ float s_warp_sums[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_warp_sums[warp_id] = warp_sum;
            }}
            __syncthreads();

            if (warp_id == 0) {{
                float block_val = 0.0f;

                if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {{
                    block_val = s_warp_sums[lane_id];
                }}
                block_val = warp_reduce_sum(block_val);


                if (lane_id == 0) {{
                    atomicAdd(&out[n_idx], block_val);
                }}
            }}
        }}

        void bhattacharyya_sum_cuda(
            torch::Tensor x,
            torch::Tensor y,
            torch::Tensor out,
            float eps,
            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);

            bhattacharyya_split_kernel<<<blocks, threads>>>(
                x.data_ptr<float>(),
                y.data_ptr<float>(),
                out.data_ptr<float>(),
                eps,
                D_vec
            );
        }}
        """

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

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

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

        out = torch.zeros(N, device=x.device, dtype=torch.float32)

        self.op.bhattacharyya_sum_cuda(
            x,
            y,
            out,
            self.eps,
            N,
            D
        )

        return -torch.log(out + self.eps)