import torch
import numpy as np
import pytest

from ding.model import ConvEncoder, FCEncoder, IMPALAConvEncoder
from ding.torch_utils import is_differentiable

B = 4
C, H, W = 3, 128, 128


@pytest.mark.unittest
class TestEncoder:

    def output_check(self, model, outputs):
        loss = outputs.sum()
        is_differentiable(loss, model)

    def test_conv_encoder(self):
        inputs = torch.randn(B, C, H, W)
        model = ConvEncoder((C, H, W), hidden_size_list=[32, 48, 64, 64, 128], activation=torch.nn.Tanh())
        print(model)
        outputs = model(inputs)
        self.output_check(model, outputs)
        assert outputs.shape == (B, 128)

    def test_dreamer_conv_encoder(self):
        inputs = torch.randn(B, C, H, W)
        model = ConvEncoder(
            (C, H, W),
            hidden_size_list=[32, 64, 128, 256, 128],
            activation=torch.nn.SiLU(),
            kernel_size=[4, 4, 4, 4],
            layer_norm=True
        )
        print(model)
        outputs = model(inputs)
        self.output_check(model, outputs)
        assert outputs.shape == (B, 128)

    def test_fc_encoder(self):
        inputs = torch.randn(B, 32)
        hidden_size_list = [128 for _ in range(3)]
        model = FCEncoder(32, hidden_size_list, res_block=True, activation=torch.nn.Tanh())
        print(model)
        outputs = model(inputs)
        self.output_check(model, outputs)
        assert outputs.shape == (B, hidden_size_list[-1])

        hidden_size_list = [64, 128, 256]
        model = FCEncoder(32, hidden_size_list, res_block=False, activation=torch.nn.Tanh())
        print(model)
        outputs = model(inputs)
        self.output_check(model, outputs)
        assert outputs.shape == (B, hidden_size_list[-1])

    def test_impalaconv_encoder(self):
        inputs = torch.randn(B, 3, 64, 64)
        model = IMPALAConvEncoder(obs_shape=(3, 64, 64))
        print(model)
        outputs = model(inputs)
        self.output_check(model, outputs)
        assert outputs.shape == (B, 256)