# This file is part of Jlsca, license is GPLv3, see https://www.gnu.org/licenses/gpl-3.0.en.html
#
# Author: Cees-Bart Breunesse

using RecipesBase

export PlotRanksEvolution

"""
Plots the evolution of the rank of the correct key bytes for all phases and targets. You need to know the key to use this plot.

# Example

```
using Plots

params.knownKey = hex2bytes("00112233445566778899aabbccddeeff")
params.updateInterval = 1000
rankdata = sca(trs,params)

plot(PlotRanksEvolution(),rankdata,params)
```
"""
struct PlotRanksEvolution end

@recipe function f(::PlotRanksEvolution, rankdata::RankData, params)
    if ismissing(params.knownKey) 
        error("need params.knownKey (or if you don't know the key, look at PlotScoresEvolution)")
    end
    
    xxrows = Vector[]
    rrankings = Vector[]
    labels = String[]
    
    for phase in getPhases(rankdata)
        xrows = getNrConsumedRowsEvolution(rankdata,phase)
        for target in getTargets(rankdata, phase)
            kb = getCorrectKey(params,phase,target)
            rankings = getRankingsEvolution(rankdata, phase, target, kb)            
            push!(xxrows,xrows)
            push!(rrankings,rankings)
            push!(labels, "$phase / $target / $(rankings[end])")
        end
    end
    
    color := :auto
    label := reshape(labels, (1, length(labels)))

    legendtitlefonthalign = :left
    legendtitlefontsize := 8
    legendtitle := "phase / target / final rank"
    ylabel := "rank"
    xlabel --> "#traces"
    title --> "rank evolution"
    xxrows,rrankings
end

export PlotScoresEvolution

"""
Plots the evolution of scores for a given sca result, parameter set, phase and target.

# Example
Plot the score evolution for phase 1, target 2
```
using Plots

params.knownKey = hex2bytes("00112233445566778899aabbccddeeff")
params.updateInterval = 1000
rankdata = sca(trs,params)

plot(PlotScoresEvolution(),rankdata,params,1,2)
```
If you don't know the key and params.knownKey is not set (or set to missing), this plot will show you the evolution of the 5 highest ranked key candidates.
```
params.knownKey = missing
params.updateInterval = 1000
plot(PlotScoresEvolution(),rankdata,params,1,2)
```    

"""
struct PlotScoresEvolution end

@recipe function f(::PlotScoresEvolution, rankdata::RankData, params, phase, target, combined=true, leakage=1)
    xrows = getNrConsumedRowsEvolution(rankdata,phase)
    if combined
        scores = getScoresEvolution(rankdata, phase, target)
    else
        scores = getScoresEvolution(rankdata, phase, target, leakage)
    end
    cands = size(scores)[1]

    if !ismissing(params.knownKey) 
        kb = getCorrectKey(params,phase,target)
        if combined
            rankings = getRankingsEvolution(rankdata, phase, target, kb)
        else
            rankings = getRankingsEvolution(rankdata, phase, target, leakage, kb)
        end

        color := reshape([x == kb + 1 ? :red : :grey for x in 1:cands],(1,cands))
        label := reshape([x == kb + 1 ? "correct ($(rankings[end]))" : x == kb + 2 ? "incorrect" : "" for x in 1:cands],(1,cands))
    else
        if combined
            finalscores = getScores(rankdata, phase, target)
        else
            finalscores = getScores(rankdata, phase, target, leakage)
        end
        ranked = sortperm(finalscores, rev=true)[1:5]
        color := reshape([x in ranked ? :auto : :grey for x in 1:cands],(1,cands))
        label := reshape([x in ranked ? "0x$(string(x-1, base=16))" : "" for x in 1:cands],(1,cands))
    end

    if rankdata.nrLeakages > 1
        if combined
            ylabel := "$(params.leakageCombinator) of scores"
        else
            if isdefined(params.analysis,:leakages)
                ylabel := "$(params.analysis.leakages[leakage]) score"
            else
                ylabel := "score $leakage"
            end
        end
    else
        ylabel := "score"
    end

    xlabel := "#traces"
    title := "$(params.analysis) evolution phase $phase, target $target"
    xrows,scores'
end

