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

class Model(nn.Module):
    def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
        """
        :param in_channels: Number of input channels
        :param out_1x1: Number of output channels for the 1x1 convolution
        :param reduce_3x3: Number of output channels for the 1x1 reduction before 3x3 convolution
        :param out_3x3: Number of output channels for the 3x3 convolution
        :param reduce_5x5: Number of output channels for the 1x1 reduction before 5x5 convolution
        :param out_5x5: Number of output channels for the 5x5 convolution
        :param pool_proj: Number of output channels for the pooling projection
        """
        super(Model, self).__init__()
        
        # 1x1 convolution branch
        self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
        
        # 3x3 convolution branch
        self.branch3x3 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
            nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1)
        )
        
        # 5x5 convolution branch
        self.branch5x5 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
            nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2)
        )
        
        # Max pooling branch
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1)
        )
    
    def forward(self, x):
        """
        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels, height, width)
        """
        branch1x1 = self.branch1x1(x)
        branch3x3 = self.branch3x3(x)
        branch5x5 = self.branch5x5(x)
        branch_pool = self.branch_pool(x)
        
        outputs = [branch1x1, branch3x3, branch5x5, branch_pool]
        return torch.cat(outputs, 1)

# Test code
in_channels = 480
out_1x1 = 192
reduce_3x3 = 96
out_3x3 = 208
reduce_5x5 = 16
out_5x5 = 48
pool_proj = 64
batch_size = 10
height = 224
width = 224

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

def get_init_inputs():
    return [in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj]

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 cat_inception_cuda(torch::Tensor y, torch::Tensor b3, torch::Tensor b5, torch::Tensor bp);

torch::Tensor cat_inception(torch::Tensor y, torch::Tensor b3, torch::Tensor b5, torch::Tensor bp) {
    return cat_inception_cuda(y, b3, b5, bp);
}
'''

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

__global__ void cat_inception_kernel(const float* __restrict__ y,
                                     const float* __restrict__ b3,
                                     const float* __restrict__ b5,
                                     const float* __restrict__ bp,
                                     float* __restrict__ out,
                                     long long total, int HW) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const int COUT = 512;
    const int CY = 304;
    for (; idx < total; idx += stride) {
        int hw = (int)(idx % HW);
        long long t = idx / HW;
        int c = (int)(t % COUT);
        long long n = t / COUT;
        float v;
        if (c < 192) {
            v = y[(n * CY + c) * (long long)HW + hw];
        } else if (c < 400) {
            int cc = c - 192;
            v = b3[(n * 208 + cc) * (long long)HW + hw];
        } else if (c < 448) {
            int cc = c - 400;
            v = b5[(n * 48 + cc) * (long long)HW + hw];
        } else {
            int cc = c - 448;
            v = bp[(n * 64 + cc) * (long long)HW + hw];
        }
        out[idx] = v;
    }
}

torch::Tensor cat_inception_cuda(torch::Tensor y, torch::Tensor b3, torch::Tensor b5, torch::Tensor bp) {
    const long long B = y.size(0);
    const int H = (int)y.size(2);
    const int W = (int)y.size(3);
    const int HW = H * W;
    auto out = torch::empty({B, 512, H, W}, y.options());
    long long total = B * 512LL * HW;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    cat_inception_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        y.data_ptr<float>(), b3.data_ptr<float>(), b5.data_ptr<float>(), bp.data_ptr<float>(),
        out.data_ptr<float>(), total, HW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(
    name='inception_cat_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['cat_inception'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_1x1, reduce_3x3, out_3x3, reduce_5x5, out_5x5, pool_proj):
        super(ModelNew, self).__init__()
        self.branch1x1 = nn.Conv2d(in_channels, out_1x1, kernel_size=1)
        self.branch3x3 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_3x3, kernel_size=1),
            nn.Conv2d(reduce_3x3, out_3x3, kernel_size=3, padding=1)
        )
        self.branch5x5 = nn.Sequential(
            nn.Conv2d(in_channels, reduce_5x5, kernel_size=1),
            nn.Conv2d(reduce_5x5, out_5x5, kernel_size=5, padding=2)
        )
        self.branch_pool = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, pool_proj, kernel_size=1)
        )
        self._combined_w = None
        self._combined_b = None
        self._cache_ptrs = None

    def _get_combined_1x1(self):
        ptrs = (self.branch1x1.weight.data_ptr(), self.branch3x3[0].weight.data_ptr(),
                self.branch5x5[0].weight.data_ptr(), self.branch1x1.bias.data_ptr(),
                self.branch3x3[0].bias.data_ptr(), self.branch5x5[0].bias.data_ptr())
        if self._combined_w is None or self._cache_ptrs != ptrs:
            self._combined_w = torch.cat((self.branch1x1.weight, self.branch3x3[0].weight, self.branch5x5[0].weight), dim=0).contiguous()
            self._combined_b = torch.cat((self.branch1x1.bias, self.branch3x3[0].bias, self.branch5x5[0].bias), dim=0).contiguous()
            self._cache_ptrs = ptrs
        return self._combined_w, self._combined_b

    def forward(self, x):
        # The three independent 1x1 convolutions on x are one larger 1x1 convolution.
        # This removes two cudnn launches and reads the large input only once.
        w, b = self._get_combined_1x1()
        y = F.conv2d(x, w, b)
        r3 = y[:, 192:288, :, :]
        r5 = y[:, 288:304, :, :]
        b3 = F.conv2d(r3, self.branch3x3[1].weight, self.branch3x3[1].bias, padding=1)
        b5 = F.conv2d(r5, self.branch5x5[1].weight, self.branch5x5[1].bias, padding=2)
        p = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
        bp = F.conv2d(p, self.branch_pool[1].weight, self.branch_pool[1].bias)
        if x.is_cuda and x.dtype is torch.float32 and y.is_contiguous() and b3.is_contiguous() and b5.is_contiguous() and bp.is_contiguous():
            return _ext.cat_inception(y, b3, b5, bp)
        return torch.cat((y[:, :192, :, :], b3, b5, bp), dim=1)