library(ggplot2)
library(pse)
library(deSolve)
library(rootSolve)

#### For jupyter notebook only
options(jupyter.plot_mimetypes = 'image/png')
####

oneRun <- function(X0, Y0, a, b,  times, runsteady=TRUE){
    ## parameters: initial sizes and competition coefficients
    parameters <- c(a = a, b = b)
    state <- c(X = X0, Y = Y0)
    ## The function to be integrated
    LV <- function(t, state, parameters){
        with(as.list(c(state, parameters)), {
            dX = X*(1 - X - a*Y)
            dY = Y*(1 - Y - b*X)
            list(c(dX, dY))
        })
    }
    ## Integrating: runsteady to run untill convergence (be carefull using that!)
    ## Or the usual integration with the function ode
    if(runsteady){
        out <- runsteady(y = state, times = times, func = LV, parms = parameters)
        return(out$y) # runsteady returns only final states
    }
    else{
        out <- ode(y = state, times = times, func = LV, parms = parameters)
        return(out[nrow(out),-1]) # indexing to get final state values
    }
}

modelRun <- function (my.pars) {
    return(mapply(oneRun, my.pars[,1], my.pars[,2], my.pars[,3], my.pars[,4],
                  MoreArgs=list(times=c(0,Inf), runsteady=TRUE)))
}

factors <- c("X0", "Y0", "a", "b")

q <- c("qunif", "qunif", "qunif", "qunif")

q.arg <- list(list(min=0, max=1), list(min=0, max=1), list(min=0, max=2), list(min=0, max=2))

myLHS <- LHS(model=modelRun, factors=factors, 
             N=200, q=q, q.arg=q.arg, nboot=100) # use nboot=0 if do not need partial correlations

plotecdf(myLHS, stack=TRUE, xlab="Population relative size")

plotscatter(myLHS, add.lm=FALSE)

plotprcc(myLHS)

## get.data export a table with parameter values
hypercube <- get.data(myLHS)

## get.results exports the outputs
hypercube <- cbind(hypercube, get.results(myLHS)) # appending columns of outputs to the hypercube matrix
names(hypercube)[5:6] <- c("X", "Y") # cosmetic

sum(hypercube$X>1e-6&hypercube$Y>1e-6)

## Create a logical variable to flag coexistence
hypercube$coexistence <- hypercube$X>1e-6&hypercube$Y>1e-6
## boxplots
par(mfrow=c(2,2))
boxplot(X0~coexistence, data=hypercube, xlab="Coexistence", main="X0")
boxplot(Y0~coexistence, data=hypercube, xlab="Coexistence", main="Y0")
boxplot(a~coexistence, data=hypercube, xlab="Coexistence", main="a")
boxplot(b~coexistence, data=hypercube, xlab="Coexistence", main="b")
par(mfrow=c(1,1))

plot(a~b, data=hypercube, type="n") ## to scale the plot for the whole parameter space
points(a~b, data=hypercube, subset=hypercube$X>1e-6&hypercube$Y>1e-6, col="blue")
points(a~b, data=hypercube, subset=!(hypercube$X>1e-6&hypercube$Y>1e-6), col="red")


oneRun.sir <- function(S0, I0, R0, r, a, time=seq(0, 50, by = 0.01)){
    ## parameters: transmission and recovering rates
    parameters <- c(r = r, a = a)
    ## initial population sizes: start with a single infected individual to test invasibility
    state <- c(S = S0, I = I0, R = R0)
    ## The function to be integrated
    sir <- function(t, state, parameters){
        with(as.list(c(state, parameters)), {
            dS = -r*S*I
            dI = r*S*I - a*I
            dR = a*I
            list(c(dS, dI, dR))
        })
    }
    ## Integrating
    out <- ode(y = state, times = time, func = sir, parms = parameters)
    return(any(out[,3]>I0)) # returns a logical: any infected number larger than initial number?
}

modelRun.sir <- function (my.pars) {
    return(mapply(oneRun.sir, my.pars[,1], 1, 0, my.pars[,2], my.pars[,3]))
}

## A string vector to name the parameters
factors <- c("S0", "r", "a")
## The probability quantile function used to draw values of each parameter
q <- c("qunif", "qunif", "qunif")
## A list of the parameters of the above probability distributions
q.arg <- list(list(min=10, max=200), list(min=0.001, max=0.01), list(min=0, max=1))

myLHS <- LHS(model=modelRun.sir, factors=factors, N=200, q=q, q.arg=q.arg)

##Table of parameter combinations with get.data
hypercube <- get.data(myLHS)
## adding the output for each combination with get.results
hypercube$Spread <- get.results(myLHS)
## box plots parameter ~ occurence of spread (logical)
par(mfrow=c(1,3))
boxplot(S0~Spread, data=hypercube, xlab="Disease spread", ylab="Initial number of susceptibles")
boxplot(r~Spread, data=hypercube, xlab="Disease spread", ylab="Transmission rate")
boxplot(a~Spread, data=hypercube, xlab="Disease spread", ylab="Recovering rate")
par(mfrow=c(1,1))


plot(a~r, data=hypercube, type="n") # to show all ranges of the parameters
points(a~r, data=hypercube, subset=S0<median(S0)&Spread==1, col="red")
points(a~r, data=hypercube, subset=S0<median(S0)&Spread==0, col="blue")

plot(I(a/r)~S0, data=hypercube, type="n") # to show all ranges of the parameters
points(I(a/r)~S0, subset=Spread==1, data=hypercube, col="red")
points(I(a/r)~S0, subset=Spread==0, data=hypercube, col="blue")
abline(0,1) ## equivalence line