import matplotlib.pyplot as plt
import numpy as np


def mobilenet_v1():
    markersize = 8
    plt.figure(figsize=(5.7, 5.2))
    GDP = np.array([0.267, 0.358101, 0.450991, 0.5037, 0.5825, 0.71424]) * 100
    GDP_ft_top1 = np.array([68.524, 69.938, 70.666, 71.308, 71.906, 72.435])
    our_top1_wo_ft = np.array([67.922, 69.274, 70.532, 70.888, 71.436, 71.744])
    autoSlim = np.array([150, 325]) / 569 * 100
    autoSlim_top1 = 100 - np.array([32.1, 28.5])
    MetaPruning = np.array([149, 281, 324]) / 569 * 100
    MetaPruning_top1 = np.array([66.1, 70.6, 70.9])
    eagle = np.array([284]) / 569 * 100
    eagle_top1 = np.array([70.9])
    amc = np.array([0.4, 0.5]) * 100
    amc_top1 = 70.6 - np.array([1.7, 0.4])
    netAdapt = np.array([284.3]) / 569 * 100
    netAdapt_top1 = np.array([69.1])
    scratch = np.array([286, 150]) / 569 * 100
    scratch_top1 = np.array([70.7, 65.5])
    dcp = 100 - np.array([51.25])
    dcp_top1 = 100 - np.array([29.5])
    plt.plot(eagle, eagle_top1, "ko", label="EagleE [ECCV 2020]", markersize=markersize)
    plt.plot(scratch, scratch_top1, 'bs-', label="PFS [AAAI 2020]", markersize=markersize)
    plt.plot(MetaPruning, MetaPruning_top1, "mh-", label="MetaP [ICCV 2019]", markersize=markersize)
    plt.plot(amc, amc_top1, "y*-", label="AMC [ECCV, 2018]", markersize=markersize)
    plt.plot(netAdapt, netAdapt_top1, "c<", label="NetAdapt [ECCV 2018]", markersize=markersize)
    plt.plot(dcp, dcp_top1, "D", color=(1.0, 0.66, 0.0), label="DCP [NeurIPS 2018]", markersize=markersize)
    plt.plot(autoSlim, autoSlim_top1, "P-", label="AutoSlim [NeurIPS 2019]", color=(0.4, 0.4, 1.0),
             markersize=markersize)
    plt.plot(GDP, our_top1_wo_ft, '>-', color=(0.8, 0.8, 0.8), linewidth='1', label="GDP w/o FT [Ours]",
             markersize=markersize * 0.5)
    plt.plot(GDP, GDP_ft_top1, 'ro-', label="GDP [Ours]", markersize=markersize * 1.2, linewidth='2.5')
    plt.xlabel("FLOPs Ratio w.r.t. Original Model [%]", fontsize=15)
    plt.ylabel("Top-1 Accuracy [%]", fontsize=15)
    plt.subplots_adjust(top=0.95, bottom=0.15)
    plt.legend(bbox_to_anchor=(0.435, -0.01), loc=3, ncol=1, prop={'size': 12})
    plt.tick_params(labelsize=12)
    plt.subplots_adjust(top=0.99, bottom=0.1, right=0.99)
    # plt.show()
    plt.savefig("mobilenet_v1.png")
    plt.close()


