#!/usr/bin/env python
# coding: utf-8

""" Learning Koopman Invariant Subspace
 (c) Naoya Takeishi, 2017.
 takeishi@ailab.t.u-tokyo.ac.jp
"""

import numpy as np
np.random.seed(1234567890)

from argparse import ArgumentParser
from os import path
import time

from lkis import TimeSeriesBatchMaker, KoopmanInvariantSubspaceLearner
from losses import combined_loss
from torch import device, save, manual_seed
from torch.optim import SGD

import matplotlib.pyplot as plt
import seaborn as sns

# -- Parse arguments
t = time.time()
parser = ArgumentParser(description='Learning Koopman Invariant Subspace (Now with PyTorch!)')
parser.add_argument("--name", "-n", type=str, default=f"lkis-{int(time.time())}", help="name of experiment")
parser.add_argument("--data-path", type=str, default="./train.npy", help="time-series data to model")
parser.add_argument("--epochs", "-e", type=int, default=1000, help="number of epochs to train for")
parser.add_argument("--num-batches", "-b", type=int, default=1, help="how many batchs for break the data up into")
parser.add_argument("--gpu", action="store_true", default=False, help="use a GPU or no")
parser.add_argument("--intermediate-observable", "-i", type=int, default=-1, help="intermediate dimensional observation space")
parser.add_argument("--save-model", "-m", action="store_true", default=False, help="whether or not you want the model saved to $name$.torch.mdl")
parser.add_argument("--save-training-plot", "-p", action="store_true", default=False, help="where to save plotting")
parser.add_argument("--max-lag", "-l", type=int, default=-1, help="maximum_lag")
parser.add_argument("--state-space", "-s", type=int, default=1, help="dimensionality of the underlying state space")
parser.add_argument("--alpha", "-a", type=float, default=1.0, help="value to score the reconstruction loss by")
parser.add_argument("--learning-rate", "-r", type=float, default=0.001, help="Optimizer learning rate")
parser.add_argument("--validation-data-path", "-v", type=str, default="")

#ToDo: Implement
parser.add_argument("--dmd", action="store_true", default=False, help="Execute and save the DMD on the training set")

if __name__ == "__main__":
    # grab the command line arguments
    cli_args = parser.parse_args()
    manual_seed(216)

    # find and load the training data
    data_path = cli_args.data_path
    print(f"Loading training data from {data_path}")
    data_train = np.load(data_path)
    if len(data_train.shape) == 1:
        data_train = data_train.reshape(-1, 1)
    print(f"Loaded a dataset with dimension: {data_train.shape}")
    validate = cli_args.validation_data_path != ""
    data_val = None
    if validate:
        data_path = cli_args.validation_data_path
        print(f"Loading validation data from {data_path}")
        data_val = np.load(data_path)

    # process the delay either set by the user or is set to one 10th of the data
    delay = cli_args.max_lag if cli_args.max_lag > 0 else (data_train.shape[0] // 10)

    # based on the number of batches, delay, and size of the data compute the samples per batch
    samples_per_batch = (data_train.shape[0] - delay) // cli_args.num_batches

    # construct the data preparer
    batch_iterator = TimeSeriesBatchMaker(
        y=data_train,
        batch_size=samples_per_batch,
        max_lag=delay
    )
    if validate:
        val_batch_iterator = TimeSeriesBatchMaker(
            y=data_val,
            max_lag=delay
        )

    # construct the end-to-end model
    lkis = KoopmanInvariantSubspaceLearner(
        observable_dim=data_train.shape[1],
        latent_dim=cli_args.state_space,
        intermediate_observable=cli_args.intermediate_observable,
        delay=delay
    )

    if cli_args.gpu:
        device = device("cuda")

    # initialize the optimizer
    optimizer = SGD(lkis.parameters(), lr=cli_args.learning_rate)
    losses = []
    val_losses = []
    for epoch in range(cli_args.epochs):
        loss = 0
        for b in range(cli_args.num_batches):
            optimizer.zero_grad()
            time_delayed_ys, y_true = next(batch_iterator)

            if cli_args.gpu:
                time_delayed_ys.to(device)
                y_true.to(device)

            g_pred, y_pred = lkis(time_delayed_ys)
            g_0 = g_pred[:-1]
            g_1 = g_pred[1:]

            batch_loss = combined_loss(y_pred=y_pred, y_true=y_true, g_0=g_0, g_1=g_1)

            batch_loss.backward()

            optimizer.step()

            loss += batch_loss.item()

        # display the epoch training loss
        print(f"epoch : {epoch + 1}/{cli_args.epochs}, loss = {loss:.6f}")
        losses.append(loss)

        if validate:
            y_time_delayed_val, y_true = next(val_batch_iterator)
            if cli_args.gpu:
                y_time_delayed_val.to(device)
                y_true.to(device)

            g_pred, y_pred = lkis(y_time_delayed_val)
            g_0 = g_pred[:-1]
            g_1 = g_pred[1:]

            batch_loss = combined_loss(y_pred=y_pred, y_true=y_true, g_0=g_0, g_1=g_1)
            val_loss = batch_loss.item()
            print(f"\tval-loss  = {val_loss:.6f}")
            val_losses.append(val_loss)

    if cli_args.save_model:
        save(lkis, f"{cli_args.name}.torch.mdl")

    if cli_args.save_training_plot:
        sns.lineplot(x=list(range(cli_args.epochs)), y=losses, label="training loss")
        if validate:
            sns.lineplot(x=list(range(cli_args.epochs)), y=val_losses, label="validation loss")
        plt.xlabel("Epochs")
        plt.ylabel("Combined Reconstruction and DMD Loss")
        plt.title(f"Training Loss for {cli_args.name}")
        plt.savefig(f"{cli_args.name}-training-loss.png")

[(10, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (25, 'time.time', 'time.time', False, 'import time\n'), (26, 'argparse.ArgumentParser', 'ArgumentParser', False, 'from argparse import ArgumentParser\n'), (47, 'torch.manual_seed', 'manual_seed', False, 'from torch import device, save, manual_seed\n'), (52, 'numpy.load', 'np.load', True, 'import numpy as np\n'), (70, 'lkis.TimeSeriesBatchMaker', 'TimeSeriesBatchMaker', False, 'from lkis import TimeSeriesBatchMaker, KoopmanInvariantSubspaceLearner\n'), (82, 'lkis.KoopmanInvariantSubspaceLearner', 'KoopmanInvariantSubspaceLearner', False, 'from lkis import TimeSeriesBatchMaker, KoopmanInvariantSubspaceLearner\n'), (61, 'numpy.load', 'np.load', True, 'import numpy as np\n'), (76, 'lkis.TimeSeriesBatchMaker', 'TimeSeriesBatchMaker', False, 'from lkis import TimeSeriesBatchMaker, KoopmanInvariantSubspaceLearner\n'), (138, 'torch.save', 'save', False, 'from torch import device, save, manual_seed\n'), (144, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (145, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (146, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (147, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n'), (110, 'losses.combined_loss', 'combined_loss', False, 'from losses import combined_loss\n'), (132, 'losses.combined_loss', 'combined_loss', False, 'from losses import combined_loss\n'), (27, 'time.time', 'time.time', False, 'import time\n')]