#! /usr/bin/env python
# wdecoster
'''
The main purpose of this script is to create plots for nanopore data.
Input data can be given as one or multiple of:
-compressed, standard or streamed fastq file
-compressed, standard or streamed fastq file,
 with additional information added by albacore or MinKNOW
-a bam file
-a summary file generated by albacore
'''


from __future__ import division, print_function
from argparse import ArgumentParser
import sys
import os
from time import time
import logging
from datetime import datetime as dt
import pysam
import nanoget
import nanomath
import pandas as pd
import numpy as np
from scipy import stats
from .version import __version__
import matplotlib
matplotlib.use('Agg')
import seaborn as sns
import nanoplotter


def main():
    '''
    Organization function
    -setups logging
    -gets inputdata
    -calls plotting function
    '''
    args = get_args()
    try:
        make_output_dir(args.outdir)
        init_logs(args)
        args.format = nanoplotter.check_valid_format(args.format)
        datadf, length_prefix, logBool, read_lengths_pointer = get_input(args)
        make_plots(datadf, length_prefix, logBool, read_lengths_pointer, args)
        logging.info("Succesfully processed all input.")
    except Exception as e:
        logging.error(e, exc_info=True)
        raise


def get_args():
    parser = ArgumentParser(description="Perform diagnostic plotting and QC analysis for \
                                                  Oxford Nanopore sequencing data and alignments.")
    parser.add_argument("-v", "--version",
                        help="Print version and exit.",
                        action="version",
                        version='NanoPlot {}'.format(__version__))
    parser.add_argument("-t", "--threads",
                        help="Set the allowed number of threads to be used by the script",
                        default=4,
                        type=int)
    parser.add_argument("--maxlength",
                        help="Drop reads longer than length specified.",
                        type=int)
    parser.add_argument("--drop_outliers",
                        help="Drop outlier reads with extreme long length.",
                        action="store_true")
    parser.add_argument("--downsample",
                        help="Reduce dataset to N reads by random sampling.",
                        type=int)
    parser.add_argument("--loglength",
                        help="Logarithmic scaling of lengths in plots.",
                        action="store_true")
    parser.add_argument("--readtype",
                        help="Which read type to extract information about from summary. \
                             Options are 1D, 2D, 1D2",
                        default="1D",
                        choices=['1D', '2D', '1D2'])
    parser.add_argument("--alength",
                        help="Use aligned read lengths rather than sequenced length (bam mode)",
                        action="store_true")
    parser.add_argument("-c", "--color",
                        help="Specify a color for the plots, must be a valid matplotlib color",
                        default="#4CB391")
    parser.add_argument("-o", "--outdir",
                        help="Specify directory in which output has to be created.",
                        default=".")
    parser.add_argument("-p", "--prefix",
                        help="Specify an optional prefix to be used for the output files.",
                        default="",
                        type=str)
    parser.add_argument("-f", "--format",
                        help="Specify the output format of the plots.",
                        default="png",
                        type=str,
                        choices=['eps', 'jpeg', 'jpg', 'pdf', 'pgf', 'png', 'ps',
                                 'raw', 'rgba', 'svg', 'svgz', 'tif', 'tiff'])
    parser.add_argument("--plots",
                        help="Specify which bivariate plots have to be made.",
                        default=['kde', 'hex', 'dot'],
                        type=str,
                        nargs='*',
                        choices=['kde', 'hex', 'dot'])
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--fastq",
                        help="Data is in default fastq format.",
                        nargs='*')
    target.add_argument("--fastq_rich",
                        help="Data is in fastq format generated by albacore or MinKNOW \
                             with additional information concerning channel and time.",
                        nargs='*')
    target.add_argument("--summary",
                        help="Data is a summary file generated by albacore.",
                        nargs='*')
    target.add_argument("--bam",
                        help="Data as a sorted bam file.",
                        nargs='*')
    target.add_argument("--listcolors",
                        help="List the colors which are available for plotting",
                        action="store_true")
    args = parser.parse_args()
    if args.listcolors:
        list_colors()
    return args


