####  R code for functions to calculate sensitivities for transition rates.
####  Functions for calculating sensitivities
####  All functions output:
####  1) the stationary state distribution
####  2) an array of sensitivities.  the array is indexed [i,j,k] for 
####  the sensitivity of the kth state in the stationary state distribution (e_k) 
####  to changes in the i,jth element of the transition matrix (r_i,j)

####  The first functions deal with no environmental variation and a single matrix R

###  sens calculates Eq. 6 

sens <- function(R){
	out = c()
	L = eigen(R)$values
	w = eigen(R)$vector
	v = solve(t(w))
	w[,1] = w[,1]/sum(w[,1])
	v[,1] = v[,1]/v[1,1]
	d = nrow(R)
	DW = array(0,c(d,d,d))
	for (i in 1:d){
		for (j in 1:d){
			dw = 0
			for (m in 2:d) dw = dw + w[j,1]*v[i,m]*w[,m]/(L[1]-L[m])   
			DW[i,j,] = dw
		}	
	}
	out$sensitivities = DW
	out$SSD = w[,1]
	return(out)
}

### markov.sens calculates Eq. 7 and 8  -  methods include uniform and proportional compensation

markov.sens <- function(R){
	S <- sens(R)$sensitivities
	SSD <- sens(R)$SSD
	k <- dim(R)[1]
	unif <- array(0,rep(k,3))
	prop <- array(0,rep(k,3))
  	if (k == 2){
    		for (a in 1:k){
	    		for (b in 1:k){
		    		unif[a,,b] <- S[a,,b] - (1/(k-1))*((S[-a,,b]))
	    		}
	  	}  
    		prop <- unif
  	}
  	if (k>2){
	  	for (a in 1:k){
	    		for (b in 1:k){
		    		unif[a,,b] <- S[a,,b] - (1/(k-1))*(colSums(S[-a,,b]))
		    		prop[a,,b] <- S[a,,b]-colSums(t((t(R[-a,])/colSums(R[-a,])))*(S[-a,,b]))
	    		}
	  	}
  	}
	out <- list(unif=unif,prop=prop, SSD = SSD)
	return(out)
}

##########################################################
### These functions can be used when temporal variation is included


###  Tsens calculates sensitivities by numerical perturbation
###  R is a k by k by t array where k is the number of states and t the number of time steps to be drawn from
###  The transition matrix at each time step is randomly drawn from the t possible matrices stored in R
###  Each cell of the matrix is perturbed by a value q and after burn time steps values are stored for (iter-burn) time steps and values calculated from these
###  compensation must be imposed when calculating sensitivities to maintain the MArkovian constraint


## only the stationary state distibution e(T)

ssd.T <- function(R, burn = 10000, iter = 100000){
  	t <- dim(R)[3]  ### number of time periods to sample from
  	k <- dim(R)[1]  ### number of occupancy states
  	t.rand <- sample(c(1:t),iter, replace = TRUE)  
  	H <- matrix(0,k,iter)
  	H[,1] <- 1/k  ### starting state distribution - default to uniform
  	for (a in 2:iter) H[,a] <- R[,,t.rand[a]]%*%H[,(a-1)]
  	e.T <- rowMeans(H[,(burn+1):iter])
  	return(e.T)
}

## uniform compensation
 
unif.Tsens <- function(R,q = 0.00001, burn = 10000, iter = 100000){
 	t <- dim(R)[3]  ### number of time periods to sample from
  	k <- dim(R)[1]  ### number of occupancy states
  	t.rand <- sample(c(1:t),iter, replace = TRUE)  
  	H <- matrix(0,k,iter)
  	H[,1] <- 1/k  ### starting state distribution - default to uniform
  	for (a in 2:iter) H[,a] <- R[,,t.rand[a]]%*%H[,(a-1)]
  	e.T <- rowMeans(H[,(burn+1):iter])
  	p.T <- array(0,c(k,k,k))
  	for (i in 1:k){
    		for (j in 1:k){
      			Rtemp <- R
      			for (d in 1:t){
        				Rtemp[i,j,d] <- Rtemp[i,j,d] + q
        				Rtemp[-i,j,d] <- Rtemp[-i,j,d] - (1/(k-1))*q
      			}
      			H <- matrix(0,k,iter)
      			H[,1] <- 1/k  ### starting state distribution - default to uniform
      			for (a in 2:iter) H[,a] <- Rtemp[,,t.rand[a]]%*%H[,(a-1)] 
        			p.T[i,j,] <- (rowMeans(H[,(burn+1):iter])- e.T)/q
    		}
  	}
  	out <- list(ssd = e.T, sens = p.T)
  	return(out)
}
 
#### proportional compensation
  
prop.Tsens <- function(R,q = 0.00001, burn = 10000, iter = 100000){
  	t <- dim(R)[3]  ### number of time periods to sample from
  	k <- dim(R)[1]  ### number of occupancy states
  	t.rand <- sample(c(1:t),iter, replace = TRUE)  
  	H <- matrix(0,k,iter)
  	H[,1] <- 1/k  ### starting state distribution - default to uniform
  	for (a in 2:iter) H[,a] <- R[,,t.rand[a]]%*%H[,(a-1)]
  	e.T <- rowMeans(H[,(burn+1):iter])
  	p.T <- array(0,c(k,k,k))
  	for (i in 1:k){
    	for (j in 1:k){
      		Rtemp <- R
      		for (d in 1:t){
        			Rtemp[i,j,d] <- Rtemp[i,j,d] + q
        			Rtemp[-i,j,d] <- Rtemp[-i,j,d] - Rtemp[-i,j,d]*q/sum(Rtemp[-i,j,d])
      		}
      		H <- matrix(0,k,iter)
      		H[,1] <- 1/k  ### starting state distribution - default to uniform
      		for (a in 2:iter) H[,a] <- Rtemp[,,t.rand[a]]%*%H[,(a-1)] 
        		p.T[i,j,] <- (rowMeans(H[,(burn+1):iter])- e.T)/q
    		}
  	}
  	out <- list(ssd = e.T, sens = p.T)
  	return(out)
}
 


## general perturbation routine for lower level parameters where 
## parameters making up R are perturbed by q to create R2
## example included in the arroyo toad analysis

lower.Tsens <- function(R, R2, q, burn = 10000, iter = 100000){
  	t <- dim(R)[3]  ### number of time periods to sample from
  	k <- dim(R)[1]  ### number of occupancy states
  	t.rand <- sample(c(1:t),iter, replace = TRUE)  
  	H <- matrix(0,k,iter)
  	H[,1] <- 1/k  ### starting state distribution - default to uniform
  	for (a in 2:iter) H[,a] <- R[,,t.rand[a]]%*%H[,(a-1)]
  	e.T <- rowMeans(H[,(burn+1):iter])
  	H2 <- matrix(0,k,iter)
	H2[,1] <- 1/k  ### starting state distribution - default to uniform
 	for (a in 2:iter) H2[,a] <- R2[,,t.rand[a]]%*%H2[,(a-1)]
 	p.T <- rowMeans(H2[,(burn+1):iter])
 	sens <- (p.T-e.T)/q
 	out <- list(sens = sens,e.T=e.T,p.T=p.T)
  	return(out)
}