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

class Model(nn.Module):
    def __init__(self, in_channels, out_channels, groups=3):
        """
        ShuffleNet unit implementation.

        :param in_channels: Number of input channels.
        :param out_channels: Number of output channels.
        :param groups: Number of groups for group convolution.
        """
        super(Model, self).__init__()
        
        # Ensure the output channels are divisible by groups
        assert out_channels % 4 == 0
        mid_channels = out_channels // 4
        
        # First 1x1 group convolution
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        
        # Depthwise 3x3 convolution
        self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, groups=mid_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        
        # Second 1x1 group convolution
        self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        
        # Shuffle operation
        self.shuffle = ChannelShuffle(groups)
        
        # Shortcut connection if input and output channels are the same
        if in_channels == out_channels:
            self.shortcut = nn.Sequential()
        else:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        """
        Forward pass for ShuffleNet unit.

        :param x: Input tensor, shape (batch_size, in_channels, height, width)
        :return: Output tensor, shape (batch_size, out_channels, height, width)
        """
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.shuffle(out)
        out = F.relu(self.bn3(self.conv3(out)))
        
        out += self.shortcut(x)
        return out

class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        """
        Channel shuffle operation.

        :param groups: Number of groups for shuffling.
        """
        super(ChannelShuffle, self).__init__()
        self.groups = groups
    
    def forward(self, x):
        """
        Forward pass for channel shuffle.

        :param x: Input tensor, shape (batch_size, channels, height, width)
        :return: Output tensor, shape (batch_size, channels, height, width)
        """
        batch_size, channels, height, width = x.size()
        channels_per_group = channels // self.groups
        
        # Reshape
        x = x.view(batch_size, self.groups, channels_per_group, height, width)
        
        # Transpose
        x = x.transpose(1, 2).contiguous()
        
        # Flatten
        x = x.view(batch_size, -1, height, width)
        
        return x
    
batch_size = 10
input_channels = 240
out_channels = 480
groups = 3
height = 224
width = 224
num_classes = 1000

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

def get_init_inputs():
    return [input_channels, out_channels, groups]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <vector>

torch::Tensor shuffle_cuda(torch::Tensor x, int64_t groups);

torch::Tensor shufflenet_forward(
    torch::Tensor x,
    torch::Tensor conv1_w, torch::Tensor bn1_w, torch::Tensor bn1_b, torch::Tensor bn1_rm, torch::Tensor bn1_rv, double bn1_eps, double bn1_mom,
    torch::Tensor conv2_w, torch::Tensor bn2_w, torch::Tensor bn2_b, torch::Tensor bn2_rm, torch::Tensor bn2_rv, double bn2_eps, double bn2_mom,
    torch::Tensor conv3_w, torch::Tensor bn3_w, torch::Tensor bn3_b, torch::Tensor bn3_rm, torch::Tensor bn3_rv, double bn3_eps, double bn3_mom,
    torch::Tensor sc_w, torch::Tensor sc_bn_w, torch::Tensor sc_bn_b, torch::Tensor sc_bn_rm, torch::Tensor sc_bn_rv, double sc_bn_eps, double sc_bn_mom,
    int64_t groups, bool training) {

    std::vector<int64_t> stride = {1,1};
    std::vector<int64_t> pad0 = {0,0};
    std::vector<int64_t> pad1 = {1,1};
    std::vector<int64_t> dil = {1,1};
    auto none = c10::optional<torch::Tensor>();

    auto main_stream = at::cuda::getCurrentCUDAStream();
    auto side_stream = at::cuda::getStreamFromPool();

    // The shortcut branch is independent of the main branch; launch it on a
    // side stream while the three-convolution ShuffleNet path runs.
    torch::Tensor sc;
    {
        c10::cuda::CUDAStreamGuard guard(side_stream);
        sc = at::conv2d(x, sc_w, none, stride, pad0, dil, 1);
        sc = at::batch_norm(sc, sc_bn_w, sc_bn_b, sc_bn_rm, sc_bn_rv,
                            training, sc_bn_mom, sc_bn_eps, true);
    }
    c10::cuda::CUDAStreamGuard guard_main(main_stream);

    auto o1 = at::conv2d(x, conv1_w, none, stride, pad0, dil, groups);
    o1 = at::batch_norm(o1, bn1_w, bn1_b, bn1_rm, bn1_rv,
                        training, bn1_mom, bn1_eps, true);
    o1 = at::relu_(o1);

    int64_t mid = conv2_w.size(0);
    auto o2 = at::conv2d(o1, conv2_w, none, stride, pad1, dil, mid);
    o2 = at::batch_norm(o2, bn2_w, bn2_b, bn2_rm, bn2_rv,
                        training, bn2_mom, bn2_eps, true);
    auto sh = shuffle_cuda(o2, groups);

    auto o3 = at::conv2d(sh, conv3_w, none, stride, pad0, dil, groups);
    o3 = at::batch_norm(o3, bn3_w, bn3_b, bn3_rm, bn3_rv,
                        training, bn3_mom, bn3_eps, true);
    o3 = at::relu_(o3);

    cudaEvent_t ev;
    cudaEventCreateWithFlags(&ev, cudaEventDisableTiming);
    cudaEventRecord(ev, side_stream.stream());
    cudaStreamWaitEvent(main_stream.stream(), ev, 0);
    cudaEventDestroy(ev);
    o3.add_(sc);
    return o3;
}
'''

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

__global__ void shuffle_kernel(const float* __restrict__ x, float* __restrict__ y,
                               int64_t total, int C, int HW, int groups) {
    int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int spatial = idx % HW;
    int co = (idx / HW) % C;
    int n = idx / ((int64_t)C * HW);
    int cpg = C / groups;
    int ci = (co % groups) * cpg + (co / groups);
    y[idx] = x[((int64_t)n * C + ci) * HW + spatial];
}

torch::Tensor shuffle_cuda(torch::Tensor x, int64_t groups) {
    auto y = torch::empty_like(x);
    int64_t total = x.numel();
    int C = x.size(1);
    int HW = x.size(2) * x.size(3);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    shuffle_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), total, C, HW, (int)groups);
    return y;
}
'''

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

class ChannelShuffle(nn.Module):
    def __init__(self, groups):
        super().__init__()
        self.groups = groups
    def forward(self, x):
        n, c, h, w = x.shape
        g = self.groups
        return x.view(n, g, c // g, h, w).transpose(1, 2).contiguous().view(n, c, h, w)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, groups=3):
        super().__init__()
        assert out_channels % 4 == 0
        mid_channels = out_channels // 4
        self.groups = groups
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, groups=mid_channels, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1, stride=1, padding=0, groups=groups, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        self.shuffle = ChannelShuffle(groups)
        if in_channels == out_channels:
            self.shortcut = nn.Sequential()
        else:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        return _ext.shufflenet_forward(
            x,
            self.conv1.weight, self.bn1.weight, self.bn1.bias, self.bn1.running_mean, self.bn1.running_var, self.bn1.eps, self.bn1.momentum,
            self.conv2.weight, self.bn2.weight, self.bn2.bias, self.bn2.running_mean, self.bn2.running_var, self.bn2.eps, self.bn2.momentum,
            self.conv3.weight, self.bn3.weight, self.bn3.bias, self.bn3.running_mean, self.bn3.running_var, self.bn3.eps, self.bn3.momentum,
            self.shortcut[0].weight, self.shortcut[1].weight, self.shortcut[1].bias, self.shortcut[1].running_mean, self.shortcut[1].running_var, self.shortcut[1].eps, self.shortcut[1].momentum,
            self.groups, self.training
        )