import numpy as np
import megengine
import megengine.module as M
import megengine.functional as F
import math
from . import default_init_weights

class ShuffleV2Block(M.Module):
    def __init__(self, inp, oup, mid_channels, *, ksize, stride):
        super().__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.mid_channels = mid_channels
        self.ksize = ksize
        pad = ksize // 2
        self.pad = pad
        self.inp = inp

        outputs = oup - inp

        branch_main = [
            # pw
            M.Conv2d(inp, mid_channels, 1, 1, 0, bias=False),
            M.ReLU(),
            # dw
            M.Conv2d(
                mid_channels, mid_channels, ksize, stride, pad,
                groups=mid_channels, bias=False,
            ),
            # pw-linear
            M.Conv2d(mid_channels, outputs, 1, 1, 0, bias=False),
            M.ReLU(),
        ]
        self.branch_main = M.Sequential(*branch_main)

        if stride == 2:
            branch_proj = [
                # dw
                
M.Conv2d(inp, inp, ksize, stride, pad, groups=inp, bias=False)