from IPython import get_ipython
get_ipython().magic('reset -sf')
import matplotlib.pyplot as plt
import numpy
import scipy.io as sio
from tvb.simulator. lab import *
import pandas
import tvb.datatypes.projections as projections
import random
import statistics
import scipy.stats
import os
from parfor import parfor
import pickle
import time
files = os.listdir("../Preprocessing/Preprocessed EEG CSV")
for file in files:
    print(file)
    original_EEG = pandas.read_excel("../Preprocessing/Preprocessed EEG CSV/" + file, header = None)
    original_EEG = original_EEG.loc[:,1:20] # Remove AFz
    original_FC = numpy.corrcoef(original_EEG,rowvar=False)
    original_FC = numpy.triu(original_FC)
    sampling_rate = 1200
    t = list(range(1,16,1))
    G = list(range(0,62,2))
    for iteration_t in list(range(0,len(t))):
        @parfor(list(range(0,len(G))))
        def fun(iteration_G):
            result = pandas.DataFrame(columns = ['tau', 'p_value'],index=range(0,1))
            random.seed(99)
            def_conn = connectivity.Connectivity.from_file('F:\OneDrive - Kormányzati Informatikai Fejlesztési Ügynökség\Charite\EEG-EMG Parkinson\Simulations\connectivity_76.zip')
            def_conn.scaled_weights(mode='tract')
            def_conn.configure()
            ss = sensors.SensorsEEG.from_file('F:\OneDrive - Kormányzati Informatikai Fejlesztési Ügynökség\Charite\EEG-EMG Parkinson\Simulations\eeg_unitvector_19.txt')
            pr = projections.ProjectionSurfaceEEG(projection_data=sio.loadmat('projection_eeg_19_surface_16k.mat')['newProjectionMatrix'])
            rm = region_mapping.RegionMapping(array_data=numpy.array(numpy.loadtxt('regionMapping_16k_76.txt',dtype='int64')))
            mon_EEG = monitors.EEG(projection=pr, sensors=ss, region_mapping=rm, period=1000/sampling_rate)
            mon_tavg = monitors.TemporalAverage(period=100.)
            what_to_watch = (mon_tavg,mon_EEG)
            integrator = integrators.EulerDeterministic(dt=0.1)
            oscillator = models.Generic2dOscillator(tau = numpy.array([t[iteration_t]]))
            global_coupling = coupling.Linear(a=numpy.array(G[iteration_G]))
            def_conn.weights = numpy.divide(def_conn.weights, numpy.abs(numpy.max(def_conn.weights)))
            weights = numpy.mean(def_conn.weights)
            sim = simulator.Simulator(model=oscillator, connectivity=def_conn,coupling=global_coupling,integrator=integrator,
            monitors=what_to_watch,simulation_length=1000*(seconds+1),conduction_speed = weights/t[iteration_t])
            sim.configure()
            (ttavg, tavg),(tEEG,EEG) = sim.run(simulation_length=1000*(seconds+1))
            EEG  = EEG[sampling_rate:EEG.shape[0]:,0,:,0] # Remove first second as Stefanovski 2019
            simulated_FC = numpy.corrcoef(EEG,rowvar=False)
            simulated_FC = numpy.triu(simulated_FC)
            tau, p_value = scipy.stats.kendalltau(original_FC, simulated_FC)
            result['tau'] = tau
            result['p_value'] = p_value
            fun = numpy.array([tau,p_value])
            return fun
        name = "Parameter Exploration Results/" + file.replace(".xlsx","_t_is_") + str(t[iteration_t]) + ".pkl"
        with open(name, 'wb') as f:  # Python 3: open(..., 'wb')
            pickle.dump(fun, f)