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

N, C, H, W = 16, 16, 64, 64
LAMBDA = 0.5
BLOCK_SIZE = 256


class ModelNew(nn.Module):

    def __init__(self, lambd=LAMBDA):
        super().__init__()
        self.lambd = lambd
        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.")

        lambd_str = f"{self.lambd:.8f}f"

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

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

        #define BLOCK_SIZE {self.block_size}

        /*
         * Softshrink Kernel with float4 Vectorization
         * Implements: y = sign(x) * max(0, |x| - lambda)
         */
        __global__ void softshrink_kernel_vec(
            const float* __restrict__ input,
            float* __restrict__ output,
            int total_vec_elements,
            float lambd // lambda value passed as argument
        ) {{
            int idx = blockIdx.x * blockDim.x + threadIdx.x;

            for (; idx < total_vec_elements; idx += gridDim.x * blockDim.x) {{

                // 1. Load float4 vector
                const float4* in_ptr = reinterpret_cast<const float4*>(input);
                float4 in_vec = in_ptr[idx];

                // 2. Apply Softshrink element-wise
                float4 out_vec;

                // CUDA Intrinsics: fabsf (abs), copysignf (sign)
                // Note: fabsf(x) - lambd is the core logic.
                // We use copysignf to restore the sign, which is robust near zero.

                // --- X component ---
                float magnitude_x = fabsf(in_vec.x);
                float thresholded_x = fmaxf(0.0f, magnitude_x - lambd); // max(0, |x| - lambda)
                out_vec.x = copysignf(thresholded_x, in_vec.x); // sign(x) * thresholded_x

                // --- Y component ---
                float magnitude_y = fabsf(in_vec.y);
                float thresholded_y = fmaxf(0.0f, magnitude_y - lambd);
                out_vec.y = copysignf(thresholded_y, in_vec.y);

                // --- Z component ---
                float magnitude_z = fabsf(in_vec.z);
                float thresholded_z = fmaxf(0.0f, magnitude_z - lambd);
                out_vec.z = copysignf(thresholded_z, in_vec.z);

                // --- W component ---
                float magnitude_w = fabsf(in_vec.w);
                float thresholded_w = fmaxf(0.0f, magnitude_w - lambd);
                out_vec.w = copysignf(thresholded_w, in_vec.w);


                // 3. Store float4 vector
                float4* out_ptr = reinterpret_cast<float4*>(output);
                out_ptr[idx] = out_vec;
            }}
        }}


        torch::Tensor softshrink_cuda_vec(torch::Tensor input, float lambd) {{

            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;

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

            softshrink_kernel_vec<<<blocks, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                total_vec_elements,
                lambd // Passing lambda
            );

            return output;
        }}
        """

        nvcc_flags = ['-O3']

        self.op = load_inline(
            name='softshrink_opt_vec',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['softshrink_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.softshrink_cuda_vec(input_cont, self.lambd)