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

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

torch::Tensor mmrelu_cuda_forward(const torch::Tensor& input, float t, float sqrt_3t, float coeff);
"""

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

#define BLOCK_SIZE 256

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

// MMReLU Logic
__device__ __forceinline__ float compute_mmrelu(float x, float t, float neg_sqrt_3t, float coeff) {
    if (x >= 0.0f) {
        return x;
    } else if (x > neg_sqrt_3t) { // -sqrt(3t) < x < 0
        return coeff * x * x + x;
    } else { // x <= -sqrt(3t)
        return t / x;
    }
}

__global__ void mmrelu_kernel(
    float* __restrict__ output,
    const float* __restrict__ input,
    const int n,
    const float t,
    const float neg_sqrt_3t,
    const float coeff) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n / 4;
    
    int i = idx;
    const int stride = blockDim.x * gridDim.x;

    for (; i < vec_n; i += stride) {
        Float4 in_vec = reinterpret_cast<const Float4*>(input)[i];
        Float4 out_vec;

        out_vec.x = compute_mmrelu(in_vec.x, t, neg_sqrt_3t, coeff);
        out_vec.y = compute_mmrelu(in_vec.y, t, neg_sqrt_3t, coeff);
        out_vec.z = compute_mmrelu(in_vec.z, t, neg_sqrt_3t, coeff);
        out_vec.w = compute_mmrelu(in_vec.w, t, neg_sqrt_3t, coeff);

        reinterpret_cast<Float4*>(output)[i] = out_vec;
    }

    int start_scalar = vec_n * 4;
    int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_threads = gridDim.x * gridDim.x;
    
    int current_idx = start_scalar + global_tid;
    while (current_idx < n) {
        output[current_idx] = compute_mmrelu(input[current_idx], t, neg_sqrt_3t, coeff);
        current_idx += total_threads;
    }
}

torch::Tensor mmrelu_cuda_forward(const torch::Tensor& input, float t, float sqrt_3t, float coeff) {
    TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    
    const int n = input.numel();
    auto output = torch::empty_like(input);

    const int vec_n = n / 4;
    const int grid_size = (vec_n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    
    int final_grid = (grid_size < 1) ? 1 : grid_size;
    if (final_grid > 65535) final_grid = 65535;
    
    float neg_sqrt_3t = -sqrt_3t;

    mmrelu_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        input.data_ptr<float>(),
        n,
        t,
        neg_sqrt_3t,
        coeff
    );

    return output;
}
"""

mmrelu_op_module = load_inline(
    name='mmrelu_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['mmrelu_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, t=1.0):
        super(ModelNew, self).__init__()
        self.t = t
        # Pre-compute constants
        self.sqrt_3t = math.sqrt(3 * t)
        self.coeff = (2 * self.sqrt_3t) / (9 * t)
        self.op = mmrelu_op_module

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.mmrelu_cuda_forward(
            input_tensor.contiguous(), self.t, self.sqrt_3t, self.coeff
        )