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, 24, 16]  # Corresponding layers in DenseNet121

        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 DenseNet121 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 cat2_cuda(torch::Tensor a, torch::Tensor b);
torch::Tensor final_eval_cuda(torch::Tensor x,
                              torch::Tensor bn_weight,
                              torch::Tensor bn_bias,
                              torch::Tensor bn_mean,
                              torch::Tensor bn_var,
                              torch::Tensor fc_weight,
                              torch::Tensor fc_bias,
                              double eps);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void cat2_kernel(const float* __restrict__ a,
                            const float* __restrict__ b,
                            float* __restrict__ out,
                            long long total,
                            int C1, int C2, int HW) {
    long long t = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    int C = C1 + C2;
    for (; t < total; t += stride) {
        int inner = (int)(t % ((long long)C * HW));
        int c = inner / HW;
        int hw = inner - c * HW;
        long long n = t / ((long long)C * HW);
        if (c < C1) {
            out[t] = a[(n * (long long)C1 + c) * HW + hw];
        } else {
            out[t] = b[(n * (long long)C2 + (c - C1)) * HW + hw];
        }
    }
}

torch::Tensor cat2_cuda(torch::Tensor a, torch::Tensor b) {
    const int N = (int)a.size(0);
    const int C1 = (int)a.size(1);
    const int H = (int)a.size(2);
    const int W = (int)a.size(3);
    const int C2 = (int)b.size(1);
    auto out = torch::empty({N, C1 + C2, H, W}, a.options());
    long long total = (long long)N * (C1 + C2) * H * W;
    const int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 4096) blocks = 4096;
    cat2_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), total, C1, C2, H * W);
    return out;
}

__global__ void final_eval_kernel(const float* __restrict__ x,
                                  const float* __restrict__ bn_w,
                                  const float* __restrict__ bn_b,
                                  const float* __restrict__ bn_m,
                                  const float* __restrict__ bn_v,
                                  const float* __restrict__ fc_w,
                                  const float* __restrict__ fc_b,
                                  float* __restrict__ out,
                                  int C, int HW, int K, float eps) {
    // one block computes one (batch, class) output
    int nk = blockIdx.x;
    int n = nk / K;
    int k = nk - n * K;

    float acc = 0.0f;
    int total = C * HW;
    for (int idx = threadIdx.x; idx < total; idx += blockDim.x) {
        int c = idx / HW;
        float v = x[(n * (long long)C * HW) + idx];
        float y = (v - bn_m[c]) * rsqrtf(bn_v[c] + eps) * bn_w[c] + bn_b[c];
        y = y > 0.0f ? y : 0.0f;
        acc += y * fc_w[k * (long long)C + c];
    }

    __shared__ float smem[256];
    smem[threadIdx.x] = acc;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out[n * K + k] = smem[0] / (float)HW + fc_b[k];
    }
}

torch::Tensor final_eval_cuda(torch::Tensor x,
                              torch::Tensor bn_weight,
                              torch::Tensor bn_bias,
                              torch::Tensor bn_mean,
                              torch::Tensor bn_var,
                              torch::Tensor fc_weight,
                              torch::Tensor fc_bias,
                              double eps) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int HW = (int)(x.size(2) * x.size(3));
    const int K = (int)fc_weight.size(0);
    auto out = torch::empty({N, K}, x.options());
    final_eval_kernel<<<N * K, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
        bn_mean.data_ptr<float>(), bn_var.data_ptr<float>(), fc_weight.data_ptr<float>(),
        fc_bias.data_ptr<float>(), out.data_ptr<float>(), C, HW, K, (float)eps);
    return out;
}
"""

_ext = load_inline(
    name="densenet121_cat_final_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["cat2_cuda", "final_eval_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


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

    def forward(self, x):
        for layer in self.layers:
            new_feature = layer(x)
            if x.is_cuda and x.dtype is torch.float32 and new_feature.is_contiguous() and x.is_contiguous():
                x = _ext.cat2_cuda(x, new_feature)
            else:
                x = torch.cat((x, new_feature), 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):
        return self.transition(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, 24, 16]
        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 forward(self, x: torch.Tensor) -> torch.Tensor:
        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)

        # In eval/inference mode fuse final BN + ReLU + global average pool + linear.
        # In training mode BatchNorm must use batch statistics and update buffers, so keep
        # the exact PyTorch path for correctness.
        if (not self.training and x.is_cuda and x.dtype is torch.float32 and x.is_contiguous()
                and self.classifier.bias is not None):
            return _ext.final_eval_cuda(
                x,
                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,
            )

        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