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

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        """
        EfficientNetB0 architecture implementation in PyTorch.

        :param num_classes: The number of output classes (default is 1000 for ImageNet).
        """
        super(Model, self).__init__()
        
        # Initial convolutional layer
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(32)
        
        # MBConv blocks
        self.blocks = nn.Sequential(
            # MBConv1 (32, 16, 1, 1)
            MBConv(32, 16, kernel_size=3, stride=1, expand_ratio=1),
            # MBConv6 (16, 24, 2, 6)
            MBConv(16, 24, kernel_size=3, stride=2, expand_ratio=6),
            # MBConv6 (24, 24, 1, 6)
            MBConv(24, 24, kernel_size=3, stride=1, expand_ratio=6),
            # MBConv6 (24, 40, 2, 6)
            MBConv(24, 40, kernel_size=5, stride=2, expand_ratio=6),
            # MBConv6 (40, 40, 1, 6)
            MBConv(40, 40, kernel_size=5, stride=1, expand_ratio=6),
            # MBConv6 (40, 80, 2, 6)
            MBConv(40, 80, kernel_size=3, stride=2, expand_ratio=6),
            # MBConv6 (80, 80, 1, 6)
            MBConv(80, 80, kernel_size=3, stride=1, expand_ratio=6),
            # MBConv6 (80, 112, 1, 6)
            MBConv(80, 112, kernel_size=5, stride=1, expand_ratio=6),
            # MBConv6 (112, 112, 1, 6)
            MBConv(112, 112, kernel_size=5, stride=1, expand_ratio=6),
            # MBConv6 (112, 192, 2, 6)
            MBConv(112, 192, kernel_size=5, stride=2, expand_ratio=6),
            # MBConv6 (192, 192, 1, 6)
            MBConv(192, 192, kernel_size=5, stride=1, expand_ratio=6),
            # MBConv6 (192, 192, 1, 6)
            MBConv(192, 192, kernel_size=5, stride=1, expand_ratio=6),
            # MBConv6 (192, 320, 1, 6)
            MBConv(192, 320, kernel_size=3, stride=1, expand_ratio=6)
        )
        
        # Final convolutional layer
        self.conv2 = nn.Conv2d(320, 1280, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(1280)
        
        # Fully connected layer
        self.fc = nn.Linear(1280, num_classes)
    
    def forward(self, x):
        """
        Forward pass of the EfficientNetB0 model.

        :param x: The input tensor, shape (batch_size, 3, 224, 224)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.blocks(x)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

class MBConv(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(MBConv, 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
num_classes = 1000

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

def get_init_inputs():
    return [num_classes]

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

_cpp = r"""
#include <torch/extension.h>
torch::Tensor bn_act_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps, int act);
torch::Tensor bn_add_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, torch::Tensor res, double eps);
torch::Tensor bn_relu_avg_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor bn_act_train_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor rmean, torch::Tensor rvar, double eps, double momentum, int act);
torch::Tensor bn_add_train_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor rmean, torch::Tensor rvar, torch::Tensor res, double eps, double momentum);
torch::Tensor bn_relu_avg_train_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor rmean, torch::Tensor rvar, double eps, double momentum);
"""
_cuda = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float activate(float v, int act){ if(act==1) return fmaxf(v,0.f); if(act==2) return fminf(fmaxf(v,0.f),6.f); return v; }
__global__ void bn_act_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ mean,const float* __restrict__ var,float* __restrict__ y,long total,int C,int HW,float eps,int act){ long idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return; int c=(idx/HW)%C; float v=(x[idx]-mean[c])*rsqrtf(var[c]+eps)*w[c]+b[c]; y[idx]=activate(v,act); }
__global__ void bn_add_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ mean,const float* __restrict__ var,const float* __restrict__ res,float* __restrict__ y,long total,int C,int HW,float eps){ long idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=total)return; int c=(idx/HW)%C; float v=(x[idx]-mean[c])*rsqrtf(var[c]+eps)*w[c]+b[c]; y[idx]=v+res[idx]; }
__global__ void bn_relu_avg_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,const float* __restrict__ mean,const float* __restrict__ var,float* __restrict__ y,int N,int C,int HW,float eps){ int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=N*C)return; int c=idx%C,n=idx/C; const float* p=x+((long)n*C+c)*HW; float sc=w[c]*rsqrtf(var[c]+eps), bs=b[c]-mean[c]*sc, sum=0.f; for(int i=0;i<HW;i++){ float v=p[i]*sc+bs; sum+=fmaxf(v,0.f);} y[idx]=sum/(float)HW; }

__global__ void bn_train_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ rmean,float* __restrict__ rvar,float* __restrict__ y,int N,int C,int HW,float eps,float mom,int act){
    int c=blockIdx.x, tid=threadIdx.x, M=N*HW; __shared__ float ss[256]; __shared__ float ss2[256];
    float sum=0.f, sum2=0.f;
    for(int i=tid;i<M;i+=blockDim.x){ int n=i/HW, h=i-n*HW; float v=x[((long)n*C+c)*HW+h]; sum+=v; sum2+=v*v; }
    ss[tid]=sum; ss2[tid]=sum2; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){ if(tid<s){ ss[tid]+=ss[tid+s]; ss2[tid]+=ss2[tid+s]; } __syncthreads(); }
    float mean=ss[0]/(float)M; float var=fmaxf(ss2[0]/(float)M-mean*mean,0.f); float sc=w[c]*rsqrtf(var+eps); float bs=b[c]-mean*sc;
    if(tid==0){ rmean[c]=rmean[c]*(1.f-mom)+mom*mean; if(M>1) rvar[c]=rvar[c]*(1.f-mom)+mom*var*((float)M/(float)(M-1)); }
    for(int i=tid;i<M;i+=blockDim.x){ int n=i/HW, h=i-n*HW; long idx=((long)n*C+c)*HW+h; y[idx]=activate(x[idx]*sc+bs,act); }
}
__global__ void bn_add_train_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ rmean,float* __restrict__ rvar,const float* __restrict__ res,float* __restrict__ y,int N,int C,int HW,float eps,float mom){
    int c=blockIdx.x, tid=threadIdx.x, M=N*HW; __shared__ float ss[256]; __shared__ float ss2[256]; float sum=0.f,sum2=0.f;
    for(int i=tid;i<M;i+=blockDim.x){ int n=i/HW,h=i-n*HW; float v=x[((long)n*C+c)*HW+h]; sum+=v; sum2+=v*v; }
    ss[tid]=sum; ss2[tid]=sum2; __syncthreads(); for(int s=128;s>0;s>>=1){ if(tid<s){ss[tid]+=ss[tid+s]; ss2[tid]+=ss2[tid+s];} __syncthreads(); }
    float mean=ss[0]/(float)M; float var=fmaxf(ss2[0]/(float)M-mean*mean,0.f); float sc=w[c]*rsqrtf(var+eps); float bs=b[c]-mean*sc;
    if(tid==0){ rmean[c]=rmean[c]*(1.f-mom)+mom*mean; if(M>1) rvar[c]=rvar[c]*(1.f-mom)+mom*var*((float)M/(float)(M-1)); }
    for(int i=tid;i<M;i+=blockDim.x){ int n=i/HW,h=i-n*HW; long idx=((long)n*C+c)*HW+h; y[idx]=x[idx]*sc+bs+res[idx]; }
}
__global__ void bn_relu_avg_train_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ rmean,float* __restrict__ rvar,float* __restrict__ y,int N,int C,int HW,float eps,float mom){
    int c=blockIdx.x, tid=threadIdx.x, M=N*HW; __shared__ float ss[256]; __shared__ float ss2[256]; float sum=0.f,sum2=0.f;
    for(int i=tid;i<M;i+=blockDim.x){ int n=i/HW,h=i-n*HW; float v=x[((long)n*C+c)*HW+h]; sum+=v; sum2+=v*v; }
    ss[tid]=sum; ss2[tid]=sum2; __syncthreads(); for(int s=128;s>0;s>>=1){ if(tid<s){ss[tid]+=ss[tid+s]; ss2[tid]+=ss2[tid+s];} __syncthreads(); }
    float mean=ss[0]/(float)M; float var=fmaxf(ss2[0]/(float)M-mean*mean,0.f); float sc=w[c]*rsqrtf(var+eps); float bs=b[c]-mean*sc;
    if(tid==0){ rmean[c]=rmean[c]*(1.f-mom)+mom*mean; if(M>1) rvar[c]=rvar[c]*(1.f-mom)+mom*var*((float)M/(float)(M-1)); }
    for(int n=0;n<N;n++){ float s=0.f; for(int i=tid;i<HW;i+=blockDim.x){ float v=x[((long)n*C+c)*HW+i]*sc+bs; s+=fmaxf(v,0.f);} ss[tid]=s; __syncthreads(); for(int off=128;off>0;off>>=1){ if(tid<off) ss[tid]+=ss[tid+off]; __syncthreads(); } if(tid==0) y[n*C+c]=ss[0]/(float)HW; __syncthreads(); }
}

torch::Tensor bn_act_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor mean,torch::Tensor var,double eps,int act){ auto y=torch::empty_like(x); long total=x.numel(); int C=x.size(1),HW=x.size(2)*x.size(3); int th=256,bl=(total+th-1)/th; bn_act_kernel<<<bl,th>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),y.data_ptr<float>(),total,C,HW,(float)eps,act); return y; }
torch::Tensor bn_add_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor mean,torch::Tensor var,torch::Tensor res,double eps){ auto y=torch::empty_like(x); long total=x.numel(); int C=x.size(1),HW=x.size(2)*x.size(3); int th=256,bl=(total+th-1)/th; bn_add_kernel<<<bl,th>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),res.data_ptr<float>(),y.data_ptr<float>(),total,C,HW,(float)eps); return y; }
torch::Tensor bn_relu_avg_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor mean,torch::Tensor var,double eps){ int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); auto y=torch::empty({N,C},x.options()); int th=256,bl=(N*C+th-1)/th; bn_relu_avg_kernel<<<bl,th>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),y.data_ptr<float>(),N,C,HW,(float)eps); return y; }
torch::Tensor bn_act_train_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor rmean,torch::Tensor rvar,double eps,double momentum,int act){ auto y=torch::empty_like(x); int C=x.size(1); bn_train_kernel<<<C,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),rmean.data_ptr<float>(),rvar.data_ptr<float>(),y.data_ptr<float>(),x.size(0),C,x.size(2)*x.size(3),(float)eps,(float)momentum,act); return y; }
torch::Tensor bn_add_train_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor rmean,torch::Tensor rvar,torch::Tensor res,double eps,double momentum){ auto y=torch::empty_like(x); int C=x.size(1); bn_add_train_kernel<<<C,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),rmean.data_ptr<float>(),rvar.data_ptr<float>(),res.data_ptr<float>(),y.data_ptr<float>(),x.size(0),C,x.size(2)*x.size(3),(float)eps,(float)momentum); return y; }
torch::Tensor bn_relu_avg_train_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor rmean,torch::Tensor rvar,double eps,double momentum){ int N=x.size(0),C=x.size(1),HW=x.size(2)*x.size(3); auto y=torch::empty({N,C},x.options()); bn_relu_avg_train_kernel<<<C,256>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),rmean.data_ptr<float>(),rvar.data_ptr<float>(),y.data_ptr<float>(),N,C,HW,(float)eps,(float)momentum); return y; }
"""
_ext=load_inline(name="effnet_fused_bn_ext_v2",cpp_sources=_cpp,cuda_sources=_cuda,functions=["bn_act_cuda","bn_add_cuda","bn_relu_avg_cuda","bn_act_train_cuda","bn_add_train_cuda","bn_relu_avg_train_cuda"],extra_cuda_cflags=["-O3"],extra_cflags=["-O3"],verbose=False)

