run:
    shell("{ldak} --cut-folds {params.folds} \
    --bfile {params.bfile} \
    --num-folds {params.numfolds}")

run:
    shell("{ldak} --calc-kins-direct {params.outdir} \
    --bfile {params.bfile} \
    --ignore-weights YES \
    --power 0 \
    --kinship-raw YES")

run:
    shell("{ldak} --reml {params.out} \
    --pheno {input.pheno} \
    --pheno-name {params.trait} \
    --grm {params.grm} \
    --covar {input.covar} \
    --constrain YES \
    --reml-iter 500")

run:
    for i in CV:
        for j in INDEX:
            shell("{ldak} --reml {params.out}{i}.{j} \
            --pheno {input.pheno} \
            --pheno-name {params.trait} \
            --grm {params.grm} \
            --keep {params.keep}{i}.train{j} \
            --covar {input.covar} \
            --constrain YES \
            --reml-iter 500")

run:
    for i in CV:
        for j in INDEX:
            shell("{ldak} --calc-blups {params.out}{i}.{j} \
            --grm {params.grm} \
            --remlfile {params.reml}{i}.{j}.reml \
            --bfile {params.bfile} \
            --covar {input.covar}")
            shell("{ldak} --calc-scores {params.out}{i}.{j} \
            --bfile {params.bfile} \
            --power 0 \
            --scorefile {params.out}{i}.{j}.blup \
            --keep {params.keep}{i}.test{j} \
            --pheno {input.pheno} \
            --pheno-name {params.trait}")

run:
    shell("Rscript src/04_summarize_gblup.R {input.blues}")

run:
    shell("Rscript src/05_select_pathway_snps.R {params.pathway} \"{params.bincode}\"")

run:
    # partition kinship matrix
    # --ignore-weights YES & --power 0 based on LDAK recommendations for prediction
    shell("{ldak} --cut-kins {params.outdir} \
    --bfile {params.bfile} \
    --partition-number 2 \
    --partition-prefix {params.prefix}")
    shell("{ldak} --calc-kins {params.outdir} \
    --bfile {params.bfile} \
    --partition 1 \
    --ignore-weights YES \
    --power 0 \
    --kinship-raw YES")
    shell("{ldak} --calc-kins {params.outdir} \
    --bfile {params.bfile} \
    --partition 2 \
    --ignore-weights YES \
    --power 0 \
    --kinship-raw YES")

run:
    for p in PATHWAYS.keys():
        shell("{ldak} --reml {params.out}{p}/{params.trait}.multiblup \
        --pheno {input.pheno} \
        --pheno-name {params.trait} \
        --mgrm {params.mgrm}{p}/partition.list \
        --covar {input.covar} \
        --constrain YES \
        --reml-iter 500")

run:
    for p in PATHWAYS.keys():
        for i in CV:
            for j in INDEX:
                shell("{ldak} --reml {params.out}{p}/{params.trait}.cv{i}.{j} \
                --pheno {input.pheno} \
                --pheno-name {params.trait} \
                --mgrm {params.mgrm}{p}/partition.list \
                --keep {params.keep}{i}.train{j} \
                --covar {input.covar} \
                --constrain YES \
                --reml-iter 100")

run:
    for p in PATHWAYS.keys():
        for i in CV:
            for j in INDEX:
                shell("{ldak} --calc-blups {params.out}{p}/{params.trait}.cv{i}.{j} \
                --mgrm {params.mgrm}{p}/partition.list \
                --remlfile {params.out}{p}/{params.trait}.cv{i}.{j}.reml \
                --bfile {params.bfile} \
                --covar {input.covar}")
                shell("{ldak} --calc-scores {params.out}{p}/{params.trait}.cv{i}.{j} \
                --bfile {params.bfile} \
                --power 0 \
                --scorefile {params.out}{p}/{params.trait}.cv{i}.{j}.blup \
                --keep {params.keep}{i}.test{j} \
                --pheno {input.pheno} \
                --pheno-name {params.trait}")

run:
    shell("Rscript src/06_summarize_multiblup.R {input.blues}")

shell:
    "Rscript src/07_null_group_sizes.R"

shell:
    "Rscript src/08_generate_null_gene_groups.R {params.num}"

