%% Make fig 2 e, f
close all;
clear all;
clc;

%% Define search path.
parentDir = fileparts(fileparts(mfilename('fullpath')));
saveFolder = [parentDir filesep 'data'];
load([saveFolder filesep 'Fig2-Tm9-Bars-Noise-processedTableWithBackgroundOnOffNoiseCropped'], 'ProcessedTable');

%%
figDir = [parentDir filesep 'figures' filesep 'Fig2' filesep];
if ~exist(figDir, 'dir')
    mkdir(figDir); 
end

%% Define search path.
datasetInd = 3;
[~, ~, stimSetCellStr] = parseDatasetVars(datasetInd);

%%
ProcessedTable = ProcessedTable(ProcessedTable.typeOfProblem == 0, :);
% Check only bars.
nStims = 8; % 8 for fff, bars, and noise; 5 for fff and bars.
stimSetCellStr = stimSetCellStr(1: nStims);
%% Split table into cell types.
tokens = arrayfun(@(x)  regexp(x, '.*Tm(\d).*', 'tokens'), ProcessedTable.timeSeriesPath);
cellTypeNumArray = cellfun(@(x) str2num(x{1}{1}), tokens);
cellTypeNum = unique(cellTypeNumArray);
for iNeuron = 1: numel(cellTypeNum)
    StimTableNeuronCell{iNeuron} = ProcessedTable(cellTypeNumArray == cellTypeNum(iNeuron), :);
end
%%
stimInd = 1;
nStims = numel(stimSetCellStr);
for iNeuron = 1: numel(cellTypeNum)
    for jStim = 1: nStims
        StimTableNeuronStim{iNeuron, jStim} = getStimSubsetTable(StimTableNeuronCell{iNeuron}, stimSetCellStr, jStim);
    end
end

for iFly = 1: numel(unique(ProcessedTable.flyInd))
    flyInds = ProcessedTable.flyInd == iFly;
    flyStims = ProcessedTable(flyInds, :).stimParamFileName;
    % Find if stim is within the chosen subset.
    flyStimInds{iFly} = (cellfun(@(x) find(strcmp(x, stimSetCellStr)), flyStims, 'uni', 0));
    % If it contains stimuli, check that it includes the 5 chosen.
    if ~isempty(flyStimInds{iFly})
        setDiff{iFly} = setdiff(1: nStims, [flyStimInds{iFly}{:}]);
    else
        setDiff{iFly} = -1;
    end
    % For flies with all chosen stimuli order the table accordingly.
    if isempty(setDiff{iFly})
        sortFlyIndsTemp = find(flyInds);
        [~, sortInd] = sort([flyStimInds{iFly}{:}]);
        sortFlyInds{iFly} = sortFlyIndsTemp(sortInd);
    end
end

StimOrderedTable = ProcessedTable(cat(1, sortFlyInds{:}), :);
%% Split table into cell types.

tokens = arrayfun(@(x)  regexp(x, '.*Tm(\d).*', 'tokens'), StimOrderedTable.timeSeriesPath);
cellTypeNumArray = cellfun(@(x) str2num(x{1}{1}), tokens);
cellTypeNum = unique(cellTypeNumArray);
for iNeuron = 1: numel(cellTypeNum)
    StimTableNeuronCell{iNeuron} = StimOrderedTable(cellTypeNumArray == cellTypeNum(iNeuron), :);
end
%%
nStims = numel(stimSetCellStr);
for iNeuron = 1: numel(cellTypeNum)
    for jStim = 1: nStims
        StimTableNeuronStim{iNeuron, jStim} = getStimSubsetTable(StimTableNeuronCell{iNeuron}, stimSetCellStr, jStim);
    end
end

OriginalStimTableNeuronStim = StimTableNeuronStim;

%% Delete invalid rois for all flies and stimuli.

