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=dilation, bias=True
        )
        self.conv2 = M.Conv2d(channels, channels, 3, 1, padding=1, bias=True)
        if in_channels == channels and stride == 1:
            self.downsample = M.Identity()
        elif stride == 1:
            self.downsample = M.Conv2d(in_channels, channels, 1, stride, bias=False)
        else:
            self.downsample = M.Sequential(
                    M.AvgPool2d(kernel_size=stride, stride=stride),
                    M.Conv2d(in_channels, channels, 1, 1, bias=False)
                )
        self.fc1 = M.Conv2d(channels, 16, kernel_size=1) 
        self.fc2 = 
M.Conv2d(16, channels, kernel_size=1)