using DataFrames: DataFrames
using PrettyTables: pretty_table

using ReCo: ReCo

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

    return nothing
end

"""
    export_q_matrix(env_helper::ReCo.RL.EnvHelper, filename_without_extension::String)

Generate a LaTeX table to the Q-matrix of `env_helper`.

The output is `ReCo.jl/exports/filename_without_extension.tex`. `env_helper` has to be an environment helper with the abstract type `EnvHelper`.

Return `nothing`.
"""
function export_q_matrix(env_helper::ReCo.RL.EnvHelper, filename_without_extension::String)
    table = copy(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 = DataFrames.DataFrame(table, states)
    DataFrames.insertcols!(df, 1, :Actions => actions)

    latex_table(df, "$filename_without_extension.tex")

    return nothing
end