run:
    for r in RANDOM:
        shell("{ldak} --cut-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition-number 2 \
        --partition-prefix {params.prefix}{r}/list")
        shell("{ldak} --calc-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition 1 \
        --ignore-weights YES \
        --power 0")
        shell("{ldak} --calc-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition 2 \
        --ignore-weights YES \
        --power 0")

run:
    for r in RANDOM:
        shell("{ldak} --reml {params.prefix}{r}/{params.trait}.h2 \
        --pheno {input.pheno} \
        --pheno-name {params.trait} \
        --mgrm {params.mgrm}{r}/partition.list \
        --covar {input.covar} \
        --constrain YES \
        --reml-iter 500")

run:
    shell("Rscript src/09_summarize_null.R")

shell:
    "Rscript src/08a_generate_null_gene_groups_pathways.R {params.pathway}"

run:
    for r in range(1,1001):
        shell("{ldak} --cut-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition-number 2 \
        --partition-prefix {params.prefix}{r}/list")
        shell("{ldak} --calc-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition 1 \
        --ignore-weights YES \
        --power 0")
        shell("{ldak} --calc-kins {params.prefix}{r} \
        --bfile {params.bfile} \
        --partition 2 \
        --ignore-weights YES \
        --power 0")

run:
    for r in range(1,1001):
        shell("{ldak} --reml {params.prefix}{r}/{params.trait}.h2 \
        --pheno {input.pheno} \
        --pheno-name {params.trait} \
        --mgrm {params.mgrm}{r}/partition.list \
        --covar {input.covar} \
        --constrain YES \
        --reml-iter 500")

run:
    shell("Rscript src/09a_summarize_null_pathways.R {params.pathway}")

shell:
    "Rscript src/02_extract_gene_ids.R {input}"

run:
    shell("{ldak} --cut-weights {params.outdir} \
    --bfile {params.bfile} \
    --section-kb 0.25")
    shell("{ldak} --calc-weights-all {params.outdir} \
    --bfile {params.bfile}")

library(tidyverse)
library(broom)

# use command line args
args <- commandArgs(trailingOnly = TRUE)

# phenotypes
pheno <- read.table(paste0(args[1]), header = TRUE) %>%
  gather(key = "trait",
         value = "phenotype",
         -FID, -IID)

# heritability
multiblup_h2 <- tibble(file = list.files(path = "models/multiblup_h2",
                                     pattern = ".share$",
                                     full.names = TRUE, recursive = TRUE)) %>%
  separate(file, sep = "/|[.]", into = c("source", "method", "pathway", "trait", "method2", "fill"), remove = FALSE) %>%
  # mutate(data = lapply(file, read.table, header = TRUE, skip = 13)) %>%
  mutate(data = lapply(file, read.table, header = TRUE)) %>%
  unnest(data) %>%
  # filter(Component %in% c("Her_K1", "Her_K2")) %>%
  # select(trait, pathway, Component, Heritability, Her_SD)
  select(trait, pathway, Component, Share, SD)

head(multiblup_h2)

# likelihood ratio
multiblup_llik <- tibble(file = list.files(path = "models/multiblup_h2",
                                     pattern = ".reml$",
                                     full.names = TRUE, recursive = TRUE)) %>%
  separate(file, sep = "/|[.]", into = c("source", "method", "pathway", "trait", "method2", "fill"), remove = FALSE) %>%
  mutate(data = lapply(file, read.table, header = TRUE)) %>%
  unnest(data) %>%
  filter(Num_Kinships %in% "Alt_Likelihood") %>%
  mutate(multiblup_llik = as.numeric(as.character(X2))) %>%
  select(trait, pathway, multiblup_llik)

head(multiblup_llik)

# predictive ability

# individuals used for test set
cross_val <- tibble(filename = list.files(path = "data/processed/cross_validation",
                    pattern = ".test", full.names = TRUE)) %>%
             mutate(cvnum = word(filename, start = 4, end = 4, sep = "/|[.]"),
                    fold = word(filename, start = 5, end = 5, sep = "/|[.]"),
                    fold = gsub("test", "", fold)) %>%
             mutate(file_contents = map(filename, ~read.table(., header = FALSE))) %>%
             unnest(cols = c("file_contents")) %>%
             rename("FID" = "V1",
                    "IID" = "V2")

