import numpy as np
import pandas as pd



data_folder = "../data/"
stim_file = "Stiminfo_PVCre_2021_0012_s06_e14.csv"
# Stimulus DataFrame
stim = pd.read_csv(data_folder+stim_file)
num_trial = len(stim)
stim_val = {}
trial_stim_id = {}
stim_val['pair'], trial_stim_id['pair'] = np.unique(stim[['grat_orientation', 'grat_phase']], return_inverse=True, axis=0) # 50 trials per orientation-phase pair
stim_val['orientation'], trial_stim_id['orientation'] = np.unique(stim['grat_orientation'], return_inverse=True) # 1000 trials per orientation
stim_val['phase'], trial_stim_id['phase'] = np.unique(stim['grat_phase'], return_inverse=True) # 1000 trials per phase
key_list = ['pair', 'orientation', 'phase']
stim_id_trial = {}
num_stim = {}
for key in key_list:
    stim_id_trial[key] = [np.where(trial_stim_id[key] == i)[0] for i in range(len(stim_val[key]))]
    num_stim[key] = len(stim_val[key])

# stim_id_trial[i] returns an array of trial indices where the stimulus has the value of the ith stimulus.
def pair_id(i):
    return (int(i//num_stim['orientation']), i%num_stim['orientation'])

trial_pair_id = np.array([pair_id(i) for i in trial_stim_id['pair']], dtype=object)
pair_val = stim_val['pair'].reshape(num_stim['orientation'], num_stim['phase'],2)

# for each (orientation_id = i, phase_id = j) find the trial indices
pair_trial_id = np.ndarray((num_stim['orientation'], num_stim['phase']), dtype=object)
for i in range(num_stim['pair']):
    pair_trial_id[pair_id(i)] = stim_id_trial['pair'][i]

# s.close()

stim_data = {'stim_val':stim_val, 'trial_stim_id':trial_stim_id, 'key_list':key_list, 'num_trial':num_trial,
             'trial_pair_id':trial_pair_id, 'pair_val':pair_val, 'pair_trial_id':pair_trial_id, 'stim_id_trial':stim_id_trial, 'num_stim':num_stim}

import pickle
file = open('stim_data.pkl', 'wb')
pickle.dump(stim_data, file)
file.close()