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

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

        :param num_classes: The number of output classes. Default is 1000.
        """
        super(Model, self).__init__()
        
        def _make_divisible(v, divisor, min_value=None):
            """
            This function ensures that the number of channels is divisible by the divisor.
            """
            if min_value is None:
                min_value = divisor
            new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
            # Make sure that round down does not go down by more than 10%.
            if new_v < 0.9 * v:
                new_v += divisor
            return new_v

        def _inverted_residual_block(inp, oup, stride, expand_ratio):
            """
            Inverted Residual Block for MobileNetV2.
            """
            hidden_dim = int(inp * expand_ratio)
            use_res_connect = stride == 1 and inp == oup

            layers = []
            if expand_ratio != 1:
                # Pointwise convolution
                layers.append(nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False))
                layers.append(nn.BatchNorm2d(hidden_dim))
                layers.append(nn.ReLU6(inplace=True))

            layers.extend([
                # Depthwise convolution
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                # Pointwise linear convolution
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            ])

            if use_res_connect:
                return nn.Sequential(*layers), True
            else:
                return nn.Sequential(*layers), False

        # MobileNetV2 architecture
        input_channel = 32
        last_channel = 1280
        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        # Building first layer
        features = [nn.Conv2d(3, input_channel, 3, 2, 1, bias=False),
                    nn.BatchNorm2d(input_channel),
                    nn.ReLU6(inplace=True)]

        # Building inverted residual blocks
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c, 8)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(_inverted_residual_block(input_channel, output_channel, stride, expand_ratio=t)[0])
                input_channel = output_channel

        # Building last several layers
        features.append(nn.Conv2d(input_channel, last_channel, 1, 1, 0, bias=False))
        features.append(nn.BatchNorm2d(last_channel))
        features.append(nn.ReLU6(inplace=True))

        # Final layer
        features.append(nn.AdaptiveAvgPool2d((1, 1)))

        self.features = nn.Sequential(*features)

        # Linear layer
        self.classifier = nn.Sequential(
            nn.Dropout(0.0),
            nn.Linear(last_channel, num_classes),
        )

        # Weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        """
        Forward pass of the MobileNetV2 model.

        :param x: The input tensor, shape (batch_size, 3, 224, 224)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        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
from torch.utils.cpp_extension import load_inline

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

torch::Tensor avg_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor bn_relu_avg_linear_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                                      torch::Tensor mean, torch::Tensor var,
                                      torch::Tensor weight, torch::Tensor bias, double eps);

torch::Tensor avg_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    return avg_linear_cuda(x, weight, bias);
}

