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],
                   "hLUCB" => pal_colors[17],
                   "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]);

Ks = [10, 50, 100, 250, 500, 750, 1000];


# Sparse Benchmark
# values for N=100 runs

μs_cst_EBTCa_15 = [2111, 11161, 22918, 58899, 123858, 180105, 247135];
σs_cst_EBTCa_15 = [523, 2055, 2577, 5890, 23651, 17467, 19287];
μs_cst_EBTCa_1 = [2258, 10377, 21842, 51851, 106180, 152291, 207832];
σs_cst_EBTCa_1 = [1000, 3386, 7601, 11062, 33885, 16512, 35927];
μs_cst_EBTCa_05 = [2859, 13081, 29499, 77433, 142777, 230658, 293749];
σs_cst_EBTCa_05 = [2030, 6082, 15892, 58217, 56066, 129209, 138873];
μs_cst_T3C = [1923, 8422, 16349, 40847, 83475, 128895, 175315];
σs_cst_T3C = [543, 1427, 2016, 3686, 5780, 7036, 10099];
μs_cst_EBTCI = [2139, 10465, 21609, 54842, 113785, 166365, 234551];
σs_cst_EBTCI = [506, 1267, 2815, 5983, 17481, 17167, 25907];
μs_cst_TTUCB = [2148, 9707, 18932, 47790, 95133, 141658, 190180];
σs_cst_TTUCB = [668, 1818, 2994, 4770, 7273, 9752, 10925];
μs_EBTCa_15 = [1876, 7786, 15214, 36518, 70948, 110698, 142363];
σs_EBTCa_15 = [555, 1557, 2907, 6316, 8179, 15171, 14354];
μs_EBTCa_1 = [1749, 7048, 13848, 32870, 61996, 99337, 123398];
σs_EBTCa_1 = [638, 1787, 2833, 10007, 10005, 32092, 19037];
μs_EBTCa_05 = [2130, 9681, 22468, 43507, 81935, 131822, 173797];
σs_EBTCa_05 = [1261, 7110, 16793, 22047, 26780, 58276, 68340];
μs_LUCB = [3196, 12718, 23533, 55044, 105724, 157473, 211244];
σs_LUCB = [673, 1391, 2288, 4808, 8907, 13948, 16192];
μs_Uniform = [3493, 22429, 50323, 127715, 279476, 437637, 584650];
σs_Uniform = [1138, 6678, 12538, 36612, 69908, 101641, 137851];

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

