"""
Code is from https://github.com/facebookresearch/habitat-lab/blob/main/habitat_baselines/rl/models/simple_cnn.py
"""

from typing import Dict, List, Union

import numpy as np
import torch
from torch import nn as nn


class SimpleCNN(nn.Module):
    """A Simple 3-Conv CNN followed by a fully connected layer

    Takes in observations and produces an embedding of the rgb and/or depth
        components. Note the observation spaces should be in format (W, H, C) but
        the input to the model should be in format (C, W, H).

    :param obs_shape: Either a dictionary of observation space shapes or the
        observation space shape. This will only select the keys from the
        observation space have 3 tensor dimensions.
    :param output_size: The hidden dimension output size
    """

    def __init__(
        self,
        obs_shape: Union[Dict[str, List[int]], List[int]],
        output_size: int,
    ):
        super().__init__()
        if not isinstance(obs_shape, dict):
            obs_shape = {None: obs_shape}

        obs_input_keys = [k for k, v in obs_shape.items() if len(v) == 3]

        self._n_input = sum(obs_shape[k][2] for k in obs_input_keys)
        self._obs_input_keys = obs_input_keys

        # kernel size for different CNN layers
        self._cnn_layers_kernel_size = [(8, 8), (4, 4), (3, 3)]

        # strides for different CNN layers
        self._cnn_layers_stride = [(4, 4), (2, 2), (1, 1)]

        cnn_dims = np.array(obs_shape[obs_input_keys[0]][:2], dtype=np.float32)
        for kernel_size, stride in zip(
            self._cnn_layers_kernel_size, self._cnn_layers_stride
        ):
            cnn_dims = self._conv_output_dim(
                dimension=cnn_dims,
                padding=np.array([0, 0], dtype=np.float32),
                dilation=np.array([1, 1], dtype=np.float32),
                kernel_size=np.array(kernel_size, dtype=np.float32),
                stride=np.array(stride, dtype=np.float32),
            )

        self.cnn = nn.Sequential(
            nn.Conv2d(
                in_channels=self._n_input,
                out_channels=32,
                kernel_size=self._cnn_layers_kernel_size[0],
                stride=self._cnn_layers_stride[0],
            ),
            nn.ReLU(True),
            nn.Conv2d(
                in_channels=32,
                out_channels=64,
                kernel_size=self._cnn_layers_kernel_size[1],
                stride=self._cnn_layers_stride[1],
            ),
            nn.ReLU(True),
            nn.Conv2d(
                in_channels=64,
                out_channels=32,
                kernel_size=self._cnn_layers_kernel_size[2],
                stride=self._cnn_layers_stride[2],
            ),
            #  nn.ReLU(True),
            nn.Flatten(),
            nn.Linear(32 * cnn_dims[0] * cnn_dims[1], output_size),
            nn.ReLU(True),
        )

        self.layer_init()

    def _conv_output_dim(self, dimension, padding, dilation, kernel_size, stride):
        r"""Calculates the output height and width based on the input
        height and width to the convolution layer.

        ref: https://pytorch.org/docs/master/nn.html#torch.nn.Conv2d
        """
        assert len(dimension) == 2
        out_dimension = []
        for i in range(len(dimension)):
            out_dimension.append(
                int(
                    np.floor(
                        (
                            (
                                dimension[i]
                                + 2 * padding[i]
                                - dilation[i] * (kernel_size[i] - 1)
                                - 1
                            )
                            / stride[i]
                        )
                        + 1
                    )
                )
            )
        return tuple(out_dimension)

    def layer_init(self):
        for layer in self.cnn:  # type: ignore
            if isinstance(layer, (nn.Conv2d, nn.Linear)):
                nn.init.kaiming_normal_(layer.weight, nn.init.calculate_gain("relu"))
                if layer.bias is not None:
                    nn.init.constant_(layer.bias, val=0)

    def forward(self, obs: Union[Dict[str, torch.Tensor], torch.Tensor]):
        if isinstance(obs, dict):
            cnn_inputs = torch.cat([obs[k] for k in self._obs_input_keys], dim=1)
        else:
            cnn_inputs = obs
        return self.cnn(cnn_inputs)
