###########################################################################################
## Functions for simulation code in:
## "Dynamics of tree-cavity occupancy support the hypothesis that low use of 
##  avian-excavated cavities is driven by resource suitability"
## D.J. Zavala, K.L. Cockle, M.R. Gomez, C.A. Ferreyra, E.B. Bonaparte, 
## F.G. DiSallo, and G. Ferraz
## DOI: 10.5281/zenodo.14262784
###########################################################################################

###########################################################################################
# F1. Define function that simulates one matrix of cavity states with as many rows as 
#     cavities and ny years 

simCavUse <- function(cvtp,birds,ne,ny) {
  
  # setup matrix to receive cavity states through time
  CVStt <- matrix(0,nrow=length(cvtp),ncol=ny)
  # unpack numbers of birds
  nEXC <- birds[1]
  nSCN <- birds[2]
  # unpack number of cavities
  Nexc <- sum(cvtp=="E")
  Nnat <- sum(cvtp=="N")
  
  # distribute birds through cavities, one year at a time
  for(t in 1:(ny)){
    # find out how many excavators are going to take non-excavated cavities
    nEXCn <- sum(rbinom(nEXC,1,ne*Nnat / nEXC)) # number of EXC in trans
    if(nEXCn>0) {
      CVStt[sample(which(cvtp=="N"),nEXCn), t] <- 1 # 1 for occupied by EXC 
    }
    # all but nEXCn excavators take up unused excavated cavities (they actually
    # excavate them)
    CVStt[sample(which(cvtp=="E"),(nEXC-nEXCn)), t] <- 1
    
    # next distribute all SCN birds through the available cavities regardless of 
    # cavity type
    CVStt[sample(which(CVStt[,t]==0),nSCN), t] <- 2 # 2 for occupied by SCN
  }
  
  if(any(rowSums(CVStt)==0)) {
    cvtp <- cvtp[-which(rowSums(CVStt)==0)]
    CVStt <- CVStt[-which(rowSums(CVStt)==0),]
  }
  
  return(list(STATES=CVStt,types=cvtp))
  
}


###########################################################################################
# F2. Define function that simulates cavity loss

