import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# import colormap from sns
import seaborn as sns
# cmap = sns.cubehelix_palette(as_cmap=True)
cmap = sns.color_palette("coolwarm", as_cmap=True)

data = pd.read_csv('property-mocked.csv')

xlabel = "PB"
ylabel = "VA"
zlabels = ["mocked-pbpd3_GM", "mocked-liquid_GM", "mocked-fcc_GM"]

temperatures = [453]

for temperature in temperatures:
    for zlabel in zlabels:
        # Select data for a given temperature
        data_temp = data[data["T"] == temperature]

        # Plot the data
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        # ax.plot_trisurf(data_temp[xlabel], data_temp[ylabel], data_temp[zlabel], color='red', linewidth=0.2)
        ax.plot_trisurf(data_temp[xlabel], data_temp[ylabel], data_temp[zlabel], cmap=cmap, linewidth=0.2)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_xlim([np.min(data_temp[xlabel]), np.max(data_temp[xlabel])])
        ax.set_ylim([np.min(data_temp[ylabel]), np.max(data_temp[ylabel])])

        # limit x and y axis to use only 3 ticks
        ax.set_xticks(ax.get_xticks()[::2])
        ax.set_yticks(ax.get_yticks()[::2])

        # roate the view
        ax.view_init(10, -30)

        # plt.legend()
        plt.tight_layout()

        # plt.show()
        plt.savefig("surface-{}-{}K.png".format(zlabel, temperature), dpi=300)