save.image("logs/build_library.RData")

log <- file(snakemake@log[[1]], open="wt")
sink(log, type = "message")
sink(log, type = "output")


library(BSgenome.Hsapiens.NCBI.GRCh38)
library(Biostrings)
library(plyranges)
library(magrittr)
library(tidyverse)

source("lib/helpers.R")

set.seed(snakemake@config$seed)

n_sub_libraries <- snakemake@config$sub_libraries

max_oligo_length <- snakemake@config$oligo$max_len

re_seqs <- DNAStringSet(unlist(snakemake@config$restriction_enzymes))

primer_5p <- snakemake@config$oligo$primer_5p
primer_3p <- snakemake@config$oligo$primer_3p

cloning_site_1 <- re_seqs[snakemake@config$cloning_site$enzyme1]
cloning_site_2 <- re_seqs[snakemake@config$cloning_site$enzyme2]

stuffer_bp <- snakemake@config$cloning_site$stuffer_bp
stuffer_gc <- snakemake@config$cloning_site$stuffer_GC

cloning_site_length <- str_length(cloning_site_1) + stuffer_bp + str_length(cloning_site_2)

barcodes_per_frag <- snakemake@config$bc_per_frag

bc_length <- snakemake@config$oligo$bc_len

max_indel_size <- snakemake@config$oligo$max_indel

random_control_n <- snakemake@config$random_controls
random_control_w_mut_n <- snakemake@config$random_controls_with_mutation

wiggle_room <- max_indel_size + 1

final_frag_len <- max_oligo_length - str_length(primer_5p) - cloning_site_length - str_length(primer_3p) - bc_length - wiggle_room

window_flank <- floor((final_frag_len - 1) / 2)

snps <- read_tsv(snakemake@input$filtered_snps)

snps_cleaned <- snps %>%
    filter(!is.na(ref), !is.na(alt),
           abs(str_count(ref, "[ACGT]") - str_count(alt, "[ACGT]")) <= max_indel_size) %>%
    # mutate(alt = str_extract(alt, "^[ACGT-]+"), # only use first alternate allele if multiple listed
    mutate(alt = str_split(alt, ",")) %>%
    unnest(alt) %>%
    mutate(is_indel = str_count(ref, "[ACGT]") != str_count(alt, "[ACGT]"),  # annotate indels
           is_indel_fixed_base = is_indel & !str_detect(ref, "-") & !str_detect(alt, "-"), # annotate indels with the preceding fixed base
           ref = ifelse(is_indel_fixed_base, str_sub(ref, 2), str_replace(ref, "-", "")), # remove fixed first base or the "-" annotation for indels
           alt = ifelse(is_indel_fixed_base, str_sub(alt, 2), str_replace(alt, "-", ""))) %>%
    mutate(pos = ifelse(is_indel, as.numeric(pos) + 1, as.numeric(pos))) %>%
    distinct(snp, chr, pos, ref, alt, .keep_all = T) %>%
    group_by(snp, chr, pos, ref) %>%
    mutate(alt_n = as.integer(table(alt)[alt])) %>%
    slice_max(alt_n, n = 1, with_ties = F) %>%
    ungroup() %>%
    select(-alt_n) %>%
    arrange(chr, pos, snp) %>%
    mutate(fragment = row_number()) # account for removing fixed base


snps_gr <- snps_cleaned %>%
    transmute(seqnames = chr, start = pos, end = pos + str_length(ref) - 1,
              fragment, chr, pos, snp, ref, alt, is_indel, is_indel_fixed_base, source) %>%
    as_granges() %>%
    mutate()


snps_windows_gr <- resize(snps_gr, width = 2 * window_flank + 1, fix = "center")

## Get a window of sequence around SNP

snp_fragments_raw <- getSeq(BSgenome.Hsapiens.NCBI.GRCh38, `seqlevelsStyle<-`(snps_windows_gr, "NCBI"))


snps_frags <- snps_cleaned %>%
    select(fragment, chr, pos, snp, ref, alt, is_indel, is_indel_fixed_base, source) %>%
    mutate(seq = as.character(snp_fragments_raw),
           window_start = start(snps_windows_gr),
           window_end = end(snps_windows_gr))



# remove 10bp from each flank
ignore_edge_bps <- 10
overlapping_snps_gr <- resize(snps_windows_gr, width = width(snps_windows_gr) - 2*ignore_edge_bps,
                              fix = "center") %>%
    select(fragment) %>%
    join_overlap_left(select(snps_gr, -fragment))



# interim_data <- "./outs/LDhap"
# dir.create(interim_data)
# haplo_results_file <- "./outs/haplo_results.rds"


# if (!file.exists(haplo_results_file)) {

#     haplotypes_df <- overlapping_snps_gr %>%
#         as_tibble() %>%
#         filter(!is.na(snp), str_detect(snp, "rs\\d+")) %>%
#         # filter(!is.na(rsID_ldlink), str_detect(rsID_ldlink, "rs\\d+")) %>%
#         group_by(fragment) %>%
#         filter(dplyr::n() > 1) %>%
#         nest() %>%
#         mutate(out_file = file.path(interim_data, paste0("fragment_", fragment, ".txt")),
#                query = map(data, ~ pull(., snp)))
#                # query = map(data, ~ ifelse(str_detect(.x$rsID_ldlink, "rs\\d+"), .x$rsID_ldlink, paste(.x$chr, .x$pos, sep = ":"))))

