
%% Code to generate wavelet coherence
%% S1

load('s1.mat')

% Meta information
n_days = 3;
electrode_list = [60, 69];
e_1 = electrode_list(1); e_2 = electrode_list(2);

% Loop through all days
for a = 1:n_days
    
    day_struct = c41.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        temp_data_notch = multi_iirnotch_filtering(temp_data, 1200, 60:60:240);

        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate coherence 
        % mscohere: do 70-90Hz as control for numerous reasons
        % cpsd: look @ phase lag between selected electrodes
        % plv: try to bin to get a shorter temporal resolution, not average
        % across 1s intervals
        % wcoherence: try varying time intervals to get this working?

        for f = 1:temp_n_trials
            
            cue_start_t = cue_segment(f, 1);
            cue_end_t = cue_segment(f, 2);
            bp_start_t = bp_segment(f, 1);
            bp_end_t = bp_segment(f, 2);
            post_bp_start_t = post_bp_segment(f, 1);
            post_bp_end_t = post_bp_segment(f, 2); 

            [wc_hg, ~, frq_hg, coi_hg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                'FrequencyLimits', [61, 300], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);  
            [wc_hg_raw, ~, ~, ~] = wcoherence(temp_data(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                'FrequencyLimits', [61, 300], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);
            [wc_lg, ~, frq_lg, coi_lg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                'FrequencyLimits', [1, 4], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);                          
            wccs.day(a).run(b).wc_hg(:, :, f) = wc_hg;
            wccs.day(a).run(b).wc_lg(:, :, f) = wc_lg;
            wccs.day(a).run(b).wc_hg_raw(:, :, f) = wc_hg_raw;
            wccs.hf_frq = frq_hg;
            wccs.lf_frq = frq_lg;

        end   
        
    end
    
end

%% S2

load('s2.mat')

% Meta information
n_days = 3;
electrode_list = [2, 3, 7, 8];
electrode_pairs = nchoosek(electrode_list, 2);

% % Select electrode pairs (will RUN OUT OF RAM IF I LOOP)
% y = 1;
% 
% %Set electrodes
% e_1 = electrode_pairs(y, 1);
% e_2 = electrode_pairs(y, 2);

for y = 1:size(electrode_pairs, 1)
    e_1 = electrode_pairs(y, 1);
    e_2 = electrode_pairs(y, 2);

    % Loop through all days
    for a = 1:n_days
        
        day_struct = c45.day(a);
        n_run = size(day_struct.data_out, 2);
        
        for b = 1:n_run
    
            temp_data = double(day_struct.data_out{1, b});
            temp_data_notch = multi_iirnotch_filtering(temp_data, 1200, [60:60:240]);
    
            % Load indices from specific run
            temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
            temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
            temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
            temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
    
            % Get total number of trials
            temp_n_trials = size(temp_stim_idxs, 2);
    
            % Get trials w/ specific cue/stimulus modalities, 
            tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
            no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
    
            % Loop across all trials and get time segments for PSD
    
            cue_segment = zeros(temp_n_trials, 2);
            bp_segment = zeros(temp_n_trials, 2);
            post_bp_segment = zeros(temp_n_trials, 2);
    
            for c = 1:temp_n_trials
    
                cue_segment(c, 1) = temp_cue_idxs(c);
                cue_segment(c, 2) = temp_cue_idxs(c) + 1200;
    
                bp_segment(c, 1) = temp_button_idxs(c) - 600;
                bp_segment(c, 2) = temp_button_idxs(c) + 600;
    
                post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
                post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
    
            end
    
            % Loop across all trials to calculate coherence 
            % mscohere: do 70-90Hz as control for numerous reasons
            % cpsd: look @ phase lag between selected electrodes
            % plv: try to bin to get a shorter temporal resolution, not average
            % across 1s intervals
            % wcoherence: try varying time intervals to get this working?
    
            for f = 1:temp_n_trials
    
                cue_start_t = cue_segment(f, 1);
                cue_end_t = cue_segment(f, 2);
                bp_start_t = bp_segment(f, 1);
                bp_end_t = bp_segment(f, 2);
                post_bp_start_t = post_bp_segment(f, 1);
                post_bp_end_t = post_bp_segment(f, 2);   
    
                [wc_hg, ~, frq_hg, coi_hg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [61, 300], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);  
                [wc_lg, ~, frq_lg, coi_lg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [1, 4], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);                          
                wccs.pair(y).day(a).run(b).wc_hg(:, :, f) = wc_hg;
                wccs.pair(y).day(a).run(b).wc_lg(:, :, f) = wc_lg;
                wccs.hf_frq = frq_hg;
                wccs.lf_frq = frq_lg;
    
            end   
    
        end

    end
end

% Store averaged TF plot (mean across trial)
wc_hf_tf_all = zeros; wc_lf_tf_all = zeros;
for g = 1:6

    wc_hf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_hg), cat(3, wccs.pair(g).day(2).run.wc_hg), cat(3, wccs.pair(g).day(3).run.wc_hg));
    wc_lf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_lg), cat(3, wccs.pair(g).day(2).run.wc_lg), cat(3, wccs.pair(g).day(3).run.wc_lg));
    wc_hf_tf_all = wc_hf_tf_all + mean(wc_hf_tf_temp, 3);
    wc_lf_tf_all = wc_lf_tf_all + mean(wc_lf_tf_temp, 3);

