# set the library path
#.libPaths("/secondary/projects/bbc/tools/kin_R_packages/R-3.6.0_20190701_pkges")

knitr::opts_chunk$set(echo=TRUE, warning=TRUE, message=TRUE, cache=TRUE, dev=c('png','pdf'), fig.width=8, fig.height=8)

# define function for reading from 10x-style output directory and renaming by Celseq 192 barcode number as listed in Sagar...Gruen paper
read10x_and_bc_num_as_colnames <- function(tenx_dir, decoder){
  counts <- Seurat::Read10X(tenx_dir)

  # check that barcode sequences match
  stopifnot(identical(sort(colnames(counts)), sort(decoder$bc)))

  # rename columns with bc_num
  colnames(counts) <- decoder$bc_num[match(colnames(counts), decoder$bc)]

  tibble::as_tibble(counts, rownames="Gene") %>% dplyr::select(Gene, as.character(1:192))
}

# get input, output, params from the Snakemake pipeline
out_file <- snakemake@output[[1]] %>% str_replace(".html", ".txt")
out_file
starsolo_raw_dir <- dirname(snakemake@input[[1]])
starsolo_raw_dir # print out starsolo dir
decoder <- snakemake@params[["decoder"]]
decoder

# read barcode decoder from file in snakePipes installation
bc_nums_decoder <- read_tsv(decoder, col_names = c("bc_num","bc"))

# read the counts and rename to BC number
counts <- read10x_and_bc_num_as_colnames(starsolo_raw_dir, bc_nums_decoder)
str(counts) # print out str()

# output the counts to files
write_tsv(counts, out_file)

sessionInfo()

# save start time for script
start_tm <- Sys.time()
start_tm

# set the library path
knitr::opts_chunk$set(echo=TRUE, warning=TRUE, message=TRUE, cache=FALSE, dev=c('pdf'), fig.width=8, fig.height=8)

#outdir <- "./snprelate_out"
#if(!outdir %in% list.dirs()) dir.create(outdir, recursive = TRUE)
vcf.fn <- snakemake@input[[1]]
# convert to GDS
gds_file <- snakemake@params[["gds"]]
snpgdsVCF2GDS(vcf.fn, gds_file, method="biallelic.only")

# Figures
old_figs_folder <- snakemake@params[["figures_dir"]]
new.folder <- snakemake@params[["new_figures_dir"]]

# Set maximum missing rate
missingrate <- 0.05

genofile <- snpgdsOpen(gds_file)

set.seed(1000)
# Try different LD thresholds for sensitivity analysis
snpset <- snpgdsLDpruning(genofile, ld.threshold=0.5, missing.rate=missingrate) # using LD threshold default used in SNPhylo
# Get all selected snp id
snpset.id <- unlist(snpset)

set.seed(100)
ibs <- snpgdsIBS(genofile, num.thread=1, autosome.only=TRUE, missing.rate=missingrate, snp.id = snpset.id)
# we find 1- ibs so that the values become 'distances' where higher values represent more dissimilar pairs of samples
one_minus_ibs <- 1 - ibs$ibs
colnames(one_minus_ibs) <- ibs$sample.id
rownames(one_minus_ibs) <- ibs$sample.id

loc <- cmdscale(one_minus_ibs, k = 2)
#x <- loc[, 1]; y <- loc[, 2]
df <- tibble::as_tibble(loc, rownames="sample.id") %>%
  dplyr::rename(Dimension1=V1, Dimension2=V2)
#df$sample.id <- ibs$sample.id
# plot(x, y, xlab = "", ylab = "",
#     main = "Multidimensional Scaling Analysis (IBS)")
ggplot(df, aes(x=Dimension1, y=Dimension2, label=sample.id)) + 
  geom_point() + ggrepel::geom_label_repel() + ggtitle("MDS (based on IBS)")

ibs.hc <- snpgdsHCluster(ibs)
rv <- snpgdsCutTree(ibs.hc)

plot(rv$dendrogram, main="Dendrogram based on IBS")

#pca <- snpgdsPCA(genofile, snp.id=snpset.id, num.thread=1)
pca <- snpgdsPCA(genofile, num.thread=1, autosome.only=TRUE, missing.rate=missingrate, snp.id = snpset.id)
# variance proportion (%)
pc.percent <- pca$varprop*100
head(round(pc.percent, 2))
tab <- data.frame(sample.id = pca$sample.id,
    EV1 = pca$eigenvect[,1],    # the first eigenvector
    EV2 = pca$eigenvect[,2],    # the second eigenvector
    EV3 = pca$eigenvect[,3],
    EV4 = pca$eigenvect[,4],
    EV5 = pca$eigenvect[,5],
    stringsAsFactors = FALSE)
head(tab)
#plot(tab$EV1, tab$EV2, xlab="eigenvector 1", ylab="eigenvector 2")
ggplot(tab, aes(x=EV1, y=EV2, label=sample.id)) + geom_point() + ggrepel::geom_label_repel() + ggtitle("PC1 and 2")
ggplot(tab, aes(x=EV3, y=EV4, label=sample.id)) + geom_point() + ggrepel::geom_label_repel() + ggtitle("PC3 and 4")

set.seed(100)
ibs <- snpgdsIBS(genofile, num.thread=1, autosome.only=TRUE, missing.rate=missingrate)
# we find 1- ibs so that the values become 'distances' where higher values represent more dissimilar pairs of samples
one_minus_ibs <- 1 - ibs$ibs
colnames(one_minus_ibs) <- ibs$sample.id
rownames(one_minus_ibs) <- ibs$sample.id

