% Script 3
% goals of this script:
% - comparison of optical flow velocity magnitude for chosen ROI (across groups, across conditions, etc.)
% if you have questions, feel free to contact Briana R. McRae: brm2144@columbia.edu

% NOTE: this script requires the most specialization to your experiment,
% including accessing and reading event onsets, as well as carrying out the
% group/condition comparisons of interest

% edit all sections!

%% Section 1: process raw optical flow

clc; clear all; close all;

% inputs: edit!
outputFolder = '/path/to/supplementalMaterials/exampleData/outputs/';
vidPath = '/path/to/supplementalMaterials/exampleData/';
vidOnsetsDir = '/path/to/supplementalMaterials/exampleData/';

% set up folders
cd(outputFolder)
optflowOutputPath = [outputFolder, 'optflowOutputs/'];

% inputs: edit!
load('allRawVelMagData.mat') % access the output from script 2, section 3
eventDurSec = 5; % duration of stimulus events, in seconds
vidFrameRate = 29.97; % in frames per second

% analysis options: edit!
zsc = 1; % flag where 1 = z-score data
blrm = 1; % flag where 1 = mean baseline removal
blzsc = 0; % flag where 1 = baseline zscore per trial
demeanPerTrial = 0; % flag where 1 = demean per trial
applyLPfilter = 1; % flag where 1 = apply lowpass filter
LPfilter = 5; % in Hz, up to experimenter's discretion

% calculate standard timing info
eventDur = round(eventDurSec * vidFrameRate); % in frame idxs, not seconds
bufferPre = round(eventDurSec * vidFrameRate); % edit! frames before event onset you want to collect
bufferPost = round(eventDurSec * vidFrameRate); % edit! frames after event offset you want to collect
epochSize = eventDur + bufferPre + bufferPost;
timeBounds = [1:epochSize]; % 1:epochSize = use full time bounds of pre during and post stim to plot optflow over time
bounds_pre = 1:bufferPre; % take the whole pre-stim period
bounds_stim_full = bufferPre:(bufferPre + eventDur); % full stimulus period
stim_analysis = 2.5; % edit! can analyze subset of stimulus period, such as the first 2.5s
bounds_stim_analysis = bufferPre:(bufferPre + round(stim_analysis*vidFrameRate));

% process raw data
NmaskedFrames = 600; % edit! for example data, we masked the first 600 frames (20s) to allow for brightness to stabilize
for ii=1:length(allRawVelMagData)
    optFlow = allRawVelMagData(ii).optFlow;
    optFlow(1:NmaskedFrames) = [];
    if zsc
        optFlow = zscore(optFlow); % z-score for whole video
    end
    if applyLPfilter
        optFlow = lowpass(optFlow, LPfilter, vidFrameRate); % apply filter for whole video optflow
    end
    allRawVelMagData(ii).procOptFlow = [NaN(NmaskedFrames,1); optFlow; repmat(optFlow(end),150,1)]; % add back masked frames at beginning as NaNs
end

% chunk into epochs: edit the code block below to find/parse videos and event timing files!
allVelMagVids = dir([optflowOutputPath, '*_flow.avi']);
for ii=1:length(allVelMagVids)
    % identify event onsets for this video
    thisVelMagID = allVelMagVids(ii).name;
    thisSubj_split = strsplit(thisVelMagID,'_');
    thisSubj = thisSubj_split{2};
    thisVidID = erase(thisVelMagID, '_flow.avi');
    thisVidOnsets = [thisVidID, '_videoTimestamps.mat'];
    onsets = load([vidOnsetsDir, thisVidOnsets]).outputStruct; % assumes that your onsets are provided in seconds
    stim1onsets = round((onsets.pupCalls_onsets) * vidFrameRate); % convert to frame idxs, not seconds
    stim2onsets = round((onsets.pureTonesIn_onsets) * vidFrameRate);

    % identify frames of interest for stimulus 1
    stim1frames = NaN(epochSize, length(stim1onsets));
    for aa=1:size(stim1frames,2)
        stim1frames(:,aa) = [(stim1onsets(aa)-bufferPre):(stim1onsets(aa)+eventDur+bufferPost-1)];
    end
    
    % identify frames of interest for stimulus 2
    stim2frames = NaN(epochSize, length(stim2onsets));
    for bb=1:size(stim2frames,2)
        stim2frames(:,bb) = [(stim2onsets(bb)-bufferPre):(stim2onsets(bb)+eventDur+bufferPost-1)];
    end

    % pull out stimulus 1 epochs
    try
        allRawVelMagData(ii).stim1optflow = allRawVelMagData(ii).procOptFlow(stim1frames);
    catch % in case first event was too close to start of video, this should catch negative frame indices
        allRawVelMagData(ii).stim1optflow = allRawVelMagData(ii).procOptFlow(stim1frames(:,2:end));
    end

    % pull out stimulus 2 epochs
    try
        allRawVelMagData(ii).stim2optflow = allRawVelMagData(ii).procOptFlow(stim2frames);
    catch % in case first event was too close to start of video, this should catch negative frame indices
        allRawVelMagData(ii).stim2optflow = allRawVelMagData(ii).procOptFlow(stim2frames(:,2:end));
    end
