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

class Bottleneck(nn.Module):
    expansion = 4

    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(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        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 * expansion, height, width)
        """
        identity = x

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

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

        out = self.conv3(out)
        out = self.bn3(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, layers, num_classes=1000):
        """
        :param block: Type of block to use (BasicBlock or Bottleneck)
        :param layers: List of integers specifying the number of blocks in each layer
        :param num_classes: Number of output classes
        """
        super(Model, self).__init__()
        self.in_channels = 64

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

        block = Bottleneck

        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.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 = 10
height = 224
width = 224
layers = [3, 4, 23, 3]
num_classes = 1000

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

def get_init_inputs():
    return [layers, num_classes]

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

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

torch::Tensor add_relu_inplace(torch::Tensor a, torch::Tensor b);
torch::Tensor relu_maxpool2d_3x3s2p1(torch::Tensor x);
torch::Tensor global_avgpool_flat(torch::Tensor x);
'''

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

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

torch::Tensor add_relu_inplace(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "add_relu_inplace: CUDA tensors expected");
    TORCH_CHECK(a.scalar_type() == torch::kFloat && b.scalar_type() == torch::kFloat, "add_relu_inplace: float32 only");
    TORCH_CHECK(a.numel() == b.numel(), "add_relu_inplace: size mismatch");
    const long long n = a.numel();
    if (n == 0) return a;
    const int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    add_relu_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), n);
    return a;
}

__global__ void relu_maxpool_kernel(const float* __restrict__ x, float* __restrict__ y,
                                    int N, int C, int H, int W, int OH, int OW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * OH * OW;
    int stride = blockDim.x * gridDim.x;
    for (int t = idx; t < total; t += stride) {
        int ow = t % OW;
        int q = t / OW;
        int oh = q % OH;
        q /= OH;
        int c = q % C;
        int n = q / C;
        int hs = oh * 2 - 1;
        int ws = ow * 2 - 1;
        float m = -FLT_MAX;
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            int ih = hs + kh;
            if ((unsigned)ih < (unsigned)H) {
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    int iw = ws + kw;
                    if ((unsigned)iw < (unsigned)W) {
                        float v = x[((n * C + c) * H + ih) * W + iw];
                        m = v > m ? v : m;
                    }
                }
            }
        }
        y[t] = m > 0.0f ? m : 0.0f;
    }
}

torch::Tensor relu_maxpool2d_3x3s2p1(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "relu_maxpool2d_3x3s2p1: CUDA tensor expected");
    TORCH_CHECK(x.scalar_type() == torch::kFloat, "relu_maxpool2d_3x3s2p1: float32 only");
    TORCH_CHECK(x.dim() == 4, "relu_maxpool2d_3x3s2p1: 4D tensor expected");
    x = x.contiguous();
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    int OH = (H + 1) / 2;
    int OW = (W + 1) / 2;
    auto y = torch::empty({N, C, OH, OW}, x.options());
    int total = N * C * OH * OW;
    if (total == 0) return y;
    const int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 65535) blocks = 65535;
    relu_maxpool_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, OH, OW);
    return y;
}

__global__ void global_avgpool_kernel(const float* __restrict__ x, float* __restrict__ y,
                                      int N, int C, int H, int W) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C;
    int HW = H * W;
    int stride = blockDim.x * gridDim.x;
    for (int t = idx; t < total; t += stride) {
        int c = t % C;
        int n = t / C;
        const float* base = x + ((n * C + c) * H * W);
        float s = 0.0f;
        for (int i = 0; i < HW; ++i) s += base[i];
        y[t] = s / (float)HW;
    }
}

torch::Tensor global_avgpool_flat(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "global_avgpool_flat: CUDA tensor expected");
    TORCH_CHECK(x.scalar_type() == torch::kFloat, "global_avgpool_flat: float32 only");
    TORCH_CHECK(x.dim() == 4, "global_avgpool_flat: 4D tensor expected");
    x = x.contiguous();
    int N = (int)x.size(0), C = (int)x.size(1), H = (int)x.size(2), W = (int)x.size(3);
    auto y = torch::empty({N, C}, x.options());
    int total = N * C;
    if (total == 0) return y;
    const int threads = 256;
    int blocks = (total + threads - 1) / threads;
    if (blocks > 65535) blocks = 65535;
    global_avgpool_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W);
    return y;
}
'''

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


class _BottleneckNew(nn.Module):
    expansion = 4

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

    def forward(self, x):
        identity = x

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

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

        out = self.conv3(out)
        out = self.bn3(out)

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

        return _ext.add_relu_inplace(out, identity)


class ModelNew(nn.Module):
    def __init__(self, layers, num_classes=1000):
        super(ModelNew, self).__init__()
        self.in_channels = 64

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

        block = _BottleneckNew
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.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):
        x = self.conv1(x)
        x = self.bn1(x)
        x = _ext.relu_maxpool2d_3x3s2p1(x)

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

        x = _ext.global_avgpool_flat(x)
        x = self.fc(x)
        return x