<gh_stars>0
from typing import Tuple

import pytorch_lightning as pl
import torch
from torch import nn, optim


class SimpleEncoder(nn.Module):
    def __init__(self, state_size: int, hidden_size: int, latent_size: int) -> None:
        super(SimpleEncoder, self).__init__()

        self.encoder = nn.Sequential(
            nn.Linear(state_size, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, latent_size),
        )

    def forward(self, x):
        return self.encoder(x)


class SimpleDecoder(nn.Module):
    def __init__(self, state_size: int, hidden_size: int, latent_size: int) -> None:
        super(SimpleDecoder, self).__init__()

        self.decoder = nn.Sequential(
            nn.Linear(latent_size, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, state_size),
        )

    def forward(self, x):
        return self.decoder(x)


class SimpleLitAE(pl.LightningModule):
    def __init__(self, hparams_dict: dict) -> None:
        super().__init__()

        self.save_hyperparameters(hparams_dict)

        self.encoder = SimpleEncoder(
            hparams_dict["state_size"],
            hparams_dict["hidden_size"],
            hparams_dict["latent_size"],
        )
        self.decoder = SimpleDecoder(
            hparams_dict["state_size"],
            hparams_dict["hidden_size"],
            hparams_dict["latent_size"],
        )

        self.lr = hparams_dict["lr"]
        self.num_epochs = hparams_dict["num_epochs"]

        self.loss = nn.MSELoss()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        embedding = self.encoder(x)
        return embedding

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.lr)

        # init and fill dict
        lr_scheduler = {}
        lr_scheduler["name"] = "learning_rate"
        lr_scheduler["interval"] = "step"
        lr_scheduler["frequency"] = 1
        lr_scheduler["scheduler"] = optim.lr_scheduler.OneCycleLR(
            optimizer=optimizer,
            max_lr=self.lr,
            steps_per_epoch=len(self.train_dataloader()),
            epochs=self.num_epochs,
        )
        return [optimizer], [lr_scheduler]

    def training_step(self, train_batch: Tuple, batch_idx):
        x, y = train_batch
        x, y = x.float(), y.float()

        x = self.encoder(x)
        x = self.decoder(x)

        loss = self.loss(x, y)

        self.log("loss", loss)
        return {"loss": loss}

    def validation_step(self, val_batch: Tuple, batch_idx):
        x, y = val_batch
        x, y = x.float(), y.float()

        x = self.encoder(x)
        x = self.decoder(x)

        loss = self.loss(x, y)

        self.log("val_loss", loss)

from typing import Tuple

import pytorch_lightning as pl
import torch
from torch import nn, optim


class SimpleEncoder(nn.Module):
    def __init__(self, state_size: int, hidden_size: int, latent_size: int) -> None:
        super(SimpleEncoder, self).__init__()

        self.encoder = nn.Sequential(
            nn.Linear(state_size, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, latent_size),
        )

    def forward(self, x):
        return self.encoder(x)


class SimpleDecoder(nn.Module):
    def __init__(self, state_size: int, hidden_size: int, latent_size: int) -> None:
        super(SimpleDecoder, self).__init__()

        self.decoder = nn.Sequential(
            nn.Linear(latent_size, hidden_size),
            nn.BatchNorm1d(hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, state_size),
        )

    def forward(self, x):
        return self.decoder(x)


class SimpleLitAE(pl.LightningModule):
    def __init__(self, hparams_dict: dict) -> None:
        super().__init__()

        self.save_hyperparameters(hparams_dict)

        self.encoder = SimpleEncoder(
            hparams_dict["state_size"],
            hparams_dict["hidden_size"],
            hparams_dict["latent_size"],
        )
        self.decoder = SimpleDecoder(
            hparams_dict["state_size"],
            hparams_dict["hidden_size"],
            hparams_dict["latent_size"],
        )

        self.lr = hparams_dict["lr"]
        self.num_epochs = hparams_dict["num_epochs"]

        self.loss = nn.MSELoss()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        embedding = self.encoder(x)
        return embedding

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=self.lr)

        # init and fill dict
        lr_scheduler = {}
        lr_scheduler["name"] = "learning_rate"
        lr_scheduler["interval"] = "step"
        lr_scheduler["frequency"] = 1
        lr_scheduler["scheduler"] = optim.lr_scheduler.OneCycleLR(
            optimizer=optimizer,
            max_lr=self.lr,
            steps_per_epoch=len(self.train_dataloader()),
            epochs=self.num_epochs,
        )
        return [optimizer], [lr_scheduler]

    def training_step(self, train_batch: Tuple, batch_idx):
        x, y = train_batch
        x, y = x.float(), y.float()

        x = self.encoder(x)
        x = self.decoder(x)

        loss = self.loss(x, y)

        self.log("loss", loss)
        return {"loss": loss}

    def validation_step(self, val_batch: Tuple, batch_idx):
        x, y = val_batch
        x, y = x.float(), y.float()

        x = self.encoder(x)
        x = self.decoder(x)

        loss = self.loss(x, y)

        self.log("val_loss", loss)
# init and fill dict