end

if blrm 
    for ii=1:length(allRawVelMagData)
        a = allRawVelMagData(ii).stim1optflow;
        allRawVelMagData(ii).stim1optflow = (a-mean( a(bounds_pre,:), 1));

        b = allRawVelMagData(ii).stim2optflow;
        allRawVelMagData(ii).stim2optflow = (b-mean( b(bounds_pre,:), 1));
    end
end

if blzsc 
    for ii=1:length(allRawVelMagData)
        a = allRawVelMagData(ii).stim1optflow;
        allRawVelMagData(ii).stim1optflow = (a-mean( a(bounds_pre,:), 1))./std( a(bounds_pre,:));

        b = allRawVelMagData(ii).stim2optflow;
        allRawVelMagData(ii).stim2optflow = (b-mean( b(bounds_pre,:), 1))./std( b(bounds_pre,:));
    end
end

disp('done with script 3, section 1!')

%% Section 2: get average framewise optical flow velocity magnitude vectors per subject, broken down by event type

allSubjs = unique({allRawVelMagData.subj});
disp('order of subjects: ')
disp(allSubjs) % use this to determine group indices in next section
stim1optflow_subjAvg = zeros(epochSize, length(allSubjs));
stim2optflow_subjAvg = zeros(epochSize, length(allSubjs));
counts = zeros(1, length(allSubjs));

for ii=1:length(allRawVelMagData)
    thisSubj = allRawVelMagData(ii).subj;
    subjN = find(strcmp(allSubjs,thisSubj));

    if demeanPerTrial ~= 1
        stim1mean = mean(allRawVelMagData(ii).stim1optflow,2, 'omitnan'); % getting 1 avg optflow vector per video
        stim2mean = mean(allRawVelMagData(ii).stim2optflow,2, 'omitnan');
    else
        stim1mean = mean(allRawVelMagData(ii).stim1optflow - mean(allRawVelMagData(ii).stim1optflow(bounds_pre,:), 1) ,2, 'omitnan'); % getting 1 avg optflow vector per video
        stim2mean = mean(allRawVelMagData(ii).stim2optflow  - mean(allRawVelMagData(ii).stim2optflow(bounds_pre,:), 1) ,2, 'omitnan');
    end

    stim1optflow_subjAvg(:,subjN) = stim1optflow_subjAvg(:,subjN) + stim1mean;
    stim2optflow_subjAvg(:,subjN) = stim2optflow_subjAvg(:,subjN) + stim2mean;
    counts(subjN) = counts(subjN) + 1; % keep track of how many videos per subject
end

% get 1 average optflow vector per event type for each subj
stim1optflow_subjAvg = stim1optflow_subjAvg./counts;
stim2optflow_subjAvg = stim2optflow_subjAvg./counts;

disp('done with script 3, section 2!')

%% Section 3: compare optical flow during stimulus presentation between groups/conditions

% inputs: edit!
% NOTE: for example data, group 1 is mothers, group 2 is virgins, 
% stimulus 1 is pup calls, and stimulus 2 is pure tones
group1idx = [2];
group2idx = [1];

% determine comparisons (either group or condition comparison): edit!
% NOTE: default example is comparing mothers' vs. virgins' pup call responses
comp1_stim = stim1optflow_subjAvg;
comp1_group = group1idx;
comp2_stim = stim1optflow_subjAvg;
comp2_group = group2idx;
myTitle = "Pup calls";
group1label = "Mothers";
group2label = "Virgins";

% plot group average optflow over time
y_group1 = mean(comp1_stim(timeBounds,comp1_group),2);
sem_group1 = std(comp1_stim(timeBounds,comp1_group),0,2)./sqrt(length(comp1_group));
y_group2 = mean(comp2_stim(timeBounds,comp2_group),2);
sem_group2 = std(comp2_stim(timeBounds,comp2_group),0,2)./sqrt(length(comp2_group));

