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


class ModelNew(nn.Module):
    def __init__(self, bins=32):
        super().__init__()
        self.bins = bins
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor vi_cuda(torch::Tensor x, torch::Tensor y, int bins);
        """

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

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

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

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

        __global__ void vi_kernel(
            const float* __restrict__ x,
            const float* __restrict__ y,
            float* __restrict__ output,
            int feature_dim,
            int batch_size) 
        {
            const int BINS = 32;

            __shared__ int s_hist[BINS][BINS];
            __shared__ float s_px[BINS];
            __shared__ float s_py[BINS];

            // Min/Max reduction buffer
            static __shared__ float s_reduce_buf[8][4]; 

            int bid = blockIdx.x;
            if (bid >= batch_size) return;

            int tid = threadIdx.x;
            int lane = tid % 32;
            int wid = tid / 32;
            for (int i = tid; i < BINS * BINS; i += blockDim.x) {
                reinterpret_cast<int*>(s_hist)[i] = 0;
            }
            if (tid < BINS) {
                s_px[tid] = 0.0f;
                s_py[tid] = 0.0f;
            }
            __syncthreads();

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

            float l_min_x = FLT_MAX, l_max_x = -FLT_MAX;
            float l_min_y = FLT_MAX, l_max_y = -FLT_MAX;

            int vec_loops = feature_dim / 4;
            int vec_remainder = feature_dim % 4;
            const float4* x_vec = reinterpret_cast<const float4*>(x_row);
            const float4* y_vec = reinterpret_cast<const float4*>(y_row);

            for (int i = tid; i < vec_loops; i += blockDim.x) {
                float4 vx = x_vec[i];
                float4 vy = y_vec[i];

                l_min_x = fminf(l_min_x, fminf(vx.x, fminf(vx.y, fminf(vx.z, vx.w))));
                l_max_x = fmaxf(l_max_x, fmaxf(vx.x, fmaxf(vx.y, fmaxf(vx.z, vx.w))));
                l_min_y = fminf(l_min_y, fminf(vy.x, fminf(vy.y, fminf(vy.z, vy.w))));
                l_max_y = fmaxf(l_max_y, fmaxf(vy.x, fmaxf(vy.y, fmaxf(vy.z, vy.w))));
            }
            l_min_x = warp_reduce_min(l_min_x);
            l_max_x = warp_reduce_max(l_max_x);
            l_min_y = warp_reduce_min(l_min_y);
            l_max_y = warp_reduce_max(l_max_y);

            if (lane == 0) {
                s_reduce_buf[wid][0] = l_min_x;
                s_reduce_buf[wid][1] = l_max_x;
                s_reduce_buf[wid][2] = l_min_y;
                s_reduce_buf[wid][3] = l_max_y;
            }
            __syncthreads();

            // Block Reduce (by thread 0)
            float min_x, max_x, min_y, max_y;
            if (tid == 0) {
                min_x = FLT_MAX; max_x = -FLT_MAX;
                min_y = FLT_MAX; max_y = -FLT_MAX;
                for (int i = 0; i < blockDim.x / 32; ++i) {
                    min_x = fminf(min_x, s_reduce_buf[i][0]);
                    max_x = fmaxf(max_x, s_reduce_buf[i][1]);
                    min_y = fminf(min_y, s_reduce_buf[i][2]);
                    max_y = fmaxf(max_y, s_reduce_buf[i][3]);
                }
                // Store back to shared for broadcast
                s_reduce_buf[0][0] = min_x; s_reduce_buf[0][1] = max_x;
                s_reduce_buf[0][2] = min_y; s_reduce_buf[0][3] = max_y;
            }
            __syncthreads();

            min_x = s_reduce_buf[0][0]; max_x = s_reduce_buf[0][1];
            min_y = s_reduce_buf[0][2]; max_y = s_reduce_buf[0][3];

            float range_x = max_x - min_x + 1e-6f;
            float range_y = max_y - min_y + 1e-6f;

            for (int i = tid; i < vec_loops; i += blockDim.x) {
                float4 vx = x_vec[i];
                float4 vy = y_vec[i];

                #pragma unroll
                for (int k = 0; k < 4; ++k) {
                    float val_x = (k==0?vx.x:k==1?vx.y:k==2?vx.z:vx.w);
                    float val_y = (k==0?vy.x:k==1?vy.y:k==2?vy.z:vy.w);

                    int bx = (int)((val_x - min_x) / range_x * BINS);
                    int by = (int)((val_y - min_y) / range_y * BINS);
                    bx = min(max(bx, 0), BINS - 1);
                    by = min(max(by, 0), BINS - 1);

                    atomicAdd(&s_hist[bx][by], 1);
                }
            }

            if (tid == 0 && vec_remainder > 0) {
                for (int i = 0; i < vec_remainder; ++i) {
                    int idx = vec_loops * 4 + i;
                    int bx = (int)((x_row[idx] - min_x) / range_x * BINS);
                    int by = (int)((y_row[idx] - min_y) / range_y * BINS);
                    bx = min(max(bx, 0), BINS - 1);
                    by = min(max(by, 0), BINS - 1);
                    atomicAdd(&s_hist[bx][by], 1);
                }
            }
            __syncthreads();

            
            if (tid < BINS) {
                int px = 0;
                int py = 0;
                for (int i = 0; i < BINS; ++i) {
                    px += s_hist[tid][i];
                    py += s_hist[i][tid];
                }
                s_px[tid] = (float)px / feature_dim;
                s_py[tid] = (float)py / feature_dim;
            }
            __syncthreads();

            

            double sum_mi = 0.0;
            double sum_hx = 0.0;
            double sum_hy = 0.0;
            const float eps = 1e-12f;

            // Compute MI (over 32x32 grid)
            for (int i = tid; i < BINS * BINS; i += blockDim.x) {
                int r = i / BINS;
                int c = i % BINS;

                int count = s_hist[r][c];
                if (count > 0) {
                    float p_xy = (float)count / feature_dim;
                    float px = s_px[r];
                    float py = s_py[c];

                    sum_mi += (double)(p_xy * logf((p_xy + eps) / (px * py + eps)));
                }
            }

            // Compute H(X) and H(Y) (over 32 bins)
            // Each thread takes one bin if tid < 32
            if (tid < BINS) {
                float px = s_px[tid];
                float py = s_py[tid];
                if (px > 0) sum_hx += (double)(-px * logf(px + eps));
                if (py > 0) sum_hy += (double)(-py * logf(py + eps));
            }

            // Reductions
            sum_mi = warp_reduce_sum(sum_mi);
            sum_hx = warp_reduce_sum(sum_hx);
            sum_hy = warp_reduce_sum(sum_hy);

            static __shared__ double s_vals[3][32]; // [0]:MI, [1]:HX, [2]:HY

            if (lane == 0) {
                s_vals[0][wid] = sum_mi;
                s_vals[1][wid] = sum_hx;
                s_vals[2][wid] = sum_hy;
            }
            __syncthreads();

            if (tid == 0) {
                double total_mi = 0.0, total_hx = 0.0, total_hy = 0.0;
                int n_warps = blockDim.x / 32;

                for (int i = 0; i < n_warps; ++i) {
                    total_mi += s_vals[0][i];
                    total_hx += s_vals[1][i];
                    total_hy += s_vals[2][i];
                }

                // VI = H(X) + H(Y) - 2 * MI
                output[bid] = (float)(total_hx + total_hy - 2.0 * total_mi);
            }
        }

        torch::Tensor vi_cuda(torch::Tensor x, torch::Tensor y, int bins) {
            auto x_c = x.contiguous();
            auto y_c = y.contiguous();

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

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

            int threads = 256; 
            int blocks = batch_size;

            if (bins != 32) {
                 // Fallback logic could be added here
            }

            vi_kernel<<<blocks, threads>>>(
                x_c.data_ptr<float>(),
                y_c.data_ptr<float>(),
                output.data_ptr<float>(),
                feature_dim,
                batch_size
            );

            return output;
        }
        """

        self.op = load_inline(
            name="vi_opt_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["vi_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=False
        )

    def forward(self, x, y):
        return self.op.vi_cuda(x, y, self.bins)