#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import megengine as mge
import megengine.module as M
import pytest

from basecls.models.regnet import RegBottleneckBlock
from basecls.models.resnet import (
    AnyStage,
    ResBasicBlock,
    ResBottleneckBlock,
    ResDeepStem,
    ResStem,
    SimpleStem,
)


@pytest.mark.parametrize("Block", [RegBottleneckBlock, ResBasicBlock, ResBottleneckBlock])
@pytest.mark.parametrize("w_in", [32])
@pytest.mark.parametrize("w_out", [32, 64])
@pytest.mark.parametrize("stride", [1, 2])
@pytest.mark.parametrize("bot_mul", [1.0, 0.25])
@pytest.mark.parametrize("group_w", [8])
@pytest.mark.parametrize("se_r", [0.0, 0.25])
@pytest.mark.parametrize("avg_down", [True, False])
@pytest.mark.parametrize("drop_path_prob", [0.05, 0.1])
@pytest.mark.parametrize("norm_name", ["BN"])
@pytest.mark.parametrize("act_name", ["relu"])
def test_block(
    Block,
    w_in,
    w_out,
    stride,
    bot_mul,
    group_w,
    se_r,
    avg_down,
    drop_path_prob,
    norm_name,
    act_name,
):
    m = Block(
        w_in,
        w_out,
        stride,
        bot_mul=bot_mul,
        group_w=group_w,
        se_r=se_r,
        avg_down=avg_down,
        drop_path_prob=drop_path_prob,
        norm_name=norm_name,
        act_name=act_name,
    )
    assert isinstance(m, M.Module)

    m(mge.random.normal(size=(2, 32, 8, 8)))


@pytest.mark.parametrize("Stem", [ResDeepStem, ResStem, SimpleStem])
@pytest.mark.parametrize("w_in", [3])
@pytest.mark.parametrize("w_out", [8, 16])
@pytest.mark.parametrize("norm_name", ["BN"])
@pytest.mark.parametrize("act_name", ["relu"])
def test_stem(Stem, w_in, w_out, norm_name, act_name):
    m = Stem(w_in, w_out, norm_name=norm_name, act_name=act_name)
    assert isinstance(m, M.Module)

    m(
mge.random.normal(size=(2, 3, 8, 8))