import abc
import torch
import torch.nn as nn

from typing import Union, List, Sequence, Iterator, Dict


class VQVAEBase(nn.Module, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def forward(self, images: torch.Tensor) -> Dict[str, List[torch.Tensor]]:
        raise NotImplementedError

    @abc.abstractmethod
    def encode(self, images: torch.Tensor) -> List[torch.Tensor]:
        """
            Taking the batch it passes it through the encoder and returns the required
            list of tensors to be passed to self.quantize method.

            :param images: torch.Tensor holding the images to be reconstructed
            :return: Lists of torch.Tensor with the encodings.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def quantize(self, encodings: List[torch.Tensor]) -> List[torch.Tensor]:
        """
            Taking a List of torch.Tensor it passes them through the quantization process
            taking care of any conditionings between quantization levels.

            :param encodings: List of torch.Tensor encodings to be quantized.
            :return: List of torch.Tensor with the quantized tensors.
        """
        raise NotImplemented

    @abc.abstractmethod
    def decode(self, quantizations: List[torch.Tensor]) -> torch.Tensor:
        """
            Taking a List of torch.Tensor it passes them through the decoder and returns the
            reconstructed images.

            :param quantizations: List of torch.Tensor quantizations to be decoded.
            :return: torch.Tensor reconstructions of the images
        """
        raise NotImplemented

    @abc.abstractmethod
    def index_quantize(self, images: torch.Tensor) -> List[torch.Tensor]:
        """

            Taking torch.Tensor of images it passes them through the encoding and  quantization process
            taking care of any conditionings between quantization levels and returns the indexes of the atomic element.
            This method is aimed at generating the required input for autoregressive models.

            :param images: torch.Tensor holding the images for which the index quantization will be generated
            :return: List of torch.Tensor with the quantized tensors where values are the indexes of the atomic
                elements.
        """
        raise NotImplemented

    @abc.abstractmethod
    def decode_samples(self, embedding_indices: List[torch.Tensor]) -> torch.Tensor:
        """
            Taking a List of torch.Tensor autoregressive samples it passes them through the
            embedding process and then it reconstructs the images by passing the through the
            decoder.

            :param embedding_indices: Samples generated by an autoregressive method
            :return: samples projected in the image space
        """
        raise NotImplemented

    @abc.abstractmethod
    def construct_encoder(self) -> nn.ModuleList:
        """
            Method that aims to return the final structure of the encoder part of the network.

            This method needs to be called in the __init__ method to set the value of self.encoder
            which needs to have the type Final[List[nn.Sequential]] to make it TorchScript compatible.

            The structure is a List of nn.Sequential where each element returns a level of encodings
            which will be further passed to the self.quantizer method.

            :return: ModuleList of nn.Sequential encoding elements
        """
        raise NotImplemented

    @abc.abstractmethod
    def construct_quantizer(self) -> nn.ModuleList:
        """
            Method that aims to return the final structure of the quantization part of the network.

            This method needs to be called in the __init__ method to set the value of self.quantizer
            which needs to have the type Final[List[nn.Sequential]] to make it TorchScript compatible.

            Structure of the List[nn.Sequential] should be implementation dependent.

            :return: ModuleList of nn.Sequential quantization elements
        """
        raise NotImplemented

    @abc.abstractmethod
    def construct_decoder(self) -> nn.ModuleList:
        """
            Method that aims to return the final structure of the decoder part of the network.

            This method needs to be called in the __init__ method to set the value of self.decoder
            which needs to have the type Final[List[nn.Sequential]] to make it TorchScript compatible.

            Each element should take either a quantization element, concatenation of latent representation
            concatenated with a quantization element or their summation.

            :return: ModuleList of nn.Sequential decoding elements
        """
        raise NotImplemented

    @abc.abstractmethod
    def get_ema_decay(self) -> Sequence[float]:
        raise NotImplemented

    @abc.abstractmethod
    def set_ema_decay(self, decay: Union[Sequence[float], float]) -> Sequence[float]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_commitment_cost(self) -> Sequence[float]:
        raise NotImplementedError

    @abc.abstractmethod
    def set_commitment_cost(
        self, commitment_factor: Union[Sequence[float], float]
    ) -> Sequence[float]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_perplexity(self) -> Sequence[float]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_last_layer(self) -> nn.parameter.Parameter:
        raise NotImplementedError

    def trigger_code_restarts(self) -> None:
        """
            Triggers the codebook restarting procedure.

            The method should be used if the network has the capabilities of point
            III.B Batch data-dependent codebook updates from the first reference.

            Reference:
                Adrian Lancucki et al.
                Robust Training of Vector Quantized Bottleneck Models.
                In 2020 International Joint Conference on Neural Networks,
                IJCNN 2020, Glasgow, United Kingdom, July 19-24, 2020 (pp. 1–7). IEEE.
        """
        pass

    def get_parameters(
        self, with_codebook: bool = True
    ) -> Iterator[nn.parameter.Parameter]:
        """
            Fetches the module parameters with or without the codebooks ones.

            The method should be used if the network has the capabilities of point
            III.C EMA: an alternative training rule

            :param with_codebook: Boolean that dictates if we include or not the codebooks parameters
            :return: Iterator containing the requested parameters

            Reference:
                Adrian Lancucki et al.
                Robust Training of Vector Quantized Bottleneck Models.
                In 2020 International Joint Conference on Neural Networks,
                IJCNN 2020, Glasgow, United Kingdom, July 19-24, 2020 (pp. 1–7). IEEE.
        """
        pass

    def get_codebook_parameters(self) -> Iterator[nn.parameter.Parameter]:
        """
            Fetches the codebooks parameters.

            The method should be used if the network has the capabilities of point
            III.C EMA: an alternative training rule

            :return: Iterator containing the codebooks parameters

            Reference:
                Adrian Lancucki et al.
                Robust Training of Vector Quantized Bottleneck Models.
                In 2020 International Joint Conference on Neural Networks,
                IJCNN 2020, Glasgow, United Kingdom, July 19-24, 2020 (pp. 1–7). IEEE.
        """
        pass
