######################################################################
#Python Programm "Spatial Games"
######################################################################
import networkx as nx
import matplotlib.pyplot as plt
from random import randint, uniform
import numpy as np
import matplotlib
import matplotlib.gridspec as gridspec
from matplotlib.ticker import NullFormatter
from scipy.integrate import solve_ivp
import os

# Erstelle den Ausgabeordner
os.makedirs("./output", exist_ok=True)

#Festlegung der Auszahlungsmatrix des symmetrischen (2x2)-Spiels
# Anti-Koordinationsspiel
a = 1
b = 5
c = 4
d = 3

# Vergleich mit der analytischen Loesung
# Definition der Funktion g
def g(x,a,b,c,d):
    g = ((a-c)*(x-x*x) + (b-d)*(1-2*x+x*x))*x
    return g

# Definition der DGL
def DGL(t, x):
    dxdt = g(x,a,b,c,d)
    return dxdt

# Definition der symmetrischen Spielmatrix (in reinen (sA,sB=0,1) und gemischten Strategien)
def Dollar(sA,sB,a,b,c,d):
    Dollar = a*sA*sB + b*sA*(1-sB) + c*(1-sA)*sB + d*(1-sA)*(1-sB)
    return Dollar

#Erzeugung des regulaeren Netzwerkes zum Anfangszeitpunkt t=0 
#Jeder Knoten ist mit seinem naechsten Nachbarn auf einem 2-D Gitter verbunden
#Festlegung der Anfangsstrategien (alle spielen s=1 (Blau) ausser Knoten 12: s=0 (Rot))
G = nx.Graph()
Nx = 205
Ny = Nx
NKn = Nx*Ny #Anzahl der Knoten
Nit = 30 #Anzahl der Zeit/Iterationsschritte
P = np.zeros([NKn,10])#Erzeugung einer numpy-Matrix zum Speichern der Knoten/Spieler Eigenschaften

k = 0
for i in range(0,Nx):
    for j in range(0,Ny):
        G.add_node(k)
        P[k,0] = k #Knotennummer
        P[k,1] = i #x-Koordinate auf dem 2-D Gitter
        P[k,2] = j #y-Koordinate auf dem 2-D Gitter
# 80 Prozent spielen Strategie s_1
        zufall = uniform(0, 100) #Zufall
        if zufall >= 0.8*100:
            strategy = 0 #entspricht Strategie s_2
        else:
            strategy = 1 #entspricht Strategie s_1

        P[k,3] = strategy #aktuelle Strategie
        P[k,5] = strategy #zukuenftige Strategie
        
        center = i*Ny + j
        upperleft = ((i-1)%Nx)*Ny + (j-1)%Ny
        up = ((i  )%Nx)*Ny + (j-1)%Ny
        upperright = ((i+1)%Nx)*Ny + (j-1)%Ny
        right = ((i+1)%Nx)*Ny + (j  )%Ny
        lowerright = ((i+1)%Nx)*Ny + (j+1)%Ny
        low = ((i  )%Nx)*Ny + (j+1)%Ny
        lowerleft = ((i-1)%Nx)*Ny + (j+1)%Ny
        left = ((i-1)%Nx)*Ny + (j  )%Ny
        G.add_edge(center, upperleft)
        G.add_edge(center, up)
        G.add_edge(center, upperright)
        G.add_edge(center, right)
        G.add_edge(center, lowerright)
        G.add_edge(center, low)
        G.add_edge(center, lowerleft)
        G.add_edge(center, left)
        k = k + 1

#Erzeugung einer Liste zum Speichern der mittleren Knoten/Spieler Eigenschaften
AvStrat = [np.sum(P[:,3])/NKn]

# Analytische Loesung der DGL
Loes = solve_ivp(DGL, [0, Nit], AvStrat, t_eval=np.linspace(0,Nit,500))
Loes_seq = solve_ivp(DGL, [0, Nit], [0.9,0.7,0.5,0.3,0.1], t_eval=np.linspace(0,Nit,500))

#plot settings
params = {
    'figure.figsize'    : [7.5, 10],
    'text.usetex'       : True,
    'legend.fontsize'   : 12
}
matplotlib.rcParams.update(params)

