import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed convolution, followed by max pooling, hardtanh activation, mean operation, and tanh activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, maxpool_kernel_size, maxpool_stride, hardtanh_min, hardtanh_max):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.maxpool = nn.MaxPool2d(kernel_size=maxpool_kernel_size, stride=maxpool_stride)
        self.hardtanh = nn.Hardtanh(min_val=hardtanh_min, max_val=hardtanh_max)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.maxpool(x)
        x = self.hardtanh(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)
        x = torch.tanh(x)
        return x

batch_size = 128
in_channels  = 64  
out_channels = 64  
height = width = 256  
kernel_size  = 3
stride = 1
padding = 1
maxpool_kernel_size = 2
maxpool_stride = 2
hardtanh_min = -1
hardtanh_max = 1

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, maxpool_kernel_size, maxpool_stride, hardtanh_min, hardtanh_max]

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor pool_clamp_mean_tanh(torch::Tensor x, double clamp_min, double clamp_max);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>

__inline__ __device__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

__global__ void pool_clamp_mean_tanh_kernel(const float* __restrict__ x, float* __restrict__ out,
                                            int H, int W, float clamp_min, float clamp_max) {
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    int pooledW = W >> 1;
    int total = (H >> 1) * pooledW;
    const float* base = x + ((long long)idx) * H * W;
    float sum = 0.0f;
    for (int p = tid; p < total; p += blockDim.x) {
        int py = p / pooledW;
        int px = p - py * pooledW;
        const float* ptr = base + (long long)(py << 1) * W + (px << 1);
        float2 a = *reinterpret_cast<const float2*>(ptr);
        float2 b = *reinterpret_cast<const float2*>(ptr + W);
        float m = fmaxf(fmaxf(a.x, a.y), fmaxf(b.x, b.y));
        m = fminf(fmaxf(m, clamp_min), clamp_max);
        sum += m;
    }
    sum = warp_sum(sum);
    __shared__ float smem[8];
    int lane = tid & 31;
    int warp = tid >> 5;
    if (lane == 0) smem[warp] = sum;
    __syncthreads();
    if (warp == 0) {
        sum = (tid < (blockDim.x >> 5)) ? smem[lane] : 0.0f;
        sum = warp_sum(sum);
        if (tid == 0) out[idx] = tanhf(sum / (float)total);
    }
}

torch::Tensor pool_clamp_mean_tanh(torch::Tensor x, double clamp_min, double clamp_max) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    x = x.contiguous();
    const int64_t N = x.size(0), C = x.size(1);
    const int H = (int)x.size(2), W = (int)x.size(3);
    auto out = torch::empty({N, C, 1, 1}, x.options());
    int NC = (int)(N*C);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    pool_clamp_mean_tanh_kernel<<<NC, 256, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(), H, W, (float)clamp_min, (float)clamp_max);
    return out;
}
'''

_ext = load_inline(name='fused_post_ext_submit', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['pool_clamp_mean_tanh'], verbose=False, extra_cuda_cflags=['-O3', '--use_fast_math'])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, maxpool_kernel_size, maxpool_stride, hardtanh_min, hardtanh_max):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.hardtanh_min = float(hardtanh_min)
        self.hardtanh_max = float(hardtanh_max)
    def forward(self, x):
        y = self.conv_transpose(x)
        return _ext.pool_clamp_mean_tanh(y, self.hardtanh_min, self.hardtanh_max)