%
% Copyright 2021 Renjie Wu
%
% Licensed under the Apache License, Version 2.0 (the "License");
% you may not use this file except in compliance with the License.
% You may obtain a copy of the License at
%
%     http://www.apache.org/licenses/LICENSE-2.0
%
% Unless required by applicable law or agreed to in writing, software
% distributed under the License is distributed on an "AS IS" BASIS,
% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
% See the License for the specific language governing permissions and
% limitations under the License.
%

function TEASERResultPlotter(outputJsonPath)
    global json orderedPredictions sortedId resultFigure;
    json = jsondecode(fileread(outputJsonPath));
    
    earlinesses = [json.predictions.earliness];
    uniqueEarlinesses = unique(earlinesses);
    probabilities = [json.predictions.probability];
    maxProbablities = nonzeros(accumarray(earlinesses', probabilities', [], @max));
    
    probablityDistribution = zeros(length(uniqueEarlinesses), 6);
    for i=1:length(uniqueEarlinesses)
        probabilitiesPerEarliness = probabilities(earlinesses == uniqueEarlinesses(i));
        probablityDistribution(i,:) = histcounts(probabilitiesPerEarliness, [0.5 0.7 0.8 0.9 0.95 0.98 1], 'normalization', 'probability');
        probablityDistribution(i,:) = probablityDistribution(i,:) * maxProbablities(i);
    end

    figure;
    histX = 1:length(uniqueEarlinesses);
    
    subplot(2, 1, 1);
    bar(histX, histc(earlinesses, uniqueEarlinesses));
    xticks(histX);
    xticklabels(uniqueEarlinesses);
    xlabel('Number of datapoints seen when early classification made');
    ylabel('Count');
    title('Histogram on Earlinesses');
    
    subplot(2, 1, 2);
    bar(histX, probablityDistribution, 'stacked');
    legend({'50%\leq{P}<70%','70%\leq{P}<80%','80%\leq{P}<90%','90%\leq{P}<95%','95%\leq{P}<98%','98%\leq{P}\leq{100%}'}, 'location', 'eastoutside');
    
    xticks(histX);
    xticklabels(uniqueEarlinesses);
    yticks(0:0.2:1);
    yticklabels({'0%','20%','40%','60%','80%','100%'});
    xlabel('Number of datapoints seen when early classification made');
    ylabel('Highest probability for predicted label');
    title('Probability Distribution');
    
    suptitle(sprintf('Statistics for %s on %s', json.classifierName, json.datasetName));
    
    orderedPredictions = struct2table(json.predictions);
    orderedPredictions.exemplarId = [1:length(json.predictions)]';
    orderedPredictions = sortrows(orderedPredictions, [2 5], {'ascend', 'descend'});
    orderedPredictions = table2struct(orderedPredictions);
    
    sortedId = 1;
    resultFigure = figure('pos', [42 42 733 962], 'keypressfcn', @onKeyPress);
    updateFigure;
end

function onKeyPress(~, event)
    global orderedPredictions sortedId;
    
    if event.Key == "leftarrow" && sortedId > 1
        sortedId = sortedId - 1;
        updateFigure;
    elseif event.Key == "rightarrow" && sortedId < length(orderedPredictions)
        sortedId = sortedId + 1;
        updateFigure;
    elseif event.Key == "r"
        sortedId = randi([1 length(orderedPredictions)]);
        updateFigure;
    end
end

function updateFigure()
    global json orderedPredictions sortedId resultFigure;
    
    drawFigure(orderedPredictions(sortedId));
    suptitle(sprintf('TEASER on %s testing #%d', json.datasetName, orderedPredictions(sortedId).exemplarId));
    set(resultFigure, 'name', sprintf('%s on %s testing #%d', json.classifierName, json.datasetName, orderedPredictions(sortedId).exemplarId));
end

function drawFigure(predict)
    exemplar = predict.exemplar;
    earlyPredictions = predict.earlyPredictions;
    X = zeros(length(earlyPredictions), 1);
    Y = zeros(length(earlyPredictions), 1);

    for j = 1:length(earlyPredictions)
        X(j) = earlyPredictions(j).snapshotLength;
        Y(j) = earlyPredictions(j).probability;

        if earlyPredictions(j).label == 2
            Y(j) = -Y(j);
        end
    end

    hold on;
    subplot(2, 1, 1);
    plot(exemplar);
    xline(predict.earliness, '--r');

    xlim([-10 X(end)+10]);
    xticks(unique(sort([0:50:X(end) predict.earliness])));

    title('Testing Time Series');

    subplot(2, 1, 2);
    plot(X, Y);
    yline(0);
    xl = xline(predict.earliness, '--r', {'Early Classification Made'});
    xl.LabelVerticalAlignment = 'middle';

    xlabel('Number of datapoints seen')
    xlim([-10 X(end)+10]);
    xticks(unique(sort([0:50:X(end) predict.earliness])));

    ylabel({'Probability for predicted label'; 'Predicted as 2            Predicted as 1'})
    ylim([-1.2 1.2]);
    yticks(-1:0.2:1);
    yticklabels({'100%','80%','60%','40%','20%','0%','20%','40%','60%','80%','100%'});

    title(sprintf('Actual: %d, Predicted: %d, Confidence: %.2f%%', ...
        predict.truthLabel, predict.predictedLabel, predict.probability * 100));
    
    hold off;
end