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

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

        :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.mbconv1 = self._make_mbconv_block(32, 16, 1, 1)
        self.mbconv2 = self._make_mbconv_block(16, 24, 2, 6)
        self.mbconv3 = self._make_mbconv_block(24, 40, 2, 6)
        self.mbconv4 = self._make_mbconv_block(40, 80, 2, 6)
        self.mbconv5 = self._make_mbconv_block(80, 112, 1, 6)
        self.mbconv6 = self._make_mbconv_block(112, 192, 2, 6)
        self.mbconv7 = self._make_mbconv_block(192, 320, 1, 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 _make_mbconv_block(self, in_channels, out_channels, stride, expand_ratio):
        """
        Creates a MBConv block.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param stride: Stride of the depthwise convolution.
        :param expand_ratio: Expansion ratio for the hidden layer.
        :return: A sequential MBConv block.
        """
        hidden_dim = round(in_channels * expand_ratio)
        return 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),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=stride, padding=1, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True),
            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 EfficientNetB1 model.

        :param x: Input tensor, shape (batch_size, 3, 240, 240)
        :return: Output tensor, shape (batch_size, num_classes)
        """
        x = F.relu(self.bn1(self.conv1(x)))
        
        x = self.mbconv1(x)
        x = self.mbconv2(x)
        x = self.mbconv3(x)
        x = self.mbconv4(x)
        x = self.mbconv5(x)
        x = self.mbconv6(x)
        x = self.mbconv7(x)
        
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = torch.flatten(x, 1)
        x = self.fc(x)
        
        return x

# Test code
batch_size = 10
input_shape = (3, 240, 240)
num_classes = 1000

def get_inputs():
    return [torch.rand(batch_size, *input_shape)]

def get_init_inputs():
    return [num_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 avg_pool_forward(torch::Tensor x);
torch::Tensor bn_act_eval(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps, int act);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void avg_pool_nhwc_kernel(const float* __restrict__ x, float* __restrict__ y, int N, int C, int H, int W) {
    int n = blockIdx.x;
    int c = blockIdx.y * blockDim.x + threadIdx.x;
    if (c >= C) return;
    float sum = 0.0f;
    const int HW = H * W;
    const float* base = x + n * HW * C + c;
    #pragma unroll
    for (int i = 0; i < 64; ++i) {
        if (i < HW) sum += base[i * C];
    }
    y[n * C + c] = sum / (float)HW;
}

torch::Tensor avg_pool_forward(torch::Tensor x) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    auto y = torch::empty({N, C}, x.options());
    dim3 block(256);
    dim3 grid(N, (C + 255) / 256);
    avg_pool_nhwc_kernel<<<grid, block>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}

__global__ void bn_act_eval_kernel(const float* __restrict__ x, float* __restrict__ y,
                                   const float* __restrict__ weight, const float* __restrict__ bias,
                                   const float* __restrict__ mean, const float* __restrict__ var,
                                   int total, int C, float eps, int act) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int c = idx % C;
        float v = (x[idx] - mean[c]) * rsqrtf(var[c] + eps) * weight[c] + bias[c];
        if (act == 1) {
            v = v > 0.0f ? v : 0.0f;
        } else if (act == 2) {
            v = v > 0.0f ? v : 0.0f;
            v = v < 6.0f ? v : 6.0f;
        }
        y[idx] = v;
    }
}

torch::Tensor bn_act_eval(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps, int act) {
    auto y = torch::empty_like(x, torch::MemoryFormat::ChannelsLast);
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int total = N * C * H * W;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 4096) blocks = 4096;
    bn_act_eval_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), total, C, (float)eps, act);
    return y;
}
"""

_ext = load_inline(name="effb1_ext_v3", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["avg_pool_forward", "bn_act_eval"], extra_cuda_cflags=["-O3"], extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(32)
        self.mbconv1 = self._make_mbconv_block(32, 16, 1, 1)
        self.mbconv2 = self._make_mbconv_block(16, 24, 2, 6)
        self.mbconv3 = self._make_mbconv_block(24, 40, 2, 6)
        self.mbconv4 = self._make_mbconv_block(40, 80, 2, 6)
        self.mbconv5 = self._make_mbconv_block(80, 112, 1, 6)
        self.mbconv6 = self._make_mbconv_block(112, 192, 2, 6)
        self.mbconv7 = self._make_mbconv_block(192, 320, 1, 6)
        self.conv2 = nn.Conv2d(320, 1280, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(1280)
        self.fc = nn.Linear(1280, num_classes)
        self.to(memory_format=torch.channels_last)
    def _make_mbconv_block(self, in_channels, out_channels, stride, expand_ratio):
        hidden_dim = round(in_channels * expand_ratio)
        return 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),
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=stride, padding=1, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True),
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
        )
    def _bn_eval(self, x, bn, act):
        return _ext.bn_act_eval(x, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps, act)
    def _mb_eval(self, x, mb):
        x = F.conv2d(x, mb[0].weight, None, 1, 0, 1, 1)
        x = self._bn_eval(x, mb[1], 2)
        x = F.conv2d(x, mb[3].weight, None, mb[3].stride, mb[3].padding, mb[3].dilation, mb[3].groups)
        x = self._bn_eval(x, mb[4], 2)
        x = F.conv2d(x, mb[6].weight, None, 1, 0, 1, 1)
        x = self._bn_eval(x, mb[7], 0)
        return x
    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last)
        if not self.training:
            x = F.conv2d(x, self.conv1.weight, None, self.conv1.stride, self.conv1.padding, self.conv1.dilation, self.conv1.groups)
            x = self._bn_eval(x, self.bn1, 1)
            x = self._mb_eval(x, self.mbconv1)
            x = self._mb_eval(x, self.mbconv2)
            x = self._mb_eval(x, self.mbconv3)
            x = self._mb_eval(x, self.mbconv4)
            x = self._mb_eval(x, self.mbconv5)
            x = self._mb_eval(x, self.mbconv6)
            x = self._mb_eval(x, self.mbconv7)
            x = F.conv2d(x, self.conv2.weight, None, 1, 0, 1, 1)
            x = self._bn_eval(x, self.bn2, 1)
        else:
            x = F.relu(self.bn1(self.conv1(x)))
            x = self.mbconv1(x)
            x = self.mbconv2(x)
            x = self.mbconv3(x)
            x = self.mbconv4(x)
            x = self.mbconv5(x)
            x = self.mbconv6(x)
            x = self.mbconv7(x)
            x = F.relu(self.bn2(self.conv2(x)))
        x = _ext.avg_pool_forward(x)
        x = self.fc(x)
        return x