def mobilenet_v2():
    meta_pruning = np.array([43, 84, 105, 124, 140, 217]) / 313 * 100
    meta_pruning_top1 = np.array([58.3, 63.8, 65, 67.3, 68.2, 71.2])
    good = np.array([107, 138, 152, 170, 201, 220, 258]) / 314 * 100
    good_top1 = np.array([66.9, 68.8, 69.7, 70.4, 71.2, 71.6, 71.9])
    emc = np.array([50, 60, 70])
    emc_top1 = np.array([69.4, 70.8, 71.4])
    amc = np.array([0.7]) * 100
    amc_top1 = 71.8 - np.array([1.0])
    dmcp = np.array([43, 59, 87, 97, 211]) / 300 * 100
    dmcp_top1 = np.array([59.1, 62.7, 66.1, 67, 72.2])
    scratch = np.array([210, 300]) / 300 * 100
    scratch_top1 = np.array([70.9, 72.1])

    GDP = np.array(
        [0.130367, 0.164592, 0.219775, 0.26394, 0.304017, 0.36110, 0.437698, 0.52104, 0.624057, 0.7764, 0.96107]) * 100
    GDP_ft_top1 = np.array([60.254, 62.616, 65.404, 66.752, 67.790, 68.924, 70.150, 70.916, 71.834, 72.466, 73.114])
    our_top1_wo_ft = np.array([59.824, 62.272, 64.914, 66.352, 67.516, 68.820, 69.962, 70.668, 71.402, 72.268, 72.762])

    markersize = 8
    plt.figure(figsize=(5.7, 5.0))
    plt.plot(meta_pruning, meta_pruning_top1, 'bs-', label="MetaP [ICCV 2019]", markersize=markersize)
    plt.plot(good, good_top1, "mh-", label="GFS [PMLR 2020]", markersize=markersize)
    plt.plot(emc, emc_top1, "y*-", label="LeGR [CVPR 2020]", markersize=markersize)
    plt.plot(dmcp, dmcp_top1, "kp-", label="DMCP [CVPR 2020]", markersize=markersize)
    plt.plot(amc, amc_top1, "c<", label="AMC [ECCV 2018]", markersize=markersize)
    plt.plot(scratch, scratch_top1, "g8-", label="PFS [AAAI 2020]")
    plt.plot(GDP, our_top1_wo_ft, '>-', color=(0.8, 0.8, 0.8), linewidth='1', label="GDP w/o FT [Ours]",
             markersize=markersize * 0.5)
    plt.plot(GDP, GDP_ft_top1, 'ro-', label="GDP [Ours]", markersize=markersize * 1.2)
    plt.yticks([58, 60, 62, 64, 66, 68, 70, 72])
    plt.xlabel("FLOPs Ratio w.r.t. Original Model [%]", fontsize=15)
    plt.ylabel("Top-1 Accuracy [%]", fontsize=15)

    plt.legend(bbox_to_anchor=(0.41, -0.01), loc=3, ncol=1, prop={'size': 15})
    plt.tick_params(labelsize=15)  # 刻度字体大小20
    plt.subplots_adjust(top=0.99, bottom=0.11, right=0.98)
    # plt.show()
    plt.savefig("mobilenet_v2.png")
    plt.close()


