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

class Model(nn.Module):
    def __init__(self, num_classes=1000, input_channels=3, alpha=1.0):
        """
        MobileNetV1 architecture implementation.

        :param num_classes: The number of output classes (default: 1000)
        :param input_channels: The number of input channels (default: 3 for RGB images)
        :param alpha: Width multiplier (default: 1.0)
        """
        super(Model, self).__init__()
        
        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )
        
        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),
                
                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )
        
        self.model = nn.Sequential(
            conv_bn(input_channels, int(32 * alpha), 2),
            conv_dw(int(32 * alpha), int(64 * alpha), 1),
            conv_dw(int(64 * alpha), int(128 * alpha), 2),
            conv_dw(int(128 * alpha), int(128 * alpha), 1),
            conv_dw(int(128 * alpha), int(256 * alpha), 2),
            conv_dw(int(256 * alpha), int(256 * alpha), 1),
            conv_dw(int(256 * alpha), int(512 * alpha), 2),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(1024 * alpha), 2),
            conv_dw(int(1024 * alpha), int(1024 * alpha), 1),
            nn.AvgPool2d(7),
        )
        self.fc = nn.Linear(int(1024 * alpha), num_classes)
    
    def forward(self, x):
        """
        :param x: The input tensor, shape (batch_size, input_channels, height, width)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        x = self.model(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Test code
batch_size = 10
input_channels = 3
height = 224
width = 224
num_classes = 1000
alpha = 1.0

def get_inputs():
    return [torch.rand(batch_size, input_channels, height, width)]

def get_init_inputs():
    return [num_classes, input_channels, alpha]

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 affine_relu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, torch::Tensor mean, torch::Tensor var, double eps);
torch::Tensor depthwise_conv_bn_relu(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 avgpool_linear(torch::Tensor x, torch::Tensor w, torch::Tensor b);
'''

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

__global__ void affine_relu_kernel(float* __restrict__ x,
                                   const float* __restrict__ gamma,
                                   const float* __restrict__ beta,
                                   const float* __restrict__ mean,
                                   const float* __restrict__ var,
                                   int total, int C, int HW, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (idx / HW) % C;
    float scale = gamma[c] * rsqrtf(var[c] + eps);
    float y = x[idx] * scale + (beta[c] - mean[c] * scale);
    x[idx] = y > 0.0f ? y : 0.0f;
}

torch::Tensor affine_relu(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta,
                          torch::Tensor mean, torch::Tensor var, double eps) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int total = N * C * H * W;
    const int threads = 256;
    affine_relu_kernel<<<(total + threads - 1) / threads, threads>>>(
        x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        mean.data_ptr<float>(), var.data_ptr<float>(), total, C, H * W, (float)eps);
    return x;
}

__global__ void depthwise_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 total, int C, int H, int Wd, int Ho, int Wo,
                                 int stride, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % Wo;
    int t = idx / Wo;
    int oh = t % Ho;
    t /= Ho;
    int c = t % C;
    int n = t / C;

    int ih0 = oh * stride - 1;
    int iw0 = ow * stride - 1;
    const float* xp = x + ((n * C + c) * H) * Wd;
    const float* wp = w + c * 9;
    float acc = 0.0f;
#pragma unroll
    for (int ky = 0; ky < 3; ++ky) {
        int ih = ih0 + ky;
        if ((unsigned)ih < (unsigned)H) {
#pragma unroll
            for (int kx = 0; kx < 3; ++kx) {
                int iw = iw0 + kx;
                if ((unsigned)iw < (unsigned)Wd) {
                    acc += xp[ih * Wd + iw] * wp[ky * 3 + kx];
                }
            }
        }
    }
    float scale = gamma[c] * rsqrtf(var[c] + eps);
    acc = acc * scale + (beta[c] - mean[c] * scale);
    y[idx] = acc > 0.0f ? acc : 0.0f;
}

torch::Tensor depthwise_conv_bn_relu(torch::Tensor x, torch::Tensor w,
                                     torch::Tensor gamma, torch::Tensor beta,
                                     torch::Tensor mean, torch::Tensor var,
                                     int64_t stride, double eps) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), Wd = (int)x.size(3);
    int s = (int)stride;
    int Ho = (H - 1) / s + 1;
    int Wo = (Wd - 1) / s + 1;
    auto y = torch::empty({N, C, Ho, Wo}, x.options());
    int total = N * C * Ho * Wo;
    const int threads = 256;
    depthwise_kernel<<<(total + threads - 1) / threads, threads>>>(
        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>(), total, C, H, Wd, Ho, Wo, s, (float)eps);
    return y;
}

__global__ void avg_kernel(const float* __restrict__ x, float* __restrict__ avg,
                           int NC, int H, int Wd) {
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    int HW = H * Wd;
    float s = 0.0f;
    const float* xp = x + idx * HW;
    for (int i = tid; i < HW; i += blockDim.x) s += xp[i];
    __shared__ float sm[64];
    sm[tid] = s;
    __syncthreads();
    if (tid < 32) sm[tid] += sm[tid + 32];
    __syncthreads();
    if (tid < 16) sm[tid] += sm[tid + 16];
    __syncthreads();
    if (tid < 8) sm[tid] += sm[tid + 8];
    __syncthreads();
    if (tid < 4) sm[tid] += sm[tid + 4];
    __syncthreads();
    if (tid < 2) sm[tid] += sm[tid + 2];
    __syncthreads();
    if (tid == 0) avg[idx] = (sm[0] + sm[1]) / (float)HW;
}

__global__ void linear_kernel(const float* __restrict__ avg,
                              const float* __restrict__ w,
                              const float* __restrict__ b,
                              float* __restrict__ out,
                              int N, int C, int K) {
    int k = blockIdx.x;
    int n = blockIdx.y;
    int tid = threadIdx.x;
    float s = 0.0f;
    const float* ap = avg + n * C;
    const float* wp = w + k * C;
    for (int c = tid; c < C; c += blockDim.x) s += ap[c] * wp[c];
    __shared__ float sm[256];
    sm[tid] = s;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) sm[tid] += sm[tid + off];
        __syncthreads();
    }
    if (tid == 0) out[n * K + k] = sm[0] + b[k];
}

torch::Tensor avgpool_linear(torch::Tensor x, torch::Tensor w, torch::Tensor b) {
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), Wd = (int)x.size(3);
    int K = (int)w.size(0);
    auto avg = torch::empty({N, C}, x.options());
    auto out = torch::empty({N, K}, x.options());
    avg_kernel<<<N * C, 64>>>(x.data_ptr<float>(), avg.data_ptr<float>(), N * C, H, Wd);
    dim3 grid(K, N);
    linear_kernel<<<grid, 256>>>(avg.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), N, C, K);
    return out;
}
'''

_mobilenet_ext = load_inline(
    name='mobilenetv1_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['affine_relu', 'depthwise_conv_bn_relu', 'avgpool_linear'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000, input_channels=3, alpha=1.0):
        super(ModelNew, self).__init__()
        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True)
            )
        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                nn.ReLU(inplace=True),
                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )
        self.model = nn.Sequential(
            conv_bn(input_channels, int(32 * alpha), 2),
            conv_dw(int(32 * alpha), int(64 * alpha), 1),
            conv_dw(int(64 * alpha), int(128 * alpha), 2),
            conv_dw(int(128 * alpha), int(128 * alpha), 1),
            conv_dw(int(128 * alpha), int(256 * alpha), 2),
            conv_dw(int(256 * alpha), int(256 * alpha), 1),
            conv_dw(int(256 * alpha), int(512 * alpha), 2),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(512 * alpha), 1),
            conv_dw(int(512 * alpha), int(1024 * alpha), 2),
            conv_dw(int(1024 * alpha), int(1024 * alpha), 1),
            nn.AvgPool2d(7),
        )
        self.fc = nn.Linear(int(1024 * alpha), num_classes)

    def _eval_forward(self, x):
        x = x.contiguous()
        m0 = self.model[0]
        x = F.conv2d(x, m0[0].weight, None, stride=m0[0].stride, padding=m0[0].padding, dilation=m0[0].dilation, groups=1)
        bn = m0[1]
        x = _mobilenet_ext.affine_relu(x, bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps)
        for i in range(1, 14):
            blk = self.model[i]
            dw = blk[0]; bn1 = blk[1]
            x = _mobilenet_ext.depthwise_conv_bn_relu(x, dw.weight, bn1.weight, bn1.bias, bn1.running_mean, bn1.running_var, dw.stride[0], bn1.eps)
            pw = blk[3]; bn2 = blk[4]
            x = F.conv2d(x, pw.weight, None, stride=1, padding=0, dilation=1, groups=1)
            x = _mobilenet_ext.affine_relu(x, bn2.weight, bn2.bias, bn2.running_mean, bn2.running_var, bn2.eps)
        return _mobilenet_ext.avgpool_linear(x, self.fc.weight, self.fc.bias)

    def forward(self, x):
        if (not self.training) and x.is_cuda and x.dtype is torch.float32:
            return self._eval_forward(x)
        # Training mode must preserve BatchNorm batch-statistics and running-stat updates.
        for i in range(14):
            x = self.model[i](x)
        if x.is_cuda and x.dtype is torch.float32:
            return _mobilenet_ext.avgpool_linear(x.contiguous(), self.fc.weight, self.fc.bias)
        x = self.model[14](x)
        x = x.view(x.size(0), -1)
        return self.fc(x)