# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.5.3
#     language: julia
#     name: julia-1.5
# ---

include("table.jl")

# # ENV

ENV["COLUMNS"] = 200

# # Load the results

results = load_results()

df = results2df(results)

sort(df[(df.model .== "EQ") .& (df.train_mat .== "maxCOV"),:], [:train_d, :train_gtype, :train_mat])

sort(df[(df.model .== "EQ") .& (df.train_mat .== "COR"),:], [:train_d, :train_gtype, :train_mat])

sort(df[(df.model .== "CNN") .& (df.train_mat .== "maxCOV"),:], [:train_d, :train_gtype, :train_mat])

# # Test dataframe selection
#
# And exploring the data

df[1:4,["model","train_d"]]

# +
# sum(df.raw)

# I don't need train_k, because all of them are 1s
df = df[!, Not(All(:train_k, :train_noise))]

unique(df.train_mat)
unique(df.model)
# -

# # Pre-defined tables

table_cmp_baseline(df)

table_ER24(df)

table_raw(df)

table_transfer_k(df)

table_transfer_noise(df)

table_transfer_gtype(df)

table_ensemble(df)

# # Developing the tables

# ## Table 1: compare EQ/CNN/FC with the literature
#
# - [ ] TODO literature in baseline.py
# - [ ] TODO I want to probably generate the final layout to use in the paper.
#
# I need:
# - add sep lines
# - break into "two columns"
#
# Possibly through:
# - I need to generate nice output here
# - generate latex code to copy&paste into paper

# +
function table_cmp_literature(df)
    # comparison of EQ with FC and CNN (and random) baselines
    # only show d=10,15,20, SF and ER graphs, k=1
    sel0 = (df.train_mat .== "maxCOV")
    sel1 = in.(df.model, Ref(["FC", "EQ", "CNN"]))
    sel2 = in.(df.train_d, Ref([10, 20, 50, 100]))
    selector = (sel0 .& sel1 .& sel2
                .& (df.train_gtype .== df.test_gtype)
                .& in.(df.train_gtype, Ref(["SF", "ER"]))
                .& (df.test_noise .== "Gaussian")
                .& (df.train_mec .== "Linear")
                .& (df.test_mec .== "Linear")
#                 .& (df.train_mat .== "COR")
#                 .& (df.test_mat .== "COR")
                .& (df.test_k .== 1))

    sort(df[selector,
            All(:model, :train_d, :prec, :recall, :shd)],
         [:train_d, :model])
end

df1 = table_cmp_literature(df)
# -

# add ER and SF and report average results
#
# for every 2 column, add together and compute average
newdf = copy(df1)
for i in 1:Int(size(newdf,1) / 2)
    idx1 = i*2-1
    idx2 = i*2
    d1 = convert(Array, newdf[idx1,[:prec, :recall, :shd]])
    d2 = convert(Array, newdf[idx2,[:prec, :recall, :shd]])
    newdata = (d1 .+ d2) ./ 2
    newdf[idx1,[:prec, :recall, :shd]] .= newdata
end
newdf = newdf[1:2:size(newdf,1),:]

display(MIME("text/latex"), newdf)

using Latexify

latexify(newdf; env=:table, latex=false)

# ## Table 2: tranfer experiments
#
# - [ ] different noise model
# - [ ] different causal strength `k`
# - [ ] different types of graph
# - [ ] transfer to very large graphs (d=200,300,400), and compare to GES
# - [ ] ensemble d=10,15,20, and test on d=10,15,20,30,50,80

function table_transfer_noise(df)
    # - train on Gaussian, test on Exp, Gumbel
    selector = ((df.model .== "EQ")
        .& (df.train_mat .== "maxCOV")
                .& in.(df.train_gtype, Ref(["SF", "ER"]))
                .& (df.train_gtype .== df.test_gtype)
                .& (df.test_k .== 1)
                .& (df.train_mec .== "Linear")
                .& (df.test_mec .== "Linear"))
    sort(df[selector,
            Not(All(
                    # CAUTION I need to show this to tell that the training and
                    # testing noise are different :train_noise,
                    :test_d, :test_k,
                    :train_mec, :test_mec,
                    :test_gtype,
                    :train_mat, :test_mat))],
         [:train_gtype, :model, :train_d, :test_noise])
end
table_transfer_noise(df)

function table_transfer_gtype(df)
    # - train on ER, test on SF
    selector = ((df.model .== "EQ")
                # TODO ER2 ER4 SF2 SF4
                .& in.(df.train_gtype, Ref(["ER", "SF"]))
        # TODO this is not in the new experiments
                .& (df.train_gtype .!= df.test_gtype)
#                 .& (df.test_noise .== "Gaussian")
                .& (df.test_k .== 1)
                .& (df.train_mec .== "Linear")
                .& (df.test_mec .== "Linear")
                .& (df.train_mat .== "maxCOV"))
    sort(df[selector,
            Not(All(
                    :test_d, :test_k,
                    :train_mec, :test_mec, :test_noise,
                    :train_mat, :test_mat))],
         [:train_gtype, :model, :test_gtype, :train_d])
end
table_transfer_gtype(df)

# +

function table_transfer_k(df)
    # - train on k=1, test on k=1,2,4
    selector = ((df.model .== "EQ")
                .& (df.test_noise .== "Gaussian")
                .& (df.train_mec .== "Linear")
                .& (df.test_mec .== "Linear")
                .& (df.train_mat .== "maxCOV"))
    sort(df[selector,
            Not(All(:train_k, :train_noise,
                    :test_d,
                    :train_mec, :test_mec, :test_noise,
                    :train_mat, :test_mat))],
         [:train_gtype, :model, :test_gtype, :train_d, :test_k])
end
table_transfer_k(df)
# -

# ## Ensemble Table
#
# Training specs:
# - d=10,20
# - k=1,4,10,20
# - gtype=SF,ER
#
# Testing specs:
# - d=10,20
# - Sachs-2005

function table_ensemble(df)
    selector = (in.(df.model, Ref(["deep-EQ-ER-ensemble",
                                   "deep-EQ-SF-ensemble"]))
                .& (df.test_k .== 1)
                .& (df.test_mec .== "Linear")
                .& (df.test_mat .== "COR"))
    sort(df[selector,
            Not(All(:train_k, :train_noise,
                    :test_k,
                    r"train_",
                    :train_mec, :test_mec, :test_noise,
                    :train_mat, :test_mat))],
         [:model, :test_gtype, :test_d])
end
table_ensemble(df)

# # Testing displaysize

displaysize()

# show all of df
show(df[1:3,:], allcols=true)

IOContext(stdout, :displaysize=>(30,100)) |> displaysize

StarcoderdataJupyterScriptsDedupFiltered