from collections.abc import Sequence
from typing import Optional, Tuple, Union

import torch
from torch import Tensor

from .skip_connection import SkipConnection


class Autoencoder(torch.nn.Module):
    """
    Base autoencoder class that defines some of the
    basic methods and functionality. Autoencoders are
    defined here as a set of sequential blocks that
    have an `encode` method, which acts on the input
    data to the autoencoder, and a `decode` method, which
    acts on the encoded vector generated by the `encode`
    method. `forward` just runs these steps one after the
    other. Although it isn't explicitly enforced, a good
    rule of thumb is that the ouput of a block's `decode`
    method should have the same shape as the _input_ of its
    `encode` method.

    Accepts a `skip_connection` argument that defines how to
    combine information from the input of one block's `encode`
    layer with the output to its `decode`layer. See `skip_connections.py`
    for more info about what these classes are expected to contain
    and how they operate.
    """

    def __init__(
        self, skip_connection: Optional[SkipConnection] = None
    ) -> None:
        super().__init__()
        self.skip_connection = skip_connection
        self.blocks = torch.nn.ModuleList()

    def encode(
        self, *X: Tensor, return_states: bool = False
    ) -> Union[Tensor, Tuple[Tensor, Sequence]]:
        states = []
        for block in self.blocks:
            if isinstance(X, tuple):
                X = block.encode(*X)
            else:
                X = block.encode(X)
            states.append(X)

        # don't need to return the last
        # state, since that's just equal
        # to the output of this layer
        if return_states:
            return X, states[:-1]
        return X

    def decode(self, *X, states: Optional[Sequence[Tensor]] = None) -> Tensor:
        if self.skip_connection is not None and states is None:
            raise ValueError(
                "Must pass intermediate states when autoencoder "
                "has a skip connection function specified"
            )
        elif states is not None:
            if len(states) != len(self.blocks) - 1:
                raise ValueError(
                    "Passed {} intermediate states, expected {}".format(
                        len(states), len(self.blocks) - 1
                    )
                )

            # Don't skip connect the output layer
            states = states[::-1] + [None]

        for i, block in enumerate(self.blocks[::-1]):
            if isinstance(X, tuple):
                X = block.decode(*X)
            else:
                X = block.decode(X)

            state = states[-i - 1]
            if state is not None:
                X = self.skip_connection(X, state)
        return X

    def forward(self, *X: Tensor) -> Tensor:
        return_states = self.skip_connection is not None
        X = self.encode(*X, return_states=return_states)
        if return_states:
            *X, states = X
        else:
            states = None

        if isinstance(X, Tensor):
            X = (X,)
        return self.decode(*X, states=states)