#     haplo_results <- haplotypes_df %>%
#         mutate(LDhap_data = map2(query, out_file, query_ldhap))


#     write_rds(haplo_results, haplo_results_file)
# } else {
#     haplo_results <- read_rds(haplo_results_file)
# }


# tidy_LDhap_result <- function(x) {
#     haplotype_allele_df <- x %>% mutate(haplotype = row_number()) %>%
#         dplyr::rename(count = Count, frequency = Frequency) %>%
#         gather("snp", "allele", -haplotype, -count, -frequency) %>%
#         group_by(snp, allele) %>%
#         mutate(allele_count = sum(count)) %>%
#         group_by(snp) %>%
#         mutate(allele_freq = allele_count / sum(allele_count[!duplicated(allele)]))
#     return(haplotype_allele_df)
# }


# haplo_results_tidy <- haplo_results %>%
#     mutate(LDhap_data_tidy = map(LDhap_data, tidy_LDhap_result)) %>%
#     select(fragment, LDhap_data_tidy) %>%
#     unnest(LDhap_data_tidy) %>%
#     group_by(fragment) %>%
#     filter(n_distinct(snp) > 1) %>%
#     ungroup()


# haplo_fragments <-
#     haplo_results_tidy %>%
#     mutate(allele = str_replace(allele, "-", "")) %>%
#     left_join(select(snps_cleaned, snp, chr, pos, ref, alt, is_indel)) %>%
#     left_join(select(snps_frags, fragment, seq, window_start, window_end)) %>%
#     filter(allele == ref | allele == alt) %>%
#     mutate(seq_sub = str_sub(seq, pos - window_start + 1, pos - window_start + str_length(ref))) %>%
#     group_by(fragment, haplotype, frequency, count, seq) %>% nest() %>%
#     mutate(haplotype_seq = map2_chr(seq, data, ~ mutate_fragments(.x, .y$pos-.y$window_start+1, .y$ref, .y$allele))) %>%
#     ungroup()

# haplo_fragments_check <- haplo_fragments %>% unnest(data) %>% ungroup() %>%
#     group_by(fragment, haplotype) %>%
#     summarise(full_seq_diff = str_length(haplotype_seq[1]) - str_length(seq[1]),
#               var_seq_diff = sum(str_length(allele) - str_length(ref)))

# # This should only return haplotype fragments where SNP positions are overlapping
# haplo_fragments_check %>% filter(full_seq_diff != var_seq_diff)






snps_frags_allele <- snps_frags %>% mutate(allele = map2(ref, alt, ~ c(.x, .y))) %>% unnest(allele) %>%
    mutate(var_seq = pmap_chr(.,
                              function(seq, pos, window_start, ref, allele, ...) {
                                  mutate_fragments(seq, pos - window_start + 1, ref, allele)
                              })) %>%
    mutate(variant = ifelse(allele == ref, "ref", "alt"))



flank_width <- max(width(re_seqs)) - 1 # for checking RE sites


combined_fragments <- bind_rows(
    snps_frags_allele %>%
        mutate(haplotype = NA) %>%
        select(fragment, variant, haplotype, seq, var_seq, snp, chr, pos, allele, ref, alt, is_indel, window_start, window_end) %>%
        group_by(fragment, variant, haplotype, seq, var_seq) %>% nest() %>% ungroup #,
    # haplo_fragments %>%
    #     filter(haplotype_seq %!in% snps_frags_allele$var_seq,
    #            frequency >= 0.01) %>%
    #     unnest(data) %>%
    #     select(fragment, haplotype, seq, var_seq = haplotype_seq, snp, chr, pos, allele, ref, alt, is_indel, window_start, window_end) %>%
    #     group_by(fragment, haplotype, seq, var_seq) %>% nest() %>% ungroup
) %>%
    mutate(var_seq_flanks  = paste0(str_sub(primer_5p, -flank_width),
                                    var_seq,
                                    str_sub(cloning_site_1, 1, flank_width)))



# Randomly mutate restriction sites in ref/alt pairs
# If haplotype created unique restriction site, toss it


## Search from R.E. sites
re_matches <- combined_fragments %>%
    ungroup %>%
    mutate(re_matches = map(var_seq_flanks, function(x) {
        map_dfr(names(re_seqs), function(re_name) {
            vmatchPattern(re_seqs[[re_name]], x)[[1]] %>%
                as.data.frame() %>% mutate(RE = re_name)
        })
    }))



# re_matches %>% ungroup %>%
#     select(fragment, variant, haplotype, var_seq, re_matches) %>%
#     unnest(re_matches, keep_empty = T) %>%
#     group_by(fragment) %>%
#     filter(any(!is.na(RE))) %>%
#     filter(any(is.na(RE))) %>% View




