#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import networkx as nx
import random
import math

"""Yields edges between each node and halfk neighbors.    
    halfk: number of edges from each node
"""
def adjacent_edges(nodes, halfk):
    n = len(nodes)
    for i, u in enumerate(nodes):
        for j in range(i+1, i+halfk+1):
            v = nodes[j % n]
            yield u, v
            
"""Makes a ring lattice with n nodes and degree z.    
    Note: this only works correctly if z is even.    
    n: number of nodes
    z: degree of each node
"""           
def make_ring_lattice(N, z):
    G = nx.Graph()
    nodes = range(N)
    G.add_nodes_from(nodes)
    G.add_edges_from(adjacent_edges(nodes, z//2))
    return G

"""
This function performs the rewiring mechanism
in order to construct an homogeneous 
random graph. The parameter f gives the fraction of edges
to be rewired.
"""
def rewire(G,f):
    edges_used=set()
    edges_to_swap=set()
    num_swaps_to_be_made=int((f*G.number_of_edges())/2)
    num_swaps_done_so_far=0
    while num_swaps_done_so_far<num_swaps_to_be_made:
        for _ in range(2):
            choices=set(G.edges())-edges_used-edges_to_swap
            edge=random.choice(list(choices))
            edges_to_swap.add(edge)
        u_0,v_0=edges_to_swap.pop()
        u_1,v_1=edges_to_swap.pop()
        if not u_0==v_1 and not u_1==v_0:
            if not (u_0,v_1) in G.edges() and not (u_1,v_0) in G.edges():
                G.add_edge(u_0,v_1)
                G.add_edge(u_1,v_0)
                G.remove_edge(u_0,v_0)
                G.remove_edge(u_1,v_1)
                edges_used.add((u_0,v_1))
                edges_used.add((u_1,v_0))
                num_swaps_done_so_far+=1
    return G

def initialize_HoSW(N,z,f):
    return rewire(make_ring_lattice(N,z),f) 


def fermi_distribution(fitness_A,fitness_B,beta=0.005):
    return 1/(1+(math.exp((-beta)*(fitness_B-fitness_A))))

def get_individual_fitness(A,graph,S,T,cooperators,defectors):
    fitness=0
    R=1
    P=0
    for neighbor in graph[A].keys():
        if A in cooperators and neighbor in cooperators:
            fitness+=R
        elif A in cooperators and neighbor in defectors:
            fitness+=S
        elif A in defectors and neighbor in cooperators:
            fitness+=T
        elif A in defectors and neighbor in defectors:
            fitness+=P
    return fitness
    
def flip(p):
    return random.random() < p

"""
This function corresponds to the pairwise comparison rule.
One node A is chosen randomly among A's first neighbors. A and B accumulate
total payoffs and the strategy of B replaces that of A with a probability given
by the fermi distribution.
This is the synchronous version
"""
def strategy_update(graph,A,B,cooperators,defectors,ncooperators,ndefectors,S,T):
    if A in cooperators and B in cooperators:
        ncooperators.append(A)
    elif A in defectors and B in defectors:
        ndefectors.append(A)
    else:
        fitness_A=get_individual_fitness(A,graph,S,T,cooperators,defectors)
        fitness_B=get_individual_fitness(B,graph,S,T,cooperators,defectors)
        if flip(fermi_distribution(fitness_A,fitness_B)):
            if B in cooperators and A in defectors:
                ncooperators.append(A)
            else:
                ndefectors.append(A)
        else:
            if A in cooperators:
                ncooperators.append(A)
            else:
                ndefectors.append(A)
    
"""
Same as above but an asynchronous version
"""
def strategy_update2(graph,A,B,cooperators,defectors,S,T):
    if not (A in cooperators and B in cooperators) and not (A in defectors and B in defectors):
        fitness_A=get_individual_fitness(A,graph,S,T,cooperators,defectors)
        fitness_B=get_individual_fitness(B,graph,S,T,cooperators,defectors)
        if flip(fermi_distribution(fitness_A,fitness_B)):
            if B in cooperators and A in defectors:
                defectors.remove(A)
                cooperators.append(A)
            else:
                cooperators.remove(A)
                defectors.append(A)
    return cooperators,defectors    


def rewire_A_to_random_neigh_B(graph,A,B):
    neighbors_B=list(graph[B].keys())
    neighbors_A=list(graph[A].keys())
    neighbors_B.remove(A)
    neighbors_A.remove(B)
    cond = 0
    
    if len(neighbors_B)>1:
        while cond == 0 and len(neighbors_B)>1:
            C=random.choice(neighbors_B)
            if not(C in neighbors_A):
                graph.remove_edge(A,B)
                graph.add_edge(A,C)
                cond=1
            else:
               neighbors_B.remove(C)
    return graph
      
"""
Given an edge with individuals A and B at the extremes, we say that A(B) is satisfied
with the edge if the strategy of B(A) is a cooperator, being dissatisfied otherwise.
If A is satisfied, she will decide to maintain the link. If dissatisfied, then she may
compete with B to rewire the link, rewiring being attempted to a random neighbor of B.
"""
def structural_update(graph,A,B,cooperators,defectors,S,T):
    #If A is a cooperator and B is a cooperator, then both are satisfied and do nothing
    #Else they play the game
    if not (A in cooperators and B in cooperators):
        fitness_A=get_individual_fitness(A,graph,S,T,cooperators,defectors)
        fitness_B=get_individual_fitness(B,graph,S,T,cooperators,defectors)
        if A in defectors and B in defectors:
            if flip(fermi_distribution(fitness_A,fitness_B)):
                graph=rewire_A_to_random_neigh_B(graph,A,B)
            else:
                graph=rewire_A_to_random_neigh_B(graph,B,A)
        elif A in cooperators and B in defectors:
            if flip(fermi_distribution(fitness_A,fitness_B)):
                graph=rewire_A_to_random_neigh_B(graph,A,B)
        elif B in cooperators and A in defectors:
            if flip(fermi_distribution(fitness_A,fitness_B)):
                graph=rewire_A_to_random_neigh_B(graph,B,A)
    
    return graph


"""
Synchronous updating change the strategy update to all individuals (nodes)
"""
def synchronous_updating(graph,cooperators,defectors,S,T):
    N=len(graph.nodes())
    new_cooperators=list()
    new_defectors=list()
    for indiv in range(N):
        neigh_indiv=random.choice(list(graph[indiv].keys()))
        strategy_update(graph,indiv,neigh_indiv,cooperators,defectors,new_cooperators,new_defectors,S,T)
    
    return graph,new_cooperators,new_defectors

"""
Asynchronous updating selects a random individual and, depending on a 
probability (where W is involved) a strategy update or a structural update
take place
"""
def asynchronous_updating(graph,cooperators,defectors,S,T,W):
    indiv=random.choice(list(graph.nodes))
    neigh_indiv=random.choice(list(graph[indiv].keys()))
    
    if(flip(1/(1+W))):        
        cooperators,defectors=strategy_update2(graph,indiv,neigh_indiv,cooperators,defectors,S,T)        
    else:
        graph=structural_update(graph,indiv,neigh_indiv,cooperators,defectors,S,T)
        
    return graph,cooperators,defectors
    
"""
Implements the simulation with the Watson's article approach 
- S,T values of the game (always S>T)
- W   'wave of cooperation' (the higher it's value, the more cooperation prevails)
- N   number of nodes
- z   degree of each node
- f   percentage of rewired nodes after creating the ring lattice
- St  step between relaxations
- NR  number of relaxations
"""
def simulation_relaxations(S,T,W,N,z,f,St,NR):
    graph=initialize_HoSW(N,z,f)
    cooperators=list(random.sample(range(N),N//2))
    defectors=list(set(range(N))-set(cooperators))
     
    fraction_of_cooperators_per_generation=[]
    fraction_of_cooperators_to_print=[]
    mean_of_cooperators_per_relaxation=[]
    
    for relaxations in range(NR):
        for steps in range(St):
            #If W>0, evolution of strategy and structure proceed together 
            #under asynchronous updating  
            if W>0:
                graph,cooperators,defectors=asynchronous_updating(graph,cooperators,defectors,S,T,W)
                #fraction_of_cooperators_per_generation.append(len(cooperators)/N)  
                #a=fraction_of_cooperators_per_generation[-1]
                #print(a)           
            else:
                graph,cooperators,defectors=synchronous_updating(graph,cooperators,defectors,S,T)    
        fraction_of_cooperators_per_generation.append(len(cooperators)/N)
        a=fraction_of_cooperators_per_generation[-1]
        
                
        fraction_of_cooperators_to_print.append(fraction_of_cooperators_per_generation[-1])
        #mean_of_cooperators_per_relaxation.append(sum(fraction_of_cooperators_per_generation[-steps:])/steps)
        
        cooperators=list(random.sample(range(N),N//2))   # ******Here the relaxation take place*****
        defectors=list(set(range(N))-set(cooperators))   # ******Here the relaxation take place*****
        
    if len(cooperators)!=N:
        percentage_of_cooperators=fraction_of_cooperators_per_generation[-1]
    else:
        percentage_of_cooperators=1.0
        
    return percentage_of_cooperators,k_max(graph),fraction_of_cooperators_to_print



# Helper functions

def degrees(graph):
    degs={}
    for node in graph.nodes():
        deg=graph.degree(node)
        if deg not in degs:
            degs[deg]=0
        degs[deg]+=1
    return degs
  

# Return the max degree of a node 
def k_max(graph):
    return max(degrees(graph).keys())