# -*- coding: utf-8 -*-
"""
Based on a script from phanikumar.ds@gmail.com
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.ion()
# Read the selected file: Numpy knows rather well how to do that :)
X, Y, Z = np.loadtxt('./Plasma.csv', delimiter=",", unpack=True)
# Getting ready for the dispaly
fig = plt.figure(num="plasma")
ax = fig.add_subplot(1, 1, 1, projection='3d')
# Set the labels
ax.set_xlabel('X (mm)')
ax.set_ylabel('Y (mm)')
ax.set_zlabel('CPD (mV)')
# Setting the places properly...
#
# It would be nice if an mplot3d expert could tell if there is another way
# to set the z-axis label and tick on the left side (rather than the right
# one) that would not rely on private attributes...
#
tmp_planes = ax.zaxis._PLANES
ax.zaxis._PLANES = (tmp_planes[2], tmp_planes[3],
tmp_planes[0], tmp_planes[1],
tmp_planes[4], tmp_planes[5])
# Reshape the data from 1D to 2D because of `ax.plot_surf` signature
# Find the correct shape (could also be done manually by looking at the file)
nb_unique_x = len(np.unique(X))
nb_unique_y = len(np.unique(Y))
shape_xy = (nb_unique_y, nb_unique_x)
xx = X.reshape(shape_xy)
yy = Y.reshape(shape_xy)
zz = Z.reshape(shape_xy)
# Surface plot
surf = ax.plot_surface(xx, yy, zz, cmap="inferno")
cbar = fig.colorbar(surf, ax=ax, label="CPD (mV)")
fig.tight_layout()
fig.savefig(fig.get_label() + ".png")