using Printf
using Grids

struct TrainingSample
    # @todo make training sample 3 fields: game, position, steps
    image::AbstractImage
    to_play::Integer # This information is encoded in the image, and is only needed for logging purposes.
    actions::Vector{<:Integer}
    targets::Vector{Target}
    image_label::String
    # @todo change image_state to initial_state
    image_state::AbstractEnvState
    states::Vector{AbstractEnvState}
    target_states::Vector{AbstractEnvState}
    function TrainingSample(image, to_play, actions, targets, states, target_states; image_label, image_state)
        @assert length(targets) == length(actions) + 1 "Num targets ($(length(targets))) must be 1 more than num actions ($(length(actions)))"
        @assert size(image)[end] == 3 "Image/Observation must have 3 channels (has $(size(image)[end])"
        new(image, to_play, actions, targets, image_label, image_state, states, target_states)
    end
end

struct TrainingSampleSVG
    sample::TrainingSample
end

svg(x::TrainingSample) = TrainingSampleSVG(x)

actions(x::TrainingSample) = x.actions
targets(x::TrainingSample) = x.targets
image(x::TrainingSample) = x.image 
image_channels(x::TrainingSample) = x |> image |> channels

function TrainingSample(game::MuZeroGame, position::Integer, num_unroll_steps::Integer, td_steps::Integer)
    image = make_image(game, position)
    image_state = get_state(game, position)
    image_label = make_image_label(game, position)
    to = min(length(history(game)), position + num_unroll_steps - 1)
    actions = map(x->x.idx, history(game)[position:to])
    targets = make_target(game, position, num_unroll_steps, td_steps)
    num_targets = length(actions)+1 
    player = to_play(position)
    states = map(idx->get_state(game, idx), position:(to+1))
    game_copy = copy_at(game.game, position)
    target_states = map(target->begin
        greedy_action = argmax(policy(target))
        step!(game_copy, greedy_action)
        get_state(game_copy)
    end, targets)
    TrainingSample(image, player, actions, targets[1:num_targets], states, target_states; image_label=image_label, image_state=image_state)
end

function TrainingSample(config::MuZeroConfig, game::MuZeroGame)
    TrainingSample(game, sample_position(game), config.num_unroll_steps, config.td_steps) 
end

tuple(x::TrainingSample) = x.image, x.actions, x.targets

get_states(x::TrainingSample) = x.states
get_target_states(x::TrainingSample) = x.target_states

function getdf(x::TrainingSample)
    inftype_col = "Inf Type"
    action_col = "Action"
    reward_col = "Target Reward"
    value_col = "Target Value"
    policy_col = "Target Policy"
    d = Dict(
        inftype_col=>String[], 
        action_col=>Union{Int,Nothing}[], 
        reward_col=>Float32[], 
        value_col=>Float32[],
        policy_col=>String[],
    )
    # @todo change this to x.action_space_size
    num_policy_actions = x.targets |> first |> policy |> length
    for i in 1:num_policy_actions
        d["π[$i]"] = Float32[]
    end

    states = get_states(x)
    actions = [nothing, x.actions...]
    targets = x.targets
    for (action_idx, zipped) in enumerate(zip(states, actions, targets))
        (s,a,target) = zipped
        push!(d[inftype_col], action_idx == 1 ? "Initial" : "Recurrent")
        push!(d[action_col], a)
        push!(d[reward_col], reward(target))
        push!(d[value_col], value(target))

        pred = policy(target)
        limit = 3
        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[1:end]))"
        end, zip(actions, prs))
        result_policy = join(action_prs, ", ")
        push!(d[policy_col], result_policy)
        # π = policy(target)
        # for i in 1:num_policy_actions
        #     push!(d["π[$i]"], round(π[i], sigdigits=2))
        # end
    end

    DataFrame(
        inftype_col=>d[inftype_col],
        action_col=>d[action_col],
        reward_col=>d[reward_col],
        value_col=>d[value_col],
        policy_col=>d[policy_col],
    )
end

