# Establish chromosome vector
chr <- c(1:19)

# read in covariate file
metadata <- readr::read_csv(file = "data/4WC_covar.csv", 
                            col_types = c(sample = "c",
                                          generation = "n",
                                          Sex = "c"), 
                            show_col_types = T)

# match sample names to what is encoded in genotype files
geno_header <- c(as.character(readr::read_csv(file = "data/STITCH_data/4WC_ddRADseq/geno1.csv", 
                                              col_names = F, skip = 3, n_max = 1)))[-1]

new_sampleID <- c()
for(i in 1:length(metadata$SampleID)){
  sample_index <- grep(geno_header, pattern = gsub(metadata$SampleID[i], 
                                                   pattern = "-",
                                                   replacement = "."))
  new_sampleID[i] <- geno_header[sample_index]
}

# Write updated metadata file
updated_4WC_metadata <- metadata %>%
  dplyr::mutate(newSampID = new_sampleID) %>%
  dplyr::select(-SampleID,-Sex) %>%
  dplyr::select(newSampID, everything()) %>%
  dplyr::rename(SampleID = newSampID)
write.csv(updated_4WC_metadata, file = "data/STITCH_data/4WC_ddRADseq/4WC_ddRADseq_crossinfo.csv", quote = F, row.names = F)

# Write updated metadata file
sex_4WC <- metadata %>%
  dplyr::mutate(newSampID = new_sampleID) %>%
  dplyr::select(newSampID, Sex) %>%
  dplyr::rename(SampleID = newSampID) 
write.csv(sex_4WC, file = "data/STITCH_data/4WC_ddRADseq/sex_4WC.csv", quote = F, row.names = F)

# Write control file
qtl2::write_control_file(output_file = "data/STITCH_data/4WC_ddRADseq/4WC_ddRADseq.json",
                         crosstype="genail4",
                         description="4WC_ddRADseq",
                         founder_geno_file=paste0("foundergeno", chr, ".csv"), 
                         founder_geno_transposed=TRUE,
                         gmap_file=paste0("gmap", chr, ".csv"),
                         pmap_file=paste0("pmap", chr, ".csv"),
                         geno_file=paste0("geno", chr, ".csv"),
                         geno_transposed = TRUE,
                         geno_codes=list(A=1, H=2, B=3), 
                         sex_file = "sex_4WC.csv",
                         sex_codes=list(F="Female", M="Male"),
                         # crossinfo_covar=c("Generation","A","B","C","D"), 
                         crossinfo_file = "4WC_ddRADseq_crossinfo.csv",
                         overwrite = T)

# Load in the cross object
ddRADseq_4WC <- qtl2::read_cross2("data/STITCH_data/4WC_ddRADseq/4WC_ddRADseq.json")

# Also loading in the cross object from GigaMUGA genotypes
load("data/4WC_cross.RData")

# Drop null markers
ddRADseq_4WC <- qtl2::drop_nullmarkers(ddRADseq_4WC)

# Reordering genotypes so that most common allele in founders is first
for(chr in seq_along(ddRADseq_4WC$founder_geno)) {
  fg <- ddRADseq_4WC$founder_geno[[chr]]
  g <- ddRADseq_4WC$geno[[chr]]
  f1 <- colSums(fg==1)/colSums(fg != 0)
  
  fg[fg==0] <- NA
  g[g==0] <- NA
  
  fg[,f1 < 0.5] <- 4 - fg[,f1 < 0.5]
  g[,f1 < 0.5]  <- 4 - g[,f1 < 0.5]
  
  fg[is.na(fg)] <- 0
  g[is.na(g)] <- 0
  
  ddRADseq_4WC$founder_geno[[chr]] <- fg
  ddRADseq_4WC$geno[[chr]] <- g
}

# Calculate the percent of missing genotypes per sample
percent_missing <- qtl2::n_missing(ddRADseq_4WC, "ind", "prop")*100
missing_genos_df <- data.frame(names(percent_missing), percent_missing) %>%
  `colnames<-`(c("sample","percent_missing"))

