import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, applies Mish activation, adds a value, 
    applies Hardtanh activation, and scales the output.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.add_value = add_value
        self.scale = scale

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.nn.functional.mish(x) # Mish activation
        x = x + self.add_value
        x = torch.nn.functional.hardtanh(x, min_val=-1, max_val=1) # Hardtanh activation
        x = x * self.scale # Scaling
        return x

batch_size = 128
in_channels  = 64  
out_channels = 64  
height = width = 128  
kernel_size  = 3
stride       = 2  
padding      = 1
output_padding = 1
add_value = 0.5
scale = 2

def get_inputs():
    return [torch.rand(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale]

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

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>
torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
    std::vector<int64_t> stride, std::vector<int64_t> padding, std::vector<int64_t> output_padding,
    int64_t groups, std::vector<int64_t> dilation, double add_value, double scale);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#include <vector>

__global__ void act_kernel(float* __restrict__ x, long long n, float add_value, float scale) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    for (long long i = idx; i < n; i += stride) {
        float v = x[i];
        float mish;
        if (v > 20.0f) {
            mish = v;
        } else if (v < -20.0f) {
            mish = v * __expf(v);
        } else {
            float t = __expf(v);
            float a = t * (t + 2.0f);
            mish = v * (a / (a + 2.0f));
        }
        float y = mish + add_value;
        y = fminf(fmaxf(y, -1.0f), 1.0f);
        x[i] = y * scale;
    }
}

torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
    std::vector<int64_t> stride, std::vector<int64_t> padding, std::vector<int64_t> output_padding,
    int64_t groups, std::vector<int64_t> dilation, double add_value, double scale) {
    auto y = at::conv_transpose2d(input, weight, bias, stride, padding, output_padding, groups, dilation);
    const long long n = y.numel();
    if (n != 0) {
        const int threads = 256;
        int blocks = (int)((n + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        act_kernel<<<blocks, threads>>>(y.data_ptr<float>(), n, (float)add_value, (float)scale);
    }
    return y;
}
'''

_ext = load_inline(
    name='fused_convtranspose_mish_hardtanh_scale_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['fused_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.add_value = float(add_value)
        self.scale = float(scale)

    def forward(self, x):
        c = self.conv_transpose
        return _ext.fused_forward(x, c.weight, c.bias, list(c.stride), list(c.padding),
                                  list(c.output_padding), c.groups, list(c.dilation),
                                  self.add_value, self.scale)