import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a convolution, applies activation, and then applies Batch Normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels, eps=eps, momentum=momentum)

    def forward(self, x):
        x = self.conv(x)
        x = torch.multiply(torch.tanh(torch.nn.functional.softplus(x)), x)
        x = self.bn(x)
        return x

batch_size = 64
in_channels = 64
out_channels = 128
height, width = 128, 128
kernel_size = 3

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]

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

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor mish_bn_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                   torch::Tensor running_mean, torch::Tensor running_var,
                                   bool training, double momentum, double eps);

torch::Tensor mish_bn_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                              torch::Tensor running_mean, torch::Tensor running_var,
                              bool training, double momentum, double eps) {
    return mish_bn_forward_cuda(x, weight, bias, running_mean, running_var, training, momentum, eps);
}
'''

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

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#endif

static __device__ __forceinline__ float mish_f(float v) {
    float sp = (v > 20.0f) ? v : log1pf(expf(v));
    return v * tanhf(sp);
}

template<int TPB>
__global__ void mish_partials_kernel(const float* __restrict__ x, float* __restrict__ act,
                                     float* __restrict__ psum, float* __restrict__ psq,
                                     int N, int C, int HW, int parts) {
    int c = blockIdx.x;
    int p = blockIdx.y;
    int tid = threadIdx.x;
    int M = N * HW;
    int chunk = (M + parts - 1) / parts;
    int start = p * chunk;
    int end = start + chunk;
    if (end > M) end = M;

    float s = 0.0f, q = 0.0f;
    for (int k = start + tid; k < end; k += TPB) {
        int n = k / HW;
        int hw = k - n * HW;
        int idx = (n * C + c) * HW + hw;
        float z = mish_f(x[idx]);
        act[idx] = z;
        s += z;
        q = fmaf(z, z, q);
    }
    __shared__ float ss[TPB];
    __shared__ float qq[TPB];
    ss[tid] = s; qq[tid] = q;
    __syncthreads();
    for (int off = TPB/2; off > 0; off >>= 1) {
        if (tid < off) { ss[tid] += ss[tid + off]; qq[tid] += qq[tid + off]; }
        __syncthreads();
    }
    if (tid == 0) {
        psum[c * parts + p] = ss[0];
        psq[c * parts + p] = qq[0];
    }
}

template<int TPB>
__global__ void stats_kernel(const float* __restrict__ psum, const float* __restrict__ psq,
                             const float* __restrict__ weight, const float* __restrict__ bias,
                             float* __restrict__ running_mean, float* __restrict__ running_var,
                             float* __restrict__ scale, float* __restrict__ shift,
                             int C, int M, int parts, float momentum, float eps) {
    int c = blockIdx.x;
    int tid = threadIdx.x;
    float s = 0.0f, q = 0.0f;
    for (int p = tid; p < parts; p += TPB) {
        s += psum[c * parts + p];
        q += psq[c * parts + p];
    }
    __shared__ float ss[TPB];
    __shared__ float qq[TPB];
    ss[tid] = s; qq[tid] = q;
    __syncthreads();
    for (int off = TPB/2; off > 0; off >>= 1) {
        if (tid < off) { ss[tid] += ss[tid + off]; qq[tid] += qq[tid + off]; }
        __syncthreads();
    }
    if (tid == 0) {
        float mean = ss[0] / (float)M;
        float var = qq[0] / (float)M - mean * mean;
        var = var < 0.0f ? 0.0f : var;
        float inv = rsqrtf(var + eps);
        float g = weight[c];
        float b = bias[c];
        scale[c] = g * inv;
        shift[c] = b - mean * g * inv;
        // PyTorch BatchNorm updates running_var with the unbiased variance estimate.
        float unbiased = var * ((float)M / (float)(M - 1));
        running_mean[c] = running_mean[c] * (1.0f - momentum) + momentum * mean;
        running_var[c] = running_var[c] * (1.0f - momentum) + momentum * unbiased;
    }
}

__global__ void norm_kernel(const float* __restrict__ act, float* __restrict__ out,
                            const float* __restrict__ scale, const float* __restrict__ shift,
                            int total, int C, int HW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < total; i += stride) {
        int c = (i / HW) % C;
        out[i] = fmaf(act[i], scale[c], shift[c]);
    }
}

__global__ void mish_bn_eval_kernel(const float* __restrict__ x, float* __restrict__ out,
                                    const float* __restrict__ weight, const float* __restrict__ bias,
                                    const float* __restrict__ running_mean, const float* __restrict__ running_var,
                                    int total, int C, int HW, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < total; i += stride) {
        int c = (i / HW) % C;
        float inv = rsqrtf(running_var[c] + eps);
        float z = mish_f(x[i]);
        out[i] = (z - running_mean[c]) * inv * weight[c] + bias[c];
    }
}

torch::Tensor mish_bn_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                   torch::Tensor running_mean, torch::Tensor running_var,
                                   bool training, double momentum, double eps) {
    CHECK_CUDA(x); CHECK_CUDA(weight); CHECK_CUDA(bias); CHECK_CUDA(running_mean); CHECK_CUDA(running_var);
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    auto out = torch::empty_like(x);
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    const int HW = H * W;
    const int total = N * C * HW;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (!training) {
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 4096) blocks = 4096;
        mish_bn_eval_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), out.data_ptr<float>(),
            weight.data_ptr<float>(), bias.data_ptr<float>(), running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            total, C, HW, (float)eps);
        return out;
    }
    const int parts = 16;
    auto act = torch::empty_like(x);
    auto opts = x.options();
    auto psum = torch::empty({C, parts}, opts);
    auto psq = torch::empty({C, parts}, opts);
    auto scale = torch::empty({C}, opts);
    auto shift = torch::empty({C}, opts);
    dim3 grid1(C, parts);
    mish_partials_kernel<256><<<grid1, 256, 0, stream>>>(x.data_ptr<float>(), act.data_ptr<float>(), psum.data_ptr<float>(), psq.data_ptr<float>(), N, C, HW, parts);
    stats_kernel<32><<<C, 32, 0, stream>>>(psum.data_ptr<float>(), psq.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(), C, N*HW, parts, (float)momentum, (float)eps);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    norm_kernel<<<blocks, threads, 0, stream>>>(act.data_ptr<float>(), out.data_ptr<float>(), scale.data_ptr<float>(), shift.data_ptr<float>(), total, C, HW);
    return out;
}
'''

_ext = load_inline(
    name='mish_bn_fused_ext_v1',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['mish_bn_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.bn = nn.BatchNorm2d(out_channels, eps=eps, momentum=momentum)

    def forward(self, x):
        y = F.conv2d(x, self.conv.weight, self.conv.bias, self.conv.stride,
                     self.conv.padding, self.conv.dilation, self.conv.groups)
        return _ext.mish_bn_forward(y.contiguous(), self.bn.weight, self.bn.bias,
                                    self.bn.running_mean, self.bn.running_var,
                                    self.bn.training, float(self.bn.momentum), float(self.bn.eps))