from matplotlib import pyplot as plt
import numpy as np
from matplotlib.ticker import MultipleLocator


def slope(xs, ys):
    nans = np.argwhere(np.isnan(ys))
    if len(nans)>0:
        xs = np.delete(xs, nans)
        ys = np.delete(ys, nans)
    return np.mean(xs*ys)/np.mean(xs*xs)

def mkline(ca, xs):
    return ca*xs

# leggiamo i dati in ingresso usando la funzione genfromtxt di numpy
dati = np.genfromtxt("misure_affondamenti.dat", skip_header=1)

# i dati vengono caricati in un np.ndarray 2d e sono accessibili tramite slice multidimensionali

# riscaliamo i dati sottraendo il valore contenuto nella prima riga
dati_scalati = dati - dati[0]

# a differenza di R non e' possibile accedere alle righe/colonne tramite label
# per cui usiamo l'indirizzamento diretto, facendo una slice dell'ndarray

x  = dati_scalati[:,0] # la prima colonna contiene le ascisse
y1 = dati_scalati[:,1]
y2 = dati_scalati[:,2]
y3 = dati_scalati[:,3]
y4 = dati_scalati[:,4]

# prima di tutto un plot "complessivo"
figure, ax = plt.subplots()

ln1, = ax.plot(x, y1, 'o-', color="darkblue", label="Cilindro 1")
ln2, = ax.plot(x, y2, 'o-', color="blue", label="Cilindro 2")
ln3, = ax.plot(x, y3, 'o-', color="orange", label="Prisma")
ln4, = ax.plot(x, y4, 'o-', color="darkgrey", label="Cono")

ax.legend(handles=[ln1, ln2, ln3, ln4])

ax.grid()
ax.yaxis.set_major_locator(MultipleLocator(5.))

ax.set_title("Massa d'acqua spostata in funzione dell'affondamento")
ax.set_xlabel("affondamento (cm)")
ax.set_ylabel("massa d'acqua spostata (g)")
figure.tight_layout()

# e mostriamo sullo schermo l'immagine risultante
figure.show()

input("Premi return per continuare")

# creiamo uno spazio per i plot
figure, axs = plt.subplots(2,2)

# e facciamo un po' di grafici
ln1, = axs[0,0].plot(x, y1, color="darkblue", label="Cilindro 1")
m1 = slope(x, y1)
linea1 = mkline(m1, x)
ln1b, = axs[0,0].plot(x, linea1, linestyle="dashed", color="green", label="Slope: %4.3f"%m1)

ln2, = axs[0,1].plot(x, y2, color="blue", label="Cilindro 2")
m2 = slope(x, y2)
linea2 = mkline(m2, x)
ln2b, = axs[0,1].plot(x, linea2, linestyle="dashed", color="green", label="Slope: %4.3f"%m2)

ln3, = axs[1,0].plot(x*x, y3, color="orange", label="Prisma")
m3 = slope(x*x, y3)
linea3 = mkline(m3, x*x)
ln3b, = axs[1,0].plot(x*x, linea3, linestyle="dashed", color="green", label="Slope (vs x^2): %4.3f"%m3)

ln4, = axs[1,1].plot(x*x*x, y4, color="darkgrey", label="Cono")
m4 = slope(x*x*x, y4)
linea4 = mkline(m4, x*x*x)
ln4b, = axs[1,1].plot(x*x*x, linea4, linestyle="dashed", color="green", label="Slope (vs x^3): %4.3f"%m4)


axs[0,0].legend(handles=[ln1, ln1b])
axs[0,1].legend(handles=[ln2, ln2b])
axs[1,0].legend(handles=[ln3, ln3b])
axs[1,1].legend(handles=[ln4, ln4b])

axs[0,0].grid()
axs[0,1].grid()
axs[1,0].grid()
axs[1,1].grid()

figure.suptitle("Variazione di letture sulla bilancia in funzione dell'affondamento")
figure.tight_layout()

# e mostriamo sullo schermo l'immagine risultante
figure.show()

input("Premi return per continuare")

# prima di tutto un plot "complessivo"
figure, ax = plt.subplots()

xi = np.linspace(0, 6, int(6./0.1))

ln1,  = ax.plot(x, y1, 'o-', color="darkblue", label="Cilindro 1")
ln1b, = ax.plot(xi, m1*xi, linestyle="dashed", color="darkblue")
ln2, = ax.plot(x, y2, 'o-', color="blue", label="Cilindro 2")
ln2b, = ax.plot(xi, m2*xi, linestyle="dashed", color="blue")
ln3, = ax.plot(x, y3, 'o-', color="orange", label="Prisma")
ln3b, = ax.plot(xi, m3*xi*xi, linestyle="dashed", color="orange")
ln4, = ax.plot(x, y4, 'o-', color="darkgrey", label="Cono")
ln4b, = ax.plot(xi, m4*xi*xi*xi, linestyle="dashed", color="darkgrey")

ax.legend(handles=[ln1, ln2, ln3, ln4])

ax.grid()
ax.yaxis.set_major_locator(MultipleLocator(5.))

ax.set_title("Massa d'acqua spostata in funzione dell'affondamento")
ax.set_xlabel("affondamento (cm)")
ax.set_ylabel("massa d'acqua spostata (g)")
figure.tight_layout()

# e mostriamo sullo schermo l'immagine risultante
figure.show()

input("Premi return per continuare")