import matplotlib.pyplot as plt
import numpy as np
import Knowledge_Distillation.wasserstein as gdw

def draw_PD(PD1, save_name, PD2 = None):
    x = PD1[:, 0]
    y = PD1[:, 1]
    plt.scatter(x, y, color = 'hotpink', label = 'predict')

    x = [0, 0.5, 1]
    y = [0, 0.5, 1]
    plt.plot(x, y, color = 'black', label = 'diagonal')

    if PD2 is not None:
        x = PD2[:, 0]
        y = PD2[:, 1]
        plt.scatter(x, y, color= '#88c999', label = 'ground truth')

    plt.xlim((0, 1.1))
    plt.ylim((0, 1.1))
    plt.legend()
    if PD2 is not None:
        plt.title("Loss: {}".format(gdw.wasserstein_distance_inference(PD1, PD2, order = 2, enable_autodiff = True)[0]))
    plt.savefig(save_name)
    plt.close()


if __name__ == '__main__':
    PD0_predict = np.array([[0.0557, 1.0645],
        [0.1205, 0.9456],
        [0.2760, 0.6805],
        [0.1401, 0.8785],
        [0.1267, 0.8189],
        [0.1503, 0.8817],
        [0.1857, 0.8033],
        [0.2818, 0.7801],
        [0.2622, 0.6145],
        [0.2183, 0.4986],
        [0.1546, 0.6561],
        [0.2290, 0.5459],
        [0.2696, 0.6339],
        [0.2170, 0.5002],
        [0.2444, 0.6489],
        [0.2818, 0.7801],
        [0.2247, 0.5269],
        [0.2247, 0.5269],
        [0.2795, 0.6082],
        [0.2422, 0.5945],
        [0.2052, 0.4403],
        [0.1505, 0.9245],
        [0.2461, 0.6459],
        [0.1505, 0.9245],
        [0.2309, 0.5496],
        [0.1294, 0.8308],
        [0.1232, 0.5848],
        [0.2924, 0.6079]])

    PD0_gt = np.array([[0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0385, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 0.3462],
        [0.0769, 0.3462],
        [0.0769, 0.3462],
        [0.0769, 0.3462],
        [0.0769, 0.1923],
        [0.0769, 0.1538],
        [0.1154, 0.1923],
        [0.1154, 0.1923],
        [0.1154, 0.1923],
        [0.0385, 1.0000]])
    draw_PD(PD1 = PD0_predict, save_name = './Cora_degree_0_dim.png', PD2=PD0_gt)

    PD1_predict = np.array([[0.1451, 0.3440],
        [0.1226, 1.0115],
        [0.1211, 1.0115],
        [0.1211, 1.0115],
        [0.1046, 0.3272],
        [0.1859, 1.0148],
        [0.0811, 0.9763],
        [0.0178, 1.0145],
        [0.1211, 1.0115],
        [0.0473, 1.0136],
        [0.0535, 1.0133],
        [0.1211, 1.0115],
        [0.0807, 1.0128],
        [0.0320, 0.2771],
        [0.1588, 1.0140],
        [0.0899, 1.0114],
        [0.0514, 1.0132],
        [0.0673, 0.4324],
        [0.1211, 1.0115]])
    PD1_gt = np.array([[0.1154, 0.1923],
        [0.1154, 0.3462],
        [0.1154, 0.3462],
        [0.1538, 0.3462],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.0769, 1.0000],
        [0.1154, 1.0000],
        [0.1154, 1.0000],
        [0.1154, 1.0000],
        [0.1154, 1.0000],
        [0.1538, 1.0000],
        [0.1538, 1.0000],
        [0.1923, 1.0000]])
    draw_PD(PD1 = PD1_predict, save_name = './Cora_degree_1_dim.png', PD2=PD1_gt)
    plt.show()