#!/usr/bin/env python3

"""
Tool for detecting DNA base-modifications from kinetic signatures.
"""

import cProfile
import functools
import gc
import itertools
import argparse
import json

import os
import logging
import sys
import multiprocessing
import time
import threading
import numpy as np
import queue
import traceback
from pkg_resources import Requirement, resource_filename

from pbcommand.common_options import add_debug_option
from pbcommand.cli import get_default_argparser_with_base_opts, pacbio_args_runner
from pbcommand.utils import setup_log
from pbcore.io import AlignmentSet

from kineticsTools.KineticWorker import KineticWorkerProcess
from kineticsTools.ResultWriter import KineticsWriter
from kineticsTools.ipdModel import IpdModel
from kineticsTools import ReferenceUtils, loader

__version__ = "3.0"

log = logging.getLogger(__name__)


class Constants(object):
    PVALUE_DEFAULT = 0.01
    MAX_LENGTH_DEFAULT = int(3e12)


def _getResourcePathSpec():
    default_dir = resource_filename(Requirement.parse(
        'kineticsTools'), 'kineticsTools/resources')
    return loader.getResourcePathSpec(default_dir)


def _validateResource(func, p):
    """Basic func for validating files, dirs, etc..."""
    if func(p):
        return os.path.abspath(p)
    else:
        raise IOError("Unable to find {p}".format(p=p))


def _validateNoneOrResource(func, p):
    """
    Handle optional values. If a file or dir is explicitly provided, then
    it will validated.
    """
    if p is None:
        return p
    else:
        return _validateResource(func, p)


def validateNoneOrPathSpec(ps):
    """
    Handle optional values. If a pathspec is explicitly provided, then
    it will be validated.
    """
    if ps is None:
        return ps
    pths = []
    for p in ps.split(':'):
        pths.append(_validateResource(os.path.isdir, p))
    if not pths:
        raise ValueError("Empty pathspec!")
    return pths


validateFile = functools.partial(_validateResource, os.path.isfile)
validateDir = functools.partial(_validateResource, os.path.isdir)

validateNoneOrFile = functools.partial(_validateNoneOrResource, os.path.isfile)
validateNoneOrDir = functools.partial(_validateNoneOrResource, os.path.isdir)


