from sklearn.cross_decomposition import CCA

import numpy as np

import pickle

from scipy.io import loadmat

from sklearn.decomposition import PCA

import matplotlib.pyplot as plt

from scipy.ndimage import gaussian_filter1d

from sklearn.linear_model import LinearRegression, Ridge

from sklearn.metrics import r2_score

from scipy.spatial import procrustes

from sklearn.decomposition import PCA

import sys

sys.path.insert(0, '../SAC/')

import kinematics_preprocessing_specs

import config

parser = config.config_parser()

args, unknown = parser.parse_known_args()

#Load the test data of nusim

with open('../test_data/test_data.pkl', 'rb') as file:

    test_data = pickle.load(file)

print(test_data.keys())

#Get the timepoints of each condition per cycle

with open('../kinematics_data/kinematics.pkl', 'rb') as file:

    kin_train_test = pickle.load(file)

kin_train = kin_train_test['train']

kin_test = kin_train_test['test']

#First update the keys of self.kin_test

for cond in range(len(kin_test)):

    kin_test[len(kin_train) + cond] = kin_test.pop(cond)

kin = kin_train

kin.update(kin_test)

conds = [kin[cond].shape[-1] for cond in range(len(kin))]

total_conds = len(conds)

#Select the cycle for each condition (training conditions followed by testing): 0 for 1st cycle and so on

#The number of elements should be equal to num_train_conditions + num_test_conditions

cycles = [2, 2, 2, 2, 2, 2]

#Number of fixedsteps in the start of each condition

#Fix this {todo}: Get the values automatically from ..SAC/kinematics_preprocessing_specs.py

n_fixedsteps= args.n_fixedsteps

#Load the network activities

A_agent = []

for idx, cond_activity in test_data['rnn_activity'].items():

    act_agent = cond_activity

    act_agent = act_agent[n_fixedsteps + cycles[idx] * conds[idx] : n_fixedsteps + (cycles[idx]+1) * conds[idx]]

    print(act_agent.shape)

    A_agent.append(act_agent[:, :])

#Do the collective PCA for all speeds

nusim_pca = PCA(n_components= 3)

A_agent_c = A_agent

#concatenate the musim activity for all conditions

for i_cond in range(len(A_agent_c)):

    if i_cond == 0:

        nusim_activity_pca = A_agent_c[i_cond]

    else:

        nusim_activity_pca = np.concatenate((nusim_activity_pca, A_agent_c[i_cond]), axis=0)

nusim_activity_pca = nusim_pca.fit_transform(nusim_activity_pca)

#Plot the PCA of the activities

colors = plt.cm.ocean(np.linspace(0,1,8))

ax = plt.figure(dpi=100).add_subplot(projection='3d')

prev_cond = 0

for i_cond in range(len(A_agent_c)):

    ax.plot(nusim_activity_pca[prev_cond:prev_cond+A_agent_c[i_cond].shape[0],0], 

            nusim_activity_pca[prev_cond:prev_cond+A_agent_c[i_cond].shape[0], 1], 

            nusim_activity_pca[prev_cond:prev_cond+A_agent_c[i_cond].shape[0], 2], color= colors[i_cond])

    prev_cond += A_agent_c[i_cond].shape[0]

# Hide grid lines

ax.grid(False)

plt.grid(b=None)

# Hide axes ticks

ax.set_xticks([])

ax.set_yticks([])

ax.set_zticks([])

plt.axis('off')

plt.show()