library(MASS)
require(MCMCpack)

# Generating some multiple trait data

ntrait = 4
# Contemporary Group Covariance Matrix

CG = matrix(data=c(1200, 400, 200, 20,
                   400, 600, 80, -10,
                  200, 80, 100, -5,
                  20, -10, -5, 16),byrow=TRUE, ncol=4)
CG
CG2 = CG

# Residual Covariance Matrix
RR = matrix(data=c(8400, 3200, 1600, 160,
                   3200, 4200, 640, -100,
                   1600, 640, 700, -40,
                    160, -100, -40, 130),byrow=TRUE,ncol=4)
RR

# Genetic Covariance Matrix
GG = matrix(data=c(4500, -1099, -1055, 505,
                   -1099, 2300, -1025, 1008,
                   -1055, -1025, 370, 110,
                   505, 1008, 110, 63),byrow=TRUE, ncol=4)
GG
GG2 = GG

#  Check that each covariance matrix is positive definite

GE = eigen(GG)
GE$values
# if all are > 0, then matrix is positive definite, OK to use
nre = length(GE$values)
for(i in 1:nre){
   qp = GE$values[i]
   if(qp < 0)qp = (qp*qp)/10000
   GE$values[i] = qp  }
Gh = GE$vectors
GG2 = Gh %*% diag(GE$values) %*% t(Gh)
GG2
GG   # If variances are 'known better' then keep them unchanged and
# reduce covariances until eigenvalues are positive

#  Check residual matrix
RE = eigen(RR)
RE$values
RR2 = RR
# if one or more eigenvalues are negative, then they have to be 'repaired'
nre = length(RE$values)
for (i in 1:nre){
  qp = RE$values[i]
  if(qp < 0)qp = (qp*qp)/10000
  RE$values[i] = qp }

RE$values
Rh = RE$vectors

# New covariance matrix created using new eigenvalues
RR2 = Rh %*% diag(RE$values) %*% t(Rh)
RR2
RR

#  Check the Contemporary Group covariance matrix
CGE = eigen(CG)
CGE$values

#  Repairing the negative eigenvalues (squared and divided by 1000)
nre = length(CGE$values)
for (i in 1:nre){
  qp = CGE$values[i]
  if(qp < 0)qp = (qp*qp)/10000
  CGE$values[i] = qp
}
CGE$values
CGh = CGE$vectors
CG2 = CGh %*% diag(CGE$values) %*% t(CGh)
CG2
CG

# Get Cholesky decompositions of each covariance matrix
Gh = chol(GG2)
Rh = chol(RR2)
Ch = chol(CG2)
Rh
Ch
GGi = ginv(GG2)
CGi = ginv(CG2)

#  Simple pedigree for demonstration purposes

anim = c("G","F","H","K","L","M","P","Q","S","T")
sire = c("A","A","B","B","C","C","C","H","K","A")
dam = c("D","E","D","F","E","F","G","M","P","L")
cbind(anim,sire,dam)
sl = levels(as.factor(sire))
dl = levels(as.factor(dam))
al = levels(as.factor(anim))
cl = c(sl,dl,al)
cl
aid = levels(as.factor(cl))
aid
nanim = length(aid)
nanim

k = match(aid,anim)
sid = sire[k]
sid
did = dam[k]
pedc = data.frame(aid,sid,did)
pedc

# Number animals - for loopings
aaid = match(aid,aid)
ssid = match(sid,aid)
ddid = match(did,aid)
ped = data.frame(aaid,ssid,ddid)
ped
bii = c(1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2)

# Simulation of data for Multiple Trait model
#  Model:  Age, CG, animal, residual

# Assign contemporary groups to animals - randomly here using 'sample'
ncg = c(1:3)
cg = sample(ncg,nanim,replace=TRUE)
cg

nage = c(1:4)
age = sample(nage,nanim,replace=TRUE)
age

# CG is a random factor, covariance matrix given above 
nrn = length(ncg)*ntrait
W1 = matrix(data=c(rnorm(nrn,0,1)),ncol=ntrait)
cgfx = W1 %*% Ch 
cgfx

# Age Effects - fixed
agefx = matrix(data=c(800, 500, 300, 200,
                      900, 580, 370, 250,
                     1000, 660, 440, 300,
                     1100, 740, 510, 350),byrow=TRUE,ncol=4)
agefx
 
# True genetic values
nrn = nanim*ntrait
W1 = matrix(data=c(rnorm(nrn,0,1)),ncol=ntrait)
tbv = W1 %*% Gh
tbv
#  build in the genetic relationships among animals
for(i in 1:nanim){
   kis = ssid[i]
   kid = ddid[i]
    dii = sqrt(1/bii[i])
   if(!is.na(kis)){  
      pa = 0.5*(tbv[kis,]+tbv[kid,])+dii*tbv[i,] 
      tbv[i,] = pa
               } 
    }  
tbv # Note difference from first tbv
var(tbv)

#  Residual effects
W2 = matrix(data=c(rnorm(nrn,0,1)),ncol=ntrait)
rsd = W2 %*% Rh
var(rsd)

