# -*- 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 *


clf()
mu=5.

                
l=0.5                 #losses in the transitions between pools (needed for model stabilization) 
#herbivory======================================================
H=0.0              
def dH(H,tau):
    return H*(1-tau)
#photosynthesis==================================================
phi0=10.#.8
GB=1.

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

def Phi(x):
    return phi0*(1-x)

def Al(x,yl,yr,Nw,tau):
    alpha_tau=(alpha+tau)/(1-tau)
    return Phi(x)*Nw*yl/(1+alpha_tau*Nw*(yl+yr))

def Ar(x,yl,yr,Nw,tau):
    alpha_tau=(alpha+tau)/(1-tau)
    return Phi(x)*Nw*yr/(1+alpha_tau*Nw*(yl+yr))


def vectorfield(X,t,Nw,h):
    f=np.array([0.,0.])    
    x=X[0]
    yl=X[1]
    yr=0.
    tau=0.
    A_l=Al(x,yl,yr,Nw,tau)
    p=0.5*np.sqrt(1+4*A_l)-0.5
    f[0]=mu*(p - h)*x
    f[1]= 1-yl + x*( lbar*h*(p+1)-A_l)*GB*mu/Nw
    return f    


#effective absorption functions
#def xi(x,y,Nw):                       
#    alpha=(1-g)/g
#    return Phi(x)/(1+Nw*alpha*y)    

#def xi_mut(x,y,z,Nw,tau):
#    gamma=(1-tau)*g        #nutrient in biomass 
#    alpha_tau=(1-gamma)/gamma  #C/N
#    return Phi(x)/(1+Nw*alpha_tau*(y+z))    

#def P(x,y,Nw):
#    return (-1+np.sqrt(1+4*xi(x,y,Nw)*Nw*y))/2. 
    
#def Pmut(x,y,yr,Nw,tau):
#    return (-1+np.sqrt(1+4*xi_mut(x,y,yr,Nw,tau)*Nw*(y+yr)))/2. 

#def vectorfield(X,t,Nw,h,mu):
#    f=np.array([0.,0.])    
#    x=X[0]
#    y=X[1]
#    p=P(x,y,Nw)
#    f[0]=mu*(p - h)*x
#    f[1]= 1-y - mu*l*(p+1)*((p-h)/l + h )*B*x/Nw        
#    return f    

N=15
nw=np.linspace(0.5,1.2,N)#(np.logspace(-0.1,0.7,N)
hh=np.linspace(0.2,1.2,N)#np.logspace(-2.5,0,N)
#g=0.1#np.linspace(0.05,0.3,N)

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

Dtau=0.01
z0=.02

dW0=np.zeros([N,N])
for i in range(N):
    for j in range(N):
        h=hh[j]#+dH(H,0)/mu     #herbibory on resident
        hmut=hh[j]#+dH(H,tau)/mu  #herbivory on mutant
        Nw=nw[i]        
        
        #dtot=d+d_H #total senescence on resident
        def F(X,t):
            return vectorfield(X,t,Nw,h)
        X=odeint(F,Xinit,t)
        x=X[len(t)-1,0]  #resident equilibrium state variables
        y=X[len(t)-1,1]
        yr=mu*h*Dtau*z0/Nw
        A=Al(x,y,yr,Nw,Dtau)+Ar(x,y,yr,Nw,Dtau)        
        pmut=0.5*np.sqrt(1+4*A)-0.5#Pmut(x,y,yr,Nw,Dtau)#dtot_mut=d+d_Hmut #senecence of mutant
        #dnr=mu(g[j])*dtot_mut*dtau*xinit/D            #mutant's initial nr pool 
        #Abs_mut=absorb_mut(x,nl,dnr,dtau,g[j])  #mutant nutrient absobtion rate
        dW0[i,j]=mu*(pmut-hmut)/Dtau#(kappa*pmut(x,nl,dnr,dtau,g[j])/g[j] - dtot_mut)/dtau  #mutant fitness gardient at tau=0

clf()
level_max=1
nb=10
levels = np.linspace(0,level_max,nb)#[0.,0.01,0.02,0.03,0.04,0.06]
#levels = [0.,0.002,0.004,0.006,0.008,0.01]

cmap=contour(dW0,levels,colors=('k'),
         extend='both',
         extent=(min(hh),max(hh),min(nw),max(nw)) )
clabel(cmap,colors = 'k', fmt = '%3.2f', fontsize=12)
contourf(dW0,levels,extent=(min(hh),max(hh),min(nw),max(nw)) )
ylabel("$\mathcal{N}$",fontsize=16)
#xscale('log')
#yscale('log')
#ylabel("bare soil nutrient content",fontsize=16)
xlabel("$h$",fontsize=16)
title("Invasion fitness of tannins")
#title("herbivory, no symbiotic capacity")
#title("symbiotic capacity, no herbivory")
#savefig("invasibility_H.pdf")
#colorbar()