import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        """
        MBConv block implementation.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param kernel_size: Kernel size for the depthwise convolution.
        :param stride: Stride for the depthwise convolution.
        :param expand_ratio: Expansion ratio for the intermediate channels.
        """
        super(Model, self).__init__()
        
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = in_channels * expand_ratio
        
        if expand_ratio != 1:
            self.expand_conv = nn.Sequential(
                nn.Conv2d(in_channels, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True)
            )
        
        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True)
        )
        
        self.project_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels)
        )
    
    def forward(self, x):
        """
        Forward pass of the MBConv block.

        :param x: The input tensor, shape (batch_size, in_channels, H, W)
        :return: The output tensor, shape (batch_size, out_channels, H', W')
        """
        identity = x
        
        if hasattr(self, 'expand_conv'):
            x = self.expand_conv(x)
        
        x = self.depthwise_conv(x)
        x = self.project_conv(x)
        
        if self.use_residual:
            x += identity
        
        return x

# Test code
batch_size = 10
in_channels = 112
out_channels = 192
kernel_size = 5
stride = 2
expand_ratio = 6

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

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

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

cpp_source = r'''
#include <torch/extension.h>
torch::Tensor dw_bn_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, double eps);
'''

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void dw_conv_reduce(const float* __restrict__ x, const float* __restrict__ w, float* __restrict__ y, float* __restrict__ psum, float* __restrict__ psq, int N,int C,int H,int W,int OH,int OW){
    int nc=blockIdx.x; int n=nc / C; int c=nc - n*C; int tid=threadIdx.x; int M=OH*OW;
    float sum=0.f, sq=0.f;
    const float* xbase=x + ((n*C+c)*H*W);
    const float* wb=w + c*25;
    float* ybase=y + ((n*C+c)*OH*OW);
    for(int idx=tid; idx<M; idx+=blockDim.x){
        int oh=idx/OW; int ow=idx-oh*OW; int ih0=oh*2-2; int iw0=ow*2-2; float acc=0.f;
        #pragma unroll
        for(int kh=0; kh<5; ++kh){
            int ih=ih0+kh;
            if((unsigned)ih < (unsigned)H){
                #pragma unroll
                for(int kw=0; kw<5; ++kw){
                    int iw=iw0+kw;
                    if((unsigned)iw < (unsigned)W) acc += xbase[ih*W+iw]*wb[kh*5+kw];
                }
            }
        }
        ybase[idx]=acc; sum += acc; sq += acc*acc;
    }
    __shared__ float ss[256], qq[256]; ss[tid]=sum; qq[tid]=sq; __syncthreads();
    for(int s=128; s>0; s>>=1){ if(tid<s){ ss[tid]+=ss[tid+s]; qq[tid]+=qq[tid+s]; } __syncthreads(); }
    if(tid==0){ psum[n*C+c]=ss[0]; psq[n*C+c]=qq[0]; }
}

__global__ void stat_kernel(const float* __restrict__ psum,const float* __restrict__ psq,float* __restrict__ mean,float* __restrict__ inv,int N,int C,int OH,int OW,float eps){
    int c=blockIdx.x*blockDim.x+threadIdx.x; if(c>=C) return;
    float s=0.f,q=0.f; for(int n=0;n<N;n++){s+=psum[n*C+c];q+=psq[n*C+c];}
    float M=(float)(N*OH*OW); float m=s/M; float var=q/M-m*m; if(var<0)var=0; mean[c]=m; inv[c]=rsqrtf(var+eps);
}

__global__ void norm_relu(const float* __restrict__ y0,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ inv,long total,int C,int OH,int OW){
    long HW=(long)OH*OW;
    for(long idx=blockIdx.x*blockDim.x+threadIdx.x; idx<total; idx+=(long)blockDim.x*gridDim.x){
        int c=(idx/HW)%C;
        float v=(y0[idx]-mean[c])*inv[c]*gamma[c]+beta[c];
        y[idx]=fminf(fmaxf(v,0.f),6.f);
    }
}

torch::Tensor dw_bn_relu_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, double eps){
    int N=x.size(0), C=x.size(1), H=x.size(2), W=x.size(3); int OH=(H+1)/2, OW=(W+1)/2;
    auto y0=torch::empty({N,C,OH,OW}, x.options());
    auto y=torch::empty_like(y0);
    auto psum=torch::empty({N,C},x.options());
    auto psq=torch::empty({N,C},x.options());
    auto mean=torch::empty({C},x.options());
    auto inv=torch::empty({C},x.options());
    dw_conv_reduce<<<N*C,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),y0.data_ptr<float>(),psum.data_ptr<float>(),psq.data_ptr<float>(),N,C,H,W,OH,OW);
    stat_kernel<<<(C+255)/256,256>>>(psum.data_ptr<float>(),psq.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),N,C,OH,OW,(float)eps);
    long total=y.numel(); int blocks=(int)min((long)65535,(total+255)/256);
    norm_relu<<<blocks,256>>>(y0.data_ptr<float>(),y.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),total,C,OH,OW);
    return y;
}
'''

_ext = load_inline(name='mbconv_dw_fused_submit_v1', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['dw_bn_relu_cuda'], with_cuda=True, extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        super().__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = in_channels * expand_ratio
        if expand_ratio != 1:
            self.expand_conv = nn.Sequential(
                nn.Conv2d(in_channels, hidden_dim, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True)
            )
        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True)
        )
        self.project_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels)
        )
    def forward(self, x):
        identity = x
        if hasattr(self, 'expand_conv'):
            x = self.expand_conv(x)
        conv, bn, _ = self.depthwise_conv
        if conv.kernel_size == (5, 5) and conv.stride == (2, 2) and conv.padding == (2, 2) and x.is_cuda and x.dtype == torch.float32:
            x = _ext.dw_bn_relu_cuda(x.contiguous(), conv.weight, bn.weight, bn.bias, bn.eps)
        else:
            x = self.depthwise_conv(x)
        x = self.project_conv(x)
        if self.use_residual:
            x = x + identity
        return x