% Written by
%   Alessandro Motta <alessandro.motta@brain.mpg.de>
clear;

%% Configuration
nExc = 1800;
pelVec = linspace(0, 1, 101);
pefVec = linspace(0, 1, 101);
nlVec = 2:12;

peeBarrel = [0.15, 0.25];
reeBarrel = [0.15, 0.35];

info = Util.runInfo();
Util.showRunInfo(info);

%% Evaluate parameters
curData = cell(1, 3);
[curData{:}] = ndgrid(pelVec, pefVec, nlVec);

paramT = table;
paramT.pel = curData{1}(:);
paramT.pef = curData{2}(:);
paramT.nl = curData{3}(:);
Util.clear(curData);

paramT.pee = ...
    ... lateral connections in all layers
    (1 ./ paramT.nl .* paramT.pel) ...
    ... forward connections in all layers except last one
  + (1 - 1 ./ paramT.nl) .* (1 ./ paramT.nl .* paramT.pef);

% Reciprocity: P(reciprocal connection | connection)
% * forward connections cannot be reciprocated, by definition
% * lateral connections are reciprocated with probability pel
% Reciprocity: lateral connectvity × P(is lateral connection | connection)
paramT.ree = paramT.pel .* ((1 ./ paramT.nl .* paramT.pel) ./ paramT.pee);

nlBarrel = unique(paramT.nl( ...
    ((peeBarrel(1) <= paramT.pee) & (paramT.pee <= peeBarrel(2))) ...
  & ((reeBarrel(1) <= paramT.ree) & (paramT.ree <= reeBarrel(2)))));

%% Plotting
clear cur*;

curFig = figure();
curAx = axes(curFig);
hold(curAx, 'on');

axis(curAx, 'square');
xlim(curAx, [0, 1]);
ylim(curAx, [0, 1]);

curColors = parula(numel(nlVec));
colormap(curAx, curColors);
curAx.CLim = nlVec([1, end]) + [-0.5, +0.5];

for nlIdx = 1:numel(nlVec)
    curNl = nlVec(nlIdx);
    curColor = curColors(nlIdx, :);
    
    curT = paramT(paramT.nl == curNl, :);
    curT(isnan(curT.ree), :) = [];
    
    curX = curT.pee; curY = curT.ree;
    curIds = boundary(curX, curY);
    
    patch(curX(curIds), curY(curIds), curColor);
end

curP = plot( ...
    peeBarrel([1, 2, 2, 1, 1]), ...
    reeBarrel([1, 1, 2, 2, 1]), ...
    'Color', 'black', 'LineWidth', 2);

curCbar = colorbar(curAx);
curCbar.Ticks = nlVec;

curSuffix = {'', ' (ok)'};
curCbarLabels = arrayfun( ...
    @(n, b) sprintf('%d%s', n, curSuffix{1 + b}), ...
    nlVec(:), ismember(nlVec(:), nlBarrel), ...
    'UniformOutput', false);
curCbar.TickLabels = curCbarLabels;

ylabel(curCbar, 'n_l');

xlabel(curAx, 'p_{ee}');
ylabel(curAx, 'r_{ee} (NOT rr_{ee})');

Figure.config(curFig, info);
curFig.Position(3:4) = [360, 250];

%% Reproducing panel S4h
% NOTE(amotta): For derivation of these equations, see:
% https://gitlab.mpcdf.mpg.de/connectomics/Emmanuel/-/blob/480f0adedcf3a10d78392940f6d1d3c5a7227920/figures/ll_reparametrization.ipynb
pee = 0.2;
ree = [0.17, 0.17 + 0.06];

pdfRaw = 1 / diff(ree);
pdfMax = 11.5;
pdfMin = 9.0;

curPee = repelem(pee, 101);
curRee = linspace(ree(1), ree(2), 101);

curColors = linspace(0.975, 0, 101);
curColors = repmat(curColors(:), [1, 3]);

curFig = figure();
curAx = axes(curFig);
hold(curAx, 'on');
colormap(curAx, curColors);

for curNl = 2:4
    curPef = (curNl^2 .* curPee ...
       - sqrt(curNl^3 .* curPee .* curRee)) / (curNl - 1);
    curPel = sqrt(curNl * curPee .* curRee);
    
    curPdf = diff([curPef(:), curPel(:)], 1, 1);
    curPdf = sqrt(sum(curPdf .* curPdf, 2));
    curPdf = (diff(curRee(:)) * pdfRaw) ./ curPdf;
    
    assert(min(curPdf) > pdfMin);
    assert(max(curPdf) < pdfMax);
    
    curPdfColors = (curPdf - pdfMin) / (pdfMax - pdfMin);
    curPdfColors = round(curPdfColors * size(curColors, 1));
    curPdfColors = curColors(curPdfColors, :);
    
    for s = 1:(numel(curPdf) - 1)
        curP = plot(curAx, curPef(s + [0, 1]), curPel(s + [0, 1]));
        curP.Color = curPdfColors(s, :);
        curP.LineWidth = 2;
    end
end

curX = [0.19, 0.57];
curY = [0.26, 0.43];

plot(curAx, ...
    curX([1, 2, 2, 1, 1]), ...
    curY([1, 1, 2, 2, 1]), ...
    'k--');

curLims = [ ...
    min(curAx.XLim(1), curAx.YLim(1)), ...
    max(curAx.XLim(2), curAx.YLim(2))];

curAx.CLim = [pdfMin, pdfMax];
colorbar(curAx);

xlabel(curAx, 'p_{ef}');
xlim(curAx, [0.15, 0.6]);
xticks(curAx, 0.2:0.1:0.6);

ylabel(curAx, 'p_{el}');
ylim(curAx, [0.15, 0.6]);
yticks(curAx, 0.2:0.1:0.6);
axis(curAx, 'square');

Figure.config(curFig, info);
curFig.Position(3:4) = [285, 210];
