################################################################
################################################################
## Load functions and packages 
source("branching_functions.R")
library(ggplot2)
library(reshape2)
library(scales)
library(popbio)

################################################################
## Global parameters 
N = c(30, 100, 1000) # vector of number of individuals
nMusToUse <- 40 # Resolution of mus between 0 and 1

## matrix of mu values for two life-stage organisms 
muTwo = rbind(rep(0, nMusToUse),
    seq(0, 1.0, length = nMusToUse))

## matrix of mu values for five life-stage organisms 
muFive = rbind(
    rep(0, nMusToUse),
    seq(0, 1.0, length = nMusToUse),
    seq(0, 1.0, length = nMusToUse),
    seq(0, 1.0, length = nMusToUse),
    seq(0, 1.0, length = nMusToUse))

## Number of stochastic replicates to run
nReps = 5000
## Years to observe risk of extinction
nYears = c(10, 30, 100)

################################################################
## Enter parameters for the cave bat
sMeanCaveBat = c(
    0.922 * 0.902 * 0.898,
    0.951 * 0.948 * 0.948)
sPrecisionCaveBat = c(1000, 1000)
meanBirthCaveBat = NA
meanBirthSDCaveBat = NA
maxBirthCaveBat = 1
birthDistCaveBat = 'binom'
pBirthCaveBat = c(0.4, 0.90)* 0.5 * 0.8
## pup production * 1/2 female * survival of pups to first winter 
pBirthPrecisionCaveBat = c(1000, 1000)


## Population projection matrix for the cave bat
PPM <-
    matrix(data = c(
           pBirthCaveBat  * sMeanCaveBat,
           sMeanCaveBat),
       nrow = 2, ncol=2, byrow = TRUE)
PPM
eigen(PPM)$values
stableCaveBat <- stable.stage(PPM)

## Tree Bat
sMeanTreeBat = c(0.42, 0.65)
sPrecisionTreeBat = c(1000, 1000)
meanBirthTreeBat = NA
meanBirthSDTreeBat = NA
maxBirthTreeBat = 6
birthDistTreeBat = 'binom'
noBirths <- 0:maxBirthTreeBat
pBirthTreeBat = c(1.5/6, 2.5/6)* 0.5
pBirthPrecisionTreeBat = c(1000, 1000)

PPMtb <-
    matrix(data = c(
           pBirthTreeBat * maxBirthTreeBat *  sMeanTreeBat,
           sMeanTreeBat
               ),
       nrow = 2, ncol=2, byrow = TRUE)
PPMtb

eigen(PPMtb)
eigen(PPM)
stableTreeBat <- stable.stage(PPMtb)

## Passerine (Horned lark) parameters 
sMeanPass = c(0.255, 0.60)
sPrecisionPass = c(1000, 1000)
meanBirthPass = c(1.50, 2.25)
meanBirthSDPass = c(0.24, 0.24)
maxBirthPass = 6
birthDistPass = 'pois'

PPMp <-
    matrix(data = c(
           meanBirthPass * sMeanPass,
           sMeanPass),
       nrow = 2, ncol=2, byrow = TRUE)
PPMp
eigen(PPMp)$values
eigen(PPMtb)$values
eigen(PPM)$values
stablePass <- stable.stage(PPMp)
stableTreeBat
stableCaveBat
stablePass

### Eagle parameters 
sMeanEagle = c(0.68, 0.7, 0.8, 0.9, 0.955)
sPrecisionEagle = c(1000, 1000, 1000, 1000, 1000)
meanBirthEagle = NA
meanBirthSDEagle = NA
maxBirthEagle = 3
birthDistEagle = 'binom'
noBirths <- 0:maxBirthEagle
pBirthEagle = c(0, 0, 0, 0.3/maxBirthEagle, 1.1/maxBirthEagle) * 0.5
pBirthPrecisionEagle =  c(1000, 1000, 1000, 1000, 1000)
TransEagle <- matrix(c(0, 1, 0, 0, 0,
                       0, 0, 1, 0, 0,
                       0, 0, 0, 1, 0,
                       0, 0, 0, 0, 1,
                       0, 0, 0, 0, 1), nrow = 5, byrow = TRUE)

temp <- diag(sMeanEagle)[1:4,]
temp[4,5] <- sMeanEagle[5]

PPMe <- rbind(
    pBirthEagle*maxBirthEagle * sMeanEagle,
    temp)
               
PPMe
eigen(PPMe)$values
eigen(PPMp)$values
eigen(PPMtb)$values
eigen(PPM)$values
stableEagle <- stable.stage(PPMe)
stableTreeBat
stableCaveBat
stablePass
stableEagle
generation.time(PPM)
generation.time(PPMtb)
generation.time(PPMp)
generation.time(PPMe)
###############################################################
## Compare risk for all four speices

