# Import libraries
import numpy as np # standard library for numerical operations in python
import matplotlib.pyplot as plt # standard library for plotting in python

# Ehrenfest urn model draw (s*=0.5, Pfb=0)
def Ehrenfest(randNum, A, B):

    s = B / N # calculate fraction of B opinions

    if (randNumb > s): # draw is A

        B = B + 1 # put back one B by increasing number of B
        A = A - 1 # must decrease number of As

    else: # draw is B

        A = A + 1 # put back one A by increasing number of A
        B = B - 1 # must decrease number of Bs

    if (A > N): # check if A is out of bounds, above 1
        A = N
        B = 0

    if (A < 0): # check if A is out of bounds, below 0
        A = 0
        B = N

    return [A, B] # return updated values of A and B opinions

# Eigen urn model draw (s*=0 or s*=1, Pfb=1)
def Eigen(randNum, A, B):

    s = B / N # calculate fraction of B opinions

    if (randNumb > s): # draw is A

        A = A + 1 # add one A by increasing number of A
        B = B - 1 # must decrease number of Bs

    else: # draw is B

        B = B + 1 # add one B by increasing number of B
        A = A - 1 # must decrease number of As

    if (A > N): # check if A is out of bounds, above 1
        A = N
        B = 0

    if (A < 0): # check if A is out of bounds, below 0
        A = 0
        B = N

    return [A, B] # return updated values of A and B opinions

# Hamann urn model draw (s*=0.23 or s*=0.77, Pfb=0.75sin(pi*s))
def Hamann(randNum, A, B):

    s = B / N # calculate fraction of B opinions

    dB = 4 * (0.75 * np.sin(3.14 * s) - 0.5) * (s - 0.5)

    if (dB > 0):
        Pfb = -1 # negative feedback for B
    else:
        Pfb = 1 # positive feedback for B

    if (randNumb > s): # draw is A

        B = B + Pfb # depends on feedback
        A = A - Pfb # depends on feedback

    else: # draw is B

        A = A + Pfb # depends on feedback
        B = B - Pfb # depends on feedback

    if (A > N): # check if A is out of bounds, above 1
        A = N
        B = 0

    if (A < 0): # check if A is out of bounds, below 0
        A = 0
        B = N

    return [A, B] # return updated values of A and B opinions

Ainit = 49 # number of agents with opinion A
Binit = 51 # number of agents with opinion B
N = Ainit + Binit # total number of agents

numIterations = 1000 # number of iterations/draws from the urn
numSim = 100 # number of simulations
sArray = np.zeros(numIterations) # array for keeping results of s = B / N

# Loop over all simulations
for sim in range(numSim):

    A = Ainit # initializes the A agents at start of simulation
    B = Binit # initializes the B agents at start of simulation

    # Loop over all draws
    for iteration in range(numIterations):

        randNumb = np.random.rand() # random uniform number between 0 and 1

        # Ehrenfest s*=0.5, Pfb = 0
        #[A, B] = Ehrenfest(randNumb, A, B)

        # Eigen s*=0 or s*=1, Pfb = 1
        #[A, B] = Eigen(randNumb, A, B)

        # Hamann s*=0.23 or s*=0.77, Pfb = 0.75sin(pi*s)
        [A, B] = Hamann(randNumb, A, B)

        # keep the agent opinions fro plotting later on
        sArray[iteration] = B / N # fraction of agents with opinion B

    plt.plot(sArray)  # plot the urn draw data

# Calculate the urn dynamics using the analytic urn model
sArrayUrnModel = np.zeros(numIterations) # array for keeping data
sUrn = Binit / N # the fraction of agents with opinion B

for iteration in range(numIterations): # loop over all iteraruin

    #dsUrn = (4 * (0 - 0.5) * (sUrn - 0.5)) / N # Ehrenfest Pfb=0, s*=0.5
    #dsUrn = (4 * (1 - 0.5) * (sUrn - 0.5)) / N # Eigen Pfb=1, s*=0 or s*=1
    dsUrn = (4 * (0.75 * np.sin(3.14 * sUrn) - 0.5) * (sUrn - 0.5)) / N  # Hamann Pfb=0.75sin(pi*s), s*=0.23 or s*=0.77

    sUrn = sUrn + dsUrn # uopdate fraction of Bs
    if (sUrn > 1): sUrn = 1 # check if out of bounds, above 1
    if (sUrn < 0): sUrn = 0 # check if out of bounds, below 0
    sArrayUrnModel[iteration] = sUrn # keeep data

plt.plot(sArrayUrnModel, 'black')  # plot the analytical urn solution
plt.ylim(0,1)  #set x-axis scale
plt.xlabel("Number of iterations") # set x-axis label
plt.ylabel("Fraction s=B/N") # set y-axis label
plt.show()