############################################################################################
### R code to generate simulated trait data to reproduce results presented in Figures    ###
### 2, 3, and 4 of the Read et al. manuscript under review (and in the supplemental      ###
### figures). Code edited 14 Sep 2015 by QDR                                             ###
### Contact: qread@utk.edu                                                               ###
############################################################################################



# Function to generate a trait data set given parameters.

simulate.trait <- function(n.indiv, n.markers, h2=0.5, trait.mean=0, sd.g = 1) {
  # The code in this function is modified from an example script given in some documentation in the rrBLUP package.
  # Make set of binary markers for each individual
  M <- matrix(rep(0,n.indiv*n.markers),n.indiv,n.markers)
  for (i in 1:n.indiv) {
    M[i,] <- ifelse(runif(n.markers)<0.5,-1,1)
  }
  
  # Generate random phenotypes normally distributed around trait mean, varying more if h2 is lower
  u <- rnorm(n.markers, mean = 0, sd = sd.g)
  g <- as.vector(crossprod(t(M),u))
  y <- g + rnorm(n.indiv,mean=trait.mean,sd=(1-sqrt(h2))/sqrt(var(g)))
  
  #assign genotype to each unique marker combination
  res <- data.frame(cbind(M, trait=y, gen=NA))
  for (i in 1:nrow(unique(M))) res$gen[apply(M, 1, function(x) all(x==unique(M)[i,]))] <- i
  res$gen <- as.factor(res$gen)
  return(res)
}


# Function to generate multiple trait data sets for multiple species and format the data frame properly.

sim.species <- function(n.indiv=c(1000,1000), n.markers=c(2,2), h2=c(.5,.5), trait.mean=c(0,1), sd.g = c(1,1)) {
  res <- data.frame(matrix(nrow=0, ncol=3))
  
  for (i in 1:length(n.indiv)) {
    s <- simulate.trait(n.indiv=n.indiv[i], n.markers=n.markers[i], h2=h2[i], trait.mean=trait.mean[i], sd.g = sd.g[i])
    res <- rbind(res, data.frame(species=paste('S',i,sep=''), genotype=paste('S',i,'G',s$gen,sep=''), trait=as.numeric(s$trait)))
  }
  return(res)
}



# The code below generates the trait data sets.


h2.all <- c(0.0, seq(from=0.1, to=0.9, by=0.1))
marker.all <- 1:3
mean.diff <- c(0, 0.1, 0.25, 0.5, 1, 2, 5, 10) 
n.reps <- 100
indiv.pop <- c(10, 50, 100)
intra.sd <- c(0.1, 0.25, 0.5, 1, 2, 5, 10)
n.spp <- 2

alpha <- 0.05 # Value for significance threshold.

sim <- list()
run <- 0
library(lme4)
library(lmtest)
library(MuMIn)

for (i in h2.all) {
  cat('\nHeritability',i)
  for (j in intra.sd) { 
    cat('\nGenotype st dev',j)
    for (k in marker.all) {
      for (l in mean.diff) {
        for (m in indiv.pop) {		
          for (n in 1:n.reps) {
            run <- run+1
		# Generate trait dataset for the particular combination of parameters.
            s <- sim.species(n.indiv=rep(m, n.spp),
                             n.markers=rep(k, n.spp),
                             h2=rep(i, n.spp),
                             trait.mean=seq(from=1, by=l, length=n.spp),
                             sd.g = rep(j, n.spp))
            
            # Fit models.
            linear <- try(lm(trait ~ species, data=s), TRUE)
            mixed <- try(lmer(trait ~ species + (1|genotype), data=s), TRUE)
            if (inherits(linear, 'try-error')) {
              aic.linear <- NA
              rsq.linear <- NA
              pval.linear <- NA
            }
            else  {
		  # Calculate AIC, R2, and likelihood ratio test for the linear model.
              aic.linear <- AIC(linear)
              rsq.linear <- summary(linear)$r.sq
              linear.null <- lm(trait ~ 1, data = s)
              pval.linear <- lrtest(linear, linear.null)$Pr[2]
            }
            if (inherits(mixed, 'try-error')) {
              aic.mixed <- NA
              rsq.mixed <- c(NA,NA)
              pval.mixed <- NA
              error.type <- NA
            }
            else {
		  # Calculate AIC, marginal and conditional R2, and likelihood ratio test for the mixed model.
              aic.mixed <- AIC(mixed)
              rsq.mixed <- r.squaredGLMM(mixed)
              mixed.null <- lmer(trait ~ (1|genotype), data=s)
              pval.mixed <- suppressMessages(anova(mixed, mixed.null))$Pr[2]
              if (pval.linear <= alpha & pval.mixed <= alpha)
              { error.type <- 'Correct +' }
              else if (pval.linear > alpha & pval.mixed > alpha)
              { error.type <- 'Correct -' }
              else if (pval.linear <= alpha & pval.mixed > alpha)
              { error.type <- 'False +' }
              else { error.type <- 'False -' }
            }
                        
            sim[[run]] <- list(smeans=with(s, tapply(trait, species, mean)),
                               gmeans=with(s, tapply(trait, genotype, mean)),
                               aic.linear=aic.linear, aic.mixed=aic.mixed,
                               rsq.linear=rsq.linear, rsq.mixed=rsq.mixed,
                               pval.linear=pval.linear, pval.mixed=pval.mixed,
                               error.type=error.type,
                               h2=i, n.spp=n.spp, n.marker=k, mean.diff=l, n.indiv=m, intra.sd=j, rep=n, run=run)
          }				   
        }
      }
    }
  }

# Combine results into a data frame.
# This is updated throughout the loop so that the results are not lost.
sim_list <- lapply(sim, function(x) with(x, data.frame(mu1=smeans[1], mu2=smeans[2], aic.linear=aic.linear, aic.mixed=aic.mixed, rsq.linear=rsq.linear, rsq.marg.mixed=rsq.mixed[1], rsq.cond.mixed=rsq.mixed[2], pval.linear=pval.linear, pval.mixed=pval.mixed, error.type=error.type, h2=h2, n.marker=n.marker, mean.diff=mean.diff, intra.sd=intra.sd, n.indiv=n.indiv, rep=rep, run=run)))
sim_list <- do.call('rbind', sim_list)
save(sim_list, file= paste0('sim_output',i*10,'.r'))
write.table(sim_list, file=paste0('sim_output',i*10,'.csv'), sep = ',', row.names=FALSE)
sim <- list()
run <- 0
}

