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

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

__global__ void spatial_mean_kernel(const float* __restrict__ x, float* __restrict__ m,
                                    int N, int C, int H, int W) {
    int idx = blockIdx.x;
    int c = idx % C;
    int n = idx / C;
    int HW = H * W;
    const float* p = x + ((n * C + c) * H * W);
    float s = 0.0f;
    // H*W is normally 64 here; keep generic.
    for (int i = threadIdx.x; i < HW; i += blockDim.x) s += p[i];
    __shared__ float sm[128];
    sm[threadIdx.x] = s;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) sm[threadIdx.x] += sm[threadIdx.x + off];
        __syncthreads();
    }
    if (threadIdx.x == 0) m[n * C + c] = sm[0] / (float)HW;
}

__global__ void hidden_kernel(const float* __restrict__ m,
                              const float* __restrict__ cw,
                              const float* __restrict__ bw,
                              const float* __restrict__ bb,
                              const float* __restrict__ brm,
                              const float* __restrict__ brv,
                              float* __restrict__ h,
                              int N, int C, int K, float eps) {
    int idx = blockIdx.x;
    int k = idx % K;
    int n = idx / K;
    float sum = 0.0f;
    const float* mp = m + n * C;
    const float* wp = cw + k * C;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        sum += mp[c] * wp[c];
    }
    __shared__ float sm[256];
    sm[threadIdx.x] = sum;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) sm[threadIdx.x] += sm[threadIdx.x + off];
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        float y = (sm[0] - brm[k]) * rsqrtf(brv[k] + eps) * bw[k] + bb[k];
        h[n * K + k] = y > 0.0f ? y : 0.0f;
    }
}

__global__ void fc_kernel(const float* __restrict__ h,
                          const float* __restrict__ fw,
                          const float* __restrict__ fb,
                          float* __restrict__ out,
                          int N, int K, int O) {
    int idx = blockIdx.x;
    int o = idx % O;
    int n = idx / O;
    const float* hp = h + n * K;
    const float* wp = fw + o * K;
    float sum = 0.0f;
    for (int k = threadIdx.x; k < K; k += blockDim.x) {
        sum += hp[k] * wp[k];
    }
    __shared__ float sm[256];
    sm[threadIdx.x] = sum;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (threadIdx.x < off) sm[threadIdx.x] += sm[threadIdx.x + off];
        __syncthreads();
    }
    if (threadIdx.x == 0) out[n * O + o] = sm[0] + fb[o];
}

torch::Tensor eff_tail_forward(torch::Tensor x,
                               torch::Tensor conv_w,
                               torch::Tensor bn_w,
                               torch::Tensor bn_b,
                               torch::Tensor bn_rm,
                               torch::Tensor bn_rv,
                               torch::Tensor fc_w,
                               torch::Tensor fc_b,
                               double eps) {
    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);
    const int K = (int)conv_w.size(0);
    const int O = (int)fc_w.size(0);
    auto opts = x.options();
    auto means = torch::empty({N, C}, opts);
    auto hidden = torch::empty({N, K}, opts);
    auto out = torch::empty({N, O}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    spatial_mean_kernel<<<N * C, 128, 0, stream>>>(x.data_ptr<float>(), means.data_ptr<float>(), N, C, H, W);
    hidden_kernel<<<N * K, 256, 0, stream>>>(means.data_ptr<float>(), conv_w.data_ptr<float>(),
                                             bn_w.data_ptr<float>(), bn_b.data_ptr<float>(),
                                             bn_rm.data_ptr<float>(), bn_rv.data_ptr<float>(),
                                             hidden.data_ptr<float>(), N, C, K, (float)eps);
    fc_kernel<<<N * O, 256, 0, stream>>>(hidden.data_ptr<float>(), fc_w.data_ptr<float>(),
                                         fc_b.data_ptr<float>(), out.data_ptr<float>(), N, K, O);
    return out;
}
'''

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor eff_tail_forward(torch::Tensor x,
                               torch::Tensor conv_w,
                               torch::Tensor bn_w,
                               torch::Tensor bn_b,
                               torch::Tensor bn_rm,
                               torch::Tensor bn_rv,
                               torch::Tensor fc_w,
                               torch::Tensor fc_b,
                               double eps);
'''

_eff_ext = load_inline(
    name='effnet_b1_tail_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['eff_tail_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    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)

    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 forward(self, x):
        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)
        # In inference mode, commute the final 1x1 convolution with global average pooling
        # and fuse conv2 + bn2 + relu + avgpool + linear in custom CUDA kernels.
        if x.is_cuda and x.dtype is torch.float32 and (not self.training):
            return _eff_ext.eff_tail_forward(
                x.contiguous(), self.conv2.weight.contiguous().view(1280, 320),
                self.bn2.weight.contiguous(), self.bn2.bias.contiguous(),
                self.bn2.running_mean.contiguous(), self.bn2.running_var.contiguous(),
                self.fc.weight.contiguous(), self.fc.bias.contiguous(), float(self.bn2.eps))
        else:
            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