function [legEndpointsStore, frameFlags] = ErrorCorrectionFunction(movName,overWriteFlag,targetNumLegs,finalMoviePauseTime)
% ErrorCorrectionFunction loads the movie, allows the user to verify "correct" frames and fix
% "incorrect" frames, and saves the final movie.
%
%%%%%%% Arguments:
% movName [string]: name of movie to consider
% overWriteFlag ['y' OR 'n' OR 'default']: whether quitting the movie before completing review of all
    % frames or skipping some frames temporarily overwrites the frameFlags in the Automatically Tracking 
    % Movies folder. Selecting 'n' means you lose all of your tagging work if you do not complete
    % the movie. Note that movies are NOT saved in the final folder unless all frameFlags are 10
    % (frames checked and accepted by user).
        % NOTE: final movie with tracked legs only gets saved here if all frame flags are 10 (completely checked and accepted)
        % default = 'y'
% targetNumLegs [integer between 1 and 6 OR 'default']: target number of legs to tag.
    % If default is set, "preamputation" in move title sets to 6 legs, "postamputation" in movie
    % title sets to 5 legs
% finalMoviePauseTime [double OR 'default']: number of second pause between frame for FINAL movie
    % (sorted legs in different colors)
        % default = .03
    
    
if isequal(overWriteFlag,'default')
    overWriteFlag = 'y';
end

if isequal(targetNumLegs,'default')
    if not(isempty(strfind(movName,'preamputation'))) % if file name includes "preamputation"
        targetNumLegs = 6; % find 6 legs
    elseif not(isempty(strfind(movName,'postamputation'))) % if file name includes "postamputation"
        targetNumLegs = 5; % find 5 legs
    else
        display('Warning - neither "preamputation" nor "postamputation" appears in video. Please manually specify number of legs. Assuming 5 legs!')
        targetNumLegs = 5;
    end
end

if isequal(finalMoviePauseTime,'default')
    finalMoviePauseTime = .03;
end

close

% Create Temporary Working Directory
workingDir = tempname;
mkdir(workingDir);
mkdir(workingDir,'images');

% Load AutoTracked Movie And Create Sorted Images From Frames, load frameFlags,
% legEndpointsStore, load centroidCoordinates, load headCoordinates
frameFlagsLoadString = strrep(movName, '.avi','_frameFlags.mat');
legEndpointsStoreLoadString = strrep(movName, '.avi', '_legEndpointsStore.mat');
centroidCoordinatesLoadString = strrep(movName, '.avi', '_centroidCoordinates.mat');
headDirectionCoordinatesLoadString = strrep(movName, '.avi', '_headDirectionCoordinates.mat');
load(frameFlagsLoadString)
load(legEndpointsStoreLoadString)
load(centroidCoordinatesLoadString)
load(headDirectionCoordinatesLoadString)

if isempty(find(frameFlags ~= 10,1)) %#ok<*NODEF> % if all flags are already 10 (movie is already corrected)
    display('Movie has already been corrected - no action was taken.')
    return % end function
else % proceed with error correction

    movAutoTracked = VideoReader(movName);
    numFrames = movAutoTracked.NumberOfFrames; %#ok<*VIDREAD>

    for i = 1:numFrames
        img = read(movAutoTracked,i);
        imwrite(img,fullfile(workingDir,'images',sprintf('img%d.jpg',i)));
    end

    imageNames = dir(fullfile(workingDir,'images','*.jpg'));
    imageNames = {imageNames.name}';

    imageStrings = regexp([imageNames{:}],'(\d*)','match');
    imageNumbers = str2double(imageStrings);

    [~,sortedIndices] = sort(imageNumbers);
    sortedImageNames = imageNames(sortedIndices);



% Error Correction