# Plot with eps_0 = eps_1, and using IDS
plot(Ks, μs_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.3, label="EB-TC" * L"_{\varepsilon}", color=dict_colors["EB-TCa-e1"], markershape=:auto, linestyle=:auto, legend=:topleft, xaxis=:log, yaxis=:log, xlabel=L"K", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ks, μs_cst_T3C, ribbon=σs_cst_T3C, fillalpha=.3, label=L"\varepsilon" * "-T3C", color=dict_colors["T3C"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_EBTCI, ribbon=σs_cst_EBTCI, fillalpha=.3, label=L"\varepsilon" * "-EB-TCI", color=dict_colors["EB-TCI"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_TTUCB, ribbon=σs_cst_TTUCB, fillalpha=.3, label=L"\varepsilon" * "-TTUCB", color=dict_colors["TTUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_LUCB, ribbon=σs_LUCB, fillalpha=.3, label="LUCB", color=dict_colors["LUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_Uniform, ribbon=σs_Uniform, fillalpha=.3, label="Uniform", color=dict_colors["Uniform"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "sparse_benchmark_N100.pdf"
savefig(file);


# Alpha 0.3 Benchmark
# values for N=100 runs

μs_cst_EBTCa_15 = [357, 2047, 4361, 11615, 25302, 38878, 54179];
σs_cst_EBTCa_15 = [110, 357, 588, 1287, 2779, 4461, 6139];
μs_cst_EBTCa_1 = [357, 1927, 3910, 9913, 20357, 31043, 40882];
σs_cst_EBTCa_1 = [106, 306, 489, 1172, 2361, 2033, 3067];
μs_cst_EBTCa_05 = [389, 2068, 4306, 11320, 23127, 34597, 47727];
σs_cst_EBTCa_05 = [137, 372, 753, 1353, 1400, 1840, 3807];
μs_cst_T3C = [356, 1674, 3269, 7669, 15021, 22722, 30199];
σs_cst_T3C = [104, 269, 425, 785, 1085, 1987, 2102];
μs_cst_EBTCI = [386, 2274, 4753, 12579, 26663, 40818, 55549];
σs_cst_EBTCI = [110, 380, 618, 1163, 2182, 3073, 3448];
μs_cst_TTUCB = [360, 2240, 5310, 15456, 35313, 57867, 76545];
σs_cst_TTUCB = [132, 640, 1279, 3761, 7188, 12555, 13358];
μs_EBTCa_15 = [341, 1510, 3085, 7748, 15674, 23488, 31826];
σs_EBTCa_15 = [114, 325, 619, 1229, 2398, 3714, 4803];
μs_EBTCa_1 = [316, 1484, 2861, 6572, 12718, 18914, 25155];
σs_EBTCa_1 = [113, 416, 637, 946, 1457, 2197, 2496];
μs_EBTCa_05 = [323, 1509, 3132, 7334, 14736, 21748, 28771];
σs_EBTCa_05 = [123, 479, 949, 1351, 2070, 4090, 3729];
μs_LUCB = [587, 2563, 5003, 11937, 23478, 35126, 46423];
σs_LUCB = [128, 296, 528, 779, 1254, 2486, 2616];
μs_Uniform = [814, 8275, 24651, 86687, 234582, 418352, 611282];
σs_Uniform = [325, 3472, 10719, 40578, 97752, 173518, 213363];

plot(Ks, μs_cst_EBTCa_15, ribbon=σs_cst_EBTCa_15, fillalpha=.3, label="c-EB-TC" * L"_{0.15}", color=dict_colors_Bench_ATT["c-EB-TCa-e15"], markershape=:auto, linestyle=:auto, legend=:topleft, xaxis=:log, yaxis=:log, xlabel="number of arms " * L"K", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ks, μs_cst_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.3, label="c-EB-TC" * L"_{0.1}", color=dict_colors_Bench_ATT["c-EB-TCa-e1"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_EBTCa_05, ribbon=σs_cst_EBTCa_05, fillalpha=.3, label="c-EB-TC" * L"_{0.05}", color=dict_colors_Bench_ATT["c-EB-TCa-e05"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_EBTCa_15, ribbon=σs_EBTCa_15, fillalpha=.3, label="EB-TC" * L"_{0.15}", color=dict_colors_Bench_ATT["EB-TCa-e15"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_EBTCa_1, ribbon=σs_EBTCa_1, fillalpha=.3, label="EB-TC" * L"_{0.1}", color=dict_colors_Bench_ATT["EB-TCa-e1"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_EBTCa_05, ribbon=σs_EBTCa_05, fillalpha=.3, label="EB-TC" * L"_{0.05}", color=dict_colors_Bench_ATT["EB-TCa-e05"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "alpha3_bench_ATT_N100.pdf"
savefig(file);


# Plot with eps_0 = eps_1, and using IDS
plot(Ks, μs_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.3, label="EB-TC" * L"_{\varepsilon}", color=dict_colors["EB-TCa-e1"], markershape=:auto, linestyle=:auto, legend=:topleft, xaxis=:log, yaxis=:log, xlabel=L"K", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ks, μs_cst_T3C, ribbon=σs_cst_T3C, fillalpha=.3, label=L"\varepsilon" * "-T3C", color=dict_colors["T3C"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_EBTCI, ribbon=σs_cst_EBTCI, fillalpha=.3, label=L"\varepsilon" * "-EB-TCI", color=dict_colors["EB-TCI"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_TTUCB, ribbon=σs_cst_TTUCB, fillalpha=.3, label=L"\varepsilon" * "-TTUCB", color=dict_colors["TTUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_LUCB, ribbon=σs_LUCB, fillalpha=.3, label="LUCB", color=dict_colors["LUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_Uniform, ribbon=σs_Uniform, fillalpha=.3, label="Uniform", color=dict_colors["Uniform"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "alpha3_benchmark_N100.pdf"
savefig(file);

# Alpha 0.6 Benchmark
# values for N=100 runs

Ks = [10, 50, 100, 250, 500];


μs_cst_EBTCa_15 = [734, 5083, 11068, 32329, 71655];
σs_cst_EBTCa_15 = [324, 1164, 1849, 5617, 11007];
μs_cst_EBTCa_1 = [734, 4216, 9135, 22743, 46493];
σs_cst_EBTCa_1 = [269, 961, 1491, 2495, 4583];
μs_cst_EBTCa_05 = [894, 5770, 12358, 33405, 67632];
σs_cst_EBTCa_05 = [367, 1725, 2613, 4643, 7653];
μs_cst_T3C = [675, 4081, 7974, 19865, 39640];
σs_cst_T3C = [224, 988, 1394, 2779, 4487];
μs_cst_EBTCI = [755, 5086, 10369, 28115, 60627];
σs_cst_EBTCI = [254, 1108, 1603, 3510, 6329];
μs_cst_TTUCB = [799, 6348, 14645, 41265, 81312];
σs_cst_TTUCB = [292, 2133, 3916, 8199, 12438];
μs_EBTCa_15 = [691, 3873, 8000, 22271, 44375];
σs_EBTCa_15 = [285, 1159, 1926, 4398, 6547];
μs_EBTCa_1 = [674, 3447, 6762, 16052, 29442];
σs_EBTCa_1 = [303, 927, 1493, 2489, 4019];
μs_EBTCa_05 = [729, 4652, 9001, 23621, 46061];
σs_EBTCa_05 = [341, 1346, 2529, 5524, 7561];
μs_LUCB = [1228, 6518, 12069, 28862, 56708];
σs_LUCB = [346, 1232, 1623, 4034, 5927];
μs_Uniform = [2212, 36769, 119146, 363314, 846632];
σs_Uniform = [1004, 17134, 50511, 137599, 182074];

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


# Plot with eps_0 = eps_1, and using IDS
plot(Ks, μs_EBTCa_1, ribbon=σs_cst_EBTCa_1, fillalpha=.3, label="EB-TC" * L"_{\varepsilon}", color=dict_colors["EB-TCa-e1"], markershape=:auto, linestyle=:auto, legend=:topleft, xaxis=:log, yaxis=:log, xlabel=L"K", ylabel=L"E_{\nu}[\tau_{\varepsilon,\delta}]", linewidth=2);
plot!(Ks, μs_cst_T3C, ribbon=σs_cst_T3C, fillalpha=.3, label=L"\varepsilon" * "-T3C", color=dict_colors["T3C"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_EBTCI, ribbon=σs_cst_EBTCI, fillalpha=.3, label=L"\varepsilon" * "-EB-TCI", color=dict_colors["EB-TCI"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_cst_TTUCB, ribbon=σs_cst_TTUCB, fillalpha=.3, label=L"\varepsilon" * "-TTUCB", color=dict_colors["TTUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_LUCB, ribbon=σs_LUCB, fillalpha=.3, label="LUCB", color=dict_colors["LUCB"], markershape=:auto, linestyle=:auto, linewidth=2);
plot!(Ks, μs_Uniform, ribbon=σs_Uniform, fillalpha=.3, label="Uniform", color=dict_colors["Uniform"], markershape=:auto, linestyle=:auto, linewidth=2);
file = "alpha6_benchmark_N100.pdf"
savefig(file);


