#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#~~~~~~~~~~~~~~IMPORTS~~~~~~~~~~~~~~#
# Standard library imports
import argparse
import sys
import os
import json
import logging
import datetime
import textwrap

# Third party imports
import plotly.offline as py
from jinja2 import Environment, PackageLoader, Template

# Local imports
from pycoQC.pycoQC import pycoQC
from pycoQC.Fast5_to_seq_summary import Fast5_to_seq_summary
from pycoQC.common import *
from pycoQC import __version__ as package_version
from pycoQC import __name__ as package_name

logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger(__name__)

# List of current valid plotting methods names
PLOT_METHODS = [
    "summary",
    "barcode_summary",
    "run_id_summary",
    "reads_len_1D",
    "reads_qual_1D",
    "reads_len_qual_2D",
    "output_over_time",
    "len_over_time",
    "qual_over_time",
    "barcode_counts",
    "channels_activity"]

#~~~~~~~~~~~~~~Fast5_to_seq_summary CLI ENTRY POINT~~~~~~~~~~~~~~#
def main_Fast5_to_seq_summary (args=None):
    if args is None:
        args = sys.argv[1:]

    # Define parser object
    parser = argparse.ArgumentParser(
        description ="Fast5_to_seq_summary generate a sequencing summary like file from a directory containing Fast5 files")
    parser.add_argument('--version', '-v', action='version', version="{} v{}".format(package_name, package_version))

    # Define arguments
    parser.add_argument("--fast5_dir", "-f", required=True, type=str,
        help="""Directory containing fast5 files. Can contain multiple subdirectories""")
    parser.add_argument("--seq_summary_fn", "-s", required=True, type=str,
        help="""path of the summary sequencing file where to write the data extracted from the fast5 files""")
    parser.add_argument("--max_fast5", type=int, default=0,
        help="Maximum number of file to try to parse. 0 to deactivate (default: %(default)s)")
    parser.add_argument("--threads", "-t", type=int, default=4,
        help="Total number of threads to use. 1 thread is used for the reader and 1 for the writer. Minimum 3 (default: %(default)s)")
    parser.add_argument("--basecall_id", type=int, default=0,
        help="id of the basecalling group. By default leave to 0, but if you perfome multiple basecalling on the same fast5 files, this can be used to indicate the corresponding group (1, 2 ...) (default: %(default)s)")
    parser.add_argument("--fields", type=str, nargs="+", default=["read_id", "run_id", "channel", "start_time", "sequence_length_template", "mean_qscore_template", "calibration_strand_genome_template", "barcode_arrangement"],
        help="list of field names corresponding to attributes to try to fetch from the fast5 files (default: %(default)s)")
    parser.add_argument("--include_path", action='store_true', default=False,
        help="If given, the absolute path to the corresponding file is added in an extra column (default: %(default)s)")
    parser.add_argument("--verbose_level", type=int, default=0,
        help="Level of verbosity, from 2 (Chatty) to 0 (Nothing) (default: %(default)s)")

    # Try to parse arguments
    args = parser.parse_args()

    # Run main function
    Fast5_to_seq_summary (
        fast5_dir = args.fast5_dir,
        seq_summary_fn = args.seq_summary_fn,
        max_fast5 = args.max_fast5,
        threads = args.threads,
        basecall_id = args.basecall_id,
        fields = args.fields,
        include_path = args.include_path,
        verbose_level = args.verbose_level)

