#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
# pytest ./tests/test_super_model.py -s             #
#####################################################
import torch
import unittest

from xautodl.xlayers import super_core
from xautodl import spaces


class TestSuperLinear(unittest.TestCase):
    """Test the super linear."""

    def test_super_linear(self):
        out_features = spaces.Categorical(12, 24, 36)
        bias = spaces.Categorical(True, False)
        model = super_core.SuperLinear(10, out_features, bias=bias)
        print("The simple super linear module is:\n{:}".format(model))
        model.apply_verbose(True)

        print(model.super_run_type)
        self.assertTrue(model.bias)

        inputs = torch.rand(20, 10)
        print("Input shape: {:}".format(inputs.shape))
        print("Weight shape: {:}".format(model._super_weight.shape))
        print("Bias shape: {:}".format(model._super_bias.shape))
        outputs = model(inputs)
        self.assertEqual(tuple(outputs.shape), (20, 36))

        abstract_space = model.abstract_search_space
        abstract_space.clean_last()
        abstract_child = abstract_space.random()
        print("The abstract searc space:\n{:}".format(abstract_space))
        print("The abstract child program:\n{:}".format(abstract_child))

        model.set_super_run_type(super_core.SuperRunMode.Candidate)
        model.enable_candidate()
        model.apply_candidate(abstract_child)

        output_shape = (20, abstract_child["_out_features"].value)
        outputs = model(inputs)
        self.assertEqual(tuple(outputs.shape), output_shape)

    def test_super_mlp_v1(self):
        hidden_features = spaces.Categorical(12, 24, 36)
        out_features = spaces.Categorical(24, 36, 48)
        mlp = super_core.SuperMLPv1(10, hidden_features, out_features)
        print(mlp)
        mlp.apply_verbose(False)
        self.assertTrue(mlp.fc1._out_features, mlp.fc2._in_features)

        inputs = torch.rand(4, 10)
        outputs = mlp(inputs)
        self.assertEqual(tuple(outputs.shape), (4, 48))

        abstract_space = mlp.abstract_search_space
        print(
            "The abstract search space for SuperMLPv1 is:\n{:}".format(abstract_space)
        )
        self.assertEqual(
            abstract_space["fc1"]["_out_features"],
            abstract_space["fc2"]["_in_features"],
        )
        self.assertTrue(
            abstract_space["fc1"]["_out_features"]
            is abstract_space["fc2"]["_in_features"]
        )

        abstract_space.clean_last()
        abstract_child = abstract_space.random(reuse_last=True)
        print("The abstract child program is:\n{:}".format(abstract_child))
        self.assertEqual(
            abstract_child["fc1"]["_out_features"].value,
            abstract_child["fc2"]["_in_features"].value,
        )

        mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
        mlp.enable_candidate()
        mlp.apply_candidate(abstract_child)
        outputs = mlp(inputs)
        output_shape = (4, abstract_child["fc2"]["_out_features"].value)
        self.assertEqual(tuple(outputs.shape), output_shape)

    def test_super_mlp_v2(self):
        hidden_multiplier = spaces.Categorical(1.0, 2.0, 3.0)
        out_features = spaces.Categorical(24, 36, 48)
        mlp = super_core.SuperMLPv2(10, hidden_multiplier, out_features)
        print(mlp)
        mlp.apply_verbose(False)

        inputs = torch.rand(4, 10)
        outputs = mlp(inputs)
        self.assertEqual(tuple(outputs.shape), (4, 48))

        abstract_space = mlp.abstract_search_space
        print(
            "The abstract search space for SuperMLPv2 is:\n{:}".format(abstract_space)
        )

        abstract_space.clean_last()
        abstract_child = abstract_space.random(reuse_last=True)
        print("The abstract child program is:\n{:}".format(abstract_child))

        mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
        mlp.enable_candidate()
        mlp.apply_candidate(abstract_child)
        outputs = mlp(inputs)
        output_shape = (4, abstract_child["_out_features"].value)
        self.assertEqual(tuple(outputs.shape), output_shape)

    def test_super_stem(self):
        out_features = spaces.Categorical(24, 36, 48)
        model = super_core.SuperAlphaEBDv1(6, out_features)
        inputs = torch.rand(4, 360)

        abstract_space = model.abstract_search_space
        abstract_space.clean_last()
        abstract_child = abstract_space.random(reuse_last=True)
        print("The abstract searc space:\n{:}".format(abstract_space))
        print("The abstract child program:\n{:}".format(abstract_child))

        model.set_super_run_type(super_core.SuperRunMode.Candidate)
        model.enable_candidate()
        model.apply_candidate(abstract_child)
        outputs = model(inputs)
        output_shape = (4, 60, abstract_child["_embed_dim"].value)
        self.assertEqual(tuple(outputs.shape), output_shape)