end
wc_hf_tf_all = wc_hf_tf_all/6; 
wc_lf_tf_all = wc_lf_tf_all/6;



% Store average timecourse (mean across frequencies)
wc_hg_all = zeros; wc_del_all = zeros; wc_hfo_all = zeros; wc_all_hf_all = zeros;
for h = 1:p_no

    wc_hf_temp = cat(3, cat(3, wccs.pair(h).day(1).run.wc_hg), cat(3, wccs.pair(h).day(2).run.wc_hg), cat(3, wccs.pair(h).day(3).run.wc_hg));
    wc_lf_temp = cat(3, cat(3, wccs.pair(h).day(1).run.wc_lg), cat(3, wccs.pair(h).day(2).run.wc_lg), cat(3, wccs.pair(h).day(3).run.wc_lg));
    wc_hg_all = wc_hg_all + mean(wc_hf_temp(28:68, :, :), 1);
    wc_hfo_all = wc_hfo_all + mean(wc_hf_temp(6:27, :, :), 1);
    wc_all_hf_all = wc_all_hf_all + mean(wc_hf_temp(6:68, :, :), 1);
    wc_del_all = wc_del_all + mean(wc_lf_temp, 1);

end


%% S3

load('s3.mat')

% Meta information
n_days = 1;
electrode_list = [1, 2, 3];
electrode_pairs = nchoosek(electrode_list, 2);

for y = 1:size(electrode_pairs, 1)
    e_1 = electrode_pairs(y, 1);
    e_2 = electrode_pairs(y, 2);
% Loop through all days
    for a = 1:n_days
        
        day_struct = c57.day(a);
        n_run = size(day_struct.data_out, 2);
        
        for b = 1:n_run
    
            temp_data = double(day_struct.data_out{1, b});
            temp_data_notch = multi_iirnotch_filtering(temp_data, 1200, [60:60:240]);
            fs = 1200;
    
            % Load indices from specific run
            temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
            temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
            temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
            temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
    
            % Get total number of trials
            temp_n_trials = size(temp_stim_idxs, 2);
    
            % Get trials w/ specific cue/stimulus modalities, 
            tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
            no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
    
            % Loop across all trials and get time segments for PSD
    
            cue_segment = zeros(temp_n_trials, 2);
            bp_segment = zeros(temp_n_trials, 2);
            post_bp_segment = zeros(temp_n_trials, 2);
    
            for c = 1:temp_n_trials
    
                cue_segment(c, 1) = temp_cue_idxs(c);
                cue_segment(c, 2) = temp_cue_idxs(c) + fs;
    
                bp_segment(c, 1) = temp_button_idxs(c) - fs/2;
                bp_segment(c, 2) = temp_button_idxs(c) + fs/2;
    
                post_bp_segment(c, 1) = temp_button_idxs(c) + fs;
                post_bp_segment(c, 2) = temp_button_idxs(c) + 2*fs;  
    
            end
    
            % Loop across all trials to calculate coherence 
    
            for f = 1:temp_n_trials
    
                cue_start_t = cue_segment(f, 1);
                cue_end_t = cue_segment(f, 2);
                bp_start_t = bp_segment(f, 1);
                bp_end_t = bp_segment(f, 2);
                post_bp_start_t = post_bp_segment(f, 1);
                post_bp_end_t = post_bp_segment(f, 2); 
                
         
                % Wavelet coherence Coh(t, frq) (w/ cross-spectrum) 
                [wc_hg, ~, frq_hg, coi_hg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [61, 300], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);  
                [wc_lg, ~, frq_lg, coi_lg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [1, 4], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32); 
                wccs.pair(y).day(a).run(b).wc_hg(:, :, f) = wc_hg;
                wccs.pair(y).day(a).run(b).wc_lg(:, :, f) = wc_lg;
                wccs.hf_frq = frq_hg;
                wccs.lf_frq = frq_lg;
            end   
    
        end
        
    end
