from torch.optim import Adam

from .modules import *
from torch.nn.functional import mse_loss
from torch import Tensor


#######################
# Basic AE-Implementation
class AE_WithAttention(AbstractNeuralNetwork, ABC):

    def __init__(self, latent_dim: int=0, features: int = 0, **kwargs):
        assert latent_dim and features
        super(AE_WithAttention, self).__init__()
        self.latent_dim = latent_dim
        self.features = features
        self.encoder = Encoder(self.latent_dim)
        self.decoder = Decoder(self.latent_dim, self.features)

    def forward(self, batch: 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_repeatet = Repeater((batch.shape[0], batch.shape[1], -1))(z)
        x_hat = self.decoder(z_repeatet)
        return z, x_hat

    def training_step(self, x, batch_nb):
        # ToDo: We need a new loss function, fullfilling all attention needs
        # z, x_hat
        _, x_hat = self.forward(x)
        loss = mse_loss(y, x_hat) if self.train_on_predictions else mse_loss(x, x_hat)
        return {'loss': loss}

    def configure_optimizers(self):
        return [Adam(self.parameters(), lr=0.02)]


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