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

class DenseBlock(nn.Module):
    def __init__(self, num_layers: int, num_input_features: int, growth_rate: int):
        """
        :param num_layers: The number of layers in the dense block
        :param num_input_features: The number of input feature maps
        :param growth_rate: The growth rate for the dense block (new features added per layer)
        """
        super(DenseBlock, self).__init__()
        layers = []
        for i in range(num_layers):
            layers.append(self._make_layer(num_input_features + i * growth_rate, growth_rate))
        self.layers = nn.ModuleList(layers)

    def _make_layer(self, in_features: int, growth_rate: int):
        """
        Creates a single layer with BatchNorm, ReLU, Conv2D, and Dropout.
        """
        return nn.Sequential(
            nn.BatchNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_features, growth_rate, kernel_size=3, padding=1, bias=False),
            nn.Dropout(0.0)
        )

    def forward(self, x):
        """
        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
        :return: Concatenated output tensor with shape (batch_size, num_output_features, height, width)
        """
        features = [x]
        for layer in self.layers:
            new_feature = layer(x)
            features.append(new_feature)
            x = torch.cat(features, 1)  # Concatenate along channel axis
        return x

class TransitionLayer(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        """
        :param num_input_features: The number of input feature maps
        :param num_output_features: The number of output feature maps
        """
        super(TransitionLayer, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_input_features, num_output_features, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        """
        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
        :return: Downsampled tensor with reduced number of feature maps
        """
        return self.transition(x)

class Model(nn.Module):
    def __init__(self, growth_rate: int = 32, num_classes: int = 1000):
        """
        :param growth_rate: The growth rate of the DenseNet (new features added per layer)
        :param num_classes: The number of output classes for classification
        """
        super(Model, self).__init__()

        # Initial convolution and pooling
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        # Each dense block is followed by a transition layer, except the last one
        num_features = 64
        block_layers = [6, 12, 48, 32]  # Corresponding layers in DenseNet201

        self.dense_blocks = nn.ModuleList()
        self.transition_layers = nn.ModuleList()

        for i, num_layers in enumerate(block_layers):
            block = DenseBlock(num_layers=num_layers, num_input_features=num_features, growth_rate=growth_rate)
            self.dense_blocks.append(block)
            num_features = num_features + num_layers * growth_rate

            if i != len(block_layers) - 1:
                transition = TransitionLayer(num_input_features=num_features, num_output_features=num_features // 2)
                self.transition_layers.append(transition)
                num_features = num_features // 2

        # Final batch norm and classifier
        self.final_bn = nn.BatchNorm2d(num_features)
        self.classifier = nn.Linear(num_features, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        :param x: Input tensor of shape (batch_size, 3, height, width)
        :return: Output tensor of shape (batch_size, num_classes)
        """
        x = self.features(x)

        for i, block in enumerate(self.dense_blocks):
            x = block(x)
            if i != len(self.dense_blocks) - 1:
                x = self.transition_layers[i](x)

        x = self.final_bn(x)
        x = F.relu(x, inplace=True)
        x = F.adaptive_avg_pool2d(x, (1, 1)).view(x.size(0), -1)
        x = self.classifier(x)
        return x

# Testing the DenseNet201 model
batch_size = 10
num_classes = 10
height, width = 224, 224  # Standard input size for DenseNet

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

def get_init_inputs():
    return [32, 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 relu_maxpool2d(torch::Tensor x);
torch::Tensor fused_relu_pool_fc(torch::Tensor x, torch::Tensor fc_w, torch::Tensor fc_b);
'''

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

__global__ void relu_maxpool_kernel(const float* __restrict__ x,
                                    float* __restrict__ y,
                                    int B, int C, int H, int W, int OH, int OW) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * C * OH * OW;
    if (idx >= total) return;

    int ow = idx % OW;
    int t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int c = t % C;
    int b = t / C;

    int ih0 = oh * 2 - 1;
    int iw0 = ow * 2 - 1;
    float m = -FLT_MAX;
    for (int kh = 0; kh < 3; ++kh) {
        int ih = ih0 + kh;
        if ((unsigned)ih < (unsigned)H) {
            for (int kw = 0; kw < 3; ++kw) {
                int iw = iw0 + kw;
                if ((unsigned)iw < (unsigned)W) {
                    float v = x[((b * C + c) * H + ih) * W + iw];
                    m = v > m ? v : m;
                }
            }
        }
    }
    y[idx] = m > 0.0f ? m : 0.0f;
}

torch::Tensor relu_maxpool2d(torch::Tensor x) {
    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int OH = (H + 1) / 2;
    int OW = (W + 1) / 2;
    auto y = torch::empty({B, C, OH, OW}, x.options());
    int total = B * C * OH * OW;
    relu_maxpool_kernel<<<(total + 255) / 256, 256>>>(x.data_ptr<float>(), y.data_ptr<float>(), B, C, H, W, OH, OW);
    return y;
}

__global__ void relu_pool_fc_kernel(const float* __restrict__ x,
                                    const float* __restrict__ fw,
                                    const float* __restrict__ fb,
                                    float* __restrict__ out,
                                    int B, int C, int H, int W, int K) {
    int b = blockIdx.x;
    int k = blockIdx.y;
    int tid = threadIdx.x;
    int HW = H * W;
    float acc = 0.0f;

    for (int c = tid; c < C; c += blockDim.x) {
        const float* xp = x + ((b * C + c) * H) * W;
        float s = 0.0f;
        #pragma unroll
        for (int p = 0; p < 49; ++p) {
            if (p < HW) {
                float v = xp[p];
                s += v > 0.0f ? v : 0.0f;
            }
        }
        acc += (s / (float)HW) * fw[k * C + c];
    }

    __shared__ float sm[256];
    sm[tid] = acc;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) sm[tid] += sm[tid + off];
        __syncthreads();
    }
    if (tid == 0) out[b * K + k] = sm[0] + fb[k];
}

torch::Tensor fused_relu_pool_fc(torch::Tensor x, torch::Tensor fc_w, torch::Tensor fc_b) {
    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int K = (int)fc_w.size(0);
    auto out = torch::empty({B, K}, x.options());
    relu_pool_fc_kernel<<<dim3(B, K), 256>>>(x.data_ptr<float>(), fc_w.data_ptr<float>(), fc_b.data_ptr<float>(),
                                             out.data_ptr<float>(), B, C, H, W, K);
    return out;
}
'''

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

class DenseBlock(nn.Module):
    def __init__(self, num_layers: int, num_input_features: int, growth_rate: int):
        super(DenseBlock, self).__init__()
        self.layers = nn.ModuleList([self._make_layer(num_input_features + i * growth_rate, growth_rate)
                                     for i in range(num_layers)])

    def _make_layer(self, in_features: int, growth_rate: int):
        return nn.Sequential(
            nn.BatchNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_features, growth_rate, kernel_size=3, padding=1, bias=False)
        )

    def forward(self, x):
        for layer in self.layers:
            y = layer[0](x)
            y = layer[1](y)
            y = layer[2](y)
            x = torch.cat((x, y), 1)
        return x

class TransitionLayer(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        super(TransitionLayer, self).__init__()
        self.transition = nn.Sequential(
            nn.BatchNorm2d(num_input_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(num_input_features, num_output_features, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        x = self.transition[0](x)
        x = self.transition[1](x)
        x = self.transition[2](x)
        x = self.transition[3](x)
        return x

class ModelNew(nn.Module):
    def __init__(self, growth_rate: int = 32, num_classes: int = 1000):
        super(ModelNew, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        num_features = 64
        block_layers = [6, 12, 48, 32]
        self.dense_blocks = nn.ModuleList()
        self.transition_layers = nn.ModuleList()

        for i, num_layers in enumerate(block_layers):
            self.dense_blocks.append(DenseBlock(num_layers=num_layers,
                                                num_input_features=num_features,
                                                growth_rate=growth_rate))
            num_features = num_features + num_layers * growth_rate
            if i != len(block_layers) - 1:
                self.transition_layers.append(TransitionLayer(num_input_features=num_features,
                                                              num_output_features=num_features // 2))
                num_features = num_features // 2

        self.final_bn = nn.BatchNorm2d(num_features)
        self.classifier = nn.Linear(num_features, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32:
            x = self.features[0](x)
            x = self.features[1](x)
            x = _ext.relu_maxpool2d(x)
        else:
            x = self.features(x)

        for i, block in enumerate(self.dense_blocks):
            x = block(x)
            if i != len(self.dense_blocks) - 1:
                x = self.transition_layers[i](x)

        x = self.final_bn(x)
        if x.is_cuda and x.dtype == torch.float32 and x.is_contiguous() and self.classifier.bias is not None:
            return _ext.fused_relu_pool_fc(x, self.classifier.weight, self.classifier.bias)

        x = F.relu(x, inplace=True)
        x = F.adaptive_avg_pool2d(x, (1, 1)).view(x.size(0), -1)
        x = self.classifier(x)
        return x