end

% Store averaged trial wccs across all HG frequencies

wc_hf_tf_all = zeros; wc_lf_tf_all = zeros;
for g = 1:3

    wc_hf_tf_temp = cat(3, wccs.pair(g).day.run.wc_hg);
    wc_lf_tf_temp = cat(3, wccs.pair(g).day.run.wc_lg);
    wc_hf_tf_all = wc_hf_tf_all + mean(wc_hf_tf_temp, 3);
    wc_lf_tf_all = wc_lf_tf_all + mean(wc_lf_tf_temp, 3);

end
wc_hf_tf_all = wc_hf_tf_all./3; 
wc_lf_tf_all = wc_lf_tf_all./3;

% Store average timecourse (mean across frequencies)
wc_hg_all = zeros; wc_del_all = zeros; wc_hfo_all = zeros; wc_all_hf_all = zeros;
for h = 1:3

    wc_hf_temp = cat(3, wccs.pair(h).day.run.wc_hg);
    wc_lf_temp = cat(3, wccs.pair(h).day.run.wc_lg);
    wc_hg_all = wc_hg_all + mean(wc_hf_temp(28:68, :, :), 1);
    wc_hfo_all = wc_hfo_all + mean(wc_hf_temp(6:27, :, :), 1);
    wc_all_hf_all = wc_all_hf_all + mean(wc_hf_temp(6:68, :, :), 1);
    wc_del_all = wc_del_all + mean(wc_lf_temp, 1);

end

wc_hg_all = squeeze(wc_hg_all)./3; wc_del_all = squeeze(wc_del_all)./3;
wc_hfo_all = squeeze(wc_hfo_all)./3; wc_all_hf_all = squeeze(wc_all_hf_all)./3;

%% S4

load('s4.mat')

% Meta information
n_days = 1;
electrode_list = [68, 69];
e_1 = electrode_list(1); e_2 = electrode_list(2);

% Loop through all days
for a = 1:n_days
    
    day_struct = c71.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        temp_data_notch = multi_iirnotch_filtering(temp_data, 1200, [60:60:240]);

        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate coherence 
        
        for f = 1:temp_n_trials
            
            cue_start_t = cue_segment(f, 1);
            cue_end_t = cue_segment(f, 2);
            bp_start_t = bp_segment(f, 1);
            bp_end_t = bp_segment(f, 2);
            post_bp_start_t = post_bp_segment(f, 1);
            post_bp_end_t = post_bp_segment(f, 2); 
            
            % Wavelet coherence Coh(t, frq) (w/ cross-spectrum) 
                [wc_hg, ~, frq_hg, coi_hg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [61, 300], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32);  
                [wc_lg, ~, frq_lg, coi_lg] = wcoherence(temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_1), temp_data_notch(bp_start_t - 2400:bp_end_t + 2399, e_2), 1200, ...
                    'FrequencyLimits', [1, 4], 'NumScalesToSmooth', 1, 'VoicesPerOctave', 32); 
                wccs.day(a).run(b).wc_hg(:, :, f) = wc_hg;
                wccs.day(a).run(b).wc_lg(:, :, f) = wc_lg;


        end   
        
    end
    
end

% Get T-F spectrogram
wc_hg = cat(3, wccs.day.run.wc_hg);
wc_lg = cat(3, wccs.day.run.wc_lg);

%% Plot: TF coherence matrix
% S1
load('time_frq_coherence\s1_v2.mat')
s1_hf_tf = mean(s1.wc_hf, 3);
s1_lf_tf = mean(s1.wc_lf, 3);

% S2
load('time_frq_coherence\s2.mat')
hf_frq = wccs.hf_frq; lf_frq = wccs.lf_frq;
for g = 1:6
    wc_hf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_hg), cat(3, wccs.pair(g).day(2).run.wc_hg), cat(3, wccs.pair(g).day(3).run.wc_hg));
    wc_lf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_lg), cat(3, wccs.pair(g).day(2).run.wc_lg), cat(3, wccs.pair(g).day(3).run.wc_lg));
    wc_hf_tf_pair(:, :, g) = mean(wc_hf_tf_temp, 3);
    wc_lf_tf_pair(:, :, g) = mean(wc_lf_tf_temp, 3);