def get_parser():
    p = get_default_argparser_with_base_opts(
        version=__version__,
        description=__doc__,
        default_level="INFO")
    p.add_argument("alignment_set", help="BAM or Alignment DataSet")
    p.add_argument("--reference", action="store",
                   required=True,
                   type=validateFile, help="Fasta or Reference DataSet")
    p.add_argument("--gff", action="store", default=None,
                   help="Output GFF file of modified bases")
    p.add_argument("--csv", action="store", default=None,
                   help="Output CSV file out per-nucleotide information")
    p.add_argument("--bigwig", action="store", default=None,
                   help="Output BigWig file encoding IpdRatio for both strands")
    # FIXME use central --nproc option
    p.add_argument('--numWorkers', '-j',
                   dest='numWorkers',
                   default=1,
                   type=int,
                   help='Number of thread to use (-1 uses all logical cpus)')
    # common options
    p.add_argument("--pvalue",
                   type=float,
                   default=Constants.PVALUE_DEFAULT,
                   help="P-value cutoff")
    p.add_argument("--maxLength",
                   type=int,
                   default=Constants.MAX_LENGTH_DEFAULT,
                   help="Maximum number of bases to process per contig")
    p.add_argument(
        "--identify",
        action="store",
        default="m6A,m4C",
        help="Specific modifications to identify (comma-separated " +
        "list).  Currrent options are m6A, m4C, m5C_TET.  Using --control " +
        "overrides this option.")
    _DESC = "In the --identify mode, add --methylFraction to " +\
            "command line to estimate the methylated fraction, along with " +\
            "95%% confidence interval bounds."
    p.add_argument("--methylFraction", action="store_true",
                   help=_DESC)
    p.add_argument('--outfile',
                   dest='outfile',
                   default=None,
                   help='Use this option to generate all possible output files. Argument here is the root filename of the output files.')

    # FIXME: Need to add an extra check for this; it can only be used if
    # --useLDA flag is set.
    p.add_argument('--m5Cgff',
                   dest='m5Cgff',
                   default=None,
                   help='Name of output GFF file containing m5C scores')

    # FIXME: Make sure that this is specified if --useLDA flag is set.
    p.add_argument('--m5Cclassifier',
                   dest='m5Cclassifier',
                   default=None,
                   help='Specify csv file containing a 127 x 2 matrix')

    p.add_argument('--pickle',
                   dest='pickle',
                   default=None,
                   help='Name of output pickle file.')

    p.add_argument('--ms_csv',
                   dest='ms_csv',
                   default=None,
                   help='Multisite detection CSV file.')

    # Calculation options:
    p.add_argument('--control',
                   dest='control',
                   default=None,
                   type=validateNoneOrFile,
                   help='AlignmentSet or mapped BAM file containing a control sample. Tool will perform a case-control analysis')

    # Temporary addition to test LDA for Ca5C detection:
    p.add_argument('--useLDA',
                   action="store_true",
                   dest='useLDA',
                   default=False,
                   help='Set this flag to debug LDA for m5C/Ca5C detection')

    # Parameter options:
    defaultParamsPathSpec = _getResourcePathSpec()
    p.add_argument('--paramsPath',
                   dest='paramsPath',
                   default=defaultParamsPathSpec,
                   type=validateNoneOrPathSpec,
                   help='List of :-delimited directory paths containing in-silico trained models (default is "%s")' % defaultParamsPathSpec)

    # XXX hacky workaround for running tests using obsolete chemistry inputs
    p.add_argument("--useChemistry",
                   dest="useChemistry",
                   default=None,
                   help=argparse.SUPPRESS)

    p.add_argument('--minCoverage',
                   dest='minCoverage',
                   default=3,
                   type=int,
                   help='Minimum coverage required to call a modified base')

    p.add_argument('--maxQueueSize',
                   dest='maxQueueSize',
                   default=20,
                   type=int,
                   help='Max Queue Size')

    p.add_argument('--maxCoverage',
                   dest='maxCoverage',
                   type=int, default=-1,
                   help='Maximum coverage to use at each site')

    p.add_argument('--mapQvThreshold',
                   dest='mapQvThreshold',
                   type=float,
                   default=-1.0)

    p.add_argument('--ipdModel',
                   dest='ipdModel',
                   default=None,
                   type=validateNoneOrFile,
                   help='Alternate synthetic IPD model HDF5 file')

    p.add_argument('--modelIters',
                   dest='modelIters',
                   type=int,
                   default=-1,
                   help='[Internal] Number of GBM model iteration to use')

    p.add_argument('--cap_percentile',
                   dest='cap_percentile',
                   type=float,
                   default=99.0,
                   help='Global IPD percentile to cap IPDs at')

    p.add_argument("--methylMinCov",
                   type=int,
                   dest='methylMinCov',
                   default=10,
                   help="Do not try to estimate methylFraction unless coverage is at least this.")

    p.add_argument("--identifyMinCov",
                   type=int,
                   dest='identifyMinCov',
                   default=5,
                   help="Do not try to identify the modification type unless coverage is at least this.")

    p.add_argument("--maxAlignments",
                   type=int,
                   dest="maxAlignments",
                   default=1500,
                   help="Maximum number of alignments to use for a given window")

    # Computation management options:

    p.add_argument("-w", "--referenceWindow", "--referenceWindows",
                   "--refContigs",  # backwards compatibility
                   type=str,
                   dest='referenceWindowsAsString',
                   default=None,
                   help="The window (or multiple comma-delimited windows) of the reference to " + \
                   "be processed, in the format refGroup[:refStart-refEnd] " + \
                   "(default: entire reference).")

    def slurpWindowFile(fname):
        return ",".join(map(str.strip, open(fname).readlines()))

    p.add_argument("--refContigIndex", type=int, dest='refContigIndex', default=-1,
                   help="For debugging purposes only - rather than enter a reference contig name, simply enter an index")

    p.add_argument("-W", "--referenceWindowsFile",
                   "--refContigsFile",  # backwards compatibility
                   type=slurpWindowFile,
                   dest='referenceWindowsAsString',
                   default=None,
                   help="A file containing reference window designations, one per line")

    p.add_argument("--skipUnrecognizedContigs",
                   type=bool,
                   default=False,
                   help="Whether to skip, or abort, unrecognized contigs in the -w/-W flags")
    # FIXME shouldn't it always do this?
    p.add_argument("--alignmentSetRefWindows",
                   action="store_true",
                   dest="referenceWindowsFromAlignment",
                   help="Use refWindows in dataset")

    # Debugging help options:
    p.add_argument("--profile",
                   action="store_true",
                   dest="doProfiling",
                   default=False,
                   help="Enable Python-level profiling (using cProfile).")

    add_debug_option(p)

    p.add_argument("--seed",
                   action="store",
                   dest="randomSeed",
                   type=int,
                   default=None,
                   help="Random seed (for development and debugging purposes only)")

    p.add_argument("--referenceStride", action="store", type=int,
                   default=1000,
                   help="Size of reference window in internal " +
                   "parallelization.  For testing purposes only.")

    return p


