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

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

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

__global__ void total_correlation_elu_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ target, 
    float* __restrict__ y, 
    int batch_size, 
    int width) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    if (row >= batch_size) return;

    extern __shared__ float s_mem[];
    float* s_x = s_mem;
    float* s_t = s_mem + width;

    for (int i = tid; i < width; i += blockDim.x) {
        s_x[i] = x[row * width + i];
        s_t[i] = target[i];
    }
    __syncthreads();

    float sum_x = 0.0f;
    float sum_t = 0.0f;

    for (int i = tid; i < width; i += blockDim.x) {
        sum_x += s_x[i];
        sum_t += s_t[i];
    }

    sum_x = warp_reduce(sum_x);
    sum_t = warp_reduce(sum_t);

    static __shared__ float shared_sum_x[32];
    static __shared__ float shared_sum_t[32];

    int lane = tid % 32;
    int wid = tid / 32;

    if (lane == 0) {
        shared_sum_x[wid] = sum_x;
        shared_sum_t[wid] = sum_t;
    }
    __syncthreads();

    sum_x = (tid < blockDim.x / 32) ? shared_sum_x[lane] : 0.0f;
    sum_t = (tid < blockDim.x / 32) ? shared_sum_t[lane] : 0.0f;

    if (wid == 0) {
        sum_x = warp_reduce(sum_x);
        sum_t = warp_reduce(sum_t);
    }
    __syncthreads();

    // Broadcast means
    if (tid == 0) {
        shared_sum_x[0] = sum_x / width; 
        shared_sum_t[0] = sum_t / width;
    }
    __syncthreads();

    float mean_x = shared_sum_x[0];
    float mean_t = shared_sum_t[0];

    float sum_cov = 0.0f;
    float sum_var_x = 0.0f;
    float sum_var_t = 0.0f;

    for (int i = tid; i < width; i += blockDim.x) {
        float dx = s_x[i] - mean_x;
        float dt = s_t[i] - mean_t;
        sum_cov += dx * dt;
        sum_var_x += dx * dx;
        sum_var_t += dt * dt;
    }

    sum_cov = warp_reduce(sum_cov);
    sum_var_x = warp_reduce(sum_var_x);
    sum_var_t = warp_reduce(sum_var_t);

    if (lane == 0) {
        shared_sum_x[wid] = sum_cov;
        shared_sum_t[wid] = sum_var_x;
        // reuse shared memory slot
        // store sum_var_t in s_x buffer temporarily or add another shared array?
        // simple: use float* s_var_t = (float*)&shared_sum_x[0] + 32? No.
        // Let's assume blockDim <= 1024, max 32 warps.
    }
    __syncthreads();

    // Need 3 values reduced. 
    // Just use atomic or simpler serial reduction for the last 32 values if precision allows,
    // or allocate 3 shared arrays.
    // Let's use 3 dedicated shared arrays for reduction to avoid conflict.
    static __shared__ float s_r_cov[32];
    static __shared__ float s_r_vx[32];
    static __shared__ float s_r_vt[32];

    if (lane == 0) {
        s_r_cov[wid] = sum_cov;
        s_r_vx[wid] = sum_var_x;
        s_r_vt[wid] = sum_var_t;
    }
    __syncthreads();

    sum_cov = (tid < blockDim.x / 32) ? s_r_cov[lane] : 0.0f;
    sum_var_x = (tid < blockDim.x / 32) ? s_r_vx[lane] : 0.0f;
    sum_var_t = (tid < blockDim.x / 32) ? s_r_vt[lane] : 0.0f;

    if (wid == 0) {
        sum_cov = warp_reduce(sum_cov);
        sum_var_x = warp_reduce(sum_var_x);
        sum_var_t = warp_reduce(sum_var_t);
    }

    if (tid == 0) {
        float corr = sum_cov / (sqrtf(sum_var_x) * sqrtf(sum_var_t) + 1e-6f);
        // ELU: x > 0 ? x : exp(x) - 1
        float val = corr;
        if (val > 0.0f) {
            y[row] = val;
        } else {
            y[row] = expf(val) - 1.0f;
        }
    }
}

torch::Tensor launch_total_correlation_elu(torch::Tensor x, torch::Tensor target) {
    auto batch_size = x.size(0);
    auto width = x.size(1);
    auto y = torch::empty({batch_size}, x.options());

    const int threads = 256;
    const int blocks = batch_size;
    int shared_mem = 2 * width * sizeof(float);

    total_correlation_elu_kernel<<<blocks, threads, shared_mem>>>(
        x.data_ptr<float>(), 
        target.data_ptr<float>(), 
        y.data_ptr<float>(), 
        batch_size, 
        width
    );
    return y;
}
"""

cpp_source = """
torch::Tensor launch_total_correlation_elu(torch::Tensor x, torch::Tensor target);
"""

total_correlation_elu_module = load_inline(
    name='total_correlation_elu_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_total_correlation_elu'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, target):
        super(ModelNew, self).__init__()
        self.target = nn.Parameter(target)
        self.op = total_correlation_elu_module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op.launch_total_correlation_elu(x.contiguous(), self.target.contiguous())