import numpy as np

from DeepKoopmanModel import DeepKoopManConfig, DeepKoopMan
from AutoEncoderModel import AutoEncoder
import torch
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

torch.random.manual_seed(42)

config = DeepKoopManConfig(
    state_dim=13,
    latent_dim=64,
    hidden_dim=64
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torch.load('ModelLib/DeepKoopmanModel.pt', map_location="cpu")

Data = np.loadtxt('./DataLib/Data.csv', delimiter=',')
Data = Data[:, :config.state_dim]


DataMaxMin = np.loadtxt('./DataLib/DataMaxMin.csv', delimiter=',')
DataMax = DataMaxMin[0, :]
DataMin = DataMaxMin[1, :]

DataNor = 2*(Data-DataMin)/(DataMax-DataMin)-1
InputState = DataNor[:-1, :config.state_dim]
OutputState = DataNor[1:, :config.state_dim]

InputState = torch.from_numpy(InputState).to(torch.float32)

_, OutputState_pre = model(InputState)

OutputState_pre = OutputState_pre.detach().numpy()


for feature in np.arange(config.state_dim):
    plt.figure(feature)
    time_step = np.arange(OutputState.shape[0])
    plt.plot(time_step, OutputState[:, feature], color='blue', linestyle='-')
    plt.plot(time_step, OutputState_pre[:, feature], color='red', linestyle='--')
    plt.show()