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

%% Configuration
discRnnPerfCsvFile = fullfile( ...
    '/tmpscratch', 'amotta', 'klinger', ...
    '2019-07-23-inh-exc-relu-rnn-with-pruning', ...
    'exported-csvs', '20190926T160505', '20190723T111932.csv');

memRnnPerfCsvFiles = { ...
    '20200301T185156.csv', '20200301T185220.csv', ...
    '20200301T185923.csv', '20200301T185946.csv'};
memRnnPerfCsvFiles = fullfile( ...
    '/tmpscratch', 'amotta', 'klinger', ...
    '2020-03-01-inh-exc-relu-memory-rnn-with-pruning', ...
    'tfsummaries', '20200318T175925', memRnnPerfCsvFiles);

% See "circuit constraints" section of methods
numExc = 1800; connExc = 0.2;
numInh = 200; connInh = 0.6;
numAll = numExc + numInh;

numConns = ...
    numExc * (connExc * (numAll - 1)) ...
  + numInh * (connInh * (numAll - 1));
fracConn = ...
    numConns / (numAll * (numAll - 1));

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

%% Loading data
[discTrainT, discValidT] = ...
    RNN.loadPerformance( ...
        discRnnPerfCsvFile, {'loss', 'accuracy'});

[~, memValidT] = cellfun( ...
    @(f) RNN.loadPerformance(f, {'loss'}), ...
    memRnnPerfCsvFiles, 'UniformOutput', false);

%% Pruning
clear cur*;
pruningIdx = diff(discTrainT.weightCount);
pruningIdx = find([false; pruningIdx < 0]);
pruningSteps = discTrainT.step(pruningIdx);

%% Accuracy and connectivity over training
clear cur*;

curY = [-1; 0] + reshape(pruningIdx, 1, []);
curY = discTrainT.weightCount(curY(:));

curX = [0; repelem(pruningSteps, 2)];
curY = [curY(1); curY] / (numAll * (numAll - 1));

curFig = figure();
curAx = axes(curFig);

yyaxis(curAx, 'left');
hold(curAx, 'on');

% Accuracy
plot(curAx, discValidT.step, discValidT.accuracy, 'k-');

% Pruning events
for s = reshape(pruningSteps, 1, [])
    plot(curAx, repelem(s, 2), 1 + [0.05, 0.1], 'k-');
end

curIdx = discTrainT.weightCount(pruningIdx);
curIdx = curIdx / (numAll * (numAll - 1));
[~, curIdx] = min(abs(curIdx - fracConn));

plot(curAx, repelem(pruningSteps(curIdx + 1), 2), [0, 1], 'k-');

xlabel(curAx, 'Training steps');
ylabel(curAx, 'Performance');

curAx.Children = circshift(curAx.Children, 1);
curAx.YTick = [0, 0.5, 1];
curAx.YLim(end) = 1.1;

curFig.Color = 'white';
curFig.Position(3:4) = [325, 160];
curAx.TickDir = 'out';

% Connectivity
yyaxis(curAx, 'right');
plot(curAx, curX, curY, 'LineWidth', 2);

curAx.YScale = 'log';
curLimY = log10(curAx.YLim);
curTicksY = 10 .^ (curLimY(1):curLimY(end));
curLimY(end) = curLimY(end) + (1.1 - 1) * diff(curLimY);

ylim(curAx, 10 .^ curLimY);
yticks(curAx, curTicksY);
ylabel(curAx, 'Connectivity');

%% Zoom-in
curZoomFig = curFig;
curZoomAx = curZoomFig.Children;

yyaxis(curZoomAx, 'left');
curZoomPlots = curZoomAx.Children;

xlim(curZoomAx, [3.15, 3.35] * 1E6);
curZoomFig.Position(3:4) = [205, 160];

%% Accuracy delta vs. connection count
clear cur*;

curX = discTrainT.weightCount(pruningIdx - 1);
curX = curX / (numAll * (numAll - 1));

