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_sources = r'''
#include <torch/extension.h>
torch::Tensor mbconv_forward_cuda(
    torch::Tensor x,
    torch::Tensor ew, torch::Tensor ebnw, torch::Tensor ebnb, torch::Tensor erm, torch::Tensor erv, double eeps,
    torch::Tensor dw,
    torch::Tensor dbnw, torch::Tensor dbnb, torch::Tensor drm, torch::Tensor drv, double deps,
    torch::Tensor pw, torch::Tensor pbnw, torch::Tensor pbnb, torch::Tensor prm, torch::Tensor prv, double peps,
    bool residual, bool training);
torch::Tensor mbconv_forward(
    torch::Tensor x,
    torch::Tensor ew, torch::Tensor ebnw, torch::Tensor ebnb, torch::Tensor erm, torch::Tensor erv, double eeps,
    torch::Tensor dw,
    torch::Tensor dbnw, torch::Tensor dbnb, torch::Tensor drm, torch::Tensor drv, double deps,
    torch::Tensor pw, torch::Tensor pbnw, torch::Tensor pbnb, torch::Tensor prm, torch::Tensor prv, double peps,
    bool residual, bool training) {
  return mbconv_forward_cuda(x, ew, ebnw, ebnb, erm, erv, eeps, dw, dbnw, dbnb, drm, drv, deps, pw, pbnw, pbnb, prm, prv, peps, residual, training);
}
'''

cuda_sources = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/ops/conv2d.h>
#include <ATen/ops/batch_norm.h>
#include <cuda_runtime.h>
#include <optional>

__global__ void relu6_kernel(float* y, long total) {
  long idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < total) { float v = y[idx]; y[idx] = fminf(fmaxf(v,0.0f),6.0f); }
}
__global__ void affine_relu6_kernel(float* __restrict__ y, const float* __restrict__ gamma,
                                    const float* __restrict__ beta, const float* __restrict__ mean,
                                    const float* __restrict__ var, float eps,
                                    long total, int C, int HW) {
  long idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= total) return;
  int c = (idx / HW) % C;
  float v = y[idx];
  v = (v - mean[c]) * (gamma[c] * rsqrtf(var[c] + eps)) + beta[c];
  y[idx] = fminf(fmaxf(v, 0.0f), 6.0f);
}
__global__ void affine_kernel(float* __restrict__ y, const float* __restrict__ gamma,
                              const float* __restrict__ beta, const float* __restrict__ mean,
                              const float* __restrict__ var, const float* __restrict__ ident,
                              float eps, long total, int C, int HW, bool residual) {
  long idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= total) return;
  int c = (idx / HW) % C;
  float v = (y[idx] - mean[c]) * (gamma[c] * rsqrtf(var[c] + eps)) + beta[c];
  if (residual) v += ident[idx];
  y[idx] = v;
}
__global__ void add_kernel(float* y, const float* ident, long total) {
  long idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < total) y[idx] += ident[idx];
}
__global__ void depthwise5s2_kernel(const float* __restrict__ x, const float* __restrict__ w,
                                    float* __restrict__ y, long total, int C, int H, int Wd, int OH, int OW) {
  long idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= total) return;
  int ow = idx % OW;
  int oh = (idx / OW) % OH;
  int c = (idx / (OW * OH)) % C;
  int n = idx / (OW * OH * C);
  int ih0 = oh * 2 - 2;
  int iw0 = ow * 2 - 2;
  const float* xbase = x + ((long)n * C + c) * H * Wd;
  const float* wbase = w + (long)c * 25;
  float acc = 0.0f;
#pragma unroll
  for (int kh = 0; kh < 5; ++kh) {
    int ih = ih0 + kh;
    if ((unsigned)ih < (unsigned)H) {
#pragma unroll
      for (int kw = 0; kw < 5; ++kw) {
        int iw = iw0 + kw;
        if ((unsigned)iw < (unsigned)Wd) acc = fmaf(xbase[ih * Wd + iw], wbase[kh * 5 + kw], acc);
      }
    }
  }
  y[idx] = acc;
}

