import h5py
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import math as math

f = h5py.File('/Users/Krzysztof/PycharmProjects/untitled/venv/4studs.hdf5', 'r')

# f1 = f.get('eeg_subject_am1_trial_0_attended_female')
# f2 = f.get('envelope_subject_am1_trial_0_sex_female')
# f3 = f.get('eeg_subject_am1_trial_0_attended_male')
# f4 = f.get('envelope_subject_am1_trial_0_sex_male')
#
# f5 = f.get('eeg_subject_am1_trial_1_attended_female')
# f6 = f.get('envelope_subject_am1_trial_1_sex_female')
# f7 = f.get('envelope_subject_am1_trial_1_sex_male')
# eeg3 = np.array(f5)
# env3 = np.array(f6)
# env4 = np.array(f7)
#
# eeg2 = np.array(f3)
# env2 = np.array(f4)
#
# eeg1 = np.array(f1)
# env1 = np.array(f2)
# eeg1_w1 = tf.slice(eeg1, [1, 0], [1, 6250])
#
# array1 = np.zeros(((6211, 40)))
# for i in range(6211):
#     array1[i] = tf.slice(eeg1_w1, [0, i], [1, 40])
#
# array2 = np.zeros((6211,40))
# for i in range(6211):
#     array2[i] = tf.slice(env1, [i], [40])
# ###########
#
# ##TRIAL0
# eeg_tab = np.zeros((32,6250))
# for i in range(32):
#     eeg_tab[i] = tf.slice(eeg1, [i, 0], [1,6250])
# eeg = tf.keras.utils.normalize(eeg_tab)
#
# eeg_split = np.zeros((32,6211,40))
# for a in range(32):
#     for i in range(6211):
#         eeg_split[a,i] = tf.slice(eeg[a], [i], [40])
#
# ##TRIAL1
# eeg_tab1 = np.zeros((32,6250))
# for i in range(32):
#     eeg_tab1[i] = tf.slice(eeg2, [i, 0], [1,6250])
# eeg2_2 = tf.keras.utils.normalize(eeg_tab1)
#
# eeg_split2 = np.zeros((32,6211,40))
# for a in range(32):
#     for i in range(6211):
#         eeg_split2[a,i] = tf.slice(eeg2_2[a], [i], [40])
# out2 = tf.slice(env2, [0], [6211])
# new_out2 = np.zeros(((32,6211)))
# for i in range(32):
#     new_out2[i] = out2
#
# ##TRIAL3
# eeg_tab3 = np.zeros((32,6250))
# for i in range(32):
#     eeg_tab3[i] = tf.slice(eeg3, [i, 0], [1,6250])
# eeg3_3 = tf.keras.utils.normalize(eeg_tab3)
#
# eeg_split3 = np.zeros((32,6211,40))
# for a in range(32):
#     for i in range(6211):
#         eeg_split3[a,i] = tf.slice(eeg3_3[a], [i], [40])
# out3 = tf.slice(env3, [0], [6211])
# new_out3 = np.zeros(((32,6211)))
# for i in range(32):
#     new_out3[i] = out3
# out4 = tf.slice(env4, [0], [6211])
# new_out4 = np.zeros(((32,6211)))
# for i in range(32):
#     new_out4[i] = out4
# ###############
# eeg_w2 = tf.slice(eeg1, [2, 0], [1, 6250])
# eeg_w3 = tf.slice(eeg1, [3, 0], [1, 6250])
# array3 = np.zeros(((6211, 40)))
# array4 = np.zeros((6211,40))
#
# for i in range(6211):
#     array3[i] = tf.slice(eeg_w2, [0, i], [1, 40])
#     array4[i] = tf.slice(eeg_w3, [0, i], [1, 40])
#
# ex_1 = tf.keras.utils.normalize(array1)
# ex_2 = tf.keras.utils.normalize(array3)
# ex_3 = tf.keras.utils.normalize(array4)
# out1 = tf.slice(env1, [0], [6211])
# new_out = np.zeros(((32,6211)))
# for i in range(32):
#     new_out[i] = out1
#
# ###########
# output_1 = np.zeros((6211,1))
# for i in range(6211):
#     output_1[i] = tf.slice(env1, [i], [1])
#
# #dataset_train = tf.data.Dataset.from_tensor_slices((array1,output_1))
# #dataset_val = tf.data.Dataset.from_tensor_slices((ex_2,new_out))
#
# model = tf.keras.Sequential([
#     tf.keras.layers.Dense(64, activation='elu'),
#     tf.keras.layers.Dense(32, activation='elu'),
#     tf.keras.layers.Dense(1)
# ])
#
# model.compile(optimizer='sgd',loss='mean_squared_error')
#
#
# model.fit(x=eeg_split,y=new_out,epochs=2)
# model.fit(x=eeg_split2,y=new_out2,epochs=2)
#
# prediction = model.predict(eeg_split3[0])


