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

N, C, H, W = 32, 64, 56, 56
GROUPS = 4

assert (H * W) % 4 == 0, "Spatial size (H*W) must be a multiple of 4"


class ModelNew(nn.Module):

    def __init__(self, groups=GROUPS):
        super().__init__()
        self.groups = groups
        self.block_size = 256
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor channel_shuffle_cuda(torch::Tensor input, int groups);
        """

        cuda_source = f"""
        #include <cuda_runtime.h>

        #define BLOCK_SIZE {self.block_size}


        __global__ void channel_shuffle_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int N,
            int C,
            int S_vec, // H * W / 4
            int groups,
            int channels_per_group
        ) {{

            int total_threads = N * C * S_vec;
            int idx = blockIdx.x * blockDim.x + threadIdx.x;

            for (; idx < total_threads; idx += gridDim.x * blockDim.x) {{


                int s = idx % S_vec;
                int tmp = idx / S_vec;
                int c_in = tmp % C;         
                int n = tmp / C;           



                int g_idx = c_in / channels_per_group;
                int c_idx = c_in % channels_per_group;
                int c_out = c_idx * groups + g_idx;


                int out_global_idx = (n * C + c_out) * S_vec + s;


                const float4* in_ptr = reinterpret_cast<const float4*>(input);
                float4* out_ptr = reinterpret_cast<float4*>(output);

                out_ptr[out_global_idx] = in_ptr[idx];
            }}
        }}


        torch::Tensor channel_shuffle_cuda(torch::Tensor input, int groups) {{

            TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
            TORCH_CHECK(input.dim() == 4, "Input must be 4D (N, C, H, W)");

            int N = input.size(0);
            int C = input.size(1);
            int H = input.size(2);
            int W = input.size(3);

            TORCH_CHECK(C % groups == 0, "Channels must be divisible by groups");
            TORCH_CHECK((H * W) % 4 == 0, "Spatial size must be divisible by 4 for float4 optimization");


            input = input.contiguous();

            auto output = torch::empty_like(input);

            int S = H * W;
            int S_vec = S / 4;
            int channels_per_group = C / groups;

            int total_threads = N * C * S_vec;
            int blocks = std::min((total_threads + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);

            channel_shuffle_kernel<<<blocks, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                N,
                C,
                S_vec,
                groups,
                channels_per_group
            );

            return output;
        }}
        """

        self.op = load_inline(
            name='channel_shuffle_cuda_opt_v1',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['channel_shuffle_cuda'],
            extra_cuda_cflags=['-O3'],
            verbose=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_cuda: x = x.cuda()
        return self.op.channel_shuffle_cuda(x, self.groups)