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

class Model(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        """
        MBConv block implementation.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param kernel_size: Kernel size for the depthwise convolution.
        :param stride: Stride for the depthwise convolution.
        :param expand_ratio: Expansion ratio for the intermediate channels.
        """
        super(Model, self).__init__()
        
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = in_channels * expand_ratio
        
        if expand_ratio != 1:
            self.expand_conv = 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)
            )
        
        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True)
        )
        
        self.project_conv = nn.Sequential(
            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 MBConv block.

        :param x: The input tensor, shape (batch_size, in_channels, H, W)
        :return: The output tensor, shape (batch_size, out_channels, H', W')
        """
        identity = x
        
        if hasattr(self, 'expand_conv'):
            x = self.expand_conv(x)
        
        x = self.depthwise_conv(x)
        x = self.project_conv(x)
        
        if self.use_residual:
            x += identity
        
        return x

# Test code
batch_size = 10
in_channels = 112
out_channels = 192
kernel_size = 5
stride = 2
expand_ratio = 6

def get_inputs():
    return [torch.rand(batch_size, in_channels, 224, 224)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, expand_ratio]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

void relu6_cuda(torch::Tensor x);
void add_cuda(torch::Tensor y, torch::Tensor x);

torch::Tensor mbconv_forward(torch::Tensor x,
                            torch::Tensor ew, torch::Tensor eb,
                            torch::Tensor dw, torch::Tensor db,
                            torch::Tensor pw, torch::Tensor pb,
                            bool has_expand,
                            int64_t stride, int64_t padding,
                            int64_t groups,
                            bool residual) {
    torch::Tensor y;
    if (has_expand) {
        y = at::conv2d(x, ew, eb, {1, 1}, {0, 0}, {1, 1}, 1);
        relu6_cuda(y);
    } else {
        y = x;
    }

    auto z = at::conv2d(y, dw, db, {stride, stride}, {padding, padding}, {1, 1}, groups);
    relu6_cuda(z);

    auto out = at::conv2d(z, pw, pb, {1, 1}, {0, 0}, {1, 1}, 1);
    if (residual) {
        add_cuda(out, x);
    }
    return out;
}
'''

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

__global__ void relu6_kernel(float* __restrict__ p, int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t step = (int64_t)blockDim.x * gridDim.x;
    for (; i < n; i += step) {
        float v = p[i];
        v = v < 0.0f ? 0.0f : v;
        v = v > 6.0f ? 6.0f : v;
        p[i] = v;
    }
}

__global__ void add_kernel(float* __restrict__ y, const float* __restrict__ x, int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t step = (int64_t)blockDim.x * gridDim.x;
    for (; i < n; i += step) {
        y[i] += x[i];
    }
}

void relu6_cuda(torch::Tensor x) {
    const int64_t n = x.numel();
    if (n == 0) return;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 4096) blocks = 4096;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    relu6_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void add_cuda(torch::Tensor y, torch::Tensor x) {
    const int64_t n = y.numel();
    if (n == 0) return;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 4096) blocks = 4096;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    add_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), x.data_ptr<float>(), n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name='mbconv_folded_bn_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['mbconv_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)


def _fold_conv_bn(conv, bn):
    # Evaluation-mode BatchNorm folded into convolution weight and bias.
    w = conv.weight
    scale = bn.weight * torch.rsqrt(bn.running_var + bn.eps)
    wf = w * scale.reshape(-1, 1, 1, 1)
    bf = bn.bias - bn.running_mean * scale
    return wf.contiguous(), bf.contiguous()


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
        super(ModelNew, self).__init__()
        self.use_residual = (stride == 1 and in_channels == out_channels)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.expand_ratio = expand_ratio
        hidden_dim = in_channels * expand_ratio
        self.hidden_dim = hidden_dim

        if expand_ratio != 1:
            self.expand_conv = 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)
            )

        self.depthwise_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=kernel_size, stride=stride,
                      padding=(kernel_size - 1) // 2, groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True)
        )

        self.project_conv = nn.Sequential(
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels)
        )
        self._cache = None
        self._cache_device = None

    def load_state_dict(self, *args, **kwargs):
        r = super().load_state_dict(*args, **kwargs)
        self._cache = None
        self._cache_device = None
        return r

    def _prepare(self, x):
        dev = x.device
        if self._cache is not None and self._cache_device == dev:
            return self._cache
        with torch.no_grad():
            if hasattr(self, 'expand_conv'):
                ew, eb = _fold_conv_bn(self.expand_conv[0], self.expand_conv[1])
                has_expand = True
            else:
                # Unused placeholders for the no-expand case.
                ew = torch.empty(0, device=dev, dtype=x.dtype)
                eb = torch.empty(0, device=dev, dtype=x.dtype)
                has_expand = False
            dw, db = _fold_conv_bn(self.depthwise_conv[0], self.depthwise_conv[1])
            pw, pb = _fold_conv_bn(self.project_conv[0], self.project_conv[1])
        self._cache = (ew, eb, dw, db, pw, pb, has_expand)
        self._cache_device = dev
        return self._cache

    def forward(self, x):
        ew, eb, dw, db, pw, pb, has_expand = self._prepare(x)
        return _ext.mbconv_forward(
            x.contiguous(), ew, eb, dw, db, pw, pb,
            bool(has_expand), int(self.stride), int((self.kernel_size - 1) // 2),
            int(self.hidden_dim), bool(self.use_residual)
        )