#~~~~~~~~~~~~~~pycoQC CLI ENTRY POINT~~~~~~~~~~~~~~#
def main_pycoQC (args=None):
    if args is None:
        args = sys.argv[1:]

    # Define parser object
    parser = argparse.ArgumentParser(
        formatter_class = argparse.RawDescriptionHelpFormatter,
        description = textwrap.dedent("""
            pycoQC computes metrics and generates interactive QC plots from the sequencing summary report generated by Oxford Nanopore technologies basecallers\n
            * Minimal usage
                pycoQC -f sequencing_summary.txt -o pycoQC_output.html
            * Including Guppy barcoding file and json output
                pycoQC -f sequencing_summary.txt -b barcoding_sequencing.txt -o pycoQC_output.html -j pycoQC_output.json"""))

    parser.add_argument('--version', action='version', version="{} v{}".format(package_name, package_version))

    # Define arguments
    parser_io = parser.add_argument_group('Input/output options')
    parser_io.add_argument("--summary_file", "-f", default=None, type=str, nargs='+',
        help=textwrap.dedent("""Path to the sequencing_summary generated by Albacore 1.0.0 + (read_fast5_basecaller.py) / Guppy 2.1.3+ (guppy_basecaller).
        One can also pass multiple space separated file paths or a UNIX style regex matching multiple files (Required)"""))
    parser_io.add_argument("--barcode_file", "-b", default=None, type=str, nargs='+',
        help=textwrap.dedent("""Path to the barcode_summary_file generated by Guppy 2.1.3+ (guppy_barcoder)..
        One can also pass multiple space separated file paths or a UNIX style regex matching multiple files (optional)"""))
    parser_io.add_argument("--html_outfile", "-o", default=None, type=str,
        help="Path to an output html file (required if json_outfile not given)")
    parser_io.add_argument("--json_outfile", "-j", default=None, type=str,
        help="Path to an output json file (required if html_outfile not given)")
    parser_filt = parser.add_argument_group('Filtering options')
    parser_filt.add_argument("--min_pass_qual", "-q", default=7, type=int,
        help="Minimum quality to consider a read as 'pass' (default: %(default)s)")
    parser_filt.add_argument("--filter_calibration", default=False, action='store_true',
        help="If given reads flagged as calibration strand by the basecaller are removed (default: %(default)s)")
    parser_filt.add_argument("--min_barcode_percent", default=0.1, type=float,
        help="Minimal percent of total reads to retain barcode label. If below the barcode value is set as `unclassified` (default: %(default)s)")
    parser_html = parser.add_argument_group('HTML report options')
    parser_html.add_argument("--title", "-t", default=None, type=str,
        help="A title to be used in the html report (default: %(default)s)")
    parser_html.add_argument("--template_file", type=str, default=None,
        help="Jinja2 html template for the html report (default: %(default)s)")
    parser_html.add_argument("--config", "-c", type=str, default=None,
        help=textwrap.dedent("""Path to a JSON configuration file for the html report.
        If not provided, looks for it in ~/.pycoQC and ~/.config/pycoQC/config. If it's still not found, falls back to default parameters.
        The first level keys are the names of the plots to be included.
        The second level keys are the parameters to pass to each plotting function (default: %(default)s)")"""))
    parser_other = parser.add_argument_group('Other options')
    parser_other.add_argument("--default_config", "-d", action='store_true',
        help="Print default configuration file. Can be used to generate a template JSON file (default: %(default)s)")
    parser_other.add_argument("--list_plots", "-l", default=None, action='store_true',
        help="Print the list of available plotting functions and exit (default: %(default)s)")
    parser_other.add_argument("--verbose_level", "-v", choices=[2,1,0], type=int, default=1,
        help="Level of verbosity, from 2 (Chatty) to 0 (Nothing) (default: %(default)s)")

    # Try to parse arguments
    args = parser.parse_args()

    # Set logging level
    logLevel_dict = {2:logging.DEBUG, 1:logging.INFO, 0:logging.WARNING}
    logger.setLevel (logLevel_dict.get (args.verbose_level, logging.INFO))

    # Print the default config parameters and exit
    if args.default_config:
        json.dump (_default_config(), sys.stdout, indent=2)

    # Print the names of valid pycoQC plotting functions
    elif args.list_plots:
        logger.warning ("Available pycoQC methods")
        for i in PLOT_METHODS:
            logger.warning ("* {}".format(i))

    elif not args.summary_file:
        logger.warning ("ERROR: `--summary_file` is a required argument")
        parser.print_help()

    elif not args.html_outfile and not args.json_outfile:
        logger.warning ("ERROR: At least one output file required `--html_outfile` or `--json_outfile`")
        parser.print_help()

    # Run pycoQC
    else:
        # Initiate pycoQC
        logger.warning("PARSING DATA FILES")
        pycoqc_res = pycoQC (
            seq_summary_file=args.summary_file,
            barcode_summary_file=args.barcode_file,
            verbose_level=args.verbose_level,
            min_pass_qual=args.min_pass_qual,
            filter_calibration=args.filter_calibration,
            min_barcode_percent=args.min_barcode_percent)

        # Run html output function
        if args.html_outfile:
            logger.warning("GENERATING HTML REPORT")
            generate_html_report(
                pycoqc_res=pycoqc_res,
                outfile=args.html_outfile,
                config = args.config,
                template_file = args.template_file,
                title=args.title)

        # Run json output function
        if args.json_outfile:
            logger.warning("GENERATING JSON REPORT")
            generate_json_report(
                pycoqc_res=pycoqc_res,
                outfile=args.json_outfile)

