#test example for multiple breed genomic prediction using GLS + Selection Index
# AL, 5 Oct 2021
rm(list=ls())
set.seed(1234)
# number of animals per population
n1=100;n2=120;n3=90
# number of markers
m=1000
# draw allele frequencies from "not-too-different populations" with correlated allele frequencies
CorFreq=diag(3)*0.5+0.5
require(MASS)
fq=mvrnorm(m,c(0,0,0),CorFreq)
# put back to frequences
fq=qbeta(pnorm(fq),2,2)
cor(fq)


# sample genotypes and assign {-1,0,1} coding
Z1=matrix(NA,ncol=m,nrow=n1)
Z2=matrix(NA,ncol=m,nrow=n2)
Z3=matrix(NA,ncol=m,nrow=n3)
for (i in 1:m){
  Z1[,i]=
    #1st allele
    sample(0:1,n1,replace=TRUE,prob=c(1-fq[i,1],fq[i,1]))+
    #2nd allele
    sample(0:1,n1,replace=TRUE,prob=c(1-fq[i,1],fq[i,1]))+
    (-1)
  Z2[,i]=
    #1st allele
    sample(0:1,n2,replace=TRUE,prob=c(1-fq[i,2],fq[i,2]))+
    #2nd allele
    sample(0:1,n2,replace=TRUE,prob=c(1-fq[i,2],fq[i,2]))+
    (-1)
  Z3[,i]=
    #1st allele
    sample(0:1,n3,replace=TRUE,prob=c(1-fq[i,3],fq[i,3]))+
    #2nd allele
    sample(0:1,n3,replace=TRUE,prob=c(1-fq[i,3],fq[i,3]))+
    (-1)
  
}
# assign {0,1,0} coding
W1=1-Z1^2
W2=1-Z2^2
W3=1-Z3^2

X=matrix(0,ncol=6,nrow=n1+n2+n3)
X[1:n1,1]=1
X[(n1+1):(n1+n2),2]=1
X[(n1+n2+1):(n1+n2+n3),3]=1
# genomic inbreeding (homozygosity)
X[1:n1,4]=apply(W1,1,mean)
X[(n1+1):(n1+n2),5]=apply(W2,1,mean)
X[(n1+n2+1):(n1+n2+n3),6]=apply(W3,1,mean)

# variances (these numbers do not pretend to make biological sense)
vara=rWishart(1,5,diag(3))[,,1]/m
vard=rWishart(1,5,diag(3))[,,1]/m
vare=c(2,1,3)

a=mvrnorm(m,c(0,0,0),vara)
d=mvrnorm(m,c(0,0,0),vard)

# simulated phenotypes
y=rep(0,n1+n2+n3)
#1st pop
y[1:n1]=Z1%*%a[,1]
#2nd pop
y[(n1+1):(n1+n2)]=Z2%*%a[,2]
#3rd pop
y[(n1+n2+1):(n1+n2+n3)]=Z3%*%a[,3]
# same, with dominance
y[1:n1]=y[1:n1]+W1%*%d[,1]
y[(n1+1):(n1+n2)]=y[(n1+1):(n1+n2)]+W2%*%d[,2]
y[(n1+n2+1):(n1+n2+n3)]=y[(n1+n2+1):(n1+n2+n3)]+W3%*%d[,3]

var(y)

#arbitrary means and depressions
beta=matrix(c(10,20,30,-.1,-.2,-.3),ncol=1)

y=y+X%*%beta

#residual 
y[1:n1]=y[1:n1]+rnorm(n1)*sqrt(vare[1])
y[(n1+1):(n1+n2)]=y[(n1+1):(n1+n2)]+rnorm(n2)*sqrt(vare[2])
y[(n1+n2+1):(n1+n2+n3)]=y[(n1+n2+1):(n1+n2+n3)]+rnorm(n3)*sqrt(vare[3])



# write MME for SNP-BLUP
#y=L * sol
# weight incidence matrix by 1/sqrt(vare) to get MME properly
Xw=X
Xw[,1]=Xw[,1]/sqrt(vare[1])
Xw[,4]=Xw[,4]/sqrt(vare[1])
Xw[,2]=Xw[,2]/sqrt(vare[2])
Xw[,5]=Xw[,5]/sqrt(vare[2])
Xw[,3]=Xw[,3]/sqrt(vare[3])
Xw[,6]=Xw[,6]/sqrt(vare[3])
Z=rbind(
  cbind(Z1,Z1*0,Z1*0,W1,W1*0,W1*0)/sqrt(vare[1]),
  cbind(Z2*0,Z2,Z2*0,W2*0,W2,W2*0)/sqrt(vare[2]),
  cbind(Z3*0,Z3*0,Z3,W3*0,W3*0,W3)/sqrt(vare[3])
) 
L=cbind(Xw,Z)

lhs=crossprod(L)
# here I need to consider properly (again) vare
w=c(rep(vare[1],n1),rep(vare[2],n2),rep(vare[3],n3))
rhs=crossprod(L,y/sqrt(w))


# add variance components, inverted
varai=solve(vara)
vardi=solve(vard)

