import matplotlib.pyplot as plt
import numpy as np
from sklearn.pipeline import Pipeline
from sklearn.metrics import mean_squared_error
from swimnetworks import (Dense, Linear)

from swim_rnn import RNN
from utils import D_stsp
from datafold import TSCDataFrame
import time
from utilities.pse import *

np.random.seed(0)

# hyperparams
layer_width = 300
regularization_constant = 1e-4

rng = np.random.default_rng(5)

dt = 0.01

# data-loading
training_data_npz = np.load('datasets/Rossler_normalized_train.npy')
validation_data_npz = np.load('datasets/Rossler_normalized_valid.npy')
test = True

training_tsc_list = []
num_training_traj = training_data_npz.shape[2]
for i in range(num_training_traj):
    training_tsc_list.append(TSCDataFrame.from_array(training_data_npz[:, :, i], feature_names=list(str(i) for i in range(training_data_npz.shape[1]))))
training_data = TSCDataFrame.from_frame_list(training_tsc_list)

validation_tsc_list = []
num_validation_traj = validation_data_npz.shape[2]
for i in range(num_validation_traj):
    validation_tsc_list.append(TSCDataFrame.from_array(validation_data_npz[:, :, i], feature_names=list(str(i) for i in range(validation_data_npz.shape[1]))))
validation_data = TSCDataFrame.from_frame_list(validation_tsc_list)

steps = [
    ("hidden", Dense(layer_width=layer_width, activation='tanh', parameter_sampler='tanh',
                     random_seed=rng.integers(100, size=1))),
]
network_dictionary = Pipeline(steps=steps)

swim_rnn = RNN(network_dictionary, regularization_constant)

start_time = time.time()
swim_rnn.fit(training_data)
end_time = time.time()

elapsed_time = end_time - start_time
print("Fit time: ", elapsed_time)

# # evaluation on training set
# training_pred = []
# training_mse = []
# training_D_stsp = []
# training_D_h = []
# for i in range(num_training_traj):
#     pred = swim_rnn.predict(training_tsc_list[i].initial_states(),
#                             time_values=training_tsc_list[i].time_values())
#     training_pred.append(pred)
#     training_mse.append(mean_squared_error(pred, training_tsc_list[i]))
#     training_D_stsp.append(D_stsp(X_pred=pred, X_true=training_tsc_list[i].to_numpy()))
#     training_D_h.append(power_spectrum_error(x_gen=pred, x_true=training_tsc_list[i].to_numpy()))
# training_D_stsp_np = np.array(training_D_stsp)
# training_D_stsp_np[training_D_stsp_np < 1e-6] = 0
# print(f'MSE averaged over training trajectories: \t {np.mean(training_mse)}')
# print(f'D_stsp averaged over training trajectories: \t {np.mean(training_D_stsp_np)}')
# print(f'D_h averaged over training trajectories: \t {np.mean(training_D_h)}')

# evaluation on validation set
validation_pred = []
validation_mse = []
validation_D_stsp = []
validation_D_h = []
for i in range(num_validation_traj):
    pred = swim_rnn.predict(validation_tsc_list[i].initial_states(),
                            time_values=validation_tsc_list[i].time_values())
    validation_pred.append(pred)
    validation_mse.append(mean_squared_error(pred, validation_tsc_list[i]))
    validation_D_stsp.append(D_stsp(X_pred=pred, X_true=validation_tsc_list[i].to_numpy()))
    validation_D_h.append(power_spectrum_error(x_gen=pred, x_true=validation_tsc_list[i].to_numpy()))
validation_D_stsp_np = np.array(validation_D_stsp)
validation_D_stsp_np[validation_D_stsp_np < 1e-8] = 0
print('validation dstsp max=',max(validation_D_stsp))
print(f'MSE averaged over validation trajectories: \t {np.mean(validation_mse)}')
print(f'D_stsp averaged over validation trajectories: \t {np.mean(validation_D_stsp_np)}')
print(f'D_h averaged over validation trajectories: \t {np.mean(validation_D_h)}')
print(f'Max Dstsp={np.max(validation_D_stsp_np)}')

# num_plot_traj = 1
#
# fig, ax = plt.subplots(1, 1, figsize=(8, 4))
# for k in rng.integers(0, num_validation_traj, size=num_plot_traj):
#     ax.plot(np.arange(validation_tsc_list[i].shape[0])[:20000] * dt, validation_tsc_list[k].to_numpy()[:20000], '-', label='true')
#     ax.plot(np.arange(validation_tsc_list[i].shape[0])[:20000] * dt, validation_pred[k][:20000], '--', label='prediction')
# ax.set_xlabel('$t$')
# ax.legend()
# plt.savefig(f'Rossler_validation_{num_plot_traj}_trajectories.pdf')
# plt.show()

# fig = plt.figure(figsize=(4, 4))
# ax = fig.add_subplot(1, 1, 1, projection='3d')
# for k in rng.integers(0, num_validation_traj, size=num_plot_traj):
#     ax.plot(*validation_tsc_list[k].to_numpy().T, '-', label='true')
#     ax.plot(*validation_pred[k].T, '--', label='prediction')
# ax.set_xlabel('$h_1$')
# ax.set_ylabel('$h_2$')
# ax.set_zlabel('$h_3$')
# # ax.set_aspect("equal")
# ax.legend()
# plt.savefig(f'Rossler_validation_{num_plot_traj}_phase_space.pdf')
# plt.show()


