from matplotlib import pyplot as plt
import torch
import os
from utils.loadcolor import loadcolor
from matplotlib.ticker import ScalarFormatter

def getfig(plott, plotf, config_yaml, logpath, truncationnum):
    truncationnum = truncationnum*1e2
    savename = config_yaml["log"]["savename"]
    histt1 = plott[0]; histt2 = plott[1]; histt3 = plott[2]; histt4 = plott[3]; histt5 = plott[4]
    histf1 = plotf[0]; histf2 = plotf[1]; histf3 = plotf[2]; histf4 = plotf[3]; histf5 = plotf[4]

    pcolor = loadcolor()
    plt.figure(figsize=(8,6), dpi=80, facecolor='w')
    plt.subplots_adjust(left=0.15, right=0.9, bottom=0.15, top=0.9)


    plt.plot(histt4.detach().numpy(), histf4.detach().numpy(), '-', linewidth=5, markersize=3, color=pcolor['red'].numpy(),label='GS-JOBCD')
    plt.plot(histt5.detach().numpy(), histf5.detach().numpy(), '-', linewidth=5, markersize=3, color=pcolor['black'].numpy(),label='J-JOBCD')
    plt.plot(histt3.detach().numpy(), histf3.detach().numpy(), ':', linewidth=5, markersize=3, color=pcolor['blue'].numpy(),label='CSDM')
    plt.plot(histt2.detach().numpy(), histf2.detach().numpy(), '-.', linewidth=5, markersize=3, color=pcolor['green'].numpy(), label='ADMM')
    plt.plot(histt1.detach().numpy(), histf1.detach().numpy(), '--', linewidth=5, markersize=3, color=pcolor['purple'].numpy(), label='UMCM')

    plt.gca().yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
    plt.gca().ticklabel_format(style='sci', axis='y', scilimits=(0, 0), useOffset=False, useLocale=False,
                               useMathText=True)
    plt.gca().yaxis.get_offset_text().set_size('x-small')
    plt.gca().yaxis.get_major_formatter().set_powerlimits((0, 1))
    plt.gca().yaxis.get_major_formatter().set_scientific(True)
    plt.gca().yaxis.get_major_formatter().set_useOffset(False)
    plt.tick_params(axis='y', which='major', labelsize=10)

    hleg = plt.legend(fontsize=15,frameon=True, framealpha=1)
    hleg.set_title('', prop={'size': 15, 'weight': 'normal', 'family': 'times new Roman'})
    hleg.get_title().set_fontsize('15')

    plt.xlabel('Time (seconds)', fontsize=20)
    plt.ylabel('Objective', fontsize=20)
    plt.xticks(fontsize=17)
    plt.yticks(fontsize=17)

    plt.grid(True)
    plt.savefig(os.path.join(logpath, '{}.png'.format(savename)), format='png', dpi=300)
    plt.savefig(os.path.join(logpath, '{}.eps'.format(savename)), format='eps', dpi=300)
    plt.show()
