# The function requires imput for starting allele frequencies for A1 (here A2) and S1 (Here S2)
# Also requires imput for migration rate, probability of assortative mating, probability of selfing, selection against maladpated alleles, selection against hybrids, and inbreeding depression 

#Current output at equilibrium includes allele frequencies, time until fixation (in generations) and theta, which is a measure of purebred matings/reproductive isolation

assort<-function(A2,S2,m,a,s,sl,sh,si){


## Variables that are free to change
#population size, actually 1/2 of pop size
Ninit<-10000
K<-10000
gen<-1000
#Description of alleles
# L is locally adapted allele L1 is adapted to environment 1 L2 to the opposite environment. Some fitness cost to having wrong L
# M and N are DMI genes that when not paired N1M1 suffer some disadvantage
# A is positive assorative mating locus at some probability A2 means assortative mating occurs
# S is selfing or not. Probability of selfing by S2 individuals
# 0 allele is L1, 1 allele is L2
freqL<-c(0.99,0.01)
# 0 allele is N1, 1 allele is N2
freqN<-c(0.99,0.01)
# 0 allele is M1, 1 allele is M2
freqM<-c(0.99,0.01)
# 0 allele is A1, 1 allele is A2
freqA<-c(1-A2,A2)
# 0 allele is S1, 1 allele is S2
freqS<-c(1-S2,S2)
L1<-sample(0:1,Ninit, replace=T, prob=freqL)
N1<-sample(0:1,Ninit, replace=T, prob=freqN)
M1<-sample(0:1,Ninit, replace=T, prob=freqM)
L2<-sample(0:1,Ninit, replace=T, prob=1-freqL)
N2<-sample(0:1,Ninit, replace=T, prob=1-freqN)
M2<-sample(0:1,Ninit, replace=T, prob=1-freqM)
A<-sample(0:1,Ninit, replace=T, prob=freqA)
S<-sample(0:1,Ninit, replace=T, prob=freqS)
pop1_init<-cbind(L1,N1,M1,A,S)
pop2_init<-cbind(L2,N2,M2,A,S)
## End of variables


#Start loop here:
pop1<-pop1_init
pop2<-pop2_init


#matrix for allele frequencies
pop1freq<-c(sum(pop1[,1])/length(pop1[,1]),sum(pop1[,2])/length(pop1[,2]),sum(pop1[,3])/length(pop1[,3]),sum(pop1[,4])/length(pop1[,4]),sum(pop1[,5])/length(pop1[,5]), length(pop1[,1]))

pop2freq<-c(sum(pop2[,1])/length(pop2[,1]),sum(pop2[,2])/length(pop2[,2]),sum(pop2[,3])/length(pop2[,3]),sum(pop2[,4])/length(pop2[,4]),sum(pop2[,5])/length(pop2[,5]),length(pop2[,1]))

for (i in 1:gen){

#Sample individuals that will migrate before mating
m1<-rep()
rm1<-rep()
for (l1 in 1:length(pop1[,1])){
	mig<-sample(0:1,1,prob=c(m,1-m))
	 if (mig==0) {m1<-rbind(m1,pop1[l1,]); rm1<-c(rm1,l1)}
	}
if (length(rm1)==0) {pop1<-pop1} else {pop1<-pop1[-rm1,]}
m2<-rep()
rm2<-rep()
for (l2 in 1:length(pop2[,1])){
	mig<-sample(0:1,1,prob=c(m,1-m))
	if (mig==0) {m2<-rbind(m2,pop2[l2,]); rm2<-c(rm2,l2)} 
	}
if (length(rm2)==0) {pop2<-pop2} else {pop2<-pop2[-rm2,]}

#New populations after migration has occured
pop1m<-rbind(pop1,m2)
pop2m<-rbind(pop2,m1)

#Calculate total number of "local homozygotes" to calcualte number of "purebred matings"
dt1<-0
for (i in 1:length(pop1m[,1])){
	if (pop1m[i,2]+pop1m[i,3]==0) {dt1<-dt1+1}
}
dt2<-0
for (i in 1:length(pop2m[,1])){
	if (pop2m[i,2]+pop2m[i,3]==2) {dt2<-dt2+1}
}



#Sort populations into mating units
#Remove individuals that have the potential to self
pop1s2<-rep()
pop2s2<-rep()
out1<-rep()
out2<-rep()
for (k1 in 1:length(pop1m[,1])){
	ind1<-pop1m[k1,]
	if (ind1[5]==1) {pop1s2<-rbind(ind1,pop1s2); out1<-c(k1,out1) }
}
for (k2 in 1:length(pop2m[,1])){
	ind2<-pop2m[k2,]
	if (ind2[5]==1) {pop2s2<-rbind(ind2,pop2s2) ; out2<-c(k2,out2)}
	}
if (length(out1)==0) {pop1ms<-pop1m} else {pop1ms<-pop1m[-out1,]}
if (length(out2)==0) {pop2ms<-pop2m} else {pop2ms<-pop2m[-out2,]}

# Sample individuals that will self

rms1<-rep()
for (s1 in 1:length(pop1s2[,1])){
	self<-sample(0:1,1,prob=c(1-s,s))
	 if (self==0) {rms1<-c(rms1,s1)}
	}
# these are the individuals that do not self and go back into general population
if (length(rms1)>0) {
pop1ns<-pop1s2[rms1,]
# these are the individuals that will self
pop1s2<-pop1s2[-rms1,]
}
if (length(rms1)==0){
pop1ns<-rep()
pop1s2<-pop1s2
}


rms2<-rep()
for (s1 in 1:length(pop2s2[,1])){
	self<-sample(0:1,1,prob=c(1-s,s))
	 if (self==0) {rms2<-c(rms2,s1)}
	}
# these are the individuals that do not self and go back into general population
if (length(rms2)>0){
pop2ns<-pop2s2[rms2,]
# these are the individuals that will self
pop2s2<-pop2s2[-rms2,]
}
if (length(rms2)==0){
pop2ns<-rep()
pop2s2<-pop2s2
}


#population after non selfing individuals are added back in
pop1ms<-rbind(pop1ms,pop1ns)
pop2ms<-rbind(pop2ms,pop2ns)

#Remove individuals who have the potential to mate assoratively
#assortative mating individuals stored in the first two vectors
pop1a2<-rep()
pop2a2<-rep()
#place holders for individuals being removed
out1<-rep()
out2<-rep()
for (k1 in 1:length(pop1ms[,1])){
	ind1<-pop1ms[k1,]
	assort1<-sample(0:1,1,prob=c(1-a,a))
	if (ind1[4]==1 & assort1==1) {pop1a2<-rbind(ind1,pop1a2); out1<-c(k1,out1) }
}
for (k2 in 1:length(pop2ms[,1])){
	ind2<-pop2ms[k2,]
	assort2<-sample(0:1,1,prob=c(1-a,a))
	if (ind2[4]==1 & assort2==1) {pop2a2<-rbind(ind2,pop2a2) ; out2<-c(k2,out2)}
	}
#population that will not mate assoratively
if (length(out1)==0) {pop1msa<-pop1ms} else {pop1msa<-pop1ms[-out1,]}
if (length(out2)==0) {pop2msa<-pop2ms} else {pop2msa<-pop2ms[-out2,]}

#Need to distinguish between different DM combinations throw hybrids into random mating group?
pop1a2N0M0<-rep()
pop1a2N1M1<-rep()
pop1a2h<-rep()
pop2a2N0M0<-rep()
pop2a2N1M1<-rep()
pop2a2h<-rep()
if (length(pop1a2)>0){
for (j1 in 1:length(pop1a2[,1])){
	ind<-pop1a2[j1,]
	if (ind[2]+ind[3]==0) {pop1a2N0M0<-rbind(ind,pop1a2N0M0)}
	if (ind[2]+ind[3]==2) {pop1a2N1M1<-rbind(ind,pop1a2N1M1)}
	if (ind[2]+ind[3]==1) {pop1a2h<-rbind(ind,pop1a2h)}
	}
} else {pop1a2N0M0<-rep() ; pop1a2N1M1<-rep() ; pop1a2h<-rep()}
if (length(pop2a2)>0){
for (j1 in 1:length(pop2a2[,1])){
	ind<-pop2a2[j1,]
	if (ind[2]+ind[3]==0) {pop2a2N0M0<-rbind(ind,pop2a2N0M0)}
	if (ind[2]+ind[3]==2) {pop2a2N1M1<-rbind(ind,pop2a2N1M1)}
	if (ind[2]+ind[3]==1) {pop2a2h<-rbind(ind,pop2a2h)}
	}	
} else {pop2a2N0M0<-rep() ; pop2a2N1M1<-rep() ; pop2a2h<-rep()}
if (length(pop1a2h)!=0) {pop1msa<-rbind(pop1msa,pop1a2h)}
if (length(pop2a2h)!=0) {pop2msa<-rbind(pop2msa,pop2a2h)}

#Calculate the number of purebreds that randomly mate and their frequency
dtr1<-0
for (i in 1:length(pop1msa[,1])){
	if (pop1msa[i,2]+pop1msa[i,3]==0) {dtr1<-dtr1+1}
	}
lmsa1<-length(pop1msa[,1])
fr1<-(lmsa1*(dtr1/lmsa1)^2+(dt1-dtr1))/dt1


dtr2<-0
for (i in 1:length(pop2msa[,1])){
	if (pop2msa[i,2]+pop2msa[i,3]==2) {dtr2<-dtr2+1}
	}
lmsa2<-length(pop2msa[,1])
fr2<-(lmsa2*(dtr2/lmsa2)^2+(dt2-dtr2))/dt2

#recap of breeding "groups" for each population
# pop 1= pop1s2 are individuals that will self
#		pop1a2N0M0, pop1a2N1M1 will mate assortatively
#		pop1msa will mate randomly
# have the same population groups for population 2 as well pop2s2, pop2a2N0M0,pop2a2N1M1,pop2msa

#Selfing groups will just be recapitulated as themselves just need to put in death. They will make density dependent number of offspring. It will be 2X what outcrosses make
#Because it is assumed outcrossers make one progeny as mom and one as dad (or they make 2 progeny as well)

birth1<-1-(length(pop1[,1])-K)/length(pop1[,1])
if (birth1<1) {birth1<-1}
#if (birth1>3) {birth1<-3}

birth2<-1-(length(pop2[,1])-K)/length(pop2[,1])
if (birth2<1) {birth2<-1}
#if (birth2>3) {birth2<-3}


#double the length of the selfers
rms1<-rep()
if (length(pop1s2)>0){
if (length(pop1[,1])>2*K) {pop1s2<-pop1s2} else {pop1s2<-rbind(pop1s2,pop1s2)}
for (x1 in 1:length(pop1s2)){
	d<-sample(0:1,1,prob=c(1-si,si))
	if (d==1) {rms1<-rbind(rms1,x1)}
	}
}
if (length(rms1)>0){pop1sng<-pop1s2[-rms1,]} else {pop1sng<-pop1s2}

rms2<-rep()
if (length(pop2s2)>0){
if (length(pop2[,1])>2*K) {pop2s2<-pop2s2} else {pop2s2<-rbind(pop2s2,pop2s2)}
for (x2 in 1:length(pop2s2)){
	d<-sample(0:1,1,prob=c(1-si,si))
	if (d==1) {rms2<-rbind(rms2,x2)}
	}
}
if (length(rms2)>0) {pop2sng<-pop2s2[-rms2,]} else {pop2sng<-pop2s2}	

#Mating within assortative groups
#Calculate allel frequencies for pop1a2N0M0
# need an if statement in case this genotype is not present
if (length(pop1a2N0M0[,1])==0) {pop1a2N0M0ng<-rep()}
if (length(pop1a2N0M0[,1]>0)){
	fL2<-sum(pop1a2N0M0[,1])/length(pop1a2N0M0[,1])
	fA<-sum(pop1a2N0M0[,4])/length(pop1a2N0M0[,1])
	fS<-sum(pop1a2N0M0[,5])/length(pop1a2N0M0[,1])
	#Make frequencies into list
	fL<-c(1-fL2,fL2)
	fN<-c(1,0)
	fM<-c(1,0)
	fA<-c(1-fA,fA)
	fS<-c(1-fS,fS)

	#Generate allele vectors
	n1L<-sample(0:1,length(pop1a2N0M0[,1])*birth1, replace=T, prob=fL)
	n1N<-sample(0:1,length(pop1a2N0M0[,1])*birth1, replace=T, prob=fN)
	n1M<-sample(0:1,length(pop1a2N0M0[,1])*birth1, replace=T, prob=fM)
	n1A<-sample(0:1,length(pop1a2N0M0[,1])*birth1, replace=T, prob=fA)
	n1S<-sample(0:1,length(pop1a2N0M0[,1])*birth1, replace=T, prob=fS)
	#join together alleles	
	pop1a2N0M0ng<-cbind(n1L,n1N,n1M,n1A,n1S)
}
#Calculate frequencies for pop1a2N1M1
if (length(pop1a2N1M1[,1])==0) {pop1a2N1M1ng<-rep()}
if (length(pop1a2N1M1[,1]>0)){
	fL2<-sum(pop1a2N1M1[,1])/length(pop1a2N1M1[,1])
	fA<-sum(pop1a2N1M1[,4])/length(pop1a2N1M1[,1])
	fS<-sum(pop1a2N1M1[,5])/length(pop1a2N1M1[,1])
	fL<-c(1-fL2,fL2)
	fN<-c(0,1)
	fM<-c(0,1)
	fA<-c(1-fA,fA)
	fS<-c(1-fS,fS)
	n2L<-sample(0:1,length(pop1a2N1M1[,1])*birth1, replace=T, prob=fL)
	n2N<-sample(0:1,length(pop1a2N1M1[,1])*birth1, replace=T, prob=fN)
	n2M<-sample(0:1,length(pop1a2N1M1[,1])*birth1, replace=T, prob=fM)
	n2A<-sample(0:1,length(pop1a2N1M1[,1])*birth1, replace=T, prob=fA)
	n2S<-sample(0:1,length(pop1a2N1M1[,1])*birth1, replace=T, prob=fS)
	pop1a2N1M1ng<-cbind(n2L,n2N,n2M,n2A,n2S)
}

#Calculate allel frequencies for pop2a2N0M0
if (length(pop2a2N0M0[,1])==0) {pop2a2N0M0ng<-rep()}
if (length(pop2a2N0M0[,1]>0)){
	fL2<-sum(pop2a2N0M0[,1])/length(pop2a2N0M0[,1])
	fA<-sum(pop2a2N0M0[,4])/length(pop2a2N0M0[,1])
	fS<-sum(pop2a2N0M0[,5])/length(pop2a2N0M0[,1])
	fL<-c(1-fL2,fL2)
	fN<-c(1,0)
	fM<-c(1,0)
	fA<-c(1-fA,fA)
	fS<-c(1-fS,fS)
	n3L<-sample(0:1,length(pop2a2N0M0[,1])*birth2, replace=T, prob=fL)
	n3N<-sample(0:1,length(pop2a2N0M0[,1])*birth2, replace=T, prob=fN)
	n3M<-sample(0:1,length(pop2a2N0M0[,1])*birth2, replace=T, prob=fM)
	n3A<-sample(0:1,length(pop2a2N0M0[,1])*birth2, replace=T, prob=fA)
	n3S<-sample(0:1,length(pop2a2N0M0[,1])*birth2, replace=T, prob=fS)
	pop2a2N0M0ng<-cbind(n3L,n3N,n3M,n3A,n3S)
}

#Calculate frequencies for pop2a2N1M1
if (length(pop1a2N1M1[,1])==0) {pop2a2N1M1ng<-rep()}
if (length(pop2a2N1M1[,1]>0)){
	fL2<-sum(pop2a2N1M1[,1])/length(pop2a2N1M1[,1])
	fA<-sum(pop2a2N1M1[,4])/length(pop2a2N1M1[,1])
	fS<-sum(pop2a2N1M1[,5])/length(pop2a2N1M1[,1])
	fL<-c(1-fL2,fL2)
	fN<-c(0,1)
	fM<-c(0,1)
	fA<-c(1-fA,fA)
	fS<-c(1-fS,fS)
	n4L<-sample(0:1,length(pop2a2N1M1[,1])*birth2, replace=T, prob=fL)
	n4N<-sample(0:1,length(pop2a2N1M1[,1])*birth2, replace=T, prob=fN)
	n4M<-sample(0:1,length(pop2a2N1M1[,1])*birth2, replace=T, prob=fM)
	n4A<-sample(0:1,length(pop2a2N1M1[,1])*birth2, replace=T, prob=fA)
	n4S<-sample(0:1,length(pop2a2N1M1[,1])*birth2, replace=T, prob=fS)
	pop2a2N1M1ng<-cbind(n4L,n4N,n4M,n4A,n4S)
}

#mating within large population

#Calculate frequencies for pop1msa
fL2<-sum(pop1msa[,1])/length(pop1msa[,1])
fN2<-sum(pop1msa[,2])/length(pop1msa[,1])
fM2<-sum(pop1msa[,3])/length(pop1msa[,1])
fA<-sum(pop1msa[,4])/length(pop1msa[,1])
fS<-sum(pop1msa[,5])/length(pop1msa[,1])
fL<-c(1-fL2,fL2)
fN<-c(1-fN2,fN2)
fM<-c(1-fM2,fM2)
fA<-c(1-fA,fA)
fS<-c(1-fS,fS)
p1L<-sample(0:1,length(pop1msa[,1])*birth1, replace=T, prob=fL)
p1N<-sample(0:1,length(pop1msa[,1])*birth1, replace=T, prob=fN)
p1M<-sample(0:1,length(pop1msa[,1])*birth1, replace=T, prob=fM)
p1A<-sample(0:1,length(pop1msa[,1])*birth1, replace=T, prob=fA)
p1S<-sample(0:1,length(pop1msa[,1])*birth1, replace=T, prob=fS)
pop1msang<-cbind(p1L,p1N,p1M,p1A,p1S)

#Calculate frequencies for pop2msa
fL2<-sum(pop2msa[,1])/length(pop2msa[,1])
fN2<-sum(pop2msa[,2])/length(pop2msa[,1])
fM2<-sum(pop2msa[,3])/length(pop2msa[,1])
fA<-sum(pop2msa[,4])/length(pop2msa[,1])
fS<-sum(pop2msa[,5])/length(pop2msa[,1])
fL<-c(1-fL2,fL2)
fN<-c(1-fN2,fN2)
fM<-c(1-fM2,fM2)
fA<-c(1-fA,fA)
fS<-c(1-fS,fS)
p2L<-sample(0:1,length(pop2msa[,1])*birth2, replace=T, prob=fL)
p2N<-sample(0:1,length(pop2msa[,1])*birth2, replace=T, prob=fN)
p2M<-sample(0:1,length(pop2msa[,1])*birth2, replace=T, prob=fM)
p2A<-sample(0:1,length(pop2msa[,1])*birth2, replace=T, prob=fA)
p2S<-sample(0:1,length(pop2msa[,1])*birth2, replace=T, prob=fS)
pop2msang<-cbind(p2L,p2N,p2M,p2A,p2S)

# bring together all progeny
if (length(pop1a2N0M0ng)==0) {pop1a2N0M0ng<-rep()}
if (length(pop1a2N1M1ng)==0) {pop1a2N1M1ng<-rep()}
if (length(pop1sng)==0) {pop1sng<-rep()}
pop1ng<-rbind(pop1a2N0M0ng,pop1a2N1M1ng,pop1msang,pop1sng)
if (length(pop2a2N0M0ng)==0) {pop2a2N0M0ng<-rep()}
if (length(pop2a2N1M1ng)==0) {pop2a2N1M1ng<-rep()}
if (length(pop2sng)==0) {pop2sng<-rep()}
pop2ng<-rbind(pop2a2N0M0ng,pop2a2N1M1ng,pop2msang,pop2sng)


#Death in large population after adding back in selfing individuals and combining other selfing groups
#Death due to hybrid status, or local maladaptation
dth1<-rep()
for (r1 in 1:length(pop1ng[,1])){
	ind<-pop1ng[r1,]
	sl1<-sample(0:1,1,prob=c(1-sl,sl))
	sh1<-sample(0:1,1,prob=c(1-sh,sh))
	if (ind[1]==1 & sl1==1) {dth1<-rbind(dth1,r1)}
	if (ind[2]+ind[3]==1 & sh1==1) {dth1<-rbind(dth1,r1)}
}
if (length(dth1)==0) {pop1ngd<-pop1ng} else {pop1ngd<-pop1ng[-dth1,]}

dth2<-rep()
for (r2 in 1:length(pop2ng[,1])){
	ind<-pop2ng[r2,]
	sl2<-sample(0:1,1,prob=c(1-sl,sl))
	sh2<-sample(0:1,1,prob=c(1-sh,sh))
	if (ind[1]==0 & sl2==1) {dth2<-rbind(dth2,r2)}
	if (ind[2]+ind[3]==1 & sh2==1) {dth2<-rbind(dth2,r2)}
}
if (length(dth2)==0) {pop2ngd<-pop2ng} else {pop2ngd<-pop2ng[-dth2,]}

pop1<-pop1ngd
pop2<-pop2ngd

pop1freq_i<-c(sum(pop1[,1])/length(pop1[,1]),sum(pop1[,2])/length(pop1[,2]),sum(pop1[,3])/length(pop1[,3]),sum(pop1[,4])/length(pop1[,4]),sum(pop1[,5])/length(pop1[,5]),fr1)

pop2freq_i<-c(sum(pop2[,1])/length(pop2[,1]),sum(pop2[,2])/length(pop2[,2]),sum(pop2[,3])/length(pop2[,3]),sum(pop2[,4])/length(pop2[,4]),sum(pop2[,5])/length(pop2[,5]),fr2)

pop1freq<-rbind(pop1freq,pop1freq_i)
pop2freq<-rbind(pop2freq,pop2freq_i)

}
print(pop1freq_i)
print(pop2freq_i)
#closing the function bracket
}