using PlotlyJS
using PlotlyJS: savefig
using Statistics: mean, std
using DataFrames
include("aux_func.jl")

data = load_pickle("efficiency_per_size_sum.pkl");

#all_data_ene["mni"][s]["50"]["bay"]

data_type = ["mni", "cif"]
model_type = ["bay", "frq"]
experiment_type = ["100", "est", "acc", "wat"]

experiment_100=["50","50","50","50","50"]
experiment_est=["est","est","est","est","est"]
experiment_wat=["wat","wat","wat","wat","wat"]
experiment_acc=["acc","acc","acc","acc","acc"]

model_bayes = ["BCNN","BCNN","BCNN","BCNN","BCNN"]
model_lenet = ["LeNet","LeNet","LeNet","LeNet","LeNet"]

data_mnist = ["MNIST","MNIST","MNIST","MNIST","MNIST"]
data_cifar = ["CIFAR","CIFAR","CIFAR","CIFAR","CIFAR"]

sizes = [1,2,3,4,5]

# MNIST 100 efficency
bayes_100_mnist_eff = []
for s = 1:5
    push!(
        bayes_100_mnist_eff,
        data[model_type[1]][data_type[1]][experiment_type[1]][string(s)],
    )
end
bayes_100_mnist_eff = DataFrame(Efficiency=bayes_100_mnist_eff,Model=model_bayes,Experiment=experiment_100,Dataset=data_mnist,Size=sizes)

freqs_100_mnist_eff = []
for s = 1:5
    push!(
        freqs_100_mnist_eff,
        data[model_type[2]][data_type[1]][experiment_type[1]][string(s)],
    )
end
freqs_100_mnist_eff = DataFrame(Efficiency=freqs_100_mnist_eff,Model=model_lenet,Experiment=experiment_100,Dataset=data_mnist,Size=sizes)


# MNIST est efficency
bayes_est_mnist_eff = []
for s = 1:5
    push!(
        bayes_est_mnist_eff,
        data[model_type[1]][data_type[1]][experiment_type[2]][string(s)],
    )
end
bayes_est_mnist_eff = DataFrame(Efficiency=bayes_est_mnist_eff,Model=model_bayes,Experiment=experiment_est,Dataset=data_mnist,Size=sizes)

freqs_est_mnist_eff = []
for s = 1:5
    push!(
        freqs_est_mnist_eff,
        data[model_type[2]][data_type[1]][experiment_type[2]][string(s)],
    )
end
freqs_est_mnist_eff = DataFrame(Efficiency=freqs_est_mnist_eff,Model=model_lenet,Experiment=experiment_est,Dataset=data_mnist,Size=sizes)


# MNIST wat efficency
bayes_wat_mnist_eff = []
for s = 1:5
    push!(
        bayes_wat_mnist_eff,
        data[model_type[1]][data_type[1]][experiment_type[4]][string(s)],
    )
end
bayes_wat_mnist_eff = DataFrame(Efficiency=bayes_wat_mnist_eff,Model=model_bayes,Experiment=experiment_wat,Dataset=data_mnist,Size=sizes)

freqs_wat_mnist_eff = []
for s = 1:5
    push!(
        freqs_wat_mnist_eff,
        data[model_type[2]][data_type[1]][experiment_type[4]][string(s)],
    )
end
freqs_wat_mnist_eff = DataFrame(Efficiency=freqs_wat_mnist_eff,Model=model_lenet,Experiment=experiment_wat,Dataset=data_mnist,Size=sizes)


# MNIST acc efficency
bayes_acc_mnist_eff = []
for s = 1:5
    push!(
        bayes_acc_mnist_eff,
        data[model_type[1]][data_type[1]][experiment_type[3]][string(s)],
    )
end
bayes_acc_mnist_eff = DataFrame(Efficiency=bayes_acc_mnist_eff,Model=model_bayes,Experiment=experiment_acc,Dataset=data_mnist,Size=sizes)

freqs_acc_mnist_eff = []
for s = 1:5
    push!(
        freqs_acc_mnist_eff,
        data[model_type[2]][data_type[1]][experiment_type[3]][string(s)],
    )
end
freqs_acc_mnist_eff = DataFrame(Efficiency=freqs_acc_mnist_eff,Model=model_lenet,Experiment=experiment_acc,Dataset=data_mnist,Size=sizes)


# CIFAR 100 efficency
bayes_100_cifar_eff = []
for s = 1:5
    push!(
        bayes_100_cifar_eff,
        data[model_type[1]][data_type[2]][experiment_type[1]][string(s)],
    )
end
bayes_100_cifar_eff = DataFrame(Efficiency=bayes_100_cifar_eff,Model=model_bayes,Experiment=experiment_100,Dataset=data_cifar,Size=sizes)
#for i = 1:5
#    t_std = std(bayes_100_cifar_eff)
#    if (bayes_100_cifar_eff[i] > 2 * t_std) || (bayes_100_cifar_eff[i] < 2 * t_std)
#        bayes_100_cifar_eff[i] = mean(bayes_100_cifar_eff)
#    end
#end

