### R code from vignette source 'homework_and_coursenotes6.Rnw'

###################################################
### code chunk number 1: homework_and_coursenotes6.Rnw:189-243
###################################################
reed.frost <- function(pp, nn, c1, t.end, cumul.only = FALSE) { 
  if (t.end > 1000) {
    stop("t.end too big")
  }
  if (t.end <= 0) {
    stop("negative or zero ending time")
  }
  if (c1 < 0 || abs(round(c1) - c1) > 1e-07) {
    stop("invalid starting number of cases")
  }
  if (nn < 0 || abs(round(nn) - nn) > 1e-07) { 
    stop("invalid population size")
  }
  if (pp > 1 || pp < 0) {
    stop("invalid transmission probability")
  }
  if (nn < c1) {
    stop("more cases than people")
  }
  ss <- rep(0, t.end)
  cc <- rep(0, t.end)
  cumul <- 0
  current.cc <- c1
  current.ss <- nn-c1
  if (!cumul.only) {
    cc[1] <- current.cc
    ss[1] <- current.ss
  }
  for (ii in 2:t.end) {
    rr <- 1-(1-pp)^current.cc
    current.cc.new <- rbinom(1,size=current.ss,prob=rr)
    current.ss.new <- current.ss - current.cc.new
    if (!cumul.only) {
      ss[ii] <- current.ss.new
      cc[ii] <- current.cc.new
    }
    current.ss <- current.ss.new
    current.cc <- current.cc.new
    cumul <- cumul + current.cc.new
    if (current.cc.new == 0) {
      if (!cumul.only) {
        for (jj in (ii+1):t.end) {
          ss[jj] <- current.ss
        }
      }
      break
    }
  }
  if (cumul.only) {
    cumul
  } else {
    list(susc=ss,cases=cc,cumul.new.cases=cumul)
  }
}


###################################################
### code chunk number 2: homework_and_coursenotes6.Rnw:251-272
###################################################
reed.frost.average <- function(pp,nn,c1,t.end,nreps=1,cumul.only=FALSE) {
  if (cumul.only) {
    cumul <- reed.frost(pp,nn,c1,t.end,cumul.only) 
    for (ii in 2:nreps) {
      cumul <- cumul + reed.frost(pp,nn,c1,t.end,cumul.only)
    }
    cumul/nreps
  } else {
    run <- reed.frost(pp,nn,c1,t.end)
    ss <- run$susc
    cc <- run$cases
    if (nreps >= 2) {
      for (ii in 2:nreps) {
        run <- reed.frost(pp,nn,c1,t.end)
        ss <- ss+ run$susc
        cc <- cc + run$cases
      }
    }
    list(susc=ss/nreps,cases=cc/nreps)
  }
}


###################################################
### code chunk number 3: homework_and_coursenotes6.Rnw:276-295
###################################################
colrz <- c("black","red","blue","green","cyan","gray","magenta","orange")
multi.rf <- function(nsims,pp,nn,endtime,pcolors=colrz) {
  if (nsims<1) {
    stop("bad nsims")
  }
  ans <- reed.frost(pp=pp,nn=nn,c1=1,t.end=endtime) 
  sims <- list(ans)
  for (jj in 2:nsims) {
    sims[[jj]] <- reed.frost(pp=pp,nn=nn,c1=1,t.end=endtime) 
  }
  timez <- 1:endtime
  plot(timez,sims[[1]]$cases)
  points(timez,sims[[1]]$cases,type="l")
  for (ii in 2:nsims) {
    points(timez,sims[[ii]]$cases,col=pcolors[[(ii %% length(pcolors))+1]])
    points(timez,sims[[ii]]$cases,col=pcolors[[(ii %% length(pcolors))+1]],type="l")
  }
}
multi.rf(nsims=8,pp=1/16,nn=64,endtime=16)


###################################################
### code chunk number 4: homework_and_coursenotes6.Rnw:581-604
###################################################
require(deSolve)
km <- function(t,sys,params) {
  dsys <- sys[]
  c1 <- params[["beta"]] * sys[1] * sys[2]
  c2 <- sys[2]*params[["rho"]]
  dsys[1] <- -c1
  dsys[2] <- c1-c2
  dsys[3] <- c2
  list(dsys,NULL)
}
dosim <- function(pars,init,endtime) {
  tmp <- lsoda(init,seq(0,endtime,length=512),km,pars)
  ans <- as.data.frame(tmp)
  names(ans) <- c("time","X","Y","Z")
  ans 
}
u <- dosim(pars=c(beta=0.001,rho=0.1),init=c(1999,1,0),endtime=40)
kmplot <- function(kmout) {
  plot(u$time,u$X,col="black",type="l")
  points(u$time,u$Y,col="red",type="l")
  points(u$time,u$Z,col="blue",type="l")
}
kmplot(u)