curDrop = arrayfun(@(i) ...
    discValidT.accuracy(discValidT.step == discTrainT.step(i - 1)) ...
  - discValidT.accuracy(discValidT.step == discTrainT.step(i)), ...
    pruningIdx);

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

plot(curAx, curX, curDrop, 'k', 'LineWidth', 2);

curAx.XDir = 'reverse';
curAx.XScale = 'log';
curAx.XTick = 10 .^ (log10(curAx.XLim(1)):0);
curAx.XTickLabel = arrayfun( ...
    @num2str, curAx.XTick, ...
    'UniformOutput', false);

xlabel(curAx, 'Connectivity');
ylabel(curAx, ['Accuracy drop', newline, 'by pruning']);

curFig.Color = 'white';
curFig.Position(3:4) = [325, 160];
curAx.TickDir = 'out';

%% Illustration of connectome pruning
clear cur*;
rng(1);

N = 10;
wLim = [-2, +2];
excN = numExc / numAll * N;
inhN = numInh / numAll * N;

nonNan = @(d) d(not(isnan(d)));
conn = genConn(N, excN, wLim);
cmap = redBlue(128);

plotConnCirc(conn);

for i = 1:10
    % Training
    modConn = rand(N);
    conn = ...
        modConn .* genConn(N, excN, wLim) ...
      + (1 - modConn) .* conn;
  
    % After training
    plotConnCirc(conn);
    
    % Pruning
    thresh = conn(not(isnan(conn)));
    thresh = prctile(abs(thresh), 10);
    conn(abs(conn) < thresh) = nan;
    
    % After pruning
    plotConnCirc(conn);
end

%% Utilities
function conn = genConn(N, excN, wLim)
    inhN = N - excN;
    excM = reshape(1:N <= excN, [], 1);
    inhM = reshape(1:N > excN, [], 1);

    conn = randn(N);
    conn(excM, :) = conn(excM, :) * inhN / excN;
    conn = max(min(conn, wLim(2)), wLim(1));
    
    conn(excM, :) = +abs(conn(excM, :));
    conn(inhM, :) = -abs(conn(inhM, :));
    conn(1:(N + 1):end) = nan;
end

function plotConnCirc(conn)
    N = size(conn, 1);
    circRadius = 0.4;
    neuronRadius = 0.05;
    connShift = neuronRadius / 4;

    fig = figure();
    fig.Position(3:4) = 300;

    for n = 1:N
        nPos = 0.5 + circRadius .* [ ...
            cos(2 * pi / N * n), sin(2 * pi / N * n)];
        annotation('ellipse', [ ...
            nPos - neuronRadius, ...
            2 * neuronRadius, 2 * neuronRadius]);

        for m = 1:N
            if isnan(conn(n, m)); continue; end
            mPos = 0.5 + circRadius .* [ ...
                cos(2 * pi / N * m), sin(2 * pi / N * m)];
            mMinusN = mPos - nPos;

            mMinusNLen = sqrt(sum(mMinusN .^ 2));
            mMinusN = mMinusN / mMinusNLen;

            mTheta = acos(mMinusN(1));
            if mMinusN(2) < 0; mTheta = 2 * pi - mTheta; end

            mTheta = mTheta - pi / 2;
            mShift = connShift .* [cos(mTheta), sin(mTheta)];

            sPos = nPos + 1.5 * neuronRadius * mMinusN;
            tPos = mPos - 1.5 * neuronRadius * mMinusN;

            ann = annotation('line', ...
                [sPos(1), tPos(1)] + mShift(1), ...
                [sPos(2), tPos(2)] + mShift(2));
            ann.LineWidth = 3 * max(eps, abs(conn(n, m)));
        end
    end
end

function cmap = redBlue(n)
    m = round(0.5 * n);
    cmap = nan(n, 3);
    
    cmap(:, 1) = min((1:n) / m, 1);
    cmap(:, 2) = 1 - abs(1 - (1:n) / m);
    cmap(:, 3) = min(2 - (1:n) / m, 1);
end
