using DataFrames
using Grids
using PrettyTables

struct RecurrentRVPTest <: TrainingTest
    key::String
    state::AbstractEnvState
    state_next::AbstractEnvState
    image::AbstractImage
    action::AbstractEnvAction
    action_image::AbstractImage
    # @todo - "actions" should be named better
    # @todo - enable rvp
    # rvp::PolicyValueReward 
    # @todo - should be "actions::Vector{<:AbstractEnvAction}"
    actions::Vector{<:Integer}
    value::Float32
    reward::Float32
end

function Base.show(io::IO, m::MIME"text/plain", x::RecurrentRVPTest)
    println("\nRECURRENT RVP TEST #$(x.key)")
    println("\nTo Play:\n$(x.state.to_play.name)")
    println("\nAction:\n$(x.action.idx)")
    println("\nReward:\n$(x.reward)")
    println("\nNext To Play:\n$(x.state_next.to_play.name)")
    println("\nNext Optimal Action(s):\n$(join(map(y->y.idx, x.actions), ","))")
    println("\nNext Value:\n$(x.value)")
end

struct RecurrentRVPTestSVG
    test::RecurrentRVPTest
end
svg(x::RecurrentRVPTest) = RecurrentRVPTestSVG(x)

function Base.show(io::IO, m::MIME"image/svg+xml", x::RecurrentRVPTestSVG)
    test = x.test
    s = test.state
    a = test.action
    r = test.reward
    s_next = test.state_next
    actions = test.actions
    state_grid = getgrid(
        s; 
        header="to play: $(s.to_play.name)", 
        footer="action: $(a.idx)", 
        margin_left=0
    )
    state_next_grid = getgrid(
        s_next; 
        header="to play: $(s_next.to_play.name)", 
        margin_left=0
    )
    next_states = Rect[]
    from = 1
    len = length(actions)
    to = min(len, 3)
    for (idx,a) in enumerate(actions[from:to])
        push!(
            next_states,
            getgrid(s_next, a, r, getnextstate(s_next, a); header="action: $(a.idx)", margin_left=0)
        )
    end
    if length(actions) > 3
        push!(next_states, Grid(1,1, (r,c)->"..."; line_thickness=0))
    end
    hiddenstate_grid = Grid(
        [
            getgrid(
                s, 
                a, 
                0f0, 
                s_next; 
                header="to play: $(s_next.to_play.name)", 
                margin_left=0
            )
        ], 
        :, 
        1, 
        header="HIDDEN STATE", 
        margin_left=0, 
        line_thickness=0
    )
    policy_grid_header = "ACTION TEST"
    policy_grid = Grid(
        next_states, 
        :, 
        1, 
        header=(text=policy_grid_header, height=40), 
        margin_left=0, 
        line_thickness=0,
        bg_color=RandColor()
    )
    g = Grid(
        [
            hiddenstate_grid,
            policy_grid
        ]; 
        cell_valign=:top, 
        header="RECURRENT RVP TEST #$(test.key)", 
        line_thickness=0
    )
    show(io, m, g)
end