add <- function(eff,pop){
  # returns last and first address of eff x for this pop
  # legth of each effect
  place=(eff-1)*3+pop
  pos=c(1,1,1,1,1,1,m,m,m,m,m,m)
  first=sum(pos[0:(place-1)])+1
  last=sum(pos[0:place])
  c(first,last)  
}

#a
for (j in 1:3){
  for (k in 1:3){
    # 3 populations
    toto1=add(3,j)
    toto2=add(3,k)
    cat(toto1,toto2,"\n")
    lhs[toto1[1]:toto1[2],toto2[1]:toto2[2]]=
      lhs[toto1[1]:toto1[2],toto2[1]:toto2[2]]+diag(varai[j,k],m,m)
  }
}
# d

for (j in 1:3){
  for (k in 1:3){
    # 3 populations
    toto1=add(4,j)
    toto2=add(4,k)
    cat(toto1,toto2,"\n")
    lhs[toto1[1]:toto1[2],toto2[1]:toto2[2]]=
      lhs[toto1[1]:toto1[2],toto2[1]:toto2[2]]+diag(vardi[j,k],m,m)
  }
}
sol=solve(lhs,rhs)

# now use equivalent method
add2 <- function(pop){
  # gives first and last position in V
  place=pop
  pos=c(n1,n2,n3)
  first=sum(pos[0:(pop-1)])+1
  last=sum(pos[0:pop])
  c(first,last)  
  
} 

V=diag(
  c(rep(vare[1],n1),rep(vare[2],n2),rep(vare[3],n3))
)
#a
for (j in 1:3){
  for (k in 1:3){
    toto1=add2(j)
    toto2=add2(k)
    cat(toto1,toto2,"\n")
    dum1=get(paste("Z",j,sep=""))
    dum2=get(paste("Z",k,sep=""))
    cat(j,k,paste("Z",j,sep=""),paste("Z",k,sep=""),"\n")
    V[toto1[1]:toto1[2],toto2[1]:toto2[2]]=
      V[toto1[1]:toto1[2],toto2[1]:toto2[2]]+dum1%*%t(dum2)*vara[j,k]    
  }
}
#d
for (j in 1:3){
  for (k in 1:3){
    toto1=add2(j)
    toto2=add2(k)
    cat(toto1,toto2,"\n")
    dum1=get(paste("W",j,sep=""))
    dum2=get(paste("W",k,sep=""))
    cat(j,k,paste("W",j,sep=""),paste("W",k,sep=""),"\n")
    V[toto1[1]:toto1[2],toto2[1]:toto2[2]]=
      V[toto1[1]:toto1[2],toto2[1]:toto2[2]]+dum1%*%t(dum2)*vard[j,k]    
  }
}

lhs=t(X)%*%solve(V)%*%X
rhs=t(X)%*%solve(V)%*%y
beta=solve(lhs,rhs)
# beta is identical to previous head(sol)

yc=solve(V)%*%(y-X%*%beta)
v=matrix(ncol=1,nrow=3*m)

add3 <- function(pop){
    c((i-1)*m+1,i*m)  
}

# extract ahat
for (i in 1:3){
  toto1=add3(i)
  toto2=add2(i)
  cat(toto1,toto2,"\n")
  dum1=get(paste("Z",i,sep=""))
  v[toto1[1]:toto1[2]]=t(dum1)%*%yc[toto2[1]:toto2[2]]
}

ahat=matrix(0,ncol=1,nrow=3*m)

ahat[1:m]=          v[1:m]*vara[1,1]+v[(m+1):(2*m)]*vara[1,2]+v[(2*m+1):(3*m)]*vara[1,3]
ahat[(m+1):(2*m)]=  v[1:m]*vara[2,1]+v[(m+1):(2*m)]*vara[2,2]+v[(2*m+1):(3*m)]*vara[2,3]
ahat[(2*m+1):(3*m)]=v[1:m]*vara[3,1]+v[(m+1):(2*m)]*vara[3,2]+v[(2*m+1):(3*m)]*vara[3,3]

# correlation with solution from SNP-BLUP
cor(ahat,sol[7:(6+3*m)])

# extract dhat
for (i in 1:3){
  toto1=add3(i)
  toto2=add2(i)
  cat(toto1,toto2,"\n")
  dum1=get(paste("W",i,sep=""))
  v[toto1[1]:toto1[2]]=t(dum1)%*%yc[toto2[1]:toto2[2]]
}

dhat=matrix(0,ncol=1,nrow=3*m)

dhat[1:m]=          v[1:m]*vard[1,1]+v[(m+1):(2*m)]*vard[1,2]+v[(2*m+1):(3*m)]*vard[1,3]
dhat[(m+1):(2*m)]=  v[1:m]*vard[2,1]+v[(m+1):(2*m)]*vard[2,2]+v[(2*m+1):(3*m)]*vard[2,3]
dhat[(2*m+1):(3*m)]=v[1:m]*vard[3,1]+v[(m+1):(2*m)]*vard[3,2]+v[(2*m+1):(3*m)]*vard[3,3]

# correlation with solution from SNP-BLUP
cor(dhat,sol[(6+3*m+1):(6+6*m)])