# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------
# -----------------------------------------------------------------------
# This script plots the descendant distributions obtained using the script
# Descendant_distributions.py , published as part of the paper
#
# --------
# Scaling law for the impact of mutant contagion
# by Jonas S. Juul and Steven H. Strogatz
# --------
#
# The script was written by Jonas S. Juul, jonas.juul  at  nbi.ku.dk.
# 
# If you wish to use this script, please consider citing the corresponding 
# paper.
#
# Pease do not distribute this script without the consent of the authors.
# -----------------------------------------------------------------------
# -----------------------------------------------------------------------


import matplotlib.pyplot as plt
import numpy as np 
# Define plot settings
plt.style.use('bmh')
plt.rcParams.update({'font.size': 14})
plt.rc('xtick', labelsize=12) 
plt.rc('ytick', labelsize=12) 
# ---------------------------------------------------------
# DEFINITIONS 
# ---------------------------------------------------------

# Define analytical solution
from scipy.special import beta

# Analytical solution for configuration-model and erdos-renyi graphs
def rand_analytical(k,d) : #Inputs: average degree of the network, k, and d, the number of descendants at which the function should be evaluated..
    return  (k-1)/float(k-2)*beta(float(k-1)/float(k-2)+float(d),2.)
def Smallworld_analytical(p,d) : 
    return (2.*p+1.)/(2.*p)*beta((2*p+1.)/(2.*p)+float(d),2.)
def Complete_analytical(d) :
    return (1/((d+1.)*(d+2.)))

# ---------------------------------------------------------
networks = ['Regular','ER','Small-world','Complete','Facebook','Twitter']

# --------------------------------------------------------
# PLEASE MAKE THE FOLLOWING CHOICES
# --------------------------------------------------------
# Choose network type: k-regular, Erdős-Rényi, Small-world or Complete
network = networks[0]
# Choose value of name_var. name_var corresponds to k, p, p and NAN for the 4 network choices respectively (see Descendant_distributions.py for elaboration)
name_var = 10
# Choose value of N (number of nodes). This is necessary because N is part of file names containing the simulated results
N = 1000
# Choose value of Nexp (number of realizations). This is necessary because Nexp is part of file names containing the simulated results
Nexp = 10
# --------------------------------------------------------
# IMPORT SIMULATION RESULTS
# --------------------------------------------------------
# Define file_name
f = open('Distribution_%s_var%s_N%s_Nexp%s.txt'%(network,name_var,N,Nexp),'r')
descendants_simulation = []
f.readline()
for line in f :
    descendants_simulation.append(float(line))
# --------------------------------------------------------
# ANALYTICAL SOLUTIONS, PLOT
# --------------------------------------------------------
reg_solution = []

for i in range (N) :
    if (network == 'Regular') :
        reg_solution.append(N*rand_analytical(name_var,i))
    elif (network=='ER') :
        reg_solution.append(N*rand_analytical(name_var*N,i))
    elif (network == 'Small-world') :
        reg_solution.append(N*Smallworld_analytical(name_var,i))
    elif (network == 'Complete') :
        reg_solution.append(N*Complete_analytical(i))
print reg_solution
plt.figure(1,figsize=(3.8,3.8))
# Plot simulation results
plt.plot(np.log(np.arange(0,len(descendants_simulation))),np.log(descendants_simulation)+np.log(N),'<',alpha=0.3,markersize=4,label='%s\nvar'%(network)+r'$=%s$'%name_var)#, $N_{\rm exp}=10^3$')
# Plot analytical solution
if (network != 'Facebook' and network != 'Twitter') :
    plt.plot(np.log(np.arange(0,len(reg_solution))),np.log(reg_solution),'k--',label='Analytical')#, $N_{\rm exp}=10^3$')
else : 
    plt.plot(np.log(np.arange(0,len(descendants_simulation))[int(N/20.):N])-np.log(int(N/20.)/2.),np.log(N)*1.5-2*np.log(np.arange(0,len(descendants_simulation))[int(N/20.):N]),'k--',label=r'$\propto d^{-2}$')#, $N_{\rm exp}=10^3$')

# Additional settings for plot
ax = plt.gca() 
ax.set_facecolor('white')
ax.tick_params(top='on',right='on')
plt.legend(loc=3,fontsize=12,frameon=False)
ax.grid(None)
ax.set_xlim([-0.4,10.4])
plt.xlabel(r'$\ln (d)$',fontsize=14)
plt.ylabel(r'$ \ln (NP_{d})$',fontsize=14)
plt.tight_layout()
plt.savefig('Plot_%s_var%s.pdf'%(network,name_var))

print sum(descendants_simulation),sum(reg_solution)

plt.show()