# Random regression model

# Fun5.dat - small example data.
# Fun5.ped - pedigree for data.

library(MASS)
require(MCMCpack)
# Retrieve pedigrees for the small example  Fun5.ped
zfil=file.choose()
rrped <- read.table(file=zfil,header=FALSE,
col.names=c("anim","sire","dam","bii"))
rrped[1:10,]
nrow(rrped)

# Retrieve data for animals in pedigree  Fun5.dat
zfil=file.choose()
rrdat = read.table(file=zfil,header=FALSE,
col.names=c("aid","T1","T2","T3","T4","T5","T6","T7"))
rrdat[1:10,]
nrow(rrdat)
summary(rrdat)
plot(rrdat)
PV=var(rrdat)
PV=PV[2:8,2:8]
PV

# A look at covariance functions
# Function for Legendre polynomials, given x between -1 to 1
Legdre = function(x) {
     y = rep(0,7)
       y[1] = sqrt(.5)
       y[2] = sqrt(1.5)*x
       x2 = x*x; y[3] = sqrt(2.5)*(1.5*x2 - 0.5)
       x3 = x2*x; y[4] = sqrt(3.5)*((15*x3/6)-(9*x/6))
       x4 = x3*x; y[5]=sqrt(4.5)*((105*x4/24)-(90*x2/24)+(3/8))
       x5 = x4*x; y[6]=sqrt(5.5)*((945*x5/120)-(105*x3/12)+(225*x/120))
       x6 = x5*x
       y[7]=sqrt(6.5)*((10395*x6/720)-(2835*x4/144)+(4725*x2/720)-(5/16))
       y  }

# Example use
p =  Legdre(0.4)[1:5]
p

#  Data example
tt = c(4,11,18,25,32,39,46)  # mid week days between 1 to 7 weeks(1 to 49 days)

zz = (2*(tt - 1)/48) - 1
zz
kk = matrix(data=c(0),nrow=7,ncol=7)
for(i in 1:7){
  p = zz[i]
  kk[i,] = Legdre(p) }

kk1=ginv(kk)
hh = kk1 %*% PV %*% t(kk1)

#  Assuming Order 2 Legendre is sufficient, then
kks = kk[ ,1:3]
kks
SS = t(kks) %*% kks
SSi = ginv(SS)
PVr = t(kks) %*% PV %*% kks
PVr
SSi
hh = SSi %*% PVr %*% SSi
hh
PV2 = kks %*% hh %*% t(kks)
DD = PV - PV2
PV2
PV

# Back to the Assignment - Generating some data
#  Genetic covariance matrix, GV = 0.2*PV
#  Remainder = PV - GV = PEV + diag(RV)
GV = 0.2*PV
RM = PV - GV
RV = diag(7)*3
PEV = RM - RV
eigen(PEV)
eigen(GV)
# Change to 3x3 matrices for RR model
GVr = t(kks) %*% GV %*% kks
PEVr = t(kks) %*% PEV %*% kks
HGV = SSi %*% GVr %*% SSi     # Prior values for Parameter estimation
HPEV = SSi %*% PEVr %*% SSi   # Prior values
HGVi = ginv(HGV)
HPEi = ginv(HPEV)

# Now this becomes similar to multiple trait model 
hg = chol(HGV)
hp = chol(HPEV)
b1 = 0.50
b2 = -0.0081225
mu = 24
T = c(1:49)
TT = T*T
mut = mu + b1*T + b2*TT  # mut = vector of means for each day from 1 to 49
plot(mut)

#  Pedigree preparations
aaid = rrped$anim
ssid = rrped$sire
ddid = rrped$dam
bii = rrped$bii
nanim = nrow(rrped)
ddid[ddid < 1] = nanim + 1
ssid[ssid < 1] = nanim + 1

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


#  Set up coded pedigree files
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] 


nrr = 3               # Number of random regression covariates
nrn = nanimp*nrr
W1 = matrix(data=c(rnorm(nrn,0,1)),ncol=nrr)
W2 = matrix(data=c(rnorm(nrn,0,1)),ncol=nrr)
tbv = W1 %*% hg
pefx = W2 %*% hp
tbv[nanimp, ]=0
pefx[nanimp, ]=0
#  Build in the genetic relationships among animals

for(i in 1:nanim){
   kis = ssid[i]
   kid = ddid[i]
    dii = sqrt(1/bii[i])
   if(kis > 0 ){  
      pa = 0.5*(tbv[kis,]+tbv[kid,])+dii*tbv[i,] 
      tbv[i,] = pa
               } 
    }  

#  Compute diagonals of inverse of A
adiag = rep(0,nanimp)       # 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
adiag[nanimp] = 0

# For animals 121 to 320, generate
#     from 1 to 7 records (between days 1 to 49)
#  
rsd = sqrt(3)      # Temporary Error Variance = 3
trec = 0
aid = rep(0,nanim*7)
day = aid
obs = aid

for(i in 121:nanim){
  nrec = round(runif(1,1,7))      # A number between 1 and 7
  T = sample(c(1:49),nrec) # Days observed (sampled without replacement)
  for(j in 1:nrec){
      trec = trec + 1
      aid[trec] = i
      day[trec] = T[j]
      zz = (2*(T[j] - 1)/48) - 1
      p =  Legdre(zz)[1:3]
      arr = tbv[i,] %*% p
      perr = pefx[i,] %*% p
      terr = rnorm(1,0,rsd)
      obs[trec] =  mut[T[j]]+arr+perr + terr
   }
}

aid = aid[1:trec]
day = day[1:trec]
obs = round(obs[1:trec])
rrdat2 = data.frame(aid,day,obs)
trec
rrdat2[1:10, ]

#  Estimate parameters, G, PE, and R

source("ranregGS.R")

# Heritability for days 1 to 49

# Correlation of day 10 with days 1 to 49

# Solve MME

source("ranregMME.R")

# Plot results for a few animals (different genetic curves)

# Fix up the assignment sheet to reflect new material - LRS TO DO