## Code to run the IBM in

### NOTES:
###
###  In the paper, the focus is on evaluating fitness associated with static thresholds dividing self and non-self, rather than allowing them to evolven
###  However, the code is written to allow evolution to occur, which can be achieved here by setting some part of the mutation vector to -1 and 1.
###
###  The code also currently constrains mean and variance of offspring epitopes to overall pop features, rather than allowing heritability in this characteristic
###  It is possible to unhash some code to allow this to evolve as well
###
###  In this framing mu.d is added to an individual's hazard for every epitope it experience that is below its threshold
###



##' Generate this year's pathogens
##'
##' @param npath      number of pathogens in the year
##' @param mean       their mean, along the continuous epitope scale
##' @param sd         their sd, likewise
##' @param bins       the bins into which the values will be discretized
##'
##' @return a vector of length the number of bins containing numbers of pathogens in each bin that year
##'

## Returns a vector of length Bins
genPathogen <- function(npath, mean,sd, bins=seq(-10,10,length=20)) {
    store <- rep(0,length(bins))
    u <- rnorm(npath,mean,sd)
    n.per <- table(findInterval(u,bins, all.inside=TRUE))
    store[as.numeric(names(n.per))] <- n.per
    return(store)
}


##' Generate offspring distribution based on a random deviate from overall
##' pop characteristics; could be modulated to allow inheritance of epitope distribution, but here assume constrained.
##'
##' @param nepitope   number of 'self' epitopes
##' @param mean.mother  mother's mean epitopes
##' @param mean.pop     population mean
##' @param sd.pop       population sd
##' @param bins       the bins into which the values will be discretized
##'
##' @return a matrix with as many rows as mothers, and columns as the number of bins
##'

genKid <- function(nepitope, mean.mother,
                   mean.pop,sd.pop, bins=seq(-10,10,length=20)) {

    ## storage
    store <- matrix(0,length(mean.mother),length(bins))

    ## pull out fathers
    mean.father <- rnorm(1,mean.pop,sd.pop)

    ## generate distribution (will iterate over each set of means
    u <- matrix(rnorm(nepitope*length(mean.mother),0.5*(mean.mother+mean.father),sd.pop),
                length(mean.mother),nepitope,byrow = FALSE)

    for (j in 1:length(mean.mother)) {
        n.per <- table(findInterval(u[j,],bins, all.inside=TRUE))
        #print(c(n.per))
        store[j,as.numeric(names(n.per))] <- as.numeric(n.per)

    }

    return(store)
}


##' Generate survival probability given epitope distribution and pathogen distribution
##' for the current population of 'mothers'
##'
##' @param mother.dist       matrix with rows equal to n mothers; and columns set by nbins; contains number of 'self' epitopes
##' @param mother.threshold  mother's mean threshold, stored as an index in bins
##' @param pathogen.dist     vector of number of pathogens in each bin (output of genPathogen)
##' @param mother.memory     matrix with rows equal to n mothers; and columns set by nbins; contains memory to each pathogen epitope
##' @param params            a list containing mortality hazards
##'
##' @return a vector with as many values as mothers, containing the probability of survival
##'

survProb <- function(mother.dist,
                     mother.threshold,  ## this as an index along the vector
                     pathogen.dist,
                     mother.memory=rep(0,length=20),
                     params=c(mu.b=1/20,mu.i=0.1,mu.d=0.5,mu.di=0.3)) {

    lts <- length(mother.dist[1,])
    ix <- mother.dist*0
    surv <- rep(0,nrow(mother.dist))
    for (j in 1:nrow(mother.dist)) {

        ## First, for every bin, if pathogen present
        ix[j,] <- (pathogen.dist>0)*(1-mother.memory[j,])

        haz.disease <- sum(ix[j,1:mother.threshold[j]]*params["mu.d"])
        haz.id <- sum(ix[j,min(mother.threshold[j]+1, lts):lts]*params["mu.di"])
        haz.immunopath <- sum((mother.dist[j,]>1)[min(mother.threshold[j]+1,lts):lts]*params["mu.i"])

        ## Probabilty of survival
        surv[j] <- exp(-(params["mu.b"]+haz.disease+haz.id+haz.immunopath))
    }

    return(surv)

}


##' Generate 'outcomes' given epitope distribution and pathogen distribution
##' for the current population of 'mothers'
##'
##' @param mother.dist       matrix with rows equal to n mothers; and columns set by nbins; contains number of 'self' epitopes
##' @param mother.threshold  mother's mean threshold, stored as an index in bins
##' @param pathogen.dist     vector of number of pathogens in each bin (output of genPathogen)
##' @param mother.memory     matrix with rows equal to n mothers; and columns set by nbins; contains memory to each pathogen epitope
##' @param params            a list containing mortality hazards
##'
##' @return a vector of length three, indicating the incidence of mu.d, mu.id, and mu.i in that order
##'


propOutcomes <- function(mother.dist,
                     mother.threshold,  ## this as an index along the vector
                     pathogen.dist,
                     mother.memory=rep(0,length=20),
                     params=c(mu.b=1/20,mu.i=0.1,mu.d=0.5,mu.di=0.3)) {

    lts <- length(mother.dist[1,])
    ix <- mother.dist*0

    haz.disease <- haz.id <- haz.immunopath <- 0

    for (j in 1:nrow(mother.dist)) {

        ## First, for every bin, if pathogen present
        ix[j,] <- (pathogen.dist>0)*(1-mother.memory[j,])

        haz.disease <- haz.disease+sum(ix[j,1:mother.threshold[j]]*params["mu.d"])
        haz.id <- haz.id + sum(ix[j,min(mother.threshold[j]+1, lts):lts]*params["mu.di"])
        haz.immunopath <- haz.immunopath + sum((mother.dist[j,]>1)[min(mother.threshold[j]+1,lts):lts]*params["mu.i"])

    }

    return(c(haz.disease,haz.id,haz.immunopath)/nrow(mother.dist))

}



