##################################################################################################################################
# R code for large-scale genome scan of associations between SNPs and defense-related phenotypes in Arabidopsis
#                   Last change: 08/15/2017
#
# The code performs the following two tasks:
#
# Task 1. Load the 1001genome snp data by chunks, convert it to the 0,1,2 format, and save the new genotype code into seperate text files.
# Task 2. Run single marker association analysis and calcuate FDR-adjusted p-values
# 
#
# Data files
# A) 1001 genomes SNP data 
#File: 1001genomes_snp-short-indel_only_ACGTN_v3.1.vcf.snpeff
#   Rows 1-12: descriptive information
#   Row 12: column names for SNP data (1144 columns)
#   Rows 13 and after: Each row provides data for one SNP
#     column 1.	CHROM: chromosome ID
#     column 2.	POS:  physical position (bp)
#     column 3.	ID:  ()
#     column 4.	REF:  reference allele
#     column 5.	ALT:  alternative allele
#     column 6.	QUAL: phred-scaled quality score for the assertion made in ALT (numeric)
#     column 7.	FILTER: PASS if this position has passed all filters
#     column 8.	INFO:  additional information (SNP Annotation etc.)
#     column 9.	FORMAT (GT:GQ:DP): format of the genotype data in the VCF file
#     Columns 10 ~ 1144: genotyping data defined by VCF
#               {allele1}/{allel2}:{genotype quality}:{read depth} )
#
# B) AtPolyDB phenotype data for 23 Defense-related traits
# File: 1001Genomes.phen.from.AtPolyDB.Defense-related.csv
#   column 1. FID:
#   column 2. IID: IDs of arabidopsis lines were used to align with genotype data
#   Columns 10-25: numberic values for phenotype observations for the 23 Defense-related traits
# Notes: AtPolyDB phenotype data file has been pre-ordered to match the 1001 Genomes SNP data
#######################################################################################################################################

###########################  Settings  ###########################
work.dir = "~"

# path to the 1001genome snp data file
vcf.file="v:/PubData/1001genomics_snpeff_v31/1001genomes_snp-short-indel_only_ACGTN_v3.1.vcf.snpeff"

#path to the AtPolyDB phenotype data for the 23 Defense-related traits
phe.file="v:/PubData/1001genomics_snpeff_v31/output/dat/1001Genomes.phen.from.AtPolyDB.Defense-related.csv"

## Specify the number of chunks (i.e., how many parts of the full data set need to be split) and chunk size.
## This depends on the computing power (CPU speed and memory) of the computer. We chose chunk size of 10000. 
 chunk.size = 10000  # an R object of 10k SNPs from the vcf file uses about 110MB of memory
 chunk.n=Inf     


library(stats)
setwd(work.dir)

######  Task 1: Read data by chunks, convert genotype data, and save new data to text files  ######

# R function to extract genotype data from the string of one row (SNP) in the VCF file
VCF.to.012<-function(x, sep="|") {
  x=strsplit(x, "\t")[[1]]
  

  
  # Keep columns (CHROM, POS, REF, ALT, and INFO)
  info=((x[c(1,2,4,5, 8)]))
  # info=((x[c(1,2,4,5)]))
  
  # Keep only genotype values and conver them to 0,1,2 format for each SNP 
  (suppressWarnings(SNPs<-matrix(as.numeric(do.call(rbind, sapply(gsub("[:].*?$", "", x[-(1:9)]), strsplit, "[|]"))), ncol=2)))
  gen=apply(SNPs, 1, sum)
  
  # sep="|", format of genotyping data for MYSQL database;
  # sep="\n", format of genotyping data for to saved as csv files;
  genotype=gsub("NA", "", paste(gen, collapse = sep))
  
  
  ## number of alleles, MAF, number of genotyped Arabidopsis lines
  cnt.alleles=table(SNPs)
  N.alleles=length(cnt.alleles)
  N.genotyped=sum(table(gen))
  MAF=ifelse(N.alleles==1, 0, min(cnt.alleles/sum(cnt.alleles)))
  
  
  unlist(c(info, N.genotyped, N.alleles, MAF, genotype))
}



# create a file connection to access the 1001 genome SNP data file
con=file(vcf.file)
open(con, "r")
# skip the descriptive rows (the first 11 rows) in the SNP file
tmp<-readLines(con, 11)

# read column names from the 12th row in the raw VCF file
varname = strsplit(readLines(con, 1)[1], "\t")[[1]]

# IDs of Arabidopsis lines in the SNP data
LineID.SNP=varname[-(1:9)]

chunk_id=1

