'''
Skrevet av: Vegard G. Jervell
Hensikt: Hjelpe folk med å lære alt de ikke lærte i ITGK
Spesifikt: Matplotlib
'''

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D #til 3D-plotting, du må importe denne selv om PyCharm sier at den ikke er i bruk
import matplotlib.cm as cm #for å autogenerere fine farger
import numpy as np

#lager litt dummy data
x_akse = np.linspace(1, 10, 1000) #pleier ikke å trenge mer enn 100 punkter
y_akse = np.linspace(1, 5, 1000) #men det ene eksempeldatasettet (y4_list) får litt for få punkter med 100, så jeg bruker 1000

x_punkter = np.linspace(1, 10, 5)
y_punkter = np.linspace(1, 5, 5)
y1_list = x_akse ** 2
y2_list = np.exp(0.1 * x_akse ** 2)
y3_list = np.sin(x_akse)
y4_list = np.cos(x_akse ** 2)

z_func = lambda x,y: np.exp(np.sin(x * y)) #en fænsi one-liner funksjon. det samme som:
#def z_func(x, y):
    #return np.exp(np.sin(x * y))


#Skal lage to plots i samme figur, med 1 rad og 2 kolonner
fig, axs = plt.subplots(1,2, figsize=(10,5)) #axs er en liste med to "plots", figsize=(10,5) gjør plottet større

axs[0].plot(x_akse, y1_list, label='data 1')
axs[1].plot(x_akse, y2_list, label='data 2')

axs[0].set_xlabel('en x-akse') #merk at det er en annen kommando for å sette x- og y-label
axs[0].set_ylabel('noen datapunkter')

#i stedet for å bruke axs[i].gjør_noe() kan vi si til matplotlib at vi vil jobbe med et besemt plot

plt.sca(axs[1]) #nå skal jeg jobbe med plot nr. 1

plt.xlabel('den samme x-aksen')
plt.ylabel('men andre data')
plt.legend() #denne lager bare legend på det plottet jeg jobber med nå

plt.suptitle('En tittel for hele figuren') #merk "suptitle()" for å lage en tittel til hele figuren

plt.savefig('plots/plot1')
plt.show()
plt.close(fig) #lukker figuren til slutt for å være snille med minnet.

#skal lage tre plots fordelt på tre rader, med felles x-akse
fig, axs = plt.subplots(3, 1, sharex=True, figsize = (10,5))

z_list = z_func(x_akse, y_akse)

axs[0].plot(x_akse, y3_list, label='y3_data')
axs[1].plot(x_akse, y4_list, label='y4_data')
axs[2].plot(x_akse, z_list, label='z_data')

for i in range(3):
    axs[i].legend(loc='lower left') #viser legend for alle tre med en for-løkke

plt.suptitle('Wow! Dette ser fænsi ut!')
plt.savefig('plots/plot2')
plt.show()
plt.close(fig) #dette er lurt å gjøre hvis du har et script som genererer veldig mange plots

#Nå skal vi lage et 3D plot!

x_3D, y_3D = np.meshgrid(x_akse, y_akse) #lager et "aksesystem" vi kan bruke i 3D
x_punkter_3D, y_punkter_3D = np.meshgrid(x_punkter, y_punkter)
fig = plt.figure() #lager en figur
ax = fig.add_subplot(111, projection='3d') #legger til et plot som er sentrert i figuren og viser et 3d plot. Putter plottet i variabelen 'ax'.

z_data_3D = z_func(x_3D, y_3D)
z_punkter_3D = z_func(x_punkter_3D, y_punkter_3D)

ax.scatter(x_punkter_3D, y_punkter_3D, z_punkter_3D, color='red') #scatterplot
ax.plot_wireframe(x_3D, y_3D, z_data_3D - 5, color='black', rcount=10, ccount=10) #wireframe, rcount og ccount angir antall linjer i hver retning.
ax.plot_surface(x_3D, y_3D, z_data_3D + 5, color='green') #lager en flate

ax.set_xlabel('En retning')
ax.set_ylabel('En annen retning')
ax.set_zlabel('verdier')
plt.title('Lykke til med å lage dette på 3 min. med 9 linjer i Excel ;)')
plt.savefig('plots/plot3')
plt.show()
plt.close(fig) #ikke glem denne :)

#Nå skal vi lage fænsi plots med masse fine farger
# Vi begynner med et enkelt eksempel
fine_farger = cm.get_cmap('autumn') #google 'matplotlib colormap' for å finne en fullstendig liste med mulige ting å gi denne