##' Update the memory of each mother, using the pathogens on that year, and the thresholds correspond to
##' the current population of 'mothers'
##'
##' @param mother.dist       matrix with rows equal to n mothers; and columns set by nbins; contains number of 'self' epitopes
##' @param mother.threshold  mother's mean threshold, stored as an index in bins
##' @param pathogen.dist     vector of number of pathogens in each bin (output of genPathogen)
##' @param mother.memory     matrix with rows equal to n mothers; and columns set by nbins; contains current memory to each pathogen epitope
##' @param params            a list containing a value named 'memory' - which is the prob that the individuals retains memory to epitopes
##'
##' @return a matrix with as many rows as mothers, and columns as the number of bins
##'
##'
motherMemory <- function(mother.dist, mother.threshold,
                         pathogen.dist,
                         mother.memory=rep(0,length=20),
                         params=c(memory=1)){

    lts <- length(mother.dist[1,])

    for (k in 1:nrow(mother.dist)) {
        #higher chance memory if more pathogens in that psace
        #tmp <- 1-exp(-params["memory"]*pathogen.dist[mother.threshold[k]:lts])
        #mother.memory[k,mother.threshold[k]:lts] <- rbinom(1,1,tmp)  #stoch

        #or just have it if any pathogens in that space
        tmp <- params["memory"]*(pathogen.dist[mother.threshold[k]:lts]>0)
        mother.memory[k,mother.threshold[k]:lts] <- tmp

    }

    return(mother.memory)

}





##' Evalute the probability of miscarriage
##'
##' @param mother.dist       matrix with rows equal to n mothers; and columns set by nbins; contains number of 'self' epitopes
##' @param mother.threshold  mother's mean threshold, stored as an index in bins
##' @param kid.dist          matrix with rows equal to n kids; and columns set by nbins; contains number of 'self' epitopes of kdis
##' @param params            a list containing a value named 'miscarry' - sets the prob of miscarriage given existence of 1 kid epitope below mother thrreshold
##'
##' @return a vector with as many values as mothers, containing the probability of miscarriage
##'
##'
miscarriage <- function(mother.dist, mother.threshold,kid.dist,params=c(miscarry=1)) {

    lts <- length(mother.dist[1,])

    p.miscarriage <- rep(0,nrow(kid.dist))
    for (j in 1:nrow(mother.dist)) {
        #p.miscarriage[j] <- 1-exp(-sum((kid.dist[j,]>1)[min(mother.threshold[j]+1,lts):lts]*params["miscarry"]))
        p.miscarriage[j] <- params["miscarry"]*(sum((kid.dist[j,]>1)[min(mother.threshold[j]+1,lts):lts])>0)
    }
    return(p.miscarriage)
}



##' Generate survival probability given epitope distribution and pathogen distribution
##' for the current population of 'offspring'
##'
##' @param kid.dist       matrix with rows equal to n offspring; and columns set by nbins; contains number of 'self' epitopes
##' @param kid.threshold  kid's threshold, stored as an index in bins
##' @param pathogen.dist     vector of number of pathogens in each bin (output of genPathogen)
##' @param mother.memory     matrix with rows equal to n kids; and columns set by nbins; contains memory to each pathogen epitope from mother
##' @param params            a list containing mortality hazards, and a named value of 'transfer' which defines probability mother is transferred
##'
##' @return a vector with as many values as kids, containing the probability of survival
##'
survProbFirstYear <- function(kid.dist,kid.threshold,pathogen.dist,
                              mother.memory=rep(0,length=20),
                              params=c(mu.b=1/30,mu.i=0.1,mu.d=0.5,mu.di=0.3, transfer=1, no.immunity=0)) {


    lts <- length(kid.dist[1,])
    ix <- mother.memory*0
    surv <- rep(0,nrow(kid.dist))
    for (j in 1:nrow(kid.dist)) {

        #print(j)
        ## First, for every bin, if pathogen present
        ix[j,] <- (pathogen.dist>0)*(1-params["transfer"]*mother.memory[j,])

        if (params["no.immunity"]==0) {
            haz.disease <- sum(ix[j,1:kid.threshold[j]]*params["mu.d"])
            haz.id <- sum(ix[j,min(kid.threshold[j]+1,lts):lts]*params["mu.di"])
            haz.immunopath <- sum((kid.dist[j,]>1)[min(kid.threshold[j]+1,lts):lts]*params["mu.i"])
        } else {
            haz.disease <- sum(ix[j,1:lts]*params["mu.d"])
            haz.id <- 0
            haz.immunopath <- 0
        }

	## Probabilty of survival
	surv[j] <- exp(-min(sum(params["mu.b"]+haz.disease+haz.id+haz.immunopath),100))

    }

    return(surv)

}