# Plot missing genotypes per sample
missing_genos_plot <- ggplot(data = missing_genos_df, mapping = aes(x = reorder(sample, percent_missing),
                                                                    y = percent_missing)) +
  theme_bw() +
  geom_point(shape = 21) +
  labs(title = "4WC Missing Genotypes") +
  theme(legend.position = "bottom",
        panel.grid = element_blank(),
        axis.text.x = element_blank(),
        axis.ticks.x = element_blank())
plotly::ggplotly(missing_genos_plot)

percent_missing_cutoff <- 10

# Determine if any samples are duplicates based on genetic similarity
cg <- qtl2::compare_geno(ddRADseq_4WC, cores=0)
qtl2::plot_compare_geno(x = cg, rug = T, main = "4WC - ddRADseq Data")

cgGM <- qtl2::compare_geno(X4WC_cross, cores=0)
qtl2::plot_compare_geno(x = cgGM, rug = T, main = "4WC - GigaMUGA Data")

# Insert pseudomarkers
map <- qtl2::insert_pseudomarkers(ddRADseq_4WC$gmap, step = 1)

# Calculate genotype probs
dir.create("results/4WC_ddRADseq_pr")

fpr <- suppressWarnings(qtl2fst::calc_genoprob_fst(cross = ddRADseq_4WC, 
                                                   map = map,
                                                   fbase = "pr",
                                                   fdir = "results/4WC_ddRADseq_pr",
                                                   error_prob=0.002, 
                                                   overwrite=TRUE,cores = parallel::detectCores()))

# Make viterbi
m <- maxmarg(fpr, minprob=0.5)

# Phase genotypes
ph <- qtl2::guess_phase(cross = ddRADseq_4WC, geno = m)

# Write Plots
X4WCcolors <- c(qtl2::CCcolors[6],
  "#5ADBFF",
  "#153B50",
  qtl2::CCcolors[7])
names(X4WCcolors)[2:3] <- c("POHN","GOR")
for(i in 1:nrow(ddRADseq_4WC$cross_info)){
  png(file=paste0("output/plot_onegeno_",rownames(ddRADseq_4WC$cross_info)[i],".png"))
  qtl2::plot_onegeno(geno = ph, 
                     map = map,
                     ind = i, 
                     col = X4WCcolors) # add legends here
  legend(17, 85, 
         legend=c("A", "B", "C", "D"),
         fill=X4WCcolors, cex=0.8)
  dev.off()
}

# GigaMUGA haplotypes
load("data/4WC_genoprobs.RData")
gm_m <- maxmarg(pr, minprob=0.5)
gm_ph <- qtl2::guess_phase(cross = X4WC_cross, geno = gm_m)
gm_map <- qtl2::insert_pseudomarkers(X4WC_cross$gmap, step = 1)

# Example plots
indivs <- sample(seq(1:nrow(ddRADseq_4WC$cross_info)), size = 5, replace = F)

# ddRADseq
plot_indiv <- rownames(ddRADseq_4WC$cross_info)[indivs[1]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = X4WCcolors,
                   main = plot_indiv) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

# DO
do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(X4WC_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = X4WCcolors,
                   main = rownames(X4WC_cross$cross_info)[do_sample_index]) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

plot_indiv <- rownames(ddRADseq_4WC$cross_info)[indivs[2]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = X4WCcolors,
                   main = plot_indiv) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(X4WC_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = X4WCcolors,
                   main = rownames(X4WC_cross$cross_info)[do_sample_index]) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

plot_indiv <- rownames(ddRADseq_4WC$cross_info)[indivs[3]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = X4WCcolors,
                   main = plot_indiv) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(X4WC_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = X4WCcolors,
                   main = rownames(X4WC_cross$cross_info)[do_sample_index]) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

plot_indiv <- rownames(ddRADseq_4WC$cross_info)[indivs[4]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = X4WCcolors,
                   main = plot_indiv) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(X4WC_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = X4WCcolors,
                   main = rownames(X4WC_cross$cross_info)[do_sample_index]) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

plot_indiv <- rownames(ddRADseq_4WC$cross_info)[indivs[5]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = X4WCcolors,
                   main = plot_indiv) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(X4WC_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = X4WCcolors,
                   main = rownames(X4WC_cross$cross_info)[do_sample_index]) # add legends here
legend(17, 85, legend=c("A", "B", "C", "D"), fill=X4WCcolors, cex=0.8)

# Establish chromosome vector
chr <- c(1,3:19)