## Create objects to store outputs
riskAll <- NULL
riskPolygonCaveBat <- NULL
riskPolygonTreeBat <- NULL
riskPolygonPass <- NULL
riskPolygonEagle <- NULL

## Start for loop across number of individuals
for(sizeSim in 1:length(N)){
    outCaveBat <- uqStochMultiStageBranch(
        sMean = sMeanCaveBat,
        sPrecision = sPrecisionCaveBat,
        meanBirth = meanBirthCaveBat,
        meanBirthSD = meanBirthSDCaveBat,
        birthDist = birthDistCaveBat,
        mu = muTwo,
        nYears = max(nYears),
        nReps = nReps,
        maxBirth = maxBirthCaveBat,
        pBirth = pBirthCaveBat,
        pBirthPrecision = pBirthPrecisionCaveBat,
        N = N[sizeSim] * stableCaveBat,
        saveYears = nYears
        )

    outTreeBat <- uqStochMultiStageBranch(
        sMean = sMeanTreeBat,
        sPrecision = sPrecisionTreeBat,
        meanBirth = meanBirthTreeBat,
        meanBirthSD = meanBirthSDTreeBat,
        birthDist = birthDistTreeBat,
        mu = muTwo,
        nYears = max(nYears),
        nReps = nReps,
        maxBirth = maxBirthTreeBat,
        pBirth = pBirthTreeBat,
        pBirthPrecision =
        pBirthPrecisionTreeBat,
        N = N[sizeSim] * stableTreeBat,
        saveYears = nYears
        )
    
    outPass <- uqStochMultiStageBranch(
        sMean = sMeanPass,
        sPrecision = sPrecisionPass,
        meanBirth = meanBirthPass,
        meanBirthSD = meanBirthSDPass,
        birthDist = birthDistPass,
        mu = muTwo,
        nYears = max(nYears),
        nReps = nReps,
        maxBirth = maxBirthPass,
        pBirth = NA,
        pBirthPrecision = NA,
        N = N[sizeSim] * stablePass,
        saveYears = nYears
        )

    outEagle <- uqStochMultiStageBranch(
        sMean = sMeanEagle,
        sPrecision = sPrecisionEagle,
        meanBirth = meanBirthEagle,
        meanBirthSD = meanBirthSDEagle,
        birthDist = birthDistEagle,
        mu = muFive,
        nYears = max(nYears),
        nReps = nReps,
        maxBirth = maxBirthEagle,
        pBirth = pBirthEagle,
        pBirthPrecision = pBirthPrecisionEagle,
        Trans = TransEagle,
        N = N[sizeSim] * stableEagle,
        saveYears = nYears
        )

    riskAll <- rbind(
        riskAll,
        data.frame(
            melt(
                data.frame(
                    outCaveBat$quantiles[grep(x =
                                              rownames(
                                                  outCaveBat$quantiles),
                                              pattern = "50%", perl = TRUE), ]) ,
                id.vars = "Year", value.name = 'ext', variable.name = "mu"),
            Taxa = "Cave Bat",
            N = N[sizeSim]
            ),
        data.frame(
            melt(data.frame(
                outTreeBat$quantiles[grep(x = rownames(outTreeBat$quantiles),
                                          pattern = "50%", perl = TRUE), ]) ,
                 id.vars = "Year", value.name = 'ext', variable.name = "mu"),
            Taxa = "Tree Bat",
            N = N[sizeSim]
            ),
        data.frame(
            melt(data.frame(
                outPass$quantiles[grep(x = rownames(outPass$quantiles),
                                       pattern = "50%", perl = TRUE), ]) ,
                 id.vars = "Year", value.name = 'ext', variable.name = "mu"),
            Taxa = "Grassland Songbird",
            N = N[sizeSim]
            ),
        data.frame(
            melt(data.frame(
                outEagle$quantiles[grep(x = rownames(outEagle$quantiles),
                                        pattern = "50%", perl = TRUE), ]) ,
                 id.vars = "Year", value.name = 'ext', variable.name = "mu"),
            Taxa = "Eagle",
            N = N[sizeSim]
            )
        )
    
    riskPolygonCaveBat <- rbind(
        riskPolygonCaveBat,
        cbind(
            Taxa = "Cave Bat",
            N = N[sizeSim],
            melt(
                t(
                    cbind(
                        outCaveBat$quantiles[grep(
                            x =rownames(outCaveBat$quantiles),
                            pattern =
                            "2.5%", perl = TRUE), -dim(outTreeBat$quantiles)[2]],
                        outCaveBat$quantiles[grep(
                            x = rownames(outCaveBat$quantiles),
                            pattern = "97.5%",
                            perl = TRUE), (dim(outCaveBat$quantiles)[2] - 1):1]
                        )
                    ),
                value.name = 'ext', varnames = c("mu", "Year"))
            )
        )
    
    riskPolygonTreeBat <- rbind(
        riskPolygonTreeBat,
        cbind(
            Taxa = "Tree Bat",
            N = N[sizeSim],
            melt(
                t(cbind(
                    outTreeBat$quantiles[grep(
                        x =rownames(outTreeBat$quantiles),
                        pattern =
                        "2.5%", perl = TRUE), -dim(outTreeBat$quantiles)[2]],
                    outTreeBat$quantiles[grep(
                        x = rownames(outTreeBat$quantiles),
                        pattern = "97.5%",
                        perl = TRUE), (dim(outTreeBat$quantiles)[2]-1):1])),
                value.name = 'ext', varnames = c("mu", "Year"))
            )
        )
    
    riskPolygonPass <- rbind(
        riskPolygonPass,
        cbind(
            Taxa = "Grassland Songbird",
            N = N[sizeSim],
            melt(
                t(cbind(
                    outPass$quantiles[grep(
                        x =rownames(outPass$quantiles),
                        pattern =
                        "2.5%", perl = TRUE), -dim(outPass$quantiles)[2]],
                    outPass$quantiles[grep(
                        x = rownames(outPass$quantiles),
                        pattern = "97.5%", perl = TRUE),
                                      (dim(outPass$quantiles)[2]-1):1])),
                value.name = 'ext', varnames = c("mu", "Year"))
            )
        )

    riskPolygonEagle <- rbind(
        riskPolygonEagle,
        cbind(
            Taxa = "Eagle",
            N = N[sizeSim],
            melt(
                t(cbind(
                    outEagle$quantiles[grep(
                        x =rownames(outEagle$quantiles),
                        pattern =
                        "2.5%", perl = TRUE), -dim(outEagle$quantiles)[2]],
                    outEagle$quantiles[grep(
                        x = rownames(outEagle$quantiles),
                        pattern = "97.5%", perl = TRUE),
                                       (dim(outEagle$quantiles)[2]-1):1])),
                value.name = 'ext', varnames = c("mu", "Year"))
            )
        )
}

