import matplotlib.pyplot as plt
import pickle
import sys

if __name__ == "__main__":
    k = int(sys.argv[1])
    with open(f'./results/results_ourszo_k_{k}.pkl', 'rb') as file:
        results_ourszo = pickle.load(file)
    with open(f'./results/results_theirszo_k_{k}.pkl', 'rb') as file:
        results_theirszo = pickle.load(file)
    plt.figure()
    plt.plot(results_ourszo['izo'], results_ourszo['losses'], color='green', label='HZO-HT', marker='x', markersize=5, alpha=0.5)
    plt.plot(results_theirszo['izo'], results_theirszo['losses'], color='red', label='ZOHT', marker='>', markersize=5, alpha=0.5)
    plt.legend()
    plt.title('IZO')
    plt.savefig(f'./results/izo_comp_k_{k}.png')

    plt.figure()
    plt.plot(results_ourszo['nht'], results_ourszo['losses'], color='green', label='HZO-HT', marker='x', markersize=5, alpha=0.5)
    plt.plot(results_theirszo['nht'], results_theirszo['losses'], color='red', label='ZOHT', marker='>', markersize=5, alpha=0.5)
    plt.legend()
    plt.title('NHT')
    plt.savefig(f'./results/nht_comp_k_{k}.png')