# predictive accuracy
multiblup_pa <- tibble(file = list.files(path = "models/multiblup",
                                     pattern = ".pred$",
                                     full.names = TRUE,
                                     recursive = TRUE)) %>%
  separate(file, sep = "/|[.]", into = c("source", "method", "pathway", "trait", "cvnum", "fold"), remove = FALSE) %>%
  mutate(data = lapply(file, read.table, header = TRUE)) %>%
  unnest(data) %>%
  left_join(pheno, ., by = c("FID" = "ID1", "IID" = "ID2", "trait")) %>%
  inner_join(., cross_val)

# training_coef <- tibble(file = list.files(path = "models/multiblup",
#                                           pattern = ".coeff$",
#                                           full.names = TRUE, recursive = TRUE)) %>%
#   separate(file, sep = "/|[.]", into = c("source", "method", "pathway", "trait", "cvnum", "fold", "fill"), remove = FALSE) %>%
#   mutate(data = lapply(file, read.table, header = TRUE)) %>%
#   unnest(data) %>%
#   filter(Component %in% "Intercept") %>%
#   select(trait, pathway, cvnum, fold, Effect)
#
# head(training_coef)
#
# multiblup_pa <- tibble(file = list.files(path = "models/multiblup",
#                                          pattern = ".profile$",
#                                          full.names = TRUE, recursive = TRUE)) %>%
#   separate(file, sep = "/|[.]", into = c("source", "method", "pathway", "trait", "cvnum", "fold"), remove = FALSE) %>%
#   mutate(data = lapply(file, read.table, header = TRUE)) %>%
#   unnest(data) %>%
#   left_join(., training_coef, by = c("pathway", "trait", "cvnum", "fold"))

head(multiblup_pa)

# number of snps/genes in each pathway
files <- list.files(path = "data/processed/pathways", pattern = "*.txt", full.names = TRUE,
                    recursive = TRUE)

pathways <- tibble(file = files) %>%
  separate(file, sep = "/", into = c("dir", "dir2", "dir3", "pathway", "pathway2"),
           remove = FALSE) %>%
  mutate(data = lapply(file, read.table, header = TRUE),
         data = lapply(data, mutate_all, as.character)) %>%
  unnest(data) %>%
  group_by(pathway) %>%
  summarise(size = length(unique(snp_id)),
            n_genes = length(unique(ensembl_gene_id.x))) %>%
  select(pathway, n_genes, size) %>%
  ungroup() %>%
  write_csv(., "reports/pathways_summary.csv")

head(pathways)

save(multiblup_h2, multiblup_llik, multiblup_pa, pathways, file = "reports/multiblup.RData")

gene_group_sizes <- round(runif(5000, 1, 50000))
hist(gene_group_sizes)
write.table(gene_group_sizes, "data/interim/null_group_sizes.txt", row.names = FALSE)

# write sampling file (control_sampling.txt) for sbatch
# this can be used if submitting with sbatch to run groups of 50 random sets at a time
sampling <- as.data.frame(seq(1, length(gene_group_sizes), by = 50))
colnames(sampling) <- "start"
sampling$end <- sampling$start + 49

write.table(sampling, "data/interim/control_sampling.txt", row.names = FALSE, col.names = FALSE)

library(tidyverse)

# use command line arguments
args = commandArgs(trailingOnly = TRUE)

# name of pathway
pathway <- args[1]

# pathway SNPs
snp_ids <- read.table(paste0("data/processed/pathways/", pathway, "/list1"))

# number of SNPs to sample
nsnps <- nrow(snp_ids)

# number of random subsets to export
end <- 1000

cat("pathway:", pathway, "- sampling", end, "random subsets of", nsnps, "snps each\n")

# read in SNP and gene information for all SNPs
snps <- read.table("data/processed/snp_gene_ids_tair10.txt", header = TRUE) %>%
        arrange(chromosome, position) %>%
        filter(!snp_id %in% snp_ids$V1) # remove pathway SNPs

str(snps)