end
s2_hf_tf = mean(wc_hf_tf_pair, 3);
s2_lf_tf = mean(wc_lf_tf_pair, 3);

% S3
load('time_frq_coherence\s3.mat')
hf_frq = wccs.hf_frq; lf_frq = wccs.lf_frq;
for g = 1:3

    wc_hf_tf_temp = cat(3, wccs.pair(g).day.run.wc_hg);
    wc_lf_tf_temp = cat(3, wccs.pair(g).day.run.wc_lg);
    wc_hf_tf_pair(:, :, g) = mean(wc_hf_tf_temp, 3);
    wc_lf_tf_pair(:, :, g) = mean(wc_lf_tf_temp, 3);

end
s3_hf_tf = mean(wc_hf_tf_pair, 3);
s3_lf_tf = mean(wc_lf_tf_pair, 3);

% S4
load('time_frq_coherence\s4_v2.mat')
s4_hf_tf = mean(s4.wc_hf, 3);
s4_lf_tf = mean(s4.wc_lf, 3);

% Plot HF and LF wavelet coherence
figure;
h = pcolor(linspace(-2.5, 2.5, 6000), hf_frq, (s1_hf_tf + s2_hf_tf + s3_hf_tf + s4_hf_tf)/4); h.EdgeColor = 'none';
box off; axis square;
xlim([-1, 1]); ylim([70, 170]); 
colorbar; clim([0.4, 0.6]); set(colorbar, 'TickDir', 'out'); 
set(gca,'FontSize', 14); set(gca, 'TickDir', 'out'); set(gca, 'TickLength', [0.02, 0.05]);
xlabel('Time (s)'); ylabel('Frequency (Hz)'); title('Average High Gamma Wavelet Coherence')
xticks(-1:1:1); yticks([70, 120, 170]);

figure;
h2 = pcolor(linspace(-2.5, 2.5, 6000), lf_frq, (s1_lf_tf + s2_lf_tf + s3_lf_tf + s4_lf_tf)/4); h2.EdgeColor = 'none';
box off; axis square;
xlim([-1, 1]); ylim([1, 4]); 
colorbar; clim([0.4, 0.7]); set(colorbar, 'TickDir', 'out');
set(gca,'FontSize', 14); set(gca, 'TickDir', 'out'); set(gca, 'TickLength', [0.02, 0.05]);
xlabel('Time (s)'); ylabel('Frequency (Hz)'); title('Average Delta Wavelet Coherence')
xticks(-1:1:1); yticks(1:1:4);

%% Plot: TF coherence timecourse
% S1
load('time_frq_coherence\s1_v2.mat')
% Frequencies
hf_frq = s1.frq_hg; lf_frq = s1.frq_lf;

% Get timecourse for every trial
s1_hg_trials = squeeze(mean(s1.wc_hf(28:68, :, :), 1)); 
s1_hfo_trials = squeeze(mean(s1.wc_hf(5:27, :, :), 1)); 
s1_all_hf_trials = squeeze(mean(s1.wc_hf(5:68, :, :), 1));

clear s1

% S2
load('time_frq_coherence\s2.mat')
hf_frq = wccs.hf_frq; lf_frq = wccs.lf_frq;

wc_hg_all = zeros; wc_hfo_all = zeros; wc_hf_all = zeros;
for g = 1:6
    wc_hf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_hg), cat(3, wccs.pair(g).day(2).run.wc_hg), cat(3, wccs.pair(g).day(3).run.wc_hg));
    wc_lf_tf_temp = cat(3, cat(3, wccs.pair(g).day(1).run.wc_lg), cat(3, wccs.pair(g).day(2).run.wc_lg), cat(3, wccs.pair(g).day(3).run.wc_lg));
    wc_hg_all = wc_hg_all + squeeze(mean(wc_hf_tf_temp(28:68, :, :), 1));
    wc_hfo_all = wc_hfo_all + squeeze(mean(wc_hf_tf_temp(5:27, :, :), 1));
    wc_hf_all = wc_hf_all + squeeze(mean(wc_hf_tf_temp(5:68, :, :), 1));
end
s2_hg_trials = wc_hg_all./6; s2_hfo_trials = wc_hfo_all./6; s2_all_hf_trials = wc_hf_all./6;

clear wccs wc_hf_tf_temp wc_lf_tf_temp

% S3
load('time_frq_coherence\s3.mat')
hf_frq = wccs.hf_frq; lf_frq = wccs.lf_frq;

