library(MASS)
require(MCMCpack)

#  Maternal Genetics effect model - creation of data
#  Simple pedigree for demonstration purposes

anim = c("G","F","H","K","L","M","P")
sire = c("A","A","B","B","C","C","C")
dam = c("D","E","D","F","E","F","G")
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)

# Simulation of data for an MATERNAL EFFECTS model
#  Model:  CG, animal, dam, dam PE, residual

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

# CG is a random factor with variance = 100
cgsd = 10
cgfx = rnorm(length(ncg),0,cgsd)
cgfx
mean(cgfx)
var(cgfx)

#  Assign animal genetic values 
#  for direct and maternal effects - random effects
g = matrix(data=c(300, -30, -30, 180),byrow=TRUE, ncol=2)
g
gi = ginv(g)
gi
gsd = chol(g)
gsd

rnd = matrix(data=rnorm(2*nanim,0,1),ncol=2)
rnd
tbv = rnd %*% gsd
tbv
var(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)

# Maternal Permanent Environmental effects
pesd = sqrt(66)
pefx = rnorm(nanim,0,pesd)
pefx = c(pefx, 0)
pefx

#  Residual effects
rsd = sqrt(700)
resfx = rnorm(nanim,0,rsd)

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

# With Maternal Effects - add direct of animal and maternal of dam
yyy = cgfx[cg] + tbv[aaid,1] + aa[ddid,2] + pefx[ddid] + resfx
yyy
obs = round(yyy)
obs
mean(obs)
var(obs)
obs[ssid > nanim] = NA
obs

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

# Constructing and Solving Mixed Model Equations by Iteration on Data

#  Preliminary preparations
#  Set up some arrays before iteration

b2 = bats[ !is.na(obs), ]
b2   # data.frame only for animals with data
cgfac = b2$cg   # cg indicator array
anfac = b2$aaid # animal indicator array
damfac = b2$ddid # maternal effects  and mat PE array
y = b2$obs              # observation vector
cgfac
anfac
damfac
y

# 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] 

zz = rep(0,nanim)          # diagonals of Z'Z for direct effects
zz
k = b2$aaid
zz[k] = 1
zz

mm = rep(0,nanim)          # diagonals of M'M for maternal effects
k=b2$ddid
kfac=factor(k)
kfac
kfl = as.numeric(levels(kfac))
kfl
nk = tapply(y,k,length)
nk
mm[kfl]=nk
mm

adiag = rep(0,nanim+1)       # diagonals of A-inverse
bii = c(bii,1)               # bii values for animals


#  initialize solutions
scg = tapply(y,cgfac,mean)   # simple means for CG - forced to zero sum
scg
sanm = rep(0,nanimp)         # Initial animal solutions set to 0
smat = rep(0,nanimp)         # Maternal
smpe = rep(0,nanim)         # Maternal PE
gi = 700*ginv(g)               # Variance ratios for direct and maternal
ratpe = 700/66
gi
ratpe

#  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
zold = c(scg,sanm,smat,smpe)

#  Estimate parameters

source("matGS01.R")
#  Plot convergence
plot(VRhd2,type="l", col="blue", lty="solid", ylim=c(-1,1))
plot(VRhm2,type="l", col="red", lty="solid", ylim=c(-1,1),add=TRUE)
plot(VRdmc,type="l", col="green", lty="solid", ylim=c(-1,1),add=TRUE)
plot(PEV,type="l",col="blue",lty="solid",ylim=c(0,700))
plot(AA,type="l",col="blue",lty="solid",ylim=c(0,700),add=TRUE)
plot(AM,type="l",col="blue",lty="solid",ylim=c(0,700),add=TRUE)
plot(MM,type="l",col="blue",lty="solid",ylim=c(0,700),add=TRUE)

#  Do the Iterations

source("matMME01.R")

# If you have birth years, then compute average EBV by year of birth
# Plot the averages - do for both direct and maternal on one graph.