simpleBranch <- function(s, p, mu = 0, nYears){
    ext0 = 0
    ext <- numeric(nYears + 1)
    ext[1] <- ext0

    for(year in 1:nYears){
        ext[year+1] <- 1 - s * (1 - mu) +
            sum(s *  (1 - mu) * p * (ext[year] ^ c(1:length(p))))
    }
    names(ext) <- paste("Year ",0:nYears)
    return(list(ext = ext, final = ext[nYears+1]))
}

uqStocSimpleBranch <- function(sMean, sPrecision, meanBirth, meanBirthSD,
                                birthDist = 'pois', mu = 0, nYears, nReps,
                                N = 1){
    s = rbeta(n = nReps,
        shape1 = sMean * sPrecision,
        shape2 = (1 - sMean) * sPrecision)
    
    if(birthDist == 'pois'){
        xMax <- max(
            which(
                dpois(x = 0:{meanBirth* 100},
                      lambda = meanBirth) >= 10^{-8}
                )
            ) + 2
        x <- 0:xMax
        p <- matrix(0, nrow = nReps, ncol = length(x))
        meanLambda <- rnorm(nReps, meanBirth, meanBirthSD)
        for(reps in 1:nReps){
            p[reps, ] <- dpois(x = x, lambda = meanLambda[reps])
        }   
    }

    final <- matrix(data = 0, nrow = nReps, ncol = length(mu))

    for(reps in 1:nReps){
        for(exposure in 1:length(mu)){
            final[reps, exposure] =
                simpleBranch(s = s[reps],
                     p = p[reps,],
                     mu = mu[exposure],
                     nYears = nYears)$final
    }
    }
    finalN = prod(final^N)
    return(list(
        p = p,
        s = s,
        eigenValue = {p %*% x} * s,
        final = finalN,
        quantiles = apply(finalN, 2,
            quantile, c(0.025, 0.50, 0.975))
        )
           )
}

multiStageBranch <- function(s, p,
                             J = c(1, rep(0 , (length(s) - 1))),
                             Trans = matrix(c(0, 1, 0, 1), nrow = 2, byrow = 2),
                             mu= c(0, 1.09), nYears
                             ){
  
    extStage <- function(x, maxP){
        y <- x^(0:(maxP-1))
        return(y)
    }

    maxS <- length(s)
    maxP <- dim(p)[2]
    ext  <- matrix(0, nrow = maxS, ncol = nYears + 1)
    EXT  <- matrix(0, nrow = maxP, ncol = maxS)
    S    <- diag(s * (1 - mu))
    Id   <- diag(1, length(s))
    
    for(year in 1:nYears){
        for(stage in 1:(maxS)){
            EXT[, stage] <-
                extStage( ext[stage, year], maxP)
        }
        ext[,year + 1] = 
            (Id - S) %*% rep(1, length(s)) +
                ((S %*% Trans) %*% diag(ext[,year]) %*% p %*% EXT %*% J )
    }

    colnames(ext) <- paste("Year ",0:nYears)
    rownames(ext) <- paste("Stage ", 1:maxS)
    return(list(ext = ext, final =
                ext[,nYears+1]
                ))
}

