using PrettyTables

struct PolicyTestResults
    results::Vector{PolicyTestResult}
end

function PolicyTestResults(x::MuZeroConfig)
    getpolicytests()
    d = policy_testset(x.game_type)
    tests = []
    for (idx, kv) in enumerate(d)
        game, actions = kv
        push!(tests, PolicyTest(idx, game, actions))
    end
    PolicyTestResults(tests)
end

struct PolicyTestResultsSVG
    tests::PolicyTestResults
end

using Random

svg(x::PolicyTestResults) = PolicyTestResultsSVG(x)
Base.rand(x::PolicyTestResults) = rand(x.results)
Base.getindex(x::PolicyTestResults, idx::Integer) = x.results[idx]
Base.getindex(x::PolicyTestResults, idx::Vector{<:Integer}) = x.results[idx]
Base.length(x::PolicyTestResults) = length(x.results)
Base.lastindex(x::PolicyTestResults) = length(x)
Base.iterate(x::PolicyTestResults, state=1) = state > length(x) ? nothing : (x[state], state+1)

function df(x::PolicyTestResults; range=1:length(x), limit=3)
    test_num = "Test #"
    test_pf = "P/F"
    test_result = "Pred (π*)"
    optimal_action = "Actual (π*)"
    d = Dict(
        test_num=>[],
        test_pf=>[],
        test_result=>[],
        optimal_action=>[],
    )
    for result in x.results[range]
        idx = result.test.idx
        optimal = result.test.actions
        pred = result.prediction
        _, greedy_action = findmax(pred)
        pass = in(greedy_action, optimal)
        from = 1 
        to = min(limit, length(pred))
        actions = reverse(sortperm(pred))[from:to]
        prs = pred[actions]
        action_prs = map(x->begin
            action, pr = x
            pr_s = @sprintf "%.2f" pr
            "$action ($(pr_s[2:end]))"
        end, zip(actions, prs))
        result_policy = join(action_prs, ", ")
        push!(d[test_num], idx)
        push!(d[test_pf], pass ? "Pass" : "Fail")
        push!(d[test_result], result_policy)
        push!(d[optimal_action], join(optimal, ","))
    end
    DataFrame(
        test_num=>d[test_num],
        test_pf=>d[test_pf],
        test_result=>d[test_result],
        optimal_action=>d[optimal_action],
    )
end

function Base.show(io::IO, m::MIME"text/plain", x::PolicyTestResults)
    len = length(x.results)
    println()
    pretty_results(
        io,
        df(x); 
        title="Policy Test Results (all $len)"
    )
end

function pretty_results(io::IO, df::DataFrame; title)
    pf_hl = Highlighter((data, i, j)->(j==2) && data[i,j] == "Pass", crayon"green");
    pretty_table(
        io,
        df; 
        show_row_number=false, 
        nosubheader=true, 
        title=title,
        highlighters=(pf_hl),
    )
end

function Base.getindex(x::PolicyTestResults, r::UnitRange{<:Integer})
    println()
    pretty_results(
        stdout,
        df(x; range=r), 
        title="Policy Test Results ($(r))"
    )
end