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

class Model(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(Model, 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
    
batch_size = 10
num_layers = 6
num_input_features = 32
growth_rate = 32
height, width = 224, 224

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

def get_init_inputs():
    return [num_layers, num_input_features , growth_rate]

import torch
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r'''
#include <torch/extension.h>
#include <vector>
void dense_forward_cuda(torch::Tensor x, torch::Tensor out, torch::Tensor temp, torch::Tensor mean, torch::Tensor var,
                        std::vector<torch::Tensor> conv_w, std::vector<torch::Tensor> bn_w, std::vector<torch::Tensor> bn_b,
                        int num_layers, int c0, int growth);

torch::Tensor dense_forward(torch::Tensor x, std::vector<torch::Tensor> conv_w, std::vector<torch::Tensor> bn_w,
                            std::vector<torch::Tensor> bn_b, int num_layers, int c0, int growth) {
    x = x.contiguous();
    int64_t N=x.size(0), H=x.size(2), W=x.size(3);
    int ctot = c0 + num_layers*growth;
    auto out = torch::empty({N, ctot, H, W}, x.options());
    auto temp = torch::empty({N, ctot-growth, H, W}, x.options());
    auto mean = torch::empty({ctot-growth}, x.options());
    auto var = torch::empty({ctot-growth}, x.options());
    dense_forward_cuda(x, out, temp, mean, var, conv_w, bn_w, bn_b, num_layers, c0, growth);
    return out;
}
'''

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

__global__ void copy_input_kernel(const float* __restrict__ x, float* __restrict__ out, int C0, int Ctot, int HW, int total){
    int idx = blockIdx.x*blockDim.x + threadIdx.x;
    if(idx<total){
        int p=idx%HW; int c=(idx/HW)%C0; int n=idx/(C0*HW);
        out[(n*Ctot+c)*HW+p] = x[idx];
    }
}

__global__ void mean_kernel(const float* __restrict__ out, float* __restrict__ mean, int N, int Ctot, int HW){
    int c=blockIdx.x, tid=threadIdx.x;
    float sum=0.f; int M=N*HW;
    for(int i=tid;i<M;i+=256){
        int n=i/HW; int p=i-n*HW;
        sum += out[(n*Ctot+c)*HW+p];
    }
    __shared__ float s[256]; s[tid]=sum; __syncthreads();
    for(int off=128; off>0; off>>=1){ if(tid<off) s[tid]+=s[tid+off]; __syncthreads(); }
    if(tid==0) mean[c]=s[0]/(float)M;
}

__global__ void var_kernel(const float* __restrict__ out, const float* __restrict__ mean, float* __restrict__ var, int N, int Ctot, int HW){
    int c=blockIdx.x, tid=threadIdx.x;
    float m=mean[c], sum=0.f; int M=N*HW;
    for(int i=tid;i<M;i+=256){
        int n=i/HW; int p=i-n*HW;
        float d=out[(n*Ctot+c)*HW+p]-m;
        sum += d*d;
    }
    __shared__ float s[256]; s[tid]=sum; __syncthreads();
    for(int off=128; off>0; off>>=1){ if(tid<off) s[tid]+=s[tid+off]; __syncthreads(); }
    if(tid==0) var[c]=s[0]/(float)M;
}

__global__ void bn_relu_kernel(const float* __restrict__ out, float* __restrict__ temp,
                               const float* __restrict__ gamma, const float* __restrict__ beta,
                               const float* __restrict__ mean, const float* __restrict__ var,
                               int total, int C, int Ctot, int HW){
    int idx=blockIdx.x*blockDim.x+threadIdx.x;
    if(idx<total){
        int p=idx%HW; int c=(idx/HW)%C; int n=idx/(C*HW);
        float v = out[(n*Ctot+c)*HW+p];
        float y = (v - mean[c]) * (1.0f/sqrtf(var[c] + 1e-5f)) * gamma[c] + beta[c];
        temp[idx] = y > 0.f ? y : 0.f;
    }
}

__global__ void conv3x3_kernel(const float* __restrict__ temp, float* __restrict__ out,
                               const float* __restrict__ w,
                               int C, int G, int H, int W, int Ctot, int outCoff, int total){
    int idx = blockIdx.x*blockDim.x + threadIdx.x;
    int HW=H*W;
    if(idx>=total) return;
    int p = idx % HW;
    int oc = (idx / HW) % G;
    int n = idx / (G*HW);
    int h = p / W;
    int wi = p - h*W;
    float sum=0.f;
    for(int ic=0; ic<C; ++ic){
        const float* base = temp + ((n*C + ic)*H)*W;
        const float* ww = w + ((oc*C + ic)*3)*3;
        if(h>0){
            int row=(h-1)*W;
            if(wi>0) sum += base[row+wi-1]*ww[0];
            sum += base[row+wi]*ww[1];
            if(wi+1<W) sum += base[row+wi+1]*ww[2];
        }
        int row=h*W;
        if(wi>0) sum += base[row+wi-1]*ww[3];
        sum += base[row+wi]*ww[4];
        if(wi+1<W) sum += base[row+wi+1]*ww[5];
        if(h+1<H){
            int row2=(h+1)*W;
            if(wi>0) sum += base[row2+wi-1]*ww[6];
            sum += base[row2+wi]*ww[7];
            if(wi+1<W) sum += base[row2+wi+1]*ww[8];
        }
    }
    out[(n*Ctot + outCoff + oc)*HW + p] = sum;
}

void dense_forward_cuda(torch::Tensor x, torch::Tensor out, torch::Tensor temp, torch::Tensor mean, torch::Tensor var,
                        std::vector<torch::Tensor> conv_w, std::vector<torch::Tensor> bn_w, std::vector<torch::Tensor> bn_b,
                        int num_layers, int c0, int growth){
    int N=x.size(0), H=x.size(2), W=x.size(3), HW=H*W;
    int Ctot=c0+num_layers*growth;
    int threads=256;
    int copy_total=N*c0*HW;
    copy_input_kernel<<<(copy_total+threads-1)/threads,threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), c0, Ctot, HW, copy_total);
    for(int l=0; l<num_layers; ++l){
        int C=c0 + l*growth;
        mean_kernel<<<C,256>>>(out.data_ptr<float>(), mean.data_ptr<float>(), N, Ctot, HW);
        var_kernel<<<C,256>>>(out.data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), N, Ctot, HW);
        int total_bn=N*C*HW;
        bn_relu_kernel<<<(total_bn+threads-1)/threads,threads>>>(out.data_ptr<float>(), temp.data_ptr<float>(), bn_w[l].data_ptr<float>(), bn_b[l].data_ptr<float>(), mean.data_ptr<float>(), var.data_ptr<float>(), total_bn, C, Ctot, HW);
        int total_conv=N*growth*HW;
        conv3x3_kernel<<<(total_conv+threads-1)/threads,threads>>>(temp.data_ptr<float>(), out.data_ptr<float>(), conv_w[l].data_ptr<float>(), C, growth, H, W, Ctot, C, total_conv);
    }
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, num_layers: int, num_input_features: int, growth_rate: int):
        super().__init__()
        self.num_layers = num_layers
        self.num_input_features = num_input_features
        self.growth_rate = growth_rate
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.BatchNorm2d(num_input_features + i * growth_rate),
                nn.ReLU(inplace=True),
                nn.Conv2d(num_input_features + i * growth_rate, growth_rate, kernel_size=3, padding=1, bias=False),
                nn.Dropout(0.0)
            ) for i in range(num_layers)
        ])

    def forward(self, x):
        conv_w = [self.layers[i][2].weight.contiguous() for i in range(self.num_layers)]
        bn_w = [self.layers[i][0].weight.contiguous() for i in range(self.num_layers)]
        bn_b = [self.layers[i][0].bias.contiguous() for i in range(self.num_layers)]
        return _ext.dense_forward(x, conv_w, bn_w, bn_b, self.num_layers, self.num_input_features, self.growth_rate)