uqStochMultiStageBranch <- function(
    sMean = c(0.20, 0.50),
    sPrecision = c(100, 100),
    meanBirth = c(0.70, 1.20),
    meanBirthSD = c(0.05, 0.05),
    birthDist = 'pois',
    pBirth = NA,
    pBirthPrecision = NA,
    maxBirth = NA,
    mu = matrix(c(0, 0.5), nrow = 2),
    nYears = 30,
    nReps = 10,
    N = c(15, 15),
    J = c(1, rep(0 , (length(sMean) - 1))),
    Trans = matrix(c(0, 1, 0, 1), nrow = 2, byrow = 2),
    saveYears = nYears
    ){

    library(popbio)    
    nStages = length(sMean)
    s = matrix(NA, nrow = nStages, ncol = nReps)

    for(i in 1:nStages){
        s[i,] = rbeta(n = nReps,
             shape1 = sMean[i] * sPrecision[i],
             shape2 = (1 - sMean[i]) * sPrecision[i])
    }

    if(birthDist == 'pois'){
        xMax <- max(
            which(
                dpois(x = 0:{max(meanBirth)* 100},
                      lambda = max(meanBirth)) >= 10^{-8}
                )
            ) + 2
        x <- 0:xMax

        meanLambda <- matrix(0, nrow = nStages, ncol = nReps)
        for(i in 1:nStages){
            meanLambda[i,] <-
                rnorm(nReps, meanBirth[i], meanBirthSD[i])
        }

        p = list()
        for(reps in 1:nReps){
            p[[reps]] <- matrix(0, nrow = nStages,
                                ncol = (xMax + 1))
            for(i in 1:nStages){
                p[[reps]][i,] <-
                    dpois(x = x,
                          lambda = meanLambda[i, reps])
                }
            }
    }

    if(birthDist == 'binom'){
        pBirths = matrix(0, nrow = nStages, ncol = nReps)
        for(i in 1:nStages){
            ifelse(pBirth[i] == 0,
                   pBirths[i,] <-0, 
                   pBirths[i,] <-
                   rbeta(
                       n = nReps,
                       shape1 = pBirth[i] * 1000,
                       shape2 = (1 - pBirth[i]) * 1000)
                   )
        }

        p = list()
        for(reps in 1:nReps){
            p[[reps]] <- matrix(0, nrow = nStages,
                                ncol = (maxBirth +1 ))
            for(stage in 1:nStages){
                p[[reps]][stage,] <-
                    dbinom(x = 0:maxBirth, size = maxBirth,
                           prob = pBirths[stage, reps])
            }
        }    
    }
    ## Define needed function
    extStage <- function(x, maxP){
        y <- x^(0:(maxP-1))
        return(y)
    }

    ## Define variables that are only created once
    ext <- matrix(0, nrow = nStages, ncol = nYears + 1)
    final <- list()
    for(save in 1:length(saveYears)){
        final[[save]] <- matrix(data = 0, nrow = nReps, ncol = dim(mu)[2])
    }
    eigenValue <- numeric(nReps)
    genTime <- numeric(nReps)

    for(reps in 1:nReps){
        maxP <- dim(p[[reps]])[2]
        maxS <- dim(p[[reps]])[1] 
        EXT <- matrix(0, nrow = maxP, ncol = maxS)
        Id <- diag(1, maxS)
        for(exposure in 1:dim(mu)[2]){
            S <- diag(s[,reps] * (1 - mu[, exposure]))
            for(year in 1:nYears){
                for(stage in 1:(maxS)){
                    EXT[, stage] <- extStage(
                        ext[stage, year], maxP)
                }    
                ext[,year + 1] =
                    (Id - S) %*% rep(1, length(s[ ,reps])) +
                        ((S %*% Trans) %*% diag(ext[,year]) %*%
                         p[[reps]] %*% EXT %*% J)
            } ## end years
            for(save in 1:length(saveYears)){
                final[[save]][reps, exposure] <- prod(ext[,saveYears[save]+ 1]^N)
            }
        } ## end mu
        projMatrix <- diag(s[, reps])
        projMatrix[nStages - 1,nStages] <-
            projMatrix[nStages,nStages]
        projMatrix <- rbind(
            apply(
                p[[reps]] %*% 0:(maxP - 1),
                1, sum),
            projMatrix[1:(nStages-1),]
            )
        eigenValue[reps] <-
            max(Re(eigen(projMatrix)$values))
        genTime[reps] <- 
            generation.time(round(projMatrix, 2))
   
    } ## end reps 
    quantiles <- NULL
    for(save in 1:length(saveYears)){
        quantiles <-
            rbind(quantiles,
                  data.frame( apply(final[[save]], 2,
                                    quantile, c(0.025, 0.50, 0.975)),
                             Year = saveYears[save])
                )
    }
    colnames(quantiles)[1:dim(mu)[2]] <- mu[dim(mu)[1], ] 


    return(
        list(
            p = p,
            s = s,
            final = final,
            eigenValue = eigenValue,
            genTime= genTime,
            quantiles = quantiles
        )
        )


}



cleanPplot <- function(x = outCaveBat$p[[1]],
                       nStages = length(sMeanCaveBat),
                       sMeanCaveBat = sMeanCaveBat,
                       species = 'Snuffleupagus'){

    maxBirthCaveBat = dim(x)[2]
    pCaveBatGG <- data.frame(x)
    colnames(pCaveBatGG) <- 0:(maxBirthCaveBat-1)
    rownames(pCaveBatGG) <- 1:nStages
    pCaveBatGG <- pCaveBatGG[,apply(pCaveBatGG, 2, max) >= 10^-4]
    pCaveBatGG <- melt(pCaveBatGG)
    colnames(pCaveBatGG) <- c("Births", "Probability")
    pCaveBatGG$Survival <- sMeanCaveBat
    pCaveBatGG$LifeStage <- 1:nStages
    pCaveBatGG$Species <- species
    return(pCaveBatGG)
}