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

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

torch::Tensor first_conv_bn_relu6_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor dw_conv_bn_relu6_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, int64_t stride, double eps);
torch::Tensor avgpool2d_forward(torch::Tensor x);
"""

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

static inline int div_up_int64(int64_t a, int b) { return (int)((a + b - 1) / b); }

__global__ void first_conv_bn_relu6_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mean,
    const float* __restrict__ var,
    float* __restrict__ y,
    int N, int H, int W, int OH, int OW, float eps)
{
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)N * 32 * OH * OW;
    if (idx >= total) return;

    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int oc = t % 32;
    int n = t / 32;

    int ih0 = oh * 2 - 1;
    int iw0 = ow * 2 - 1;
    float acc = 0.0f;
    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            int ih = ih0 + kh;
            if ((unsigned)ih >= (unsigned)H) continue;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                int iw = iw0 + kw;
                if ((unsigned)iw < (unsigned)W) {
                    float xv = x[((n * 3 + ic) * H + ih) * W + iw];
                    float wv = w[((oc * 3 + ic) * 3 + kh) * 3 + kw];
                    acc = fmaf(xv, wv, acc);
                }
            }
        }
    }
    float scale = gamma[oc] * rsqrtf(var[oc] + eps);
    float out = fmaf(acc - mean[oc], scale, beta[oc]);
    out = fminf(fmaxf(out, 0.0f), 6.0f);
    y[idx] = out;
}

__global__ void dw_conv_bn_relu6_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mean,
    const float* __restrict__ var,
    float* __restrict__ y,
    int N, int C, int H, int W, int OH, int OW, int stride, float eps)
{
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)N * C * OH * OW;
    if (idx >= total) return;

    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int c = t % C;
    int n = t / C;

    int ih0 = oh * stride - 1;
    int iw0 = ow * stride - 1;
    float acc = 0.0f;
    const float* wp = w + c * 9;
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
        int ih = ih0 + kh;
        if ((unsigned)ih >= (unsigned)H) continue;
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
            int iw = iw0 + kw;
            if ((unsigned)iw < (unsigned)W) {
                float xv = x[((n * C + c) * H + ih) * W + iw];
                float wv = wp[kh * 3 + kw];
                acc = fmaf(xv, wv, acc);
            }
        }
    }
    float scale = gamma[c] * rsqrtf(var[c] + eps);
    float out = fmaf(acc - mean[c], scale, beta[c]);
    out = fminf(fmaxf(out, 0.0f), 6.0f);
    y[idx] = out;
}

__global__ void avgpool2d_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;
    const float* xp = x + ((n * C + c) * H) * W;
    float s = 0.0f;
    int HW = H * W;
    for (int i = 0; i < HW; ++i) s += xp[i];
    y[idx] = s / (float)HW;
}

torch::Tensor first_conv_bn_relu6_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, double eps) {
    x = x.contiguous();
    w = w.contiguous();
    int N = (int)x.size(0);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int OH = (H + 1) / 2;
    int OW = (W + 1) / 2;
    auto y = torch::empty({N, 32, OH, OW}, x.options());
    int threads = 256;
    int64_t total = (int64_t)N * 32 * OH * OW;
    first_conv_bn_relu6_kernel<<<div_up_int64(total, threads), threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), y.data_ptr<float>(), N, H, W, OH, OW, (float)eps);
    return y;
}

torch::Tensor dw_conv_bn_relu6_forward(torch::Tensor x, torch::Tensor w, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, int64_t stride, double eps) {
    x = x.contiguous();
    w = w.contiguous();
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int s = (int)stride;
    int OH = (H + 2 - 3) / s + 1;
    int OW = (W + 2 - 3) / s + 1;
    auto y = torch::empty({N, C, OH, OW}, x.options());
    int threads = 256;
    int64_t total = (int64_t)N * C * OH * OW;
    dw_conv_bn_relu6_kernel<<<div_up_int64(total, threads), threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, OH, OW, s, (float)eps);
    return y;
}

torch::Tensor avgpool2d_forward(torch::Tensor x) {
    x = x.contiguous();
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    auto y = torch::empty({N, C}, x.options());
    int threads = 256;
    int total = N * C;
    avgpool2d_kernel<<<div_up_int64(total, threads), threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}
"""

_mobilenet_ext = load_inline(
    name="mobilenetv2_fused_depthwise_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["first_conv_bn_relu6_forward", "dw_conv_bn_relu6_forward", "avgpool2d_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__()
        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)
            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), (stride == 1 and inp == oup)

        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_inner_seq_eval(self, seq, x):
        i = 0
        L = len(seq)
        while i < L:
            m = seq[i]
            # Fuse depthwise 3x3 conv + following BatchNorm + ReLU6 in eval mode.
            if (i + 2 < L and isinstance(m, nn.Conv2d) and m.groups == m.in_channels and
                m.in_channels == m.out_channels and m.kernel_size == (3, 3) and
                m.padding == (1, 1) and isinstance(seq[i+1], nn.BatchNorm2d) and
                isinstance(seq[i+2], nn.ReLU6) and x.is_cuda and x.dtype is torch.float32):
                bn = seq[i+1]
                x = _mobilenet_ext.dw_conv_bn_relu6_forward(
                    x, m.weight, bn.weight, bn.bias, bn.running_mean, bn.running_var,
                    int(m.stride[0]), float(bn.eps))
                i += 3
            else:
                x = m(x)
                i += 1
        return x

    def forward(self, x):
        # Training-mode BatchNorm uses batch statistics; keep the exact PyTorch path there.
        if self.training or (not x.is_cuda) or x.dtype is not torch.float32:
            x = self.features(x)
            x = x.view(x.size(0), -1)
            return self.classifier(x)

        # First conv + BN + ReLU6 fused.
        conv0 = self.features[0]
        bn0 = self.features[1]
        x = _mobilenet_ext.first_conv_bn_relu6_forward(
            x, conv0.weight, bn0.weight, bn0.bias, bn0.running_mean, bn0.running_var, float(bn0.eps))

        # Run all layers except the final adaptive average pool; depthwise conv blocks are fused.
        for j in range(3, len(self.features) - 1):
            layer = self.features[j]
            if isinstance(layer, nn.Sequential):
                x = self._forward_inner_seq_eval(layer, x)
            else:
                x = layer(x)

        x = _mobilenet_ext.avgpool2d_forward(x)
        # Dropout probability is zero in the reference, so this is exactly the classifier's Linear.
        lin = self.classifier[1]
        return F.linear(x, lin.weight, lin.bias)