using JLD2, Printf, ArgParse;

include("runit.jl");
include("helpers_experiments.jl");
include("helpers.jl");

function parse_commandline()
    s = ArgParseSettings();

    @add_arg_table! s begin
        "--exp_dir"
            help = "Directory for loading the experiment's data, while saving the plots in the corresponding folder."
            arg_type = String
            default = "experiments/"
        "--format"
            help = "Format to save the figures."
            arg_type = String
            default = "png"
        "--no_plot_samp"
            help = "Don't plot the sample complexity"
            action = :store_true
        "--no_plot_arms"
            help = "Don't plot arms empirical allocation (in proportion)"
            action = :store_true
        "--no_plot_cput"
            help = "Don't plot the CPU running time"
            action = :store_true
        "--plot_hist"
            help = "Plot the Histograms"
            action = :store_true
        "--print_table"
            help = "Plot the Histograms"
            action = :store_true
    end

    parse_args(s);
end

parsed_args = parse_commandline();
exp_dir = parsed_args["exp_dir"];
format = parsed_args["format"];
no_plot_samp = parsed_args["no_plot_samp"];
no_plot_arms = parsed_args["no_plot_arms"];
no_plot_cput = parsed_args["no_plot_cput"];
plot_hist = parsed_args["plot_hist"];
print_table = parsed_args["print_table"];

experiments = filter(x -> isdir("$(exp_dir)$x"), readdir(exp_dir));
for experiment in experiments
    # Check for existing plots in the corresponding format
    existing_plots = length(filter(x -> occursin(".$format", x), readdir("$(exp_dir)$(experiment)/"))) != 0;

    if ! existing_plots
        println("Creating plots in .$format based on $(exp_dir)$(experiment) ...");

        name_experiment = split(experiment, ":")[2];
        @load "$(exp_dir)$(experiment)/$(name_experiment).dat" dists μs Tstar wstar pep iss data δs Nruns seed;

        # Plot histograms of the arm distrbution
        if plot_hist
            file_plot_hist = "$(exp_dir)$(experiment)/plot_hist_$(name_experiment).$format"
            plot_histogram(dists, file_plot_hist)
        end

        if print_table
            for i in 1:length(δs)
                delta = split(string(δs[i]), ".")[2];
                file_table_timeout = "$(exp_dir)$(experiment)/table_timeout_$(name_experiment)_delta$(delta).txt";
                plot_table_timeout(pep, μs, Tstar, wstar, δs[i], iss, getindex.(data, i), Nruns, file_table_timeout);
            end
        end

        # Plot the sample complexity per learner in a boxplot fashion
        if ! no_plot_samp
            for i in 1:length(δs)
                delta = split(string(δs[i]), ".")[2];
                file_plot_samp = "$(exp_dir)$(experiment)/plot_samp_$(name_experiment)_delta$(delta).$format";
                plot_samp(pep, μs, Tstar, wstar, δs[i], iss, getindex.(data, i), Nruns, file_plot_samp);
            end
        end

        # Plot the empirical proportion of allocation over arms per learner in a boxplot fashion (one per arm)
        if ! no_plot_arms
            for i in 1:length(δs)
                delta = split(string(δs[i]), ".")[2];
                file_plot_arms = "$(exp_dir)$(experiment)/plot_arms_$(name_experiment)_delta$(delta).$format";
                plot_arms(pep, μs, Tstar, wstar, δs[i], iss, getindex.(data, i), Nruns, file_plot_arms);
            end
        end

        # Plot the CPU time per learner in a boxplot fashion
        if ! no_plot_cput
            for i in 1:length(δs)
                delta = split(string(δs[i]), ".")[2];
                file_plot_cput = "$(exp_dir)$(experiment)/plot_cput_$(name_experiment)_delta$(delta).$format";
                plot_cput(pep, μs, Tstar, wstar, δs[i], iss, getindex.(data, i), Nruns, file_plot_cput);
            end
        end
    else
        println("Plots in .$format based on $(exp_dir)$(experiment) already exist.");
    end
    println("");
end