bgRois = cellfun(@(y) cellfun(@(x) x.backgroundRois, y.roiMeta), StimTableNeuronStim, 'uni', 0);
%Test if all background rois are the same for all stimuli.
assert(isequal(bgRois{:}));
invalidRois = zeros(size(StimTableNeuronStim{1}, 1), size(StimTableNeuronStim, 2));
for iFly = 1: size(StimTableNeuronStim{1}, 1)
    for jStim = 1: size(StimTableNeuronStim, 2)
        invalidInds = StimTableNeuronStim{jStim}.roiMeta{iFly}.invalidRois;
        if ~isempty(invalidInds)
            invalidRois(iFly, jStim) = invalidInds;
        end
    end
end

roisToDeletePerFly = sum(unique(invalidRois', 'rows'))';

%% Extract the parameters for the analysis of spatial RFs.
jBar = 1;
for iStim = 2:5
    dummyTable = StimTableNeuronStim{iStim};
    paramsCell = cellfun(@(x) x(1: end - 1*0), dummyTable.paramTuning, 'uni', 0);
    responseIndex = cellfun(@(x, y) x(1: end - 1*0), dummyTable.responseIndex, 'uni', 0);
    for jFly = find(roisToDeletePerFly ~= 0)
        if invalidRois(jFly, iStim) == 0        
            paramsCell{jFly}(roisToDeletePerFly(jFly) - 1*0) = [];
            responseIndex{jFly}(roisToDeletePerFly(jFly) - 1*0) = [];
        end
    end
    amp(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.a1, x, 'uni', 1), paramsCell, 'uni', 0)');
    pos(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.b1, x, 'uni', 1), paramsCell, 'uni', 0)');
    width(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.c1, x, 'uni', 1), paramsCell, 'uni', 0)');
    rsquare(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.gof.rsquare, x, 'uni', 1), paramsCell, 'uni', 0)');
    jBar = jBar + 1;
end


for iStim = [7, 8]
    dummyTable = StimTableNeuronStim{iStim};
    paramsCell = cellfun(@(x) x(1: end - 1*0), dummyTable.paramTuning, 'uni', 0);
%     responseIndex = cellfun(@(x, y) x(1: end - 1), dummyTable.responseIndex, 'uni', 0);                    
    for jFly = find(roisToDeletePerFly ~= 0)
        if invalidRois(jFly, iStim) == 0
            paramsCell{jFly}(roisToDeletePerFly(jFly) - 1*0) = [];