## Reformat data for plotting 
head(riskAll)
levels(riskAll$mu) <- muTwo[2,]
riskAll$mu <- as.numeric(as.character(riskAll$mu))

colnames(riskAll)
riskAll$Year <- as.factor(riskAll$Year)
levels(riskAll$Year) <- paste("Year", levels(riskAll$Year), sep = " ")
riskAll$N <- as.factor(riskAll$N)
levels(riskAll$N) <- paste("N = ", levels(riskAll$N), sep = "")

head(riskAll)
str(riskAll)

dim(riskPolygonCaveBat)
riskPolygon <- rbind(
    riskPolygonCaveBat,
    riskPolygonTreeBat,
    riskPolygonPass,
    riskPolygonEagle)

dim(riskPolygon)
str(riskPolygon)
head(riskPolygon)
levels(riskPolygon$mu)
colnames(riskPolygon)
riskPolygon$mu <- as.numeric(as.character(riskPolygon$mu))
levels(riskPolygon$Year) <- paste("Year", nYears, sep = " ")
riskPolygon$N <- as.factor(riskPolygon$N)
levels(riskPolygon$N) <- paste("N = ", levels(riskPolygon$N), sep = "")


allEigen <- round(Re(c(
    eigen(PPM)$values[1],
    eigen(PPMtb)$values[1],    
    eigen(PPMp)$values[1],
    eigen(PPMe)$values[1])), 2)
allEigen
    
## Plot risk for all 4 species
example <-
    ggplot(data = riskAll, aes(x = mu, y = ext, color = Taxa)) +
    geom_polygon(data = riskPolygon, aes(x = mu, y = ext, fill = Taxa),
                 color = alpha('white', 0)) +
    facet_grid(Year ~ N) +
    geom_line(size = 1.05) +
    ylab("P(extinction) from\ndemographic stochasticity") +
    xlab(
        expression("Decrease in survival "*(mu))
        ) +
    scale_x_continuous(breaks =
                       seq(0, max(muTwo), by = 0.2)) +
    scale_y_continuous(breaks = seq(0,1, by = 0.2)) + theme_bw() +
    scale_colour_manual(
        labels = c(
            bquote(.(levels(riskAll$Taxa)[1]) * " " * lambda == .(allEigen[1])),
            bquote(.(levels(riskAll$Taxa)[2]) * " " * lambda == .(allEigen[2])),
            bquote(.(levels(riskAll$Taxa)[3]) * " " * lambda == .(allEigen[3])),
            bquote(.(levels(riskAll$Taxa)[4]) * " " * lambda == .(allEigen[4]))),
        values = c("black", "blue", "red", "orange")) +
    scale_fill_manual(
        labels = c(
            bquote(.(levels(riskAll$Taxa)[1]) * " " * lambda == .(allEigen[1])),
            bquote(.(levels(riskAll$Taxa)[2]) * " " * lambda == .(allEigen[2])),
            bquote(.(levels(riskAll$Taxa)[3]) * " " * lambda == .(allEigen[3])),
            bquote(.(levels(riskAll$Taxa)[4]) * " " * lambda == .(allEigen[4]))),
        values = alpha(c("black", "blue", "red", "orange"), 0.25)) +
    coord_cartesian(xlim = c(0, max(muTwo)+0.2), ylim = c(0, 1.05))
    

