####################################################################################################
####################### this program is used to determine stopping tolerance #######################
####################################################################################################

rm(list=ls())

### load package ###
library(BGLR)
library(MASS)

### program paarameters ###
h2=0.5 ### heredity ###

n0=10 ### number of initial sample points ###
nsel=10 ### number of sequentially added sample points ###
nsim=5000 ### number of Monte-Carlo simulations ###

mu=100  ### mean of multivariate normal distribution ###
sg=25 ### genetic variance ###
se=sg*(1-h2)/h2 ### noise variance ###

### kinship matrix ###
kin=as.matrix(read.table(file="kin.txt",sep="",header=FALSE))
n=ncol(kin)

### simulate data ###
g=mvrnorm(nsim,rep(0,n),sg*kin)
p=mu+g+rnorm(n,0,sqrt(se))

### initial training sets ###
ini=matrix(0,nsim,n0)
for (i in 1:nsim)
{
  best.id=which(g[i,]==max(g[i,]))
  ini[i,]=sample(setdiff(1:n,best.id),n0)
}

ns=NULL
d1=NULL
d2=NULL
d3=NULL
d4=NULL
for (i in 1:nsim)
{
  ### phenotypic and genotypic values ###
  y=p[i,]
  best.id=which(g[i,]==max(g[i,]))
  
  ### training and testing sets ###
  tr=ini[i,]
  te=setdiff(1:n,tr)
  
  ### search start ###
  mtc=TRUE
  tot.ei=NULL
  while (mtc==TRUE)
  {
    ### split kinship matrix ###
    k11=kin[tr,tr]
    k12=kin[tr,te]
    k21=kin[te,tr]
    k22=kin[te,te]
    
    ### gblup model ###
    fit=suppressWarnings(BGLR(y=y[tr],ETA=list(MRK=list(K=k11,model="RKHS")),nIter=5000,burnIn=10^3,verbose=FALSE))
    sg=fit$ETA$MRK$varU 
    se=fit$varE         
    g1=fit$ETA$MRK$u
    
    ### conditional mean and conditional variance ###
    mu.g2=k21%*%solve(k11)%*%g1
    var.g2=diag(sg*(k22-k21%*%solve(k11)%*%k12))
    
    ### expected improvement ###
    uf=g1-diag(sqrt(sg)*k11)
    muf=max(uf)
    mx=g1[which(uf==muf)]
    z=(mu.g2-mx)/sqrt(var.g2)
    ei=((mu.g2-mx)*pnorm(z)+sqrt(var.g2)*dnorm(z))*(1-sqrt(se)/sqrt(var.g2+se))
    
    ### update training set and testing set ###
    out=cbind(te,round(ei,12))
    out=out[order(out[,2],decreasing=TRUE),]
    if (nrow(out)>=nsel) add=out[1:nsel,1]
    if (nrow(out)<nsel) add=out[,1]
    tr=c(tr,add)
    te=setdiff(1:n,tr)
    
    ### average ei ###
    aei=mean(ei)
    if (is.na(aei) || is.infinite(aei)) aei=10^8
    tot.ei=c(tot.ei,aei)
    
    ### stooping criterion ###
    if (best.id%in%tr && aei<0.0001) mtc=FALSE
    if (length(tr)==ncol(kin)) mtc=FALSE
  }
  ns=c(ns,which(tr==best.id))
  d1=c(d1,min(which(tot.ei<0.1)))
  d2=c(d2,min(which(tot.ei<0.01)))
  d3=c(d3,min(which(tot.ei<0.001)))
  d4=c(d4,min(which(tot.ei<0.0001)))
}

### output ###
d1=nsel*d1
d2=nsel*d2
d3=nsel*d3
d4=nsel*d4
d1[which(d1==Inf)]=round(ncol(kin),-1)
d2[which(d2==Inf)]=round(ncol(kin),-1)
d3[which(d3==Inf)]=round(ncol(kin),-1)
d4[which(d4==Inf)]=round(ncol(kin),-1)
out=cbind(ns,d1,d2,d3,d4)
apply(out,2,mean)