# # read in random uniform distribution for number of SNPs
# snp_number <- read.table("data/interim/gene_groups.txt", header = TRUE)
#
# a function to sample genes randomly up to a specified number of SNPs
# n = numer of SNPs to sample
# bprange = number of basepairs to include before/after a gene

sample_genes <- function(data, nsnps, bprange) {
  out <- data.frame()
  while (nrow(out) < nsnps) {
    # subset snps by genes already selected in output
    set <- na.omit(data[!data$ensembl_gene_id %in% out$ensembl_gene_id,])
    gene <- sample(set$ensembl_gene_id, 1, replace = FALSE) # select a random gene
    snp_sub <- set[set$ensembl_gene_id==gene,] # select all snps within gene
    chr <- snp_sub$chromosome[1]
    mn <- min(snp_sub$position) - bprange
    mx <- max(snp_sub$position) + bprange
    positions <- data[data$position > mn & data$position < mx & data$chromosome==chr,] # select SNPs in gene with X kb buffer
    out <- rbind(out, positions) # combine until threshold number of snps are met
    }
    return(out)
  }

# loop through gene sets specified by commandArgs

for (i in 1:end) {
  if(file.exists(paste0("data/processed/random_sets_pathways/", pathway, "/null_", i, ".txt"))) {
    cat("file already exists for sample", i, "\n")
    next
  }
  tryCatch({

  cat("exporting random gene group", i, "of", end, "\n")

  random <- sample_genes(snps, n = nsnps, bprange = 2500) %>%
  arrange(chromosome, position)

  dir.create(paste0("data/processed/random_sets_pathways/", pathway), showWarnings = FALSE, recursive = TRUE)

  write.table(random, paste0("data/processed/random_sets_pathways/", pathway, "/null_", i, ".txt"))

  # export snplists for LDAK
  dir <- paste0("data/processed/random_sets_pathways/", pathway, "/c_", i)
  dir.create(dir, recursive = TRUE, showWarnings = FALSE)

  # SNPs in random gene group
  random %>%
    select(snp_id) %>%
    unique() %>%
    write.table(., paste0(dir, "/list1"), quote = FALSE, row.names = FALSE, col.names = FALSE)

  # remaining SNPs
  snps %>%
    select(snp_id) %>%
    filter(!snp_id %in% random$V1) %>%
    write.table(., paste0(dir, "/list2"), quote = FALSE, row.names = FALSE, col.names = FALSE)

  }, error = function(e){cat("ERROR:", conditionMessage(e), "\n")})
}

library(tidyverse)

# use command line arguments
args = commandArgs(trailingOnly = TRUE)
# end <- as.numeric(args[1]) + 49
end <- 5000
cat("random subsets from ", args[1], "to", end, "\n")

# read in SNP and gene information
snps <- read.table("data/processed/snp_gene_ids_tair10.txt", header = TRUE) %>%
arrange(chromosome, position)

# read in random uniform distribution for number of SNPs
snp_number <- read.table("data/interim/gene_groups.txt", header = TRUE)

# a function to sample genes randomly up to a specified number of SNPs
# n = numer of SNPs to sample
# bprange = number of basepairs to include before/after a gene

sample_genes <- function(data, n, bprange) {
  out <- data.frame()
  while (nrow(out) < n) {
    # subset snps by genes already selected in output
    set <- na.omit(data[!data$ensembl_gene_id %in% out$ensembl_gene_id,])
    gene <- sample(set$ensembl_gene_id, 1, replace = FALSE) # select a random gene
    snp_sub <- set[set$ensembl_gene_id==gene,] # select all snps within gene
    chr <- snp_sub$chromosome[1]
    mn <- min(snp_sub$position) - bprange
    mx <- max(snp_sub$position) + bprange
    positions <- data[data$position > mn & data$position < mx & data$chromosome==chr,] # select SNPs in gene with X kb buffer
    out <- rbind(out, positions) # combine until threshold number of snps are met
    }
    return(out)
  }

# loop through gene sets specified by commandArgs