wc_hg_all = zeros; wc_hfo_all = zeros; wc_hf_all = zeros;
for g = 1:3

    wc_hf_tf_temp = cat(3, wccs.pair(g).day.run.wc_hg);
    wc_lf_tf_temp = cat(3, wccs.pair(g).day.run.wc_lg);
    wc_hg_all = wc_hg_all + squeeze(mean(wc_hf_tf_temp(28:68, :, :), 1)); % 70.3 - 167.2
    wc_hfo_all = wc_hfo_all + squeeze(mean(wc_hf_tf_temp(5:27, :, :), 1));
    wc_hf_all = wc_hf_all + squeeze(mean(wc_hf_tf_temp(5:68, :, :), 1));

end
s3_hg_trials = wc_hg_all./3; s3_hfo_trials = wc_hfo_all./3; s3_all_hf_trials = wc_hf_all./3;

clear wc_hf_tf_temp wc_lf_tf_temp wccs

% S4
load('time_frq_coherence\s4_v2.mat')
hf_frq = s4.frq_hf; lf_frq = s4.frq_lf;

% Get timecourse for every trial
s4_hg_trials = squeeze(mean(s4.wc_hf(28:68, :, :), 1)); 
s4_hfo_trials = squeeze(mean(s4.wc_hf(5:27, :, :), 1)); 
s4_all_hf_trials = squeeze(mean(s4.wc_hf(5:68, :, :), 1));

% PLOT: Z-score timecourse of HG, HFO, and all HF ACROSS TRIALS
all_hg_trials = [s1_hg_trials, s2_hg_trials, s3_hg_trials, s4_hg_trials];
all_hfo_trials = [s1_hfo_trials, s2_hfo_trials, s3_hfo_trials, s4_hfo_trials];
% all_hf_trials = [s1_all_hf_trials, s2_all_hf_trials, s3_all_hf_trials, s4_all_hf_trials];

% Get baselines

% hg_mean = mean(mean(all_hg_trials(50:5950, :), 2));
% hfo_mean = mean(mean(all_hfo_trials(50:5950, :), 2));
hg_mean = mean(mean(all_hg_trials(1000:2000, :)));
hfo_mean = mean(mean(all_hfo_trials(1000:2000, :)));

hfo_percent = (all_hfo_trials/hfo_mean - 1)*100;
hg_percent = (all_hg_trials/hg_mean - 1)*100;

t_start = 3500; t_end = 4200;
[h, p] = ttest(mean(hfo_percent(t_start:t_end, :), 1), mean(hg_percent(t_start:t_end, :), 1));

% FIGURE: % change
bmap = brewermap(6, 'Blues');
rmap = brewermap(6, 'Reds');
m_hg = mean(hg_percent, 2); s_hg = std(hg_percent, 0, 2)/sqrt(3457);
m_hfo = mean(hfo_percent, 2); s_hfo = std(hfo_percent, 0, 2)/sqrt(3457);
hg_1 = m_hg + s_hg; hg_2 = m_hg - s_hg;
hfo_1 = m_hfo + s_hfo; hfo_2 = m_hfo - s_hfo;
x_plot = linspace(-2.5, 2.5, 6000);
x_all = [x_plot, fliplr(x_plot)];
inBetween1 = [hg_1', fliplr(hg_2')]; inBetween2 = [hfo_1', fliplr(hfo_2')];
figure;
fill(x_all, inBetween1, bmap(2, :), 'FaceAlpha', 0.8, 'EdgeAlpha', 0);
hold on;
fill(x_all, inBetween2, rmap(2, :), 'FaceAlpha', 0.8, 'EdgeAlpha', 0);
plot(linspace(-2.5, 2.5, 6000), mean(hg_percent, 2), 'Color', bmap(5, :));
plot(linspace(-2.5, 2.5, 6000), mean(hfo_percent, 2), 'Color', rmap(5, :));
box off; axis square;
xlim([-0.5, 0.5]); ylim([-1.5, 15]);
set(gca, 'FontSize', 14); set(gca, 'TickDir', 'out'); set(gca, 'TickLength', [0.02, 0.05]);
xlabel('Time (s)'); ylabel('% Change');


% FIGURE: z-score
figure; 
plot(linspace(-2.5, 2.5, 6000), zscore(mean(all_hg_trials, 2)), 'LineWidth', 0.5, 'Color', 'r') 
hold on;
plot(linspace(-2.5, 2.5, 6000), zscore(mean(all_hfo_trials, 2)), 'LineWidth', 0.5, 'Color', 'b')
box off; axis square;
xlim([-0.5, 0.5]); 
set(gca, 'FontSize', 14); set(gca, 'TickDir', 'out');
xlabel('Time (s)'); ylabel('Z-score');

