#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Last update, 16/01/2018
Continuous Replicator Dynamic in Asexual model
check the stability
@author: Shota Shibasaki
"""

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

a1=1
e1=0.0001
a2=0.5
e2=0.0001
h=0.5
e0=0.000001

def Matrix(a1,b1,a2,b2,theta):
    """
   Defining patoff matrixes
   A1: payoff matrix for fruiting body formation
   A2: payoff matrix for macrocyst formation
   A: patoff matrix generated by the mixture of A1 and A2.
      see Eq [2] in the main text for more detail.
   Parameters
   a2: benefit of cmutual ooperation in macrocyst formation
   b2: benefit of exploitation for defectors in macrocyst formation
   theta: prob. of macrocyst fromation
      
   """
    A1=np.array([[a1,h*a1,0],\
                 [e0,e0,e0],\
                 [b1,h*e1,e1]])
    
    A2=np.array([[a2,0,a2],\
                 [b2,e2,b2],\
                 [a2,0,a2]])
    A=(1-theta)*A1+(theta)*A2
    return A

def func(x, t, a1,b1,a2,b2,theta):
    """
    Continuous replicator dynamics
    x: array of the frequency of the three strategy
    t: time
    a2, b2, theta: defined in function Matrix
    """
    #x is array and A is matrix
    A=Matrix(a1,b1,a2,b2,theta)
    phi=np.dot(x,A)
    phi=np.dot(phi,x.T)#mean fitness in the population
    #calculate replicator dynamics
    f0=np.dot(A[0,:],x.T)-phi
    f1=np.dot(A[1,:],x.T)-phi
    f2=np.dot(A[2,:],x.T)-phi     
    f=np.array([x[0]*f0,x[1]*f1,x[2]*f2])
    return(f)

def main():    
    # condition of the linear stability in the discrete replicator
    THETA=np.array([])#theta
    DEL1=np.array([])#b1-a1
    DEL2=np.array([])#b2-a2
    step=26#number of points in each coordinate axe
    the=np.linspace(0,1,step)
    del1=np.linspace(0,1,step)
    del2=np.linspace(0,1,step)
    for t in range(np.size(the)):
        theta=the[t]
        #print(theta)
        for d1 in range(np.size(del1)):
            b1=a1+del1[d1]
            #print(b1)
            for d2 in range(np.size(del2)):
                b2=a2+del2[d2]
                #check there is a global attractor in replicator dynamics
                #print(b2)
                #original payoff matrix CC,DC,CD
                A=Matrix(a1,b1,a2,b2,theta)
                #arranged payoff matrix for cntinuous replicator equation
                Ac=np.array([[A[0,0]-A[0,0],A[0,1]-A[1,1],A[0,2]-A[2,2]],\
                           [A[1,0]-A[0,0],A[1,1]-A[1,1],A[1,2]-A[2,2]],\
                           [A[2,0]-A[0,0],A[2,1]-A[1,1],A[2,2]-A[2,2]]])
                odds=(1-theta)/theta
                if (b2-a2)/(a1-e1)<odds and odds<e2/(h*e1-e0):
                 #condition for RSP game   
                    if np.linalg.det(Ac)>0:
                    #unique global attractor in continuous replicator dynamics
                        #analytical solution
                        p1 = A[2,1]*A[0,2]+A[1,2]*A[0,1]+A[2,1]*A[1,2]
                        p2 = A[1,0]*A[0,2]+A[1,2]*A[2,0]+A[0,2]*A[2,0]
                        p3 = A[2,1]*A[1,0]+A[0,1]*A[2,0]+A[0,1]*A[1,0]
                        pa=np.array([p1,p2,p3])
                        p=pa/sum(pa)
                        #calculate the linearized discrete replicator
                        Ave=(p.dot(A)).dot(p.T)
                        ev1=np.array([1,0,0]).T
                        ev2=np.array([0,1,0]).T
                        ev3=np.array([0,0,1]).T
                        L=np.zeros([3,3])
                        L[0,0]=1+p[0]*(A[0,0]-(p.dot(A)).dot(ev1))/Ave
                        L[0,1]=0+p[0]*(A[0,1]-(p.dot(A)).dot(ev2))/Ave
                        L[0,2]=0+p[0]*(A[0,2]-(p.dot(A)).dot(ev3))/Ave
                        L[1,0]=0+p[1]*(A[1,0]-(p.dot(A)).dot(ev1))/Ave
                        L[1,1]=1+p[1]*(A[1,1]-(p.dot(A)).dot(ev2))/Ave
                        L[1,2]=0+p[1]*(A[1,2]-(p.dot(A)).dot(ev3))/Ave
                        L[2,0]=0+p[2]*(A[2,0]-(p.dot(A)).dot(ev1))/Ave
                        L[2,1]=0+p[2]*(A[2,1]-(p.dot(A)).dot(ev2))/Ave
                        L[2,2]=1+p[2]*(A[2,2]-(p.dot(A)).dot(ev3))/Ave
                        #conditions where |eigenvalues of L| < 1
                        (eig,v)=np.linalg.eig(L)
                        meig=abs(eig)
                        if max(meig)<1:
                            THETA=np.append(THETA,theta)
                            DEL1=np.append(DEL1,del1[d1])
                            DEL2=np.append(DEL2,del2[d2])
    #3D scatter plot
    fig=plt.figure(figsize=(6,6))
    ax  = Axes3D(fig)
    ax.plot(DEL1, DEL2, THETA, "o", color="b", ms= 8, mew=0.5)
    ax.set_xlim(0,1.0)
    ax.set_ylim(0,1.0)
    ax.set_zlim(0,1.0)
    ax.invert_xaxis()
    ax.invert_yaxis()
    #ax.grid(False)
    ax.patch.set_facecolor('white')
    ax.set_xlabel('$b_1-a_1$',fontsize=12)
    ax.set_ylabel('$b_2-a_2$',fontsize=12)
    ax.set_zlabel(r'$\theta$',fontsize=12) 
    #plt.savefig("CondDisRep_asexual.pdf")
    #add the boundary of conditions for continuous replicator dyanmics
    z1=a1/(a1+del2)
    z2=(h*e1-e0)/(h*e1-e0+e2+del2*0)
    ax.plot(del2,z1,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z1,zs=1,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z2,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z2,zs=1,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z1,zs=0,zdir='x',color="k",linestyle="--")
    z3=a1/(a1+1+0*del1)
    z4=a1/(a1+0+0*del1)
    z5=(h*e1-e0)/(h*e1-e0+e2+1*+del1*0)
    z6=(h*e1-e0)/(h*e1-e0+e2+0+del1*0)
    ax.plot(del1,z3,zs=1,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z4,zs=0,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z5,zs=0,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z6,zs=1,zdir='y',color="k",linestyle="--")
    zz=np.linspace((h*e1-e0)/(h*e1-e0+e2),1,step)
    y1=zz*0+0
    ax.plot(y1,zz,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(y1,zz,zs=1,zdir='x',color="k",linestyle="--")
    zz=np.linspace((h*e1-e0)/(h*e1-e0+e2),a1/(a1+1),step)
    y1=zz*0+1
    ax.plot(y1,zz,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(y1,zz,zs=1,zdir='x',color="k",linestyle="--")
    plt.savefig("CondRSPandDisc_Merge_asexual.pdf")
    #plt.show()              
    
if __name__ == '__main__':
    main()