# R program for ROMPrev
#
# MH Roy-Gagnon
# 3/05
#
# last modified: 10/3/06

# Function ROMPrevPrep: Reads in data and reformats them for ROMPrev

ROMPrevPrep<-function(filename) {
  data<-read.table(filename,sep=",",header=F,na.strings="    ")
  names(data)[1:6]<-c("fam","ind","fa","mo","sex","trait")

  # Sort data
  index<-order(data$fam,data$ind)
  data<-data[index,]

  # Recode SNP genotype for ROMP
  nloci<-ncol(data)-6
  snp<-matrix(NA,nrow=nrow(data),ncol=nloci)

  for (i in 1:nloci) {
   snp[,i][data[,i+6]==11]<-0
   snp[,i][data[,i+6]==12]<-1
   snp[,i][data[,i+6]==21]<-1
   snp[,i][data[,i+6]==22]<-2
  }

# Add trait value for mother and father and mid-parental value
  ftrait<-rep(NA,nrow(data))
  mtrait<-rep(NA,nrow(data))
  ptrait<-rep(NA,nrow(data))
  indic<-rep(NA,nrow(data))

  i<-0
  repeat{
   i<-i+1
   if (i>nrow(data)) break
   if (is.na(data$fa[i])) data$fa[i]<-0
   if (is.na(data$mo[i])) data$mo[i]<-0
   if (data$fa[i]==0 & data$mo[i]==0) next
   j<-i
   repeat{
    if (j>=nrow(data)) break
    j<-j+1
    if (data$fam[j]!=data$fam[i]) break
    if (data$ind[j]==data$fa[i]) ftrait[i]<-data$trait[j]
    if (data$ind[j]==data$mo[i]) mtrait[i]<-data$trait[j]
   }
   repeat{
    if (j==1) break
    j<-j-1
    if (data$fam[j]!=data$fam[i]) break
    if (data$ind[j]==data$fa[i]) ftrait[i]<-data$trait[j]
    if (data$ind[j]==data$mo[i]) mtrait[i]<-data$trait[j]
   }
  }

  mptrait<-(ftrait+mtrait)/2

  indic[is.na(mptrait)]<-1
  indic[!is.na(mptrait)]<-2
  indic[is.na(ftrait) & is.na(mtrait) & is.na(mptrait)]<-0

  ptrait[indic==1 & !is.na(ftrait)]<-ftrait[indic==1 & !is.na(ftrait)]
  ptrait[indic==1 & !is.na(mtrait)]<-mtrait[indic==1 & !is.na(mtrait)]

  # Pick which parent to use for ROOP when 2 parents are available
  who<-rep(NA,length(indic[indic==2]))
  for (i in seq(1,length(indic[indic==2]))){
   who[i]<-sample(1:2,1)
  }

  ptrait[indic==2][who==1]<-ftrait[indic==2][who==1]
  ptrait[indic==2][who==2]<-mtrait[indic==2][who==2]

# Creates centered variables without missing values

  y<-data$trait-mean(data$trait,na.rm=T)
  xmp<-mptrait-mean(mptrait,na.rm=T)
  xmo<-mtrait-mean(mtrait,na.rm=T)
  xfa<-ftrait-mean(ftrait,na.rm=T)
  xop<-ptrait-mean(ptrait,na.rm=T)
  z<-matrix(NA,length(y),ncol=nloci)
  for (i in 1:nloci) {
   z[,i]<-snp[,i]-mean(snp[,i],na.rm=T)
  }

  val<-list(xmp=xmp,xmo=xmo,xfa=xfa,xop=xop,y=y,z=z)
  val

}


# Function ROMPrev: ROMPrev parametric analysis

ROMPrev<-function(y,x,z) {

  y1<-y[!is.na(y) & !is.na(x) & !is.na(z)]
  x1<-x[!is.na(y) & !is.na(x) & !is.na(z)]
  z1<-z[!is.na(y) & !is.na(x) & !is.na(z)]
  y<-y1
  x<-x1
  z<-z1

  slmyx<-summary(lm(y~x))
  slmyxz<-summary(lm(y~x+z))

  n<-length(y)
  h2<-slmyx$coef[2,1]
  seh2<-slmyx$coef[2,2]
  cih2<-c(h2-(qt(0.975,n-2)*seh2),h2+(qt(0.975,n-2)*seh2))
  th2<-slmyx$coef[2,3]
  ph2<-slmyx$coef[2,4]

  b2<-slmyxz$coef[2,1]
  aa<-1-(h2/2)
  bb<-1-(b2/2)
  h2l<-(h2-b2)/bb
  vb1<-(slmyxz$sigma**2)/sum(x**2)
  vb2<-slmyxz$coef[2,2]**2
  varh2l<-((bb*(h2-(b2/2)-1)*vb1)+(aa**2*vb2))/(bb**4)
  seh2l<-sqrt(varh2l)
  cih2l<-c(h2l-(qt(0.975,n-3)*seh2l),h2l+(qt(0.975,n-3)*seh2l))
  th2l<-h2l/seh2l
  ph2l<-pt(abs(th2l),n-3,lower.tail=F)*2

  val<-list(n=n,h2=h2,seh2=seh2,cih2=cih2,th2=th2,ph2=ph2,
            h2l=h2l,seh2l=seh2l,cih2l=cih2l,th2l=th2l,ph2l=ph2l)
  class(val)<-"romp"
  val
}

