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

N, C, H, W = 16, 16, 64, 64
BETA = 1.0
THRESHOLD = 20.0
BLOCK_SIZE = 256


class ModelNew(nn.Module):

    def __init__(self, beta=BETA, threshold=THRESHOLD):
        super().__init__()
        self.beta = beta
        self.threshold = threshold
        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        total_elements = N * C * H * W
        if total_elements % 4 != 0:
            raise ValueError("Total elements must be divisible by 4 for float4 vectorization.")

        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor softplus_cuda_vec(torch::Tensor input, float beta, float threshold);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <math.h>

        #define BLOCK_SIZE {self.block_size}

        __device__ __forceinline__ float softplus_elem(float x, float beta, float threshold, float inv_beta) {{
            float scaled_x = x * beta;
            return (scaled_x > threshold) ? x : inv_beta * __logf(1.0f + __expf(scaled_x));
        }}

        __global__ void softplus_kernel_vec(
            const float4* __restrict__ input,
            float4* __restrict__ output,
            int total_vec_elements,
            float beta,
            float threshold,
            float inv_beta
        ) {{
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = gridDim.x * blockDim.x;

            for (int i = idx; i < total_vec_elements; i += stride) {{
                float4 in_vec = input[i];
                float4 out_vec;

                out_vec.x = softplus_elem(in_vec.x, beta, threshold, inv_beta);
                out_vec.y = softplus_elem(in_vec.y, beta, threshold, inv_beta);
                out_vec.z = softplus_elem(in_vec.z, beta, threshold, inv_beta);
                out_vec.w = softplus_elem(in_vec.w, beta, threshold, inv_beta);

                output[i] = out_vec;
            }}
        }}

        torch::Tensor softplus_cuda_vec(torch::Tensor input, float beta, float threshold) {{
            TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
            TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");

            int total_elements = input.numel();
            TORCH_CHECK(total_elements % 4 == 0, "Total elements must be divisible by 4 for float4 vectorization.");

            input = input.contiguous();
            auto output = torch::empty_like(input);

            int total_vec_elements = total_elements / 4;
            float inv_beta = 1.0f / beta;

            int blocks = std::min((total_vec_elements + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);

            softplus_kernel_vec<<<blocks, BLOCK_SIZE>>>(
                reinterpret_cast<const float4*>(input.data_ptr<float>()),
                reinterpret_cast<float4*>(output.data_ptr<float>()),
                total_vec_elements,
                beta,
                threshold,
                inv_beta
            );

            return output;
        }}
        """

        nvcc_flags = ['-O3', '--use_fast_math']

        self.op = load_inline(
            name='softplus_opt_vec',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['softplus_cuda_vec'],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if not input.is_cuda: input = input.cuda()
        input_cont = input.contiguous()
        return self.op.softplus_cuda_vec(input_cont, self.beta, self.threshold)