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

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

__global__ void bn_relu_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,
                               const float* __restrict__ mean, const float* __restrict__ var,
                               float* __restrict__ y, long total, int C, int H, int W, float eps){
    long idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx >= total) return;
    int hw = H*W;
    int c = (idx / hw) % C;
    float v = x[idx];
    float a = w[c] * rsqrtf(var[c] + eps);
    float z = (v - mean[c]) * a + b[c];
    y[idx] = z > 0.0f ? z : 0.0f;
}

__global__ void avgpool2_kernel(const float* __restrict__ x, float* __restrict__ y, int N, int C, int H, int W){
    long total = (long)N*C*(H/2)*(W/2);
    long idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx >= total) return;
    int Wo = W/2, Ho = H/2;
    int ow = idx % Wo;
    int oh = (idx / Wo) % Ho;
    int c = (idx / (Wo*Ho)) % C;
    int n = idx / (Wo*Ho*C);
    long base = ((long)n*C + c)*H*W + (2*oh)*W + 2*ow;
    y[idx] = 0.25f*(x[base] + x[base+1] + x[base+W] + x[base+W+1]);
}

__global__ void final_linear_kernel(const float* __restrict__ x,
                                    const float* __restrict__ bn_w, const float* __restrict__ bn_b,
                                    const float* __restrict__ mean, const float* __restrict__ var,
                                    const float* __restrict__ fc_w, const float* __restrict__ fc_b,
                                    float* __restrict__ out,
                                    int N, int C, int H, int W, int K, float eps){
    int nk = blockIdx.x;
    int tid = threadIdx.x;
    int n = nk / K;
    int k = nk - n*K;
    extern __shared__ float s[];
    float sum = 0.0f;
    int HW = H*W;
    int tot = C*HW;
    for(int t = tid; t < tot; t += blockDim.x){
        int c = t / HW;
        int r = t - c*HW;
        float v = x[((long)n*C + c)*HW + r];
        float z = (v - mean[c]) * (bn_w[c] * rsqrtf(var[c] + eps)) + bn_b[c];
        z = z > 0.0f ? z : 0.0f;
        sum += z * (fc_w[(long)k*C + c] / (float)HW);
    }
    s[tid] = sum;
    __syncthreads();
    for(int stride = blockDim.x/2; stride > 0; stride >>= 1){
        if(tid < stride) s[tid] += s[tid+stride];
        __syncthreads();
    }
    if(tid == 0) out[nk] = s[0] + fc_b[k];
}

torch::Tensor bn_relu(torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean, torch::Tensor var, double eps){
    auto y = torch::empty_like(x);
    int N=x.size(0), C=x.size(1), H=x.size(2), W=x.size(3);
    long total = (long)N*C*H*W;
    int threads=256;
    int blocks=(total+threads-1)/threads;
    bn_relu_kernel<<<blocks, threads>>>(x.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), y.data_ptr<float>(), total, C, H, W, (float)eps);
    return y;
}

torch::Tensor avgpool2(torch::Tensor x){
    int N=x.size(0), C=x.size(1), H=x.size(2), W=x.size(3);
    auto y = torch::empty({N,C,H/2,W/2}, x.options());
    long total=(long)N*C*(H/2)*(W/2);
    int threads=256;
    int blocks=(total+threads-1)/threads;
    avgpool2_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), N,C,H,W);
    return y;
}

torch::Tensor final_linear(torch::Tensor x, torch::Tensor bn_w, torch::Tensor bn_b, torch::Tensor mean, torch::Tensor var, torch::Tensor fc_w, torch::Tensor fc_b, double eps){
    int N=x.size(0), C=x.size(1), H=x.size(2), W=x.size(3), K=fc_w.size(0);
    auto out=torch::empty({N,K}, x.options());
    int threads=256;
    final_linear_kernel<<<N*K, threads, threads*sizeof(float)>>>(x.data_ptr<float>(), bn_w.data_ptr<float>(), bn_b.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), fc_w.data_ptr<float>(), fc_b.data_ptr<float>(), out.data_ptr<float>(), N,C,H,W,K,(float)eps);
    return out;
}
'''

_ext = load_inline(name='densenet201_fused_ext_v1', cpp_sources='', cuda_sources=cuda_src,
                   functions=['bn_relu','avgpool2','final_linear'], verbose=False,
                   extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

class DenseBlock(nn.Module):
    def __init__(self, num_layers: int, num_input_features: int, growth_rate: int):
        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):
        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)
        )

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)
        )

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):
            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
        self.final_bn = nn.BatchNorm2d(num_features)
        self.classifier = nn.Linear(num_features, num_classes)

    def _bn_relu(self, x, bn):
        # The benchmark uses fp32 CUDA inference. For training or CPU, preserve exact PyTorch semantics.
        if x.is_cuda and x.dtype is torch.float32 and (not self.training):
            return _ext.bn_relu(x.contiguous(), bn.weight, bn.bias, bn.running_mean, bn.running_var, bn.eps)
        return F.relu(bn(x), inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_cuda or x.dtype is not torch.float32 or self.training:
            x = self.features(x)
            for i, block in enumerate(self.dense_blocks):
                features=[x]
                for layer in block.layers:
                    new_feature = layer(x)
                    features.append(new_feature)
                    x = torch.cat(features, 1)
                if i != len(self.dense_blocks)-1:
                    x = self.transition_layers[i].transition(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)
            return self.classifier(x)

        # initial stem: fuse batchnorm+relu, leave highly optimized conv/maxpool to cuDNN/ATen
        x = self.features[0](x)
        x = self._bn_relu(x, self.features[1])
        x = self.features[3](x)

        for bi, block in enumerate(self.dense_blocks):
            feats = [x]
            for layer in block.layers:
                z = self._bn_relu(x, layer[0])
                nf = layer[2](z)
                feats.append(nf)
                x = torch.cat(feats, 1)
            if bi != len(self.dense_blocks) - 1:
                tr = self.transition_layers[bi].transition
                z = self._bn_relu(x, tr[0])
                x = tr[2](z)
                x = _ext.avgpool2(x.contiguous())

        return _ext.final_linear(x.contiguous(), self.final_bn.weight, self.final_bn.bias,
                                 self.final_bn.running_mean, self.final_bn.running_var,
                                 self.classifier.weight, self.classifier.bias, self.final_bn.eps)