combined_fragments_corrected <- re_matches %>% # filter(Fragment == 340) %>%
    group_by(fragment) %>%
    group_modify(~ mutate_restriction_sites(.x, .y$fragment))

tmp <- combined_fragments_corrected %>% filter(!fail) %>%
    mutate(var_seq_corrected_flanks  = paste0(str_sub(primer_5p, - flank_width),
                                              var_seq_corrected,
                                              str_sub(cloning_site_1, 1, flank_width))) %>%
    ungroup %>%
    mutate(re_matches = map(var_seq_corrected_flanks, function(x) {
        map_dfr(names(re_seqs), function(re_name) {
            vmatchPattern(re_seqs[[re_name]], x)[[1]] %>%
                as.data.frame() %>% mutate(RE = re_name)
        })
    }))


fragments_final <- combined_fragments_corrected %>% ungroup() %>%
    filter(!fail) %>%
    select(fragment_initial = fragment, variant, haplotype_initial = haplotype, seq, var_seq, var_seq_corrected, data) %>%
    mutate(fragment = dense_rank(fragment_initial)) %>%
    group_by(fragment) %>%
    mutate(haplotype = dense_rank(haplotype_initial)) %>%
    ungroup() %>%
    mutate(frag_id = ifelse(!is.na(variant),
                            paste0('fragment-', str_pad(fragment, width = 4, pad = 0), "_", "allele-", variant),
                            paste0('fragment-', str_pad(fragment, width = 4, pad = 0), "_", "haplotype-", haplotype)))


random_controls <- fragments_final %>% filter(!is.na(variant)) %>%
    distinct(seq) %>%
    # sample_n(random_control_n*2, replace = random_control_n*2 > nrow(.)) %>%
    mutate(fragment_initial = row_number(),
           GC = str_count(seq, "[GC]")/str_length(seq),
           shuffle_seq = map_chr(seq, ~ str_split(., "") %>% unlist %>% sample() %>% str_c(collapse = "")),
           shuffle_GC = str_count(shuffle_seq, "[GC]")/str_length(shuffle_seq),
           pos = floor((str_length(shuffle_seq) + 1)/2),
           ref = str_sub(shuffle_seq, pos, pos),
           alt = ref) %>%
    mutate(variant = map2("ref", "alt", ~ c(.x, .y))) %>% unnest(variant) %>%
    mutate(allele = ifelse(variant == "ref", ref, alt)) %>%
    mutate(var_seq = shuffle_seq) %>%
    mutate(var_seq_flanks  = paste0(str_sub(primer_5p, - flank_width),
                                    var_seq,
                                    str_sub(cloning_site_1, 1, flank_width))) %>%
    ungroup %>%
    mutate(re_matches = map_lgl(var_seq_flanks, function(x) {
        map_dfr(names(re_seqs), function(re_name) {
            vmatchPattern(re_seqs[[re_name]], x)[[1]] %>%
                as.data.frame() %>% mutate(RE = re_name)
        }) %>% nrow %>% is_greater_than(0)
    })) %>% group_by(seq) %>%
    filter(!any(re_matches)) %>%
    ungroup()


random_controls_w_mut <- fragments_final %>% filter(!is.na(variant)) %>%
    distinct(seq) %>%
    # sample_n(random_control_w_mut_n*2, replace = random_control_w_mut_n*2 > nrow(.)) %>%
    mutate(fragment_initial = row_number(),
           GC = str_count(seq, "[GC]")/str_length(seq),
           shuffle_seq = map_chr(seq, ~ str_split(., "") %>% unlist %>% sample() %>% str_c(collapse = "")),
           shuffle_GC = str_count(shuffle_seq, "[GC]")/str_length(shuffle_seq),
           pos = floor((str_length(shuffle_seq) + 1)/2),
           ref = str_sub(shuffle_seq, pos, pos),
           alt = map2_chr(shuffle_seq, ref, ~ str_split(.x, "") %>% unlist %>% str_subset(.y, negate = T) %>% sample(1))) %>%
    mutate(variant = map2("ref", "alt", ~ c(.x, .y))) %>% unnest(variant) %>%
    mutate(allele = ifelse(variant == "ref", ref, alt)) %>%
    mutate(var_seq = pmap_chr(.,
                              function(shuffle_seq, pos, ref, allele, ...) {
                                  mutate_fragments(shuffle_seq, pos, ref, allele)
                              })) %>%
    mutate(var_seq_flanks  = paste0(str_sub(primer_5p, - flank_width),
                                    var_seq,
                                    str_sub(cloning_site_1, 1, flank_width))) %>%
    ungroup %>%
    mutate(re_matches = map_lgl(var_seq_flanks, function(x) {
        map_dfr(names(re_seqs), function(re_name) {
            vmatchPattern(re_seqs[[re_name]], x)[[1]] %>%
                as.data.frame() %>% mutate(RE = re_name)
        }) %>% nrow %>% is_greater_than(0)
    })) %>% group_by(seq) %>%
    filter(!any(re_matches)) %>%
    ungroup()


