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

class Model(nn.Module):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1):
        """
        :param in_channels: Number of input channels
        :param out_channels: Number of output channels
        :param stride: Stride for the first convolutional layer
        :param downsample: Downsample layer for the shortcut connection
        """
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = nn.Sequential(
            nn.Conv2d(in_channels, out_channels * self.expansion, kernel_size=1, stride=stride, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion),
        )
        self.stride = stride

    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels, height, width)
        """
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out
    
# Test code
in_channels = 3
out_channels = 64
stride = 1
batch_size = 10
num_classes = 1000

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

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

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

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

void bn_relu_cuda(torch::Tensor inp, torch::Tensor out, torch::Tensor w, torch::Tensor b,
                  torch::Tensor rm, torch::Tensor rv, double eps);
void bn_add_relu_cuda(torch::Tensor conv, torch::Tensor ident, torch::Tensor out,
                      torch::Tensor w2, torch::Tensor b2, torch::Tensor rm2, torch::Tensor rv2,
                      torch::Tensor wi, torch::Tensor bi, torch::Tensor rmi, torch::Tensor rvi,
                      double eps);
void add_relu_inplace_cuda(torch::Tensor a, torch::Tensor b);

torch::Tensor resblock_forward(torch::Tensor x,
    torch::Tensor conv1_w, torch::Tensor bn1_w, torch::Tensor bn1_b, torch::Tensor bn1_rm, torch::Tensor bn1_rv,
    torch::Tensor conv2_w, torch::Tensor bn2_w, torch::Tensor bn2_b, torch::Tensor bn2_rm, torch::Tensor bn2_rv,
    torch::Tensor down_w, torch::Tensor down_bn_w, torch::Tensor down_bn_b, torch::Tensor down_bn_rm, torch::Tensor down_bn_rv,
    double eps, double momentum, bool training) {
    std::optional<torch::Tensor> bias = std::nullopt;
    std::vector<int64_t> s = {1,1}, p1 = {1,1}, p0 = {0,0};

    auto out1 = at::conv2d(x, conv1_w, bias, s, p1, s, 1);
    if (training) {
        auto b1 = at::batch_norm(out1, bn1_w, bn1_b, bn1_rm, bn1_rv, true, momentum, eps, true);
        at::relu_(b1);
        auto out2 = at::conv2d(b1, conv2_w, bias, s, p1, s, 1);
        auto b2 = at::batch_norm(out2, bn2_w, bn2_b, bn2_rm, bn2_rv, true, momentum, eps, true);
        auto ident = at::conv2d(x, down_w, bias, s, p0, s, 1);
        auto bi = at::batch_norm(ident, down_bn_w, down_bn_b, down_bn_rm, down_bn_rv, true, momentum, eps, true);
        add_relu_inplace_cuda(b2, bi);
        return b2;
    } else {
        auto act1 = torch::empty_like(out1);
        bn_relu_cuda(out1, act1, bn1_w, bn1_b, bn1_rm, bn1_rv, eps);
        auto out2 = at::conv2d(act1, conv2_w, bias, s, p1, s, 1);
        auto ident = at::conv2d(x, down_w, bias, s, p0, s, 1);
        auto out = torch::empty_like(out2);
        bn_add_relu_cuda(out2, ident, out, bn2_w, bn2_b, bn2_rm, bn2_rv,
                         down_bn_w, down_bn_b, down_bn_rm, down_bn_rv, eps);
        return out;
    }
}
'''

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

__global__ void bn_relu_kernel(const float* __restrict__ inp, float* __restrict__ out,
                               const float* __restrict__ w, const float* __restrict__ b,
                               const float* __restrict__ rm, const float* __restrict__ rv,
                               int64_t total, int C, int HW, float eps) {
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (idx / HW) % C;
    float val = (inp[idx] - rm[c]) * (w[c] * rsqrtf(rv[c] + eps)) + b[c];
    out[idx] = val > 0.0f ? val : 0.0f;
}

__global__ void bn_add_relu_kernel(const float* __restrict__ conv, const float* __restrict__ ident,
                                   float* __restrict__ out,
                                   const float* __restrict__ w2, const float* __restrict__ b2,
                                   const float* __restrict__ rm2, const float* __restrict__ rv2,
                                   const float* __restrict__ wi, const float* __restrict__ bi,
                                   const float* __restrict__ rmi, const float* __restrict__ rvi,
                                   int64_t total, int C, int HW, float eps) {
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int c = (idx / HW) % C;
    float v2 = (conv[idx] - rm2[c]) * (w2[c] * rsqrtf(rv2[c] + eps)) + b2[c];
    float vi = (ident[idx] - rmi[c]) * (wi[c] * rsqrtf(rvi[c] + eps)) + bi[c];
    float val = v2 + vi;
    out[idx] = val > 0.0f ? val : 0.0f;
}

__global__ void add_relu_inplace_kernel(float* __restrict__ a, const float* __restrict__ b, int64_t total) {
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    float v = a[idx] + b[idx];
    a[idx] = v > 0.0f ? v : 0.0f;
}

void bn_relu_cuda(torch::Tensor inp, torch::Tensor out, torch::Tensor w, torch::Tensor b,
                  torch::Tensor rm, torch::Tensor rv, double eps) {
    int64_t total = inp.numel();
    int C = inp.size(1);
    int HW = inp.size(2) * inp.size(3);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    bn_relu_kernel<<<blocks, threads>>>(inp.data_ptr<float>(), out.data_ptr<float>(),
        w.data_ptr<float>(), b.data_ptr<float>(), rm.data_ptr<float>(), rv.data_ptr<float>(),
        total, C, HW, (float)eps);
}

void bn_add_relu_cuda(torch::Tensor conv, torch::Tensor ident, torch::Tensor out,
                      torch::Tensor w2, torch::Tensor b2, torch::Tensor rm2, torch::Tensor rv2,
                      torch::Tensor wi, torch::Tensor bi, torch::Tensor rmi, torch::Tensor rvi,
                      double eps) {
    int64_t total = conv.numel();
    int C = conv.size(1);
    int HW = conv.size(2) * conv.size(3);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    bn_add_relu_kernel<<<blocks, threads>>>(conv.data_ptr<float>(), ident.data_ptr<float>(), out.data_ptr<float>(),
        w2.data_ptr<float>(), b2.data_ptr<float>(), rm2.data_ptr<float>(), rv2.data_ptr<float>(),
        wi.data_ptr<float>(), bi.data_ptr<float>(), rmi.data_ptr<float>(), rvi.data_ptr<float>(),
        total, C, HW, (float)eps);
}

void add_relu_inplace_cuda(torch::Tensor a, torch::Tensor b) {
    int64_t total = a.numel();
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    add_relu_inplace_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), total);
}
'''

_ext = load_inline(
    name='resblock_ext_final_submit',
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=['resblock_forward'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = nn.Sequential(
            nn.Conv2d(in_channels, out_channels * self.expansion, kernel_size=1, stride=stride, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion),
        )
        self.stride = stride

    def forward(self, x):
        return _ext.resblock_forward(
            x.contiguous(),
            self.conv1.weight, self.bn1.weight, self.bn1.bias, self.bn1.running_mean, self.bn1.running_var,
            self.conv2.weight, self.bn2.weight, self.bn2.bias, self.bn2.running_mean, self.bn2.running_var,
            self.downsample[0].weight, self.downsample[1].weight, self.downsample[1].bias,
            self.downsample[1].running_mean, self.downsample[1].running_var,
            float(self.bn1.eps), float(self.bn1.momentum), bool(self.training)
        )