import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np


class AutoEncoder(nn.Module):
    def __init__(self, state_dim, hidden_dim, latent_dim):
        super().__init__()

        self.state_dim = state_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim

        # 编码器 Encoder
        self.encoder = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, latent_dim)
        )

        # 解码器 Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, state_dim)
        )

    def forward(self, state):
        latent = self.encoder(state)
        state_pre = self.decoder(latent)
        return state_pre

    def autoencoder_train(self, batch_size, epochs, lr, datax_train, datay_train,
                          datax_val, datay_val):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.to(device)
        datax_train, datay_train = datax_train.to(device), datay_train.to(device)
        datax_val, datay_val = datax_val.to(device), datay_val.to(device)

        train_dataset = torch.utils.data.TensorDataset(datax_train, datay_train)
        val_dataset = torch.utils.data.TensorDataset(datax_val, datay_val)

        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=True)

        lossfunc = nn.MSELoss()
        op = optim.Adam(self.parameters(), lr=lr)

        TrainLoss = []
        ValLoss = []

        for epoch in range(epochs):
            train_loss = 0
            for X_batch, Y_batch in train_loader:
                X_batch, Y_batch = X_batch.to(device), Y_batch.to(device)

                outputs = self.forward(X_batch)
                loss = lossfunc(outputs, Y_batch)
                train_loss += loss.item()

                op.zero_grad()

                loss.backward()

                op.step()

            val_loss = 0
            with torch.no_grad():
                for X_batch, Y_batch in val_loader:
                    X_batch, Y_batch = X_batch.to(device), Y_batch.to(device)

                    outputs = self.forward(X_batch)
                    loss = lossfunc(outputs, Y_batch)
                    val_loss += loss.item()

            TrainLoss.append(train_loss / len(train_loader))
            ValLoss.append(val_loss / len(val_loader))

            train_log_loss = np.log10(train_loss / len(train_loader))
            val_log_loss = np.log10(val_loss / len(val_loader))

            print(
                f"LatentDim[{self.latent_dim}], HiddenDim[{self.hidden_dim}], LR[{lr}], Epoch [{epoch + 1}/{epochs}], "
                f"Train Loss: {train_log_loss}, Validation Loss: {val_log_loss}")

        return TrainLoss, ValLoss