random_controls_sample <- random_controls %>%
    filter(fragment_initial %in% sample(unique(fragment_initial), random_control_n)) %>%
    mutate(fragment = dense_rank(fragment_initial)) %>%
    mutate(frag_id = paste0('randomfragment-', str_pad(fragment, width = 4, pad = 0), "_", "allele-", variant))

random_controls_w_mut_sample <- random_controls_w_mut %>%
    filter(fragment_initial %in% sample(unique(fragment_initial), random_control_w_mut_n)) %>%
    mutate(fragment = dense_rank(fragment_initial)) %>%
    mutate(frag_id = paste0('randomfragmentmut-', str_pad(fragment, width = 4, pad = 0), "_", "allele-", variant))

random_controls_allele  <- bind_rows(random_controls_sample, random_controls_w_mut_sample)

total_frags <- nrow(fragments_final) + nrow(random_controls_allele)

barcode_df <- read_tsv(snakemake@config$barcodes, col_names = F) %>%
    set_colnames("barcode") %>%
    mutate(barcode_w_flanks = paste0(str_sub(cloning_site_2, start = -flank_width),
                                     barcode,
                                     str_sub(primer_3p, start = 1, end = flank_width))) %>%
    bind_cols(map_dfc(as.list(re_seqs), vcountPattern, subject = .$barcode_w_flanks)) %>%
    mutate(match = purrr::reduce(select(., names(re_seqs)), `+`)) %>%
    filter(match == 0)


barcode_lib <- barcode_df %>%
    sample_n(total_frags * barcodes_per_frag) %>% pull(barcode) %>%
    split(rep(1:total_frags, each = 10)) %>%
    map(sort)




barcoded_library <-
    bind_rows(fragments_final %>% select(frag_id, frag_seq = var_seq_corrected),
              random_controls_allele %>% select(frag_id, frag_seq = var_seq)) %>%
    mutate(barcode = barcode_lib) %>%
    unnest(barcode) %>%
    group_by(frag_id) %>%
    mutate(barcode_id = row_number()) %>%
    ungroup

if (stuffer_bp > 0) {
    gc_bp <- round(stuffer_gc * stuffer_bp)
    at_bp <- stuffer_bp - gc_bp
    stuffer_df <- tibble(
        stuffer = map_chr(1:nrow(barcoded_library),
        ~ paste0(sample(c(sample(c("G", "C"), gc_bp, replace = T),
                            sample(c("A", "T"), at_bp, replace = T)), replace = F), collapse = ""))) %>%
        mutate(stuffer_w_flanks = paste0(str_sub(cloning_site_1, start = 2),
                                     stuffer,
                                     str_sub(cloning_site_2, end = -2)))  %>%
        bind_cols(map_dfc(as.list(re_seqs), vcountPattern, subject = .$stuffer_w_flanks)) %>%
        mutate(match = purrr::reduce(select(., names(re_seqs)), `+`))


    while(any(stuffer_df$match > 0)) {
        stuffer_df_match <- stuffer_df %>% filter(match > 0)

        stuffer_df_match <- tibble(
            stuffer = map_chr(1:nrow(stuffer_df_match),
                              ~ paste0(sample(c(sample(c("G", "C"), gc_bp, replace = T),
                                                sample(c("A", "T"), at_bp, replace = T)), replace = F), collapse = ""))) %>%
            mutate(stuffer_w_flanks = paste0(str_sub(cloning_site_1, start = 2),
                                             stuffer,
                                             str_sub(cloning_site_2, end = -2)))  %>%
            bind_cols(map_dfc(as.list(re_seqs), vcountPattern, subject = .$stuffer_w_flanks)) %>%
            mutate(match = purrr::reduce(select(., names(re_seqs)), `+`))

        stuffer_df <- bind_rows(stuffer_df %>% filter(match == 0),
                                stuffer_df_match)
    }

    stuffer_seqs <- stuffer_df$stuffer

} else {
    stuffer_seqs <- ""
}


final_library <- barcoded_library %>%
    mutate(stuffer = stuffer_seqs) %>%
    mutate(oligo_id = paste0(frag_id, "_barcode-", str_pad(barcode_id, width = 2, pad = 0)),
           oligo = paste0(primer_5p, frag_seq, cloning_site_1, stuffer, cloning_site_2, barcode, primer_3p))



if (n_sub_libraries > 1) {

    final_library <- final_library %>%
        mutate(frag_pair_id = str_extract(frag_id, "^fragment-\\d+"))

    sub_libraries <- final_library %>%
        filter(!is.na(frag_pair_id)) %>%
        distinct(frag_pair_id) %>%
        mutate(sub_library = sample(row_number() %% n_sub_libraries) + 1) %>%
        left_join(final_library %>% filter(!is.na(frag_pair_id)))

    random_controls_sub <- final_library %>%
        filter(str_detect(frag_id, "^randomfragment"))

    sub_libraries <- sub_libraries %>%
        group_by(sub_library) %>% group_nest() %>%
        mutate(data = map(data, ~ bind_rows(., random_controls_sub)))


    walk2(sub_libraries$data, sub_libraries$sub_library,
          ~ .x %>% select(oligo_id, oligo) %>%
              write_csv(str_replace(snakemake@output$oligos, "\\.csv",
                                    paste0("_sublib-", .y, ".csv"))))


    walk2(sub_libraries$data, sub_libraries$sub_library,
          ~ .x %>% select(oligo_id, barcode) %>%
              write_csv(str_replace(snakemake@output$barcode_ref, "\\.csv",
                                    paste0("_sublib-", .y, ".csv"))))

}

