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

class Model(nn.Module):
    def __init__(self, input_channels, stages, block_widths, output_classes):
        """
        :param input_channels: int, Number of input channels for the first layer
        :param stages: int, Number of stages in the RegNet architecture
        :param block_widths: List[int], Width (number of channels) for each block in the stages
        :param output_classes: int, Number of output classes for classification
        """
        super(Model, self).__init__()

        self.stages = stages
        self.block_widths = block_widths
        
        layers = []
        current_channels = input_channels
        
        # Construct the stages with their respective blocks
        for i in range(stages):
            layers.append(self._make_stage(current_channels, block_widths[i]))
            current_channels = block_widths[i]
        
        self.feature_extractor = nn.Sequential(*layers)
        
        # Final fully connected layer for classification
        self.fc = nn.Linear(block_widths[-1], output_classes)
    
    def _make_stage(self, in_channels, out_channels):
        """
        Creates a simple block for each stage.
        :param in_channels: int, number of input channels
        :param out_channels: int, number of output channels
        :return: nn.Sequential block with convolutional layers
        """
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        """
        Forward pass through the RegNet model.
        :param x: torch.Tensor of shape (batch_size, input_channels, height, width)
        :return: torch.Tensor of shape (batch_size, output_classes)
        """
        x = self.feature_extractor(x)
        x = torch.mean(x, dim=[2, 3])  # Global Average Pooling
        x = self.fc(x)
        return x

# Test code for the RegNet model
batch_size = 8
input_channels = 3
image_height, image_width = 224, 224
stages = 3
block_widths = [64, 128, 256]
output_classes = 10

def get_inputs():
    """ Generates random input tensor of shape (batch_size, input_channels, height, width) """
    return [torch.rand(batch_size, input_channels, image_height, image_width)]

def get_init_inputs():
    """ Initializes model parameters """
    return [input_channels, stages, block_widths, output_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>
#include <ATen/cuda/CUDAContext.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")

void relu_maxpool_launcher(const at::Tensor& input, at::Tensor& output);
void avg_fc_launcher(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, at::Tensor& output);

torch::Tensor relu_maxpool_forward(torch::Tensor input) {
    CHECK_CUDA(input); CHECK_CONTIGUOUS(input); CHECK_FLOAT(input);
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");
    const auto N = input.size(0);
    const auto C = input.size(1);
    const auto H = input.size(2);
    const auto W = input.size(3);
    TORCH_CHECK((H >= 2) && (W >= 2), "spatial dims too small");
    auto output = torch::empty({N, C, H / 2, W / 2}, input.options());
    relu_maxpool_launcher(input, output);
    return output;
}

torch::Tensor avg_fc_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias) {
    CHECK_CUDA(input); CHECK_CONTIGUOUS(input); CHECK_FLOAT(input);
    CHECK_CUDA(weight); CHECK_CONTIGUOUS(weight); CHECK_FLOAT(weight);
    CHECK_CUDA(bias); CHECK_CONTIGUOUS(bias); CHECK_FLOAT(bias);
    TORCH_CHECK(input.dim() == 4, "input must be NCHW");
    TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
    TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
    TORCH_CHECK(input.size(1) == weight.size(1), "channel mismatch");
    TORCH_CHECK(weight.size(0) == bias.size(0), "class mismatch");
    auto output = torch::empty({input.size(0), weight.size(0)}, input.options());
    avg_fc_launcher(input, weight, bias, output);
    return output;
}
'''

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

__global__ void relu_maxpool_kernel(const float* __restrict__ x,
                                    float* __restrict__ y,
                                    long total, int C, int H, int W, int Ho, int Wo) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % Wo;
    long t = idx / Wo;
    int oh = t % Ho;
    t /= Ho;
    int c = t % C;
    int n = t / C;

    long base = ((long)n * C + c) * (long)H * W + (long)(oh * 2) * W + (ow * 2);
    float v0 = x[base];
    float v1 = x[base + 1];
    float v2 = x[base + W];
    float v3 = x[base + W + 1];
    float m = fmaxf(fmaxf(v0, v1), fmaxf(v2, v3));
    y[idx] = fmaxf(m, 0.0f);
}

void relu_maxpool_launcher(const at::Tensor& input, at::Tensor& output) {
    const int C = (int)input.size(1);
    const int H = (int)input.size(2);
    const int W = (int)input.size(3);
    const int Ho = (int)output.size(2);
    const int Wo = (int)output.size(3);
    const long total = output.numel();
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    relu_maxpool_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), total, C, H, W, Ho, Wo);
}

__global__ void avg_fc_kernel(const float* __restrict__ x,
                              const float* __restrict__ w,
                              const float* __restrict__ b,
                              float* __restrict__ y,
                              int N, int C, int H, int W, int K) {
    int nk = blockIdx.x;
    int k = nk % K;
    int n = nk / K;
    int tid = threadIdx.x;
    int HW = H * W;
    int CHW = C * HW;
    float sum = 0.0f;
    const float* xp = x + (long)n * CHW;
    const float* wp = w + (long)k * C;
    for (int idx = tid; idx < CHW; idx += blockDim.x) {
        int c = idx / HW;
        sum += xp[idx] * wp[c];
    }
    __shared__ float smem[256];
    smem[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        y[(long)n * K + k] = b[k] + smem[0] / (float)HW;
    }
}

void avg_fc_launcher(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, at::Tensor& output) {
    const int N = (int)input.size(0);
    const int C = (int)input.size(1);
    const int H = (int)input.size(2);
    const int W = (int)input.size(3);
    const int K = (int)weight.size(0);
    const int threads = 256;
    const int blocks = N * K;
    avg_fc_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        input.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        output.data_ptr<float>(), N, C, H, W, K);
}
'''