def _mom(bn):
    return 0.1 if bn.momentum is None else float(bn.momentum)
def _bn_act(x,bn,act,training):
    return _ext.bn_act_train_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps,_mom(bn),act) if training else _ext.bn_act_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps,act)
def _bn_add(x,bn,res,training):
    return _ext.bn_add_train_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,res,bn.eps,_mom(bn)) if training else _ext.bn_add_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,res,bn.eps)
def _bn_avg(x,bn,training):
    return _ext.bn_relu_avg_train_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps,_mom(bn)) if training else _ext.bn_relu_avg_cuda(x,bn.weight,bn.bias,bn.running_mean,bn.running_var,bn.eps)
class MBConvNew(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,1,1,0,bias=False),nn.BatchNorm2d(hidden_dim),nn.ReLU6(inplace=True))
        self.depthwise_conv=nn.Sequential(nn.Conv2d(hidden_dim,hidden_dim,kernel_size,stride,(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,1,1,0,bias=False),nn.BatchNorm2d(out_channels))
    def forward(self,x):
        identity=x; tr=self.training
        if hasattr(self,'expand_conv'):
            x=self.expand_conv[0](x); x=_bn_act(x,self.expand_conv[1],2,tr)
        x=self.depthwise_conv[0](x); x=_bn_act(x,self.depthwise_conv[1],2,tr)
        x=self.project_conv[0](x)
        x=_bn_add(x,self.project_conv[1],identity,tr) if self.use_residual else _bn_act(x,self.project_conv[1],0,tr)
        return x
class ModelNew(nn.Module):
    def __init__(self,num_classes=1000):
        super().__init__(); self.conv1=nn.Conv2d(3,32,3,2,1,bias=False); self.bn1=nn.BatchNorm2d(32)
        self.blocks=nn.Sequential(MBConvNew(32,16,3,1,1),MBConvNew(16,24,3,2,6),MBConvNew(24,24,3,1,6),MBConvNew(24,40,5,2,6),MBConvNew(40,40,5,1,6),MBConvNew(40,80,3,2,6),MBConvNew(80,80,3,1,6),MBConvNew(80,112,5,1,6),MBConvNew(112,112,5,1,6),MBConvNew(112,192,5,2,6),MBConvNew(192,192,5,1,6),MBConvNew(192,192,5,1,6),MBConvNew(192,320,3,1,6))
        self.conv2=nn.Conv2d(320,1280,1,1,0,bias=False); self.bn2=nn.BatchNorm2d(1280); self.fc=nn.Linear(1280,num_classes)
    def forward(self,x):
        tr=self.training; x=self.conv1(x); x=_bn_act(x,self.bn1,1,tr); x=self.blocks(x); x=self.conv2(x); x=_bn_avg(x,self.bn2,tr); return self.fc(x)