##' Generate 'outcomes' given epitope distribution and pathogen distribution
##' for the current population of 'offspring'
##'
##' @param kid.dist       matrix with rows equal to n offspring; and columns set by nbins; contains number of 'self' epitopes
##' @param kid.threshold  kid's threshold, stored as an index in bins
##' @param pathogen.dist     vector of number of pathogens in each bin (output of genPathogen)
##' @param mother.memory     matrix with rows equal to n kids; and columns set by nbins; contains memory to each pathogen epitope from mother
##' @param params            a list containing mortality hazards, and a named value of 'transfer' which defines probability mother is transferred
##'
##' @return a vector with as many values as kids, containing the probability of survival
##'

propOutcomesFirstYear <- function(kid.dist,kid.threshold,pathogen.dist,
                              mother.memory=rep(0,length=20),
                              params=c(mu.b=1/30,mu.i=0.1,mu.d=0.5,mu.di=0.3, transfer=1, no.immunity=0)) {


    lts <- length(kid.dist[1,])
    ix <- mother.memory*0

   haz.disease <- haz.id <- haz.immunopath <- 0

    for (j in 1:nrow(kid.dist)) {

        #print(j)
        ## First, for every bin, if pathogen present
        ix[j,] <- (pathogen.dist>0)*(1-params["transfer"]*mother.memory[j,])

        if (params["no.immunity"]==0) {
            haz.disease <- haz.disease+sum(ix[j,1:kid.threshold[j]]*params["mu.d"])
            haz.id <- haz.id + sum(ix[j,min(kid.threshold[j]+1,lts):lts]*params["mu.di"])
            haz.immunopath <- haz.immunopath +sum((kid.dist[j,]>1)[min(kid.threshold[j]+1,lts):lts]*params["mu.i"])
        } else {
            haz.disease <- haz.disease+sum(ix[j,1:lts]*params["mu.d"])
            haz.id <- haz.id + 0
            haz.immunopath <- haz.immunopath +0
        }

    }

     return(c(haz.disease,haz.id,haz.immunopath)/nrow(kid.dist))

}


##' Iterate a population through time
##'
##' @param nepitope        number of 'self' epitopes
##' @param mean.pop        population mean
##' @param sd.pop          population sd
##' @param npath           number of pathogens in the year
##' @param mean.path       their mean, along the continuous epitope scale
##' @param sd.path         their sd, likewise
##' @param params          a names vector containing mortality hazards, 'memory' (probability previously seen epitopes are 'remembered', reducing future burden)
##'                        'transfer' (probability that maternal memory is transferred), 'miscarry' (probability of miscarriage if one epitope of child exceeds
##'                        the mothers' threshold), 'no.immunity' indicates whether <1 year olds have a functional immune system or not
##' @param threshold.start desired starting threshold, or index of mother's threshold, defaults to null (in which case generated randomly)
##' @param mutation        vector from which change for each child's threshold is sampled (proportion of zeros defines how rare mutation is)
##' @param nstart          number of starting individuals
##' @param nind            max pop size
##' @param bins            the bins into which the values will be discretized
##' @param Tmax            desired number of time-steps
##'
##' @return a list including a data-frame, named dff, containing time-specific values;
##'

