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

class Model(nn.Module):
    def __init__(self, input_channels, stages, block_widths, output_classes):
        """
        :param input_channels: int, Number of input channels for the first layer
        :param stages: int, Number of stages in the RegNet architecture
        :param block_widths: List[int], Width (number of channels) for each block in the stages
        :param output_classes: int, Number of output classes for classification
        """
        super(Model, self).__init__()

        self.stages = stages
        self.block_widths = block_widths
        
        layers = []
        current_channels = input_channels
        
        # Construct the stages with their respective blocks
        for i in range(stages):
            layers.append(self._make_stage(current_channels, block_widths[i]))
            current_channels = block_widths[i]
        
        self.feature_extractor = nn.Sequential(*layers)
        
        # Final fully connected layer for classification
        self.fc = nn.Linear(block_widths[-1], output_classes)
    
    def _make_stage(self, in_channels, out_channels):
        """
        Creates a simple block for each stage.
        :param in_channels: int, number of input channels
        :param out_channels: int, number of output channels
        :return: nn.Sequential block with convolutional layers
        """
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        """
        Forward pass through the RegNet model.
        :param x: torch.Tensor of shape (batch_size, input_channels, height, width)
        :return: torch.Tensor of shape (batch_size, output_classes)
        """
        x = self.feature_extractor(x)
        x = torch.mean(x, dim=[2, 3])  # Global Average Pooling
        x = self.fc(x)
        return x

# Test code for the RegNet model
batch_size = 8
input_channels = 3
image_height, image_width = 224, 224
stages = 3
block_widths = [64, 128, 256]
output_classes = 10

def get_inputs():
    """ Generates random input tensor of shape (batch_size, input_channels, height, width) """
    return [torch.rand(batch_size, input_channels, image_height, image_width)]

