function [alignedRFs, lagDiffs] = alignRFsToAbsMax(receptiveFields, padWithNaNs, varargin)
    if isempty(receptiveFields)
        warning('Empty array, no RF to align!!')
        return
    end
    switch nargin
        case 2
            lagDiffs = calculateRFShifts(receptiveFields);
        case 3
            lagDiffs = varargin{1};
        case 4
            lagDiffs = varargin{1};
            polarity = varargin{2};
            if isempty(lagDiffs)
                lagDiffs = calculateRFShifts(receptiveFields, polarity);
            end
        otherwise
            lagDiffs = calculateRFShifts(receptiveFields);
    end
    alignedRFs = shiftReceptiveFields(receptiveFields, lagDiffs, padWithNaNs);
end

function lagDiffs = calculateRFShifts(receptiveFields, varargin)
    nEpochs = size(receptiveFields, 2);
    if ~isempty(varargin)
        polarity = varargin{1};
    else
        polarity = '';
    end
    % Get reference ROIs with center close to the center of screen.
    switch polarity
        case 'on'
            [~, maxInd] = min((receptiveFields), [], 2);
        case 'off'
            [~, maxInd] = max((receptiveFields), [], 2);
        otherwise
            [~, maxInd] = max(abs(receptiveFields), [], 2);
    end
    % Consider signal to noise ration to avoid biasing towards noisy RFs.
    %%
    medianToStd = @ (x) median(x, 2) ./ std(x, [], 2);
    signalToNoise = medianToStd(abs(receptiveFields));
    % Take only data above 90 percentile.
    snrCutoff = quantile(signalToNoise, 0.7);
    %%
    % Check distribution of center locations.
    distanceToCenter = abs(maxInd - nEpochs / 2);
    distanceCutoff = max(quantile(distanceToCenter, 0.5), 6);
    validRFCenters = distanceToCenter <= distanceCutoff & signalToNoise >= snrCutoff;
    referenceRF = mean(receptiveFields(validRFCenters, :), 1);
    % Align center RFs to reference RF.
    alignedRefRFs = nan * ones(size(receptiveFields, 1), nEpochs);
    for iCell = find(validRFCenters)'
        [alignedRefRFs(iCell, :), ~] = crossCorrShiftRF(receptiveFields(iCell, :), referenceRF);
    end
    referenceRF = nanmean(alignedRefRFs, 1);
    alignedRFs = ones(size(receptiveFields)) * nan;
    % Shift single 
    for iCell = 1: size(receptiveFields, 1)
        [alignedRFs(iCell, :), lagDiffs(iCell)] = crossCorrShiftRF(receptiveFields(iCell, :), referenceRF);
    end
    % Shift reference to center to get total shift.
    [~, maxInd] = max(abs(nanmean(alignedRFs, 1)));
    globalShift = floor(maxInd - nEpochs / 2);
%     alignedRFs = circshift(alignedRFs, -globalShift, 2);
    lagDiffs = lagDiffs + globalShift;
end

function alignedRFs = shiftReceptiveFields(receptiveFields, lagDiffs, padWithNaNs)
    nEpochs = size(receptiveFields, 2);
    for iCell = 1: size(receptiveFields, 1)
        alignedRFs(iCell, :) = circshift(receptiveFields(iCell, :), -lagDiffs(iCell));
        if lagDiffs(iCell) > nEpochs
            alignedRFs(iCell, :) = nan * receptiveFields(iCell, :);
        end
    end
    
    lagDiffs = -lagDiffs;
    for iCell = 1: size(receptiveFields, 1)
     % Set the shifted portion of tuning curve to nan.
        if padWithNaNs
            if sign(lagDiffs(iCell)) > 0
                alignedRFs(iCell, 1: lagDiffs(iCell)) = nan;
            elseif abs(lagDiffs(iCell)) < nEpochs || sign(lagDiffs(iCell)) < 0
                alignedRFs(iCell, end + lagDiffs(iCell): end) = nan;
            else
                alignedRFs(iCell, :) = nan;
            end
        end
    end
end