%% Calculate power timecourse 
%% S1

load('s1.mat')

% Meta information
n_days = 3;
electrode_list = [60, 69];
f_1 = 70; f_i = 1; f_2 = 230;
fs = 1200;
[bfilt, afilt] = butter(2, [70, 230]/600, 'bandpass');

% Loop through all days
for a = 1:n_days
    
    day_struct = c41.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        bandpassed_data = filtfilt(bfilt, afilt, temp_data);
        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        stim_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;
            
            stim_segment(c, 1) = temp_stim_idxs(c);
            stim_segment(c, 2) = temp_stim_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate power timecourse per trial 

        for d = 1:numel(electrode_list)
            for e = 1:temp_n_trials
                                            
                cue_start_t = cue_segment(e, 1); cue_end_t = cue_segment(e, 2);
                bp_start_t = bp_segment(e, 1); bp_end_t = bp_segment(e, 2);
                post_bp_start_t = post_bp_segment(e, 1); post_bp_end_t = post_bp_segment(e, 2); 
                
                highgamma_data = bandpassed_data(bp_start_t - 2400:bp_end_t + 2399, electrode_list(d));
                highgamma_envelop = abs(hilbert(highgamma_data));
                highgamma_power = highgamma_envelop.^2;
                                

                hgtc.day(a).run(b).elec(d).hgp(:, :, e) = highgamma_power;
                
            end
        end         
    end    
end

c41_hgp_temp = [[hgtc.day(1).run.elec], [hgtc.day(2).run.elec], [hgtc.day(3).run.elec]];
c41_hgp_all_trials = cat(3, c41_hgp_temp.hgp);
c41_hgp_all_trials = squeeze(c41_hgp_all_trials);

clearvars -except c41_hgp_all_trials 

%% S2

load('s2.mat')

% Meta information
n_days = 3;
electrode_list = [2, 3, 7, 8];
f_1 = 70; f_i = 1; f_2 = 230;
fs = 1200;
[bfilt, afilt] = butter(2, [70, 230]/600, 'bandpass');

% Loop through all days
for a = 1:n_days
    
    day_struct = c45.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        bandpassed_data = filtfilt(bfilt, afilt, temp_data);
        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        stim_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;
            
            stim_segment(c, 1) = temp_stim_idxs(c);
            stim_segment(c, 2) = temp_stim_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate power timecourse per trial 

        for d = 1:numel(electrode_list)
            for e = 1:temp_n_trials
                                            
                cue_start_t = cue_segment(e, 1); cue_end_t = cue_segment(e, 2);
                bp_start_t = bp_segment(e, 1); bp_end_t = bp_segment(e, 2);
                post_bp_start_t = post_bp_segment(e, 1); post_bp_end_t = post_bp_segment(e, 2); 
                
                highgamma_data = bandpassed_data(bp_start_t - 2400:bp_end_t + 2399, electrode_list(d));
                highgamma_envelop = abs(hilbert(highgamma_data));
                highgamma_power = highgamma_envelop.^2;
                                

                hgtc.day(a).run(b).elec(d).hgp(:, :, e) = highgamma_power;
                
            end
        end         
    end    
end

c45_hgp_temp = [[hgtc.day(1).run.elec], [hgtc.day(2).run.elec], [hgtc.day(3).run.elec]];
c45_hgp_all_trials = cat(3, c45_hgp_temp.hgp);
c45_hgp_all_trials = squeeze(c45_hgp_all_trials);

clearvars -except c41_hgp_all_trials c45_hgp_all_trials 

%% S3

load('s3.mat')

% Meta information
n_days = 1;
electrode_list = [1, 2, 3];
f_1 = 70; f_i = 1; f_2 = 230;
fs = 1200;
[bfilt, afilt] = butter(2, [70, 230]/600, 'bandpass');

% Loop through all days
for a = 1:n_days
    
    day_struct = c57.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        bandpassed_data = filtfilt(bfilt, afilt, temp_data);
        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        stim_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;
            
            stim_segment(c, 1) = temp_stim_idxs(c);
            stim_segment(c, 2) = temp_stim_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate power timecourse per trial 

        for d = 1:numel(electrode_list)
            for e = 1:temp_n_trials
                                            
                cue_start_t = cue_segment(e, 1); cue_end_t = cue_segment(e, 2);
                bp_start_t = bp_segment(e, 1); bp_end_t = bp_segment(e, 2);
                post_bp_start_t = post_bp_segment(e, 1); post_bp_end_t = post_bp_segment(e, 2); 
                
                highgamma_data = bandpassed_data(bp_start_t - 2400:bp_end_t + 2399, electrode_list(d));
                highgamma_envelop = abs(hilbert(highgamma_data));
                highgamma_power = highgamma_envelop.^2;
                                

                hgtc.day(a).run(b).elec(d).hgp(:, :, e) = highgamma_power;
                
            end
        end         
    end    
