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__ void warpReduceMinMax(float &min_val, float &max_val) {
    for (int offset = 16; offset > 0; offset /= 2) {
        min_val = fminf(min_val, __shfl_down_sync(0xffffffff, min_val, offset));
        max_val = fmaxf(max_val, __shfl_down_sync(0xffffffff, max_val, offset));
    }
}

__global__ void minmax_scale_shift_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    float scale,
    float shift,
    float eps,
    int rows,
    int cols
) {
    __shared__ float s_min;
    __shared__ float s_max;

    int bid = blockIdx.x;
    int tid = threadIdx.x;

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

    float local_min = INFINITY;
    float local_max = -INFINITY;

    for (int i = tid; i < cols; i += blockDim.x) {
        float val = row_in[i];
        local_min = fminf(local_min, val);
        local_max = fmaxf(local_max, val);
    }

    // Warp Reduction
    warpReduceMinMax(local_min, local_max);

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

    __shared__ float shared_min[32];
    __shared__ float shared_max[32];

    if (lane == 0) {
        shared_min[wid] = local_min;
        shared_max[wid] = local_max;
    }
    __syncthreads();

    if (wid == 0) {
        local_min = (tid < (blockDim.x / 32)) ? shared_min[tid] : INFINITY;
        local_max = (tid < (blockDim.x / 32)) ? shared_max[tid] : -INFINITY;
        warpReduceMinMax(local_min, local_max);

        if (tid == 0) {
            s_min = local_min;
            s_max = local_max;
        }
    }
    __syncthreads();

    float min_x = s_min;
    float max_x = s_max;
    float range_x_inv = rsqrtf(max_x - min_x + eps) * rsqrtf(max_x - min_x + eps); // 1 / (range + eps)

    // Pass 2: Normalize, Scale, Shift
    for (int i = tid; i < cols; i += blockDim.x) {
        float val = row_in[i];

        float norm = (val - min_x) * range_x_inv;

        // Scale and Shift: fma(norm, scale, shift)
        row_out[i] = fmaf(norm, scale, shift);
    }
}

torch::Tensor minmax_scale_shift_cuda(torch::Tensor input, float scale, float shift, float eps) {
    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;
    }

    minmax_scale_shift_kernel<<<rows, block_size>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        scale,
        shift,
        eps,
        rows,
        cols
    );

    return output;
}
"""

cpp_source = """
torch::Tensor minmax_scale_shift_cuda(torch::Tensor input, float scale, float shift, float eps);
"""

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


class ModelNew(nn.Module):
    def __init__(self, scale, shift, eps=1e-5):
        super(ModelNew, self).__init__()
        self.scale = scale
        self.shift = shift
        self.eps = eps
        self.module = module

    def forward(self, x):
        return self.module.minmax_scale_shift_cuda(x, self.scale, self.shift, self.eps)