iterateTime <- function(nepitope=100, mean.pop=-5,sd.pop=1,
                        npath=10, mean.path=5,sd.path=1,
                        params=c(mu.b=1/20,mu.i=0.1,mu.d=0.5,mu.di=0.3,
                        miscarry=1,memory=0, transfer=0,no.immunity=0),
                        threshold.start=NULL,
                        mutation=c(-1,1,rep(0,98)),
                        nstart=1000,nind=10000,bins=seq(-8,8,length=20),Tmax=200){


    ## store individuals
    store.kids <- store.mothers <- store.mothers.memory <- matrix(0,nind,length(bins))

    ## store through tie
    pop.size <- threshold.mean<-threshold.var <- mother.survive <- child.miscarried <- child.survive <-rep(NA,Tmax)
    outcomes.kids <- outcomes <- matrix(NA,Tmax,3) ## in order disease, id, i

    ## initiate the population
    store.mothers[1:nstart,]  <- genKid(nepitope=nepitope, mean.mother=rnorm(nstart,mean.pop,sd.pop),
                                        mean.pop=mean.pop,sd.pop=sd.pop, bins=bins)
    threshold <- rep(NA,nind)
    if (is.null(threshold.start)){
        threshold[1:nstart] <- sample(1:length(bins),size=nstart,replace=TRUE)
    } else {
        threshold[1:nstart] <- threshold.start
    }

    threshold.mean[1] <- mean(threshold[1:nstart])
    threshold.var[1] <- var(threshold[1:nstart])
    pop.size[1] <- nstart

    store.mothers.memory[1:nstart,] <- 0
    n.now <- nstart

    for (t in 2:Tmax) {

       # print(t)

        ## generate this years' pathogens
        path <- genPathogen(npath=npath, mean=mean.path,sd=sd.path, bins=bins)

        ## see what mothers survival probability is (updating adaptive immunity)
        psurv <- survProb(mother.dist=store.mothers[1:n.now,],mother.threshold=threshold[1:n.now],
                          pathogen.dist=path,
                          mother.memory=store.mothers.memory[1:n.now,],
                          params=params)

        ## sanity check
      #  if (sum(is.na(psurv))>0) {
       #     print('bad survival')
        #    tt <- which(is.na(psurv),arr.ind=TRUE)
         #   print(store.mothers.memory[tt,])
          #  break()
        #}


        outcomes[t,] <- propOutcomes(mother.dist=store.mothers[1:n.now,],mother.threshold=threshold[1:n.now],
                                     pathogen.dist=path,
                                     mother.memory=store.mothers.memory[1:n.now,],
                                     params=params)

        ## update mother's memory
        store.mothers.memory[1:n.now,] <- motherMemory(mother.dist=store.mothers[1:n.now,], mother.threshold=threshold[1:n.now],
                                                        pathogen.dist=path,
                                                        mother.memory=store.mothers.memory[1:n.now,],params=params)

        ## see what mothers chance carry kid to term (first make kid, etc)
        kid.dist <- genKid(nepitope=nepitope,
                           mean.mother=rep(mean.pop,n.now),mean.pop=mean.pop,sd.pop=sd.pop,
                           bins=bins)
        #allow mutation
        kid.threshold <- pmin(pmax(threshold[1:n.now]+sample(mutation, size=n.now, replace=TRUE),1),length(bins))

        ## see if kids survive first year (depends on maternal antibodies)
        pmiscarry <- miscarriage(mother.dist=store.mothers[1:n.now,],
                                 mother.threshold=threshold[1:n.now],kid.dist=kid.dist,params=params)

        #if (sum(is.na(pmiscarry))>0) {
         #   print('bad miscarry')
          #  print(store.mothers[1:n.now,])
           # break()
        #}

        psurv.kids <- survProbFirstYear(kid.dist=kid.dist,kid.threshold=kid.threshold,
                                        pathogen.dist=path,mother.memory=store.mothers.memory[1:n.now,],
                                        params=params)

        outcomes.kids[t,] <- propOutcomesFirstYear(kid.dist=kid.dist,kid.threshold=kid.threshold,
                                                   pathogen.dist=path,mother.memory=store.mothers.memory[1:n.now,],
                                                   params=params)


        ## remove dead mothers
        survived <- rbinom(length(psurv),1,psurv)
        mother.survive[t] <- sum(survived)/length(survived)  #store before adjust
        if (sum(survived)==0) survived[sample(1:length(survived),size=min(20,length(survived)),replace=FALSE)] <- 1
        store.mothers[1:sum(survived),] <- store.mothers[which(survived==1,arr.ind=TRUE),]
        store.mothers.memory[1:sum(survived),] <- store.mothers.memory[which(survived==1,arr.ind=TRUE),]
        threshold[1:sum(survived)] <- threshold[which(survived==1,arr.ind=TRUE)]

        ## add kids to population of mothers
        miscarried <- rbinom(length(pmiscarry),1,pmiscarry)
        child.miscarried[t] <- sum(miscarried)/length(pmiscarry)  #store before adjust
        surviving.kids <- rbinom(length(psurv),1,psurv.kids)*(1-miscarried)
        child.survive[t] <- sum(surviving.kids)/length(surviving.kids) #store before adjust
        if (sum(surviving.kids)==0) surviving.kids[sample(1:length(surviving.kids),size=min(20,length(survived)),replace=FALSE)] <- 1

        ## fix so doesn't exceed pop threshold
        if ((sum(survived)+sum(surviving.kids))>nind) {
            #print((sum(survived)+sum(surviving.kids)))
            #print('excess')
            excess <- (sum(survived)+sum(surviving.kids))-nind+1
            #print(excess)
            idx <- sample(1:sum(surviving.kids),size=excess,replace=FALSE)
            surviving.kids[which(surviving.kids==1,arr.ind=TRUE)[idx]] <- 0
            #print((sum(survived)+sum(surviving.kids)))

        }

        if (sum(surviving.kids)>0) {
            store.mothers[(sum(survived)+1):(sum(survived)+sum(surviving.kids)),] <- kid.dist[which(surviving.kids==1,arr.ind=TRUE),]
            threshold[(sum(survived)+1):(sum(survived)+sum(surviving.kids))] <- kid.threshold[which(surviving.kids==1,arr.ind=TRUE)]
            store.mothers.memory[(sum(survived)+1):(sum(survived)+sum(surviving.kids)),] <- 0
        }
        ## update the number
        n.now <- (sum(survived)+sum(surviving.kids))
        if(n.now>nind) n.now <- nind

        ## store threshold moments
        threshold.mean[t] <- mean(threshold[1:n.now])
        threshold.var[t] <- var(threshold[1:n.now])

        ## pop size; and mother dist
        pop.size[t] <- n.now


    }

    ## time series
    dff <- data.frame(threshold.mean=threshold.mean,
                      threshold.var=threshold.var,child.survive=child.survive,
                      child.miscarried=child.miscarried,mother.survive=mother.survive,pop.size=pop.size)

    ## current pop dist
    final.mean <- sum(colSums(store.mothers[1:n.now,])*bins)/sum(store.mothers[1:n.now,])
    final.dist <- colSums(store.mothers[1:n.now,])/sum(store.mothers[1:n.now,])

    final.memory <- sum(colSums(store.mothers.memory[1:n.now,])/length(bins))

    return(list(dff=dff,params=params, bins=bins, final.mean=final.mean, final.dist=final.dist,final.memory=final.memory,
                outcomes=outcomes,outcomes.kids=outcomes.kids))
}



############ BELOW LOOPS that run evaluations of these functions, assuming kids i) do and ii) don't have timmunity ###################################



## Start at different thresholds and identify hazards, find survivors, etc.
##

