import megengine as mge
import megengine.module as M
import megengine.functional as F
from megengine.core import Parameter
from utils import *

def addLeakyRelu(x):
    return M.Sequential(x, M.LeakyReLU(0.1))

def addSig(x):
    return M.Sequential(x, M.Sigmoid())

def up_block(x, ic, oc):
    return M.ConvTranspose2d(ic, oc, 4, stride=2, padding=1)

def down_block(x, ic, oc):
    return M.Conv2d(ic, oc, 3, padding=1, stride=2)

class BasicBlock(M.Module):
    expansion = 1

    def __init__(
        self,
        in_channels,
        channels,
        stride=1,
        groups=1,
        base_width=64,
        dilation=1,
        norm=M.BatchNorm2d,
    ):
        super().__init__()
        if groups != 1 or base_width != 64:
            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        self.conv1 = 
M.Conv2d(
            in_channels, channels, 3, stride, padding=dilatio