if test:
    test_data_npz = np.load('datasets/Rossler_normalized_test.npy')

    test_tsc_list = []
    num_test_traj = test_data_npz.shape[2]
    for i in range(num_test_traj):
        test_tsc_list.append(TSCDataFrame.from_array(test_data_npz[:, :, i], feature_names=list(str(i) for i in range(test_data_npz.shape[1]))))
    test_data = TSCDataFrame.from_frame_list(test_tsc_list)

    # evaluation on test set
    test_pred = []
    test_mse = []
    test_D_stsp = []
    test_D_h = []
    for i in range(num_test_traj):
        pred = swim_rnn.predict(test_tsc_list[i].initial_states(), time_values=test_tsc_list[i].time_values())
        test_pred.append(pred)
        test_mse.append(mean_squared_error(pred, test_tsc_list[i]))
        test_D_stsp.append(D_stsp(X_pred=pred, X_true=test_tsc_list[i].to_numpy()))
        test_D_h.append(power_spectrum_error(x_gen=pred, x_true=test_tsc_list[i].to_numpy()))
    test_D_stsp_np = np.array(test_D_stsp)
    test_D_stsp_np[test_D_stsp_np < 1e-8] = 0
    print(f'MSE averaged over test trajectories: \t {np.mean(test_mse)}')
    print(f'D_stsp averaged over test trajectories: \t {np.mean(test_D_stsp_np)}')
    print(f'D_h averaged over test trajectories: \t {np.mean(test_D_h)}')
    print(f'Max Dstsp={np.max(test_D_stsp_np)}')

# seed = 1:
# Fit time:  5.140249967575073
# validation dstsp max= 0.0016573760554785823
# MSE averaged over validation trajectories: 	 0.02940165592784143
# D_stsp averaged over validation trajectories: 	 0.00043070953598153063
# D_h averaged over validation trajectories: 	 0.22717388872982058
# Max Dstsp=0.0016573760554785823
# MSE averaged over test trajectories: 	 0.03132171539452584
# D_stsp averaged over test trajectories: 	 0.0003819032044585863
# D_h averaged over test trajectories: 	 0.2250928177344255
# Max Dstsp=0.002352728724678885
#
# seed = 2:
# Fit time:  4.650840520858765
# validation dstsp max= 0.0007327413054236326
# MSE averaged over validation trajectories: 	 0.020732687055289287
# D_stsp averaged over validation trajectories: 	 0.00012204299519220334
# D_h averaged over validation trajectories: 	 0.1958096841462998
# Max Dstsp=0.0007327413054236326
# MSE averaged over test trajectories: 	 0.01717028125450918
# D_stsp averaged over test trajectories: 	 0.00011852507068670404
# D_h averaged over test trajectories: 	 0.18399414158402116
# Max Dstsp=0.0010622615779596375
#
# seed = 3:
# Fit time:  4.391016960144043
# validation dstsp max= 0.0010099983404533728
# MSE averaged over validation trajectories: 	 0.026307814573230802
# D_stsp averaged over validation trajectories: 	 0.0001402812143837675
# D_h averaged over validation trajectories: 	 0.15190443958772723
# Max Dstsp=0.0010099983404533728
# MSE averaged over test trajectories: 	 0.030581061540880805
# D_stsp averaged over test trajectories: 	 0.00013350165817390712
# D_h averaged over test trajectories: 	 0.15270656208549396
# Max Dstsp=0.0008438150976152998
#
# seed = 4:
# Fit time:  6.3937599658966064
# validation dstsp max= 0.000738871472425535
# MSE averaged over validation trajectories: 	 0.021223911963957152
# D_stsp averaged over validation trajectories: 	 0.00013320260882066198
# D_h averaged over validation trajectories: 	 0.1337950369782515
# Max Dstsp=0.000738871472425535
# MSE averaged over test trajectories: 	 0.018416397544973498
# D_stsp averaged over test trajectories: 	 9.314530620796075e-05
# D_h averaged over test trajectories: 	 0.12625251318801883
# Max Dstsp=0.0007898208744880956
#
# seed = 5:
# Fit time:  6.206273794174194
# validation dstsp max= 0.0005567238364252874
# MSE averaged over validation trajectories: 	 0.02092555369935451
# D_stsp averaged over validation trajectories: 	 8.798172165121078e-05
# D_h averaged over validation trajectories: 	 0.10897360207036628
# Max Dstsp=0.0005567238364252874
# MSE averaged over test trajectories: 	 0.01711909140947611
# D_stsp averaged over test trajectories: 	 5.8553638579771554e-05
# D_h averaged over test trajectories: 	 0.1053761889889552
# Max Dstsp=0.0006089801950183752
#
#
# mean fit time: 5.356428241729736
# mean d_stsp on test data: 0.00015712577562138598
# mean d_h on test data: 0.16017891537385517
