import sys
import gzip
import re

DEBUG = 0

ASCII_OFFSET = ord("!")
INDEL_PATTERN = re.compile(r"[+-](\d+)")
INDEL_STRING = "[+-]INTEGER[ACGTN]{INTEGER}"
ROUND = 3  # places to right of decimal point

min_coverage, min_proportion, min_qual = snakemake.params
pileup_file = snakemake.input[0]
out_file = snakemake.output[0]

min_coverage = int(min_coverage)
min_proportion = float(min_proportion)
min_qual = int(min_qual)

## function to process each line of pileup

def parse_pileup(line):

    # read fields from line of pileup
    chromosome, position, reference, coverage, pileup, quality = line.rstrip("\r\n").split("\t")

    # proportion is 0 if no consensus base, top base is N by default
    parsed = {
        "proportion": 0.0,
        "chromosome": chromosome,
        "position": int(position),
        "reference": reference,
        "coverage": int(coverage),
        "pileup": pileup,
        "quality": quality,
        "top_base": "N"}

    # if the base coverage is below the limit or above our acceptable max, call N
    if parsed["coverage"] < min_coverage:
        return parsed

    # uppercase pileup string for processing
    pileup = pileup.upper()

    # Remove start and stop characters from pileup string
    pileup = re.sub(r"\^.|\$", "", pileup)

    # Remove indels from pileup string
    start = 0

    while True:
        match = INDEL_PATTERN.search(pileup, start)

        if match:
            integer = match.group(1)
            pileup = re.sub(INDEL_STRING.replace("INTEGER", integer), "", pileup)
            start = match.start()
        else:
            break

    # get total base count and top base count
    total = 0
    top_base = "N"
    top_base_count = 0

    # uppercase reference base for comparison
    reference = reference.upper()

    base_counts = {"A": 0, "C": 0, "G": 0, "T": 0}

    quality_length = len(quality)

    for i in range(quality_length):

        # convert ASCII character to phred base quality
        base_quality = ord(quality[i]) - ASCII_OFFSET

        # only count high-quality bases
        if base_quality >= min_qual:

            currBase = pileup[i]

            if currBase in base_counts:
                base_counts[currBase] += 1

            else:
                base_counts[reference] += 1

            total += 1

    parsed["total"] = total

    # find top base
    for base in base_counts:
        if base_counts[base] > top_base_count:
            top_base = base
            top_base_count = base_counts[base]

    # if more that 0 bases processed
    if total > 0:
        prop = top_base_count / float(total)

        if prop >= min_proportion:
            parsed["proportion"] = prop
            parsed["top_base"] = top_base

    return parsed

## process pileup file

# read in input file
if not DEBUG:
    out_file = open(out_file, "w")

with open(pileup_file, "rt") as pileup_file:
    for pileup_file_line in pileup_file:

        parsed = parse_pileup(pileup_file_line)

        proportion = parsed["proportion"]

        # print to output file
        if DEBUG:
            print("\t".join([parsed["chromosome"], str(parsed["position"]), parsed["reference"], parsed["top_base"], str(round(proportion, ROUND)),
                parsed["pileup"], parsed["quality"]]))
        else:
            print("\t".join([parsed["chromosome"], str(parsed["position"]), parsed["reference"], parsed["top_base"], str(round(proportion, ROUND)),
                parsed["pileup"], parsed["quality"]]), file=out_file)

if not DEBUG:
    out_file.close()

import sys

# snakemake input and output files
sample_file = snakemake.input[0]
out_file = snakemake.output[0]

# minimum coverage threshold -- report percentage of bases covered at or
# beyond this depth
minCvg = int(snakemake.params[0])

totalBases = 0
coveredBases = 0
weightedAvg = 0
with open(sample_file, 'r') as sample:
    for line in sample:
        if line.startswith("genome"):
            chr, depth, numBases, size, fraction = line.rstrip('\n').split('\t')

            depth = int(depth)
            numBases = int(numBases)

            totalBases += numBases
            weightedAvg += depth * numBases

            if depth >= minCvg:
                coveredBases += numBases
avgCvg = 0
percCovered = 0

if totalBases > 0:
    avgCvg = float(weightedAvg) / float(totalBases)
    percCovered = float(coveredBases) / float(totalBases)

with open(out_file, 'w') as out:
    print("\t".join([str(round(avgCvg, 2)), str(round(percCovered, 2))]), file = out)

library(ggtree)
library(phangorn)
library(ggplot2)

# input files
tree_file <- snakemake@input[[1]]
out_file <- snakemake@output[[1]]

# read tree file
tree <- read.tree(tree_file)

# midpoint root tree if there is more than 1 node
if (tree$Nnode > 1){
  tree <- midpoint(tree)
}

# draw tree
p <- ggtree(tree, branch.length = 1) +
  geom_tiplab(offset = .05) +
  geom_tippoint(size=3) +
  xlim(0, 4) +
  geom_treescale(width = 0.1, offset=0.25)

# save output file
ggsave(out_file, plot = p, height = 1*length(tree$tip.label))

shell:
	"bwa index {input}"

shell:
	"bwa mem -t {threads} {params.ref} {input.r} | "\
	"samtools view -b -q {params.qual} | "\
	"bamtools filter -tag 'NM:<={params.nm}' | "\
	"samtools sort --threads {threads} -o {output}"

shell:
	"bedtools genomecov -ibam {input} > {output}"

shell:
	"fasttree -nt {input} > {output}"