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

cpp_source = """
#include <torch/extension.h>

torch::Tensor fake_quant_cuda_forward(
    const torch::Tensor& input,
    const torch::Tensor& scale,
    const torch::Tensor& zero_point,
    int axis,
    int64_t quant_min,
    int64_t quant_max);
"""

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

struct __align__(16) Float4 {
    float x, y, z, w;
};

// Core Math Logic
// Matches PyTorch implementation: x * inv_scale
__device__ __forceinline__ float fake_quant_op(
    float x, float s, int zp_val, float qmin, float qmax) 
{
    float zp = (float)zp_val;
    
    // 1. Scale (Multiplication by inverse, crucial for precision alignment)
    // Avoid division x / s
    float inv_s = 1.0f / s;
    float val = x * inv_s;
    
    // 2. Round (std::nearbyint -> rintf)
    val = rintf(val);
    
    // 3. Add ZeroPoint
    val = val + zp;
    
    // 4. Clamp
    val = fminf(fmaxf(val, qmin), qmax);
    
    // 5. Dequantize
    return (val - zp) * s;
}

__global__ void fake_quant_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ scale,
    const int* __restrict__ zero_point,
    const int n_elements,
    const int channel_count,
    const int inner_stride, 
    const float qmin,
    const float qmax) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_stride = blockDim.x * gridDim.x;

    int vec_loops = n_elements / 4;
    const Float4* vec_input = reinterpret_cast<const Float4*>(input);
    Float4* vec_output = reinterpret_cast<Float4*>(output);

    for (int i = idx; i < vec_loops; i += total_stride) {
        Float4 in_v = vec_input[i];
        Float4 out_v;
        
        int idx0 = i * 4;
        
        // Calculate channel for each element (robust against unaligned strides)
        int c0 = (idx0 / inner_stride) % channel_count;
        out_v.x = fake_quant_op(in_v.x, scale[c0], zero_point[c0], qmin, qmax);
        
        int c1 = ((idx0 + 1) / inner_stride) % channel_count;
        out_v.y = fake_quant_op(in_v.y, scale[c1], zero_point[c1], qmin, qmax);
        
        int c2 = ((idx0 + 2) / inner_stride) % channel_count;
        out_v.z = fake_quant_op(in_v.z, scale[c2], zero_point[c2], qmin, qmax);
        
        int c3 = ((idx0 + 3) / inner_stride) % channel_count;
        out_v.w = fake_quant_op(in_v.w, scale[c3], zero_point[c3], qmin, qmax);

        vec_output[i] = out_v;
    }

    int tail_start = vec_loops * 4;
    for (int i = tail_start + idx; i < n_elements; i += total_stride) {
        int c = (i / inner_stride) % channel_count;
        output[i] = fake_quant_op(input[i], scale[c], zero_point[c], qmin, qmax);
    }
}

torch::Tensor fake_quant_cuda_forward(
    const torch::Tensor& input,
    const torch::Tensor& scale,
    const torch::Tensor& zero_point,
    int axis,
    int64_t quant_min,
    int64_t quant_max) 
{
    TORCH_CHECK(input.is_cuda(), "Input must be CUDA");
    TORCH_CHECK(scale.is_cuda(), "Scale must be CUDA");
    TORCH_CHECK(zero_point.is_cuda(), "ZeroPoint must be CUDA");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");

    int ndim = input.dim();
    if (axis < 0) axis += ndim;
    
    int channel_count = input.size(axis);
    int inner_stride = 1;
    for (int i = axis + 1; i < ndim; ++i) {
        inner_stride *= input.size(i);
    }

    int n_elements = input.numel();
    auto output = torch::empty_like(input);

    const int block_size = 256;
    int num_vectors = n_elements / 4;
    int grid_size = (num_vectors + block_size - 1) / block_size;
    if (grid_size > 65535) grid_size = 65535;

    fake_quant_kernel<<<grid_size, block_size>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        scale.data_ptr<float>(),
        zero_point.data_ptr<int>(),
        n_elements,
        channel_count,
        inner_stride,
        (float)quant_min,
        (float)quant_max
    );

    return output;
}
"""

fq_op = load_inline(
    name='fq_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fake_quant_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3'] 
)

class ModelNew(nn.Module):
    def __init__(self, axis, q_min, q_max):
        super(ModelNew, self).__init__()
        self.axis = axis
        self.q_min = q_min
        self.q_max = q_max

    def forward(self, x, scale, zero_point):
        return fq_op.fake_quant_cuda_forward(
            x.contiguous(), scale, zero_point, self.axis, self.q_min, self.q_max
        )