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_train_relu_inplace(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, double eps, double momentum);
torch::Tensor bn_train_relu_pool(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, double eps, double momentum);
torch::Tensor gap_fc(torch::Tensor x, torch::Tensor w, torch::Tensor b);
"""
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void bn_stats1_kernel(const float* __restrict__ x,float* __restrict__ psum,float* __restrict__ psq,int B,int C,int H,int W,int parts){int c=blockIdx.x,p=blockIdx.y,tid=threadIdx.x;int HW=H*W,N=B*HW;float sum=0,ss=0;for(int n=p*blockDim.x+tid;n<N;n+=blockDim.x*parts){int b=n/HW;int hw=n-b*HW;float v=x[((int64_t)b*C+c)*HW+hw];sum+=v;ss+=v*v;}__shared__ float ssum[256],ssqq[256];ssum[tid]=sum;ssqq[tid]=ss;__syncthreads();for(int off=128;off>0;off>>=1){if(tid<off){ssum[tid]+=ssum[tid+off];ssqq[tid]+=ssqq[tid+off];}__syncthreads();}if(tid==0){psum[c*parts+p]=ssum[0];psq[c*parts+p]=ssqq[0];}}
__global__ void bn_stats_finalize_kernel(const float* __restrict__ psum,const float* __restrict__ psq,float* __restrict__ mean,float* __restrict__ invstd,float* __restrict__ running_mean,float* __restrict__ running_var,int B,int C,int H,int W,int parts,float eps,float momentum){int c=blockIdx.x,tid=threadIdx.x;float sum=0,ss=0;for(int p=tid;p<parts;p+=blockDim.x){sum+=psum[c*parts+p];ss+=psq[c*parts+p];}__shared__ float ssum[32],ssqq[32];ssum[tid]=sum;ssqq[tid]=ss;__syncthreads();for(int off=16;off>0;off>>=1){if(tid<off){ssum[tid]+=ssum[tid+off];ssqq[tid]+=ssqq[tid+off];}__syncthreads();}if(tid==0){int N=B*H*W;float m=ssum[0]/(float)N;float var=ssqq[0]/(float)N-m*m;var=var>0?var:0;mean[c]=m;invstd[c]=rsqrtf(var+eps);running_mean[c]=running_mean[c]*(1.0f-momentum)+momentum*m;running_var[c]=running_var[c]*(1.0f-momentum)+momentum*var*((float)N/(float)(N-1));}}
__global__ void bn_stats_kernel(const float* __restrict__ x,float* __restrict__ mean,float* __restrict__ invstd,float* __restrict__ running_mean,float* __restrict__ running_var,int B,int C,int H,int W,float eps,float momentum){int c=blockIdx.x,tid=threadIdx.x;int HW=H*W,N=B*HW;float sum=0,ss=0;for(int n=tid;n<N;n+=blockDim.x){int b=n/HW;int hw=n-b*HW;float v=x[((int64_t)b*C+c)*HW+hw];sum+=v;ss+=v*v;}__shared__ float ssum[256],ssq[256];ssum[tid]=sum;ssq[tid]=ss;__syncthreads();for(int off=128;off>0;off>>=1){if(tid<off){ssum[tid]+=ssum[tid+off];ssq[tid]+=ssq[tid+off];}__syncthreads();}if(tid==0){float m=ssum[0]/(float)N;float var=ssq[0]/(float)N-m*m;var=var>0?var:0;mean[c]=m;invstd[c]=rsqrtf(var+eps);running_mean[c]=running_mean[c]*(1.0f-momentum)+momentum*m;running_var[c]=running_var[c]*(1.0f-momentum)+momentum*var*((float)N/(float)(N-1));}}
__global__ void bn_relu_inplace_kernel(float* __restrict__ x,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,int64_t total,int C,int H,int W){int64_t idx=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;int64_t stride=(int64_t)blockDim.x*gridDim.x;int HW=H*W;for(;idx<total;idx+=stride){int c=(idx/HW)%C;float sc=invstd[c]*gamma[c];float v=x[idx]*sc+(beta[c]-mean[c]*sc);x[idx]=v>0?v:0;}}
__global__ void bn_relu_pool_kernel(const float* __restrict__ x,float* __restrict__ y,const float* __restrict__ gamma,const float* __restrict__ beta,const float* __restrict__ mean,const float* __restrict__ invstd,int64_t out_total,int C,int H,int W,int OH,int OW){int64_t idx=(int64_t)blockIdx.x*blockDim.x+threadIdx.x;int64_t stride=(int64_t)blockDim.x*gridDim.x;int HW=H*W;for(;idx<out_total;idx+=stride){int ow=idx%OW;int t=idx/OW;int oh=t%OH;t/=OH;int c=t%C;int b=t/C;float sc=invstd[c]*gamma[c];float off=beta[c]-mean[c]*sc;int64_t base=((int64_t)b*C+c)*HW+(oh*2)*W+(ow*2);float v0=x[base]*sc+off;v0=v0>0?v0:0;float v1=x[base+1]*sc+off;v1=v1>0?v1:0;float v2=x[base+W]*sc+off;v2=v2>0?v2:0;float v3=x[base+W+1]*sc+off;v3=v3>0?v3:0;y[idx]=fmaxf(fmaxf(v0,v1),fmaxf(v2,v3));}}
__global__ void gap_kernel(const float* __restrict__ x,float* __restrict__ feat,int B,int C,int H,int W){int bc=blockIdx.x;int b=bc/C;int c=bc-b*C;int HW=H*W;float acc=0;for(int i=threadIdx.x;i<HW;i+=blockDim.x)acc+=x[((int64_t)b*C+c)*HW+i];__shared__ float s[256];s[threadIdx.x]=acc;__syncthreads();for(int off=128;off>0;off>>=1){if(threadIdx.x<off)s[threadIdx.x]+=s[threadIdx.x+off];__syncthreads();}if(threadIdx.x==0)feat[bc]=s[0]/(float)HW;}
__global__ void fc_kernel(const float* __restrict__ feat,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ y,int B,int C,int O){int bo=blockIdx.x;int b=bo/O;int o=bo-b*O;float acc=0;for(int c=threadIdx.x;c<C;c+=blockDim.x)acc+=feat[b*C+c]*w[o*C+c];__shared__ float s[256];s[threadIdx.x]=acc;__syncthreads();for(int off=128;off>0;off>>=1){if(threadIdx.x<off)s[threadIdx.x]+=s[threadIdx.x+off];__syncthreads();}if(threadIdx.x==0)y[bo]=bias[o]+s[0];}
static inline int choose_parts(int N){return (N>=100000)?8:((N>=30000)?4:1);}static inline void launch_stats(torch::Tensor x,torch::Tensor mean,torch::Tensor invstd,torch::Tensor rmean,torch::Tensor rvar,int B,int C,int H,int W,float eps,float momentum,cudaStream_t stream){int N=B*H*W;int parts=choose_parts(N);if(parts==1){bn_stats_kernel<<<C,256,0,stream>>>(x.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),rmean.data_ptr<float>(),rvar.data_ptr<float>(),B,C,H,W,eps,momentum);}else{auto psum=torch::empty({C,parts},x.options());auto psq=torch::empty({C,parts},x.options());dim3 grid(C,parts);bn_stats1_kernel<<<grid,256,0,stream>>>(x.data_ptr<float>(),psum.data_ptr<float>(),psq.data_ptr<float>(),B,C,H,W,parts);bn_stats_finalize_kernel<<<C,32,0,stream>>>(psum.data_ptr<float>(),psq.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),rmean.data_ptr<float>(),rvar.data_ptr<float>(),B,C,H,W,parts,eps,momentum);}}
torch::Tensor bn_train_relu_inplace(torch::Tensor x,torch::Tensor gamma,torch::Tensor beta,torch::Tensor rmean,torch::Tensor rvar,double eps,double momentum){int B=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3);auto mean=torch::empty({C},x.options());auto invstd=torch::empty({C},x.options());cudaStream_t stream=at::cuda::getCurrentCUDAStream();launch_stats(x,mean,invstd,rmean,rvar,B,C,H,W,(float)eps,(float)momentum,stream);int threads=256;int64_t total=x.numel();int blocks=(int)std::min<int64_t>((total+threads-1)/threads,4096);bn_relu_inplace_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),total,C,H,W);return x;}
torch::Tensor bn_train_relu_pool(torch::Tensor x,torch::Tensor gamma,torch::Tensor beta,torch::Tensor rmean,torch::Tensor rvar,double eps,double momentum){int B=x.size(0),C=x.size(1),H=x.size(2),W=x.size(3),OH=H/2,OW=W/2;auto mean=torch::empty({C},x.options());auto invstd=torch::empty({C},x.options());auto y=torch::empty({B,C,OH,OW},x.options());cudaStream_t stream=at::cuda::getCurrentCUDAStream();launch_stats(x,mean,invstd,rmean,rvar,B,C,H,W,(float)eps,(float)momentum,stream);int threads=256;int64_t total=y.numel();int blocks=(int)std::min<int64_t>((total+threads-1)/threads,4096);bn_relu_pool_kernel<<<blocks,threads,0,stream>>>(x.data_ptr<float>(),y.data_ptr<float>(),gamma.data_ptr<float>(),beta.data_ptr<float>(),mean.data_ptr<float>(),invstd.data_ptr<float>(),total,C,H,W,OH,OW);return y;}
torch::Tensor gap_fc(torch::Tensor x,torch::Tensor w,torch::Tensor b){int B=x.size(0),C=x.size(1),H=x.size(2),Wd=x.size(3),O=w.size(0);auto feat=torch::empty({B,C},x.options());auto y=torch::empty({B,O},x.options());cudaStream_t stream=at::cuda::getCurrentCUDAStream();gap_kernel<<<B*C,256,0,stream>>>(x.data_ptr<float>(),feat.data_ptr<float>(),B,C,H,Wd);fc_kernel<<<B*O,256,0,stream>>>(feat.data_ptr<float>(),w.data_ptr<float>(),b.data_ptr<float>(),y.data_ptr<float>(),B,C,O);return y;}
"""
_ext=load_inline(name="regnet_fused_ext_v4",cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=["bn_train_relu_inplace","bn_train_relu_pool","gap_fc"],extra_cuda_cflags=["-O3"],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,input_channels,stages,block_widths,output_classes):
        super().__init__(); self.stages=stages; self.block_widths=block_widths; layers=[]; cur=input_channels
        for i in range(stages): layers.append(self._make_stage(cur,block_widths[i])); cur=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,3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(),nn.Conv2d(out_channels,out_channels,3,padding=1),nn.BatchNorm2d(out_channels),nn.ReLU(),nn.MaxPool2d(2,2))
    def forward(self,x):
        for st in self.feature_extractor:
            c1,bn1,_,c2,bn2,_,_=st
            x=F.conv2d(x,c1.weight,c1.bias,stride=1,padding=1); x=_ext.bn_train_relu_inplace(x,bn1.weight,bn1.bias,bn1.running_mean,bn1.running_var,bn1.eps,bn1.momentum)
            x=F.conv2d(x,c2.weight,c2.bias,stride=1,padding=1); x=_ext.bn_train_relu_pool(x,bn2.weight,bn2.bias,bn2.running_mean,bn2.running_var,bn2.eps,bn2.momentum)
        return _ext.gap_fc(x,self.fc.weight,self.fc.bias)