import json
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import copy



def main():
    data = {}
    with open("results/cifar10_fgsm.json", "r") as infile:
        data = json.load(infile)

    attack = data["attack"]
    epsilons = data["epsilons"]
    filters = data["filters"]
    dataset = data["dataset"]
    strength_count = len(filters[list(filters.keys())[0]][0])

    strengths = np.arange(strength_count)
    epsilons = np.array(epsilons)

    # Assume constant step of strength and epsilon
    dstrength = strengths[1] - strengths[0]
    depsilon = epsilons[1] - epsilons[0]

    # Make a grid from strengths and epsilons
    strengths, epsilons = np.meshgrid(strengths,epsilons)
    #strengths, epsilons = strengths.ravel(), epsilons.ravel()

    colors = ('blue', 'orange', 'red', 'purple', 'green', 'yellow', 'brown')
    
    z = np.zeros_like(strengths)
    best_performance = np.zeros_like(strengths)
    bottoms = np.zeros_like(strengths)

    for i, filter in enumerate(filters):
        performance = np.array(filters[filter])
        best_performance = np.where(performance > best_performance, i, best_performance)
        z = np.where(performance > z, performance, z)
        print(best_performance)

    for i, filter in enumerate(filters): 
        fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
        tops = np.where(best_performance == i, best_performance, 0)
        print(strengths.shape, epsilons.shape, bottoms.shape, dstrength.shape, depsilon.shape, tops.shape)
        ax.bar3d(strengths, epsilons, bottoms, dstrength, depsilon, tops, color=colors[i])
        plt.show()

    ax.view_init(90, -90, 0)

    plt.legend()
    plt.title(f"{filter} Performance")
    plt.xlabel(f"{attack} Attack Strength ($\\epsilon$)")
    plt.ylabel(f"{dataset} Classification Accuracy")
    plt.show()

if __name__ == "__main__":
    main()