def list_colors():
    parent_directory = os.path.dirname(os.path.abspath(os.path.dirname(__file__)))
    colours = open(os.path.join(parent_directory, "extra/color_options.txt")).readlines()
    print("{}".format(", ".join([c.strip() for c in colours])))
    sys.exit(0)


def make_output_dir(path):
    try:
        if not os.path.exists(path):
            os.makedirs(path)
    except IOError:
        sys.exit("ERROR: No writing permission to the output directory.")


def init_logs(args):
    '''
    Initiate log file
    Log arguments and module versions
    '''
    start_time = dt.fromtimestamp(time()).strftime('%Y%m%d_%H%M')
    logging.basicConfig(
        format='%(asctime)s %(message)s',
        filename=os.path.join(args.outdir, args.prefix + "Nanoplot_" + start_time + ".log"),
        level=logging.INFO)
    logging.info('Nanoplot {} started with arguments {}'.format(__version__, args))
    logging.info('Python version is: {}'.format(sys.version.replace('\n', ' ')))
    logging.info('Versions of key modules are:')
    for module in [np, sns, pd, pysam, matplotlib, nanoget, nanoplotter, nanomath]:
        logging.info('{}: {}'.format(module, module.__version__))


def get_input(args):
    '''
    Get input and process accordingly.
    Data can be:
    - a uncompressed, bgzip, bzip2 or gzip compressed fastq file
    - a sorted bam file
    - a sequencing_summary.txt file generated by albacore
    Handle is passed to the proper functions to get DataFrame with metrics
    Multiple files of the same type can be used to extract info from
    The resulting output DataFrames are concatenated
    '''
    if args.fastq:
        datadf = pd.concat(
            [nanoget.process_fastq_plain(inp, args.threads) for inp in args.fastq],
            ignore_index=True)
    elif args.bam:
        datadf = pd.concat(
            [nanoget.process_bam(inp, args.threads) for inp in args.bam],
            ignore_index=True)
    elif args.fastq_rich:
        datadf = pd.concat(
            [nanoget.process_fastq_rich(inp) for inp in args.fastq_rich],
            ignore_index=True)
    elif args.summary:
        datadf = pd.concat(
            [nanoget.process_summary(inp, args.readtype) for inp in args.summary],
            ignore_index=True)
    logging.info("Gathered metrics for plotting")
    nanomath.write_stats(datadf, os.path.join(args.outdir, args.prefix + "NanoStats.txt"))
    logging.info("Calculated statistics")
    datadf, length_prefix, logBool, read_lengths_pointer = filter_data(datadf, args)
    return (datadf, length_prefix, logBool, read_lengths_pointer)


def filter_data(datadf, args):
    '''
    Perform filtering on the data based on arguments set on commandline
    - use aligned length or sequenced length (bam mode only)
    - drop outliers
    - drop reads longer than args.maxlength
    - use log10 scaled reads
    - downsample reads to args.downsample
    Return an accurate prefix which is added to plotnames using this filtered data
    '''
    length_prefix = []
    if args.alength and args.bam:
        read_lengths_pointer = "aligned_lengths"
        length_prefix.append("Aligned_")
        logging.info("Using aligned read lengths for plotting.")
    else:
        read_lengths_pointer = "lengths"
        logging.info("Using sequenced read lengths for plotting.")
    if args.drop_outliers:
        datadf = nanomath.remove_length_outliers(datadf, read_lengths_pointer)
        length_prefix.append("OutliersRemoved_")
        logging.info("Removing length outliers for plotting.")
    if args.maxlength:
        datadf = datadf[datadf[read_lengths_pointer] < args.maxlength]
        length_prefix.append("MaxLength-" + str(args.maxlength) + '_')
        logging.info("Removing reads longer than {}.".format(str(args.maxlength)))
    if args.loglength:
        datadf["log_" + read_lengths_pointer] = np.log10(datadf[read_lengths_pointer])
        read_lengths_pointer = "log_" + read_lengths_pointer
        length_prefix.append("Log_")
        logging.info("Using Log10 scaled read lengths.")
        logBool = True
    else:
        logBool = False
    if args.downsample:
        newNum = min(args.downsample, len(datadf.index))
        length_prefix.append("Downsampled_")
        logging.info("Downsampling the dataset from {} to {} reads".format(
            len(datadf.index), newNum))
        datadf = datadf.sample(newNum)
    logging.info("Processed the metrics, optionally performed filtering.")
    return(datadf, ''.join(length_prefix), logBool, read_lengths_pointer)