# Load data from the VCF file by chunks 
while (NROW(dat <- readLines(con, n = chunk.size))>0 & chunk_id <= chunk.n) {
  # the "dat" from the above is a vector of string, in which each element is for one SNP (row) in the VCF file 

  # extract genotype values of the SNPs from the VCF data and convert them to 0,1,2 format
  snp.dat=do.call(rbind, lapply(dat, VCF.to.012, sep="|"))
  ## snp.dat has the following columns: 
  #CHROM: chromosome ID
  #POS:  position of SNP in bp
  #REF:  reference allele
  #ALT:  alternative allele
  #INFO: SNP annotation in the raw VCF file
  #genotype: a string concateneated the genotype values of  the 1135 Arabidopsis lines in the 1001 genomes SNP data 
  #          the genotypes of different Arabidopsis lines are delimited by "|"
  
  
  ####  save the new genotype data to seperate files for each chunk
  subset.file=paste0(getwd(), "/SNP_subset_", chunk_id, ".txt")
  
  # save the new genotype data a txt file for the current data chunk
  # the operation will require about 5GB of free spaces at the working folder
  write.table(file=subset.file, x = snp.dat, sep = "\t", row.names = F, col.names = F, quote = F)
 
  # increate the chunk id by one for the next data chunk
  chunk_id=chunk_id+1
}



######  Part 2: Run GWAS  #################################################################

# run GWAS analysis for kth trait with all SNPs in the current data chunk
GWAS.get.pval<-function(k) {
  #### use gid.lst, gen, phe in the global enviroment
  OneSNP.pval.ttest.fisher<-function(xi, y) {
    xy=na.exclude(cbind(xi, y))
    
    #### set results as missing if the SNP is not diallelic
    if (NROW(table(xy[,1])) != 2) return (-1)
    
    #### set results as missing if there is no variation for the current trait
    if (NROW(table(xy[,2])) < 2) return (-1)
    
    if (NROW(table(xy[,2]))==2) {
      #### Fisher's exact test for binary traits 
      tb=table(xy[,2],xy[,1])
      rr=fisher.test(tb, alternative = "two.side")
      pval=rr$p.value
    } else {
      #### t.test for quantitative traits
      tb=split(xy[,2],xy[,1])
      if (any(sapply(tb, NROW)<2)) return (-1)
      rr=t.test(tb[[1]], tb[[2]])
      pval=rr$p.value
    }
    pval
  }
  
  # List of IDs of arabidopsis lines with valid phenotype values
  ids=gid.lst[[k]]
  
  x=gen[,ids, drop=F]
  y=phe[ids,2+k]
  
  as.matrix(apply(x, 1, OneSNP.pval.ttest.fisher, y))
}


# load phenotype data for the 23 Defense-related traits from AtPolyDB data
phe=read.csv(phe.file, stringsAsFactors = F, check.names = F)
gid.lst=apply(phe[,-(1:2)], 2, function(x) which(!is.na(x)))

# list all subset files for genotype data
file.lst=list.files(pattern = "SNP_subset_[0-9]+[.]txt")

for (file in file.lst) {
  # data from one subset file
  gen.dat=read.table(file, stringsAsFactors = F)
  
  suppressWarnings(gen <- do.call(rbind, lapply(sapply(gen.dat$V9, strsplit, "[|]"), as.numeric)))
  row.names(gen)<-NULL
  
  # Get p-values by using single SNP association analysis for the 23 Defense-related traits in the AtPolyDB data
  #     -1 is assigned if no valid data for the SNP
  p.mat=sapply(1:23, GWAS.get.pval)
  if (NROW(gen.dat)==1) p.mat=t(p.mat)
  output=cbind(gen.dat[,1:2, drop=F], p.mat)
  
  # save GWAS results to files  'SNP_subset_[0-9]_pval.txt'
  output.file=gsub(".txt", "_pval.txt", file)
  write.table(file = output.file, x = output, sep="\t", col.names = c("CHROM", "POS", paste0("t", 1:23)), row.names = F)
}



## calcuate FDR-adjusted p-values
file.lst=list.files(pattern = "SNP_subset_[0-9]+_pval[.]txt")
pval=do.call(rbind, lapply(file.lst, read.table, head=T, colClasses = rep("double", 25)))

col.lst=which(apply(phe[,-(1:2)], 2, function(x) NROW(table(x))>2))+2
pval[pval==-1]=NA

# FDR adjustment for quantitative traits
# No FDR correction for the binary traits because the Fisher's exact test would give non-uniform p-values across all SNPs
for (col in col.lst) {
  idx=which(!is.na(pval[,col]))
  pval[idx,col]=p.adjust(pval[idx,col])
}

## obtain the numbers of significant SNPs (adjusted-p <= 0.01) the for the 23 Defense-related traits in the AtPolyDB data
sigSNP.lst=apply(pval[,-c(1:2)], 2, function(x) (which(x<0.01)))
sigSNP.cnt=table(unlist(sigSNP.lst))

# The number of the SNPs that signifcant for at least one defense-related traits (Adj-p <= 0.01)
sigSNP.N=NROW(sigSNP.cnt)
cat("Number of Sig. SNPs: ", sigSNP.N, "\n")

## Save adjusted p-values
# write.table(file="GWAS_Adj_pval.txt", pval, sep="\t", col.names = c("Chrom", "pos", colnames(phe)[-(1:2)]), row.names = F)