#############################################
# znorm = tf.keras.utils.normalize(f1)
# t1 = tf.reshape(znorm, [1, 34, 6250, 1])
# k = tf.constant(1, shape=(34,40), dtype = 'float64')
# k1 = tf.reshape(k, [34, 40, 1, 1])
# k2 = tf.reshape(tf.constant(1, shape=(1,40), dtype='float64'), [1, 40, 1, 1])
# conv = tf.nn.conv2d(t1,k2,strides = [1, 1, 1, 1], padding='VALID')
# print('conv: ',conv.shape)
# conv_2 = tf.slice(conv, [0, 0, 0, 0], [1,32, 6211,1])
# print('conv2: ',conv_2.shape)
# shape = conv_2.get_shape().as_list()
# conv3 = tf.reshape(conv_2, [32, 6211,-1])
# print('conv3: ',conv3.shape)
# conv4 = tf.reshape(conv3, [32,6211])
# print(conv4.shape)

kern = tf.reshape(tf.constant(1, shape=(1,40), dtype='float64'), [1, 40, 1, 1])
#AM_TR0_FEMALE
f1 = f.get('eeg_subject_am1_trial_0_attended_female')
f1_norm = tf.keras.utils.normalize(f1)
f1_final = tf.reshape(f1_norm, [1, 34, 6250, 1])
f1_conv = tf.nn.conv2d(f1_final, kern, strides=[1,1,1,1], padding='VALID')
f1_conv2 = tf.slice(f1_conv, [0, 0, 0, 0], [1,32, 6211,1])
f1_conv3 = tf.reshape(f1_conv2, [32, 6211,-1])
f1_in = tf.reshape(f1_conv3, [32,6211])
f1_input = tf.transpose(f1_in)

f1_out1 = f.get('envelope_subject_am1_trial_0_sex_female')
f1_out = tf.slice(f1_out1, [0],[6211])

#AM_TR0_MALE
f2 = f.get('eeg_subject_am1_trial_0_attended_male')
f2_norm = tf.keras.utils.normalize(f2)
f2_final = tf.reshape(f2_norm, [1, 34, 6250, 1])
f2_conv = tf.nn.conv2d(f2_final, kern, strides=[1,1,1,1], padding='VALID')
f2_conv2 = tf.slice(f2_conv, [0, 0, 0, 0], [1,32, 6211,1])
f2_conv3 = tf.reshape(f2_conv2, [32, 6211,-1])
f2_in = tf.reshape(f2_conv3, [32,6211])
f2_input = tf.transpose(f2_in)

f2_out1 = f.get('envelope_subject_am1_trial_0_sex_male')
f2_out = tf.slice(f2_out1, [0],[6211])

#MODEL
model = tf.keras.Sequential([
     tf.keras.layers.Dense(128, activation='elu'),
     tf.keras.layers.Dense(1, activation='elu')
 ])
model.compile(optimizer='sgd',loss='mean_squared_error')
model.fit(input,output,epochs=10)
f.close()