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>

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

__inline__ __device__ float blockReduceSum(float val) {
    static __shared__ float shared[32];
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;

    val = warpReduceSum(val);

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;

    if (wid == 0) val = warpReduceSum(val);

    return val;
}

__global__ void exp_normalize_log_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int rows,
    int cols,
    float eps
) {
    int bid = blockIdx.x; 
    int tid = threadIdx.x;

    if (bid >= rows) return;

    const float* row_in = input + bid * cols;
    float* row_out = output + bid * cols;

    // Pass 1: Compute Y = exp(X) and Sum of Squares of Y
    float sum_sq = 0.0f;
    for (int i = tid; i < cols; i += blockDim.x) {
        float y = expf(row_in[i]);
        sum_sq += y * y;
    }

    sum_sq = blockReduceSum(sum_sq);

    __shared__ float inv_norm;
    if (tid == 0) {
        inv_norm = rsqrtf(sum_sq + eps);
    }
    __syncthreads();

    float norm_factor = inv_norm;

    // Pass 2: Normalize (Y * inv_norm) and Log
    for (int i = tid; i < cols; i += blockDim.x) {
        float y = expf(row_in[i]);
        float z = y * norm_factor;
        row_out[i] = logf(z);
    }
}

torch::Tensor exp_normalize_log_cuda(torch::Tensor input) {
    auto output = torch::empty_like(input);

    int cols = input.size(input.dim() - 1);
    int rows = input.numel() / cols;

    int block_size = 256;
    while (block_size < cols && block_size < 1024) {
        block_size *= 2;
    }

    exp_normalize_log_kernel<<<rows, block_size>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        rows,
        cols,
        1e-12f
    );

    return output;
}
"""

cpp_source = """
torch::Tensor exp_normalize_log_cuda(torch::Tensor input);
"""

module = load_inline(
    name="exp_normalize_log",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["exp_normalize_log_cuda"],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.module = module

    def forward(self, x):
        return self.module.exp_normalize_log_cuda(x)