# wdecoster
"""
This module provides a few simple math and statistics functions
for other scripts processing Oxford Nanopore sequencing data




# FUNCTIONS
* Calculate read N50 from a set of lengths
get_N50(readlenghts)
* Remove extreme length outliers from a dataset
remove_length_outliers(dataframe, columname)
* Calculate the average Phred quality of a read
ave_qual(qualscores)
* Write out the statistics report after calling readstats function
write_stats(dataframe, outputname)
* Compute a number of statistics, return a dictionary
calc_read_stats(dataframe)
"""

import numpy as np
from math import log
import sys


class Stats(object):
    def __init__(self, df):
        self.number_of_reads = len(df)
        self.number_of_bases = np.sum(df["lengths"])
        if "aligned_lengths" in df:
            self.number_of_bases_aligned = np.sum(df["aligned_lengths"])
        self.median_read_length = np.median(df["lengths"])
        self.mean_read_length = np.mean(df["lengths"])
        self.n50 = get_N50(np.sort(df["lengths"]))
        if "percentIdentity" in df:
            self.average_identity = np.mean(df["percentIdentity"])
            self.median_identity = np.median(df["percentIdentity"])
        if "channelIDs" in df:
            self.active_channels = np.unique(df["channelIDs"]).size
        if "quals" in df:
            self.qualgroups = [5, 7, 10, 12, 15]  # needs 5 elements in current implementation
            self.mean_qual = np.mean(df["quals"])
            self.median_qual = np.median(df["quals"])
            self.top5_lengths = get_top_5(df=df,
                                          col="lengths",
                                          values=["lengths", "quals"])
            self.top5_quals = get_top_5(df=df,
                                        col="quals",
                                        values=["quals", "lengths"])
            self.reads_above_qual = [reads_above_qual(df, q) for q in self.qualgroups]


def get_N50(readlengths):
    """Calculate read length N50.

    Based on https://github.com/PapenfussLab/Mungo/blob/master/bin/fasta_stats.py
    """
    return readlengths[np.where(np.cumsum(readlengths) >= 0.5 * np.sum(readlengths))[0][0]]


def remove_length_outliers(df, columnname):
    """Remove records with length-outliers above 3 standard deviations from the median."""
    return df[df[columnname] < (np.median(df[columnname]) + 3 * np.std(df[columnname]))]


def phred_to_percent(phred):
    return 100 * (1 - 10 ** (phred / -10))


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:
        return None


def median_qual(quals):
    """Receive the integer quality scores of a read and return the median quality for that read."""
    return np.median(quals)


def get_top_5(df, col, values):
    if "readIDs" in df:
        values.append("readIDs")
    res = df.sort_values(col, ascending=False) \
        .head(5)[values] \
        .reset_index(drop=True) \
        .itertuples(index=False, name=None)
    if "readIDs" in df:
        return [str(round(i, ndigits=1)) + " (" +
                str(round(j, ndigits=1)) + "; " + k + ")" for i, j, k in res]
    else:
        return [str(round(i, ndigits=1)) + " (" +
                str(round(j, ndigits=1)) + ")" for i, j in res]


def reads_above_qual(df, qual):
    numberAboveQ = np.sum(df["quals"] > qual)
    megAboveQ = np.sum(df.loc[df["quals"] > qual, "lengths"]) / 1e6
    return "{} ({}%) {}Mb".format(numberAboveQ,
                                  round(100 * (numberAboveQ / len(df.index)), ndigits=1),
                                  round(megAboveQ, ndigits=1))


def feature_list(stats, feature, index=None, padding=15):
    if index is None:
        return ' '.join(['{:>{},.1f}'.format(s.__dict__[feature], padding) for s in stats])
    else:
        return '\t'.join([str(s.__dict__[feature][index]) if len(s.__dict__[feature]) > index
                          else "NA"
                          for s in stats])


def write_stats(datadfs, outputfile, names=[]):
    """Call calculation functions and write stats file.

    This function takes a list of DataFrames,
    and will create a column for each in the tab separated output.
    """
    if outputfile == 'stdout':
        output = sys.stdout
    else:
        output = open(outputfile, 'wt')

    stats = [Stats(df) for df in datadfs]
    features = {
        "Number of reads": "number_of_reads",
        "Total bases": "number_of_bases",
        "Total bases aligned": "number_of_bases_aligned",
        "Median read length": "median_read_length",
        "Mean read length": "mean_read_length",
        "Read length N50": "n50",
        "Average percent identity": "average_identity",
        "Median percent identity": "median_identity",
        "Active channels": "active_channels",
        "Mean read quality": "mean_qual",
        "Median read quality": "median_qual",
    }
    max_len = max([len(k) for k in features.keys()])
    try:
        max_num = max(max([len(str(s.number_of_bases)) for s in stats]),
                      max([len(str(n)) for n in names])) + 6
    except ValueError:
        max_num = max([len(str(s.number_of_bases)) for s in stats]) + 6
    output.write("{:<{}}{}\n".format('General summary:', max_len,
                                     " ".join(['{:>{}}'.format(n, max_num) for n in names])))
    for f in sorted(features.keys()):
        try:
            output.write("{f:{pad}}{v}\n".format(
                f=f + ':',
                pad=max_len,
                v=feature_list(stats, features[f], padding=max_num)))
        except KeyError:
            pass
    if all(["quals" in df for df in datadfs]):
        long_features = {
            "Top 5 longest reads and their mean basecall quality score":
            ["top5_lengths", range(1, 6)],
            "Top 5 highest mean basecall quality scores and their read lengths":
            ["top5_quals", range(1, 6)],
            "Number, percentage and megabases of reads above quality cutoffs":
            ["reads_above_qual", [">Q" + str(q) for q in stats[0].qualgroups]],
        }
        for lf in sorted(long_features.keys()):
            output.write(lf + "\n")
            for i in range(5):
                output.write("{}:\t{}\n".format(
                    long_features[lf][1][i], feature_list(stats, long_features[lf][0], index=i)))
