####################################################################################################
########## this program is used to identify the best genotype using Bayesian optimization ##########
####################################################################################################

rm(list=ls())

### load package ###
library("BGLR")

### search parameters ###
n0=10 ### number of initial samples ###
nsel=20 ### number of selected samples ###
delta=0.01  ### stopping tolerance ###

### read data ###
y=as.matrix(read.table(file=paste("phenotype.txt",sep=""),sep="",header=TRUE))  ### phenotype data ###
kin=as.matrix(read.table(file="kinship.txt",sep="",header=FALSE)) ### kinship matrix ###

### training set and testing set ###
n=nrow(y)
tr=sample(1:n,n0)
te=setdiff(1:n,tr)

### search start ###
ns=NULL
mtc=TRUE
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:nrow(out),1]
  tr=c(tr,add)
  te=setdiff(1:n,tr)
    
  ### stopping criterion ### 
  if (mean(ei)<delta) mtc=FALSE
}

### best genotype ###
k11=kin[tr,tr]
fit=suppressWarnings(BGLR(y=y[tr],ETA=list(MRK=list(K=k11,model="RKHS")),nIter=5000,burnIn=10^3,verbose=FALSE))
g=fit$ETA$MRK$u
best.id=tr[which(g==max(g))]
best.id