% Script 2
% goals of this script:
% - optical flow estimation per video (retain spatial and temporal information)
% - optical flow velocity magnitude extraction per video (averaged across chosen ROI)
% if you have questions, feel free to contact Briana R. McRae: brm2144@columbia.edu

%% Section 1: conduct optical flow analysis on all transformed videos
% NOTE: this section only needs to be run once on entire video dataset

clc; clear all; close all;

% inputs: edit!
outputFolder = '/path/to/supplementalMaterials/exampleData/outputs/';
vidPath = '/path/to/supplementalMaterials/exampleData/';
subjs = {'Br011', 'Br014'};

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

% edit the code block below!
for ss=1:length(subjs)
    subj = subjs{ss};
    allVidsForSubj = dir([vidPath,'/*', subj, '*.avi']);
    for ii=1:length(allVidsForSubj)
        opticFlow = opticalFlowHS;
        thisVidFile = [vidPath, allVidsForSubj(ii).name];
        thisVid = VideoReader(thisVidFile);
        numFrames = thisVid.NumFrames;
        jj=1; % frame idx
        baseFileName = erase(allVidsForSubj(ii).name,'.avi');
        
        % set up videowriter to save optical flow magnitudes as a video (more efficient than saving as matrix)
        newvidname = [optflowOutputPath, erase(allVidsForSubj(ii).name,'.avi'), '_flow.avi'];
        vidwriter = VideoWriter(newvidname);
        vidwriter.FrameRate = 29.97;
        open(vidwriter);

        while hasFrame(thisVid)
            if mod(jj,50)==0
            clc;
            disp(['video :: ' baseFileName ' progress :: ' num2str(100*jj/numFrames) ' %'])
            end
            frameRGB = readFrame(thisVid);
            frameGray = im2gray(frameRGB);
            flow = estimateFlow(opticFlow, frameGray);
            img = flow.Magnitude;
            writeVideo(vidwriter, img(:,:,[1 1 1])) % save as grayscale
            jj = jj + 1; % move onto next frame idx
        end

        close(vidwriter)
    end
end

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

%% Section 2: visualize variability of optical flow of each pixel in full FOV

% loop through optflow outputs, which were saved as videos
allVelMagVids = dir([optflowOutputPath, '*_flow.avi']);
optflowSTDsum = zeros(480,720);
for ii=1:length(allVelMagVids)
    % set up video reader
    thisVelMagID = allVelMagVids(ii).name;
    thisSubj_split = strsplit(thisVelMagID,'_');
    thisSubj = thisSubj_split{2};
    flowmagvid = VideoReader([optflowOutputPath, thisVelMagID]);
    ROIwidth = 1:flowmagvid.Width;
    ROIheight = 1:flowmagvid.Height;
    % read in all frames
    allFramesRGB = read(flowmagvid,[1,Inf]);
    allFramesGray = squeeze(allFramesRGB(ROIheight,ROIwidth,1,:)); % only take 1 dimension of RGB values, to get grayscale
    optflowSTDsum = optflowSTDsum + std(double(allFramesGray),0,3);
    disp(['done with video #', num2str(ii)])
end

optflowSTDavg = optflowSTDsum/ii;
figure()
imagesc(imrotate(optflowSTDavg, -90)); % edit! rotate according to orientation of videos
ax = gca;
ax.CLim = [0,2.5];
colorbar
set(gcf, 'position',[558 254 486 532])
title('stdev of optical flow, averaged across all videos')

% save output
save('avgOptflowSTD.mat','optflowSTDavg');

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

%% Section 3: get framewise optical flow velocity magnitude vectors per ROI
% NOTE: this section only needs to be run once per ROI

% ROI inputs: edit!
selectROI = 1; % flag where 1 = you have to draw an ROI
provideROI = 0; % flag where 1 = you have to provide predetermined coordinates of an ROI
if selectROI
    newCoordsName = 'ROIcoords'; % edit! enter name of ROI such that it's coordinates can be saved later
end
if provideROI
    inputCoords = load('ROIcoords.mat'); % edit! where to find existing coordinates
end

% other inputs: edit!
vidFrameRate = 29.97; % in frames per second (Hz)

% calculate standard timing info: edit!
bufferPre = round(5 * vidFrameRate); % frames before event onset to collect (5sec)
bufferPost = round(5 * vidFrameRate); % frames after event offset to collect (5sec)

% identify ROI (based on average aligned image)
if selectROI
    avgimg = load('avgimg_aligned.mat');
    imshow(uint8(avgimg.avgimg))
    h = drawrectangle;
    coords = get(h,'position');
    ROIwidth = [round(coords(1)):round(coords(1)+coords(3))];
    ROIheight = [round(coords(2)):round(coords(2)+coords(4))];
    save([newCoordsName, '.mat'],'coords'); % save coordinates in case it's needed later
elseif provideROI
    coords = inputCoords.coords;
    ROIwidth = [round(coords(1)):round(coords(1)+coords(3))];
    ROIheight = [round(coords(2)):round(coords(2)+coords(4))];
end

% loop through optflow outputs, which were saved as videos
allVelMagVids = dir([optflowOutputPath, '*_flow.avi']);
allRawVelMagData = struct();
for ii=1:length(allVelMagVids) % loop through each video
    % get optflow video info
    thisVelMagID = allVelMagVids(ii).name;
    thisSubj_split = strsplit(thisVelMagID,'_');
    thisSubj = thisSubj_split{2};
    thisVidID = erase(thisVelMagID, '_optFlowVelMag.mat');

    % set up video reader
    flowmagvid = VideoReader([optflowOutputPath, thisVelMagID]);
    
    % get optflow velocity magnitude vector
    allFramesRGB = read(flowmagvid,[1,Inf]);
    allFramesGray = squeeze(allFramesRGB(ROIheight,ROIwidth,1,:)); % only take 1 dimension of RGB values, to get grayscale
    optFlow = squeeze(mean(allFramesGray,[1,2])); % get 1 value per frame, representing optflow velocity magnitude averaged over all pixels

    clear flowmagvid; % close video reader

    allRawVelMagData(ii).subj = thisSubj;
    allRawVelMagData(ii).optFlow = optFlow;
    
    disp(['done with video #', num2str(ii)])
end

% save output
save('allRawVelMagData.mat','allRawVelMagData');
if selectROI
    save([newCoordsName, '.mat'],'coords');
end

close all;

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