import time
import math
import concurrent.futures
import numpy as np
import matplotlib.pyplot as plt
from equations import get_eps_audit


def compute_y(x_values, p, delta, proportion_correct, key):
    return key, [get_eps_audit(x, x, math.floor(x *proportion_correct),  delta, p) for x in x_values]


def get_plots():
    final_values = dict()
    mul = 1.5 #1.275 #1.5
    max = 60000 #2000 #60000

    x_values = np.floor((mul)**np.arange(30)).astype(int)
    x_values = np.concatenate([x_values[x_values < max], [max]])

    with concurrent.futures.ProcessPoolExecutor(max_workers=16) as executor:
        start_time = time.time()
        futures = [
            executor.submit(compute_y, x_values, 0.05,  0.0, 1.0, "y11"),
            executor.submit(compute_y, x_values, 0.05, 1e-6, 1.0, "y12"),
            executor.submit(compute_y, x_values, 0.05, 1e-4, 1.0, "y13"),
            executor.submit(compute_y, x_values, 0.05, 1e-2, 1.0, "y14"),
            executor.submit(compute_y, x_values, 0.01,  0.0, 1.0, "y21"),
            executor.submit(compute_y, x_values, 0.01, 1e-6, 1.0, "y22"),
            executor.submit(compute_y, x_values, 0.01, 1e-4, 1.0, "y23"),
            executor.submit(compute_y, x_values, 0.01, 1e-2, 1.0, "y24"),
            executor.submit(compute_y, x_values, 0.05,  0.0, 0.9, "y31"),
            executor.submit(compute_y, x_values, 0.05, 1e-6, 0.9, "y32"),
            executor.submit(compute_y, x_values, 0.05, 1e-4, 0.9, "y33"),
            executor.submit(compute_y, x_values, 0.05, 1e-2, 0.9, "y34"),
            executor.submit(compute_y, x_values, 0.01,  0.0, 0.9, "y41"),
            executor.submit(compute_y, x_values, 0.01, 1e-6, 0.9, "y42"),
            executor.submit(compute_y, x_values, 0.01, 1e-4, 0.9, "y43"),
            executor.submit(compute_y, x_values, 0.01, 1e-2, 0.9, "y44"),
        ]

        for future in concurrent.futures.as_completed(futures):
            k, v = future.result()
            final_values[k] = v
        print(f"Took: {time.time()-start_time}s")

    return final_values, x_values


def plot_to(value_set, x_values, title, fig_name):
    plt.xscale('log')
    plt.plot(x_values, value_set[0], marker='o', label='δ=0')
    plt.plot(x_values, value_set[1], marker='o', label='δ=1e-6')
    plt.plot(x_values, value_set[2], marker='o', label='δ=1e-4')
    plt.plot(x_values, value_set[3], marker='o', label='δ=1e-2')

    plt.xlabel("Number of samples attacked")
    plt.ylabel("Maximum ε lower-bound from audit")
    plt.title(title)
    plt.legend()
    plt.savefig(fig_name, dpi=300, bbox_inches='tight')


def main():
    final_values, x_values = get_plots()

    plot_to(
        [final_values[f"y1{i}"] for i in range(1,5)],
        x_values,
        "Maximum ε audit with p-value=0.05 and 100% MIA accuracy",
        "/dev/shm/plot_05_100.png"
    )
    plot_to(
        [final_values[f"y1{i}"] for i in range(1,5)],
        x_values,
        "Maximum ε audit with p-value=0.01 and 100% MIA accuracy",
        "/dev/shm/plot_01_100.png"
    )
    plot_to(
        [final_values[f"y1{i}"] for i in range(1,5)],
        x_values,
        "Maximum ε audit with p-value=0.05 and 90% MIA accuracy",
        "/dev/shm/plot_05_90.png"
    )
    plot_to(
        [final_values[f"y1{i}"] for i in range(1,5)],
        x_values,
        "Maximum ε audit with p-value=0.01 and 90% MIA accuracy"
        "/dev/shm/plot_01_90.png"
    )


if __name__ == '__main__':
    main()