using DataFrames
using PrettyTables: pretty_table

function latex_table(dataframe, filename::String; path="exports/$filename")
    open(path, "w") do f
        pretty_table(f, dataframe; backend=:latex, nosubheader=true, alignment=:c)
    end

    return nothing
end

function latex_rl_table(env_helper, filename)
    table = env_helper.shared.agent.policy.learner.approximator.table

    for col in 1:size(table)[2]
        table[:, col] ./= sum(table[:, col])
    end

    table .= round.(table, digits=2)

    state_spaces_labels = env_helper.shared.env.shared.state_spaces_labels
    states = AbstractString[]
    for i in state_spaces_labels[1]
        for j in state_spaces_labels[2]
            push!(states, i * " ; " * j)
        end
    end

    action_spaces_labels = env_helper.shared.env.shared.action_spaces_labels
    actions = AbstractString[]

    for i in action_spaces_labels[1]
        for j in action_spaces_labels[2]
            push!(actions, i * "; " * j)
        end
    end

    df = DataFrame(table, states)
    insertcols!(df, 1, :Actions => actions)

    latex_table(df, filename)

    return nothing
end