simCavLoss <- function(CVUse, cvtp, eloss, nloss, ny) {
  
  ## Insert NA before first use
  # find out the year of first observation for each cavity
  yfo <- apply(CVUse,1,function(x) min(which(x>0))) # yr of fist use
  ilo <- which(yfo>1) # row indices of cavities not used in the first year
  yfo <- yfo[ilo]     # year of first use for those not used in frist year
  # loop through cavities with late first obs and fill with NA accordingly
  for(i in 1:length(yfo)){
    CVUse[ilo[i],1:(yfo[i]-1)] <- NA
  }
  
  # Break down cavity use matrix in two parts, for excavated and non-excavated cavities
  eStt <- CVUse[which(cvtp=="E"),] # Excavated
  nStt <- CVUse[which(cvtp=="N"),] # non-Excavated
  
  # Set up empty matrices of cavity loss
  eLss <- matrix(0,nrow=dim(eStt)[1],ncol=ny)
  nLss <- matrix(0,nrow=dim(nStt)[1],ncol=ny)
  
  # Insert losses in _Lss matrices according to the three possible non-loss states
  for(i in 1:3) { # loop through current states 0, 1, 2
    
    # handle excavated cavities
    ecs <- which(eStt==(i-1)) # get 1-D indices of curr state in Excavated cavs
    elost <- rbinom(length(ecs),size=1,prob=eloss[i]) # toss coin to see who's lost
    ielost <- ecs[which(elost==1)]              # find out indices of lost cavities
    ielost2d <- arrayInd(ielost,.dim=dim(eStt)) # get corresponding 2-D indices
    # advance col index by one to indicate first year that cavity is completely no more
    ielost2d[,2] <- ielost2d[,2]+1 # for excavated cavities
    # get rid of those losses that would happen after the end of the data period 
    if(any(ielost2d[,2]>dim(eStt)[2])) {
      ielost2d <- ielost2d[-which(ielost2d[,2]>dim(eStt)[2]),]
    }
    # convert 2-D indices back to 1-D
    ielost <- (ielost2d[,2]-1) * dim(eStt)[1] + ielost2d[,1]
    # register losses in the Lss matrix
    eLss[ielost] <- 1
    
    # handle non-excavated cavities
    ncs <- which(nStt==(i-1)) # get 1-D indices of curr state in non-Exc cavs
    nlost <- rbinom(length(ncs),size=1,prob=nloss[i]) # toss coin to see who's lost
    inlost <- ncs[which(nlost==1)]              # find out indices of lost cavities
    inlost2d <- arrayInd(inlost,.dim=dim(nStt)) # get corresponding 2-D indices
    # advance col index by one to indicate first year that cavity is completely no more
    inlost2d[,2] <- inlost2d[,2]+1 # same for non-excavated cavities
    # get rid of those losses that would happen after the end of the data period
    if(any(inlost2d[,2]>dim(nStt)[2])) {
      inlost2d <- inlost2d[-which(inlost2d[,2]>dim(nStt)[2]),]
      if(length(inlost2d)<3) { 
        if(length(inlost2d==2)) {inlost2d <- t(as.matrix(inlost2d))}
        if(length(inlost2d==0)) next
      }
    }
    # convert 2-D indices back to 1-D
    inlost <- (inlost2d[,2]-1) * dim(nStt)[1] + inlost2d[,1] # non-excavated
    # register losses in the Lss matrix
    nLss[inlost] <- 1
    
  }
  # bind the two loss matrices
  CVLss <- rbind(eLss,nLss)
  
  ## Fill with 3 from first loss onward 
  # find out the year of loss for each cavity that is lost
  yloss <- apply(CVLss,1,function(x) ifelse(any(x==1),min(which(x==1)),0))
  # get indices of cavities with year of loss greater than zero
  iyloss <- which(yloss>0)
  # keep only the years of loss that are greater than zero
  yloss <- yloss[iyloss]
  # loop through cavities lost and fill in 3's to the end of the row
  for(i in 1:length(yloss)){
    CVUse[iyloss[i],yloss[i]:dim(CVUse)[2]] <- 3
  }
  
  ## End
  return(CVUse)
  
}


###########################################################################################
# F3. Define function that computes transition frequencies based on matrix of 
#     cavity use

getTrans <- function(CVUse,cvtp) {
  
  ## first, setup empty transition matrices for both types of cavities
  nTM <- eTM <- matrix(NA,ncol=4,nrow=4)
  eStt <- CVUse[which(cvtp=="E"),]
  nStt <- CVUse[which(cvtp=="N"),]
  
  ## loop through states and compute the frequency of transitions to 
  ## the other states
  for(s in 1:4) {
    
    ise <- which(eStt==(s-1))                        # get 1-D indices of current state
    isn <- which(nStt==(s-1))
    ise2d <- arrayInd(ise,.dim=dim(eStt))            # convert to 2-d indices
    isn2d <- arrayInd(isn,.dim=dim(nStt))
    ise <- ise[-which(ise2d[,2]==dim(eStt)[2])]      # remove terminal start from 1d
    isn <- isn[-which(isn2d[,2]==dim(nStt)[2])]
    # remove terminal start from 2-d indices
    if(any(ise2d[,2]==dim(eStt)[2])){ ise2d <- ise2d[-which(ise2d[,2]==dim(eStt)[2]),] }
    if(any(isn2d[,2]==dim(nStt)[2])){ isn2d <- isn2d[-which(isn2d[,2]==dim(nStt)[2]),] }
    
    ieto2d <- cbind(ise2d[,1],1+ise2d[,2])
    into2d <- cbind(isn2d[,1],1+isn2d[,2])
    ieto <- (ieto2d[,2]-1)*dim(eStt)[1] + ieto2d[,1]
    into <- (into2d[,2]-1)*dim(nStt)[1] + into2d[,1]
    eto <- eStt[ieto]
    nto <- nStt[into]
    
    eTM[s,] <- c(sum(eto==0),sum(eto==1),sum(eto==2),sum(eto==3)) / length(eto)
    nTM[s,] <- c(sum(nto==0),sum(nto==1),sum(nto==2),sum(nto==3)) / length(nto)
    
  }
  
  return(list(eTM,nTM))
}