# Libraries
using PGFPlotsX, LaTeXStrings;
using MessyTimeSeries: no_combinations;
using Statistics;

# Custom functions for AJK and BJK probabilities

# - AJK
compute_ajk_cases(nT::Int64, d::Int64, k::Int64) = no_combinations(d, k)*no_combinations(nT-d, d-k);
compute_ajk_probability(nT::Int64, d::Int64, k::Int64) = compute_ajk_cases(nT, d, k)/no_combinations(nT, d);

# - BJK
compute_bjk_ith_cases(i::Int64, no_subsamples::Int64, c::Int64, k::Int64) = (no_subsamples-i >= c-k) + (i-1 >= c-k);
compute_bjk_probability(no_subsamples::Int64, c::Int64, k::Int64) = mapreduce(i -> compute_bjk_ith_cases(i, no_subsamples, c, k), +, 1:no_subsamples)/(no_subsamples^2);

#=
Custom functions to compute the expected value of the Jaccard indices
=#

# - Main methods
expected_jaccard_ajk(nT::Int64, d::Int64) = mapreduce(k -> k/(2*d-k)*compute_ajk_probability(nT, d, k), +, 1:d-1);
expected_jaccard_bjk(T::Int64, c::Int64)  = mapreduce(k -> k/(2*c-k)*compute_bjk_probability(T-c+1, c, k), +, 1:c-1);

# - Promotion rules
expected_jaccard_ajk(nT::Real, d::Real) = expected_jaccard_ajk(Int64(nT), Int64(floor(d)));
expected_jaccard_bjk(T::Real, c::Real) = expected_jaccard_bjk(Int64(T), Int64(floor(c)));

#=
Surface chart
=#

# Plot: grid
push!(PGFPlotsX.CUSTOM_PREAMBLE, 
        raw"\usepgfplotslibrary{colorbrewer}",
        raw"\usepgfplotslibrary{colormaps}",
        raw"\usepgfplotslibrary{patchplots}",
        raw"\pgfplotsset
            {   
                tick label style = {font = {\fontsize{12 pt}{12 pt}\selectfont}},
                label style = {font = {\fontsize{12 pt}{12 pt}\selectfont}},
                legend style = {font = {\fontsize{12 pt}{12 pt}\selectfont}},
                title style = {font = {\fontsize{12 pt}{12 pt}\selectfont}},
            }"
)

# GroupPlot axs
gpr = Array{Any,1}(undef, 5);