def resnet50():
    eagle_eye = np.array([1, 2, 3]) / 4.1 * 100
    eagle_eye_top1 = np.array([74.2, 76.4, 77.1])

    meta_pruning = np.array([1, 2, 3]) / 4.1 * 100
    meta_pruning_top1 = np.array([73.4, 75.4, 76.2])

    operataion_aware = 100 - np.array([54.3])
    operataion_aware_top1 = 75.89 - np.array([0.62])

    stable_low_rank = 100 - np.array([2.64]) / 4.1 * 100
    stable_low_rank_top1 = 76.13 - np.array([1.47])

    emc = np.array([47, 58, 73])
    emc_top1 = np.array([75.3, 75.7, 76.2])

    dmcp = np.array([1.1, 2.2, 2.8]) / 4.1 * 100
    dmcp_top1 = np.array([74.4, 76.2, 76.7])

    important = np.array([1.34, 2.25, 2.66, 3.27]) / 4.09 * 100
    important_top1 = 100 - np.array([28.31, 25.50, 24.52, 23.57])

    OICSR = 100 - np.array([50, 44.43, 37.30])
    OICSR_top1 = 76.31 - np.array([0.37, 0.01, -0.22])

    Collaborative = 100 - np.array([48.8, 54.1, 54.1])
    Collaborative_top1 = 76.15 - np.array([0.65, 0.83, 0.94])

    gate = 100 - np.array([40.54, 55.06])
    gate_top1 = np.array([76.19, 75.18])

    group = np.array([46.55])
    group_top1 = 100 - np.array([25.3])

    RCLD = np.array([1.11]) / 4.09 * 100
    RCLD_top1 = np.array([73.39])

    HRrank = np.array([2.3, 1.55, 0.98]) / 4.09 * 100
    HRrank_top1 = np.array([74.98, 71.98, 69.1])

    Criteria = 100 - np.array([60.8])
    Criteria_top1 = np.array([74.46])

    scratch = np.array([1, 2, 3]) / 4.1 * 100
    scratch_top1 = np.array([72.8, 75.6, 76.7])

    Centripetal_SGD = 100 - np.array([36.75, 46.24, 55.76])
    Centripetal_SGD_top1 = np.array([75.27, 74.93, 74.54])

    SCOP = 100 - np.array([45.3, 54.6])
    SCOP_top1 = 100 - np.array([24.05, 24.75])

    autoSlim = np.array([0.57, 1, 2, 3]) / 4.1 * 100
    autoSlim_top1 = 100 - np.array([27.8, 26, 24.4, 24])

    ABCPruner = np.array([1794.45, 1890.60]) / 4135.70 * 100
    ABCPruner_top1 = np.array([73.52, 73.86])

    DCPS = 100 - np.array([46.2, 55.4])  # Differentiable Channel Pruning Search
    DCPS_top1 = np.array([75.51, 75.37])
    our = np.array([0.1888611, 0.2539, 0.30421, 0.344074, 0.38667, 0.49003]) * 100
    our_top1 = np.array([74.080, 75.106, 75.560, 76.070, 76.414, 77.052])
    our_top1_wo_ft = np.array([73.844, 74.980, 75.582, 76.018, 76.136, 76.792])
    step = np.array([0.6069, 0.4762, 0.3545, 0.2785]) * 100
    step_top1 = np.array([76.258, 75.836, 75.312, 74.434])
    markersize = 8
    plt.figure(figsize=(6, 8))
    plt.plot(eagle_eye, eagle_eye_top1, "cd-", label="EagleE [ECCV 2020]", markersize=markersize)
    plt.plot(meta_pruning, meta_pruning_top1, 'bs-', label="MetaP [ICCV 2019]", markersize=markersize)
    plt.plot(operataion_aware, operataion_aware_top1, 'g^', label="SCP [PMLR 2020]", markersize=markersize)
    plt.plot(stable_low_rank, stable_low_rank_top1, "mh", label="SLRTD [ECCV 2020]", markersize=markersize)
    plt.plot(emc, emc_top1, "y*-", label="LeGR [CVPR 2020]", markersize=markersize)
    plt.plot(dmcp, dmcp_top1, "kp-", label="DMCP [CVPR 2020]", markersize=markersize)
    plt.plot(important, important_top1, "c<-", label="Talyor [CVPR 2019]", markersize=markersize)
    plt.plot(OICSR, OICSR_top1, ".-", color=(1.0, 0.66, 0.0), label="OICSR [CVPR 2019]", markersize=markersize)
    plt.plot(Collaborative, Collaborative_top1, "D-", color=(0.0, 0.66, 1.0), label="CCP [PMLR 2019]",
             markersize=markersize)
    plt.plot(gate, gate_top1, 'h-', label="GD [NeurIPS 2019]", color=(0.1, 0.3, 0.5), markersize=markersize)
    plt.plot(group, group_top1, 'ko', label="Hinge [CVPR 2020]", markersize=markersize)
    plt.plot(RCLD, RCLD_top1, 'gh', label="CURL [CVPR 2020]", markersize=markersize)
    plt.plot(HRrank, HRrank_top1, 'mX-', label="HRank [CVPR 2020]", markersize=markersize)
    plt.plot(scratch, scratch_top1, "8-", label="PFS [AAAI 2020]", color=(0.5, 0.01, 0.5))
    plt.plot(Criteria, Criteria_top1, 'rp', label="LFPC [CVPR 2020]", markersize=markersize)

    plt.plot(Centripetal_SGD, Centripetal_SGD_top1, "p-", label="C-SGD [CVPR 2019]", color=(0.8, 0.0, 0.1),
             markersize=markersize)
    plt.plot(SCOP, SCOP_top1, "p-", label="SCOP [NeurIPS 2020]", color=(0.3, 0.9, 0.1), markersize=markersize)
    plt.plot(ABCPruner, ABCPruner_top1, "v-", label="ABCPruner [IJCAI 2020]",
             color=(0.2, 0.1, 0.8), markersize=markersize)
    plt.plot(autoSlim, autoSlim_top1, "P-", label="AutoSlim [NeurIPS 2019]", color=(0.4, 0.4, 1.0),
             markersize=markersize)
    plt.plot(DCPS, DCPS_top1, "*-", color=(0.7, 0.1, 0.4), label="DCPS")
    plt.plot(step, step_top1, 'rd-', label="GDP [Ours]", markersize=markersize)
    plt.plot(our, our_top1_wo_ft, '>-', color=(0.8, 0.8, 0.8), linewidth='1', label="GDP* w/o FT [Ours]",
             markersize=markersize * 0.5)
    plt.plot(our, our_top1, 'ro-', label="GDP* [Ours]", markersize=markersize * 1.2)

    plt.xlabel("FLOPs Ratio w.r.t. Original Model [%]", fontsize=15)
    plt.ylabel("Top-1 Accuracy [%]", fontsize=15)
    plt.ylim(69, 77.2)

    plt.legend(bbox_to_anchor=(0.4, -0.01), loc=3, ncol=1, prop={'size': 15})
    plt.tick_params(labelsize=15)
    plt.subplots_adjust(top=0.995, bottom=0.07, right=0.99)
    plt.legend()
    # plt.show()
    plt.savefig("resnet50.png")
    plt.close()


