#Implement equations 1.40 on p. 29 of Rice's book for fertility selection
#use "AA", "Aa", and "aa" in place of "11", "12", or "22", resp.
#By R. Gomulkiewicz (16 Sep 2013)

#number of time steps to evolve
steps <- 50

#initial allele and genotype frequencies
p0 <- 0.8
q0 <- 1-p0
g0 <- list(AA = p0^2, Aa = 2*p0*q0, aa = q0^2)

#fertilities
w <- list(AAAA = 0.025, AAAa = 0.025, AaAA = 0.025, AAaa = 1, aaAA = 0.025, AaAa = 0.025, Aaaa = 0.025, aaAa = 0.025, aaaa = 0.025)

#mean fitness given genotype frequencies and fertilities
wbar <- function(g,w){g$AA^2*w$AAAA + g$AA*g$Aa*(w$AAAa + w$AaAA) + g$AA*g$aa*(w$AAaa + w$aaAA) + g$Aa^2*w$AaAa + g$Aa*g$aa*(w$Aaaa + w$aaAa) + g$aa^2*w$aaaa }


#create a matrix that holds the genotype frequencies, allele frequencies, and mean fitness at each time step
#separate columns for g$AA, g$Aa, g$aa, and p
out <-matrix(0,nrow=steps+1,ncol=4)

#initialize genotype and allele frequencies
out[1,1]<- g0$AA
out[1,2]<- g0$aa
out[1,3]<- g0$Aa
out[1,4] <- p0

#also keep track of the mean fertility (fitness)
fit <- numeric(steps)
fit[1] <- wbar(g0, w)

for(i in 1:steps){
	g <- list(AA = out[i,1],Aa = out[i,2],aa = out[i,3])
	w.bar <- wbar(g, w)
	fit[i] <- w.bar
	out[i+1,1] <- (g$AA^2*w$AAAA + (1/2)*g$AA*g$Aa*(w$AAAa + w$AaAA) + (1/4)*g$Aa^2*w$AaAa)/w.bar #new freq of AA
	out[i+1,2] <- (g$aa^2*w$aaaa + (1/2)*g$aa*g$Aa*(w$Aaaa + w$aaAa) + (1/4)*g$Aa^2*w$AaAa)/w.bar #new freq of aa
	out[i+1,3] <- 1 - out[i+1,1] - out[i+1,2]   #new freq of Aa
	out[i+1,4] <- out[i+1,1] + out[i+1,3]/2  #new freq of A
	}

par(mfrow=c(2,1))
plot(0:(steps-1), fit, type = "l", xlab="step", ylab="mean fitness")	#plot the mean fertility (fitness)

matplot(0:steps,out,type="l",xlab="step", ylab="frequency",ylim=c(0,1))  #plot genotype and allele frequencies
legend("topright",c("g.AA","g.aa","g.Aa","p"),lty=1:4,bty="n",col=1:4)