#~~~~~~~~~~~~~~PUBLIC FUNCTION~~~~~~~~~~~~~~#
def generate_html_report(
    pycoqc_res,
    outfile,
    config=None,
    template_file=None,
    title=None):
    """"""
    # Parse configuration file
    logger.info("\tParsing html config file")
    config_dict = _parse_config_file(config)
    logger.debug(config_dict)

    # Loop over configuration file and run the pycoQC functions defined
    plots = list()
    titles = list()
    for method_name, method_args in config_dict.items ():

        # Check if method exists and is callable
        if not method_name in PLOT_METHODS:
            logger.info("\tMethod `{}` is not defined in pycoQC".format(method_name))

        try:
            logger.info("\tRunning method {}".format(method_name))
            logger.debug ("\t{} ({})".format(method_name, method_args))

            # Store plot title for HTML title and remove from data passed to plotly
            plot_title = method_args["plot_title"]
            method_args["plot_title"]=""

            # Get method and generate plot
            method = getattr(pycoqc_res, method_name)
            fig = method(**method_args)
            plot = py.plot(
                fig,
                output_type='div',
                include_plotlyjs=False,
                image_width='',
                image_height='',
                show_link=False,
                auto_open=False)

            plots.append(plot)
            titles.append(plot_title)

        except pycoQCError as E:
            logger.info("\t\t{}".format(E))

    # Load HTML template for Jinja
    logger.info("\tLoading HTML template")
    template = _get_jinja_template(template_file)

    # Set a title for the HTML report
    report_title=""
    if title:
        report_title+=title+"<br>"
    report_title+="Generate on {} with {} {}".format(
        datetime.datetime.now().strftime("%d/%m/%y"),
        package_name,
        package_version)

    # Render plots
    logger.info("\tRendering plots in d3js")
    rendering = template.render(
        plots=plots,
        titles=titles,
        plotlyjs=py.get_plotlyjs(),
        report_title=report_title)

    # Write to HTML file
    logger.info("\tWriting to HTML file")
    with open(outfile, "w") as f:
        f.write(rendering)

def generate_json_report(pycoqc_res, outfile):
    """"""
    logger.info("\tRunning summary_stats_dict method")
    res_dict = pycoqc_res.summary_stats_dict (barcode_split=True, run_id_split=True)

    logger.info("\tWriting to JSON file")
    with open (outfile, "w") as fp:
        json.dump(res_dict, fp, indent=2)

#~~~~~~~~~~~~~~PRIVATE FUNCTION~~~~~~~~~~~~~~#