# Loop over c
for (i, c) in enumerate([4, 8, 12, 16, 20])

    # n,T grids
    n_grid = range(5, stop=50, length=10);
    T_grid = range(50, stop=500, length=10);

    # Compute z_grid
    title_string = L"c=%$(c)";
    z_grid = [expected_jaccard_bjk(T, c)-expected_jaccard_ajk(n*T, n*c) for n in n_grid, T in T_grid];

    # Plots
    p1 = @pgf Plot3({surf, point_meta="explicit"}, Coordinates(n_grid, T_grid, z_grid, meta=z_grid));

    gpr[i] = @pgf Axis(
    {
        "axis on top"="false",
        height = "10cm",
        width = "10cm",
        grid = "major",
        grid_style = "{darkgray, very thin}",
        tick_style = "{darkgray, very thin}",
        "colormap"="{turbo}{rgb=(0.18995,0.07176,0.23217), rgb=(0.19483,0.08339,0.26149), rgb=(0.19956,0.09498,0.29024), rgb=(0.20415,0.10652,0.31844), rgb=(0.20860,0.11802,0.34607), rgb=(0.21291,0.12947,0.37314), rgb=(0.21708,0.14087,0.39964), rgb=(0.22111,0.15223,0.42558), rgb=(0.22500,0.16354,0.45096), rgb=(0.22875,0.17481,0.47578), rgb=(0.23236,0.18603,0.50004), rgb=(0.23582,0.19720,0.52373), rgb=(0.23915,0.20833,0.54686), rgb=(0.24234,0.21941,0.56942), rgb=(0.24539,0.23044,0.59142), rgb=(0.24830,0.24143,0.61286), rgb=(0.25107,0.25237,0.63374), rgb=(0.25369,0.26327,0.65406), rgb=(0.25618,0.27412,0.67381), rgb=(0.25853,0.28492,0.69300), rgb=(0.26074,0.29568,0.71162), rgb=(0.26280,0.30639,0.72968), rgb=(0.26473,0.31706,0.74718), rgb=(0.26652,0.32768,0.76412), rgb=(0.26816,0.33825,0.78050), rgb=(0.26967,0.34878,0.79631), rgb=(0.27103,0.35926,0.81156), rgb=(0.27226,0.36970,0.82624), rgb=(0.27334,0.38008,0.84037), rgb=(0.27429,0.39043,0.85393), rgb=(0.27509,0.40072,0.86692), rgb=(0.27576,0.41097,0.87936), rgb=(0.27628,0.42118,0.89123), rgb=(0.27667,0.43134,0.90254), rgb=(0.27691,0.44145,0.91328), rgb=(0.27701,0.45152,0.92347), rgb=(0.27698,0.46153,0.93309), rgb=(0.27680,0.47151,0.94214), rgb=(0.27648,0.48144,0.95064), rgb=(0.27603,0.49132,0.95857), rgb=(0.27543,0.50115,0.96594), rgb=(0.27469,0.51094,0.97275), rgb=(0.27381,0.52069,0.97899), rgb=(0.27273,0.53040,0.98461), rgb=(0.27106,0.54015,0.98930), rgb=(0.26878,0.54995,0.99303), rgb=(0.26592,0.55979,0.99583), rgb=(0.26252,0.56967,0.99773), rgb=(0.25862,0.57958,0.99876), rgb=(0.25425,0.58950,0.99896), rgb=(0.24946,0.59943,0.99835), rgb=(0.24427,0.60937,0.99697), rgb=(0.23874,0.61931,0.99485), rgb=(0.23288,0.62923,0.99202), rgb=(0.22676,0.63913,0.98851), rgb=(0.22039,0.64901,0.98436), rgb=(0.21382,0.65886,0.97959), rgb=(0.20708,0.66866,0.97423), rgb=(0.20021,0.67842,0.96833), rgb=(0.19326,0.68812,0.96190), rgb=(0.18625,0.69775,0.95498), rgb=(0.17923,0.70732,0.94761), rgb=(0.17223,0.71680,0.93981), rgb=(0.16529,0.72620,0.93161), rgb=(0.15844,0.73551,0.92305), rgb=(0.15173,0.74472,0.91416), rgb=(0.14519,0.75381,0.90496), rgb=(0.13886,0.76279,0.89550), rgb=(0.13278,0.77165,0.88580), rgb=(0.12698,0.78037,0.87590), rgb=(0.12151,0.78896,0.86581), rgb=(0.11639,0.79740,0.85559), rgb=(0.11167,0.80569,0.84525), rgb=(0.10738,0.81381,0.83484), rgb=(0.10357,0.82177,0.82437), rgb=(0.10026,0.82955,0.81389), rgb=(0.09750,0.83714,0.80342), rgb=(0.09532,0.84455,0.79299), rgb=(0.09377,0.85175,0.78264), rgb=(0.09287,0.85875,0.77240), rgb=(0.09267,0.86554,0.76230), rgb=(0.09320,0.87211,0.75237), rgb=(0.09451,0.87844,0.74265), rgb=(0.09662,0.88454,0.73316), rgb=(0.09958,0.89040,0.72393), rgb=(0.10342,0.89600,0.71500), rgb=(0.10815,0.90142,0.70599), rgb=(0.11374,0.90673,0.69651), rgb=(0.12014,0.91193,0.68660), rgb=(0.12733,0.91701,0.67627), rgb=(0.13526,0.92197,0.66556), rgb=(0.14391,0.92680,0.65448), rgb=(0.15323,0.93151,0.64308), rgb=(0.16319,0.93609,0.63137), rgb=(0.17377,0.94053,0.61938), rgb=(0.18491,0.94484,0.60713), rgb=(0.19659,0.94901,0.59466), rgb=(0.20877,0.95304,0.58199), rgb=(0.22142,0.95692,0.56914), rgb=(0.23449,0.96065,0.55614), rgb=(0.24797,0.96423,0.54303), rgb=(0.26180,0.96765,0.52981), rgb=(0.27597,0.97092,0.51653), rgb=(0.29042,0.97403,0.50321), rgb=(0.30513,0.97697,0.48987), rgb=(0.32006,0.97974,0.47654), rgb=(0.33517,0.98234,0.46325), rgb=(0.35043,0.98477,0.45002), rgb=(0.36581,0.98702,0.43688), rgb=(0.38127,0.98909,0.42386), rgb=(0.39678,0.99098,0.41098), rgb=(0.41229,0.99268,0.39826), rgb=(0.42778,0.99419,0.38575), rgb=(0.44321,0.99551,0.37345), rgb=(0.45854,0.99663,0.36140), rgb=(0.47375,0.99755,0.34963), rgb=(0.48879,0.99828,0.33816), rgb=(0.50362,0.99879,0.32701), rgb=(0.51822,0.99910,0.31622), rgb=(0.53255,0.99919,0.30581), rgb=(0.54658,0.99907,0.29581), rgb=(0.56026,0.99873,0.28623), rgb=(0.57357,0.99817,0.27712), rgb=(0.58646,0.99739,0.26849), rgb=(0.59891,0.99638,0.26038), rgb=(0.61088,0.99514,0.25280), rgb=(0.62233,0.99366,0.24579), rgb=(0.63323,0.99195,0.23937), rgb=(0.64362,0.98999,0.23356), rgb=(0.65394,0.98775,0.22835), rgb=(0.66428,0.98524,0.22370), rgb=(0.67462,0.98246,0.21960), rgb=(0.68494,0.97941,0.21602), rgb=(0.69525,0.97610,0.21294), rgb=(0.70553,0.97255,0.21032), rgb=(0.71577,0.96875,0.20815), rgb=(0.72596,0.96470,0.20640), rgb=(0.73610,0.96043,0.20504), rgb=(0.74617,0.95593,0.20406), rgb=(0.75617,0.95121,0.20343), rgb=(0.76608,0.94627,0.20311), rgb=(0.77591,0.94113,0.20310), rgb=(0.78563,0.93579,0.20336), rgb=(0.79524,0.93025,0.20386), rgb=(0.80473,0.92452,0.20459), rgb=(0.81410,0.91861,0.20552), rgb=(0.82333,0.91253,0.20663), rgb=(0.83241,0.90627,0.20788), rgb=(0.84133,0.89986,0.20926), rgb=(0.85010,0.89328,0.21074), rgb=(0.85868,0.88655,0.21230), rgb=(0.86709,0.87968,0.21391), rgb=(0.87530,0.87267,0.21555), rgb=(0.88331,0.86553,0.21719), rgb=(0.89112,0.85826,0.21880), rgb=(0.89870,0.85087,0.22038), rgb=(0.90605,0.84337,0.22188), rgb=(0.91317,0.83576,0.22328), rgb=(0.92004,0.82806,0.22456), rgb=(0.92666,0.82025,0.22570), rgb=(0.93301,0.81236,0.22667), rgb=(0.93909,0.80439,0.22744), rgb=(0.94489,0.79634,0.22800), rgb=(0.95039,0.78823,0.22831), rgb=(0.95560,0.78005,0.22836), rgb=(0.96049,0.77181,0.22811), rgb=(0.96507,0.76352,0.22754), rgb=(0.96931,0.75519,0.22663), rgb=(0.97323,0.74682,0.22536), rgb=(0.97679,0.73842,0.22369), rgb=(0.98000,0.73000,0.22161), rgb=(0.98289,0.72140,0.21918), rgb=(0.98549,0.71250,0.21650), rgb=(0.98781,0.70330,0.21358), rgb=(0.98986,0.69382,0.21043), rgb=(0.99163,0.68408,0.20706), rgb=(0.99314,0.67408,0.20348), rgb=(0.99438,0.66386,0.19971), rgb=(0.99535,0.65341,0.19577), rgb=(0.99607,0.64277,0.19165), rgb=(0.99654,0.63193,0.18738), rgb=(0.99675,0.62093,0.18297), rgb=(0.99672,0.60977,0.17842), rgb=(0.99644,0.59846,0.17376), rgb=(0.99593,0.58703,0.16899), rgb=(0.99517,0.57549,0.16412), rgb=(0.99419,0.56386,0.15918), rgb=(0.99297,0.55214,0.15417), rgb=(0.99153,0.54036,0.14910), rgb=(0.98987,0.52854,0.14398), rgb=(0.98799,0.51667,0.13883), rgb=(0.98590,0.50479,0.13367), rgb=(0.98360,0.49291,0.12849), rgb=(0.98108,0.48104,0.12332), rgb=(0.97837,0.46920,0.11817), rgb=(0.97545,0.45740,0.11305), rgb=(0.97234,0.44565,0.10797), rgb=(0.96904,0.43399,0.10294), rgb=(0.96555,0.42241,0.09798), rgb=(0.96187,0.41093,0.09310), rgb=(0.95801,0.39958,0.08831), rgb=(0.95398,0.38836,0.08362), rgb=(0.94977,0.37729,0.07905), rgb=(0.94538,0.36638,0.07461), rgb=(0.94084,0.35566,0.07031), rgb=(0.93612,0.34513,0.06616), rgb=(0.93125,0.33482,0.06218), rgb=(0.92623,0.32473,0.05837), rgb=(0.92105,0.31489,0.05475), rgb=(0.91572,0.30530,0.05134), rgb=(0.91024,0.29599,0.04814), rgb=(0.90463,0.28696,0.04516), rgb=(0.89888,0.27824,0.04243), rgb=(0.89298,0.26981,0.03993), rgb=(0.88691,0.26152,0.03753), rgb=(0.88066,0.25334,0.03521), rgb=(0.87422,0.24526,0.03297), rgb=(0.86760,0.23730,0.03082), rgb=(0.86079,0.22945,0.02875), rgb=(0.85380,0.22170,0.02677), rgb=(0.84662,0.21407,0.02487), rgb=(0.83926,0.20654,0.02305), rgb=(0.83172,0.19912,0.02131), rgb=(0.82399,0.19182,0.01966), rgb=(0.81608,0.18462,0.01809), rgb=(0.80799,0.17753,0.01660), rgb=(0.79971,0.17055,0.01520), rgb=(0.79125,0.16368,0.01387), rgb=(0.78260,0.15693,0.01264), rgb=(0.77377,0.15028,0.01148), rgb=(0.76476,0.14374,0.01041), rgb=(0.75556,0.13731,0.00942), rgb=(0.74617,0.13098,0.00851), rgb=(0.73661,0.12477,0.00769), rgb=(0.72686,0.11867,0.00695), rgb=(0.71692,0.11268,0.00629), rgb=(0.70680,0.10680,0.00571), rgb=(0.69650,0.10102,0.00522), rgb=(0.68602,0.09536,0.00481), rgb=(0.67535,0.08980,0.00449), rgb=(0.66449,0.08436,0.00424), rgb=(0.65345,0.07902,0.00408), rgb=(0.64223,0.07380,0.00401), rgb=(0.63082,0.06868,0.00401), rgb=(0.61923,0.06367,0.00410), rgb=(0.60746,0.05878,0.00427), rgb=(0.59550,0.05399,0.00453), rgb=(0.58336,0.04931,0.00486), rgb=(0.57103,0.04474,0.00529), rgb=(0.55852,0.04028,0.00579), rgb=(0.54583,0.03593,0.00638), rgb=(0.53295,0.03169,0.00705), rgb=(0.51989,0.02756,0.00780), rgb=(0.50664,0.02354,0.00863), rgb=(0.49321,0.01963,0.00955), rgb=(0.47960,0.01583,0.01055)}",
        "colorbar"=ifelse(i==5, "true", "false"),
        "colorbar style"=raw"
        {
            ytick={0,0.02,...,0.12},
            yticklabel style=
            {
                text width=2.5em,
                align=right,
                /pgf/number format/.cd,
                    fixed,
                    fixed zerofill
            }
        }",
        "colorbar/width" = "0.8cm",
        title = title_string,
        view = (45, 15),
        point_meta_min = "0.0",
        point_meta_max = "0.12",
        zmin = "0",
        zmax = "0.15",
        "scaled z ticks"="false",
        "zticklabel"=raw"\pgfkeys{/pgf/number format/.cd,fixed,precision=2,zerofill}\pgfmathprintnumber{\tick}",
        xlabel = L"n",
        ylabel = L"T",
        zlabel = ifelse(i==1, "Expected Jaccard Index (difference)", ""),
    },
        p1,
    );
end

fig = @pgf TikzPicture(GroupPlot(
    { group_style = {group_size="5 by 1", vertical_sep="50pt", horizontal_sep="50pt"},
      height="200pt",
      width="225pt"
    },
    gpr...,)
)

pgfsave("./img/appendix/simulations_surface.pdf", fig);