class KineticsToolsRunner(object):
    def __init__(self, args):
        self.args = args
        self.alignments = None

    def start(self):
        self.validateArgs()
        return self.run()

    def getVersion(self):
        return __version__

    def validateArgs(self):
        parser = get_parser()
        if not os.path.exists(self.args.alignment_set):
            parser.error('Input AlignmentSet file provided does not exist')

        # Over-ride --identify if --control was specified
        if self.args.control:
            self.args.identify = ""

        if self.args.useLDA:
            if self.args.m5Cclassifier is None:
                parser.error(
                    'Please specify a folder containing forward.csv and reverse.csv classifiers in --m5Cclassifier.')

        if self.args.m5Cgff:
            if not self.args.useLDA:
                parser.error(
                    'm5Cgff file can only be generated in --useLDA mode.')

        # if self.args.methylFraction and not self.args.identify:
        #    parser.error('Currently, --methylFraction only works when the --identify option is specified.')

    def run(self):

        # Figure out what modifications to identify
        mods = self.args.identify
        modsToCall = []
        if mods:
            items = mods.split(",")

            if 'm6A' in items:
                modsToCall.append('H')

            if 'm4C' in items:
                modsToCall.append('J')

            if 'm5C_TET' in items:
                modsToCall.append('K')

            self.args.identify = True
            self.args.modsToCall = modsToCall

        self.options = self.args
        self.options.cmdLine = " ".join(sys.argv)
        self._workers = []

        # set random seed
        # XXX note that this is *not* guaranteed to yield reproducible results
        # indepenently of the number of processing cores used!
        if self.options.randomSeed is not None:
            np.random.seed(self.options.randomSeed)

        if self.args.doProfiling:
            cProfile.runctx("self._mainLoop()",
                            globals=globals(),
                            locals=locals(),
                            filename="profile.out")

        else:
            try:
                ret = self._mainLoop()
            finally:
                # Be sure to shutdown child processes if we get an exception on
                # the main thread
                for w in self._workers:
                    if w.is_alive():
                        w.terminate()

            return ret

    def _initQueues(self):
        # Work chunks are created by the main thread and put on this queue
        # They will be consumed by KineticWorker threads, stored in
        # self._workers
        self._workQueue = multiprocessing.JoinableQueue(
            self.options.maxQueueSize)

        # Completed chunks are put on this queue by KineticWorker threads
        # They are consumed by the KineticsWriter process
        self._resultsQueue = multiprocessing.JoinableQueue(
            self.options.maxQueueSize)

    def _launchSlaveProcesses(self):
        """
        Launch a group of worker processes (self._workers), the queue
        (self._workQueue) that will be used to send them chunks of
        work, and the queue that will be used to receive back the
        results (self._resultsQueue).

        Additionally, launch the result collector process.
        """
        availableCpus = multiprocessing.cpu_count()
        logging.info("Available CPUs: %d" % (availableCpus,))
        logging.info("Requested worker processes: %d" %
                     (self.options.numWorkers,))

        # Use all CPUs if numWorkers < 1
        if self.options.numWorkers < 1:
            self.options.numWorkers = availableCpus

        # Warn if we make a bad numWorker argument is used
        if self.options.numWorkers > availableCpus:
            logging.warn("More worker processes requested (%d) than CPUs available (%d);"
                         " may result in suboptimal performance."
                         % (self.options.numWorkers, availableCpus))

        self._initQueues()

        # Launch the worker processes
        self._workers = []
        for i in range(self.options.numWorkers):
            p = KineticWorkerProcess(
                self.options,
                self._workQueue,
                self._resultsQueue,
                self.ipdModel,
                sharedAlignmentSet=self.alignments)
            self._workers.append(p)
            p.start()
        logging.info("Launched worker processes.")

        # Launch result collector
        self._resultCollectorProcess = KineticsWriter(
            self.options, self._resultsQueue, self.refInfo, self.ipdModel)
        self._resultCollectorProcess.start()
        logging.info("Launched result collector process.")

        # Spawn a thread that monitors worker threads for crashes
        self.monitoringThread = threading.Thread(target=monitorChildProcesses, args=(
            self._workers + [self._resultCollectorProcess],))
        self.monitoringThread.start()

    def _queueChunksForWindow(self, refWindow):
        """
        Compute the chunk extents and queue up the work for a single reference
        """
        winId = refWindow.refId
        winStart = refWindow.start
        winEnd = refWindow.end
        pass

    def loadReferenceAndModel(self, referencePath, ipdModelFilename):
        assert self.alignments is not None and self.referenceWindows is not None
        # Load the reference contigs - annotated with their refID from the
        # alignments
        logging.info("Loading reference contigs {!r}".format(referencePath))
        contigs = ReferenceUtils.loadReferenceContigs(referencePath,
                                                      alignmentSet=self.alignments, windows=self.referenceWindows)
        self.ipdModel = IpdModel(
            contigs, ipdModelFilename, self.args.modelIters)

    def loadSharedAlignmentSet(self, alignmentFilename):
        """
        Read the input AlignmentSet so the indices can be shared with the
        slaves.  This is also used to pass to ReferenceUtils for setting up
        the ipdModel object.
        """
        logging.info("Reading AlignmentSet: %s" % alignmentFilename)
        logging.info("           reference: %s" % self.args.reference)
        self.alignments = AlignmentSet(alignmentFilename,
                                       referenceFastaFname=self.args.reference)
        # XXX this should ensure that the file(s) get opened, including any
        # .pbi indices - but need to confirm this
        self.refInfo = self.alignments.referenceInfoTable

    def _mainLoop(self):
        """
        Main loop
        First launch the worker and writer processes
        Then we loop over ReferenceGroups in the alignments.  For each contig we will:
        1. Load the sequence into the main memory of the parent process
        3. Chunk up the contig and submit the chunk descriptions to the work queue
        Finally, wait for the writer process to finish.
        """

        # This looks scary but it's not.  Python uses reference
        # counting and has a secondary, optional garbage collector for
        # collecting garbage cycles.  Unfortunately when a cyclic GC
        # happens when a thread is calling cPickle.dumps, the
        # interpreter crashes sometimes.  See Bug 19704.  Since we
        # don't leak garbage cycles, disabling the cyclic GC is
        # essentially harmless.
        # gc.disable()

        self.loadSharedAlignmentSet(self.args.alignment_set)

        # Resolve the windows that will be visited.
        if self.args.referenceWindowsAsString is not None:
            self.referenceWindows = []
            for s in self.args.referenceWindowsAsString.split(","):
                try:
                    win = ReferenceUtils.parseReferenceWindow(
                        s, self.alignments.referenceInfo)
                    self.referenceWindows.append(win)
                except BaseException:
                    if self.args.skipUnrecognizedContigs:
                        continue
                    else:
                        raise Exception("Unrecognized contig!")
        elif self.args.referenceWindowsFromAlignment:
            self.referenceWindows = ReferenceUtils.referenceWindowsFromAlignment(
                self.alignments, self.alignments.referenceInfo)
            refNames = set([rw.refName for rw in self.referenceWindows])
            # limit output to contigs that overlap with reference windows
            self.refInfo = [r for r in self.refInfo if r.Name in refNames]
        else:
            self.referenceWindows = ReferenceUtils.createReferenceWindows(
                self.refInfo)

        # Load reference and IpdModel
        chemName = ReferenceUtils.loadAlignmentChemistry(self.alignments)
        if self.args.useChemistry is not None:
            chemName = self.args.useChemistry
        ipdModelFilename = loader.getIpdModelFilename(
            ipdModel=self.args.ipdModel,
            majorityChem=chemName,
            paramsPath=self.args.paramsPath)
        self.loadReferenceAndModel(self.args.reference, ipdModelFilename)

        # Spawn workers
        self._launchSlaveProcesses()

        logging.info(
            'Generating kinetics summary for [%s]' % self.args.alignment_set)

        #self.referenceMap = self.alignments['/RefGroup'].asDict('RefInfoID', 'ID')
        #self.alnInfo = self.alignments['/AlnInfo'].asRecArray()

        # Main loop -- we loop over ReferenceGroups in the alignments.  For each contig we will:
        # 1. Load the sequence into the main memory of the parent process
        # 2. Fork the workers
        # 3. chunk up the contig and

        self.workChunkCounter = 0

        # Iterate over references
        for window in self.referenceWindows:
            logging.info('Processing window/contig: %s' % (window,))
            for chunk in ReferenceUtils.enumerateChunks(
                    self.args.referenceStride, window):
                self._workQueue.put((self.workChunkCounter, chunk))
                self.workChunkCounter += 1

        # Shutdown worker threads with None sentinels
        for i in range(self.args.numWorkers):
            self._workQueue.put(None)

        for w in self._workers:
            w.join()

        # Join on the result queue and the resultsCollector process.
        # This ensures all the results are written before shutdown.
        self.monitoringThread.join()
        self._resultsQueue.join()
        self._resultCollectorProcess.join()
        logging.info("ipdSummary.py finished. Exiting.")
        self.alignments.close()
        return 0


