
help_message <- "This script is used to separately associate phenotypes with the cis- and trans-components of gene expression to identify candidate genes that influence the phenotype.
Usage: Rscript cistrans_twas.R [options]

Options:
  --gffs_file=value   An RData file containing gene annotation information, including a data frame named 'gffs' with columns Gene, chr, start, end, strand, Qsymbols, and Note.
  --K_file=value   An RData file containing kinship matrix, including a matrix named 'K'.
  --pheno_file=value   An RData file containing phenotype data, including a data frame named 'pheno' with row names as individual IDs and column names as traits.
  --vc_file=value   An RData file generated by 'merge_cis_res.R' that includes the variance explained by cis-genetic variants for each gene and the corresponding p-value.
  --cis_file=value   An RData file generated by 'merge_cis_res.R' that includes a matrix of the cis-component of gene expression, named 'res_cis_m', with genes as rows and individuals as columns.
  --trans_file=value   An RData file generated by 'merge_cis_res.R' that includes a matrix of the trans-component of gene expression, named 'res_trans_m', with genes as rows and individuals as columns.
  --out_dir=value    Output directory.
  --p_threshold=value   P-value threshold for cis-genetic variance.
  --help         Display this help message.
"

args <- commandArgs(trailingOnly = TRUE)

if ("--help" %in% args) {
  cat(help_message)
  quit("no", status = 0)
}

for (arg in args) {
  if (startsWith(arg, "--gffs_file=")) {
    gffs_file <- gsub("--gffs_file=", "", arg)
  }
  else if (startsWith(arg, "--K_file=")) {
    K_file <- gsub("--K_file=", "", arg)
  }
  else if (startsWith(arg, "--pheno_file=")) {
    pheno_file <- gsub("--pheno_file=", "", arg)
  }
  else if (startsWith(arg, "--vc_file=")) {
    vc_file <- gsub("--vc_file=", "", arg)
  }
  else if (startsWith(arg, "--cis_file=")) {
    cis_file <- gsub("--cis_file=", "", arg)
  }
  else if (startsWith(arg, "--trans_file=")) {
    trans_file <- gsub("--trans_file=", "", arg)
  }
  else if (startsWith(arg, "--out_dir=")) {
    out_dir <- gsub("--out_dir=", "", arg)
  }
  else if (startsWith(arg, "--p_threshold=")) {
    p_threshold <- as.numeric(gsub("--p_threshold=", "", arg))
  }
}


cis_trans_TWAS <- function(phe_mat, cis_mat, trans_mat, vc_m, KG, gff, out_path, ncore){
   library(cpgen)
   library("fdrtool")
   set_num_threads(ncore)
   
   for(trait in colnames(phe_mat)){
      out_file <- sprintf('%s/%s_cis_trans_TWAS_result.csv', out_path, trait)
      if( file.exists(out_file) ) next
      system(sprintf("touch %s",out_file))
      print(trait)
      
      vars.u <- rownames(phe_mat)[!is.na(phe_mat[,trait])]
      vars.u <- intersect(vars.u, colnames(cis_mat))
      system.time(cis_res <- as.data.frame(cGWAS.emmax(y=phe_mat[match(vars.u, rownames(phe_mat)),trait], M=t(cis_mat[,match(vars.u,colnames(cis_mat))]), A=KG[match(vars.u,rownames(KG)), match(vars.u,colnames(KG))])[c("beta","se","p_value")])); 
      rownames(cis_res) <- rownames(cis_mat); 
      colnames(cis_res) <- c("beta_cis","se_cis","pval_cis")
      cis_res$Gene <- rownames(cis_res)
      cis_res$zscore_cis <- cis_res$beta_cis/cis_res$se_cis
      cis_res$qvalue_cis <- fdrtool(cis_res$pval_cis, statistic="pvalue", cutoff.method="fndr", plot=F, verbose=F)$qval
      
      vars.u <- rownames(phe_mat)[!is.na(phe_mat[,trait])]
      vars.u <- intersect(vars.u, colnames(trans_mat))
      system.time(trans_res <- as.data.frame(cGWAS.emmax(y=phe_mat[match(vars.u, rownames(phe_mat)),trait], M=t(trans_mat[,match(vars.u,colnames(trans_mat))]), A=KG[match(vars.u,rownames(KG)), match(vars.u,colnames(KG))])[c("beta","se","p_value")])); 
      rownames(trans_res) <- rownames(trans_mat); 
      colnames(trans_res) <- c("beta_trans","se_trans","pval_trans")
      trans_res$Gene <- rownames(trans_res)
      trans_res$zscore_trans <- trans_res$beta_trans/trans_res$se_trans
      trans_res$qvalue_trans <- fdrtool(trans_res$pval_trans, statistic="pvalue", cutoff.method="fndr", plot=F, verbose=F)$qval
      
      res <- merge(cis_res, trans_res, by.x="Gene",by.y="Gene")
      res$cis_rank <- rank(res$pval_cis)
      res$trans_rank <- rank(res$pval_trans)
      res$rank_product <- (res$cis_rank * res$trans_rank)^0.5
      res <- merge(res, gff, all.x=T)
      res$cis_Vg <- vc_m[match(res$Gene,rownames(vc_m)),"Vg"]
      res$cis_p <- vc_m[match(res$Gene,rownames(vc_m)),"pvalue"]
      res <- res[,c("Gene", "chr", "start", "end", "strand", "cis_Vg", "cis_p", "beta_cis", "se_cis", "pval_cis", "zscore_cis", "qvalue_cis", "beta_trans", "se_trans", "pval_trans", "zscore_trans", "qvalue_trans", "cis_rank", "trans_rank", "rank_product", "Qsymbols","Note")]
      res <- res[order(res$rank_product,decreasing=F),]
      write.csv(res, file=out_file,row.names=F)
   }
}


(load(gffs_file))
(load(K_file))
(load(pheno_file))
(load(vc_file))
(load(cis_file))
(load(trans_file))

str(gene_all <- rownames(res_vc_m)[which(res_vc_m$pvalue < p_threshold)])
str(gene_all <- intersect(gene_all, intersect(rownames(res_cis_m), rownames(res_trans_m))))
res_vc <- res_vc_m[match(gene_all, rownames(res_vc_m)), ]
res_cis <- res_cis_m[match(gene_all, rownames(res_cis_m)), ]
res_trans <- res_trans_m[match(gene_all, rownames(res_trans_m)),]
cis_trans_TWAS(phe_mat=pheno, cis_mat=res_cis, trans_mat=res_trans, vc_m=res_vc, KG=K, gff=gffs, out_path=out_dir, ncore=2)

