import matplotlib.pyplot as plt
import numpy as np

import aux
from helpers import data_management as dm
import modules.classifier as clf
import analytics



# sensory mapping
states = ['SchliesseHand','BeugeRechtenMittelfinger', 'BeugeRechtenZeigefinger','BeugeRechtenDaumen','OeffneHand',
  'StreckeRechtenMittelfinger','StreckeRechtenZeigefinger','StreckeRechtenDaumen']

# motor mapping
# states = ['rechte_hand', 'rechter_daumen', 'linker_daumen', 'zunge', 'fuesse']



params = aux.load_config()

data_tot, tt, triggers_tot, ch_rec_list, file_names = dm.get_raw(n_triggers=params.classifier.n_classes, exploration=True)


psth_win = [-20,80]
n_trials = 10

psth_tot = np.zeros((len(states),0, np.diff(psth_win)[0], data_tot[0,0].shape[1])) 				# state x #stimuli x psth_window x #channels
psth = np.zeros((len(states),n_trials, np.diff(psth_win)[0], data_tot[0,0].shape[1])) 				# state x #stimuli x psth_window x #channels

for fids in range(data_tot.shape[0]):
	data = data_tot[fids,0]
	triggers = triggers_tot[fids,0][0]

	for state in range(len(triggers)):
		if triggers[state].size>0:
			for ii in range(triggers[state].size):
				jj = triggers[state][0,ii]
				if jj+(psth_win[0])>=0 and jj+(psth_win[1])<=data.shape[0]:
					psth[state,ii,:,:] = data[jj+psth_win[0]:jj+psth_win[1],:]

				print(state,ii,jj)
	print(file_names[fids])
	psth_tot = np.concatenate((psth_tot, psth), axis=1)

psth_xx = np.arange(psth_win[0],psth_win[1])
col = ['C0','C1','C2','C3','C4','C5','C6','C7','C8']
plt.figure(1)

# plt.subplot(321)
# for ii in range(psth.shape[0]):
# 	plt.plot(np.mean(psth[ii,0:,:,32], axis=0).T,col[ii])

# plt.subplot(322)
# for ii in range(psth.shape[0]):
# 	plt.plot(np.mean(psth[ii,0:,:,33], axis=0).T,col[ii])

plt.figure(1, figsize=[19.2 ,  9.55])
# plt.ioff()
# for ch_id in range(0,128):
for ch_id in range(87,88):
	plt.clf()
	print(f'ch_id: {ch_id}')
	for ii in range(psth_tot.shape[0]):
	# for ii in [1, 2, 3, 5, 6, 7,0,:

		ax = plt.subplot(3,4,ii+1)
		plt.gca().title.set_text(f'{states[ii]} n={psth_tot.shape[1]}')
		mu1 = np.mean(psth_tot[ii,0:,:,ch_id:ch_id+1], axis=2)		# average accross channels
		plt.plot(psth_xx, mu1.T, color=col[ii], alpha=0.2)
		plt.plot(psth_xx, np.median(mu1, axis=0), color='k',lw=2,alpha=0.5)
		plt.plot(psth_xx, np.mean(mu1, axis=0), color=col[ii],lw=2)
		plt.ylim(0,20)
		ax.set_xticks([])

	for ii in range(4):
		ax=plt.subplot(3,4,8+ii+1)
		mu1 = np.mean(psth_tot[ii,0:,:,ch_id:ch_id+1], axis=2)
		mu2 = np.mean(psth_tot[ii+4,0:,:,ch_id:ch_id+1], axis=2)
		plt.plot(psth_xx, np.mean(mu1, axis=0), color=col[ii], lw=3)
		plt.plot(psth_xx, np.mean(mu2, axis=0), color=col[ii+4], lw=3)
		plt.plot(psth_xx, np.median(mu1, axis=0), 'k--', lw=2, alpha=0.6)
		plt.plot(psth_xx, np.median(mu2, axis=0), 'k--', lw=2, alpha=0.6)
		# plt.ylim(0,20)

	plt.savefig('/kiap/src/data/results/sensor_mapping/'+f'ch_id_{ch_id}.png')
	# plt.savefig('/data/clinical/neural_new/2019-03-23/results/'+f'ch_id_{ch_id}.png')
	# plt.savefig('/data/clinical/neural_new/2019-03-26/results/'+f'ch_id_{ch_id}.png')

# plt.subplot(324)
# for ii in range(psth.shape[0]):
# 	plt.plot(np.mean(psth[ii,0:,:,0], axis=0).T,col[ii])

# plt.subplot(325)
# for ii in range(psth.shape[0]):
# 	plt.plot(np.mean(psth[ii,1:,:,0:32], axis=2).T,col[ii])

# plt.subplot(326)
# for ii in range(psth.shape[0]):
# 	plt.plot(np.mean(psth[ii,1:,:,32:96], axis=2).T,col[ii])

# print('\nsubplot1: channel 0, averaged across stimuli')
# print('subplot2: all stimuli, averaged across channels 0-64')
# print('subplot3: all stimuli, averaged across channels 64-128')

# plt.legend()
plt.show()