"""
Example for how to load data and display results on one subject's cortical surface.
"""

import matplotlib.pyplot as plt
import numpy as np
import utils
import os
from scipy.stats import zscore


# Get path to data files
fdir = os.path.abspath('../')
 
modality = "listening"
subject = "01"

# Load training data for subject 1, Listening dataset 
fname_trn = os.path.join(fdir, 'responses', 'subject{}_{}_fmri_data_trn.hdf'.format(subject, modality))
trndata = utils.load_data(fname_trn)
print(trndata.keys())

# Load validation data for subject 1, Listening dataset
# story_11 is the validation story ==> Check README
fname_val = os.path.join(fdir, 'responses', 'subject{}_{}_fmri_data_val.hdf'.format(subject, modality))
valdata = utils.load_data(fname_val, "story_11")["story_11"]

# Correlate repetitions across time points (per voxel)
valdata_corr = np.array(list(map(lambda c1,c2: (zscore(c1)*zscore(c2)).mean(0), valdata[0].T, valdata[1].T)))

# Map to subject flatmap
map_file = os.path.join(fdir, 'mappers', 'subject{}_mappers.hdf'.format(subject))
flatmap = utils.map_to_flat(valdata_corr, map_file)

# Plot flatmap
fig, ax = plt.subplots()
_ = ax.imshow(flatmap, cmap='inferno')
ax.axis('off')
plt.show()