# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# Code modified from https://github.com/marlbenchmark/on-policy
from torchvision import transforms
import torch.nn as nn
from .util import init
"""CNN Modules and utils."""


class Flatten(nn.Module):

    def forward(self, x):
        return x.view(x.size(0), -1)


class CNNLayer(nn.Module):

    def __init__(self,
                 obs_shape,
                 hidden_size,
                 use_orthogonal,
                 use_ReLU,
                 kernel_size=3,
                 stride=1):
        super(CNNLayer, self).__init__()

        active_func = [nn.Tanh(), nn.ReLU()][use_ReLU]
        init_method = [nn.init.xavier_uniform_,
                       nn.init.orthogonal_][use_orthogonal]
        gain = nn.init.calculate_gain(['tanh', 'relu'][use_ReLU])

        self.resize = transforms.Resize(84)

        def init_(m):
            return init(m,
                        init_method,
                        lambda x: nn.init.constant_(x, 0),
                        gain=gain)

        input_channel = obs_shape[0]
        input_width = obs_shape[1]
        input_height = obs_shape[2]

        self.cnn = nn.Sequential(
            init_(
                nn.Conv2d(in_channels=input_channel,
                          out_channels=hidden_size // 2,
                          kernel_size=kernel_size,
                          stride=stride)), active_func, Flatten(),
            init_(
                nn.Linear(
                    hidden_size // 2 * (input_width - kernel_size + stride) *
                    (input_height - kernel_size + stride),
                    hidden_size)), active_func,
            init_(nn.Linear(hidden_size, hidden_size)), active_func)

    def forward(self, x):
        # TODO(eugenevinitsky) hardcoding is bad
        x = self.resize(x) / 255.0
        x = self.cnn(x)
        return x


class CNNBase(nn.Module):

    def __init__(self, args, obs_shape):
        super(CNNBase, self).__init__()

        self._use_orthogonal = args.use_orthogonal
        self._use_ReLU = args.use_ReLU
        self.hidden_size = args.hidden_size

        self.cnn = CNNLayer(obs_shape, self.hidden_size, self._use_orthogonal,
                            self._use_ReLU)

    def forward(self, x):
        x = self.cnn(x)
        return x