def make_plots(datadf, length_prefix, logBool, read_lengths_pointer, args):
    '''
    Call plotting functions from nanoplotter
    read_lengths_pointer is a column in the DataFrame specifying which lengths to use
    '''
    color = nanoplotter.check_valid_color(args.color)
    plotdict = {type: args.plots.count(type) for type in ["kde", "hex", "dot"]}
    nanoplotter.length_plots(
        array=datadf[read_lengths_pointer],
        name="Read length",
        path=os.path.join(args.outdir, args.prefix + length_prefix),
        n50=nanomath.get_N50(np.sort(datadf["lengths"])),
        color=color,
        figformat=args.format,
        log=logBool)
    logging.info("Created length plots")
    nanoplotter.scatter(
        x=datadf[read_lengths_pointer],
        y=datadf["quals"],
        names=['Read lengths', 'Average read quality'],
        path=os.path.join(args.outdir, args.prefix + length_prefix + "LengthvsQualityScatterPlot"),
        color=color,
        figformat=args.format,
        plots=plotdict,
        log=logBool)
    logging.info("Created LengthvsQual plot")
    if args.fastq_rich or args.summary:
        nanoplotter.spatial_heatmap(
            array=datadf["channelIDs"],
            title="Number of reads generated per channel",
            path=os.path.join(args.outdir, args.prefix + "ActivityMap_ReadsPerChannel"),
            color="Greens",
            figformat=args.format)
        logging.info("Created spatialheatmap for succesfull basecalls.")
        nanoplotter.time_plots(
            df=datadf,
            path=os.path.join(args.outdir, args.prefix),
            color=color,
            figformat=args.format)
        logging.info("Created timeplots.")
    if args.bam:
        nanoplotter.scatter(
            x=datadf["aligned_lengths"],
            y=datadf["lengths"],
            names=["Aligned read lengths", "Sequenced read length"],
            path=os.path.join(args.outdir, args.prefix + "AlignedReadlengthvsSequencedReadLength"),
            figformat=args.format,
            plots=plotdict,
            color=color)
        logging.info("Created AlignedLength vs Length plot.")
        nanoplotter.scatter(
            x=datadf["mapQ"],
            y=datadf["quals"],
            names=["Read mapping quality", "Average basecall quality"],
            path=os.path.join(args.outdir, args.prefix + "MappingQualityvsAverageBaseQuality"),
            color=color,
            figformat=args.format,
            plots=plotdict)
        logging.info("Created MapQvsBaseQ plot.")
        nanoplotter.scatter(
            x=datadf[read_lengths_pointer],
            y=datadf["mapQ"],
            names=["Read length", "Read mapping quality"],
            path=os.path.join(args.outdir, args.prefix + length_prefix +
                              "MappingQualityvsReadLength"),
            color=color,
            figformat=args.format,
            plots=plotdict,
            log=logBool)
        logging.info("Created Mapping quality vs read length plot.")
        minPID = np.amin(datadf["percentIdentity"])
        nanoplotter.scatter(
            x=datadf["percentIdentity"],
            y=datadf["aligned_quals"],
            names=["Percent identity", "Read quality"],
            path=os.path.join(args.outdir, args.prefix + "PercentIdentityvsAverageBaseQuality"),
            color=color,
            figformat=args.format,
            plots=plotdict,
            stat=stats.pearsonr,
            minvalx=minPID)
        logging.info("Created Percent ID vs Base quality plot.")
        nanoplotter.scatter(
            x=datadf[read_lengths_pointer],
            y=datadf["percentIdentity"],
            names=["Aligned read length", "Percent identity"],
            path=os.path.join(args.outdir, args.prefix + "PercentIdentityvsAlignedReadLength"),
            color=color,
            figformat=args.format,
            plots=plotdict,
            stat=stats.pearsonr,
            log=logBool,
            minvaly=minPID)
        logging.info("Created Percent ID vs Length plot")


if __name__ == "__main__":
    main()