#Plot-Grid
plt.figure(0)
gs = gridspec.GridSpec(2,1, height_ratios=[1,3.0], hspace=0.1)
ax1 = plt.subplot(gs[0])
ax2 = plt.subplot(gs[1])
props = dict(boxstyle='round', facecolor='white', alpha=0.92)

#Achsenbeschriftung ax1
ax1.set_xlim(0,Nit-1)
ax1.set_ylim(0,1)
ax1.set_ylabel(r'$\rm x(t)$')

# Plotten der analytischen Loesung
ax1.plot(Loes.t,Loes.y[0],c="grey", linewidth=0.8, linestyle=':')
for j in range(len([0.9,0.7,0.5,0.3,0.1])):
    ax1.plot(Loes_seq.t,Loes_seq.y[j],c="lightblue", linewidth=0.8, linestyle=':')

# Groesse der Spieler Kaestchen
sgross = np.sqrt(2822400/NKn)
sklein = sgross/8

#Anfang der Schleife der zeitlichen Erzeugung des Netzwerks ( Nit Iterations-Zeit-Endpunkt )
for it in range(1,Nit):
    print(it,"---------------------------------------------------")
    #Jeder Spieler spielt das vorgegebene Spiel mit seinen naechsten Nachbarn, Speichern des gesamten Gewinns/Verlusts in P[:,4]
    for games in list(G.edges):
        P[games[0],4] += Dollar(P[games[0],3],P[games[1],3],a,b,c,d)
        P[games[1],4] += Dollar(P[games[1],3],P[games[0],3],a,b,c,d)
    
    #Vergleich der eigenen Auszahlungswerte mit den Nachbarn und Festlegung der zukuenftigen Strategie in P[:,5]
    for k in range(0,NKn):
        neig = list(G.neighbors(k))
        neig_m = np.max(P[neig,4])
        if neig_m > P[k,4]:
            j_max = np.where(np.abs(P[neig,4] - neig_m) < 1e-9 )[0][0]
            P[k,5] = P[neig[j_max],3]

    #Farbenfestlegung der aktuellen und zukuenftigen Strategienwahl
    Col=[]
    Col_new=[]
    for k in range(0,NKn):
        if P[k,3] == 0:
            Col.append("r")
        if P[k,3] == 1:
            Col.append("b")
        if P[k,5] == 0:
            Col_new.append("r")
        if P[k,5] == 1:
            Col_new.append("b")
        P[k,4] = 0
        P[k,3] = P[k,5]

    #Mittlere Strategienwahl der Spielerpopulation
    AvStrat.append(np.sum(P[:,3])/NKn)

    #Plotten der einzelnen Teildiagramme
    ax1.plot(range(it),AvStrat[0:it], c="black")
    ax2.scatter(P[:,1],P[:,2],s=sgross,c=Col,marker="s",alpha=1,edgecolor='none')
    #Visualisierung der Auszahlungen der zukuenftigen Strategie als kleines Viereck
    ax2.scatter(P[:,1]-0.25,P[:,2]-0.25,s=sklein,c=Col_new,marker="s",alpha=1,edgecolor='none')
    #Achsenbeschriftung ax2
    ax2.set_xlim(-0.5,Nx-0.5)
    ax2.set_ylim(-0.5,Ny-0.5)

    #Speicherung des Bildes als .png und .pdf Datei
    saveFig = f"./output/img-{'{:0>3d}'.format(it)}.png"
    plt.savefig(saveFig, dpi=100,bbox_inches="tight",pad_inches=0.05,format="png")
#    saveFig = "./output/img-" + "{:0>3d}".format(it) + ".pdf"
#    plt.savefig(saveFig,bbox_inches="tight",pad_inches=0.05,format="pdf")
#   Die einzelnen Bilder koennen dann in eine Animation zusammengefuegt werden.
#   z.B. mit folgende Shell-Befehl: ffmpeg -r 2 -i img-%03d.png spatialgame.avi
#   ffmpeg -r 2 -i ./output/img-%03d.png -c:v libx264 -pix_fmt yuv420p ./output/spatialgame.mp4
#   ffmpeg -r 2 -i ./output/img-%03d.png -c:v mpeg4 -q:v 2 ./output/spatialgame.avi
       
    #Loeschen des Bildes
    ax2.clear()
plt.close()
