######################################################
####### Acrescentando Selecao no modelo de HW ########
######################################################
rm(list=ls())


#Considerando selecao, como frequencias mudam de acordo com aptidoes dos diferentes genotipos
allele.deltaP<- function(s = 0, t=0 ) {
  #valores de default: s = 0, t=0, p0 = valor de uma distribuicao uniforme entre 0 e 1, 100 passos
  WAA=1-s #aptidao de AA
  WAB=1   #aptidao de Aa, ou AB aqui
  WBB=1-t #aptidao de aa, ou BB aqui
  allele.dyn<- function(WAA,WAB,WBB,p){
    (p*(p*WAA+(1-p)*WAB))/((p^2)*WAA+2*p*(1-p)*WAB+((1-p)^2)*WBB)
    } #p'em funcao de p considerando as aptidoes (considerando selecao)
  todos.p<-seq(from=0,to=1,length.out=100) #vetor de frequencias de A
  deltap=allele.dyn(WAA,WAB,WBB,todos.p) - todos.p
  
  plot(todos.p,deltap,type='l',xlab=expression(p),ylab=expression(paste(Delta,"p" )), ylim=c(-0.1,0.1)) #plot de deltap x p
  abline(a=0,b=0, lty=2) #linha que indica onde esta deltap=0
  }

################################################################
# Caso 1: s>0 e t<0 --> aa (ou BB) tem maior fitness ##########
################################################################
s=0.3
t=-0.3
p0=0.7 #comecando com 0.3 ou 0.7 vou parar no mesmo lugar, em p=0 (portanto eq p=0 eh estavel)

par(mfrow=c(1,2))
allele.deltaP(s,t)

tmax=300
p=matrix(0,tmax,1)
p[1]=p0

for(i in 1:tmax){
  WAA=1-s #aptidao de AA
  WAB=1   #aptidao de Aa, ou AB aqui
  WBB=1-t #aptidao de aa, ou BB aqui
    p[i+1]=(p[i]*(p[i]*WAA+(1-p[i])*WAB))/((p[i]^2)*WAA+2*p[i]*(1-p[i])*WAB+((1-p[i])^2)*WBB)
    
  }
  
cinzas=  gray.colors((length(p[2:tmax])), start = 0.8, end = 0, gamma = 0.7, alpha = NULL)
points(p[1],0)
points(p[2:tmax],rep(0, length(p[2:tmax])), col=cinzas,pch=19,cex=1)

plot(p~seq(0,tmax),pch=16,xlab="tempo", ylab="p", ylim=c(0,1))
 
  
#######################################################
# Caso 2: s<0 e t>0 --> AA tem maior fitness #########
#####################################################

s=-0.3
t=0.3
p0=0.3 #comecando com 0.3 ou 0.7 vou parar no mesmo lugar, em p=1 (portanto eq p=1 eh estavel)

par(mfrow=c(1,2))
allele.deltaP(s,t)

tmax=300
p=matrix(0,tmax,1)
p[1]=p0

for(i in 1:tmax){
  WAA=1-s #aptidao de AA
  WAB=1   #aptidao de Aa, ou AB aqui
  WBB=1-t #aptidao de aa, ou BB aqui
  p[i+1]=(p[i]*(p[i]*WAA+(1-p[i])*WAB))/((p[i]^2)*WAA+2*p[i]*(1-p[i])*WAB+((1-p[i])^2)*WBB)
  
}

cinzas=  gray.colors((length(p[2:tmax])), start = 0.8, end = 0, gamma = 0.7, alpha = NULL)
points(p[1],0)
points(p[2:tmax],rep(0, length(p[2:tmax])), col=cinzas,pch=19,cex=1)

plot(p~seq(0,tmax),pch=16,xlab="tempo", ylab="p", ylim=c(0,1))

###########################################################
# Caso 3: s>0 e t>0 --> Aa (AB) tem maior fitness #########
###########################################################

s=0.3
t=0.3
p0=0.7 #comecando com 0.3 ou 0.7 vou parar no mesmo lugar, em p=t/s+t (portanto eq p=t/s+t eh estavel)

par(mfrow=c(1,2))
allele.deltaP(s,t)

tmax=300
p=matrix(0,tmax,1)
p[1]=p0

for(i in 1:tmax){
  WAA=1-s #aptidao de AA
  WAB=1   #aptidao de Aa, ou AB aqui
  WBB=1-t #aptidao de aa, ou BB aqui
  p[i+1]=(p[i]*(p[i]*WAA+(1-p[i])*WAB))/((p[i]^2)*WAA+2*p[i]*(1-p[i])*WAB+((1-p[i])^2)*WBB)
  
}

cinzas=  gray.colors((length(p[2:tmax])), start = 0.8, end = 0, gamma = 0.7, alpha = NULL)
points(p[1],0)
points(p[2:tmax],rep(0, length(p[2:tmax])), col=cinzas,pch=19,cex=1)

plot(p~seq(0,tmax),pch=16,xlab="tempo", ylab="p", ylim=c(0,1))

################################################################
# Caso 4: s<0 e t<0 --> AA e aa (BB) tem maior fitness #########
###############################################################

s=-0.3
t=-0.3
p0=0.3 #comecando com 0.3 ou 0.7 de alelo A vou parar ou em p=0 ou e p=1 (portanto estes sao eq estaveis)

par(mfrow=c(1,2))
allele.deltaP(s,t)

tmax=300
p=matrix(0,tmax,1)
p[1]=p0

for(i in 1:tmax){
  WAA=1-s #aptidao de AA
  WAB=1   #aptidao de Aa, ou AB aqui
  WBB=1-t #aptidao de aa, ou BB aqui
  p[i+1]=(p[i]*(p[i]*WAA+(1-p[i])*WAB))/((p[i]^2)*WAA+2*p[i]*(1-p[i])*WAB+((1-p[i])^2)*WBB)
  
}

cinzas=  gray.colors((length(p[2:tmax])), start = 0.8, end = 0, gamma = 0.7, alpha = NULL)
points(p[1],0)
points(p[2:tmax],rep(0, length(p[2:tmax])), col=cinzas,pch=19,cex=1)

plot(p~seq(0,tmax),pch=16,xlab="tempo", ylab="p", ylim=c(0,1))



######################################################
####### Acrescentando Mutacao no modelo de HW ########
######################################################
par(mfrow=c(1,1))

pinicial=1.0 #frequencia inicial de A

tempo.Max=100 #tempo de analise da nossa equacao
p=matrix(0,tempo.Max) #vetor das frequencias de A com o passar do tempo
p[1]=pinicial #guardando o valor de p inicial.
q=1-p

mutA=0.1 #taxa de mutacao de A->B (mu)
mutB=0.3 #taxa de mutacao de B->A (ni)


for(t in 1:(tempo.Max-1)){
  p[t+1]=(p[t])*(1-mutA)+q[t]*mutB #p'=p(1-mu)+q(ni)
  q[t+1]=1-p[t+1] #q'=1-p'
}

p #nosso vetor de freq de A a cada unidade de tempo

tempo=seq(1:tempo.Max) #definicao de cada instante de tempo

plot(p~tempo, pch=19, xlab="tempo", ylab="p(t)", xlim=c(0, tempo.Max), ylim=c(-0.1,01)) #grafico de p x tempo
points(q~tempo, pch=19,col="red") #adicionando pontos de q x tempo
legend("topright", c("p", "q"), pch=c(16,16), col = c(1,2), box.lwd = 0)