#################################################################################################
#### This is a simulation to demonstrate how populations reach Hardy Weinburg equilibrium
#### under random mating.
#### Author: Corey Chivers, 2011
#################################################################################################

cross<-function(parents)
{
	offspring<-c('d','d') #initiate a child object
	offspring[1]<-sample(parents[1,],1)
	offspring[2]<-sample(parents[2,],1)

	return(offspring)
}

random_mating<-function()
{
	tmp_pop<-pop
	for(n in 1:N)
	{
		parents<-sample(1:N,2)
		tmp_pop[n,]<-cross(pop[parents,])
	}
	pop<-tmp_pop
}


#############################################
## Change the population size or 
## number of generations of random mating
## and see what happens!
N=200
num_generations=1
##############################################



genotypes<-c('A','a')
p<-0.5
q<-1-p
a_freq<-c(p,q)

#Initiate random population out of HW
pop<-array(sample(genotypes,2*N,p=a_freq,replace=T),dim=c(N,2))

# Number of populations to sim
I<-200

g_freq<-array(dim=c(I,3))
p_vec<-array(dim=I)
for(i in 1:I)
{
	p<-runif(1,0,1)
	q<-1-p
	a_freq<-c(p,q)
	
	pop[,1]<-sample(genotypes,N,p=a_freq,replace=T)
	pop[,2]<-sample(genotypes,N,p=a_freq,replace=T)

	for(g in 1:num_generations)
		random_mating()


	f_aa<-0
	f_Aa<-0
	f_AA<-0

	for(n in 1:N)
	{
		if(identical(pop[n,],c('A','A')))
			f_AA=f_AA+1
		if(identical(pop[n,],c('A','a')) || identical(pop[n,],c('a','A') ))
			f_Aa=f_Aa+1
		if(identical(pop[n,],c('a','a')))
			f_aa=f_aa+1
	
	}
	f_aa<-f_aa/N
	f_Aa<-f_Aa/N
	f_AA<-f_AA/N
	
	g_freq[i,]<-c(f_AA,f_Aa,f_aa)
	p_vec[i]<-p	
}


## Plot the sims
plot(p_vec,g_freq[,1],col='green',xlab='p, or 1-q',ylab='')
points(p_vec,g_freq[,2],col='red')
points(p_vec,g_freq[,3],col='blue')


## Theoretical Curves
curve(x^2,col='green',add=T)
curve((1-x)^2,col='blue',add=T)
curve(2*x*(1-x),col='red',add=T)

## Legend
legend('topleft',bg='white',legend=c('AA','Aa','aa'),col=c('blue','red','green'),pch=1,lty=1)