torch::Tensor mbconv_forward_cuda(
    torch::Tensor x,
    torch::Tensor ew, torch::Tensor ebnw, torch::Tensor ebnb, torch::Tensor erm, torch::Tensor erv, double eeps,
    torch::Tensor dw,
    torch::Tensor dbnw, torch::Tensor dbnb, torch::Tensor drm, torch::Tensor drv, double deps,
    torch::Tensor pw, torch::Tensor pbnw, torch::Tensor pbnb, torch::Tensor prm, torch::Tensor prv, double peps,
    bool residual, bool training) {
  x = x.contiguous(); ew = ew.contiguous(); pw = pw.contiguous(); dw = dw.contiguous();
  std::optional<at::Tensor> bias = std::nullopt;
  std::vector<int64_t> s1 = {1,1}, p0 = {0,0}, d1 = {1,1};
  const int threads = 256; cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  auto y = at::conv2d(x, ew, bias, s1, p0, d1, 1);
  if (training) {
    y = at::batch_norm(y, ebnw, ebnb, erm, erv, true, 0.1, eeps, true);
    relu6_kernel<<<(y.numel()+threads-1)/threads, threads, 0, stream>>>(y.template data_ptr<float>(), y.numel());
  } else {
    affine_relu6_kernel<<<(y.numel()+threads-1)/threads, threads, 0, stream>>>(y.template data_ptr<float>(), ebnw.template data_ptr<float>(), ebnb.template data_ptr<float>(), erm.template data_ptr<float>(), erv.template data_ptr<float>(), (float)eeps, y.numel(), y.size(1), y.size(2)*y.size(3));
  }

  int N = y.size(0), C = y.size(1), H = y.size(2), W = y.size(3);
  int OH = (H + 4 - 5) / 2 + 1, OW = (W + 4 - 5) / 2 + 1;
  auto ydw = torch::empty({N, C, OH, OW}, x.options());
  depthwise5s2_kernel<<<(ydw.numel()+threads-1)/threads, threads, 0, stream>>>(y.template data_ptr<float>(), dw.template data_ptr<float>(), ydw.template data_ptr<float>(), ydw.numel(), C, H, W, OH, OW);
  if (training) {
    ydw = at::batch_norm(ydw, dbnw, dbnb, drm, drv, true, 0.1, deps, true);
    relu6_kernel<<<(ydw.numel()+threads-1)/threads, threads, 0, stream>>>(ydw.template data_ptr<float>(), ydw.numel());
  } else {
    affine_relu6_kernel<<<(ydw.numel()+threads-1)/threads, threads, 0, stream>>>(ydw.template data_ptr<float>(), dbnw.template data_ptr<float>(), dbnb.template data_ptr<float>(), drm.template data_ptr<float>(), drv.template data_ptr<float>(), (float)deps, ydw.numel(), C, OH*OW);
  }

  auto out = at::conv2d(ydw, pw, bias, s1, p0, d1, 1);
  if (training) {
    out = at::batch_norm(out, pbnw, pbnb, prm, prv, true, 0.1, peps, true);
    if (residual) add_kernel<<<(out.numel()+threads-1)/threads, threads, 0, stream>>>(out.template data_ptr<float>(), x.template data_ptr<float>(), out.numel());
  } else {
    affine_kernel<<<(out.numel()+threads-1)/threads, threads, 0, stream>>>(out.template data_ptr<float>(), pbnw.template data_ptr<float>(), pbnb.template data_ptr<float>(), prm.template data_ptr<float>(), prv.template data_ptr<float>(), x.template data_ptr<float>(), (float)peps, out.numel(), out.size(1), out.size(2)*out.size(3), residual);
  }
  return out;
}
'''

_ext = load_inline(name='mbconv_fused_ext_final', cpp_sources=cpp_sources, cuda_sources=cuda_sources,
                   functions=['mbconv_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], verbose=False)
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)
        hidden_dim = in_channels * expand_ratio
        if expand_ratio != 1:
            self.expand_conv = nn.Sequential(nn.Conv2d(in_channels, hidden_dim, 1, 1, 0, bias=False), nn.BatchNorm2d(hidden_dim), nn.ReLU6(inplace=True))
        self.depthwise_conv = nn.Sequential(nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, (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, 1, 1, 0, bias=False), nn.BatchNorm2d(out_channels))
    def forward(self, x):
        return _ext.mbconv_forward(x, self.expand_conv[0].weight, self.expand_conv[1].weight, self.expand_conv[1].bias, self.expand_conv[1].running_mean, self.expand_conv[1].running_var, float(self.expand_conv[1].eps), self.depthwise_conv[0].weight, self.depthwise_conv[1].weight, self.depthwise_conv[1].bias, self.depthwise_conv[1].running_mean, self.depthwise_conv[1].running_var, float(self.depthwise_conv[1].eps), self.project_conv[0].weight, self.project_conv[1].weight, self.project_conv[1].bias, self.project_conv[1].running_mean, self.project_conv[1].running_var, float(self.project_conv[1].eps), bool(self.use_residual), bool(self.training))