## ------------------------------------------------------------
## MODEL FUNCTIONS
## ------------------------------------------------------------

## create a population: note, I put this inside a function after class
create.pop <- function(N) {
x <- runif(N)
y <- runif(N)
t1 <- sample(c(1,2), size=N, replace=TRUE)
t2 <- runif(N)
return(cbind(x, y, t1, t2))
}

## plot a single population in space, coloured according to their t1
## trait
plot.population <- function(pop) {
  par(oma=rep(0,4), mar=c(0.2,0.2,0.2,0.2))
  plot(NA, xlim=c(0,1), ylim=c(0,1), xlab='', ylab='', xaxt='n', yaxt='n')
  colours <- rep('red', nrow(pop))
  colours[which(pop[,'t1']==1)] <- 'blue'
  points(x=pop[,'x'], y=pop[,'y'], col=colours, pch=16)
}

## move individuals to new random locations
move.random <- function(pop) {
  pop.size <- nrow(pop)
  pop[,'x'] <- runif(pop.size)
  pop[,'y'] <- runif(pop.size)
  return(pop)
}

## move individuals in 2D space with wrap-around boundaries in both
## dimensions
move.wraparound <- function(pop) {
  ## movement distance
  d <- rnorm(nrow(pop), mean=0, sd=sigma.m)
  ## movement angle
  a <- runif(nrow(pop), min=0, max=2*pi)
  delta.x <- cos(a) * d
  delta.y <- sin(a) * d
  pop[,'x'] <- (pop[,'x'] + delta.x) %% 1
  pop[,'y'] <- (pop[,'y'] + delta.y) %% 1
  return(pop)
}

## parents produce next generation of offspring, with fitness
## specified by their t1 trait
birth <- function(pop) {
  parents <- sample(1:nrow(pop), size=nrow(pop), replace=TRUE,
                    prob=s[pop[,'t1']])
  offspring <- pop[parents,]
  pop <- offspring
  return(pop)
}

## ------------------------------------------------------------
## SPECIFY PARAM VALUES
## ------------------------------------------------------------

## model parameters:
N <- 20 ## initial population size
sigma.m <- 0.05 ## sd of movement distribution
s <- c(1.1, 1) ## fitnesses at the t1 trait

## simulation parameters:
num.gens <- 1000

## ------------------------------------------------------------
## RUN THE MODEL!
## ------------------------------------------------------------

pop <- create.pop(N)

plot.population(pop)

## run my model!
for(t in 1:num.gens) {
  cat(sprintf('\rGeneration %d', t)) ## print generation, as model runs
  plot.population(pop)
  pop <- move.wraparound(pop)
  pop <- birth(pop)
  Sys.sleep(0.1) ## pause so that plots are given time to load
}