torch::Tensor bn_relu_avg_linear(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                                 torch::Tensor mean, torch::Tensor var,
                                 torch::Tensor weight, torch::Tensor bias, double eps) {
    return bn_relu_avg_linear_cuda(x, gamma, beta, mean, var, weight, bias, eps);
}
'''

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

__global__ void avgpool_kernel(const float* __restrict__ x, float* __restrict__ avg,
                               int B, int C, int H, int W) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C;
    if (idx >= total) return;
    int c = idx % C;
    int b = idx / C;
    const float* p = x + ((b * C + c) * H * W);
    int HW = H * W;
    float s = 0.0f;
    #pragma unroll
    for (int i = 0; i < 49; ++i) {
        if (i < HW) s += p[i];
    }
    if (HW != 49) {
        s = 0.0f;
        for (int i = 0; i < HW; ++i) s += p[i];
    }
    avg[idx] = s / (float)HW;
}

__global__ void bn_relu6_avgpool_kernel(const float* __restrict__ x,
                                        const float* __restrict__ gamma,
                                        const float* __restrict__ beta,
                                        const float* __restrict__ mean,
                                        const float* __restrict__ var,
                                        float* __restrict__ avg,
                                        int B, int C, int H, int W, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C;
    if (idx >= total) return;
    int c = idx % C;
    int b = idx / C;
    float scale = gamma[c] * rsqrtf(var[c] + eps);
    float offs = beta[c] - mean[c] * scale;
    const float* p = x + ((b * C + c) * H * W);
    int HW = H * W;
    float s = 0.0f;
    #pragma unroll
    for (int i = 0; i < 49; ++i) {
        if (i < HW) {
            float v = fmaf(p[i], scale, offs);
            v = fminf(fmaxf(v, 0.0f), 6.0f);
            s += v;
        }
    }
    if (HW != 49) {
        s = 0.0f;
        for (int i = 0; i < HW; ++i) {
            float v = fmaf(p[i], scale, offs);
            v = fminf(fmaxf(v, 0.0f), 6.0f);
            s += v;
        }
    }
    avg[idx] = s / (float)HW;
}

template<int WARPS_PER_BLOCK>
__global__ void linear_warp_kernel(const float* __restrict__ avg,
                                   const float* __restrict__ weight,
                                   const float* __restrict__ bias,
                                   float* __restrict__ out,
                                   int B, int C, int K) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int ow = blockIdx.x * WARPS_PER_BLOCK + warp;
    int total = B * K;
    if (ow >= total) return;
    int k = ow % K;
    int b = ow / K;
    const float* a = avg + b * C;
    const float* w = weight + k * C;
    float sum = 0.0f;
    for (int c = lane; c < C; c += 32) {
        sum += a[c] * w[c];
    }
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
    }
    if (lane == 0) out[ow] = sum + bias[k];
}

static inline torch::Tensor linear_from_avg(torch::Tensor avg, torch::Tensor weight, torch::Tensor bias) {
    int B = (int)avg.size(0);
    int C = (int)avg.size(1);
    int K = (int)weight.size(0);
    auto out = torch::empty({B, K}, avg.options());
    constexpr int WPB = 4;
    int threads = WPB * 32;
    int blocks = (B * K + WPB - 1) / WPB;
    linear_warp_kernel<WPB><<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        avg.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), B, C, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

torch::Tensor avg_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    c10::cuda::CUDAGuard guard(x.device());
    x = x.contiguous();
    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    auto avg = torch::empty({B, C}, x.options());
    int threads = 256;
    int blocks = (B * C + threads - 1) / threads;
    avgpool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), avg.data_ptr<float>(), B, C, H, W);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return linear_from_avg(avg, weight.contiguous(), bias.contiguous());
}

torch::Tensor bn_relu_avg_linear_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                                      torch::Tensor mean, torch::Tensor var,
                                      torch::Tensor weight, torch::Tensor bias, double eps) {
    c10::cuda::CUDAGuard guard(x.device());
    x = x.contiguous();
    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    auto avg = torch::empty({B, C}, x.options());
    int threads = 256;
    int blocks = (B * C + threads - 1) / threads;
    bn_relu6_avgpool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(),
        avg.data_ptr<float>(), B, C, H, W, (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return linear_from_avg(avg, weight.contiguous(), bias.contiguous());
}
'''

_ext = load_inline(
    name="mobilenetv2_tail_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["avg_linear", "bn_relu_avg_linear"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super(ModelNew, self).__init__()
        
        def _make_divisible(v, divisor, min_value=None):
            if min_value is None:
                min_value = divisor
            new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
            if new_v < 0.9 * v:
                new_v += divisor
            return new_v

        def _inverted_residual_block(inp, oup, stride, expand_ratio):
            hidden_dim = int(inp * expand_ratio)
            use_res_connect = stride == 1 and inp == oup
            layers = []
            if expand_ratio != 1:
                layers.append(nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False))
                layers.append(nn.BatchNorm2d(hidden_dim))
                layers.append(nn.ReLU6(inplace=True))
            layers.extend([
                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
            ])
            return nn.Sequential(*layers), use_res_connect

        input_channel = 32
        last_channel = 1280
        inverted_residual_setting = [
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        features = [nn.Conv2d(3, input_channel, 3, 2, 1, bias=False),
                    nn.BatchNorm2d(input_channel),
                    nn.ReLU6(inplace=True)]
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c, 8)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(_inverted_residual_block(input_channel, output_channel, stride, expand_ratio=t)[0])
                input_channel = output_channel

        features.append(nn.Conv2d(input_channel, last_channel, 1, 1, 0, bias=False))
        features.append(nn.BatchNorm2d(last_channel))
        features.append(nn.ReLU6(inplace=True))
        features.append(nn.AdaptiveAvgPool2d((1, 1)))
        self.features = nn.Sequential(*features)
        self.classifier = nn.Sequential(
            nn.Dropout(0.0),
            nn.Linear(last_channel, num_classes),
        )

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        lin = self.classifier[1]
        if not self.training:
            # Run through the final pointwise convolution, then fuse final BN+ReLU6+global average pool+linear.
            # The last three feature modules are BatchNorm2d, ReLU6, AdaptiveAvgPool2d.
            for i in range(len(self.features) - 3):
                x = self.features[i](x)
            bn = self.features[-3]
            return _ext.bn_relu_avg_linear(x, bn.weight, bn.bias, bn.running_mean, bn.running_var,
                                           lin.weight, lin.bias, float(bn.eps))
        else:
            # Training-mode BatchNorm semantics (batch statistics and running-stat updates) are preserved by PyTorch;
            # only the 1x1 average/linear tail is replaced.
            x = self.features(x)
            return _ext.avg_linear(x, lin.weight, lin.bias)