_ext = load_inline(
    name="regnet_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["relu_maxpool_forward", "avg_fc_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, input_channels, stages, block_widths, output_classes):
        super(ModelNew, self).__init__()
        self.stages = stages
        self.block_widths = block_widths
        layers = []
        current_channels = input_channels
        for i in range(stages):
            layers.append(self._make_stage(current_channels, block_widths[i]))
            current_channels = block_widths[i]
        self.feature_extractor = nn.Sequential(*layers)
        self.fc = nn.Linear(block_widths[-1], output_classes)
        self._fold_cache_key = None
        self._fold_cache = None

    def _make_stage(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def _cache_key(self):
        key = []
        for i in range(self.stages):
            st = self.feature_extractor[i]
            for conv, bn in ((st[0], st[1]), (st[3], st[4])):
                ts = (conv.weight, conv.bias, bn.weight, bn.bias, bn.running_mean, bn.running_var)
                for t in ts:
                    key.append((t.data_ptr(), t._version, tuple(t.shape), str(t.device)))
        return tuple(key)

    def _get_folded(self):
        key = self._cache_key()
        if self._fold_cache_key == key and self._fold_cache is not None:
            return self._fold_cache
        folded = []
        with torch.no_grad():
            for i in range(self.stages):
                st = self.feature_extractor[i]
                for conv, bn in ((st[0], st[1]), (st[3], st[4])):
                    scale = bn.weight * torch.rsqrt(bn.running_var + bn.eps)
                    w = (conv.weight * scale.reshape(-1, 1, 1, 1)).contiguous()
                    if conv.bias is None:
                        cb = torch.zeros_like(bn.running_mean)
                    else:
                        cb = conv.bias
                    b = ((cb - bn.running_mean) * scale + bn.bias).contiguous()
                    folded.append((w, b))
        self._fold_cache_key = key
        self._fold_cache = folded
        return folded

    def forward(self, x):
        # Fast inference path: fold BatchNorm into convolutions and fuse ReLU+2x2 maxpool.
        if not self.training:
            folded = self._get_folded()
            p = 0
            for i in range(self.stages):
                w1, b1 = folded[p]; w2, b2 = folded[p + 1]; p += 2
                x = F.conv2d(x, w1, b1, padding=1)
                x = torch.relu_(x)
                x = F.conv2d(x, w2, b2, padding=1)
                x = _ext.relu_maxpool_forward(x.contiguous())
            return _ext.avg_fc_forward(x.contiguous(), self.fc.weight.contiguous(), self.fc.bias.contiguous())

        # Training-mode-correct path: use PyTorch conv/BN for exact batch statistics, but
        # still avoid extra ReLU allocation and fuse ReLU with maxpool and final GAP+FC.
        for i in range(self.stages):
            st = self.feature_extractor[i]
            conv1, bn1, conv2, bn2 = st[0], st[1], st[3], st[4]
            x = F.conv2d(x, conv1.weight, conv1.bias, padding=1)
            x = F.batch_norm(x, bn1.running_mean, bn1.running_var, bn1.weight, bn1.bias,
                             True, bn1.momentum, bn1.eps)
            x = torch.relu_(x)
            x = F.conv2d(x, conv2.weight, conv2.bias, padding=1)
            x = F.batch_norm(x, bn2.running_mean, bn2.running_var, bn2.weight, bn2.bias,
                             True, bn2.momentum, bn2.eps)
            x = _ext.relu_maxpool_forward(x.contiguous())
        return _ext.avg_fc_forward(x.contiguous(), self.fc.weight.contiguous(), self.fc.bias.contiguous())