using Statistics, StatsPlots, ColorSchemes;
using LaTeXStrings;


pal_colors = palette(:tab20);
blues_colors = palette(:Blues_9);
greens_colors = palette(:Greens_9);
set19_colors = palette(:Set1_9);
set28_colors = palette(:Set2_8);
cblind_colors = palette(:seaborn_colorblind6);

dict_colors = Dict("EB-TCa-e1" => cblind_colors[3],
		   "TTUCB" => set19_colors[4],
                   "EB-TCI" => set19_colors[3],
                   "T3C" => set19_colors[1],
                   "LUCB" => set19_colors[8],
                   "TaS" => cblind_colors[1],
                   "FWS" => cblind_colors[5],
                   "DKM" => cblind_colors[6],
                   "Uniform" => set19_colors[9]);

dict_colors_Bench_ATT = Dict("EB-TCa-e15" => cblind_colors[1],
                             "c-EB-TCa-e15" => cblind_colors[2],
                             "EB-TCa-e1" => cblind_colors[3],
                             "c-EB-TCa-e1" => cblind_colors[4],
                             "EB-TCa-e05" => cblind_colors[5],
                             "c-EB-TCa-e05" => cblind_colors[6]);

Ies = [1, 2, 3, 4, 5, 6, 7, 8];

#
# 2G instances
#

# Empirical stopping times for N=1000 runs

μs_cst_EBTCa_15 = [2910, 4732, 6167, 7565, 8652, 9927, 10981, 12335];
σs_cst_EBTCa_15 = [810, 1401, 1660, 1757, 2070, 2227, 2449, 2697];
μs_cst_EBTCa_1 = [2864, 4245, 5531, 6789, 7889, 9110, 10136, 11087];
σs_cst_EBTCa_1 = [1106, 1253, 1554, 1757, 1846, 2209, 2449, 2546];
μs_cst_EBTCa_05 = [3937, 9429, 13451, 16995, 20276, 23669, 26311, 29072];
σs_cst_EBTCa_05 = [3106, 4220, 5197, 5694, 6412, 7293, 7593, 8587];
μs_cst_T3C = [2665, 4260, 5547, 6758, 8044, 9132, 10299, 11376];
σs_cst_T3C = [811, 1319, 1566, 1773, 2030, 2237, 2431, 2540];
μs_cst_EBTCI = [2848, 4518, 5770, 7046, 8210, 9386, 10466, 11503];
σs_cst_EBTCI = [752, 1311, 1554, 1737, 1948, 2101, 2258, 2370];
μs_cst_TTUCB = [3108, 5587, 7297, 8780, 9982, 11215, 12245, 13181];
σs_cst_TTUCB = [923, 2007, 2157, 2135, 2126, 2147, 2280, 2387];
μs_EBTCa_15 = [2498, 4400, 5721, 6819, 7929, 8940, 9943, 11044];
σs_EBTCa_15 = [805, 1503, 1688, 1963, 2109, 2369, 2465, 2583];
μs_EBTCa_1 = [2448, 3911, 4893, 5951, 6933, 7897, 8767, 9657];
σs_EBTCa_1 = [995, 1359, 1593, 1703, 1976, 2045, 2303, 2510];
μs_EBTCa_05 = [3343, 9055, 12435, 15273, 18129, 19959, 22758, 24305];
σs_EBTCa_05 = [3166, 4274, 4947, 5545, 6057, 6995, 7388, 7854];
μs_TaS = [3484, 5397, 7071, 8699, 10584, 12111, 13872, 15506];
σs_TaS = [1077, 1592, 1993, 2327, 2810, 3064, 3548, 3811];
μs_FWS = [3657, 5587, 7163, 8785, 10352, 11930, 13657, 15276];
σs_FWS = [1248, 1707, 2061, 2323, 2747, 3158, 3591, 3903];
μs_DKM = [3391, 7857, 10106, 11467, 13038, 13924, 14571, 15483];
σs_DKM = [1082, 3114, 3320, 3487, 3462, 3342, 3364, 3517];
μs_LUCB = [4507, 9033, 11952, 14607, 17136, 19715, 22102, 24578];
σs_LUCB = [1477, 2124, 2523, 2967, 3462, 3855, 4241, 4521];
μs_Uniform = [4832, 14516, 16937, 18020, 18705, 19209, 19456, 19833];
σs_Uniform = [1772, 5886, 5460, 5090, 5039, 4924, 4647, 4456];

plot(Ies, μs_cst_EBTCa_15, ribbon=σs_cst_EBTCa_15, fillalpha=.2, label="c-EB-TC" * L"_{0.15}", color=dict_colors_Bench_ATT["c-EB-TCa-e15"], markershape=:auto, linestyle=:auto, legend=:topleft, xlabel=L"|i^{\star}(\mu)|", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ies, μs_cst_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.2, label="c-EB-TC" * L"_{0.1}", color=dict_colors_Bench_ATT["c-EB-TCa-e1"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_cst_EBTCa_05, ribbon=σs_cst_EBTCa_05, fillalpha=.2, label="c-EB-TC" * L"_{0.05}", color=dict_colors_Bench_ATT["c-EB-TCa-e05"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_EBTCa_15, ribbon=σs_EBTCa_15, fillalpha=.2, label="EB-TC" * L"_{0.15}", color=dict_colors_Bench_ATT["EB-TCa-e15"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_EBTCa_1, ribbon=σs_EBTCa_1, fillalpha=.2, label="EB-TC" * L"_{0.1}", color=dict_colors_Bench_ATT["EB-TCa-e1"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_EBTCa_05, ribbon=σs_EBTCa_05, fillalpha=.2, label="EB-TC" * L"_{0.05}", color=dict_colors_Bench_ATT["EB-TCa-e05"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "good_answers_bench_ATT_N1000.pdf"
savefig(file);

# Plot with eps_0 = eps_1, and using IDS
plot(Ies, μs_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.2, label="EB-TC" * L"_{\varepsilon}", color=dict_colors["EB-TCa-e1"], markershape=:auto, linestyle=:auto, legend=:topleft, xlabel=L"|i^{\star}(\mu)|", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ies, μs_cst_T3C, ribbon=σs_cst_T3C, fillalpha=.2, label=L"\varepsilon" * "-T3C", color=dict_colors["T3C"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_cst_EBTCI, ribbon=σs_cst_EBTCI, fillalpha=.2, label=L"\varepsilon" * "-EB-TCI", color=dict_colors["EB-TCI"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_cst_TTUCB, ribbon=σs_cst_TTUCB, fillalpha=.2, label=L"\varepsilon" * "-TTUCB", color=dict_colors["TTUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_TaS, ribbon=σs_TaS, fillalpha=.2, label=L"\varepsilon" * "-TaS", color=dict_colors["TaS"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_FWS, ribbon=σs_FWS, fillalpha=.2, label=L"\varepsilon" * "-FWS", color=dict_colors["FWS"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_DKM, ribbon=σs_DKM, fillalpha=.2, label=L"\varepsilon" * "-DKM", color=dict_colors["DKM"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_LUCB, ribbon=σs_LUCB, fillalpha=.2, label="LUCB", color=dict_colors["LUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ies, μs_Uniform, ribbon=σs_Uniform, fillalpha=.2, label="Uniform", color=dict_colors["Uniform"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "good_answers_benchmark_N1000.pdf"
savefig(file);