x11()
print(example)
ggsave("exampleTake.pdf", example, width = 8, height = 6)

#################################################################
## ## ## Run simmulations and create population size effect plot
nSize = 2^{0:09} # vector of population sizes
mu = rbind(rep(0, 10),
          seq(0, 0.2, length = 100)
          )


riskMuOnly <- matrix(0,  nrow = dim(mu)[2], ncol = length(sMeanCaveBat))
for(exposure in 1:dim(mu)[2]){
    riskMuOnly[exposure, ] <-
        multiStageBranch(
            s = sMeanCaveBat,
            p = as.matrix(cbind((1 - pBirthCaveBat), pBirthCaveBat)),
            mu = mu[, exposure],
            nYears = 1000,
            )$final
}

riskPopSize <- matrix(0,  nrow = dim(mu)[2], ncol = length(nSize))

### Loop through different population sizes 
for(muIndex in 1:dim(mu)[2]){
    for(sizeIndex in 1:length(nSize)){
        riskPopSize[ muIndex, sizeIndex] <- 
            prod(
                riskMuOnly[muIndex,] ^ (stableCaveBat *nSize[sizeIndex])
                )
    }
}

## Reformat data
riskPopSize <- as.data.frame(riskPopSize)
colnames(riskPopSize) <- nSize
riskPopSize$mu <-  mu[2,]
head(riskPopSize)

## Plot results 
riskPopSizeGG <- melt(riskPopSize, id.vars = 'mu')
head(riskPopSizeGG)
colnames(riskPopSizeGG) <- c("mu", "N", "ext")
riskPopSizeGG$N <-as.numeric(as.character(riskPopSizeGG$N))
head(riskPopSizeGG, 5)

## ## Determine which mu gives an eigenvalue of 1
muEigen <- data.frame(mu = mu[2,])
muEigen$eigenValue <- 0
for(m in 1:dim(mu)[2]){
    muEigen$eigenValue[m] <- Re(eigen(
        rbind(
            pBirthCaveBat * (1 - mu[,m]),
            sMeanCaveBat * (1 - mu[,m])
            )
        )$value[1])
}

muEigen[ min(which(muEigen$eigenValue <= 1.00)), ]

## Save the "matrix results" for plotting 
mtxResult <- data.frame(
    Matrix = factor(c(rep("< threshold", 2), rep("> threshold", 2))),
    mu = c(-2,
        mu[2 ,  min(which(muEigen$eigenValue <= 1.00))],
        mu[2 ,  min(which(muEigen$eigenValue <= 1.00))],
        2),
    ext =  c(0, 0, 1, 1))

## plot results
exampleSize <- ggplot(data = riskPopSizeGG, aes(x = mu, y = ext)) +
    geom_line(aes(color = N, group = N), size = 1.05) +
    ylab("P(extinction) from\ndemographic stochasticity") +
    xlab(
        expression("Decrease in survival "*(mu))
        ) +
    scale_x_continuous(breaks = seq(0,1.0, by = 0.1)) +
    scale_y_continuous(breaks = seq(0,1, by = 0.2)) + theme_bw() +
    scale_colour_gradient(low = "grey75", high = "black", 
                          limits = c(1, round(max(riskPopSizeGG$N),-1)),
                          name = "Population\nSize", trans = 'log10',
                          breaks = c(nSize, nSize[length(nSize)]*2)) +
    coord_cartesian(xlim = c(0, max(mu))) +
    geom_line(data = mtxResult,
              aes(x = mu, y = ext, group = Matrix),
              size = 2, color = 'gold') +
    geom_point(data = mtxResult,
              aes(x = mu, y = ext, shape = Matrix),
              size = 5, color = 'gold') +
    scale_shape_manual(values = c(16, 1), name = "Matrix\nResults")

x11()
print(exampleSize)
ggsave("exampleTakeSize.pdf", exampleSize, width = 6, height = 4)