end

c57_hgp_temp = [[hgtc.day.run.elec]];
c57_hgp_all_trials = cat(3, c57_hgp_temp.hgp);
c57_hgp_all_trials = squeeze(c57_hgp_all_trials);

clearvars -except c41_hgp_all_trials c45_hgp_all_trials c57_hgp_all_trials

%% S4

load('s4.mat')

% Meta information
n_days = 1;
electrode_list = [68, 69];
f_1 = 70; f_i = 1; f_2 = 230;
fs = 1200;
[bfilt, afilt] = butter(2, [70, 230]/600, 'bandpass');

% Loop through all days
for a = 1:n_days
    
    day_struct = c71.day(a);
    n_run = size(day_struct.data_out, 2);    
    
    for b = 1:n_run
        
        temp_data = double(day_struct.data_out{1, b});
        bandpassed_data = filtfilt(bfilt, afilt, temp_data);
        
        % Load indices from specific run
        temp_stim_idxs = day_struct.ECoG.trigger_index{1, b}.stim_idxs;
        temp_cue_idxs = day_struct.ECoG.trigger_index{1, b}.cue_idxs;   
        temp_button_idxs = day_struct.ECoG.trigger_index{1, b}.button_idxs;
        temp_stim_schedule = day_struct.ECoG.stimulus_schedule{1, b};
        
        % Get total number of trials
        temp_n_trials = size(temp_stim_idxs, 2);
        
        % Get trials w/ specific cue/stimulus modalities, 
        tac_cue_trials = find(temp_stim_schedule(:, 3) == 1); % Tactile Cue
        no_tac_cue_trials = find(temp_stim_schedule(:, 3) == 0);
        
        % Loop across all trials and get time segments for PSD
        
        cue_segment = zeros(temp_n_trials, 2);
        stim_segment = zeros(temp_n_trials, 2);
        bp_segment = zeros(temp_n_trials, 2);
        post_bp_segment = zeros(temp_n_trials, 2);
   
        for c = 1:temp_n_trials
            
            cue_segment(c, 1) = temp_cue_idxs(c);
            cue_segment(c, 2) = temp_cue_idxs(c) + 1200;
            
            stim_segment(c, 1) = temp_stim_idxs(c);
            stim_segment(c, 2) = temp_stim_idxs(c) + 1200;

            bp_segment(c, 1) = temp_button_idxs(c) - 600;
            bp_segment(c, 2) = temp_button_idxs(c) + 600;
            
            post_bp_segment(c, 1) = temp_button_idxs(c) + 1200;
            post_bp_segment(c, 2) = temp_button_idxs(c) + 2400;  
            
        end
        
        % Loop across all trials to calculate power timecourse per trial 

        for d = 1:numel(electrode_list)
            for e = 1:temp_n_trials
                                            
                cue_start_t = cue_segment(e, 1); cue_end_t = cue_segment(e, 2);
                bp_start_t = bp_segment(e, 1); bp_end_t = bp_segment(e, 2);
                post_bp_start_t = post_bp_segment(e, 1); post_bp_end_t = post_bp_segment(e, 2); 
                
                highgamma_data = bandpassed_data(bp_start_t - 2400:bp_end_t + 2399, electrode_list(d));
                highgamma_envelop = abs(hilbert(highgamma_data));
                highgamma_power = highgamma_envelop.^2;
                                

                hgtc.day(a).run(b).elec(d).hgp(:, :, e) = highgamma_power;
                
            end
        end  
        
    end    
end

c71_hgp_temp = [[hgtc.day.run.elec]];
c71_hgp_all_trials = cat(3, c71_hgp_temp.hgp);
c71_hgp_all_trials = squeeze(c71_hgp_all_trials);

clearvars -except c41_hgp_all_trials c45_hgp_all_trials c57_hgp_all_trials c71_hgp_all_trials

% Combine power timecourses across all cases and get mean plot
hgp_all_trials_all_cases = [c41_hgp_all_trials, c45_hgp_all_trials, c57_hgp_all_trials, c71_hgp_all_trials];