runDifferentThresh <- function(){

    ## control parameters
    par <- c(mu.b=1/1000,mu.i=0.02,mu.d=0.045,mu.di=0.04,
                    miscarry=0.01,memory=0, transfer=0,no.immunity=0)
    mean.pop <- -4; sd.pop <- 2.2
    mean.path <- 4; sd.path <- 2.2
    bins <-seq(-8,8,length=30)
    npath <- 20
    Tmax <- 250
    nstart <- 800
    nind <- 1000

    thresh.test <- seq(1,length(bins),length=length(bins))

    #storage
    m.survive <- k.survive <- k.miscarry <- matrix(NA,length(thresh.test),5)
    outcomes <- outcomes.kids <- array(dim=c(length(thresh.test),5,3))

    for (j in 1:length(thresh.test)) {
        print(j)
        par.now <- par
        a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                          npath=npath, mean.path=mean.path,sd.path=sd.path,
                          params=par.now,
                          threshold.start=thresh.test[j],
                          mutation=c(rep(0,2)),
                          nstart=nstart,nind=nind,
                          bins=bins,Tmax=Tmax)
        m.survive[j,1] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
        k.miscarry[j,1] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
        k.survive[j,1] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
        outcomes[j,1,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
        outcomes.kids[j,1,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])

        par.now["memory"] <- 1
        a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                          npath=npath, mean.path=mean.path,sd.path=sd.path,
                          params=par.now,
                          threshold.start=thresh.test[j],
                          mutation=c(rep(0,2)),
                          nstart=nstart,nind=nind,
                          bins=bins,Tmax=Tmax)
        m.survive[j,2] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
        k.miscarry[j,2] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
        k.survive[j,2] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
        outcomes[j,2,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
        outcomes.kids[j,2,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])

        par.now["transfer"] <- 1
        a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                          npath=npath, mean.path=mean.path,sd.path=sd.path,
                          params=par.now,
                          threshold.start=thresh.test[j],
                          mutation=c(rep(0,2)),
                          nstart=nstart,nind=nind,
                          bins=bins,Tmax=Tmax)
        m.survive[j,3] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
        k.miscarry[j,3] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
        k.survive[j,3] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
        outcomes[j,3,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
        outcomes.kids[j,3,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])

        par.now["transfer"] <- 0
        par.now["no.immunity"] <- 1
        a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                          npath=100, mean.path=mean.path,sd.path=sd.path,
                          params=par.now,
                          threshold.start=thresh.test[j],
                          mutation=c(rep(0,2)),
                          nstart=nstart,nind=nind,
                          bins=bins,Tmax=Tmax)
        m.survive[j,4] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
        k.miscarry[j,4] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
        k.survive[j,4] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
        outcomes[j,4,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
        outcomes.kids[j,4,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])


        par.now["transfer"] <- 1
        par.now["no.immunity"] <- 1
        a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                          npath=100, mean.path=mean.path,sd.path=sd.path,
                          params=par.now,
                          threshold.start=thresh.test[j],
                          mutation=c(rep(0,2)),
                          nstart=nstart,nind=nind,
                          bins=bins,Tmax=Tmax)
        m.survive[j,5] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
        k.miscarry[j,5] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
        k.survive[j,5] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
        outcomes[j,5,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
        outcomes.kids[j,5,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])


    }

    ## store originals, and create smoothed to remove stochastic wiggles.
    m.survive.original <- m.survive;
    k.survive.original <- k.survive;
    for (jj in 1:5) {
        m.survive[,jj] <- smooth.spline(m.survive[,jj])$y
        k.survive[,jj] <-  smooth.spline(k.survive[,jj])$y
    }

    #find maxima
    mm <- mk <- m  <- rep(NA,5)
    for (j in 1:5) {
        m[j] <- which(m.survive[,j]==max(m.survive[,j]))
        mk[j] <- which(k.survive[,j]==max(k.survive[,j]))
        mm[j] <- which(m.survive[,j]+k.survive[,j]==max(m.survive[,j]+k.survive[,j]))
    }

    ## plot out the surviving mothers and children across thresholds; for memory and not
    cols <- c("black","black","purple","black","purple"); ltys <- c(1,3,1,2,3)
    par(mfrow=c(3,1))
    matplot(a1$bins[thresh.test],m.survive, type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Prob mothers survive", col=cols,lty=ltys)#, ylim=c(0,1))
    abline(v=0,lty=3)
    points(a1$bins[thresh.test][m],(m.survive)[cbind(m,1:5)], pch=19,col=cols)
    matplot(a1$bins[thresh.test],k.survive, type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Prob children survive", col=cols,lty=ltys)#, ylim=c(0,1))
    points(a1$bins[thresh.test][mk],(k.survive)[cbind(mk,1:5)], pch=19,col=cols)
    abline(v=0,lty=3)
    hist(pmax(pmin(rnorm(100,mean.pop,sd.pop),max(a1$bins)),min(a1$bins)),breaks=a1$bins, col="black", main="", xlab="", ylab="")
    hist(pmax(pmin(rnorm(100,mean.path,sd.path),max(a1$bins)),min(a1$bins)), breaks=a1$bins,
         col=gray.colors(1, alpha=0.5), add=TRUE, main="", xlab="", ylab="")

    ## Integrate the two in picture for paper  - present as survival
    par(mfrow=c(2,1))
    matplot(a1$bins[thresh.test],0.5*(m.survive+k.survive), type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Survival", col=cols,lty=ltys)#, ylim=c(0,1))
    abline(v=0,lty=3)
    points(a1$bins[thresh.test][mm],0.5*(m.survive+k.survive)[cbind(mm,1:5)], pch=19,col=cols)
    hist(pmax(pmin(rnorm(100,mean.pop,sd.pop),max(a1$bins)),min(a1$bins)),breaks=a1$bins, col="black", main="", xlab="", ylab="", ylim=c(0,20))
    hist(pmax(pmin(rnorm(100,mean.path,sd.path),max(a1$bins)),min(a1$bins)), breaks=a1$bins,
         col=gray.colors(1, alpha=0.5), add=TRUE, main="", xlab="", ylab="")

    ## Integrate the two in picture for paper  - present as pop growth
    layout(matrix(c(1,1,2,1,1,2),3,2))
    par(mar=c(4,5,1,1), bty="l")
    matplot(a1$bins[thresh.test],(m.survive+k.survive), type="l", xlab=expression("Threshold, "*T[i]),
            ylab=expression("Population growth, "*lambda[t]), col=cols,lty=ltys)#, ylim=c(0,1))
    abline(v=0,lty=3)
    points(a1$bins[thresh.test][mm],(m.survive+k.survive)[cbind(mm,1:5)], pch=19,col=cols)
    hist(pmax(pmin(rnorm(200,mean.pop,sd.pop),max(a1$bins)),min(a1$bins)),breaks=a1$bins, col="black", main="", xlab="", ylab="")#, ylim=c(0,20))
    hist(pmax(pmin(rnorm(150,mean.path,sd.path),max(a1$bins)),min(a1$bins)), breaks=a1$bins, col=gray.colors(1, start=0.7,alpha=0.85),
         add=TRUE, main="", xlab="", ylab="")
    abline(v=0,lty=3)
    #box()

    ## Focus on last two - yes and no maternal immunity with no kid memory; add adult survival to show male optimals
    cols <- c("black","purple"); ltys <- c(1,1,1,2,3)
    layout(matrix(c(1,1,2,2,3,1,1,2,2,3),5,2))
    par(mar=c(4,5,1,1), bty="l")
    matplot(a1$bins[thresh.test],(m.survive)[,4:5], type="l",col=cols,lty=ltys,, xlab=expression("Threshold, "*T[i]),
            ylab=expression("Adult survival, "*s[x]))
    points(a1$bins[thresh.test][m[4:5]],(m.survive)[cbind(m[4:5],4:5)], pch=19,col=cols)
    abline(v=0,lty=1,lwd=0.5)
    matplot(a1$bins[thresh.test],(m.survive+k.survive)[,4:5], type="l", xlab=expression("Threshold, "*T[i]),
            ylab=expression("Population growth, "*lambda[t]), col=cols,lty=ltys)#, ylim=c(0,1))
    points(a1$bins[thresh.test][mm[4:5]],(m.survive+k.survive)[cbind(mm[4:5],4:5)], pch=19,col=cols)
    abline(v=0,lty=1,lwd=0.5)
    hist(pmax(pmin(rnorm(200,mean.pop,sd.pop),max(a1$bins)),min(a1$bins)),breaks=a1$bins, col="black", main="", xlab="", ylab="")#, ylim=c(0,20))
    hist(pmax(pmin(rnorm(150,mean.path,sd.path),max(a1$bins)),min(a1$bins)), breaks=a1$bins, col=gray.colors(1, start=0.7,alpha=0.85),
         add=TRUE, main="", xlab="", ylab="")
    abline(v=0,lty=1,lwd=0.5)


    ## Check out the 'outcomes'
    cols <- c("black","grey","grey"); ltys <- c(1,2,3)
    par(mfrow=c(2,4), bty="l")
    ylims <- range(outcomes)
    matplot(a1$bins[thresh.test],outcomes[,1,], type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Different hazards experienced", col=cols, lty=ltys,lwd=2, ylim=ylims)
    legend("topleft", legend=c(expression(mu[d]),expression(mu[id]),expression(mu[i])), col=cols, lty=ltys, bty="n",lwd=2)
    title("Baseline")
    matplot(a1$bins[thresh.test],outcomes[,2,], type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Different hazards experienced", col=cols, lty=ltys,lwd=2, ylim=ylims)
    title("+ memory")
    matplot(a1$bins[thresh.test],outcomes[,3,], type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Different hazards experienced", col=cols, lty=ltys,lwd=2, ylim=ylims)
    title("+ maternal immunity")
    matplot(a1$bins[thresh.test],outcomes[,5,], type="l", xlab=expression("Threshold, "*T[i]),
            ylab="Different hazards experienced", col=cols, lty=ltys,lwd=2, ylim=ylims)
    title("+ immune ontogeny")

    ylims <- range(outcomes.kids)
    matplot(a1$bins[thresh.test],outcomes.kids[,1,], type="l", xlab=expression("Threshold, "*T[i]), ylab="Different hazards experienced",
            col=cols, lty=ltys,lwd=2, ylim=ylims)
    matplot(a1$bins[thresh.test],outcomes.kids[,2,], type="l", xlab=expression("Threshold, "*T[i]), ylab="Different hazards experienced",
            col=cols, lty=ltys,lwd=2, ylim=ylims)
    matplot(a1$bins[thresh.test],outcomes.kids[,3,], type="l", xlab=expression("Threshold, "*T[i]), ylab="Different hazards experienced",
            col=cols, lty=ltys,lwd=2, ylim=ylims)
    matplot(a1$bins[thresh.test],outcomes.kids[,5,], type="l", xlab=expression("Threshold, "*T[i]), ylab="Different hazards experienced",
            col=cols, lty=ltys,lwd=2, ylim=ylims)


}



## Same as above, but now assume that kids have no immunity

effectMaternalImmIfKidsUnprotected <- function(){

    ## control parameters
    par <- c(mu.b=1/1000,mu.i=0.02,mu.d=0.045,mu.di=0.04,
                    miscarry=0.01,memory=0, transfer=0,no.immunity=0)
    mean.pop <- -4; sd.pop <- 2.2
    mean.path <- 4; sd.path <- 2.2
    bins <-seq(-5.5,4.5,length=20)
    npath <- 20
    Tmax <- 250
    nstart <- 800
    nind <- 1000

    thresh.test <- seq(1,length(bins),length=length(bins))

    ## set utp have adaptive immunity and ontogeny
    par["memory"] <- 1
    par["no.immunity"] <- 1


    ## range on mu.d and mu.di
    test.mu.d <- seq(0.01,0.1,length=5)
    test.mu.di <- seq(0.01,0.1,length=5)


    #storage
    m.survive.smooth <- k.survive.smooth <- m.survive <- k.survive <- k.miscarry <- array(dim=c(length(thresh.test),2,length(test.mu.d),length(test.mu.di)))
    outcomes <- outcomes.kids <- array(dim=c(length(thresh.test),2,length(test.mu.d),length(test.mu.di),3))
    max.mother <- max.child <- max.both <- array(dim=c(2,length(test.mu.d),length(test.mu.di)))

    for(jj in 1:length(test.mu.d)) {
        for(kk in 1:length(test.mu.di)) {
        par.now <- par
        par["mu.d"] <- test.mu.d[jj]
        par["mu.di"] <- test.mu.di[kk]

        for (j in 1:length(thresh.test)) {
            print(j)
            par.now <- par
            par.now["transfer"] <- 0

            a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                              npath=npath, mean.path=mean.path,sd.path=sd.path,
                              params=par.now,
                              threshold.start=thresh.test[j],
                              mutation=c(rep(0,2)),
                              nstart=nstart,nind=nind,
                              bins=bins,Tmax=Tmax)
            m.survive[j,1,jj,kk] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
            k.miscarry[j,1,jj,kk] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
            k.survive[j,1,jj,kk] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
            outcomes[j,1,jj,kk,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
            outcomes.kids[j,1,jj,kk,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])


            par.now["transfer"] <- 1
            a1 <- iterateTime(nepitope=100, mean.pop=mean.pop,sd.pop=sd.pop,
                              npath=npath, mean.path=mean.path,sd.path=sd.path,
                              params=par.now,
                              threshold.start=thresh.test[j],
                              mutation=c(rep(0,2)),
                              nstart=nstart,nind=nind,
                              bins=bins,Tmax=Tmax)
            m.survive[j,2,jj,kk] <-mean(a1$dff$mother.survive[floor(Tmax/2):Tmax])
            k.miscarry[j,2,jj,kk] <- mean(a1$dff$child.miscarried[floor(Tmax/2):Tmax])
            k.survive[j,2,jj,kk] <- mean(a1$dff$child.survive[floor(Tmax/2):Tmax])
            outcomes[j,2,jj,kk,] <- colMeans(a1$outcomes[floor(Tmax/2):Tmax,])
            outcomes.kids[j,2,jj,kk,] <- colMeans(a1$outcomes.kids[floor(Tmax/2):Tmax,])


        }
        m.survive.smooth[,1,jj,kk] <- smooth.spline(m.survive[,1,jj,kk])$y
        k.survive.smooth[,1,jj,kk] <- smooth.spline(k.survive[,1,jj,kk])$y

        m.survive.smooth[,2,jj,kk] <- smooth.spline(m.survive[,2,jj,kk])$y
        k.survive.smooth[,2,jj,kk] <- smooth.spline(k.survive[,2,jj,kk])$y

        max.mother[1,jj,kk] <- which(m.survive.smooth[,1,jj,kk]==max(m.survive.smooth[,1,jj,kk] ))
        max.child[1,jj,kk]  <- which(k.survive.smooth[,1,jj,kk]==max(k.survive.smooth[,1,jj,kk]))
        max.both[1,jj,kk]  <-  which(m.survive.smooth[,1,jj,kk]+k.survive.smooth[,1,jj,kk]==max(m.survive.smooth[,1,jj,kk]+k.survive.smooth[,1,jj,kk]))

        max.mother[2,jj,kk] <- which(m.survive.smooth[,2,jj,kk]==max(m.survive.smooth[,2,jj,kk] ))
        max.child[2,jj,kk]  <- which(k.survive.smooth[,2,jj,kk]==max(k.survive.smooth[,2,jj,kk]))
        max.both[2,jj,kk]  <-  which(m.survive.smooth[,2,jj,kk]+k.survive.smooth[,2,jj,kk]==max(m.survive.smooth[,2,jj,kk]+k.survive.smooth[,2,jj,kk]))

        print(c(jj,kk))
        print(max.both[,jj,kk])

    }
    }

    ## positive means that with maternal immunity, sensitivity is lower
    par(mfrow=c(1,3))
    image(test.mu.d,test.mu.di,max.both[1,,]-max.both[2,,])
    image(test.mu.d,test.mu.di,max.child[1,,]-max.child[2,,])
    image(test.mu.d,test.mu.di,max.mother[1,,]-max.mother[2,,])

    par(mfrow=c(3,2))
    matplot((outcomes.kids[,1,1,1,]), type="l", col=1:3,lty=1:3, ylim=range(outcomes.kids[,1:2,1,1,]))
    matplot((outcomes.kids[,2,1,1,]), type="l", col=1:3,lty=1:3)

    par(mfrow=c(4,2))
    for (k in 1:5) {
        matplot(k.survive[,1,k,], type="l", ylim=range(k.survive[,1:2,k,1]), col=1, xlab="Threshold", ylab="Child survival")
        matplot(k.survive[,2,k,], type="l", col=4, add=TRUE)
        title(test.mu.d[k])

        matplot(m.survive[,1,k,], type="l", ylim=range(m.survive[,1:2,k,1]), col=1, xlab="Threshold", ylab="Mother survival")
        matplot(m.survive[,2,k,], type="l", col=4, add=TRUE)
    }

    ## separate for talk
    cols <- colorRampPalette(c("grey","red"))(4)
    par(mfrow=c(2,1), mar=c(3,4,1,2), bty="l")
    matplot(k.survive[,2,1,], type="n", ylim=range(k.survive[,1:2,,1],na.rm=TRUE), col=1, xlab="", ylab="Child survival")
    for (k in 1:5)   matplot(k.survive[,2,k,], type="l", col=cols[k], add=TRUE)
    par(mar=c(5,4,1,2))
    matplot(m.survive[,2,1,], type="l", ylim=range(k.survive[,1:2,,1],na.rm=TRUE), col=1, xlab="Threshold", ylab="Mother survival")
    for (k in 1:5)   matplot(m.survive[,2,k,], type="l", col=cols[k], add=TRUE)

    par(mfrow=c(2,1), mar=c(3,4,1,2), bty="l")
    matplot(k.survive[,1,1,], type="n", ylim=range(k.survive[,1:2,,1],na.rm=TRUE), col=1, xlab="", ylab="Child survival")
    for (k in 1:5)   matplot(k.survive[,1,k,], type="l", col=cols[k], add=TRUE)
    par(mar=c(5,4,1,2))
    matplot(m.survive[,1,1,], type="l", ylim=range(k.survive[,1:2,,1],na.rm=TRUE), col=1, xlab="Threshold", ylab="Mother survival")
    for (k in 1:5)   matplot(m.survive[,1,k,], type="l", col=cols[k], add=TRUE)


    ## together for manuscript
    cols <- colorRampPalette(c(rgb(0,0,1,0), rgb(1,0,0,1)), alpha = TRUE)(5)[-1]
    par(mfrow=c(2,2), mar=c(3,4,1,2), bty="l")
    matplot(bins,k.survive.smooth[,1,1,], type="n", ylim=range(k.survive.smooth[,1:2,,1],na.rm=TRUE), col=1, xlab="", ylab="Child survival")
    for (k in 1:5)   matplot(bins,k.survive.smooth[,1,k,], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(0.5,2,length=5))
    matplot(bins,k.survive.smooth[,2,1,], type="n", ylim=range(k.survive.smooth[,1:2,,1],na.rm=TRUE),lty=1, col=1, xlab="", ylab="Child survival")
    for (k in 1:5)   matplot(bins,k.survive.smooth[,2,k,], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(0.5,2,length=5))

    par(mar=c(5,4,1,2))
    matplot(bins,m.survive.smooth[,1,1,], type="n", ylim=range(m.survive.smooth[,1:2,,1],na.rm=TRUE), col=1, xlab="Threshold", ylab="Mother survival")
    for (k in 1:5)   matplot(bins,m.survive.smooth[,1,k,], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(0.5,2,length=5))
    matplot(bins,m.survive.smooth[,2,1,], type="n", ylim=range(m.survive.smooth[,1:2,,1],na.rm=TRUE), col=1, xlab="Threshold", ylab="Mother survival")
    for (k in 1:5)   matplot(bins,m.survive.smooth[,2,k,], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(0.5,2,length=5))


    ## alternative
    layout(c(1,1,2,2,3,3,4),6,1)
    cols <- colorRampPalette(c(rgb(0,0,1,1), rgb(1,0,0,1)), alpha = TRUE)(5)[-1]
    sbst <- c(1,3,5)

    par(mar=c(3,4,2,2), bty="l")
    matplot(bins,m.survive.smooth[,1,1,], type="n", ylim=range(m.survive.smooth[,1:2,,1],na.rm=TRUE), col=1, xlab="Threshold", ylab="Mother survival")
    for (k in c(1,2,4))   matplot(bins,m.survive.smooth[,1,k,sbst], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(1,3,length=5)[sbst])
    matplot(bins,k.survive.smooth[,1,1,], type="n", ylim=range(k.survive.smooth[,1:2,,1],na.rm=TRUE), col=1, xlab="", ylab="Offspring survival")
    for (k in c(1,2,4))   matplot(bins,k.survive.smooth[,1,k,sbst], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(1,3,length=5)[sbst])
    matplot(bins,k.survive.smooth[,2,1,], type="n", ylim=range(k.survive.smooth[,1:2,,1],na.rm=TRUE),lty=1, col=1, xlab="", ylab="Offspring survival")
    for (k in c(1,2,4))   matplot(bins,k.survive.smooth[,2,k,sbst], type="l", col=cols[k], add=TRUE,lty=1,lwd=seq(1,3,length=5)[sbst])

    par(mar=c(1,4,1,2), bty="l")
    hist(rnorm(250,mean.pop,sd.pop),breaks=c(-100,a1$bins,100), col="black", main="", xlab="", ylab="", xlim=range(a1$bins), ylim=c(0,0.28))#, ylim=c(0,20))
    hist(rnorm(150,mean.path,sd.path),breaks=c(-100,a1$bins,100), col=gray.colors(1, start=0.7,alpha=0.85),
         add=TRUE, main="", xlab="", ylab="")
    abline(v=0,lty=1,lwd=0.5)


}