#object generators: 

# make a chromosome object
makeChrom <- function(length,states=1) list(length=length,breaks=length,states=states)

# make an individual object 
makeInd <- function(chrom_lengths,sex=NULL,rate=NULL,states=1){
	
	ind <- list()
	if (is.null(sex)) {
	   ind$sex <- sample(c(0,1),1)
	} else {
	  ind$sex <- sex
	}

	if (is.null(rate)){
	   ind$rate <- setRate(ind$sex)
	} else {
	  ind$rate <- rate
	}

	ind$chrom_lengths <- chrom_lengths
	ind$chroms <- list()

	for (i in 1:length(chrom_lengths)) ind$chroms[[i]] <- makeChrom(chrom_lengths[i])
	 
	return(ind)
}

## functions

# pick a recombination rate
setRate <- function(sex){

    if (sex == 1){
	mu <- 41.5/3000
	sigma <- 5.92/3000
    } else {
	 mu <- 26.9/3000
	 sigma <- 3.32/3000
    }

    return(rnorm(1,mu,sigma))

}

# preform recombination between a given chromosome and a WT chromosome
WTrecomb <- function(chrom,rate){

  newChrom <- list(length=chrom$length,sex=sample(c(0,1),1))
  newChrom$breaks <- c()
  newChrom$states <- c()
  
  strand = sample(c(0,1),1)
  lastBreak = 0
  nextBreak <- rexp(1,rate)
  
  while (nextBreak < chrom$length){
  
    if (strand == 1){

      inSplit <- which(chrom$breaks <= nextBreak && chrom$breaks > lastBreak)

      newChrom$breaks <- c(newChrom$breaks,chrom$breaks[inSplit])
      newChrom$states <- c(newChrom$states,chrom$states[inSplit])
      #newChrom$breaks <- append(newChrom$breaks,chrom$breaks[inSplit])
      #newChrom$states <- append(newChrom$states,chrom$states[inSplit])
      lastSplit <- inSplit[length(inSplit)]
      if (length(chrom$states) == 1) {
         newChrom$breaks <- c(newChrom$breaks,nextBreak)
        #newChrom$states <- c(newChrom$states,chrom$states)
      } else if (length(lastSplit) == 0){
        newChrom$breaks <- c(newChrom$breaks,nextBreak)
        newChrom$states <- c(newChrom$states,chrom$states[min(which(chrom$breaks > nextBreak))])
      } else if (chrom$breaks[lastSplit] != nextBreak){
        newChrom$breaks <- c(newChrom$breaks,nextBreak)
        newChrom$states <- c(newChrom$states,chrom$states[lastSplit + 1])
      }

    } else {

      newChrom$breaks <- c(newChrom$breaks,nextBreak)
      newChrom$states <- c(newChrom$states,0)
    }

    lastBreak <- nextBreak
    nextBreak <- nextBreak + rexp(1,rate)
    strand <- 1 - strand

  }

  newChrom$breaks <- c(newChrom$breaks,chrom$length)
  if (strand == 0){
     newChrom$states <- c(newChrom$states,0)
  } else {
     newChrom$states <- c(newChrom$states,chrom$states[length(chrom$states)])
  }

  return(newChrom)
}

# perform the trivial recombination
WTrecomb2 <- function(x) if (!any(x$states == 1)) return(list(states=0)) else return(list(states=sample(c(0,1),1)))
# produce offspring
WTcross <- function(ind,sex=Null){

  offspring <- makeInd(ind$chrom_lengths,sex)

  for (i in 1:length(ind$chroms)) {
      if (all(ind$chroms[[i]]$states == 0)) {
	 offspring$chroms[[i]] <- ind$chroms[[i]]
      } else {
	offspring$chroms[[i]] <- WTrecomb(ind$chroms[[i]],ind$rate)
      }
  }
  
  return(offspring)
}

# checks to see if any ancestual DNA is left
ancestorCheck <- function(ind){

  DNAleft <- FALSE

  for (i in 1:length(ind$chroms)) if (any(ind$chroms[[i]]$states == 1)) return(TRUE)

  return(FALSE)
}

## run the simulation

chrom_lengths <- c(247,243,199,191,181,171,159,146,140,135,143,132,114,106,100,88,79,76,64,62,47,50)

Ngen <- 30
Nruns <- 125000

staying <- matrix(0,ncol=Ngen,nrow=Nruns)

for (i in 1:Nruns){

    ind <- makeInd(chrom_lengths)

    if (i %% 10 == 0) print(i)

    staying[i,1] <- 1
    
    for (j in 2:Ngen){

      ind <- WTcross(ind)

      staying[i,j] <- ancestorCheck(ind) + 0

      if (staying[i,j] == 0) break

    }
}

# calculate sharing probability per generation 
prob <- apply(staying,2,mean)