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 = r"""
#include <torch/extension.h>
torch::Tensor final_bn_relu_pool(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor running_mean, torch::Tensor running_var, double eps, bool training);
torch::Tensor small_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""

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

__global__ void final_eval_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,
                                  const float* __restrict__ rm, const float* __restrict__ rv,
                                  float* __restrict__ out, int N, int C, int H, int W, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C;
    if (idx >= total) return;
    int c = idx % C;
    int n = idx / C;
    int HW = H * W;
    float scale = w[c] * rsqrtf(rv[c] + eps);
    float shift = b[c] - rm[c] * scale;
    const float* base = x + ((n * C + c) * HW);
    float sum = 0.0f;
    for (int i = 0; i < HW; ++i) {
        float v = base[i] * scale + shift;
        sum += v > 0.0f ? v : 0.0f;
    }
    out[n * C + c] = sum / (float)HW;
}

__global__ void final_train_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ b,
                                   float* __restrict__ out, int N, int C, int H, int W, float eps) {
    int c = blockIdx.x;
    int tid = threadIdx.x;
    int HW = H * W;
    int M = N * HW;
    extern __shared__ float sh[];
    float* ssum = sh;
    float* ssq = sh + blockDim.x;

    float sum = 0.0f;
    float sq = 0.0f;
    for (int idx = tid; idx < M; idx += blockDim.x) {
        int n = idx / HW;
        int hw = idx - n * HW;
        float v = x[(n * C + c) * HW + hw];
        sum += v;
        sq += v * v;
    }
    ssum[tid] = sum;
    ssq[tid] = sq;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            ssum[tid] += ssum[tid + stride];
            ssq[tid] += ssq[tid + stride];
        }
        __syncthreads();
    }

    float mean = ssum[0] / (float)M;
    float var = ssq[0] / (float)M - mean * mean;
    float scale = w[c] * rsqrtf(var + eps);
    float shift = b[c] - mean * scale;

    for (int n = tid; n < N; n += blockDim.x) {
        const float* base = x + ((n * C + c) * HW);
        float os = 0.0f;
        for (int i = 0; i < HW; ++i) {
            float v = base[i] * scale + shift;
            os += v > 0.0f ? v : 0.0f;
        }
        out[n * C + c] = os / (float)HW;
    }
}

__global__ void linear_kernel(const float* __restrict__ x, const float* __restrict__ w,
                              const float* __restrict__ b, float* __restrict__ y,
                              int N, int C, int K) {
    int bid = blockIdx.x;
    int n = bid / K;
    int k = bid - n * K;
    int tid = threadIdx.x;
    extern __shared__ float sh[];
    float sum = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) {
        sum += x[n * C + c] * w[k * C + c];
    }
    sh[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    if (tid == 0) y[n * K + k] = sh[0] + b[k];
}

torch::Tensor final_bn_relu_pool(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                 torch::Tensor running_mean, torch::Tensor running_var,
                                 double eps, bool training) {
    int N = x.size(0), C = x.size(1), H = x.size(2), W = x.size(3);
    auto out = torch::empty({N, C}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (training) {
        int threads = 256;
        final_train_kernel<<<C, threads, threads * 2 * sizeof(float), stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            out.data_ptr<float>(), N, C, H, W, (float)eps);
    } else {
        int threads = 256;
        int total = N * C;
        int blocks = (total + threads - 1) / threads;
        final_eval_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            out.data_ptr<float>(), N, C, H, W, (float)eps);
    }
    return out;
}

torch::Tensor small_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    int N = x.size(0), C = x.size(1), K = weight.size(0);
    auto y = torch::empty({N, K}, x.options());
    int threads = 256;
    linear_kernel<<<N * K, threads, threads * sizeof(float), at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        y.data_ptr<float>(), N, C, K);
    return y;
}
"""

_ext = load_inline(
    name="densenet_final_manual_submit",
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=["final_bn_relu_pool", "small_linear"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)

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

    def forward(self, x):
        for layer in self.layers:
            y = layer[0](x)
            y = F.relu(y, inplace=True)
            y = layer[2](y)  # dropout p=0 is an identity; skip its Python dispatch
            x = torch.cat((x, y), 1)
        return x

class TransitionLayer(nn.Module):
    def __init__(self, num_input_features: int, num_output_features: int):
        super().__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 = F.relu(x, inplace=True)
        x = self.transition[2](x)
        return F.avg_pool2d(x, kernel_size=2, stride=2)

class ModelNew(nn.Module):
    def __init__(self, growth_rate: int = 32, num_classes: int = 1000):
        super().__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):
            self.dense_blocks.append(DenseBlock(num_layers, num_features, growth_rate))
            num_features += num_layers * growth_rate
            if i != len(block_layers) - 1:
                self.transition_layers.append(TransitionLayer(num_features, num_features // 2))
                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[0](x)
        x = self.features[1](x)
        x = F.relu(x, inplace=True)
        x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)

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

        pooled = _ext.final_bn_relu_pool(
            x.contiguous(),
            self.final_bn.weight,
            self.final_bn.bias,
            self.final_bn.running_mean,
            self.final_bn.running_var,
            float(self.final_bn.eps),
            bool(self.final_bn.training),
        )
        return _ext.small_linear(pooled, self.classifier.weight, self.classifier.bias)