def monitorChildProcesses(children):
    """
    Monitors child processes: promptly exits if a child is found to
    have exited with a nonzero exit code received; otherwise returns
    zero when all processes exit cleanly (0).

    This approach is portable--catching SIGCHLD doesn't work on
    Windows.
    """
    while True:
        all_exited = all(not p.is_alive() for p in children)
        nonzero_exits = [p.exitcode for p in children if p.exitcode]
        if nonzero_exits:
            exitcode = nonzero_exits[0]
            logging.error(
                "Child process exited with exitcode=%d.  Aborting." % exitcode)

            # Kill all the child processes
            for p in children:
                if p.is_alive():
                    p.terminate()

            os._exit(exitcode)
        elif all_exited:
            return 0
        time.sleep(1)


def args_runner(args):
    kt = KineticsToolsRunner(args)
    return kt.start()


def main(argv=sys.argv, out=sys.stdout):
    setup_log_ = functools.partial(setup_log,
                                   str_formatter='%(asctime)s [%(levelname)s] %(message)s')
    try:
        return pacbio_args_runner(
            argv=argv[1:],
            parser=get_parser(),
            args_runner_func=args_runner,
            alog=logging.getLogger(__name__),
            setup_log_func=setup_log)
    # FIXME is there a more central place to deal with this?
    except Exception as e:
        type, value, tb = sys.exc_info()
        traceback.print_exc(file=sys.stderr)
        # Note: if kt.args.usePdb
        # This won't work. If an exception is raised in parseArgs,
        # then kt.args is not defined yet.
        if '--pdb' in argv:
            try:
                # this has better integration with ipython and is nicer
                # pip install ipdb
                import ipdb
                ipdb.post_mortem(tb)
            except ImportError:
                import pdb
                pdb.post_mortem(tb)
        else:
            # exit non-zero
            raise


if __name__ == "__main__":
    sys.exit(main())