shell:
    """
    bedtools intersect -a {input.gene_models} -b {input.barrnap} {input.nor} > {output.tmp}

    # grep will exit with a 1 if it finds no matches
    set +e
    grep {params.gff_feature} {output.tmp} > {output.rdna_models}
    exitcode=$?
    if [ $exitcode -eq 1 ]
    # no matches were found
    then
        exit 0
    else
        exit $exitcode
    fi
    """

run:
    # parse through the gff and extract the fields required for the BED
    # [chromosome, start, end, name]
    with open(input.gff, 'r') as IN, open(output.bed, 'w') as OUT:
        for line in IN:
            if line.startswith("#"):
                continue

            columns = line.strip().split("\t")

            chrom = columns[0]
            feature_type = columns[2]
            start = columns[3]
            end = columns[4]

            # only grab the features we want
            if feature_type == params.feature:
                # get a string like "ID=Glyma.13G020400.1.Wm82.a2.v1;Name=Glyma.13G020400.1;pacid=30501300;longest=1;Parent=Glyma.13G020400.Wm82.a2.v1"
                attrs = columns[8]

                # tokenize attrs
                tokens = {attr.split("=")[0]: attr.split("=")[1] for attr in attrs.split(";")}
                try:
                    name = tokens[params.name]
                except KeyError:
                    raise ValueError("'{name}' missing in '{line}'".format(name=params.name, line=line))

                # write the line in the BED file
                OUT.write("\t".join((chrom, start, end, name)) + "\n")

shell:
    "bedtools getfasta -name -fi {input.reference} -bed {input.models} -s > {output.fasta}"

shell:
    "dedupe.sh minidentity={params.identity} in={input.fasta} out={output.deduped}"

run:
    import re

    # gff/gtf have different ways to specify attributes
    gff_regex = re.compile('(?P<key>[^=]+) *= *(?P<value>[^;]+);? *')
    gtf_regex = re.compile('(?P<key>[^ ]+) *\"(?P<value>[^;\"]+)\"?;? *')

    filetype = None
    names = set()
    with open(input.models, 'r') as IN:
        # read the first line to determine if we are using gff/gtf
        first_line = IN.readline()
        if first_line:
            attrs = first_line.split("\t")[8]
            if gff_regex.match(attrs):
                attr_regex = gff_regex
                name_key = "ID"
            elif gtf_regex.match(attrs):
                attr_regex = gtf_regex
                name_key = "gene_id"
            else:
                raise ValueError("Couldn't filetype as either GFF/GTF")

        # reset to the beginning of the file
        IN.seek(0)

        for line in IN:
            # get a string like "ID=Glyma.13G020400.1.Wm82.a2.v1;Name=Glyma.13G020400.1;pacid=30501300;longest=1;Parent=Glyma.13G020400.Wm82.a2.v1"
            # or 'gene_id "FBgn0267517"; gene_symbol "5.8SrRNA-Psi:CR45857"';
            # see: http://uswest.ensembl.org/info/website/upload/gff.html
            attrs = line.split("\t")[8]

            # tokenize attrs
            #tokens = {attr.split("=")[0]: attr.split("=")[1] for attr in attrs.split(";")}

            for match in attr_regex.finditer(attrs):
                if match.group("key") == name_key:
                    names.add(match.group("value"))
                    break
            else:
                raise ValueError("Couldn't detect gene name in '{line}'".format(line=line))

    # write the names 1/line to output
    with open(output.names, 'w') as OUT:
        for n in names:
            OUT.write(n + "\n")

run:
    rdna = set()
    with open(input.blast, 'r') as IN:
        current_query = ""
        current_subject = ""
        current_weight_id = 0
        current_weight_n = 0
        current_cov = 0
        for line in IN:
            query, subject, start, end, pident, qcovs = line.strip().split("\t")

            # trim the bedtools names at the first ":"
            #query = query.split(":")[0]

            # gather the HSPs to each subject for each query and check them against our thresholds
            if query != current_query or subject != current_subject:
                # check the coverage and weighted perc id of the hits
                if current_query != "" and float(current_cov) >= params.qcovs and current_weight_id / current_weight_n >= params.pident:
                    rdna.add(current_query)

                # reset
                current_query = query
                current_subject = subject
                current_weight_id = 0
                current_weight_n = 0
                current_cov = qcovs

            # add this hit; keeping a running % identity for the HSPs in this query/subject pairing
            current_weight_id += float(pident) * abs(int(end) - int(start))
            current_weight_n += abs(int(end) - int(start))

    # subtract the known names
    barrnap_names = set()
    with open(input.barrnap_names, 'r') as IN:
        for line in IN:
            barrnap_names.add(line.strip())

    with open(output.names, 'w') as OUT:
        for n in rdna.difference(barrnap_names):
            OUT.write(n + "\n")

shell:
    "cat {input.barrnap} {input.homology} > {output.names}"

