% Script 4
% goals of this script:
% - PCA of optical flow magnitude
% if you have questions, feel free to contact Briana R. McRae: brm2144@columbia.edu

% NOTE: this script also requires heavy 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: conduct PCA
% NOTE: this section only needs to be run once per combination of stimnum, dsfactor, and fnfactor variables

clear all; clc; close all;

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

% set up folders
optflowOutputsPath = [outputFolder, 'optflowOutputs/'];
cd(outputFolder)
optflowVids = dir([optflowOutputsPath '*.avi']);

% ROI and timing inputs: edit!
inputCoords = load('ROIcoords.mat');
vidFrameRate = 29.97; % frame rate (Hz)
eventDurSec = 5; % duration of stimulus events, in seconds
prebuffer = 0; % how many frames before onset to look at

% analysis options: edit!
% NOTE: default example is PCA for pup call responses
stimnum = 1; % which stimulus event type to analyze
dsfactor = 1; % spatial downsample factor (default = 1, no downsampling)
fnfactor = 1; % frame number downsample factor (default = 1, no downsampling)

% set up folders
keyInfo = ['stimnum', num2str(stimnum), 'dsfactor', num2str(dsfactor), 'fnfactor', num2str(fnfactor)];
pcaOutputPath = [outputFolder, 'pcaOutputs_', keyInfo, '/']; % makes a new folder for each combination of stimnum, dsfactor, and fnfactor
mkdir(pcaOutputPath)

% calculate standard timing info
eventDur = eventDurSec*vidFrameRate; % seconds of event duration, turned into frames
postbuffer = eventDur*2; % edit! how many frames after onset to look at (i.e. length of event duration x 2)

% edit the code block below to find/parse videos and event timing files!
for vv=1:length(optflowVids)
    vidName = optflowVids(vv).name;
    fprintf(['\n analyzing video :: ' vidName])
    vidID = erase(vidName, '_flow.avi');
    vidOnsets = load([vidOnsetsDir, vidID, '_videoTimestamps.mat']); % in seconds
    vid = VideoReader([optflowOutputsPath, vidName]);

    % reading predefined ROI
    coords = inputCoords.coords;
    ROIwidth = round(coords(1)):round(coords(1)+coords(3));
    xd = ROIwidth(end)-ROIwidth(1)+1;
    ROIheight = round(coords(2)):round(coords(2)+coords(4));
    yd = ROIheight(end)-ROIheight(1)+1;

    % gathering image dimensions
    matlength = read(vid,1);
    matlength = squeeze(matlength(:,:,1));
    matlength = matlength(ROIheight,ROIwidth);
    matlength = imresize(matlength, 1/dsfactor);
    matlength = length(matlength(:));

    % vectorize each frame (1 row = 1 frame)
    framestack = zeros(floor(vid.NumFrames/fnfactor), matlength);
    for ff=1:floor(vid.NumFrames)
        frame = read(vid,ff);
        frame = reshape(frame(:,:,1),[size(frame,1) size(frame,2)]);
        tmp = frame(ROIheight,ROIwidth);
        tmp = imresize(tmp, 1/dsfactor);
        tmp = tmp(:)';
        framestack(ff,:) = double(tmp);
    end

    % get timing info & select epochs of interest
    eventtypes = fieldnames(vidOnsets.outputStruct);
    onsets = vidOnsets.outputStruct.(eventtypes{stimnum});
    onsets_frame = round(onsets*vidFrameRate);
    try
        framestack_epochs = [];
        for ii=1:length(onsets_frame)
            startframe = onsets_frame(ii)- prebuffer;
            endframe = onsets_frame(ii) + postbuffer;
            framestack_epochs = [framestack_epochs; framestack(startframe:endframe,:)];
        end
    catch % if first stimulus period started too early to capture pre-stimulus period
        framestack_epochs = [];
        for ii=2:length(onsets_frame)
            startframe = onsets_frame(ii)- prebuffer;
            endframe = onsets_frame(ii) + postbuffer;
            framestack_epochs = [framestack_epochs; framestack(startframe:endframe,:)];
        end
    end

    % run PCA on z-scored info
    [pc,~,~,~,explained] = pca(zscore(framestack_epochs));

    cmsm = cumsum(explained);
    targpc = find(cmsm>50); % save PCs until they explain 50% of variance
    targpc = targpc(1);

    pcs = pc(:,1:targpc);
    explainedvars = explained(1:targpc);
    
    % save output
    save([pcaOutputPath, erase(vidName, '.avi'), '_pcaoutputs.mat'],'pcs','explainedvars');    
end

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