final_library %>% select(oligo_id, oligo) %>%
    write_csv(snakemake@output$oligo)

final_library %>% select(oligo_id, barcode) %>%
    write_csv(snakemake@output$barcode_ref)






variant_final <- fragments_final %>% filter(!is.na(variant)) %>% unnest(data) %>%
    select(fragment, variant, snp, chr, pos, allele, ref, alt,
           ref_fragment = seq, var_fragment = var_seq_corrected,
           frag_start = window_start, frag_end = window_end, frag_id)
write_csv(variant_final, snakemake@output$variant_ref)

# haplotype_final <- fragments_final %>% filter(!is.na(haplotype)) %>% unnest(data) %>%
#     select(fragment, haplotype, snp, chr, pos, allele, ref, alt,
#            ref_fragment = seq, var_fragment = var_seq_corrected,
#            frag_start = window_start, frag_end = window_end, frag_id)
# write_csv(haplotype_final, "./data/raw/lib3_design/skin_disease_haplotype_ref.csv")

random_controls_final <-  random_controls_allele %>%
    select(fragment, variant, pos, allele, ref, alt, ref_fragment = shuffle_seq, var_fragment = var_seq, frag_id)
write_csv(random_controls_final, snakemake@output$random_ctrl_ref)

save.image("logs/clean_index_snps.RData")

log <- file(snakemake@log[[1]], open="wt")
sink(log, type = "message")
sink(log, type = "output")


library(SNPlocs.Hsapiens.dbSNP144.GRCh37)
library(SNPlocs.Hsapiens.dbSNP151.GRCh38)
library(XtraSNPlocs.Hsapiens.dbSNP141.GRCh38)
library(TxDb.Hsapiens.UCSC.hg38.knownGene)
library(magrittr)
library(tidyverse)

source("lib/helpers.R")


hg19_to_hg38_chain <- import.chain("assets/hg19ToHg38.over.chain")

# threads <- 4

# if (threads > 1) {
#     library(doMC)
#     registerDoMC(cores = threads)
#     do_parallel <- T
# } else {
#     do_parallel <- F
# }

index_snp_table <- read_tsv(snakemake@input$gwas,
                            col_types = cols(.default = col_character()), quote = "")
# index_snps <- read_tsv("./data/raw/lib3_design/skin_disease_index_snps.txt")

all(str_detect(index_snp_table$SNPS, "^rs\\d+$") |
        str_detect(index_snp_table$SNPS, "^chr[0-9XY]+:\\d+$"))


index_snps <- index_snp_table %>%
    select(disease = Disease, gwas_snp = SNPS, chr = CHR_ID, pos = CHR_POS,
           pubmed = PUBMEDID, sample = `INITIAL SAMPLE SIZE`) %>%
    mutate(coord_b38 = ifelse(is.na(chr), NA, paste0("chr", chr, ":", pos))) %>%
    mutate(coord_b38 = ifelse(is.na(coord_b38) & str_detect(gwas_snp, "chr.+:\\d+"), gwas_snp, coord_b38))


index_snps_gr <- index_snps %>%
    filter(!is.na(coord_b38)) %>%
    extract(coord_b38, c("chr", "pos"), "chr([0-9XY]+):([0-9]+)") %>%
    mutate(start = pos, end = pos) %>%
    makeGRangesFromDataFrame(keep.extra.columns = T)

snps_find_rsid_b37 <- snpsByOverlaps(SNPlocs.Hsapiens.dbSNP144.GRCh37, index_snps_gr)
snps_find_rsid_b38 <- snpsByOverlaps(SNPlocs.Hsapiens.dbSNP151.GRCh38, index_snps_gr)
snps_find_rsid_b38_xtra <- snpsByOverlaps(XtraSNPlocs.Hsapiens.dbSNP141.GRCh38,
                                          `seqlevelsStyle<-`(index_snps_gr, "dbSNP")) %>%
    `seqlevelsStyle<-`("NCBI")



snps_find_rsid_b37_tbl <-
   as.data.frame(snps_find_rsid_b37) %>%
    mutate(coord_b37 = paste0("chr", seqnames, ":", pos)) %>%
    select(rs_id_rescue_b37 = RefSNP_id, coord_b37)

snps_find_rsid_b38_tbl <-
    bind_rows(as.data.frame(snps_find_rsid_b38) %>%
                  mutate(coord_b38 = paste0("chr", seqnames, ":", pos)),
              as.data.frame(snps_find_rsid_b38_xtra) %>%
                  mutate(coord_b38 = paste0("chr", seqnames, ":", start))) %>%
    select(rs_id_rescue = RefSNP_id, coord_b38)
# snps_find_rsid_b38_tbl <- snps_find_rsid_b38 %>% as.data.frame() %>%
#     mutate(coord_b38 = paste0("chr", seqnames, ":", pos)) %>%
#     select(rs_id_rescue = RefSNP_id, coord_b38)