shell:
    "/opt/bbmap/filterbyname.sh include=t substring=name in={input.features} names={input.names} out={output.fasta}"        

run:
    from collections import defaultdict
    # count non-partial rDNA predicted by barrnap
    barrnap_results = defaultdict(lambda: defaultdict(int))
    with open(input.barrnap_predictions, 'r') as IN:
        for line in IN:
            if not line.startswith("#"):
                # detect if this is partial
                partial = "partial" in line

                # quantify each rDNA type
                if "5S_rRNA" in line:
                    barrnap_results["5S"][partial] += 1
                elif "18S_rRNA" in line:
                    barrnap_results["18S"][partial] += 1
                elif "28S_rRNA" in line:
                    barrnap_results["28S"][partial] += 1
                elif "5_8S_rRNA" in line:
                    barrnap_results["5.8S"][partial] += 1
                # needs to be expanded for prokaryotes 
                else:
                    barrnap_results["other"][partial] += 1

    # count the rDNA regions
    rDNA_regions = 0
    with open(input.rdna_regions, 'r') as IN:
        for line in IN:
            rDNA_regions += 1

    # count the models added via BLAST homology
    homology_models = 0
    with open(input.homology_models, 'r') as IN:
        for line in IN:
            homology_models += 1

    # count the rDNA genes
    rDNA_gene_models = 0
    with open(input.gene_models, 'r') as IN:
        for line in IN:
            rDNA_gene_models += 1

    # output the summary
    with open(output.summary, 'w') as OUT:
        OUT.write("riboCleaner ran barrnap to detect the following counts of rDNA features:\n")

        for elem in ["18S", "5.8S", "28S", "5S", "other"]:
            OUT.write("    {elem}: {complete} complete, {partial} partial\n".format(elem=elem, complete=barrnap_results[elem][False], partial=barrnap_results[elem][True]))

        OUT.write("\n")
        OUT.write("These rDNA features were merged into {} rDNA-rich regions.\n".format(rDNA_regions))
        OUT.write("\n")
        OUT.write("The rDNA regions were found to overlap with {overlap} features with the `{tag}` tag in the GFF file.\n".format(overlap=rDNA_gene_models - homology_models, tag=params.gff_feature))
        OUT.write("Additionally, {homology} models were identified as having significant homology to these features for a total of {total} models flagged as putative rDNA.\n".format(homology=homology_models, total=rDNA_gene_models))

        OUT.write("\n")

        # begin the warning section
        if not rDNA_regions:
            OUT.write("WARNING: There were no rDNA regions detected. This may be because there are no rDNA sequences in your reference genome (detectable by barrnap) or because they do not form regular rDNA repeats. riboCleaner works best when it is able to detect rDNA regions.\n")

        if not rDNA_gene_models:
            OUT.write("WARNING: There were no rDNA gene models detected. This is often good news but means riboCleaner will detect no rDNA for this genome. However, since there are no false gene models originating from rDNA, your counts table resulting from standard transcriptome analysis won't be biased by false rDNA gene models!\n")

shell:
    """
    salmon quant --threads {threads} -i {input.index} --libType A -r {input.r1} -o {output.tmp} 2>&1 | tee {log}
    cp -v {output.tmp}/quant.sf {output.sf}
    cp -v {output.tmp}/aux_info/meta_info.json {output.meta_info}
    """

run:
    import pandas as pd
    import json

    # get the rDNA names
    with open(input.rdna_names, 'r') as IN:
        rdna = [line.strip() for line in IN]

    sample_summary = {}

    # gather the counts for each file
    all_counts = []
    for sf, sample in zip(input.sf, params.samples):
        d = pd.read_csv(sf, sep="\t", header=0, index_col=0)
        counts = d["TPM"]

        # get the total (including rDNA)
        total = counts.sum()

        # remove rDNA
        counts = counts.filter(items=[n for n in counts.index if n not in rdna])

        # rescale counts by the rDNA ratio (to get the TPM as if the rDNA was never there)
        counts = counts * (total / counts.sum())

        counts.name = sample
        all_counts.append(counts)

    # and combine into one df
    all_counts = pd.DataFrame(all_counts).fillna(0).transpose()

    # make the index be just the gene names (not the position)
    all_counts.index = [g.split(":")[0] for g in all_counts.index]

    all_counts.to_csv(output.clean_tpm, sep="\t", index_label="gene")

run:
    # print the workflow docstring
    print("Snakefile: {p}".format(p=workflow.snakefile))
    if workflow.globals["__doc__"]:
        print(workflow.globals["__doc__"])

    # print each (public) rule
    print("\n##### RULES #####\n")
    for rule in workflow.rules:
        if not rule.name.startswith("_"):
            if rule.docstring:
                print(rule.name)
                print(rule.docstring)
                print("----------\n")

shell:
    """
    bbduk.sh in={input.r1} out={output.r1} reads={params.reads} threads={threads}
    """