def _parse_config_file(config_file):
    """"""
    # First, try to read provided configuration file if given
    if config_file:
        logger.debug ("\tTry to read provided config file")
        try:
            with open(config_file, 'r') as cf:
                return json.load(cf)
        except (FileNotFoundError, IOError,json.JSONDecodeError):
            logger.debug ("\t\tFile not found, non-readable or invalid")

    # Second, look into default system folder
    home=os.path.expanduser('~')
    for config_file in [home+"/.pycoQC", home+"/.config/pycoQC/config"]:
        logger.debug ("\tTry to read config file from system directory:{}".format(config_file))
        try:
            with open(config_file, 'r') as cf:
                return json.load(cf)
        except (FileNotFoundError, IOError,json.JSONDecodeError):
            logger.debug ("\t\tFile not found, non-readable or invalid")

    # Last use the default harcoded config_dict
    logger.debug ("\tFall back to default configuration")
    return _default_config()

def _get_jinja_template(template_file=None):
    """"""
    # First, try to read provided configuration file if given
    if template_file:
        logger.debug("\tTry to load provided template file")
        try:
            with open(template_file) as fp:
                template = Template(fp.read())
                return template
        except (FileNotFoundError, IOError):
            logger.debug ("\t\tFile not found, non-readable or invalid")

    # Last use the default harcoded config_dict
    logger.debug ("\tFall back to default template")
    env = Environment (loader=PackageLoader('pycoQC', 'templates'), autoescape=False)
    template = env.get_template('spectre.html.j2')
    return template

def _default_config():
    """"""
    config_dict = {
        'summary': {
            "plot_title":"Run summary"},
        'barcode_summary': {
            "plot_title":"Run summary by barcode"},
        'run_id_summary': {
            "plot_title":"Run summary by Run ID"},
        'reads_len_1D': {
            "plot_title":'Distribution of read length',
            "color":'lightsteelblue',
            "nbins":200,
            "smooth_sigma":2,
            "sample":100000},
        'reads_qual_1D': {
            "plot_title":'Distribution of read quality',
            "color":'salmon',
            "nbins":200,
            "smooth_sigma":2,
            "sample":100000},
        'reads_len_qual_2D': {
            "plot_title":'Mean read quality per sequence length',
            "colorscale":[[0.0, 'rgba(255,255,255,0)'], [0.1, 'rgba(255,150,0,0)'], [0.25, 'rgb(255,100,0)'], [0.5, 'rgb(200,0,0)'], [0.75, 'rgb(120,0,0)'], [1.0, 'rgb(70,0,0)']],
            "len_nbins":200,
            "qual_nbins":75,
            "smooth_sigma":2,
            "sample":100000},
        'output_over_time': {
            "plot_title":'Output over experiment time',
            "cumulative_color":"rgb(204,226,255)",
            "interval_color":"rgb(102,168,255)",
            "sample":100000},
        'len_over_time': {
            "plot_title":'Read length over time',
            "median_color":"rgb(102,168,255)",
            "quartile_color":"rgb(153,197,255)",
            "extreme_color":"rgba(153,197,255,0.5)",
            "smooth_sigma":1,
            "sample":100000},
        'qual_over_time': {
            "plot_title":'Mean read quality over time',
            "median_color":"rgb(250,128,114)",
            "quartile_color":"rgb(250,170,160)",
            "extreme_color":"rgba(250,170,160,0.5)",
            "smooth_sigma":1,
            "sample":100000},
        'barcode_counts': {
            "plot_title":'Number of reads per barcode',
            "colors":["#f8bc9c", "#f6e9a1", "#f5f8f2", "#92d9f5", "#4f97ba"]},
        'channels_activity': {
            "plot_title":'Channel activity over time',
            "colorscale":[[0.0,'rgba(255,255,255,0)'], [0.01,'rgb(255,255,200)'], [0.25,'rgb(255,200,0)'], [0.5,'rgb(200,0,0)'], [0.75,'rgb(120,0,0)'], [1.0,'rgb(0,0,0)']],
            "smooth_sigma":1,
            "sample":100000},
        }
    return(config_dict)

if __name__ == "__main__":
    # execute only if run as a script
    main_pycoQC()
