#! /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
from os import path
import logging
import nanoget
import nanomath
import numpy as np
from scipy import stats
import nanoplot.utils as utils
from .version import __version__
import nanoplotter


def main():
    '''
    Organization function
    -setups logging
    -gets inputdata
    -calls plotting function
    '''
    args = get_args()
    try:
        utils.make_output_dir(args.outdir)
        utils.init_logs(args)
        args.format = nanoplotter.check_valid_format(args.format)
        settings = dict()
        settings["path"] = path.join(args.outdir, args.prefix)
        sources = [args.fastq, args.bam, args.fastq_rich, args.fastq_minimal, args.summary]
        sourcename = ["fastq", "bam", "fastq_rich", "fastq_minimal", "summary"]
        datadf = nanoget.get_input(
            source=[n for n, s in zip(sourcename, sources) if s][0],
            files=[f for f in sources if f][0],
            threads=args.threads,
            readtype=args.readtype,
            combine="simple",
            barcoded=args.barcoded)
        nanomath.write_stats(datadf, settings["path"] + "NanoStats.txt")
        logging.info("Calculated statistics")
        datadf, settings = filter_data(datadf, args, settings)
        if args.barcoded:
            for barc in list(datadf["barcode"].unique()):
                settings["path"] = path.join(args.outdir, args.prefix + barc + "_")
                dfbarc = datadf[datadf["barcode"] == barc]
                nanomath.write_stats(dfbarc, settings["path"] + "NanoStats.txt")
                make_plots(dfbarc, settings, args)
        else:
            make_plots(datadf, settings, 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("--verbose",
                        help="Write log messages also to terminal.",
                        action="store_true")
    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', 'pauvre'])
    parser.add_argument("--barcoded",
                        help="Use if you want to split the summary file by barcode",
                        action="store_true")
    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("--fastq_minimal",
                        help="Data is in fastq format generated by albacore or MinKNOW \
                             with additional information concerning channel and time. \
                             Minimal data is extracted swiftly without elaborate checks.",
                        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:
        utils.list_colors()
    return args


def filter_data(datadf, args, settings):
    '''
    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_list = list()
    if args.alength and args.bam:
        settings["lengths_pointer"] = "aligned_lengths"
        length_prefix_list.append("Aligned_")
        logging.info("Using aligned read lengths for plotting.")
    else:
        settings["lengths_pointer"] = "lengths"
        logging.info("Using sequenced read lengths for plotting.")
    if args.drop_outliers:
        datadf = nanomath.remove_length_outliers(datadf, settings["lengths_pointer"])
        length_prefix_list.append("OutliersRemoved_")
        logging.info("Removing length outliers for plotting.")
    if args.maxlength:
        datadf = datadf[datadf[settings["lengths_pointer"]] < args.maxlength]
        length_prefix_list.append("MaxLength-" + str(args.maxlength) + '_')
        logging.info("Removing reads longer than {}.".format(str(args.maxlength)))
    if args.loglength:
        datadf["log_" + settings["lengths_pointer"]] = np.log10(datadf[settings["lengths_pointer"]])
        settings["lengths_pointer"] = "log_" + settings["lengths_pointer"]
        length_prefix_list.append("Log_")
        logging.info("Using Log10 scaled read lengths.")
        settings["logBool"] = True
    else:
        settings["logBool"] = False
    if args.downsample:
        newNum = min(args.downsample, len(datadf.index))
        length_prefix_list.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.")
    settings["length_prefix"] = ''.join(length_prefix_list)
    return(datadf, settings)


def make_plots(datadf, settings, args):
    '''
    Call plotting functions from nanoplotter
    settings["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", 'pauvre']}
    nanoplotter.length_plots(
        array=datadf[settings["lengths_pointer"]],
        name="Read length",
        path=settings["path"] + settings["length_prefix"],
        n50=nanomath.get_N50(np.sort(datadf["lengths"])),
        color=color,
        figformat=args.format,
        log=settings["logBool"])
    logging.info("Created length plots")
    if "quals" in datadf:
        nanoplotter.scatter(
            x=datadf[settings["lengths_pointer"]],
            y=datadf["quals"],
            names=['Read lengths', 'Average read quality'],
            path=settings["path"] + settings["length_prefix"] + "LengthvsQualityScatterPlot",
            color=color,
            figformat=args.format,
            plots=plotdict,
            log=settings["logBool"])
        logging.info("Created LengthvsQual plot")
    if "channelIDs" in datadf:
        nanoplotter.spatial_heatmap(
            array=datadf["channelIDs"],
            title="Number of reads generated per channel",
            path=settings["path"] + "ActivityMap_ReadsPerChannel",
            color="Greens",
            figformat=args.format)
        logging.info("Created spatialheatmap for succesfull basecalls.")
    if "start_time" in datadf:
        nanoplotter.time_plots(
            df=datadf,
            path=settings["path"],
            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=settings["path"] + "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=settings["path"] + "MappingQualityvsAverageBaseQuality",
            color=color,
            figformat=args.format,
            plots=plotdict)
        logging.info("Created MapQvsBaseQ plot.")
        nanoplotter.scatter(
            x=datadf[settings["lengths_pointer"]],
            y=datadf["mapQ"],
            names=["Read length", "Read mapping quality"],
            path=settings["path"] + settings["length_prefix"] + "MappingQualityvsReadLength",
            color=color,
            figformat=args.format,
            plots=plotdict,
            log=settings["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=settings["path"] + "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[settings["lengths_pointer"]],
            y=datadf["percentIdentity"],
            names=["Aligned read length", "Percent identity"],
            path=settings["path"] + "PercentIdentityvsAlignedReadLength",
            color=color,
            figformat=args.format,
            plots=plotdict,
            stat=stats.pearsonr,
            log=settings["logBool"],
            minvaly=minPID)
        logging.info("Created Percent ID vs Length plot")


if __name__ == "__main__":
    main()