freqs_100_cifar_eff = []
for s = 1:5
    push!(
        freqs_100_cifar_eff,
        data[model_type[2]][data_type[2]][experiment_type[1]][string(s)],
    )
end
freqs_100_cifar_eff = DataFrame(Efficiency=freqs_100_cifar_eff,Model=model_lenet,Experiment=experiment_100,Dataset=data_cifar,Size=sizes)

#for i = 1:5
#    t_std = std(freqs_100_cifar_eff)
#    if (freqs_100_cifar_eff[i] > 2 * t_std) || (freqs_100_cifar_eff[i] < 2 * t_std)
#        freqs_100_cifar_eff[i] = mean(freqs_100_cifar_eff)
#    end
#end

# CIFAR est efficency
bayes_est_cifar_eff = []
for s = 1:5
    push!(
        bayes_est_cifar_eff,
        data[model_type[1]][data_type[2]][experiment_type[2]][string(s)],
    )
end
bayes_est_cifar_eff = DataFrame(Efficiency=bayes_est_cifar_eff,Model=model_bayes,Experiment=experiment_est,Dataset=data_cifar,Size=sizes)

freqs_est_cifar_eff = []
for s = 1:5
    push!(
        freqs_est_cifar_eff,
        data[model_type[2]][data_type[2]][experiment_type[2]][string(s)],
    )
end
freqs_est_cifar_eff = DataFrame(Efficiency=freqs_est_cifar_eff,Model=model_lenet,Experiment=experiment_est,Dataset=data_cifar,Size=sizes)


# CIFAR wat efficency
bayes_wat_cifar_eff = []
for s = 1:5
    push!(
        bayes_wat_cifar_eff,
        data[model_type[1]][data_type[2]][experiment_type[4]][string(s)],
    )
end
bayes_wat_cifar_eff = DataFrame(Efficiency=bayes_wat_cifar_eff,Model=model_bayes,Experiment=experiment_wat,Dataset=data_cifar,Size=sizes)


freqs_wat_cifar_eff = []
for s = 1:5
    push!(
        freqs_wat_cifar_eff,
        data[model_type[2]][data_type[2]][experiment_type[4]][string(s)],
    )
end
freqs_wat_cifar_eff = DataFrame(Efficiency=freqs_wat_cifar_eff,Model=model_lenet,Experiment=experiment_wat,Dataset=data_cifar,Size=sizes)


# CIFAR acc efficency
bayes_acc_cifar_eff = []
for s = 1:5
    push!(
        bayes_acc_cifar_eff,
        data[model_type[1]][data_type[2]][experiment_type[3]][string(s)],
    )
end
bayes_acc_cifar_eff = DataFrame(Efficiency=bayes_acc_cifar_eff,Model=model_bayes,Experiment=experiment_acc,Dataset=data_cifar,Size=sizes)

freqs_acc_cifar_eff = []
for s = 1:5
    push!(
        freqs_acc_cifar_eff,
        data[model_type[2]][data_type[2]][experiment_type[3]][string(s)],
    )
end
freqs_acc_cifar_eff = DataFrame(Efficiency=freqs_acc_cifar_eff,Model=model_lenet,Experiment=experiment_acc,Dataset=data_cifar,Size=sizes)

mnist_dataframe = vcat(bayes_100_mnist_eff,freqs_100_mnist_eff,bayes_est_mnist_eff,freqs_est_mnist_eff,bayes_acc_mnist_eff,freqs_acc_mnist_eff,bayes_wat_mnist_eff,freqs_wat_mnist_eff)

cifar_dataframe = vcat(bayes_100_cifar_eff,freqs_100_cifar_eff,bayes_est_cifar_eff,freqs_est_cifar_eff,bayes_acc_cifar_eff,freqs_acc_cifar_eff,bayes_wat_cifar_eff,freqs_wat_cifar_eff)

#avg_rate_bcnn = (-8.266684252643054e-5 * 1000)
#avg_rate_fcnn = (0.00022035677966088333 * 1000)

