Function to model selection at a biallelic locus

This is not for phytools, but I am teaching evolutionary biology to undergrads this semester. We are presently studying population genetics and for class tomorrow I decided to write an R function to illustrate selection at one locus. The function takes a starting allele frequency and fitnesses for each of the three genotypes as input, and then it can plot various things: the frequency of A or a over time; the mean fitness through time; and the population mean fitness as a function of p. The function is as follows:

selection<-function(p0,w,time=1000,show="p",pause=0){
  p<-wbar<-vector(); p[1]<-p0
  wbar[1]<-p[1]^2*w[1]+2*p[1]*(1-p[1])*w[2]+(1-p[1])^2*w[3]
  for(i in 2:time){
    p[i]<-p[i-1]
    p[i]<-(p[i]^2*w[1]+p[i]*(1-p[i])*w[2])/wbar[i-1]
    wbar[i]<-p[i]^2*w[1]+2*p[i]*(1-p[i])*w[2]+(1-p[i])^2*w[3]
    if(show=="p")
      plot(1:i,p,type="l",xlim=c(0,time),ylim=c(0,1),
      xlab="time",main="frequency of A")
    else if(show=="q")
      plot(1:i,1-p,type="l",xlim=c(0,time),ylim=c(0,1),xlab="time",
      ylab="q",main="frequency of a")
    else if(show=="fitness")
      plot(1:i,wbar/max(w),type="l",xlim=c(0,time),ylim=c(0,1),
      xlab="time",main="mean fitness")
    else if(show=="surface")
      break
    else {
      message("not a recognized option")
      break
    }
    Sys.sleep(pause)
  }
  if(show=="surface"){
    p<-0:100/100
    wbar<-p^2*w[1]+2*p*(1-p)*w[2]+(1-p)^2*w[3]
    plot(p,wbar,type="l",ylim=c(0,1),main="mean fitness")
  }

Here is an example of a run to illustrate selection on A in which the fitness advantage of the A allele is recessive.