for (i in args[1]:end) {
  if(file.exists(paste0("data/processed/random_sets/null_", i, ".txt"))) {
    cat("file already exists for sample", i, "\n")
    next
  }
  tryCatch({
  print(snp_number[i,])
  random <- sample_genes(snps, n = snp_number[i,], bprange = 2500) %>%
  arrange(chromosome, position)

  dir.create("data/processed/random_sets", showWarnings = FALSE)

  write.table(random, paste0("data/processed/random_sets/null_", i, ".txt"))

  # export snplists for LDAK
  dir <- paste0("data/processed/random_sets/c_", i)
  dir.create(dir, recursive = TRUE, showWarnings = FALSE)

  random %>%
    select(snp_id) %>%
    unique() %>%
    write.table(., paste0(dir, "/list1"), quote = FALSE, row.names = FALSE, col.names = FALSE)

  snps %>%
    select(snp_id) %>%
    filter(!snp_id %in% random$V1) %>% # filter(!snp_id %in% random$snp_id) %>%
    write.table(., paste0(dir, "/list2"), quote = FALSE, row.names = FALSE, col.names = FALSE)

  }, error = function(e){cat("ERROR:", conditionMessage(e), "\n")})
}


for (f in files){
foo <- read.table(paste0(f), header = FALSE)
print(length(foo$V1))
foo %>% select(V1) %>% unique() %>% write.table(., paste0(f), row.names = FALSE, col.names = FALSE, quote = FALSE)
}

library(tidyverse)
library(data.table)
# library(quantreg)

# summarize random gene groups (number of SNPs & number of genes)

feature_sizes <- read_csv("reports/random_subsets_summary.csv")
head(feature_sizes)

# import log likelihood for GBLUP model
load("reports/gblup.RData")

## list files for reml output of each random gene group
files <- list.files(path = "models/null_h2", pattern = ".reml$",
                    full.names = TRUE, recursive = TRUE)

lr_null <- rbindlist(lapply(files, read.table, header = TRUE),
                      idcol = "filename", use.names = TRUE) %>%
  mutate(filename = files[filename]) %>%
  separate(filename, sep = "/|[.]",
           into = c("dir", "source", "pathway", "trait", "metric", "model"), remove = FALSE) %>%
  filter(Num_Kinships %in% "Alt_Likelihood") %>%
  mutate(null_llik = as.numeric(as.character(X2))) %>%
  left_join(., gblup_llik, by = "trait") %>%
  select(trait, pathway, Num_Kinships, null_llik, gblup_llik) %>%
  left_join(., feature_sizes, by = "pathway") %>%
  mutate(group_size = cut(size, breaks = seq(0, 50000, by = 10000),
                          include.lowest = TRUE, dig.lab = 10),
         LR = 2 * (null_llik - gblup_llik))

# lr_null <- rbindlist(lapply(files, read.table, header = TRUE),
#                       idcol = "filename", use.names = TRUE) %>%
#   mutate(filename = files[filename]) %>%
#   separate(filename, sep = "/|[.]",
#            into = c("dir", "source", "pathway", "trait", "metric", "model"), remove = FALSE) %>%
#   filter(Num_Kinships %in% "Alt_Likelihood") %>%
#   # mutate(null_llik = as.numeric(as.character(X2))) %>%
#   # left_join(., gblup_llik, by = "trait") %>%
#   # select(trait, pathway, Num_Kinships, null_llik, gblup_llik) %>%
#   rename("LRT_Stat" = X2) %>%
#   select(trait, pathway, LRT_Stat) %>%
#   left_join(., feature_sizes, by = "pathway") %>%
#   mutate(group_size = cut(size, breaks = seq(0, 50000, by = 10000),
#                           include.lowest = TRUE, dig.lab = 10))

head(lr_null)

## read in variance component estimates
files <- list.files(path = "models/null_h2", pattern = ".share$",
                    full.names = TRUE, recursive = TRUE)

h2_null <- rbindlist(lapply(files, read.table, header = TRUE),
                      idcol = "filename", use.names = TRUE) %>%
  mutate(filename = files[filename]) %>%
  separate(filename, sep = "/|[.]", into = c("dir", "source", "pathway", "trait", "metric", "model"), remove = FALSE) %>%
  left_join(., feature_sizes, by = "pathway") %>%
  select(trait, pathway, size, n_genes, Component, Share, SD)

head(h2_null)

## merge!
null_dist <- left_join(lr_null, h2_null, by = c("trait", "pathway")) %>%
  write_csv(., "reports/lr_null_results.csv")