# -*- coding: utf-8 -*-
"""
Created on Sat Jun  4 18:56:44 2016

@author: jfarnoldi
"""

import numpy as np
from scipy.integrate import odeint
from pylab import *
from matplotlib.colors import Normalize

### Adimentionalized parameters
mu=5. #relative growth rate (with respect to N leaching rate)
h=.6 #relative senecence (with respect to growth rate)
Nw=.6 #effective richness of soil
GB=1. #
Lr=.1 #relative Nr leaching (with respect to N leaching rate)
phi0=10. #relative photosythesis rate (with respect to growth rate)

h=0.5
Nw=1.


###additional parameters
l=0.5  #losses between compartments
lbar=1-l 

#stockiometry
alpha=0.5
gamma=1/(1+alpha)
Om=1.

##Functions==========================================
#photsynthese
def phi(x):
    return phi0*(1-x)#/(1.+B)

def Alpha(tau):
    return (alpha+tau)/(1.-tau)

def bl(e):
    Delta=1. #Delta>1 makes symbiotic associations possible without tannins
    return (1+Delta*e)*(1-e)

def br(e):
    br_max=2. #max effective absorption effeciency
    return 4*br_max*e*(1-e)

def Al(x,yl,yr,tau,e):
    return phi(x)*Nw*bl(e)*yl/(1+Alpha(tau)*Nw*(bl(e)*yl+br(e)*yr))

def Ar(x,yl,yr,tau,e):
    return phi(x)*Nw*br(e)*yr/(1+Alpha(tau)*Nw*(bl(e)*yl+br(e)*yr))

def Atot(x,yl,yr,tau,e):
    return Al(x,yl,yr,tau,e)+Ar(x,yl,yr,tau,e)


def vectorfield(X,t,tau,e):
    f=np.array([0.,0.,0.])    
    x=X[0]
    yl=X[1]
    yr=X[2]
    A_l=Al(x,yl,yr,tau,e)
    A_r=Ar(x,yl,yr,tau,e)
    A=A_l+A_r
    p=0.5*np.sqrt(1+4*A)-0.5
    xeff=(1-tau)*x    
    Om_eff=Om*tau/gamma/(1-tau)
    
    f[0]=mu*(p - h)*x
    f[1]= 1-yl + xeff*( lbar*h*max(0,p+1-Om_eff)-A_l)*GB*mu/Nw
    f[2]=xeff*( h*lbar*min(Om_eff,p+1)- A_r )*mu*GB/Nw - Lr*yr        
    return f    

Teco=200
res=1
Xinit=np.array([0.01,1.,0.])
t=linspace(0,Teco,res*Teco)

De=0.01
Dtau=0.01
z0=0.08


N=15
#D=np.linspace(0.1,0.9,N)
#Nw=np.linspace(0.01,0.1,N)
E=np.linspace(0.,0.5,N)
Tau=np.linspace(0.,0.5,N)
dW_tau=np.zeros([len(E),len(Tau)])
dW_e=np.zeros([len(E),len(Tau)])


#d_H=dH(H,0.)     #herbibory on resident
#d_Hmut=dH(H,tau) #herbivory on mutant
for i in range(len(E)):
    for j in range(len(Tau)):
        e=E[i]
        tau=Tau[j]
        def F(X,t):
            return vectorfield(X,t,tau,e)
        X=odeint(F,Xinit,t)
        x=X[len(t)-1,0]  #resident equilibrium state variables
        if x<=0.0001:
            print "extinction"
            yl=1.
        else:
            yl=X[len(t)-1,1]
            yr=X[len(t)-1,2]
        
        DW_tau=(Atot(x,yl,yr+z0*Dtau,tau+Dtau,e)-Atot(x,yl,yr,tau,e))/Dtau
        DW_e=(Atot(x,yl,yr,tau,e+De)-Atot(x,yl,yr,tau,e))/De
        dW_tau[i,j]=np.sign(DW_tau)*min(1.,abs(DW_tau))
        dW_e[i,j]=np.sign(DW_e)*min(1.,abs(DW_e))
clf()
levels = [0.,0.01,0.03,0.04,0.06]
levels_0 = [0]#,0.2,0.4,1,100]
levels_tau = [-1,0.,0.05,0.1,0.2,1]
levels_h = [0.,0.04,0.08,0.16]

X=np.ones(dW_tau.shape).T*E
Y=np.ones(dW_tau.shape)*Tau



def lol(x):
    return 0.6*np.sign(x)*np.log(1+abs(x))#x#10*np.clip(x,-0.01,0.01)#-0.1*np.sign(x)*np.log10(np.abs(x))#

clf()
#CS1=contourf(dW_tau.T, levels_tau, extent=(min(E),max(E),min(Tau),max(Tau)),colors=('white','0.8','0.7','0.6','0.5','0.4'))
plt.quiver(X.T.ravel(),Y.ravel(),lol(dW_e).ravel(),lol(dW_tau).ravel(),width=0.005,scale=1 / 0.15,alpha=0.5)
#CS2=contourf(dW_e.T, levels_tau, extent=(min(E),max(E),min(Tau),max(Tau)),colors=('white','0.8','0.7','0.6','0.5','0.4'))

#CS.cmap.set_under('white')
#CS.cmap.set_over('.1')

cmap=contour(dW_tau.T,levels_0,colors=('k'),alpha=0.8,linewidths=(3,),
         extend='both',
         extent=(min(E),max(E),min(Tau),max(Tau)) )
#clabel(cmap, levels_0,colors = 'k', fmt = '%3.1f', fontsize=12)

cmap=contour(dW_e.T,levels_0,colors=('k'),alpha=0.5,linewidths=(3,),
         extend='both',
         extent=(min(E),max(E),min(Tau),max(Tau)) )
#clabel(cmap, levels_tau,colors = 'k', fmt = '%3.1f', fontsize=12)


xlabel("symbiotic capacity",fontsize=16)
ylabel("tannin content",fontsize=16)
title("coevolution of tannins and symbiotic capacity",fontsize=16)
#title("with herbivory")
#title("symbiotic capacity, no herbivory")
#savefig("invasibility_e.pdf")
#colorbar()

#example==============================================
#levels = [-1.5, -1, -0.5, 0, 0.5, 1]
#CS3 = plt.contourf(X, Y, Z, levels,
#                   colors=('r', 'g', 'b'),
#                   origin=origin,
#                   extend='both')
# Our data range extends outside the range of levels; make
# data below the lowest contour level yellow, and above the
# highest level cyan:
#CS3.cmap.set_under('yellow')
#CS3.cmap.set_over('cyan')

#CS4 = plt.contour(X, Y, Z, levels,
#                  colors=('k',),
#                  linewidths=(3,),
#                  origin=origin)
#plt.title('Listed colors (3 masked regions)')
#plt.clabel(CS4, fmt='%2.1f', colors='w', fontsize=14)