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

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        """
        :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(BasicBlock, 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 = downsample
        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

class Model(nn.Module):
    def __init__(self, num_classes=1000):
        """
        :param num_classes: Number of output classes
        """
        super(Model, self).__init__()
        self.in_channels = 64

        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(BasicBlock, 64, 2, stride=1)
        self.layer2 = self._make_layer(BasicBlock, 128, 2, stride=2)
        self.layer3 = self._make_layer(BasicBlock, 256, 2, stride=2)
        self.layer4 = self._make_layer(BasicBlock, 512, 2, stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * BasicBlock.expansion, num_classes)

    def _make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if stride != 1 or self.in_channels != out_channels * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * block.expansion),
            )

        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, out_channels))

        return nn.Sequential(*layers)

    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, 3, height, width)
        :return: Output tensor, shape (batch_size, num_classes)
        """
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

# Test code
batch_size = 2
num_classes = 1000
input_shape = (batch_size, 3, 224, 224)

def get_inputs():
    return [torch.rand(input_shape)]

def get_init_inputs():
    return [num_classes]

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 add_relu_inplace_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor add_relu_inplace(torch::Tensor a, torch::Tensor b){return add_relu_inplace_cuda(a,b);}
'''

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

__global__ void add_relu_kernel(float* __restrict__ a, const float* __restrict__ b, long n) {
    long i = blockIdx.x * blockDim.x + threadIdx.x;
    long st = (long)blockDim.x * gridDim.x;
    for (; i < n; i += st) {
        float v = a[i] + b[i];
        a[i] = v > 0.0f ? v : 0.0f;
    }
}

torch::Tensor add_relu_inplace_cuda(torch::Tensor a, torch::Tensor b) {
    long n = a.numel();
    int blocks = (int)((n + 255) / 256);
    if (blocks > 4096) blocks = 4096;
    add_relu_kernel<<<blocks, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        a.data_ptr<float>(), b.data_ptr<float>(), n);
    return a;
}
'''

_ext = load_inline(
    name='resnet18_fold_addrelu_ext',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['add_relu_inplace'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3']
)

class BasicBlockNew(nn.Module):
    expansion = 1
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super().__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 = downsample
        self.stride = stride

class ModelNew(nn.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.in_channels = 64
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(BasicBlockNew, 64, 2, stride=1)
        self.layer2 = self._make_layer(BasicBlockNew, 128, 2, stride=2)
        self.layer3 = self._make_layer(BasicBlockNew, 256, 2, stride=2)
        self.layer4 = self._make_layer(BasicBlockNew, 512, 2, stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * BasicBlockNew.expansion, num_classes)
        self._fused = None

    def _make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if stride != 1 or self.in_channels != out_channels * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * block.expansion),
            )
        layers = [block(self.in_channels, out_channels, stride, downsample)]
        self.in_channels = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, out_channels))
        return nn.Sequential(*layers)

    def _fuse_conv_bn(self, conv, bn):
        scale = bn.weight * torch.rsqrt(bn.running_var + bn.eps)
        w = (conv.weight * scale.reshape(-1, 1, 1, 1)).contiguous()
        b = (bn.bias - bn.running_mean * scale).contiguous()
        return w, b

    def _build_fused(self):
        f = {}
        with torch.no_grad():
            f['conv1'] = self._fuse_conv_bn(self.conv1, self.bn1)
            for lname in ['layer1', 'layer2', 'layer3', 'layer4']:
                for i, blk in enumerate(getattr(self, lname)):
                    k = f'{lname}.{i}'
                    f[k + '.conv1'] = self._fuse_conv_bn(blk.conv1, blk.bn1)
                    f[k + '.conv2'] = self._fuse_conv_bn(blk.conv2, blk.bn2)
                    if blk.downsample is not None:
                        f[k + '.down'] = self._fuse_conv_bn(blk.downsample[0], blk.downsample[1])
        self._fused = f

    def _training_forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        for layer in [self.layer1, self.layer2, self.layer3, self.layer4]:
            for blk in layer:
                identity = x
                out = blk.conv1(x)
                out = blk.bn1(out)
                out = self.relu(out)
                out = blk.conv2(out)
                out = blk.bn2(out)
                if blk.downsample is not None:
                    identity = blk.downsample(x)
                x = self.relu(out + identity)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

    def _block_fused(self, x, blk, key):
        identity = x
        w, b = self._fused[key + '.conv1']
        out = F.conv2d(x, w, b, stride=blk.stride, padding=1)
        out = torch.relu_(out)
        w, b = self._fused[key + '.conv2']
        out = F.conv2d(out, w, b, stride=1, padding=1)
        if blk.downsample is not None:
            w, b = self._fused[key + '.down']
            identity = F.conv2d(x, w, b, stride=blk.stride, padding=0)
        return _ext.add_relu_inplace(out, identity)

    def forward(self, x):
        if self.training:
            return self._training_forward(x)
        if self._fused is None:
            self._build_fused()
        x = x.contiguous()
        w, b = self._fused['conv1']
        x = F.conv2d(x, w, b, stride=2, padding=3)
        x = torch.relu_(x)
        x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
        for lname in ['layer1', 'layer2', 'layer3', 'layer4']:
            for i, blk in enumerate(getattr(self, lname)):
                x = self._block_fused(x, blk, f'{lname}.{i}')
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = torch.flatten(x, 1)
        return self.fc(x)