using JLD2, Printf, StatsPlots, ArgParse;
using LaTeXStrings;

include("runiteps.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 = "pdf"
    end

    parse_args(s);
end

parsed_args = parse_commandline();
exp_dir = parsed_args["exp_dir"];
format = parsed_args["format"];

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;

        delta = split(string(δs[1]), ".")[2];

        if true
            # Detailed plots on Simple Regret

            # cst-EB-TC versus other anytime algos
            iss_to_keep = ["cst-EB-TC-e1-G-GK16", "DSR-G-GK16", "DSH-G-GK16", "RR-G-GK16"];
            name_iss = ["EB-TC" * L"_{\varepsilon}", "DSR", "DSH", "Uniform"];
            file_plot_sregret = "$(exp_dir)$(experiment)/plot_bench_ALL_detailed_sregret_$(name_experiment)_delta$(delta).$format";
            plot_sregret(pep, μs, Tstar, wstar, δs[1], iss, getindex.(data, 1), Nruns, file_plot_sregret, iss_to_keep, name_iss);

            # Comparison between EB-TC versions
            iss_to_keep = ["EB-TC-e15-G-GK16", "EB-TC-e1-G-GK16", "EB-TC-e05-G-GK16",
                           "cst-EB-TC-e15-G-GK16", "cst-EB-TC-e1-G-GK16", "cst-EB-TC-e05-G-GK16"];
            name_iss = ["EB-TC" * L"_{0.15}", "EB-TC" * L"_{0.1}", "EB-TC" * L"_{0.05}", "c-EB-TC" * L"_{0.15}", "c-EB-TC" * L"_{0.1}", "c-EB-TC" * L"_{0.05}"];
            file_plot_sregret = "$(exp_dir)$(experiment)/plot_bench_ATT_detailed_sregret_$(name_experiment)_delta$(delta).$format";
            plot_sregret(pep, μs, Tstar, wstar, δs[1], iss, getindex.(data, 1), Nruns, file_plot_sregret, iss_to_keep, name_iss);
        end

        if false
            # Empirical Stopping time (without fine history)

            # EB-TC versus other fixed-confidence algos
            iss_to_keep = ["EB-TC-e1-G-GK16", "cst-TS-TC-G-GK16", "cst-EB-TCI-G-GK16", "cst-UCB-TC-TB-G-GK16",
                           "TaS-C-G-GK16", "EpsFWS-C-G-GK16", "DKM-C-G-GK16", "LUCB-G-GK16", "RR-G-GK16"];
            name_iss = ["EB-TC" * L"_{\varepsilon}", L"\varepsilon" * "-T3C", L"\varepsilon" * "-EB-TCI", L"\varepsilon" * "-TTUCB",
                        L"\varepsilon" * "-TaS", L"\varepsilon" * "-FWS", L"\varepsilon" * "-DKM", "LUCB", "Uniform"];
            file_plot_samp = "$(exp_dir)$(experiment)/plot_bench_ALL_samp_$(name_experiment)_delta$(delta).$format";
            plot_samp(pep, μs, Tstar, wstar, δs[1], iss, getindex.(data, 1), Nruns, file_plot_samp, iss_to_keep, name_iss);

            # Comparison between EB-TC versions
            iss_to_keep = ["EB-TC-e15-G-GK16", "EB-TC-e1-G-GK16", "EB-TC-e05-G-GK16",
                           "cst-EB-TC-e15-G-GK16", "cst-EB-TC-e1-G-GK16", "cst-EB-TC-e05-G-GK16"];
            name_iss = ["EB-TC" * L"_{0.15}", "EB-TC" * L"_{0.1}", "EB-TC" * L"_{0.05}", "c-EB-TC" * L"_{0.15}", "c-EB-TC" * L"_{0.1}", "c-EB-TC" * L"_{0.05}"];
            file_plot_samp = "$(exp_dir)$(experiment)/plot_bench_ATT_samp_$(name_experiment)_delta$(delta).$format";
            plot_samp(pep, μs, Tstar, wstar, δs[1], iss, getindex.(data, 1), Nruns, file_plot_samp, iss_to_keep, name_iss);
        end

        if false
            # Empirical Stopping time for BAI algorithms

            # Comparison with varying slack algorithms
            iss_to_keep = ["cst-EB-TC-e1-G-GK16", "cst-EB-TC-La5-G-GK16", "cst-EB-TC-La1-G-GK16", "cst-EB-TC-La05-G-GK16",
                           "cst-EB-TC-Pa5-G-GK16", "cst-EB-TC-Pa1-G-GK16", "cst-EB-TC-Pa05-G-GK16",
                           "cst-TS-BTC-G-GK16", "cst-EB-BTCI-G-GK16", "cst-UCB-BTC-T-G-GK16",
                           "TaS-C-G-GK16", "LUCB-G-GK16", "RR-G-GK16"];
            name_iss = ["EB-TC" * L"_{\varepsilon}", "la5", "la1", "la05", "pa5", "pa1", "pa05", "T3C", "EB-TCI", "TTUCB", "TaS", "LUCB", "Uniform"];
            file_plot_samp = "$(exp_dir)$(experiment)/plot_bench_VSATT_samp_$(name_experiment)_delta$(delta).$format";
            plot_samp(pep, μs, Tstar, wstar, δs[1], iss, getindex.(data, 1), Nruns, file_plot_samp, iss_to_keep, name_iss);

        end

    else
        println("Plots in .$format based on $(exp_dir)$(experiment) already exist.");
    end
    println("");
end
