import torch
from KoopmanESNModel import KoopmanESNConfig, KoopmanESN
from BaseModel.ESNModel import ESN
from DataProcess import DataLoad
import numpy as np
import matplotlib.pyplot as plt

KoopmanModel = torch.load('./ModelLib/DeepKoopmanModel.pt', map_location="cpu")
encoder = KoopmanModel.encoder
decoder = KoopmanModel.decoder
A = KoopmanModel.A.detach().numpy()
state_dim = KoopmanModel.state_dim
latent_dim = KoopmanModel.latent_dim

lr = 0.5
sr = 0.9
sp = 0.1
ridge = 1e-7
train_start = 15000
train_len = 5000
train_warmup = 100
predict_warmup = 100
predict_len = 5000

Koopmanconfig = KoopmanESNConfig(
    units=latent_dim,
    lr=lr,
    sr=sr,
    sp=sp,
    ridge=ridge,
    train_start=train_start,
    train_len=train_len,
    train_warmup=train_warmup,
    predict_warmup=predict_warmup,
    predict_len=predict_len,
    state_dim=latent_dim
)

ESNconfig = KoopmanESNConfig(
    units=500,
    lr=lr,
    sr=sr,
    sp=sp,
    ridge=ridge,
    train_start=train_start,
    train_len=train_len,
    train_warmup=train_warmup,
    predict_warmup=predict_warmup,
    predict_len=predict_len,
    state_dim=state_dim
)

KoopESNModel = KoopmanESN(
    config=Koopmanconfig,
    A=A
)
ESNModel = ESN(
    config=ESNconfig
)
_, DataTrainKoop, DataWarmKoop, DataValKoop = DataLoad('./DataLib/LatentNor.csv', Koopmanconfig)
# KoopmanESN
KoopESNModel.koopmanesn_train(DataTrainKoop)
DataPreKoop = KoopESNModel.predict(DataWarmKoop)

DataPreKoop_Ten = torch.from_numpy(DataPreKoop).float()
DataValKoop_Ten = torch.from_numpy(DataValKoop).float()
DataWarmKoop_Ten = torch.from_numpy(DataWarmKoop).float()

StatePreKoop_Ten = decoder(DataPreKoop_Ten)
StateValKoop_Ten = decoder(DataValKoop_Ten)
StateWarmKoop_Ten = decoder(DataWarmKoop_Ten)

StatePreKoop = StatePreKoop_Ten.detach().numpy()
StateValKoop = StateValKoop_Ten.detach().numpy()
StateWarmKoop = StateWarmKoop_Ten.detach().numpy()

KoopMSE = np.linalg.norm(StatePreKoop - StateValKoop, ord='fro') ** 2 / np.prod(StateValKoop.shape)
print(KoopMSE)

t = np.arange(Koopmanconfig.predict_warmup + Koopmanconfig.predict_len)

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

StateWarmKoop = (StateWarmKoop + 1) * (DataMax-DataMin) + DataMin
StatePreKoop = (StatePreKoop + 1) * (DataMax-DataMin) + DataMin
StateValKoop = (StateValKoop + 1) * (DataMax-DataMin) + DataMin

for fea in np.arange(state_dim):
    plt.figure(fea)
    plt.plot(t[:Koopmanconfig.predict_warmup], StateWarmKoop[:, fea])
    plt.plot(t[-Koopmanconfig.predict_len:], StatePreKoop[:, fea],
             linestyle="--", color='red', label='KoopmanESNPre')
    plt.plot(t[-Koopmanconfig.predict_len:], StateValKoop[:, fea],
             linestyle="-", color='blue', label='StateReal')
    plt.legend()
    plt.show()