#Funker ved at
# fine_farger(tall) = farge, hvor 'tall' er en verdi mellom 0 og 1
# nøkkelordet 'autumn' sier oss hvilket 'sett' med farger vi skal bruke.

plt.plot(x_akse, y3_list, color=fine_farger(0.25), label='y3')
plt.plot(x_akse, y4_list, color=fine_farger(0.5), label='y4')
plt.plot(x_akse, z_list, color=fine_farger(0.75), label='z')
plt.legend()
plt.title('Basic plot Boi')
plt.savefig('plots/plot4')
plt.show()
plt.close()

#Nå skal vi gjøre det litt mer komplisert
fig, axs = plt.subplots(3,1, figsize=(15,8)) #tre plots i tre rader og en kolonne

cmap_1 = cm.get_cmap('Wistia') #Et sett med farger
cmap_2 = cm.get_cmap('summer') #et annet sett med farger
cmap_3 = cm.get_cmap('winter') #google 'matplotlib colormap' for å finne en fullstendig liste med muligheter

# Hvordan bruker jeg cmap_1, cmap_2 og cmap_3????
# Svar: De er funksjoner som fungerer ved at
# cmap_1(tall) = farge, hvor 'tall' er en verdi mellom 0 og 1
# cmap_2(tall) = en annen farge, og cmap_3(tall) = en tredje farge
# Hver av disse gir altså et unikt sett med farger
# hvis jeg har data som går fra 0-100 kan jeg bare skalere det ved å skrive
# cmap_1(data[i] / max(data[i] ) = farge

cmaps = [cmap_1, cmap_2, cmap_3] #skal bruke en for-løkke, så putter disse i en liste
cmap_names = ['Wistia', 'summer', 'winter'] #Skal bruke navnene som aksetitler, så putter navnene i en liste også
some_values = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6] #lager noen verdier vi skal bruke

# Nå skal vi plotte
for row in range(3): #iterer over rader
    cmap = cmaps[row] #henter en ny colormap for hver rad
    for val in some_values: #iterer over verdiene vi skal plotte for
        z_data = z_func(x_akse, val) #nå regner vi ut z_func(x, y) for alle verdiene i x_akse, med y = val
        axs[row].plot(x_akse, z_data, color=cmap(val/max(some_values)), label = val) #legg merke til at 'val' skaleres ved å dele på max(some_values)

    axs[row].legend(title='my_value') #viser legend og setter tittel på den
    axs[row].set_xlim(0, 10.5) #for at ikke legend skal dekke over de fine strekene våre
    axs[row].set_ylabel(cmap_names[row]) # for å vise hvilket cmap som gir hvilke farger
plt.suptitle('Du kan ikke lage dette i Excel')
plt.savefig('plots/plot5')
plt.show()
plt.close(fig)

#Til slutt skal vi iterere over både rader og kolonner mens vi bruker forskjellige cmaps.
fig, axs = plt.subplots(3,2, figsize=(15,8)) #seks plots i tre rader og to kolonner

cmap_names = [['viridis', 'cool'],
              ['plasma', 'spring'],
              ['cividis', 'inferno']]

cmaps = [[cm.get_cmap(name) for name in line]
         for line in cmap_names] #en liste med colormaps, generert fra listen med navn

some_values = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6] #lager noen verdier vi skal bruke
some_other_values = [0.5, 1, 1.5, 2, 2.5, 3] #disse skal vi også bruke
my_values = [some_values, some_other_values] #pakker alle verdiene sammen

for row in range(3): #iterer over rader
    for col in range(2): #iterer over kolonner

        cmap = cmaps[row][col]  # henter riktig colormap fra listen
        plt.sca(axs[row][col]) #nå skal vi jobbe med plottet i axs[row][col]

        for i in range(len(my_values[col])): #iterer over verdiene vi skal plotte for
            my_value = my_values[col][i]

            z_data = z_func(x_akse, my_value) # husk z_func = z_func(x, y) så nå regner vi ut z(x, y) for alle verdiene i x_akse, med y = my_value
            plt.plot(x_akse, z_data, color=cmap(my_value/max(my_values[col])), label = my_value)

        plt.legend(title='my_value')
        plt.xlim(-2, 10.5)
        plt.ylabel(cmap_names[row][col])

plt.suptitle('Holy shit det var mye farger og fine kurver!')
plt.savefig('plots/plot6')
plt.show()
plt.close(fig)