# read in covariate file
metadata <- readr::read_csv(file = "data/DO_covar.csv", 
                            col_types = c(sample = "c",
                                          Sex = "c",
                                          generation = "n"), 
                            show_col_types = T)

# match sample names to what is encoded in genotype files
geno_header <- c(as.character(readr::read_csv(file = "data/STITCH_data/DO_ddRADseq/geno1.csv", 
                                              col_names = F, skip = 3, n_max = 1)))[-1]

new_sampleID <- c()
for(i in 1:length(metadata$SampleID)){
  sample_index <- grep(geno_header, pattern = gsub(metadata$SampleID[i], 
                                                   pattern = "-",
                                                   replacement = "."))
  if(length(sample_index) == 0){
    print(paste("sample from metadata not sequenced:", metadata$SampleID[i]))
    new_sampleID[i] <- NA
  } else {
    new_sampleID[i] <- geno_header[sample_index]
  }
  
}

# Write updated metadata file
updated_DO_metadata <- metadata %>%
  dplyr::mutate(newSampID = new_sampleID) %>%
  dplyr::select(-SampleID) %>%
  dplyr::select(newSampID, everything()) %>%
  dplyr::rename(SampleID = newSampID) %>%
  dplyr::filter(!is.na(SampleID))
write.csv(updated_DO_metadata, file = "data/STITCH_data/DO_ddRADseq/DO_ddRADseq_crossinfo.csv", quote = F, row.names = F)

# Write control file
qtl2::write_control_file(output_file = "data/STITCH_data/DO_ddRADseq/DO_ddRADseq.json",
                         crosstype="do",
                         description="DO_ddRADseq",
                         founder_geno_file=paste0("foundergeno", chr, ".csv"), 
                         founder_geno_transposed=TRUE,
                         gmap_file=paste0("gmap", chr, ".csv"),
                         pmap_file=paste0("pmap", chr, ".csv"),
                         geno_file=paste0("geno", chr, ".csv"),
                         geno_transposed = TRUE,
                         geno_codes=list(A=1, H=2, B=3), 
                         sex_covar = "Sex",
                         sex_codes=list(F="Female", M="Male"), 
                         covar_file = "DO_ddRADseq_crossinfo.csv",
                         crossinfo_covar="Generation",
                         overwrite = T)

# Load in the cross object
ddRADseq_DO <- qtl2::read_cross2("data/STITCH_data/DO_ddRADseq/DO_ddRADseq.json")

# Also loading in the cross object from GigaMUGA genotypes
load("data/DO_cross.RData")

# Drop null markers
ddRADseq_DO <- qtl2::drop_nullmarkers(ddRADseq_DO)

# Reordering genotypes so that most common allele in founders is first
for(chr in seq_along(ddRADseq_DO$founder_geno)) {
  fg <- ddRADseq_DO$founder_geno[[chr]]
  g <- ddRADseq_DO$geno[[chr]]
  f1 <- colSums(fg==1)/colSums(fg != 0)
  
  fg[fg==0] <- NA
  g[g==0] <- NA
  
  fg[,f1 < 0.5] <- 4 - fg[,f1 < 0.5]
  g[,f1 < 0.5]  <- 4 - g[,f1 < 0.5]
  
  fg[is.na(fg)] <- 0
  g[is.na(g)] <- 0
  
  ddRADseq_DO$founder_geno[[chr]] <- fg
  ddRADseq_DO$geno[[chr]] <- g
}

# Calculate the percent of missing genotypes per sample
percent_missing <- qtl2::n_missing(ddRADseq_DO, "ind", "prop")*100
missing_genos_df <- data.frame(names(percent_missing), percent_missing) %>%
  `colnames<-`(c("sample","percent_missing"))

# Plot missing genotypes per sample
missing_genos_plot <- ggplot(data = missing_genos_df, mapping = aes(x = reorder(sample, percent_missing),
                                                                    y = percent_missing)) +
  theme_bw() +
  geom_point(shape = 21) +
  labs(title = "DO Missing Genotypes") +
  theme(legend.position = "bottom",
        panel.grid = element_blank(),
        axis.text.x = element_blank(),
        axis.ticks.x = element_blank())
plotly::ggplotly(missing_genos_plot)

percent_missing_cutoff <- 10

