##################################################
# Burning Grass Matrices after Silva et al. 1991 #
##################################################

A=array(0,dim=c(4,4,2))
#Burnt 
A[,,1]=matrix(c(0.08,1.63,2.42,4.4,
                0.21,0.64,0.35,0.16,
                0,0.19,0.43,0.24,
                0,0.03,0.23,0.48),4,byrow=TRUE)
#Unburnt
A[,,2]=matrix(c(0,0.706,0.391,3.59,
                    0.018,0.158,0.136,0.093,
                    0,0.08,0.07,0.21,
                    0,0,0.01,0.07),4,byrow=TRUE)
                    
#Sample P matrix

P=matrix(c(0.2,0.8,0.2,0.8),nrow=2,byrow=TRUE)

######################################
# Estimating Stochastic Growth Rates #
######################################

# INPUT VARIABLES are:
# demography: array A in which A[,,i] corresponds to the projection matrix for environmental state i
# environmental dynamics (markovian): matrix P for which P[i,j] is the probability of going from environmental state i to environmental state j
# T total run time for estimating the invasion speed, etc.
# seed the random seed used for the estimate

# RETURNS:stochastic growth rate, std. dev in growth rate

Lyap=function(A,P,T=10000,seed=1){
  # set up 
  no.states=dim(P)[1]; # number of environmental states
  k=dim(A[,,1])[1]; # number of stages
  ## ESTIMATING GROWTH RATES
  v<-matrix(1,k,1)/k; # initial vector
  state<-1; # initial state
  lyap<-numeric(T);# hold estimates for stochastic growth rate
  set.seed(seed);
  for(t in 1:T){
    temp<-(A[,,state])%*%v ;
    lyap[t]<-log(sum(temp)/sum(v)); # log growth over one time step
    v<-temp/sum(temp); # update v
    state=sample(c(1:no.states),1,prob=P[state,]); # update state
  }
  growth=mean(lyap); # estimate for Lyapunov exponent		
  # variance in stochastic growth rate. To estimate this variance, use the observation that log(lambda_S)-sigma^2/2=log(lambda_average) where lambda_average is the expected growth rate of the population size and sigma^2 is the desired variaance. 
  # use the "mega-matrix" approach to compute lambda_average 
  # create a block diagonal matrix with diagonal blooks A[,,i]
  F=matrix(0,k*no.states,k*no.states)
  for(i in 1:no.states){
    lower=(i-1)*k+1
    upper=i*k
    F[lower:upper,lower:upper]=A[,,i]
  }
  # the megamatrix B and lambda.average
  B=F%*%(P%x%diag(k))
  lambda.average=max(abs(eigen(B)$values))
  # the estimate for sigma
  growth.sd=sqrt(2*(abs(log(lambda.average)-growth)))  
  return(c(growth,growth.sd))}


################################################################################
# Estimating Stable Stage Distribution and Reproductive Values *UNTESTED CODE* #
################################################################################
# RETURNS: stable state distribution, reproductive values
# INPUT: same as Lyapunov
# OUTPUT: boxplots of stable state distribution, reproductive values
LyapDist=function(A,P,T=10000,seed=1,plotit=TRUE){  
  # set up 
  no.states=dim(P)[1]; # number of environmental states
  k=dim(A[,,1])[1]; # number of stages
  ## ESTIMATING STABLE STAGE
  v<-matrix(1,k,1)/k; # initial vector
  state<-1; # initial state
  states=numeric(T) # keeps track of all the states for the reproductive calculation
  states[1]=state
  vs=matrix(0,T,k) # holds all the vs!!
  vs[1,]=v
  lyap<-numeric(T);# hold estimates for stochastic growth rate
  set.seed(seed);
  for(t in 1:(T-1)){
    temp<-(A[,,state])%*%v ;
    lyap[t]=log(sum(temp)/sum(v)); # log growth over one time step
    v=temp/sum(temp); # update v
    state=sample(c(1:no.states),1,prob=P[state,]); # update state
    vs[t+1,]=v  
    states[t+1]=state
  } 
  ws=matrix(0,T,k)
  ws[T,]=matrix(1,1,k)/k
  for(t in (T-1):1){
    w=t(A[,,states[t]])%*%ws[t+1,]
    ws[t,]=w/sum(w)
  } 
  if(plotit){par(cex.lab=1.5,cex.axis=1.5,mfrow=c(1,2))
  	boxplot(vs,col="green",ylab="fraction",xlab="stages")
  	boxplot(ws,col="yellow",ylab="reproductive values",xlab="stages")}
  A=array(0,dim=c(T,k,2))
  A[,,1]=vs
  A[,,2]=ws		 
  return(c())}

#########################
# acorn woodpecker data #
#########################
js=c(0.56, 0.64, 0.3, 0.4, 0, 0.38, 0.18, 0.25, 0.44); # juvenile survivorship
as=c(0.53, 0.68, 0.71, 0.38, 0.54, 0.69, 0.66, 0.49, 0.61); # adult survivorship
f=c(3.38, 1.27, 2.77, 2.17, 0.05, 4.0, 2.37, 0.5, 1.6)/2; # fecundity


#####################################################################
# a spatial version of the stacey-taper model for acorn woodpeckers.#
#####################################################################
# INPUTS
# k number of patches
# rho temporal correlation
# sync spatial correlation
# d fraction dispersing
# T length of run to estimate exponent
# a strength of intraspecific competition within a patch

# OUTPUTS 
# stochastic growth rate and mean population size

woody=function(k=30,rho=0,sync=0,T=10000,d=0.9,a=0.01){
  js=c(0.56, 0.64, 0.3, 0.4, 0, 0.38, 0.18, 0.25, 0.44);
  as=c(0.53, 0.68, 0.71, 0.38, 0.54, 0.69, 0.66, 0.49, 0.61); 
  f=c(3.38, 1.27, 2.77, 2.17, 0.05, 4.0, 2.37, 0.5, 1.6)/2;
  lambda=f*js+as
  k=30 # patches
  D=(1-d)*diag(k)+d*matrix(1,k,k)/k # dispersal matrix
  # simulation code
  totals=numeric(T)
  lyap=numeric(T)
  N=matrix(1,k,1)
  v=N
  totals[1]=sum(N)
  lyap[1]=0
  temp=sample(lambda,k,replace=TRUE)*(1-sync)+sync*sample(lambda,1)-mean(lambda)
  for (t in 1:(T-1)){
    temp2=sample(lambda,k,replace=TRUE)*(1-sync)+sync*sample(lambda,1)-mean(lambda)
    temp=rho*temp+sqrt(1-rho^2)*temp2
    r=temp+mean(lambda)
    N=D%*%(r*N/(1+a*N))
    totals[t+1]=sum(N) 
    new=D%*%(r*v)
    lyap[t+1]=log(new[1]/v[1])
    v=new/new[1]
  }
  popsize=mean(totals[(T/2):T])
  growth=mean(lyap[(T/2):T])
  return(c(growth,popsize))
}