% First Pass: mark frames as right or wrong if not already done

    frameFlags(frameFlags == 2) = -10; % automatically reject frames which were deemed incorrect (uses logical indexing)
    
    for i = 1:numFrames
        if frameFlags(i) == 0 || frameFlags(i) == 1 % if frame hasn't been looked at by user before (flag = 10 or flag = -10)
                % or if we know it requires intervention (flag = 2)
            legEndpointPixelVectorHighToLow = legEndpointsStore{i};
            img = imread(fullfile(workingDir,'images',sortedImageNames{i}));
            imshow(img)
            hold on
            scatter(legEndpointPixelVectorHighToLow(:,1),legEndpointPixelVectorHighToLow(:,2),'r')
            hold off
            set(gcf, 'Position', get(0,'Screensize'));
            title('Press "y" to accept, "n" to reject, "q" to quit');

            % User key input: "y" to accept frame and set flag to 10, "n" to reject frame and set flag 
            % to -10, "q" to quit out of the function and overwrite the flags at whatever state they are 
            % (e.g. if person gets tired) if overWriteFlag == 'y', otherwise all error correction work 
            % on this movie is lost
            k=0;
            while ~k
                k = waitforbuttonpress; %#ok<NASGU>
                currkey = get(gcf,'currentcharacter');
                if strcmp(currkey,'y') % accept frame
                    frameFlags(i) = 10; %#ok<*AGROW>
                    break
                elseif strcmp(currkey,'n') % reject frame
                    frameFlags(i) = -10;
                    break
                elseif strcmp(currkey,'q') % quit before finishing
                    if isequal(overWriteFlag,'y') % save over old frameFlags and end function
                        legEndpointsSaveString = strrep(movName, '.avi', '_legEndPointsStore.mat');
                        frameFlagsSaveString = strrep(movName,'.avi','_frameFlags.mat');
                        save(legEndpointsSaveString,'legEndpointsStore')
                        save(frameFlagsSaveString,'frameFlags')
                        close
                        return
                    else % end function without saving over old frameFlags
                        return
                    end
                else
                    k = 0;
                end
            end
        end
        close
    end

    % Second Pass: correct wrong frames
    framesToCorrect = find(frameFlags == -10);
    
    for c = 1:length(framesToCorrect)
        close
        frameToCorrectIdx = framesToCorrect(c);
        legEndpointsStore{frameToCorrectIdx} = [];
        imgToCorrect = imread(fullfile(workingDir,'images',sortedImageNames{frameToCorrectIdx}));
        goodFrames = union(find(frameFlags == 10),find(frameFlags == 0)); % frames that are either correct or have high confidence
        closestGoodPriorFrame = goodFrames(find(goodFrames < frameToCorrectIdx,1,'last'));
        closestGoodPosteriorFrame = goodFrames(find(goodFrames > frameToCorrectIdx,1,'first'));

        h = tight_subplot(1,3,0,0,0); % user submitted function to remove white space between subplots (http://www.mathworks.com/matlabcentral/fileexchange/27991-tight-subplot)
        axes(h(1)) %#ok<*LAXES>
        if not(isempty(closestGoodPriorFrame))
            imgClosestGoodPriorFrame = imread(fullfile(workingDir,'images',sortedImageNames{closestGoodPriorFrame}));
            imshow(imgClosestGoodPriorFrame)
        else
            imshow([])
        end
        title(['Frame ',num2str(closestGoodPriorFrame)])
        
        axes(h(3))        
        if not(isempty(closestGoodPosteriorFrame))
            imgClosestGoodPosteriorFrame = imread(fullfile(workingDir,'images',sortedImageNames{closestGoodPosteriorFrame}));
            imshow(imgClosestGoodPosteriorFrame)
        else
            imshow([])
        end
        title(['Frame ',num2str(closestGoodPosteriorFrame)])

        while frameFlags(frameToCorrectIdx) ~= 10
            % User Clicking On Legs
            axes(h(2)) % set this last to put it in focus for clicking
            imshow(imgToCorrect)
            set(gcf, 'Position', get(0,'Screensize')); % make full-screen

            title(['Frame ',num2str(frameToCorrectIdx),' - Fix Me!'])
            xlabelString = ['Click ',num2str(targetNumLegs),' times, then press "y" to accept or "n" to start over. Press "s" to skip frame, "q" to quit'];
            xlabel(xlabelString);

            legCoordinatesUserCorrected = zeros(targetNumLegs,2);
            hold on
            
            
            [yImageSize,xImageSize] = size(imgToCorrect);
            
            for clickNum = 1:targetNumLegs
                [legCoordinatesUserCorrected(clickNum,1), legCoordinatesUserCorrected(clickNum,2)] = ginput(1);
                % clicks outside of frame are snapped to inside frame
                if legCoordinatesUserCorrected(clickNum,1) < 1
                    legCoordinatesUserCorrected(clickNum,1) = 1;
                elseif legCoordinatesUserCorrected(clickNum,1) > xImageSize
                    legCoordinatesUserCorrected(clickNum,1) = xImageSize;
                end
                
                if legCoordinatesUserCorrected(clickNum,2) < 1
                    legCoordinatesUserCorrected(clickNum,2) = 1;
                elseif legCoordinatesUserCorrected(clickNum,2) > yImageSize
                    legCoordinatesUserCorrected(clickNum,2) = yImageSize;
                end
                
                scatter(legCoordinatesUserCorrected(clickNum,1),legCoordinatesUserCorrected(clickNum,2),'r') % display the last click
            end
            hold off
            legEndpointsStore{frameToCorrectIdx} = round(legCoordinatesUserCorrected);


            % User Confirming Clicks
            waitforbuttonpress
            currkey = get(gcf,'currentcharacter');
            if strcmp(currkey,'y') % accept frame
                frameFlags(frameToCorrectIdx) = 10;
            elseif strcmp(currkey,'n') % reject frame, continue with while loop
                frameFlags(frameToCorrectIdx) = -10;
            elseif strcmp(currkey,'q') % quit before finishing
                if isequal(overWriteFlag,'y') % save over old frameFlags and end function
                    legEndpointsSaveString = strrep(movName, '.avi', '_legEndPointsStore.mat');
                    frameFlagsSaveString = strrep(movName,'.avi','_frameFlags.mat');
                    save(legEndpointsSaveString,'legEndpointsStore')
                    save(frameFlagsSaveString,'frameFlags')
                    close
                    return
                else % end function without saving over old frameFlags
                    return
                end
            elseif strcmp(currkey,'s') % skip frame (e.g. if there is a question about it)
                    % Note: this is the only way to get to the end of the function with some
                    % frameFlags not equal to 10, and it is handled at the end.
                frameFlags(frameToCorrectIdx) = -10;
                break
            else % any other button press or mouse click
                frameFlags(frameToCorrectIdx) = -10; % reject frame, continue with while loop
            end
        end
    end
    
    close
    
    if isempty(find(frameFlags ~= 10,1)) % if all frameFlags are 10 (checked and accepted by user)
        finalLegEndpointsStore = cell(1,numFrames); % for storing final sorted leg endpoints with missing legs denoted by coordinates (-1,-1)
% Order Legs
    % Note: this only occurs if all flags are 10.
        for i = 1:numFrames
            tempLegAngles = zeros(1,targetNumLegs); % angle between vectors (between 0 and pi)
            tempRotationDirection = zeros(1,targetNumLegs); % 1 = clockwise from v (right), -1 = counterclockwise from v (left)
                % (see note below).
            v = headDirectionCoordinates(i,:) - centroidCoordinates(i,:);
            tempLegEndpoints = legEndpointsStore{i};
            for j = 1:targetNumLegs
                w = tempLegEndpoints(j,:) - centroidCoordinates(i,:);
                tempLegAngles(j) = acos(dot(v,w)/(norm(v)*norm(w)));
                crossProd3D = cross([v 0],[w 0]);
                tempRotationDirection(j) = sign(crossProd3D(3));
            end
            
            % sort by increasing angle (front to back), left side first, then right side

            % left legs
            leftLegIndices = find(tempRotationDirection < 0);
            leftLegEndpoints = tempLegEndpoints(leftLegIndices,:);
            leftLegAngles = tempLegAngles(leftLegIndices);
            [~,leftSortIdx] = sort(leftLegAngles);
            leftLegEndpoints = leftLegEndpoints(leftSortIdx,:);
            
            % right legs
            rightLegIndices = find(tempRotationDirection > 0);
            rightLegEndpoints = tempLegEndpoints(rightLegIndices,:);
            rightLegAngles = tempLegAngles(rightLegIndices);
            [~,rightSortIdx] = sort(rightLegAngles);
            rightLegEndpoints = rightLegEndpoints(rightSortIdx,:);
            
            legEndpointsStore{i} = [leftLegEndpoints;rightLegEndpoints];
            
            % check for error in leg side labeling (more than three legs on one side)
            if length(rightLegEndpoints) > 3 || length(leftLegEndpoints) > 3
                frameFlags(i) = -10;
            else
                % Update to include missing leg (if any) as (-1,-1) appropriately and store in
                % finalLegEndpointStore (which we save since all frameFlags == 10 to get to this step)
                tempSortedLegEndpoints = zeros(6,2); % this has ALL 6 legs (even if there are only 5 -
                    % the sixth is marked as missing by having coordinates (-1,-1). This is for saving the
                    % final points once all frame flags == 10
                if length(leftLegIndices) < 3 % 2 left legs means front one is gone
                    tempSortedLegEndpoints(1,:) = -1;
                    tempSortedLegEndpoints(2:end,:) = [leftLegEndpoints;rightLegEndpoints];
                elseif length(rightLegIndices) < 3 % 2 right legs means front one is gone
                    tempSortedLegEndpoints(1:3,:) = leftLegEndpoints;
                    tempSortedLegEndpoints(4,:) = -1;
                    tempSortedLegEndpoints(5:end,:) = rightLegEndpoints;
                else % all legs present
                    tempSortedLegEndpoints = [leftLegEndpoints;rightLegEndpoints];
                end
                finalLegEndpointsStore{i} = tempSortedLegEndpoints;
            end
        end
    end
    
    
    if isempty(find(frameFlags ~= 10,1)) % if all frameFlags are 10 (check and accepted by user, AND user has not made an error in the labeling)
% Save frameFlags, legEndpointsStore, and Final Movie (including leg coloring).
    % Note: this only occurs if all flags are 10. Otherwise, frameFlags and legEndpointsStore is
    % simply overwritten. Also, show the final tracked movie only if showFinalMovieFlag == 'y' and frameFlags == 10


        movNameAtOutput = movName; % overwrite movie
        finalLegEndpointsStoreSaveString = strrep(movNameAtOutput, '.avi', '_finalLegEndPointsStore.mat');
        save(finalLegEndpointsStoreSaveString,'finalLegEndpointsStore')

        % overwrite in old file as well
        legEndpointsStoreSaveString = strrep(movName, '.avi', '_legEndPointsStore.mat');
        frameFlagsSaveString = strrep(movName, '.avi', '_frameFlags.mat');
        save(legEndpointsStoreSaveString,'legEndpointsStore')
        save(frameFlagsSaveString,'frameFlags')

        outputFinalTrackedMovie = VideoWriter(movNameAtOutput);
        outputFinalTrackedMovie.FrameRate = movAutoTracked.FrameRate;
        
        open(outputFinalTrackedMovie);
        for i = 1:numFrames
            finalImage = imread(fullfile(workingDir,'images',sortedImageNames{i}));
            tempLegEndpoints = finalLegEndpointsStore{i};
            
            imshow(finalImage)
            hold on
            for j = 1:6
                scatter(tempLegEndpoints(j,1),tempLegEndpoints(j,2))
                pause(finalMoviePauseTime)
            end
            scatter(centroidCoordinates(i,1),centroidCoordinates(i,2),'r')
            scatter(headDirectionCoordinates(i,1),headDirectionCoordinates(i,2),'g')
            hold off
            finalImage = getframe;
            writeVideo(outputFinalTrackedMovie,finalImage);
        end
        close(outputFinalTrackedMovie);
    else % if some frames were skipped in tagging or user made an error and labeled more than 3 legs on one side
        display(sprintf(['Note: Movie with final tracked legs not saved due to some frames being incorrect:\n',...
        'either user skipped tagging some frames or has labeled more than 3 legs on a side.\n',...
        'Leg endpoints and frame flags for corrected frames have been saved if the overWriteFlag input was "y".\n',...
        'Otherwise no action was taken.\n \n',...
        'Please re-run the error correction function for this movie.']))
        if isequal(overWriteFlag,'y')
            legEndpointsSaveString = strrep(movName, '.avi', '_legEndPointsStore.mat');
            frameFlagsSaveString = strrep(movName,'.avi','_frameFlags.mat');
            save(legEndpointsSaveString,'legEndpointsStore')
            save(frameFlagsSaveString,'frameFlags')
        end
    end
end

end