from DeepKoopmanModel import DeepKoopManConfig, DeepKoopMan
from AutoEncoderModel import AutoEncoder
import torch
from sklearn.model_selection import train_test_split

torch.random.manual_seed(42)

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

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

AutoEncoderModel = torch.load('./AutoEncoder/AutoEncoderModel.pt', map_location="cpu")
AutoEncoderModel = AutoEncoderModel.to(device)
Encoder = AutoEncoderModel.encoder
Decoder = AutoEncoderModel.decoder

model = DeepKoopMan(
    config=config,
    encoder=Encoder,
    decoder=Decoder
)

test_size = 0.2
val_size = 0.2

batch_size = 256
epochs = 100
lr = 1e-3

Data = torch.load('./DataLib/DataSeqTensor.pt')

# Data = Data[:1000, :, :]

Data_temp, Data_train = train_test_split(Data, test_size=1 - test_size - val_size, random_state=42)
Data_test, Data_val = train_test_split(Data_temp, test_size=test_size / (test_size + val_size),
                                       random_state=42)

TrainLoss, ValLoss = model.deepkoopman_train(
    batch_size=batch_size,
    epochs=epochs,
    lr=lr,
    data_train=Data_train,
    data_val=Data_val
)

torch.save(model, './ModelLib/DeepKoopmanModel.pt')