def get_init_inputs():
    """ Initializes model parameters """
    return [input_channels, stages, block_widths, output_classes]

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 bn_relu_eval_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps); torch::Tensor bn_relu_pool_eval_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps); torch::Tensor bn_relu_train_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, double eps); torch::Tensor bn_relu_pool_train_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b, double eps); torch::Tensor avg_fc_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor b);
torch::Tensor regnet_bn_relu_eval(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps){return bn_relu_eval_cuda(x,w,b,mean,var,eps);} torch::Tensor regnet_bn_relu_pool_eval(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps){return bn_relu_pool_eval_cuda(x,w,b,mean,var,eps);} torch::Tensor regnet_bn_relu_train(torch::Tensor x, torch::Tensor w, torch::Tensor b, double eps){return bn_relu_train_cuda(x,w,b,eps);} torch::Tensor regnet_bn_relu_pool_train(torch::Tensor x, torch::Tensor w, torch::Tensor b, double eps){return bn_relu_pool_train_cuda(x,w,b,eps);} torch::Tensor regnet_avg_fc(torch::Tensor x, torch::Tensor w, torch::Tensor b){return avg_fc_cuda(x,w,b);} 
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
__global__ void stats_kernel(const float* __restrict__ x,float* __restrict__ mean,float* __restrict__ invstd,int N,int C,int HW,float eps){int c=blockIdx.x,tid=threadIdx.x,M=N*HW; float s=0.f,ss=0.f; for(int m=tid;m<M;m+=blockDim.x){int n=m/HW; int p=m-n*HW; float v=x[(n*C+c)*HW+p]; s+=v; ss+=v*v;} __shared__ float sh[512]; float* shs=sh; float* shss=sh+256; shs[tid]=s; shss[tid]=ss; __syncthreads(); for(int off=128;off>0;off>>=1){if(tid<off){shs[tid]+=shs[tid+off]; shss[tid]+=shss[tid+off];} __syncthreads();} if(tid==0){float mu=shs[0]/(float)M; float var=shss[0]/(float)M-mu*mu; if(var<0.f)var=0.f; mean[c]=mu; invstd[c]=rsqrtf(var+eps);}}
__global__ void bn_relu_apply_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,int total,int C,int HW){int idx=blockIdx.x*blockDim.x+threadIdx.x,stride=blockDim.x*gridDim.x; for(int i=idx;i<total;i+=stride){int c=(i/HW)%C; float v=(x[i]-mean[c])*invstd[c]*gamma[c]+beta[c]; y[i]=v>0.f?v:0.f;}}
__global__ void bn_relu_pool_apply_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,int total_out,int C,int H,int W,int Ho,int Wo){int idx=blockIdx.x*blockDim.x+threadIdx.x,stride=blockDim.x*gridDim.x; for(int i=idx;i<total_out;i+=stride){int t=i; int ow=t%Wo; t/=Wo; int oh=t%Ho; t/=Ho; int c=t%C; int n=t/C; int base=((n*C+c)*H+(oh*2))*W+(ow*2); float a=invstd[c]*gamma[c]; float d=beta[c]-mean[c]*a; float v0=x[base]*a+d; v0=v0>0.f?v0:0.f; float v1=x[base+1]*a+d; v1=v1>0.f?v1:0.f; float v2=x[base+W]*a+d; v2=v2>0.f?v2:0.f; float v3=x[base+W+1]*a+d; v3=v3>0.f?v3:0.f; float m1=v0>v1?v0:v1; float m2=v2>v3?v2:v3; y[i]=m1>m2?m1:m2;}}
__global__ void bn_relu_eval_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ var,int total,int C,int HW,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x,stride=blockDim.x*gridDim.x; for(int i=idx;i<total;i+=stride){int c=(i/HW)%C; float a=gamma[c]*rsqrtf(var[c]+eps); float d=beta[c]-mean[c]*a; float v=x[i]*a+d; y[i]=v>0.f?v:0.f;}}
__global__ void bn_relu_pool_eval_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ var,int total_out,int C,int H,int W,int Ho,int Wo,float eps){int idx=blockIdx.x*blockDim.x+threadIdx.x,stride=blockDim.x*gridDim.x; for(int i=idx;i<total_out;i+=stride){int t=i; int ow=t%Wo; t/=Wo; int oh=t%Ho; t/=Ho; int c=t%C; int n=t/C; int base=((n*C+c)*H+(oh*2))*W+(ow*2); float a=gamma[c]*rsqrtf(var[c]+eps); float d=beta[c]-mean[c]*a; float v0=x[base]*a+d; v0=v0>0.f?v0:0.f; float v1=x[base+1]*a+d; v1=v1>0.f?v1:0.f; float v2=x[base+W]*a+d; v2=v2>0.f?v2:0.f; float v3=x[base+W+1]*a+d; v3=v3>0.f?v3:0.f; float m1=v0>v1?v0:v1; float m2=v2>v3?v2:v3; y[i]=m1>m2?m1:m2;}}
__global__ void avg_fc_kernel(const float* __restrict__ x,const float* __restrict__ w,const float* __restrict__ b,float* __restrict__ y,int N,int C,int H,int Wd,int K){int nk=blockIdx.x,n=nk/K,k=nk-n*K,tid=threadIdx.x; float sum=0.f; int HW=H*Wd; const float* xn=x+n*C*HW; for(int c=tid;c<C;c+=blockDim.x){float s=0.f; const float* xp=xn+c*HW; for(int p=0;p<HW;++p)s+=xp[p]; sum+=s*w[k*C+c];} __shared__ float sm[256]; sm[tid]=sum; __syncthreads(); for(int off=128;off>0;off>>=1){if(tid<off)sm[tid]+=sm[tid+off]; __syncthreads();} if(tid==0)y[n*K+k]=sm[0]/(float)HW+b[k];}
torch::Tensor bn_relu_train_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,double eps){int N=x.size(0),C=x.size(1),H=x.size(2),Wd=x.size(3),HW=H*Wd; auto mean=torch::empty({C},x.options()),inv=torch::empty({C},x.options()),y=torch::empty_like(x); stats_kernel<<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),N,C,HW,(float)eps); int64_t total=x.numel(); int blocks=(int)((total+255)/256); if(blocks>4096)blocks=4096; bn_relu_apply_kernel<<<blocks,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),(int)total,C,HW); return y;}
torch::Tensor bn_relu_pool_train_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,double eps){int N=x.size(0),C=x.size(1),H=x.size(2),Wd=x.size(3),HW=H*Wd,Ho=H/2,Wo=Wd/2; auto mean=torch::empty({C},x.options()),inv=torch::empty({C},x.options()); auto y=torch::empty({N,C,Ho,Wo},x.options()); stats_kernel<<<C,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),N,C,HW,(float)eps); int64_t total=(int64_t)N*C*Ho*Wo; int blocks=(int)((total+255)/256); if(blocks>4096)blocks=4096; bn_relu_pool_apply_kernel<<<blocks,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),inv.data_ptr<float>(),(int)total,C,H,Wd,Ho,Wo); return y;}
torch::Tensor bn_relu_eval_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b,torch::Tensor mean,torch::Tensor var,double eps){auto y=torch::empty_like(x); int64_t total=x.numel(); int C=x.size(1),HW=x.size(2)*x.size(3); int blocks=(int)((total+255)/256); if(blocks>4096)blocks=4096; bn_relu_eval_kernel<<<blocks,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),(int)total,C,HW,(float)eps); return y;}
torch::Tensor bn_relu_pool_eval_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),H=x.size(2),Wd=x.size(3),Ho=H/2,Wo=Wd/2; auto y=torch::empty({N,C,Ho,Wo},x.options()); int64_t total=(int64_t)N*C*Ho*Wo; int blocks=(int)((total+255)/256); if(blocks>4096)blocks=4096; bn_relu_pool_eval_kernel<<<blocks,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),y.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),mean.data_ptr<float>(),var.data_ptr<float>(),(int)total,C,H,Wd,Ho,Wo,(float)eps); return y;}
torch::Tensor avg_fc_cuda(torch::Tensor x,torch::Tensor w,torch::Tensor b){int N=x.size(0),C=x.size(1),H=x.size(2),Wd=x.size(3),K=w.size(0); auto y=torch::empty({N,K},x.options()); avg_fc_kernel<<<N*K,256,0,at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),y.data_ptr<float>(),N,C,H,Wd,K); return y;}
'''
_ext=load_inline(name='regnet_fused_ext_final',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['regnet_bn_relu_eval','regnet_bn_relu_pool_eval','regnet_bn_relu_train','regnet_bn_relu_pool_train','regnet_avg_fc'],extra_cflags=['-O3'],extra_cuda_cflags=['-O3'],verbose=False)

class ModelNew(nn.Module):
    def __init__(self,input_channels,stages,block_widths,output_classes):
        super(ModelNew,self).__init__(); self.stages=stages; self.block_widths=block_widths; layers=[]; current_channels=input_channels
        for i in range(stages): layers.append(self._make_stage(current_channels,block_widths[i])); current_channels=block_widths[i]
        self.feature_extractor=nn.Sequential(*layers); self.fc=nn.Linear(block_widths[-1],output_classes)
    def _make_stage(self,in_channels,out_channels):
        return nn.Sequential(nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(),nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(),nn.MaxPool2d(kernel_size=2,stride=2))
    def forward(self,x):
        train=self.training
        for stage in self.feature_extractor:
            conv1=stage[0]; bn1=stage[1]; conv2=stage[3]; bn2=stage[4]
            x=F.conv2d(x,conv1.weight,None if train else conv1.bias,(1,1),(1,1),(1,1),1)
            x=_ext.regnet_bn_relu_train(x.contiguous(),bn1.weight,bn1.bias,bn1.eps) if train else _ext.regnet_bn_relu_eval(x.contiguous(),bn1.weight,bn1.bias,bn1.running_mean,bn1.running_var,bn1.eps)
            x=F.conv2d(x,conv2.weight,None if train else conv2.bias,(1,1),(1,1),(1,1),1)
            x=_ext.regnet_bn_relu_pool_train(x.contiguous(),bn2.weight,bn2.bias,bn2.eps) if train else _ext.regnet_bn_relu_pool_eval(x.contiguous(),bn2.weight,bn2.bias,bn2.running_mean,bn2.running_var,bn2.eps)
        return _ext.regnet_avg_fc(x.contiguous(),self.fc.weight,self.fc.bias)