index_snps_cleaned <- left_join(index_snps, snps_find_rsid_b38_tbl) %>%
    mutate(index_snp = ifelse(str_detect(gwas_snp, "^rs\\d+"), gwas_snp,
                              ifelse(!is.na(rs_id_rescue), rs_id_rescue, NA))) %>%
    left_join(snps_find_rsid_b37_tbl, by = c("coord_b38" = "coord_b37")) %>%
    mutate(coord_b37 = ifelse(is.na(index_snp) & !is.na(rs_id_rescue_b37), coord_b38, NA),
           coord_b38 = ifelse(is.na(index_snp) & !is.na(rs_id_rescue_b37), NA, coord_b38),
           index_snp = ifelse(is.na(index_snp) & !is.na(rs_id_rescue_b37), rs_id_rescue_b37, index_snp)) %>%
    mutate(index_snp = ifelse(is.na(index_snp), gwas_snp, index_snp)) %>%
    select(disease, gwas_snp, index_snp, coord_b38, coord_b37, pubmed, sample)


write_csv(index_snps_cleaned, snakemake@output$index_snps)

readRenviron(".Renviron")

save.image("logs/get_snps_in_ld.RData")

log <- file(snakemake@log[[1]], open="wt")
sink(log, type = "message")
sink(log, type = "output")

if (! "haploR" %in% rownames(installed.packages())) {
    options(repos = list(CRAN="http://cran.rstudio.com/"))
    install.packages("haploR")
}

library(SNPlocs.Hsapiens.dbSNP144.GRCh37)
library(SNPlocs.Hsapiens.dbSNP151.GRCh38)
library(XtraSNPlocs.Hsapiens.dbSNP141.GRCh38)
library(TxDb.Hsapiens.UCSC.hg38.knownGene)
library(LDlinkR)
library(haploR)
library(VariantAnnotation)
library(magrittr)
library(tidyverse)

source("lib/helpers.R")

set.seed(snakemake@config$seed)

hg19_to_hg38_chain <- import.chain("assets/hg19ToHg38.over.chain")

# threads <- 4

# if (threads > 1) {
#     library(doMC)
#     registerDoMC(cores = threads)
#     do_parallel <- T
# } else {
#     do_parallel <- F
# }

index_snps_cleaned <- read_csv(snakemake@input$index_snps)
# index_snps <- read_tsv("./data/raw/lib3_design/skin_disease_index_snps.txt")

r2_threshold <- snakemake@config$r2_threshold
r2_threshold_pop_specific <- snakemake@config$r2_threshold_pop_spec


pops <- snakemake@config$pops
# pops <- c("EUR", "AFR", "AMR", "EAS", "SAS", "ALL")

if (!is.null(snakemake@config$gwas_pop_key)) {
    gwas_pop_key <- read_tsv(snakemake@config$gwas_pop_key)

    sample_types <- c("individuals?",
                      "cases?",
                      "controls?",
                      "men",
                      "women",
                      "boys?",
                      "girls?",
                      "adults?",
                      "adolescents?",
                      "children and adolescents",
                      "children",
                      "infants?",
                      "neonates?",
                      "mothers?",
                      "fathers?",
                      "parents?",
                      "males?",
                      "females?",
                      "users?",
                      "non-users?",
                      "families",
                      "trios?",
                      "responders?",
                      "non-responders?",
                      "attempters?",
                      "nonattempters?",
                      "alcohol drinkers?",
                      "drinkers?",
                      "non-drinkers?",
                      "smokers?",
                      "non-smokers?",
                      "donors?",
                      "twin pairs?",
                      "twins?",
                      "child sibling pairs?",
                      "fetuses",
                      "offspring",
                      "early adolescents?",
                      "remitters?",
                      "non-remitters?",
                      "athletes?",
                      "Individuals?",
                      "indivduals?",
                      "triads?",
                      "patients?",
                      "pairs?",
                      "case-parent trios?",
                      "recipients?",
                      "affected child",
                      "long sleepers?",
                      "short sleepers?",
                      "unaffected relatives?",
                      "carriers?",
                      "non-carriers?",
                      "cell lines?",
                      "indiviudals?",
                      "referents?",
                      "individuuals?",
                      "duos?",
                      "indivdiuals?",
                      "inidividuals?")

    number_regex <- "(?:(?<=(?:\\s|\\b))\\d+(?:\\,\\d+)*(?=\\s))"
    type_regex <- paste0("(?:", paste0(sample_types, collapse = "|"), ")")


    full_regex <- paste0(
        "(", number_regex, ")", # greedy match first number
        "\\s*((?:(?!.*", type_regex, ").*)|(?:.*?))\\s*", # Greedy match rest if no sample type in lookahead, or passive match
        "(", type_regex,  "?(?!.*", type_regex, "))") # Match last sample type by ensuring no sample type in lookahead

    # split_regex <- "(?<!\\d)(,[\\s\\,]*| and )(?=[\\sA-Aa-z]*[0-9]+[,0-9]*[0-9]+\\s)"
    split_regex <- paste0("((?:,+[,\\s]*\\s+)|(?:and ))(?=[\\sA-Aa-z]*", number_regex, ")")


    sample_terms <- index_snps_cleaned %>%
        distinct(pubmed, sample) %>%
        mutate(split_sample = str_split(sample, split_regex)) %>%
        unnest(split_sample)


    full_matches <- bind_cols(sample_terms,
                              str_match(sample_terms$split_sample, full_regex) %>%
                                  set_colnames(c("match", "number", "capture", "type")) %>%
                                  as_tibble())

    study_key_table <- full_matches %>%
        distinct(pubmed, sample, split_sample, capture) %>%
        rename(term = capture) %>%
        left_join(gwas_pop_key) %>%
        filter(!is.na(code))

    index_snps_pop_match <- index_snps_cleaned %>%
        left_join(study_key_table) %>%
        distinct() %>%
        group_by(disease, gwas_snp, index_snp, coord_b38, coord_b37, pubmed, sample) %>%
        summarise(pops = paste0(sort(unique(unlist(str_split(code, ",")))), collapse = ",")) %>%
        ungroup()


    write_tsv(index_snps_pop_match, "outs/gwas_study_index_snps_matched_populations.tsv")

    index_snps_pop_match %>%
        group_by(disease, pubmed, sample, pops) %>%
        summarise(n_snps = n_distinct(index_snp, na.rm = T)) %>%
        write_tsv("outs/gwas_study_matched_populations.tsv")


} else {
    index_snps_pop_match <- tibble(disease = character(),
                                   pubmed = character(),
                                   sample = character(),
                                   index_snp = character(),
                                   pops = character())
}