figure()
hold on
fill([bufferPre bufferPre+eventDur bufferPre+eventDur bufferPre],[-2 -2 2 2], [0 0 0],'FaceAlpha',.15,'edgealpha',0)
plot(y_group1,'Color',[200, 0, 161]/255,'LineWidth',2)
plot(y_group2,'Color',[255, 131, 230]/255,'LineWidth',2)
fill([timeBounds, flip(timeBounds)], [y_group1-sem_group1; flip(y_group1+sem_group1)],[200, 0, 161]/255,'FaceAlpha',.3,'edgealpha',0)
fill([timeBounds, flip(timeBounds)], [y_group2-sem_group2; flip(y_group2+sem_group2)],[255, 131, 230]/255,'FaceAlpha',.3,'edgealpha',0)
line(xlim, [0 0], 'LineStyle', '--', 'Color', 'k')
%plot(bounds_stim_analysis,ones(size(bounds_stim_analysis)),'Color','r') % edit! uncomment to visualize which segment is being analyzed for summary statistics
hold off
title(myTitle);
legend({'Stimulus presentation', group1label, group2label})
legend boxoff
xlim([75 375]) % edit! how much pre and post stimulus time to visualize
xticks([150-75 150 150+75 150+75*2 150+75*3]) % edit!
xticklabels({'-2.5','0','2.5','5','7.5'}) % edit!
xlabel('Time since stimulus onset (s)')
ylabel('Normalized optical flow')
set(gcf, 'Position',[1056 349 602 340])


% plot cumulative optflow over time
y_group1a = cumsum(comp1_stim(bounds_stim_full,comp1_group),1);
y_group1 = mean(y_group1a,2);
sem_group1 = std(y_group1a,0,2)./sqrt(length(comp1_group));

y_group2a = cumsum(comp2_stim(bounds_stim_full,comp2_group),1);
y_group2 = mean(y_group2a,2);
sem_group2 = std(y_group2a,0,2)./sqrt(length(comp2_group));

figure()
hold on
plot(bounds_stim_full,y_group1,'Color',[200, 0, 161]/255,'LineWidth',2)
plot(bounds_stim_full,y_group2,'Color',[255, 131, 230]/255,'LineWidth',2)
fill([bounds_stim_full, flip(bounds_stim_full)], [y_group1-sem_group1; flip(y_group1+sem_group1)],[200, 0, 161]/255,'FaceAlpha',.3,'edgealpha',0)
fill([bounds_stim_full, flip(bounds_stim_full)], [y_group2-sem_group2; flip(y_group2+sem_group2)],[255, 131, 230]/255,'FaceAlpha',.3,'edgealpha',0)
line(xlim, [0 0], 'LineStyle', '--', 'Color', 'k')
hold off                     
title(myTitle);
legend({group1label, group2label})
legend boxoff
xticks([150-75 150 150+75 150+75*2 150+75*3]) % edit!
xticklabels({'-2.5','0','2.5','5','7.5'}) % edit!
xlabel('Time since stimulus onset (s)')
ylabel({'Cumulative optical flow'; 'relative to pre-stimulus baseline'})
set(gcf, 'Position',[1440 1031 303 287])


% calculate summary statistics for cumulative optflow over time for stimulus analysis period (aka area under the curve, or AUC)
disp(['statistical analysis: ', group1label, ' ', group2label, ' ', myTitle])
auc_group1 = NaN(1, length(comp1_group));
auc_group2 = NaN(1, length(comp2_group));
for ii = 1:length(comp1_group) % calculate AUC, or cumulative optflow
    auc_group1(ii) = trapz(bounds_stim_analysis, comp1_stim(bounds_stim_analysis,comp1_group(ii)));
    auc_group2(ii) = trapz(bounds_stim_analysis, comp2_stim(bounds_stim_analysis,comp2_group(ii)));
end
[p,H] = ranksum(auc_group1, auc_group2); % edit! choose test as desired
disp(['p = ', num2str(p)])


% plot AUC during stimulus analysis period
figure()
hold on;
boxchart(ones(size(auc_group1)),(auc_group1),'BoxFaceColor',[200, 0, 161]/255)
boxchart(ones(size(auc_group2))*2,(auc_group2),'BoxFaceColor',[255, 131, 230]/255)
scatter(ones(1,6), auc_group1, 'filled', 'MarkerFaceColor', [240 240 240]/255,'MarkerEdgeColor','k');
scatter(2*ones(1,6), auc_group2, 'filled', 'MarkerFaceColor', [240 240 240]/255,'MarkerEdgeColor','k');
hold off;
xlim([0,3]);
ylim([-17 17])
xticks([1, 2]);
xticklabels({group1label,group2label});
title(myTitle);
ylabel({'Cumulative optical flow'; [num2str(stim_analysis), ' s from stimulus onset']})
set(gcf, 'Position',[1440 1031 303 287])


disp('done with script 3, section 3!')