function getvrgrid(target)
    value = @sprintf "%.1f" target.value
    reward = @sprintf "%.1f" target.reward
    vr = Grid(
        [
            LabelRect("value: $value", width=200, halign=:left),
            LabelRect("reward: $reward", width=200, halign=:left),
        ],
        :, 
        1,
        margin_left=0,
        line_thickness=0
    )
    Grid(
        [vr],
        :, 
        1,
        header="TARGET V/R", 
        margin_left=10,
        line_thickness=0
    )
end

function getinputgrid(s, header)
    player_name = s.to_play.name
    state_grid = getgrid(s; footer="to play: $player_name", margin_left=0)
    Grid(
        [state_grid], 
        :, 
        1, 
        header=header, 
        margin_left=0, 
        line_thickness=0
    )
end

function getpolicygrid(s, target)
    pred = target.policy
    pred_states = Rect[]
    limit = 3
    from = 1 
    to = min(limit, length(pred))
    pred_actions = reverse(sortperm(pred))[from:to]
    prs = pred[pred_actions]
    for (a,pr) in zip(pred_actions, prs)
        pr_str = @sprintf "%.2f" pr
        push!(
            pred_states,
            getgrid(
                s, 
                a, 
                0f0, 
                getnextstate(s, a); 
                margin_left=0, 
                footer="action: $a",
                subfooter=(text=pr_str, color="#666"),
            )
        )
    end
    Grid(
        pred_states, 
        :, 
        1;
        header="TARGET π",
        margin_left=0,
        line_thickness=0,
    )
end

function getinferencegrid(s, target, header, input_header)
    input_grid = getinputgrid(s, input_header)
    vr_grid = getvrgrid(target)
    policy_grid = getpolicygrid(s, target)
    policy_col = Grid([Grid(3,3), Grid(3,3)], :, 1, header=(text="TARGET π", offset=(10,0)))
    Grid(
        [input_grid, vr_grid, policy_grid],
        header=header, 
        margin_left=0,
        line_thickness=0,
    )
end

function getrecurrentinfgrid(sample)
    state_targets = zip(sample.states[2:end],sample.targets[2:end])
    map(st->begin
        s, target = st
        getinferencegrid(
            s, 
            target, 
            "RECURRENT INFERENCE", 
            "HIDDEN STATE"
        )
    end, state_targets)
end

function Base.show(io::IO, m::MIME"image/svg+xml", x::TrainingSampleSVG)
    sample = x.sample
    initial_grid = getinferencegrid(
        sample.image_state, 
        sample.targets[1], 
        "INITIAL INFERENCE", 
        "INPUT"
    )
    recurrent_grids = getrecurrentinfgrid(sample)
    g = Grid(
        [initial_grid, recurrent_grids...], 
        :, 
        1, 
        header=(text="TRAINING SAMPLE"),
        line_thickness=0,
        margin=15,
    )
    show(io, m, g)
end

function Base.show(io::IO, x::TrainingSample)
    action_hl = Highlighter((data, i, j)->(j==2) && data[i,j] == nothing , crayon"#666666");
    pretty_table(
        io,
        getdf(x); 
        show_row_number=true, 
        nosubheader=true, 
        title="Training Sample",
        highlighters=(action_hl),
    )
end

function gridsvg(x::TrainingSample)
    first_input_state, remaining_states... = get_states(x)
    first_target_state, remaining_target_states... = x.target_states
    first_target, remaining_targets... = x.targets
    state_targets = zip(get_states(x), x.targets)
    state_type = typeof(x.image_state)
    initial_grid = GridSVG([
        gridsvg(first_input_state), 
        targetsvg(first_input_state, first_target_state, first_target)
    ], dim=2)
    grids = [
        initial_grid,
    ]
    for (input_state, action, target, target_state) in zip(remaining_states, x.actions, remaining_targets, remaining_target_states)
        push!(grids, GridSVG([
            gridsvg(input_state, title="hiddenstate + action") 
            targetsvg(input_state, target_state, target)
        ], dim=2))
    end
    GridSVG(grids, main_title="Training Sample")
end

function get_player_label(x::TrainingSample, position)
    if x.to_play == 1
        return isodd(position) ? "A" : "B"
    end
    isodd(position) ? "B" : "A"
end