from torch.nn import Sequential, Linear, GRU, ReLU, Tanh
from .modules import *
from torch.nn.functional import mse_loss



#######################
# Basic AE-Implementation
class BasicAE(Module, ABC):

    @property
    def name(self):
        return self.__class__.__name__

    def __init__(self, dataParams, **kwargs):
        super(BasicAE, self).__init__()
        self.dataParams = dataParams
        self.latent_dim = kwargs.get('latent_dim', 2)
        self.encoder = self._build_encoder()
        self.decoder = self._build_decoder(out_shape=self.dataParams['features'])

    def _build_encoder(self):
        encoder = Sequential(
            Linear(6, 100, bias=True),
            ReLU(),
            Linear(100, 10, bias=True),
            ReLU()
        )
        gru = Sequential(
            TimeDistributed(encoder),
            GRU(10, 10, batch_first=True),
            RNNOutputFilter(only_last=True),
            Linear(10, self.latent_dim)
        )
        return gru

    def _build_decoder(self, out_shape):
        decoder = Sequential(
            Linear(10, 100, bias=True),
            ReLU(),
            Linear(100, out_shape, bias=True),
            Tanh()
        )

        gru = Sequential(
            GRU(self.latent_dim, 10,batch_first=True),
            RNNOutputFilter(),
            TimeDistributed(decoder)
        )
        return gru

    def forward(self, batch: torch.Tensor):
        # Encoder
        #  outputs, hidden (Batch, Timesteps aka. Size, Features / Latent Dim Size)
        z = self.encoder(batch)
        # Decoder
        # First repeat the data accordingly to the batch size
        z = Repeater((batch.shape[0], self.dataParams['size'], -1))(z)
        x_hat = self.decoder(z)
        return z, x_hat


class AELightningOverrides:

    def training_step(self, x, batch_nb):
        # z, x_hat
        _, x_hat = self.forward(x)
        loss = mse_loss(x, x_hat)
        return {'loss': loss}


if __name__ == '__main__':
    raise PermissionError('Get out of here - never run this module')