def draw_mobiV2_decay():
    ratio = np.array([0.4044, 0.4097, 0.4377, 0.4850, 0.5551, 0.6009, 0.6073]) * 100
    top1 = np.array([69.678, 69.860, 70.150, 70.774, 71.442, 71.482, 70.990])
    GDP = np.array([0.36110, 0.437698, 0.52104, 0.624057]) * 100
    GDP_ft_top1 = np.array([68.924, 70.150, 70.916, 71.834])

    markersize = 8
    plt.figure(figsize=(5.7, 5.0))
    marker = ['bs', "mh", "y*", "kp", "c<", "g8", "P-"]
    decay = [0.86, 0.88, 0.90, 0.92, 0.94, 0.96, 0.98]
    for i in range(7):
        plt.plot(ratio[i], top1[i], marker[i], markersize=markersize)
        plt.text(ratio[i], top1[i], str(decay[i]), {"fontsize": 15})
    plt.plot(GDP, GDP_ft_top1, 'ro-', label="GDP [Ours]", markersize=markersize)

    plt.yticks([68, 69, 70, 71, 72])
    plt.ylim(68, 72.5)
    plt.xlabel("FLOPs Ratio w.r.t. Original Model [%]", fontsize=15)
    plt.ylabel("Top-1 Accuracy [%]", fontsize=15)
    plt.tick_params(labelsize=15)
    plt.subplots_adjust(top=0.99, bottom=0.11, right=0.98)
    # plt.show()
    plt.savefig("epsilon_decay.png")
    plt.close()


def draw_mobiV2_init():
    ratio = np.array([0.4850, 0.5379, 0.4611, 0.5585]) * 100
    top1 = np.array([70.774, 71.152, 70.456, 71.310])
    GDP = np.array([0.36110, 0.437698, 0.52104, 0.624057]) * 100
    GDP_ft_top1 = np.array([68.924, 70.150, 70.916, 71.834])
    markersize = 8
    plt.figure(figsize=(5.7, 5.0))
    marker = ['bs', "mh", "y*", "kp"]
    decay = [0.1, 0.5, 0.05, 1.0]
    for i in range(4):
        plt.plot(ratio[i], top1[i], marker[i], markersize=markersize)
        plt.text(ratio[i], top1[i], str(decay[i]), {"fontsize": 15})
    plt.plot(GDP, GDP_ft_top1, 'ro-', label="GDP [Ours]", markersize=markersize)

    plt.yticks([68, 69, 70, 71, 72])
    plt.ylim(68, 72.5)
    plt.xlabel("FLOPs Ratio w.r.t. Original Model [%]", fontsize=15)
    plt.ylabel("Top-1 Accuracy [%]", fontsize=15)
    plt.tick_params(labelsize=15)
    plt.subplots_adjust(top=0.99, bottom=0.11, right=0.98)
    # plt.show()
    plt.savefig("epsilon_init.png")
    plt.close()


if __name__ == "__main__":
    mobilenet_v1()
    mobilenet_v2()
    resnet50()
    draw_mobiV2_decay()
    draw_mobiV2_init()