ddid[is.na(ddid)] = nanim + 1
ssid[is.na(ssid)] = nanim + 1
aa = rbind(tbv, c(0,0,0,0))
aa

# Suppose age effects are not in model for trait 3
agefx[ ,3] = 400
agefx

# Create the observations
yyy = agefx[age,] + cgfx[cg,] + tbv[aaid,] + rsd
yyy
obs = round(yyy)
obs
var(obs)
obs[ssid > nanim] = NA
obs

# Suppose some traits are missing
obs[7,4] = NA
obs[10,2] = NA
obs[13,1] = NA
obs[13,3] = NA
obs
nanimp = nanim + 1

# make a data from of all info
bats = data.frame(aaid,ssid,ddid,bii,cg,age,obs)
bats

#  Preliminary preparations
#  Set up some arrays before iteration

b2 = bats[ (ssid < nanimp), ]
b2   # data.frame only for animals with data

cgfac = b2$cg   # cg indicator array
anfac = b2$aaid # animal indicator array
agefac = b2$age # age indicator array
yy = cbind(b2$X1, b2$X2, b2$X3, b2$X4) # observation array
cgfac
anfac
agefac
yy

# Set up a coded pedigree list, and order it
#  Coded list is not needed to solve MME, but it is
#  necessary for proper Gibbs sampling

nped = 0
aaa = rep(0,3*nanim)
sss = aaa
ddd = aaa
cods = aaa
nanimp = nanim + 1

for(i in 1:nanim){
   ka = aaid[i]
   ks = ssid[i]
   kd = ddid[i]
	nped = nped + 1
      cods[nped]=0
      aaa[nped] = ka
      sss[nped] = ks
      ddd[nped] = kd
    if(ks < nanimp){
       nped = nped + 1
       cods[nped]=1
       aaa[nped] = ks
       sss[nped] = ka
       ddd[nped] = kd }
    if(kd < nanimp){
       nped = nped + 1
       cods[nped]=2
       aaa[nped] = kd
       sss[nped] = ka
       ddd[nped] = ks }
  }
aaa = aaa[1:nped]
cods = cods[1:nped]
sss = sss[1:nped]
ddd = ddd[1:nped]
kord = order(aaa)
aaa[kord] 

adiag = rep(0,nanim+1)       # diagonals of A-inverse
bii = c(bii,1)               # bii values for animals
#  Set up diagonals of A-inverse
for(i in 1:nanim){
  kis = ssid[i]
  kid = ddid[i]
  d = bii[i]
  x = 0.25*d
  adiag[i] = adiag[i] + d
  adiag[kis] = adiag[kis] + x
  adiag[kid] = adiag[kid] + x
}
adiag

nrec = length(anfac)
nrec

#  initialize solutions
sage = matrix(data=c(0),ncol=ntrait,nrow=length(nage))
sanm = matrix(data=c(0),ncol=ntrait,nrow=nanimp)
scg = matrix(data=c(0),ncol=ntrait,nrow=length(ncg))

# Function to determine appropriate R inverse, given obs on animal 8
y1 = yy[8,]
y1
rinv = function(y1,RR2) {
      ff = as.numeric(!is.na(y1)) 
      pp = diag(ff) %*% RR2 %*% diag(ff)
      ppi = diag(ff) %*% ginv(pp) %*% diag(ff)
 }
ri = rinv(y1, RR2)
ri

# Function to determine Variances of predicted missing residuals
rvar = function(ri, RR2) {
    a = RR2 %*% ri %*% RR2
    b = RR2 - a
# get the Cholesky decomposition of b
#  Note that the 'chol' function does not work when there
#  are some rows, columns that are all 0
    BH = LRShalf(b)
    }

# Cholesky decomposition, allowing for 0 rows and columns
LRShalf = function(a) {
    m = nrow(a)
    BH = a*0
    for(i in 1:m){
        im = i - 1
        ip = i + 1
        x = a[i,i]   # diagonals should always be positive
        if(x < 0.00001)x = 0
        im = i - 1
        if(im > 0){
           for( j in 1:im){
             x = x - BH[j,i]*BH[j,i]
             }  # for (j loop)
            }  # if statement
         z = 0
        if(x > 0){
           BH[i,i] = sqrt(x)
            z = 1/BH[i,i]
           }  # if statement
        if(ip < m+1){
          for(j in ip:m){
             x = a[i,j]
             if(im > 0){
               for(k in 1:im){
                 x = x - BH[k,j]*BH[k,i]
                  }  # for( k loop)
                }  # if statement
              BH[i,j] = x*z
         }   # for(j loop)
   }  # if statement
  } # for(i loop)
      BH
 }# End bracket
 




#  Begin Gibbs Sampling
# 
source("multiGS.R")
# plot results - messy
# Get average of sample values
    H = rep(0,16)
    V = rep(0,16)
     for(i in 1:16){
       H[i] = mean(VG[,i])
       V[i] = var(VG[,i]) }  

plot(VG[,1])


      




# Compute EBVs
source("multiMME.R")
# Compute index values
# Plot index value against different traits