# Function ROMPrevRS: ROMPrev permutation test and Nonparametric bootstrap
#   resampling the initial data
# Need to run the romp function first and pass the results as a parameter

ROMPrevRS<-function(y,x,z,rompres,nb) {

  y1<-y[!is.na(y) & !is.na(x) & !is.na(z)]
  x1<-x[!is.na(y) & !is.na(x) & !is.na(z)]
  z1<-z[!is.na(y) & !is.na(x) & !is.na(z)]
  y<-y1
  x<-x1
  z<-z1

# Initializes statistics for permutation test
  h20<-rep(NA,nb)
  th20<-rep(NA,nb)
  h2l0<-rep(NA,nb)
  th2l0<-rep(NA,nb)

# Initializes statistics for bootstrap
  h2s<-rep(NA,nb)
  h2ls<-rep(NA,nb)

  n<-length(y)

  for (i in 1:nb){

# Resampling without replacement for permutation test
   order<-sample(n)
   slmyx<-summary(lm(y~x[order]))
   slmyxz<-summary(lm(y~x+z[order]))
   h20[i]<-slmyx$coef[2,1]

   b2<-slmyxz$coef[2,1]
   h2l0[i]<-(rompres$h2-b2)/(1-(b2/2))


# Resampling with replacement for bootstrap
   rows<-sample(1:n,replace=T)
   xs<-x[rows]
   ys<-y[rows]
   zs<-z[rows]

   slmyx<-summary(lm(ys~xs))
   slmyxz<-summary(lm(ys~xs+zs))

   h2s[i]<-slmyx$coef[2,1]

   b2<-slmyxz$coef[2,1]
   h2ls[i]<-(h2s[i]-b2)/(1-(b2/2))
  }

  ci95h2<-quantile(h2s,c(0.025,0.975))
  ci95h2l<-quantile(h2ls,c(0.025,0.975))

  val<-list(n=n,h2=rompres$h2,ph2c=mean(h20>rompres$h2),
            bseh2=sd(h2s),ci95h2=ci95h2,bph2=mean(h2s<0),
            h2l=rompres$h2l,ph2lc=mean(h2l0>rompres$h2l),
            bseh2l=sd(h2ls),ci95h2l=ci95h2l,bph2l=mean(h2ls<0)
           )
  class(val)<-"romprs"
  val
}


# Function ROOPrev: ROOPrev parametric analysis

ROOPrev<-function(y,x,z) {

  y1<-y[!is.na(y) & !is.na(x) & !is.na(z)]
  x1<-x[!is.na(y) & !is.na(x) & !is.na(z)]
  z1<-z[!is.na(y) & !is.na(x) & !is.na(z)]
  y<-y1
  x<-x1
  z<-z1

  slmyx<-summary(lm(y~x))
  slmyxz<-summary(lm(y~x+z))

  n<-length(y)
  b1<-slmyx$coef[2,1]
  h2<-2*b1
  seh2<-2*slmyx$coef[2,2]
  cih2<-c(h2-(qt(0.975,n-2)*seh2),h2+(qt(0.975,n-2)*seh2))
  th2<-slmyx$coef[2,3]
  ph2<-slmyx$coef[2,4]

  b2<-slmyxz$coef[2,1]
  aa<-1-(b1/2)
  bb<-1-(b2/2)
  h2l<-2*(b1-b2)/bb
  vb1<-(slmyxz$sigma**2)/sum(x**2)
  vb2<-slmyxz$coef[2,2]**2
  varh2l<-4*((bb*(b1-(b2/2)-1)*vb1)+(aa**2*vb2))/(bb**4)
  seh2l<-sqrt(varh2l)
  cih2l<-c(h2l-(qt(0.975,n-3)*seh2l),h2l+(qt(0.975,n-3)*seh2l))
  th2l<-h2l/seh2l
  ph2l<-pt(abs(th2l),n-3,lower.tail=F)*2

  val<-list(n=n,h2=h2,seh2=seh2,cih2=cih2,th2=th2,ph2=ph2,
            h2l=h2l,seh2l=seh2l,cih2l=cih2l,th2l=th2l,ph2l=ph2l)
  class(val)<-"romp"
  val
}