max_pops <- snakemake@config$max_pops


index_snps_pop_match_filtered <- index_snps_pop_match %>%
    filter(!is.na(pops) & pops != "") %>%
    filter(map_lgl(str_split(pops, ","), ~ length(.) <= max_pops))


index_snps_pop_all <- crossing(index_snp = unique(index_snps_cleaned$index_snp),
                               pop = pops) %>%
    bind_rows(index_snps_pop_match_filtered %>%
                  mutate(pop = str_split(pops, ",")) %>%
                  unnest(pop) %>%
                  distinct(index_snp, pop))


snps_to_query <- index_snps_pop_all %>%
    filter(str_detect(index_snp, "rs\\d+"),
           !is.na(pop) & pop != "") %>%
    mutate(r2_threshold = ifelse(is.null(r2_threshold_pop_specific) | pop == "ALL",
                                 r2_threshold, r2_threshold_pop_specific))

out_dir <- "outs/SNPS_LDlink"

dir.create(out_dir, showWarnings = F, recursive = T)

ldlink_results <- snps_to_query %>%
    mutate(ldlink_results = pmap(list(index_snp, pop, r2_threshold),
        ~ query_ldlink(snp = ..1, pop = ..2, r2 = ..3, out_dir = out_dir, retry_errors = snakemake@config$retry_errors)))

ldlink_results_table <- ldlink_results %>%
    unnest(ldlink_results) %>%
    filter(R2 >= r2_threshold)

write_tsv(ldlink_results_table, "outs/ldlink_full_results.txt")

haploreg_pops <- c("AFR" = "AFR",
                   "AMR" = "AMR",
                   "EAS" = "ASN",
                   "EUR" = "EUR",
                   "SAS" = "ASN")

out_dir_haploreg <- "outs/SNPS_HaploReg"
dir.create(out_dir_haploreg, showWarnings = F, recursive = T)

haploreg_results <- snps_to_query %>%
    filter(pop %in% names(haploreg_pops)) %>%
    mutate(pop = haploreg_pops[pop]) %>%
    group_by(pop, r2_threshold) %>%
    summarise(index_snps = list(sort(index_snp))) %>%
    mutate(haploreg_results = pmap(list(index_snps, pop, r2_threshold),
        ~ query_haploreg(snps = ..1, pop = ..2, r2 = ..3,
                        force = T, out_dir = out_dir_haploreg))) %>%
    ungroup()

if (nrow(haploreg_results) > 0) {
    haploreg_results_table <- haploreg_results %>%
        select(pop, r2_threshold, haploreg_results) %>%
        unnest(haploreg_results) %>%
        select(index_snp = query_snp_rsid, everything()) %>%
        filter(r2 >= r2_threshold)
} else {
    haploreg_results_table <- tibble(
        index_snp = character(),
        pop = character(),
        chr = character(),
        pos_hg38 = character(),
        r2 = double(),
        D = double(),
        is_query_snp = double(),
        rsID = character(),
        ref = character(),
        alt = character()
    )
}

write_tsv(haploreg_results_table, "outs/haploreg_full_results.txt")



# Harmonize rsIDs and genomic coordinates for all LD SNPs from both sources

# ldlink_results_table <- read_tsv("./data/raw/lib3_design/ldlink_full_results.txt")
# haploreg_results_table <- read_tsv("./data/raw/lib3_design/haploreg_full_results.txt")

