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

data = load_pickle("efficiency_per_size.pkl");

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

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

# 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


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

# 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

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

# 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

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

# 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

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


# 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
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
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

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


# 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


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


# 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


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

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),
            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),
            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),
            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),
            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),
            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),
            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),
            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),
            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 = "ΔEfficency",
        xaxis_title_text = "Size";
        xaxis_range = [-1, 5],
        xaxis_type = "category",
    ),
)
savefig(en_plot, "mnist_eff_exp.png")




en_plot = plot(
    [
        scatter(
            x = ["1", "2", "3", "4", "5"],
            y = (-avg_rate_fcnn .* freqs_100_cifar_eff),
            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),
            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),
            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),
            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),
            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),
            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),
            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),
            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 = "ΔEfficency",
        xaxis_title_text = "Size";
        xaxis_range = [-1, 5],
        xaxis_type = "category",
    ),
)
savefig(en_plot, "cifar_eff_exp.png")