# coding: utf-8
from networkx import *
from random import random as rnd
from random import shuffle as shf
from random import choice as chc
import matplotlib.pyplot as plt
import math
import numpy as np
import time
from utilities import double_edge_swap_no_triangle

def initNodes(G):
    for i in G:
        G.nodes[i]['p'] = rnd()
        G.nodes[i]['q'] = rnd()
        G.nodes[i]['p0'] = G.nodes[i]['p']
        G.nodes[i]['q0'] = G.nodes[i]['q']
        G.nodes[i]['payoff'] = 0
        G.nodes[i]['total_payoff'] = 0

def reset(G):
    for i in G:
        G.nodes[i]['p'] = G.nodes[i]['p0']
        G.nodes[i]['q'] = G.nodes[i]['q0']
        G.nodes[i]['payoff'] = 0
        G.nodes[i]['total_payoff'] = 0

def  prettyPrint(G):
    nodes = list(G.nodes.keys())
    nodes.sort()
    for i in nodes:
        print(i,": p =",G.nodes[i]['p'],"q =",G.nodes[i]['q'],"payoff =",G.nodes[i]['payoff'],"total payoff =",G.nodes[i]['total_payoff'])

def simulation(G):
    step = 60000
    global printStep
    if printStep==0:
        print("Initial state :")
        print("M =",G.graph['M'])
        prettyPrint(G)
    print('-'*10)
    runs = 10000 #50000
    for i in range(runs):
        if printStep==0:
            print("Generation :",i+1,"\n#######################")
        generation(G)
        printStep = (printStep + 1) % step

def generation(G):
    def bigger(a,b):
        return int(a>=b)
    global printStep
    texts = ["fail","success"]
    proposers = list(G.nodes.keys())
    shf(proposers)
    for i in proposers:
        #subM = 0.0
        #for j in G[i]:
        #    if G.nodes[i]['p'] >= G.nodes[j]['q']:
        #        subM += 1.0
        #subM /= len(G[i])
        subM = np.mean([bigger(G.nodes[i]['p'],G.nodes[j]['q']) for j in G[i]])
        if subM >= G.graph['M']:
            G.nodes[i]['payoff'] += 1 - G.nodes[i]['p']
            for j in G[i]:
                G.nodes[j]['payoff'] += G.nodes[i]['p']/len(G[i])
        if printStep==0:
            print(i,"propose...",texts[int(subM >= G.graph['M'])],"as M' =",subM)
    if printStep==0:
        print("+"*10)
    update(G)
    if printStep==0:
        print("+"*10)
        prettyPrint(G)
        print("#####################")
    clean(G)

def update(G):
    global printStep
    nodes = list(G.nodes.keys())
    shf(nodes)
    beta = 1.0
    epsilon = 0.05
    for i in nodes:
        G.nodes[i]['total_payoff'] += G.nodes[i]['payoff']
        target = chc(list(neighbors(G,i)))
        proba = 1/(1.0+math.e**(-beta*(G.nodes[target]['payoff']-G.nodes[i]['payoff'])))
        disturbP = (rnd() * epsilon * 2) - epsilon
        disturbQ = (rnd() * epsilon * 2) - epsilon
        if rnd() <= proba:
            oldP = G.nodes[i]['p']
            oldQ = G.nodes[i]['q']
            G.nodes[i]['p'] = G.nodes[target]['p'] + disturbP
            G.nodes[i]['p'] = max(0,min(G.nodes[i]['p'],2 - G.nodes[i]['p']))
            G.nodes[i]['q'] = G.nodes[target]['q'] + disturbQ
            G.nodes[i]['q'] = max(0,min(G.nodes[i]['q'],2 - G.nodes[i]['q']))
            if printStep==0:
                print(i,"imitate",target,"( p and q from", oldP,oldQ,"to",G.nodes[i]['p'],G.nodes[i]['q'],")")

def clean(G):
    for i in G:
        G.nodes[i]['payoff'] = 0

def getMeanPMeanQ(G):
    p = 0
    q = 0
    for i in G:
        p += G.nodes[i]['p']
        q += G.nodes[i]['q']
    p /= len(G.nodes)
    q /= len(G.nodes)
    return p,q

start = time.time()
printStep = 1
M = np.array([1.0/6,2.0/6,3.0/6,4.0/6,5.0/6,1])
pMean = np.zeros(6)
qMean = np.zeros(6)
pMean2 = np.zeros(6)
qMean2 = np.zeros(6)
population = 100
RTRG = nx.watts_strogatz_graph(population, 6, 0, seed=None) # n, k, p, seed
for edge in list(RTRG.edges):
    if ((edge[0] + edge[1]) % 2 == 0):
        RTRG.remove_edge(edge[0], edge[1])
HRN = nx.watts_strogatz_graph(population, 6, 0, seed=None)
for edge in list(HRN.edges):
    if ((edge[0] + edge[1]) % 2 == 0):
        HRN.remove_edge(edge[0], edge[1])
num_edges = len(HRN.edges)
r = 0.9
swaps = r * num_edges
double_edge_swap_no_triangle(HRN, nswap=swaps, max_tries=10000)
initNodes(RTRG)
initNodes(HRN)
count = 0.0
total = len(M)*2.0
for i in range(len(M)):
    RTRG.graph["M"] = M[i]
    HRN.graph["M"] = M[i]
    simulation(RTRG)
    count+=1
    print(str(count/total)+"%")
    simulation(HRN)
    count+=1
    print(str(count/total)+"%")
    pMean[i], qMean[i] = getMeanPMeanQ(RTRG)
    pMean2[i], qMean2[i] = getMeanPMeanQ(HRN)
    reset(RTRG)
    reset(HRN)
print("Done in "+str(time.time()-start)+"sec")
plt.subplot(121)
nx.draw(RTRG, with_labels=True)
plt.subplot(122)
nx.draw(HRN, with_labels=True)
plt.subplot(121)
plt.plot(M, pMean)
plt.plot(M, pMean2, "r-")
plt.subplot(122)
plt.plot(M, qMean)
plt.plot(M, qMean2, "r-")
plt.show()