import matplotlib.pyplot as plt
import pickle
import sys

if __name__ == "__main__":
    D = float(sys.argv[1])
    k = int(sys.argv[2])
    with open(f'./results/results_fg_D_{D}_k_{k}.pkl', 'rb') as file:
        resultsFG = pickle.load(file)
    with open(f'./results/results_sg_D_{D}_k_{k}.pkl', 'rb') as file:
        resultsSG = pickle.load(file)
    with open(f'./results/results_zo_D_{D}_k_{k}.pkl', 'rb') as file:
        resultsZO = pickle.load(file)

    plt.figure()
    plt.plot(resultsFG['ifo'], resultsFG['losses'], marker='<', label='FG-HT-TSP', markersize=5, color='blue')
    plt.plot(resultsSG['ifo'], resultsSG['losses'], marker='>', label='HSG-HT-TSP', markersize=5, color='orange')
    plt.legend()
    plt.xlim(0, 20000)
    plt.title('IFO')
    plt.savefig(f'./results/ifo_k_{k}_D_{D}.png')

    plt.figure()
    plt.plot(resultsFG['nht'], resultsFG['losses'], marker='<', label='FG-HT-TSP', markersize=5, color='blue')
    plt.plot(resultsSG['nht'], resultsSG['losses'], marker='>', label='HSG-HT-TSP', markersize=5, color='orange')
    plt.plot(resultsZO['nht'], resultsZO['losses'], marker='>', label='HZO-HT-TSP', markersize=5, color='green')
    plt.legend()
    plt.title('NHT')
    plt.xlim(0, 100)
    plt.savefig(f'./results/nht_k_{k}_D_{D}.png')


    plt.figure()
    plt.plot(resultsZO['izo'], resultsZO['losses'], marker='<', label='HZO-HT-TSP', markersize=5, color='green')
    plt.legend()
    plt.title('IZO')
    plt.savefig(f'./results/izo_k_{k}_D_{D}.png')