# Function ROOPrevRS: ROOPrev permutation test and Nonparametric bootstrap
#   resampling the initial data
# Need to run the romp function first and pass the results as a parameter

ROOPrevRS<-function(y,x,z,roopres,nb) {

  y1<-y[!is.na(y) & !is.na(x) & !is.na(z)]
  x1<-x[!is.na(y) & !is.na(x) & !is.na(z)]
  z1<-z[!is.na(y) & !is.na(x) & !is.na(z)]
  y<-y1
  x<-x1
  z<-z1

# Initializes statistics for permutation test
  h20<-rep(NA,nb)
  th20<-rep(NA,nb)
  h2l0<-rep(NA,nb)
  th2l0<-rep(NA,nb)

# Initializes statistics for bootstrap
  h2s<-rep(NA,nb)
  h2ls<-rep(NA,nb)

  n<-length(y)

  for (i in 1:nb){

# Resampling without replacement for permutation test
   order<-sample(n)
   slmyx<-summary(lm(y~x[order]))
   slmyxz<-summary(lm(y~x+z[order]))
   h20[i]<-2*slmyx$coef[2,1]

   b2<-slmyxz$coef[2,1]
   h2l0[i]<-(roopres$h2-(2*b2))/(1-(b2/2))


# Resampling with replacement for bootstrap
   rows<-sample(1:n,replace=T)
   xs<-x[rows]
   ys<-y[rows]
   zs<-z[rows]

   slmyx<-summary(lm(ys~xs))
   slmyxz<-summary(lm(ys~xs+zs))

   h2s[i]<-2*slmyx$coef[2,1]

   b2<-slmyxz$coef[2,1]
   h2ls[i]<-(h2s[i]-(2*b2))/(1-(b2/2))
  }

  ci95h2<-quantile(h2s,c(0.025,0.975))
  ci95h2l<-quantile(h2ls,c(0.025,0.975))

  val<-list(n=n,h2=roopres$h2,ph2c=mean(h20>roopres$h2),
            bseh2=sd(h2s),ci95h2=ci95h2,bph2=mean(h2s<0),
            h2l=roopres$h2l,ph2lc=mean(h2l0>roopres$h2l),
            bseh2l=sd(h2ls),ci95h2l=ci95h2l,bph2l=mean(h2ls<0)
           )
  class(val)<-"romprs"
  val
}


# Function print.ROMPrev: prints objects of class romp or roop
#  from ROMPrev or ROOPrev functions

print.romp<-function(x, ...) {

  fx<-format(round(c(x$h2,x$seh2,x$cih2[1],x$cih2[2],x$h2l,x$seh2l,x$cih2l[1],
                     x$cih2l[2],-0.65271),digit=4),nsmall=4)
  fp<-format(format.pval(c(x$ph2,x$ph2l)),justify="right")

  cat("\n",
      "                    Estimate  Std. Error         95% CI 
p-value    ",
      "\n",
 
"----------------------------------------------------------------------------",
      "\n",
      "Trait h2           ",fx[1],"  ",fx[2],"    ",fx[3],"-",fx[4],"  ",
      fp[1],"\n",
      "Locus-specific h2  ",fx[5],"  ",fx[6],"    ",fx[7],"-",fx[8],"  ",
      fp[2],"\n",
 
"----------------------------------------------------------------------------",
      "\n",
      "Number of parent-offspring trios used:",x$n,"\n\n")
}


# Function print.ROMPrevRS: prints objects of class romprs or rooprs
#  from ROMPrevRS or ROOPrevRS functions

print.romprs<-function(x, ...) {

  fx<-format(round(c(x$h2,x$bseh2,x$ci95h2[1],x$ci95h2[2],x$h2l,x$bseh2l,x$ci95h2l[1],
                     x$ci95h2l[2],-0.65271),digit=4),nsmall=4)
  fp<-format(format.pval(c(x$bph2,x$ph2c,x$bph2l,x$ph2lc)),justify="right")

  cat("\n",
      "                             Bootstrap     Bootstrap 
p-value     ",
      "\n",
      "                   Estimate Std. Error      95% CI 
Bootstrap Permutation  ",
      "\n",
 
"---------------------------------------------------------------------------------",
      "\n",
      "Trait h2          ",fx[1]," ",fx[2]," ",fx[3],"-",fx[4]," ",
      fp[1]," ",fp[2],"\n",
      "Locus-specific h2 ",fx[5]," ",fx[6]," ",fx[7],"-",fx[8]," ",
      fp[3]," ",fp[4],"\n",
 
"---------------------------------------------------------------------------------",
      "\n",
      "Number of parent-offspring trios used:",x$n,"\n\n")
}