# wdecoster
"""
Script for filtering and trimming of Oxford Nanopore technologies long reads.
Filtering can be done by calculating metrics while streaming,
or alternatively using a summary file as generated by albacore while basecalling.

Filtering can be done on length and average read basecall quality.
Trimming can be done from the beginning and the end of a read.

Reads from stdin, writes to stdout.

Intended to be used:
- directly after fastq extraction
- prior to mapping
- in a stream between extraction and mapping

Example usage:
gunzip -c reads.fastq.gz | \
 NanoFilt.py -q 10 -l 500 --headcrop 50 | \
 minimap2 genome.fa - | \
 samtools sort -@24 -o alignment.bam -
"""

from __future__ import print_function
from Bio import SeqIO
import sys
import pandas as pd
import nanofilt.utils as utils
import logging
from math import log


def main():
    args = utils.get_args()
    utils.start_logging(args.logfile)
    try:
        if args.tailcrop:
            args.tailcrop = -args.tailcrop
        if args.summary:
            filter_using_summary(args.input, args)
        else:
            filter_stream(args.input, args)
        logging.info('NanoFilt finished.')
    except Exception as e:
        logging.error(e, exc_info=True)
        raise


def silent_quality_check(x):
    """When no quality check needs to be performed, simply return True"""
    return True


def filter_stream(fq, args):
    """Filter a fastq file on stdin.

    Print fastq record to stdout if it passes
    - quality filter (optional)
    - length filter (optional)
    - min/maxGC filter (optional)
    Optionally trim a number of nucleotides from beginning and end.
    Record has to be longer than args.length (default 1) after trimming
    Use a faster silent quality_check if no filtering on quality is required
    """
    if args.quality:
        quality_check = ave_qual
    else:
        quality_check = silent_quality_check
    minlen = args.length + int(args.headcrop or 0) - (int(args.tailcrop or 0))
    for rec in SeqIO.parse(fq, "fastq"):
        if args.GC_filter:
            gc = (rec.seq.upper().count("C") + rec.seq.upper().count("G")) / len(rec)
        else:
            gc = 0.50  # dummy variable
        if quality_check(rec.letter_annotations["phred_quality"]) >= args.quality \
                and minlen <= len(rec) <= args.maxlength \
                and args.minGC <= gc <= args.maxGC:
            print(rec[args.headcrop:args.tailcrop].format("fastq"), end="")


def filter_using_summary(fq, args):
    """Use quality scores from albacore summary file for filtering

    Use the summary file from albacore for more accurate quality estimate
    """
    if args.readtype == "1D":
        cols = ["read_id", "mean_qscore_template"]
    elif args.readtype in ["2D", "1D2"]:
        cols = ["read_id", "mean_qscore_2d"]
    data = pd.read_csv(
        filepath_or_buffer=args.summary,
        sep="\t",
        usecols=cols,
    ).rename(mapper={"mean_qscore_template": "quals", "mean_qscore_2d": "quals"}, axis="columns") \
        .set_index("read_id") \
        .to_dict()["quals"]
    for rec in SeqIO.parse(fq, "fastq"):
        try:
            if data[rec.id] >= args.quality and args.length <= len(rec) <= args.maxlength:
                print(rec[args.headcrop:args.tailcrop].format("fastq"), end="")
        except KeyError:
            logging.warning("mismatch between summary and fastq: \
                   {} was not found in the summary file. \
                   Falling back to calculating.".format(rec.id))
            if ave_qual(rec.letter_annotations["phred_quality"]) >= args.quality \
                    and args.length <= len(rec) <= args.maxlength:
                print(rec[args.headcrop:args.tailcrop].format("fastq"), end="")


def errs_tab(n):
    """Generate list of error rates for qualities less than equal than n."""
    return [10**(q / -10) for q in range(n + 1)]


def ave_qual(quals, qround=False, tab=errs_tab(128)):
    """Calculate average basecall quality of a read.

    Receive the integer quality scores of a read and return the average quality for that read
    First convert Phred scores to probabilities,
    calculate average error probability
    convert average back to Phred scale
    """
    if quals:
        mq = -10 * log(sum([tab[q] for q in quals]) / len(quals), 10)
        if qround:
            return round(mq)
        else:
            return mq
    else:
        sys.stderr.write("Warning: Encountered a zero-length read!")
        return 0


if __name__ == "__main__":
    main()
