## HWE example
O.AA = 89
O.Aa = 2
O.aa = 9

## G-test for HWE
HWE.Gtest = function(O.AA, O.Aa, O.aa){
    #this function takes three arguments, each of the genotypes
    n = O.AA + O.Aa + O.aa; #total number of individuals
    p = (O.AA + O.Aa/2)/n;  #allele frequency
    E.AA = n*p^2;  #expected number of hom A
    E.Aa = n*2*p*(1-p);  #expected number of hets
    E.aa = n*(1-p)^2;  #expected number of hom a

    G = 0;
    if(O.AA > 0)  #note that log(0) is undefined, and treated as 0
        G = G + 2*O.AA*log(O.AA/E.AA);
    if(O.Aa > 0)
        G = G + 2*O.Aa*log(O.Aa/E.Aa);
    if(O.aa > 0)
        G = G + 2*O.aa*log(O.aa/E.aa);
    p.value = 1-pchisq(G, 1);  #p-value comes from chi-squared distribution
    return(c(G, p.value));
}

HWE.Gtest(O.AA, O.Aa, O.aa)

## Wright-Fisher simulations

WF = function(N, p, G){ # takes 3 arguments: Population size, initial frequency, and number of generations to simulate
    t=array(,dim=G); #initialize array to empty
    t[1] = p; #initialize array with starting frequency
    for(i in 2:G){ #for loop over generations
        t[i] = rbinom(1,N,t[i-1])/N; #drift is binomial sampling!
    }
    return(t);
}

WFdemog = function(N, p, G, Gd, v){ # takes 5 arguments: Population size, initial frequency, and number of generations to simulate, generation demographic effect occurs, and magnitude of the demographic event
    t=array(,dim=G);  #initialize array to empty
    t[1] = p;  #initialize array with starting frequency
    for(i in 2:G){
        if(i == Gd){ # change the population size at this generation!
            N = N*v;
            cat("N=",N,"\n")
        }
        t[i] = rbinom(1,N,t[i-1])/N; #drift is binomial sampling!
    }
    return(t);
}

pdf("Trajectories_p=0.5.pdf",height=6, width=6, pointsize=12)
par(mfrow=c(3,3),mar=c(4,4,1,1))
for(i in 1:9){
    f = WF(100, 0.5, 200) #call WF to get trajectories
    plot(1:length(f), f, type='l', xlab="Generation", ylab="Allele frequency", main="",ylim=c(0,1))
}
dev.off()

pdf("Trajectories_exp_p=0.5.pdf",height=6, width=6, pointsize=12)
par(mfrow=c(3,3),mar=c(4,4,1,1))
for(i in 1:9){
    f = WFdemog(100, 0.5, 200, 50, 100)
    plot(1:length(f), f, type='l', xlab="Generation", ylab="Allele frequency", main="",ylim=c(0,1))
}
dev.off()


notlost=0
tries=0;
while(notlost==0){
    tries=tries+1;
    pdf("Trajectories_horse_p=0.5.pdf",height=6, width=6, pointsize=12)
    par(mfrow=c(3,3),mar=c(4,4,1,1))
    for(i in 1:9){
        f = WF(16000, 1/16000, 2500)
        plot(1:length(f), f, type='l', xlab="Generation", ylab="Allele frequency", main="")
        if(f[2500]>0){
            notlost=1;
        }
    }
    dev.off()
}
cat("tries=",tries,"\n");



fitfreq = function(q, h, s){
    p=1-q;
    return((q^2*(1+s) + p*q*(1+h*s))/( 1 + s*q*(2*h*p+q)));
}

WF.sel=function(N, q, h, s, G){
    t=array(,dim=G);
    t[1] = q;
    for(i in 2:G){
        t[i] = rbinom(1,N,fitfreq(t[i-1], h, s))/N;
    }
    return(t);
}

WF.demsel=function(N, q, h, s, G, Gd, v){
  t=array(,dim=G);
  t[1] = N*q;
  for(i in 2:G){
    if(i == Gd){
      N = N*v;
    }
    t[i] = rbinom(1,N,fitfreq(t[i-1]/N, h, s))/N;
  }
  return(t);
}

{
    pdf("Trajecotories_s=0.1_h=0.5.pdf",height=6,width=6,pointsize=12)
    par(mfrow=c(2,2), mar=c(4,4,1,1))
    N=100;
    q=1/2/N;
    h=0.5;
    s=0.1;
    G=100;
    f = matrix(,nr=100, nc=G);
    for(i in 1:nrow(f)){
        f[i,] = WF.sel(N, q, h, s, G);
    }
    ##plot 1: single trajectory
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    
    ##plot 2: 10 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:10){
        lines(1:ncol(f), f[i,], col=i)
    }

    ##plot 3: 50 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:50){
        lines(1:ncol(f), f[i,], col=i)
    }

    ##plot 4: 100 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:100){
        lines(1:ncol(f), f[i,], col=i)
    }
    dev.off();
}


{
    pdf("Trajecotories_s=0.1_h=0.5_v=100.pdf",height=6,width=6,pointsize=12)
    par(mfrow=c(2,2), mar=c(4,4,1,1))
    N=100;
    q=1/2/N;
    h=0.5;
    s=0.1;
    G=100;
    Gd=50;
    v=100;
    f = matrix(,nr=100, nc=G);
    for(i in 1:nrow(f)){
        f[i,] = WF.demsel(N, q, h, s, G, Gd, v);
    }
    ##plot 1: single trajectory
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    
    ##plot 2: 10 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:10){
        lines(1:ncol(f), f[i,], col=i)
    }

    ##plot 3: 50 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:50){
        lines(1:ncol(f), f[i,], col=i)
    }

    ##plot 4: 100 trajectories
    plot(1:ncol(f), f[1,], ylim=c(0,1), type='l', xlab="Generation", ylab="Frequency")
    for(i in 2:100){
        lines(1:ncol(f), f[i,], col=i)
    }
    dev.off();
}


## Probability of fixation

WF.fix = function(N, q, h, s){
    t=q
    g=0;
    while(t != 0 & t != 1){
        t = rbinom(1,N,fitfreq(t, h, s))/N;
        g = g+1;
    }
    return(c(t,g));
}

probFix = function(N, q, h, s, ITS){
    F=0;
    for(i in 1:ITS){
        F = F+WF.fix(N, q, h, s);
    }
    return(F/ITS);
}

N = 100;
q = 1/N;
h = 0.5;
ITS = 5000;
s = c(-0.1, -0.05, -0.01, 0, 0.01, 0.05, 0.1, 0.25, 0.5);
P = matrix(0, nc=2, nr=length(s));
for(i in 1:length(s)){
    P[i,] = probFix(N, q, h, s[i], ITS);
}

pdf("pfix.pdf", height=4, width=4, pointsize=12)
par(mar=c(4,4,1,1))
plot(s, P[,1], ylim=c(0,0.5), xlab="selection coefficient", ylab="Fixation Probability")
lines(s,(1-exp(-s))/(1-exp(-2*N*s)))
legend("topleft",c("Simulated", "(1-exp(-s))/(1-exp(-NS))"), pch=c(1,-1), lty=c(-1,1))
abline(h=0,col="grey")
abline(v=0,col="grey")
dev.off()