%             responseIndex{jFly}(roisToDeletePerFly(jFly) - 1*0) = [];
        end
    end
    amp(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.a1, x, 'uni', 1), paramsCell, 'uni', 0)');
    pos(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.b1, x, 'uni', 1), paramsCell, 'uni', 0)');
    width(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.fit.c1, x, 'uni', 1), paramsCell, 'uni', 0)');
    rsquare(jBar, :) = cell2mat(cellfun(@(x) arrayfun(@(y) y.gof.rsquare, x, 'uni', 1), paramsCell, 'uni', 0)');
    jBar = jBar + 1;
end

pos(1: 4, :) = 2 * (pos(1: 4, :) - 1);
pos(5: 6, :) = 5 * pos(5: 6, :);
width = (2 * sqrt(log(2))) * width;
width(1: 4, :) = 2 * width(1: 4, :);
width(5: 6, :) = 5 * width(5: 6, :);
validInds = all(rsquare([1, 2, 5, 6], :) > 0.5) & all(rsquare([3, 4], :) > 0.1);
%%
flyInds = repelem(1:size(StimTableNeuronStim{1}),cellfun(@numel, StimTableNeuronStim{1}.responseIndex));
flyInds(roisToDeletePerFly(1)) = [];
flyInds(validInds)
numel(unique(flyInds(validInds)))
%%
clear allRFs
for iStim = 2: 5
    tuningMethods = StimTableNeuronStim{iStim}.nonParamTuning{1}.tuningMethods;
    tuningMethodInd = strcmp(tuningMethods, 'extreme');
    tcCellArray = cellfun(@(x) squeeze(x.tc{tuningMethodInd}(:, 1: end - 1, :)), StimTableNeuronStim{iStim}.nonParamTuning, 'uni', 0);
    allRFs{iStim - 1} = cell2mat(tcCellArray);
end
%%
noiseRFs{1} = getSpaceFilterArrayFromTable(StimTableNeuronStim{7}); % Az is horizontal
noiseRFs{2} = getSpaceFilterArrayFromTable(StimTableNeuronStim{8}); % El is vertical
%%
% Delete invalid roi that shifts all images by one;                                                       
barRFs = cellfun(@(x) x(setdiff(1: length(allRFs{1}), roisToDeletePerFly(1)), 1: end), ...
                 allRFs, 'uni', 0);
noiseRFs{1}(roisToDeletePerFly(1), :) = [];

%%             
% The screen is about 48-60 deg, we center the RFs at about the center 24
% deg, using the fit positions for higher presicion.
lagsBars = round((pos(1: 4, :) - 25) / 2);
for iBar = 1: numel(barRFs)
    barRFs{iBar} = alignRFsToAbsMax(barRFs{iBar}, 0, lagsBars(iBar, :)');
end

lagsNoise = round((pos(5: 6, :) - 30) / 5);
for iBar = 1: numel(noiseRFs)
    noiseRFs{iBar} = alignRFsToAbsMax(noiseRFs{iBar}, 0, lagsNoise(iBar, :));
end

%% Blank epoch was already deleted, now delete vertical bars on edges of screen.
barRFs{2}(:, end-1: end) = [];
barRFs{4}(:, end-1: end) = [];
barRFs{2}(:, 1: 2) = [];
barRFs{4}(:, 1: 2) = [];
%% Fig2e
hFig = figure;
splitSign = 1;
doSmoothing = 1;
indsToPlot = validInds;
subplot_tight(1, 3, 1);
plotRFEnsemble(noiseRFs{1}, noiseRFs{2}, splitSign, doSmoothing, indsToPlot);
title('White noise spatial filters')
% Plot bar RFs. OFF.
subplot_tight(1, 3, 2);
plotRFEnsemble(barRFs{1}, barRFs{2}, splitSign, doSmoothing, indsToPlot);
title('OFF flashing bar filters')
% Plot bar RFs. ON.
subplot_tight(1, 3, 3);
plotRFEnsemble(barRFs{3}, barRFs{4}, splitSign, doSmoothing, indsToPlot);
title('ON flashing bar filters')
suptitle('Cartesian product (X \times Y) of one dimensional tuning curves');

figFileName = ['RFIm-ON-OFF-Noise'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')
% print(hFig, [figDir figFileName '.png'], '-dpng', '-r
%% Write data for fig2e
dataFileName = [figDir filesep 'Fig2.xls'];
DataTable = createRFsTable(noiseRFs{1}(indsToPlot, :), 'Noise_horizontal_position_deg');
% Correct spacing from 2 deg for bars to 5 deg for noise.
DataTable(:, 1).Noise_horizontal_position_deg = 5 / 2 * DataTable(:, 1).Noise_horizontal_position_deg;
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e')
DataTable = createRFsTable(noiseRFs{1}(indsToPlot, :), 'Noise_vertical_position_deg');
% Correct spacing from 2 deg for bars to 5 deg for noise.
DataTable(:, 1).Noise_vertical_position_deg = 5 / 2 * DataTable(:, 1).Noise_vertical_position_deg;
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e', 'Range', 'A15')

DataTable = createRFsTable(barRFs{1}(indsToPlot, :), 'OFF_horizontal_position_deg');
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e', 'Range', 'A29')
DataTable = createRFsTable(barRFs{2}(indsToPlot, :), 'OFF_vertical_position_deg');
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e', 'Range', 'A56')
DataTable = createRFsTable(barRFs{3}(indsToPlot, :), 'ON_horizontal_position_deg');
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e', 'Range', 'A83')
DataTable = createRFsTable(barRFs{4}(indsToPlot, :), 'ON_vertical_position_deg');
writetable(DataTable, dataFileName, 'Sheet', 'Fig2e', 'Range', 'A110')

%%
offX = barRFs{1}(validInds, :);
offY = barRFs{2}(validInds, :);
onX = barRFs{3}(validInds, :);
onY = barRFs{4}(validInds, :);
lnX = noiseRFs{1}(validInds, :);
lnY = noiseRFs{2}(validInds, :);

lnX = interp1(linspace(2.5,60, 12), lnX', linspace(2.5,60, 25))';
lnY = interp1(linspace(2.5,60, 12), lnY', linspace(2.5,60, 25))';
% Realign noise to center now at 2 deg precision.
[~, peakInd] = max(abs(lnX), [], 2);
lnX = alignRFsToAbsMax(lnX, 0, peakInd - round(length(lnX) / 2) - 1);
[~, peakInd] = max(abs(lnY), [], 2);
lnY = alignRFsToAbsMax(lnY, 0, peakInd - round(length(lnY) / 2) - 1);

maxNorm = @(x, dim) x ./ max(abs(x), [], dim);
getRow = @(x, nRow) x(nRow, :);
normRFs = cellfun(@(x) cell2mat(arrayfun(@(y) smooth(getRow(maxNorm(x,2), y))', 1: size(x, 1), 'uni', 0)'), ...
                  {offX, offY, onX, onY, lnX, lnY}, 'uni', 0);
normRFs = cellfun(@(x) maxNorm(x,2), normRFs, 'uni', 0);

hFig = createPrintFig(15 * [ 1 1 / 3]);
for iRFType = 1:3
    subplot_tight(1, 3, iRFType);
    imshow(getRFIm(mean(normRFs{2*iRFType}), mean(normRFs{2*iRFType - 1})))
end

figFileName = ['RFIm-ON-OFF-Noise-mean'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')

hFig = createPrintFig(15 * [ 1 2/3]);
hSubAx = plotLasagnaTracesFromCell(gcf, normRFs, brewermap(6, 'Paired'));
ylim(hSubAx(:, 1), [-1 1]);
arrayfun(@(x) set(get(x, 'YAxis'), 'Visible', 'on'), hSubAx(1, :))
cellfun(@(x) caxis(ZeroCenteredBounds(x, [0 100])), normRFs, 'uni', 0)
arrayfun(@prettifyAxes, hSubAx)
arrayfun(@offsetAxes, hSubAx)
set(hSubAx, 'Visible', 'off')
figFileName = ['RFcurves-ON-OFF-Noise-XY-lasagna'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')

%%
hFig = createPrintFig(15 * [1 1/2]);
pairsToAverage = {[1, 2], [3, 4], [5, 6]};
meanWidths = cellfun(@(x) mean(width(x, validInds)), pairsToAverage, 'uni', 0);
meanWidths = cellfun(@(x) mean(width(x, validInds)), pairsToAverage, 'uni', 0);

beanPlot(meanWidths, {'off', 'on', 'noise'}, flipud(brewermap(3, 'Set1')), 1, gca, [0-eps 50], 'vertical', false, true)
figFileName = ['RF-ON-OFF-Noise-FWHM-Bean'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')
%% Save data for Fig 2f
xlswrite(dataFileName, {'FHWM-OFF-bars'}, 'Fig2f', 'A1')
xlswrite(dataFileName, meanWidths{1}, 'Fig2f', 'B1')
xlswrite(dataFileName, {'FHWM-ON-bars'}, 'Fig2f', 'A2')
xlswrite(dataFileName, meanWidths{2}, 'Fig2f', 'B2')
xlswrite(dataFileName, {'FHWM-Noise-bars'}, 'Fig2f', 'A3')
xlswrite(dataFileName, meanWidths{3}, 'Fig2f', 'B3')
%% Permutation test for mean RF width difference.
[pValOFFvsON, obsDiff] = permutationTest(meanWidths{1},meanWidths{2}, 100000)
[pValONvsNoise, obsDiff] = permutationTest(meanWidths{2},meanWidths{3}, 100000)
[pValOFFvsNoise, obsDiff] = permutationTest(meanWidths{1},meanWidths{3}, 100000)
%%
hFig = createPrintFig(15 * [1/2 1/2]);

meanAmps = cellfun(@(x) mean(amp(x, validInds)), pairsToAverage(1:2), 'uni', 0);

beanPlot(meanAmps, {'off', 'on'}, flipud(brewermap(3, 'Set1')), 1, gca,'unbounded', 'vertical', false, true)
figFileName = ['RF-ON-OFF-Noise-Amp-Bean'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')

%% Save data for Fig 2f
xlswrite(dataFileName, {'Amp-OFF-bars'}, 'Fig2f', 'A5')
xlswrite(dataFileName, meanAmps{1}, 'Fig2f', 'B5')
xlswrite(dataFileName, {'Amp-ON-bars'}, 'Fig2f', 'A6')
xlswrite(dataFileName, meanAmps{2}, 'Fig2f', 'B6')
%%
% pairsToPlot = {[1 5], [2 6], [3 5], [4 6]};
pairsToPlot = {[1 2], [3 4], [1 4], [1 3], [2 4], [2 3]};
labelsCell = {'X OFF', 'Y OFF', 'X ON', 'Y ON', 'X Noise', 'Y Noise'};
indsToPlot = validInds;
% fwhmArray = [mean(width(1: 2, :)); mean(width(3: 4, :)); mean(width(5: 6, :))]';
fwhmArray = width';
rSquare = rsquare';
% hFig = createFullScreenFig;
% set(hFig, 'Units', 'Centimeters')
% set(hFig, 'Position', [hFig.Position(1:2) + 10 14 14]);
hFig = createPrintFig(15 * [ 1 2/3]);


for iPair = 1: numel(pairsToPlot)
    hScattAx(iPair) = subplot(2, 3, iPair); hold on;
    xInd = pairsToPlot{iPair}(1);
    yInd = pairsToPlot{iPair}(2);
    scatter(fwhmArray(indsToPlot, xInd), fwhmArray(indsToPlot, yInd), [], ...
            min(rSquare(indsToPlot, xInd), rSquare(indsToPlot, yInd)), 'filled');  
    [correlation, pValue] = corr(fwhmArray(indsToPlot, xInd), fwhmArray(indsToPlot, yInd), 'type', 'Pearson' );
    x = fwhmArray(indsToPlot, xInd);
    [b, bInt]= regress(fwhmArray(indsToPlot, yInd), [ones(length(x), 1) x]);
    hFit = plot(x, b(1) + x*b(2), 'Color', [1 1 1] * 0.5);
    maxVal = max([fwhmArray(indsToPlot, xInd); fwhmArray(indsToPlot, yInd)]);
    maxVal = quantile(vec(fwhmArray(indsToPlot, :)), 1);
    plot([0 maxVal], [0 maxVal], 'Color', [1 1 1] * 0.5);
    prettifyAxes(gca);
    set(gca, 'XLim', [0 maxVal], 'YLim', [0 maxVal], 'CLim', [0 1]);
    axis square;
    xlabel(labelsCell{xInd});
    ylabel(labelsCell{yInd});
    title({['\rho = ' num2str(correlation, 2) ', p = ' num2str(pValue, 2) ],...
           ['m = ' num2str(b(2), 2) ' (' num2str(bInt(2, 1), 2) ', ' num2str(bInt(2, 2), 2) ')']})
    setFavoriteColormap
end
setFontForThesis(hScattAx, gcf);

arrayfun(@prettifyAxes, hScattAx)
arrayfun(@offsetAxes, hScattAx)
% set(hSubAx, 'Visible', 'off')
figFileName = ['BarRFs-Corr-Pearson'];    
print(hFig, [figDir figFileName '.pdf'], '-dpdf')
%%
offWidth = mean(width(1: 2, indsToPlot), 1);
onWidth = mean(width(3: 4, indsToPlot), 1);
[corrWidth, pCorrWidth] = corr(offWidth(:), onWidth(:))

%%
offWidth = width(1: 2, indsToPlot);
onWidth = width(3: 4, indsToPlot);
[corrWidth, pCorrWidth] = corr(offWidth(:), onWidth(:))
%%
[corrWidth, pCorrWidth] = corr(width(1:4, indsToPlot)', 'type','Spearman' )