#=
en_plot = plot(
    [
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_100_mnist_eff),
            y = freqs_100_mnist_eff.Efficiency,
            name = "LeNet 100",
            marker = attr(symbol = 4, color = "rgb(211,120,000)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_est_mnist_eff),
            y = freqs_est_mnist_eff.Efficiency,
            name = "LeNet est",
            marker = attr(symbol = 17, color = "rgb(255,170,017)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_wat_mnist_eff),
            y = freqs_wat_mnist_eff.Efficiency,
            name = "LeNet wat",
            marker = attr(symbol = 2, color = "rgb(255,187,034)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_acc_mnist_eff),
            y = freqs_acc_mnist_eff.Efficiency,
            name = "LeNet acc",
            marker = attr(symbol = 0, color = "rgb(255,204,051)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_100_mnist_eff),
            y = bayes_100_mnist_eff.Efficiency,
            name = "BCNN 100",
            marker = attr(symbol = 4, color = "rgb(055,033,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_est_mnist_eff),
            y = bayes_est_mnist_eff.Efficiency,
            name = "BCNN est",
            marker = attr(symbol = 17, color = "rgb(033,081,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_wat_mnist_eff),
            y = bayes_wat_mnist_eff.Efficiency,
            name = "BCNN wat",
            marker = attr(symbol = 2, color = "rgb(033,115,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_acc_mnist_eff),
            y = bayes_acc_mnist_eff.Efficiency,
            name = "BCNN acc",
            marker = attr(symbol = 0, color = "rgb(151,177,255)", line_width = 1.0),
        ),
    ],
    Layout(
        mode = "overlay",
        xaxis_tickangle = -45,
        yaxis_title_text = "Efficiency",
        xaxis_title_text = "Size";
        xaxis_range = [-1, 5],
        xaxis_type = "category",
    ),
)
savefig(en_plot, "mnist_eff_exp_sum.png")


en_plot = plot(
    [
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_100_cifar_eff),
            y = freqs_100_cifar_eff.Efficiency,
            name = "LeNet 100",
            marker = attr(symbol = 4, color = "rgb(211,120,000)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_est_cifar_eff),
            y = freqs_est_cifar_eff.Efficiency,
            name = "LeNet est",
            marker = attr(symbol = 17, color = "rgb(255,170,017)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_wat_cifar_eff),
            y = freqs_wat_cifar_eff.Efficiency,
            name = "LeNet wat",
            marker = attr(symbol = 2, color = "rgb(255,187,034)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_fcnn .* freqs_acc_cifar_eff),
            y = freqs_acc_cifar_eff.Efficiency,
            name = "LeNet acc",
            marker = attr(symbol = 0, color = "rgb(255,204,051)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_100_cifar_eff),
            y = bayes_100_cifar_eff.Efficiency,
            name = "BCNN 100",
            marker = attr(symbol = 4, color = "rgb(055,033,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_est_cifar_eff),
            y = bayes_est_cifar_eff.Efficiency,
            name = "BCNN est",
            marker = attr(symbol = 17, color = "rgb(033,081,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_wat_cifar_eff),
            y = bayes_wat_cifar_eff.Efficiency,
            name = "BCNN wat",
            marker = attr(symbol = 2, color = "rgb(033,115,240)", line_width = 1.0),
        ),
        scatter(
            x = ["1", "2", "3", "4", "5"],
            #y = (-avg_rate_bcnn .* bayes_acc_cifar_eff),
            y = bayes_acc_cifar_eff.Efficiency,
            name = "BCNN acc",
            marker = attr(symbol = 0, color = "rgb(151,177,255)", line_width = 1.0),
        ),
    ],
    Layout(
        mode = "overlay",
        xaxis_tickangle = -45,
        yaxis_title_text = "Efficiency",
        xaxis_title_text = "Size";
        xaxis_range = [-1, 5],
        xaxis_type = "category",
    ),
)
savefig(en_plot, "cifar_eff_exp_sum.png")
=#




#en_plot = plot(mnist_dataframe, x=:Experiment, y=:Efficiency, boxpoints="all", kind="box")
en_plot = plot(
    mnist_dataframe, x=:Experiment, y=:Efficiency,kind="scatter",mode="markers",color=:Model,
    quartilemethod="exclusive",
    marker=attr(size=:Size, sizeref=maximum(mnist_dataframe.Size) / (10^2), sizemode="area")
    #marker=attr(size=:Size, sizeref=0.1, sizemode="area")
)
savefig(en_plot, "mnist_scatter_size.png")

en_plot = plot(
    mnist_dataframe, x=:Experiment, y=:Efficiency,kind="box", boxpoints="all",color=:Model,
    quartilemethod="exclusive",
    marker=attr(size=:Size, sizeref=maximum(mnist_dataframe.Size) / (10^2), sizemode="area")
    #marker=attr(size=:Size, sizeref=0.1, sizemode="area")
)
savefig(en_plot, "mnist_box_size.png")


en_plot = plot(
    cifar_dataframe, x=:Experiment, y=:Efficiency,kind="scatter",mode="markers",color=:Model,
    quartilemethod="exclusive",
    marker=attr(size=:Size, sizeref=maximum(mnist_dataframe.Size) / (10^2), sizemode="area")
    #marker=attr(size=:Size, sizeref=0.1, sizemode="area")
)
savefig(en_plot, "cifar_scatter_size.png")

en_plot = plot(
    cifar_dataframe, x=:Experiment, y=:Efficiency,kind="box", boxpoints="all",color=:Model,
    quartilemethod="exclusive",
    marker=attr(size=:Size, sizeref=maximum(mnist_dataframe.Size) / (10^2), sizemode="area")
    #marker=attr(size=:Size, sizeref=0.1, sizemode="area")
)
savefig(en_plot, "cifar_box_size.png")