# Determine if any samples are duplicates based on genetic similarity
cg <- qtl2::compare_geno(ddRADseq_DO, cores=0)
qtl2::plot_compare_geno(x = cg, rug = T, main = "DO - ddRADseq Data")

cgGM <- qtl2::compare_geno(DO_cross, cores=0)
qtl2::plot_compare_geno(x = cgGM, rug = T, main = "DO - GigaMUGA Data")

# Insert pseudomarkers
map <- qtl2::insert_pseudomarkers(ddRADseq_DO$gmap, step = 1)

# Calculate genotype probs
dir.create("results/DO_ddRADseq_pr")

fpr <- suppressWarnings(qtl2fst::calc_genoprob_fst(cross = ddRADseq_DO, 
                                                   map = map,
                                                   fbase = "pr",
                                                   fdir = "results/DO_ddRADseq_pr",
                                                   error_prob=0.002, 
                                                   overwrite=TRUE,cores = parallel::detectCores()))

# Make viterbi
m <- maxmarg(fpr, minprob=0.5)

# Phase genotypes
ph <- qtl2::guess_phase(cross = ddRADseq_DO, geno = m)

# Write Plots
for(i in 1:nrow(ddRADseq_DO$cross_info)){
  png(file=paste0("output/plot_onegeno_",rownames(ddRADseq_DO$cross_info)[i],".png"))
  qtl2::plot_onegeno(geno = ph, 
                     map = map,
                     ind = i, 
                     col = c(qtl2::CCcolors)) # add legends here
  legend(16, 70, 
         legend=c("A", "B", "C", "D",
                  "E", "F", "G", "H"),
         fill=c(qtl2::CCcolors), cex=0.8)
  dev.off()
}

# GigaMUGA haplotypes
load("data/DO_genoprobs.RData")
gm_m <- maxmarg(pr, minprob=0.5)
gm_ph <- qtl2::guess_phase(cross = DO_cross, geno = gm_m)
gm_map <- qtl2::insert_pseudomarkers(DO_cross$gmap, step = 1)

# Example plots
indivs <- sample(seq(1:nrow(ddRADseq_DO$cross_info)), size = 5, replace = F)

# ddRADseq
plot_indiv <- rownames(ddRADseq_DO$cross_info)[indivs[1]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = c(qtl2::CCcolors),
                   main = plot_indiv) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

# DO
do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(DO_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = c(qtl2::CCcolors),
                   main = rownames(DO_cross$cross_info)[do_sample_index]) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

plot_indiv <- rownames(ddRADseq_DO$cross_info)[indivs[2]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = c(qtl2::CCcolors),
                   main = plot_indiv) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(DO_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = c(qtl2::CCcolors),
                   main = rownames(DO_cross$cross_info)[do_sample_index]) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

plot_indiv <- rownames(ddRADseq_DO$cross_info)[indivs[3]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = c(qtl2::CCcolors),
                   main = plot_indiv) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(DO_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = c(qtl2::CCcolors),
                   main = rownames(DO_cross$cross_info)[do_sample_index]) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

plot_indiv <- rownames(ddRADseq_DO$cross_info)[indivs[4]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = c(qtl2::CCcolors),
                   main = plot_indiv) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(DO_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = c(qtl2::CCcolors),
                   main = rownames(DO_cross$cross_info)[do_sample_index]) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

plot_indiv <- rownames(ddRADseq_DO$cross_info)[indivs[5]]
qtl2::plot_onegeno(geno = ph, 
                   map = map,
                   ind = plot_indiv, 
                   col = c(qtl2::CCcolors),
                   main = plot_indiv) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)

do_sample_index  <- grep(pattern = gsub(strsplit(plot_indiv, split = "_")[[1]][1], 
                                        pattern = "ddRADseq.",
                                        replacement = ""),
                         x = gsub(rownames(DO_cross$cross_info), 
                                  pattern = "-", 
                                  replacement = "."))
qtl2::plot_onegeno(geno = gm_ph, 
                   map = gm_map,
                   ind = do_sample_index, 
                   col = c(qtl2::CCcolors),
                   main = rownames(DO_cross$cross_info)[do_sample_index]) # add legends here
legend(16, 65, legend=LETTERS[1:8], fill=c(qtl2::CCcolors), cex=0.8)