loc <- cmdscale(one_minus_ibs, k = 2)
df <- tibble::as_tibble(loc, rownames="sample.id") %>%
  dplyr::rename(Dimension1=V1, Dimension2=V2)
ggplot(df, aes(x=Dimension1, y=Dimension2, label=sample.id)) + 
  geom_point() + ggrepel::geom_label_repel() + ggtitle("MDS (based on IBS)")

ibs.hc <- snpgdsHCluster(ibs)
rv <- snpgdsCutTree(ibs.hc)

plot(rv$dendrogram, main="Dendrogram based on IBS")

#pca <- snpgdsPCA(genofile, snp.id=snpset.id, num.thread=1)
pca <- snpgdsPCA(genofile, num.thread=1, autosome.only=TRUE, missing.rate=missingrate)
# variance proportion (%)
pc.percent <- pca$varprop*100
head(round(pc.percent, 2))
tab <- data.frame(sample.id = pca$sample.id,
    EV1 = pca$eigenvect[,1],    # the first eigenvector
    EV2 = pca$eigenvect[,2],    # the second eigenvector
    EV3 = pca$eigenvect[,3],
    EV4 = pca$eigenvect[,4],
    EV5 = pca$eigenvect[,5],
    stringsAsFactors = FALSE)
head(tab)
#plot(tab$EV1, tab$EV2, xlab="eigenvector 1", ylab="eigenvector 2")
ggplot(tab, aes(x=EV1, y=EV2, label=sample.id)) + geom_point() + ggrepel::geom_label_repel() + ggtitle("PC1 and 2")
ggplot(tab, aes(x=EV3, y=EV4, label=sample.id)) + geom_point() + ggrepel::geom_label_repel() + ggtitle("PC3 and 4")

closefn.gds(genofile)

dir.create(new.folder, recursive = TRUE)

file.copy(list.files(old_figs_folder, full.names = TRUE), new.folder, recursive = TRUE)

sessionInfo()

# output time taken to run script
end_tm <- Sys.time()
end_tm
end_tm - start_tm

shell:
    """
    fastq_screen --threads {threads} --outdir analysis/fastq_screen/ {input}
    """

shell:
   """
   gunzip -c {input} > {output}
   """

script:
    "scripts/read_star_solo_raw.Rmd"

shell:
    """
    samtools reheader -c 'perl -pe "s/^(@SQ.*)(\\tSM:\S+)/\$1\$2.{wildcards.CB}/"' {params.input} 1> {output} 2> {log.stderr}
    """

shell:
    """
    java -Xms16g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp -jar $PICARD MarkDuplicates \
    --INPUT {input} \
    --BARCODE_TAG 'UB' \
    --OUTPUT {output.bam} \
    --CREATE_INDEX true \
    --VALIDATION_STRINGENCY SILENT \
    --METRICS_FILE {output.metrics}
    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -XX:+UseParallelGC -XX:ParallelGCThreads={threads} -Djava.io.tmpdir=./tmp" \
        BaseRecalibrator \
        -R {params.ref_fasta} \
        -I {input} \
        -O {output} \
        {params.known_variants}
    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g  -XX:+UseParallelGC -XX:ParallelGCThreads={threads} -Djava.io.tmpdir=./tmp" \
        ApplyBQSR \
        --add-output-sam-program-record \
        -R {params.ref_fasta} \
        -I {input.bam} \
        -O {output} \
        --bqsr-recal-file {input.recal_table}
    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" \
    HaplotypeCaller \
    -R {params.ref_fasta} \
    -I {input.bam} \
    -O {output} \
    -ERC GVCF \
    -dont-use-soft-clipped-bases \
    --native-pair-hmm-threads {threads} \
    --standard-min-confidence-threshold-for-calling 20 \
    {params.contigs} 
    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" \
    GenomicsDBImport \
    {params.sample_gvcfs} \
    --genomicsdb-workspace-path {output.genomicsdb} \
    {params.contigs} 
    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" \
    GenotypeGVCFs \
    -R {params.ref_fasta} \
    -V gendb://{params.genomicsdb} \
    -O {output.vcf}

    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" \
    SortVcf \
    -I {input.vcf} \
    -O {output.sorted_vcf} \
    -SD {params.dictionary} 

    """

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" VariantsToTable -V {input} \
            -R {params.ref_fasta} --sequence-dictionary {params.dictionary} \
            -F CHROM -F POS -F REF -F ALT -F TYPE -F ANN -GF AD -O {output.raw}

    head -n1 {output.raw} | perl -npe 's/\\tANN\\t/\\tGENE\\t/; s/\.AD(?=[\t\n])//g' > {output.parsed}

    # Parse the comma-separated ANN column (column 5, 0-based counting). Each comma-separated element contains a '|' separated string; We extract fields 3 and 4, 0-based for the gene names.
    tail -n+2 {output.raw} | perl -F'\\t' -lane 'my %hash; $hash{{join("|",(split(/\|/,$_))[3..4])}}++ foreach split(",", $F[5]); $F[5] = join(",", sort(keys(%hash))); print join("\\t", @F)' >> {output.parsed}
    """

script:
    "scripts/snprelate.Rmd"

shell:
    """
    gatk --java-options "-Xms8g -Xmx{resources.mem_gb}g -Djava.io.tmpdir=./tmp" \
            ASEReadCounter \
            -R {params.ref_fasta} \
            -I {input} \
            -V {params.vcf} \
            --min-base-quality 20 \
            --min-mapping-quality 30 \
            -O {output}
    """