## LDlink data is in hg19 coordinates
ldlink_snps <- ldlink_results_table %>%
    extract(Alleles, c("ref", "alt"), "([ACGT-]+)\\/([ACGT-]+)", remove = F) %>%
    filter(!is.na(ref), !is.na(alt))

ldlink_snps_b38 <- ldlink_snps %>%
    extract(Coord, c("chr", "start"), "(chr[0-9XY]+):(\\d+)", remove = F) %>%
    mutate(end = start) %>%
    select(seqnames = chr, start, end, snp = RS_Number, index_snp, coord_b37 = Coord, ref, alt) %>%
    makeGRangesFromDataFrame(keep.extra.columns = T) %>%
    liftOver(hg19_to_hg38_chain) %>% unlist %>%
    as_tibble() %>%
    mutate(coord_b38 = paste0(seqnames, ":", start),
           snp = ifelse(is.na(snp) | !str_detect(snp, "^rs\\d+"), coord_b38, snp)) %>%
    select(snp, coord_b38, ref, alt, index_snp, coord_b37) %>%
    distinct()



## HaploReg data is in hg38 coordinates, but not all snps returned have genome coordinates

haploreg_snps <- haploreg_results_table %>%
    mutate(coord_b38 = ifelse(is.na(chr), NA, paste0("chr", chr, ":", pos_hg38))) %>%
    select(snp = rsID, coord_b38, ref, alt, index_snp)

haploreg_snps_no_coord <- haploreg_snps %>%
    filter(is.na(coord_b38)) %>% pull(snp) %>% unique()

## Try to rescue location data from SNPlocs packages and GTEx variant info

haploreg_snps_find_locs_b38 <- snpsById(SNPlocs.Hsapiens.dbSNP151.GRCh38, haploreg_snps_no_coord, ifnotfound = "drop") %>%
    GRanges() %>% as_tibble() %>%
    mutate(chr = str_replace(as.character(seqnames), "^(chr|ch)", "")) %>%    select(chr, pos_b38 = start, snp = RefSNP_id)
haploreg_snps_find_locs_b38_xtra <- snpsById(XtraSNPlocs.Hsapiens.dbSNP141.GRCh38, haploreg_snps_no_coord, ifnotfound = "drop") %>%
    GRanges() %>% as_tibble() %>%
    mutate(chr = str_replace(as.character(seqnames), "^(chr|ch)", "")) %>%
    select(chr, pos_b38 = start, snp = RefSNP_id)


if (!is.null(snakemake@config$gtex_table)) {
    gtex_var_map <- read_tsv(snakemake@config$gtex_table,
                             col_types = "c-----cc") %>%
        dplyr::rename(rs_id = "rs_id_dbSNP151_GRCh38p7")

    haploreg_snps_find_locs_gtex <- gtex_var_map %>% filter(rs_id %in% haploreg_snps_no_coord) %>%
        extract(variant_id, c("chr", "pos_b38"), "^chr([0-9XY]+)_(\\d+)") %>%
        mutate(pos_b38 = as.numeric(pos_b38)) %>%
        select(chr, pos_b38, snp = rs_id)
} else {
    haploreg_snps_find_locs_gtex <- tibble()
}


haploreg_snps_find_locs_combined <- bind_rows(
    haploreg_snps_find_locs_b38,
    haploreg_snps_find_locs_b38_xtra,
    haploreg_snps_find_locs_gtex
) %>% distinct %>%
    mutate(coord_b38_rescue = paste0("chr", chr, ":", pos_b38)) %>%
    select(snp, coord_b38_rescue)


haploreg_snps_b38 <- haploreg_snps %>%
    left_join(haploreg_snps_find_locs_combined) %>%
    mutate(coord_b38 = as.character(ifelse(is.na(coord_b38), coord_b38_rescue, coord_b38))) %>%
    select(-coord_b38_rescue) %>%
    distinct()



## Combine LD SNPs


ld_snps_b38 <- bind_rows(
    ldlink_snps_b38 %>% mutate(source = "LDlink"),
    haploreg_snps_b38 %>% mutate(source = "HaploReg")
)



## Get TxDb annotations

ld_snps_b38_gr <- ld_snps_b38 %>%
    extract(coord_b38, c("seqnames", "start"), "(.+):(\\d+)") %>%
    filter(!is.na(seqnames), !is.na(start)) %>%
    mutate(end = start) %>%
    select(seqnames, start, end, snp) %>%
    makeGRangesFromDataFrame(keep.extra.columns = T)

ld_snps_txdb_loc <- locateVariants(ld_snps_b38_gr, TxDb.Hsapiens.UCSC.hg38.knownGene, AllVariants())

ld_snps_txdb_loc_df <- as_tibble(ld_snps_txdb_loc) %>%
    transmute(coord_b38 = paste0(seqnames, ":", start),
              txdb_annot = LOCATION) %>%
    distinct() %>%
    group_by(coord_b38) %>%
    summarise(txdb_annot = paste0(txdb_annot, collapse = ";"))

ld_snps_b38_annot <- left_join(ld_snps_b38,
                               ld_snps_txdb_loc_df, by = "coord_b38")


write_tsv(ld_snps_b38_annot, snakemake@output$ld_snps)