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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Make cuDNN pick good convolution algorithms for the fixed EfficientNet shapes.
torch.backends.cudnn.benchmark = True

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor bn_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                     torch::Tensor mean, torch::Tensor var, double eps, int64_t act);
torch::Tensor bn_add(torch::Tensor x, torch::Tensor ident, torch::Tensor weight,
                     torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor avg_pool_2d(torch::Tensor x);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdint.h>

__global__ void bn_act_kernel(const float* __restrict__ x, float* __restrict__ y,
                              const float* __restrict__ w, const float* __restrict__ b,
                              const float* __restrict__ m, const float* __restrict__ v,
                              long long total, int C, int HW, float eps, int act) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (int)((idx / HW) % C);
    float val = (x[idx] - m[c]) * rsqrtf(v[c] + eps) * w[c] + b[c];
    if (act == 1) {
        val = val > 0.0f ? val : 0.0f;
    } else if (act == 2) {
        val = val > 0.0f ? val : 0.0f;
        val = val < 6.0f ? val : 6.0f;
    }
    y[idx] = val;
}

__global__ void bn_add_kernel(const float* __restrict__ x, const float* __restrict__ ident,
                              float* __restrict__ y,
                              const float* __restrict__ w, const float* __restrict__ b,
                              const float* __restrict__ m, const float* __restrict__ v,
                              long long total, int C, int HW, float eps) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (int)((idx / HW) % C);
    float val = (x[idx] - m[c]) * rsqrtf(v[c] + eps) * w[c] + b[c];
    y[idx] = val + ident[idx];
}

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

torch::Tensor bn_act(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                     torch::Tensor mean, torch::Tensor var, double eps, int64_t act) {
    auto y = torch::empty_like(x);
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    long long total = (long long)N * C * H * W;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    bn_act_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        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, H*W, (float)eps, (int)act);
    return y;
}

torch::Tensor bn_add(torch::Tensor x, torch::Tensor ident, torch::Tensor weight,
                     torch::Tensor bias, torch::Tensor mean, torch::Tensor var, double eps) {
    auto y = torch::empty_like(x);
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    long long total = (long long)N * C * H * W;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    bn_add_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), ident.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, H*W, (float)eps);
    return y;
}

torch::Tensor avg_pool_2d(torch::Tensor x) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int H = (int)x.size(2);
    const int W = (int)x.size(3);
    auto y = torch::empty({N, C}, x.options());
    int total = N * C;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    avg_pool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}
'''

_eff_ext = load_inline(
    name="effnet_b0_fused_bn_pool_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["bn_act", "bn_add", "avg_pool_2d"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class MBConvNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        super(MBConvNew, self).__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = (kernel_size - 1) // 2
        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=self.padding, 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):
        # In training mode keep exact PyTorch BatchNorm semantics (batch statistics and running updates).
        if self.training:
            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 = x + identity
            return x

        identity = x
        if hasattr(self, 'expand_conv'):
            conv = self.expand_conv[0]
            bn = self.expand_conv[1]
            x = F.conv2d(x, conv.weight, None, 1, 0, 1, 1)
            x = _eff_ext.bn_act(x, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps, 2)

        conv = self.depthwise_conv[0]
        bn = self.depthwise_conv[1]
        x = F.conv2d(x, conv.weight, None, self.stride, self.padding, 1, conv.groups)
        x = _eff_ext.bn_act(x, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps, 2)

        conv = self.project_conv[0]
        bn = self.project_conv[1]
        x = F.conv2d(x, conv.weight, None, 1, 0, 1, 1)
        if self.use_residual:
            x = _eff_ext.bn_add(x, identity, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps)
        else:
            x = _eff_ext.bn_act(x, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps, 0)
        return x

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.blocks = nn.Sequential(
            MBConvNew(32, 16, kernel_size=3, stride=1, expand_ratio=1),
            MBConvNew(16, 24, kernel_size=3, stride=2, expand_ratio=6),
            MBConvNew(24, 24, kernel_size=3, stride=1, expand_ratio=6),
            MBConvNew(24, 40, kernel_size=5, stride=2, expand_ratio=6),
            MBConvNew(40, 40, kernel_size=5, stride=1, expand_ratio=6),
            MBConvNew(40, 80, kernel_size=3, stride=2, expand_ratio=6),
            MBConvNew(80, 80, kernel_size=3, stride=1, expand_ratio=6),
            MBConvNew(80, 112, kernel_size=5, stride=1, expand_ratio=6),
            MBConvNew(112, 112, kernel_size=5, stride=1, expand_ratio=6),
            MBConvNew(112, 192, kernel_size=5, stride=2, expand_ratio=6),
            MBConvNew(192, 192, kernel_size=5, stride=1, expand_ratio=6),
            MBConvNew(192, 192, kernel_size=5, stride=1, expand_ratio=6),
            MBConvNew(192, 320, kernel_size=3, stride=1, expand_ratio=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)

    def forward(self, x):
        if self.training:
            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

        x = F.conv2d(x, self.conv1.weight, None, 2, 1, 1, 1)
        x = _eff_ext.bn_act(x, self.bn1.weight, self.bn1.bias, self.bn1.running_mean, self.bn1.running_var, self.bn1.eps, 1)
        x = self.blocks(x)
        x = F.conv2d(x, self.conv2.weight, None, 1, 0, 1, 1)
        x = _eff_ext.bn_act(x, self.bn2.weight, self.bn2.bias, self.bn2.running_mean, self.bn2.running_var, self.bn2.eps, 1)
        x = _eff_ext.avg_pool_2d(x)
        x = self.fc(x)
        return x