%% PLOT: Power and coherence timecourse

% Manual Z-Score

hgp_baseline = mean(mean(hgp_all_trials_all_cases(1000:2000, :)));
hgp_mean = mean(hgp_all_trials_all_cases, 2);
sigma = std(hgp_mean);
hgp_zscore = (hgp_mean - hgp_baseline)/sigma; 

hg_mean_tc = mean(all_hg_trials, 2);
sigma_hg = std(hg_mean_tc);
hg_zscore = (hg_mean_tc - hg_mean)/sigma_hg;

% To get normalized z-score
hgp_z_norm = hgp_zscore/trapz(hgp_zscore); 
hg_z_norm = hg_zscore/trapz(hg_zscore);

% FWHM
m_hg = max(hg_zscore(1000:5000));
hM = m_hg/2;
index_hgc = find(hg_zscore(1000:5000) >= hM, 1, 'last');

m_hgp = max(hgp_zscore(1000:5000));
hMp = m_hgp/2;
index_hgp = find(hgp_zscore(1000:5000) >= hMp, 1, 'last');

% Z-score from biased baseline

figure;
yyaxis left
plot(linspace(-2.5, 2.5, 6000), hgp_z_norm, 'Color', [0 0.4470 0.7410], 'LineWidth', 3)
yticks([-1:1:8]); ylim([-1e-3, 4e-3])
ylabel('Z-Scored High Gamma Power')
yline(0, '--', 'Color', [0 0.4470 0.7410], 'LineWidth', 2);
hold on;
yyaxis right
plot(linspace(-2.5, 2.5, 6000), hg_z_norm, 'Color', [0.8500 0.3250 0.0980], 'LineWidth', 2);
yticks([-1:1:4]); ylim([-1e-3, 4e-3])
yline(0, '--', 'Color', [0.8500 0.3250 0.0980], 'LineWidth', 2);
ylabel('Z-Scored High Gamma Coherence');
xlim([-1, 1]); xticks([-2, -1, -0.5, 0, 0.5, 1, 2]);
xline(0, '--', 'Color', 'k', 'LineWidth', 2);
xlabel('Time relative to Button Press (s)')
box off;

figure;
yyaxis left
plot(linspace(-2.5, 2.5, 6000), zscore(mean(hgp_all_trials_all_cases, 2)), 'Color', [0 0.4470 0.7410], 'LineWidth', 3)
yticks([0:1:5:7]); ylim([-1, 8])
ylabel('Z-Scored High Gamma Power')
yline(0, '--', 'Color', [0 0.4470 0.7410], 'LineWidth', 2);
hold on;
yyaxis right
plot(linspace(-2.5, 2.5, 6000), zscore(mean(all_hfo_trials, 2)), 'Color', [0.8500 0.3250 0.0980], 'LineWidth', 3);
yticks([-1:1:4]);
yline(0, '--', 'Color', [0.8500 0.3250 0.0980], 'LineWidth', 2);
ylabel('Z-Scored HFO Coherence');
xlim([-1, 1]); xticks([-2, -1, -0.5, 0, 0.5, 1, 2]);
xline(0, '--', 'Color', 'k', 'LineWidth', 2);
xlabel('Time relative to Button Press (s)')
box off; axis square;


% PLOT: raw wc timecourse w/ confidence interval
mean_hg = mean(all_hg_trials, 2); mean_hfo = mean(all_hfo_trials);
stdev_hg = std(all_hg_trials, 0, 2); stdev_hfo = std(all_hfo_trials, 0, 2);
ci_hg = 1.96*stdev_hg/sqrt(3457); se_hg = stdev_hg/sqrt(3457);
ci_hfo = 1.96*stdev_hfo/sqrt(3457); se_hfo = stdev_hfo/sqrt(3457);

bmap = brewermap(6, 'Blues');
x_plot = linspace(-2.5, 2.5, 6000); 
x1 = [x_plot, fliplr(x_plot)];
curve_hg1 = (mean_hg./mean(mean_hg) - 1)*100 + ci_hg.*100; curve_hg2 = (mean_hg./mean(mean_hg) - 1)*100 - ci_hg.*100;
figure;
inBetween1 = [curve_hg1', fliplr(curve_hg2')];
fill(x1, inBetween1, bmap(2, :), 'FaceAlpha', 0.8, 'EdgeAlpha', 0);
hold on;
plot(x_plot, (mean_hg./mean(mean_hg) - 1)*100, 'Color', bmap(4, :), 'LineWidth', 0.5);

