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 bn_relu_avg_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta);
torch::Tensor linear_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b);
"""

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

__inline__ __device__ float warp_sum(float v) {
    for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
    return v;
}
__inline__ __device__ float block_sum(float v) {
    static __shared__ float sh[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    v = warp_sum(v);
    if (lane == 0) sh[wid] = v;
    __syncthreads();
    v = (threadIdx.x < ((blockDim.x + 31) >> 5)) ? sh[lane] : 0.0f;
    if (wid == 0) v = warp_sum(v);
    return v;
}

__global__ void stats2_kernel(const float* __restrict__ x,
                              float* __restrict__ mean,
                              float* __restrict__ inv,
                              int N, int C, int H, int W) {
    int c = blockIdx.x;
    int K = N * H * W;
    float sum = 0.0f;
    for (int idx = threadIdx.x; idx < K; idx += blockDim.x) {
        int w = idx % W;
        int t = idx / W;
        int h = t % H;
        int n = t / H;
        sum += x[((n * C + c) * H + h) * W + w];
    }
    sum = block_sum(sum);
    __shared__ float m;
    if (threadIdx.x == 0) {
        m = sum / (float)K;
        mean[c] = m;
    }
    __syncthreads();
    float ss = 0.0f;
    for (int idx = threadIdx.x; idx < K; idx += blockDim.x) {
        int w = idx % W;
        int t = idx / W;
        int h = t % H;
        int n = t / H;
        float d = x[((n * C + c) * H + h) * W + w] - m;
        ss += d * d;
    }
    ss = block_sum(ss);
    if (threadIdx.x == 0) inv[c] = rsqrtf(ss / (float)K + 1.0e-5f);
}

__global__ void avg_apply_kernel(const float* __restrict__ x,
                                 float* __restrict__ y,
                                 const float* __restrict__ mean,
                                 const float* __restrict__ inv,
                                 const float* __restrict__ gamma,
                                 const float* __restrict__ beta,
                                 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* base = x + (n * C + c) * H * W;
    float m = mean[c];
    float scale = inv[c] * gamma[c];
    float be = beta[c];
    float sum = 0.0f;
    int HW = H * W;
    for (int i = 0; i < HW; ++i) {
        float v = (base[i] - m) * scale + be;
        v = fminf(fmaxf(v, 0.0f), 6.0f);
        sum += v;
    }
    y[idx] = sum / (float)HW;
}

torch::Tensor bn_relu_avg_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta) {
    int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    auto y = torch::empty({N, C}, x.options());
    auto mean = torch::empty({C}, x.options());
    auto inv = torch::empty({C}, x.options());
    stats2_kernel<<<C, 256>>>(x.data_ptr<float>(), mean.data_ptr<float>(), inv.data_ptr<float>(), N, C, H, W);
    int threads = 256;
    int blocks = (N * C + threads - 1) / threads;
    avg_apply_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), mean.data_ptr<float>(), inv.data_ptr<float>(),
                                          gamma.data_ptr<float>(), beta.data_ptr<float>(), N, C, H, W);
    return y;
}

__global__ void linear_kernel(const float* __restrict__ x,
                              const float* __restrict__ w,
                              const float* __restrict__ b,
                              float* __restrict__ y,
                              int N, int C, int O) {
    int bid = blockIdx.x;
    int o = bid % O;
    int n = bid / O;
    float sum = 0.0f;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        sum += x[n * C + c] * w[o * C + c];
    }
    sum = block_sum(sum);
    if (threadIdx.x == 0) y[n * O + o] = sum + b[o];
}

torch::Tensor linear_forward(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    int N = x.size(0), C = x.size(1), O = w.size(0);
    auto y = torch::empty({N, O}, x.options());
    linear_kernel<<<N * O, 256>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), y.data_ptr<float>(), N, C, O);
    return y;
}
"""

_ext = load_inline(name="mobilenetv2_tail_fused_ext", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["bn_relu_avg_forward", "linear_forward"], with_cuda=True,
                   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)
            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)
        self._tail_start = len(self.features) - 4
    def forward(self, x):
        for i, layer in enumerate(self.features):
            if i >= self._tail_start:
                break
            x = layer(x)
        x = self.features[self._tail_start](x)
        bn = self.features[self._tail_start + 1]
        x = _ext.bn_relu_avg_forward(x, bn.weight, bn.bias)
        lin = self.classifier[1]
        return _ext.linear_forward(x, lin.weight, lin.bias)