import logging
import os
import json
from multiprocessing import Process, Queue, Lock

import numpy as np

from PyMaSC.core.mappability import MappableLengthCalculator
from PyMaSC.utils.progress import ProgressHook, MultiLineProgressManager
from PyMaSC.utils.compatible import tostr, xrange
from PyMaSC.utils.output import prepare_outdir
from PyMaSC.utils.calc import exec_worker_pool

logger = logging.getLogger(__name__)


class BWIOError(IOError):
    pass


class JSONIOError(IOError):
    pass


class NeedUpdate(Exception):
    pass


class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (np.long, np.float, np.float_)):
            return float(obj)
        elif isinstance(obj, (np.uint, np.int32, np.int64)):
            return int(obj)
        else:
            return super(self, NumpyEncoder).default(obj)


class MappabilityHandler(MappableLengthCalculator):
    @staticmethod
    def calc_mappable_len_required_shift_size(readlen, max_shift):
        return max_shift - readlen + 1 if max_shift > 2*readlen - 1 else readlen

    def __init__(self, path, max_shift=0, readlen=0, map_path=None, nworker=1):
        max_shift = self.calc_mappable_len_required_shift_size(readlen, max_shift)
        self.nworker = nworker

        if not os.access(path, os.R_OK):
            reason = "file is unreadable." if os.path.isfile(path) else "no such file."
            logger.critical("Failed to open '{}': {}".format(path, reason))
            raise BWIOError

        super(MappabilityHandler, self).__init__(path, max_shift)
        self.close()
        self._progress.disable_bar()
        self.need_save_stats = True

        if map_path:
            self.map_path = map_path
        else:
            self.map_path = os.path.splitext(path)[0] + "_mappability.json"

        if not os.path.exists(self.map_path):
            self._check_saving_directory_is_writable()
            logger.info("Calcurate mappable length with max shift size {}.".format(max_shift))
        elif not os.path.isfile(self.map_path):
            logger.critical("Specified path is not file: '{}'".format(self.map_path))
            raise JSONIOError
        elif not os.access(self.map_path, os.R_OK):
            logger.error("Failed to read '{}'".format(self.map_path))
        else:
            self._try_load_mappability_stats()
            if self.need_save_stats:
                self._check_stats_is_overwritable()
                logger.info("Calcurate mappable length with max shift size {}.".format(max_shift))
            else:
                logger.info("Use mappability stats read from '{}'".format(self.map_path))

    def _check_saving_directory_is_writable(self):
        dirname = os.path.dirname(self.map_path)
        dirname = dirname if dirname else '.'
        if not prepare_outdir(dirname, logger):
            raise JSONIOError

    def _try_load_mappability_stats(self):
        try:
            stats = self._read_mappability_stats()
        except IOError as e:
            logger.error("Failed to read '{}'".format(self.map_path))
            logger.error("[Errno {}] {}".format(e.errno, e.message))
        except (TypeError, OverflowError, ValueError, KeyError, IndexError) as e:
            logger.error("Failed to load json file: '{}'".format(self.map_path))
        except NeedUpdate:
            pass
        else:
            self._load_mappability_stats(stats)

    def _read_mappability_stats(self):
        with open(self.map_path) as f:
            stats = json.load(f)

        for k in ("max_shift", "__whole__", "references"):
            if k not in stats:
                logger.error("Mandatory key '{}' not found.".format(k))
                raise KeyError(k)

        if stats["max_shift"] < self.max_shift:
            logger.info("Specified shift length longer than former analysis. The stats will be updated.")
            raise NeedUpdate

        if stats["max_shift"] != len(stats["__whole__"]) - 1:
            logger.error("Max shift length for whole genome unmatched.")
            raise IndexError

        for ref in self.chromsizes:
            if ref not in stats["references"]:
                logger.error("Reference '{}' not found.".format(ref))
                raise KeyError(ref)

            if stats["max_shift"] != len(stats["references"][ref]) - 1:
                logger.error("Max shift length for 'ref' unmatched.".format(ref))
                raise IndexError

        return stats

    def _load_mappability_stats(self, stats):
        self.mappable_len = stats["__whole__"][:self.max_shift + 1]
        self.chrom2mappable_len = {ref: b[:self.max_shift + 1] for ref, b in stats["references"].items()}
        self.chrom2is_called = {ref: True for ref in self.chromsizes}
        self.is_called = True
        self.need_save_stats = False

    def _check_stats_is_overwritable(self):
        if not os.access(self.map_path, os.W_OK):
            logger.critical("Failed to overwrite '{}'".format(self.map_path))
            raise JSONIOError
        else:
            logger.warning("Existing file '{}' will be overwritten.".format(self.map_path))

    def save_mappability_stats(self):
        if not self.need_save_stats:
            return logger.info("Mappability stats updating is not required.")

        logger.info("Save mappable length to '{}'".format(self.map_path))
        try:
            with open(self.map_path, 'w') as f:
                json.dump({
                    "max_shift": self.max_shift,
                    "__whole__": self.mappable_len,
                    "references": self.chrom2mappable_len
                }, f, indent=4, sort_keys=True, cls=NumpyEncoder)
        except IOError as e:
            logger.error("Faild to output: {}\n[Errno {}] {}".format(
                         e.filename, e.errno, e.message))

        self.need_save_stats = False

    def calc_mappability(self):
        target_chroms = [tostr(c) for c, b in self.chrom2is_called.items() if b is False]
        if not target_chroms:
            return self._sumup_mappability()

        order_queue = Queue()
        report_queue = Queue()
        logger_lock = Lock()
        progress = MultiLineProgressManager()

        workers = [MappabilityCalcWorker(self.path, self.max_shift, order_queue, report_queue, logger_lock)
                   for _ in range(min(self.nworker, len(target_chroms)))]

        with exec_worker_pool(workers, target_chroms, order_queue):
            while not self.is_called:
                chrom, obj = report_queue.get()
                if chrom is None:  # update progress
                    chrom, body = obj
                    with logger_lock:
                        progress.update(chrom, body)
                else:
                    length = obj
                    self.chrom2mappable_len[chrom] = tuple(length)
                    self.chrom2is_called[chrom] = True
                    if all(self.chrom2is_called.values()):
                        self.is_called = True
                    with logger_lock:
                        progress.erase(chrom)

        progress.clean()
        self._sumup_mappability()

    def _sumup_mappability(self):
        for length in self.chrom2mappable_len.values():
            for i in xrange(self.max_shift + 1):
                self.mappable_len[i] += length[i]


class MappabilityCalcWorker(Process):
    def __init__(self, path, max_shift, order_queue, report_queue, logger_lock):
        super(MappabilityCalcWorker, self).__init__()

        self.calculator = MappableLengthCalculator(path, max_shift, logger_lock)
        self.calculator._progress.disable_bar()
        self.order_queue = order_queue
        self.report_queue = report_queue
        self.logger_lock = logger_lock
        self.calculator._progress = ProgressHook(report_queue)

    def run(self):
        with self.logger_lock:
            logger.debug("{}: Hello. My pid is {}.".format(self.name, os.getpid()))

        while True:
            chrom = self.order_queue.get()
            if chrom is None:
                break

            with self.logger_lock:
                logger.debug("{}: Process {}...".format(self.name, chrom))
            self.calculator.calc_mappability(chrom)
            self.report_queue.put((chrom, self.calculator.chrom2mappable_len[chrom]))

        with self.logger_lock:
            logger.debug("{}: Goodbye.".format(self.name))
        self.calculator.close()

[((14, 9, 14, 36), 'logging.getLogger', 'logging.getLogger', ({(14, 27, 14, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((80, 18, 80, 48), 'os.path.dirname', 'os.path.dirname', ({(80, 34, 80, 47): 'self.map_path'}, {}), '(self.map_path)', False, 'import os\n'), ((163, 22, 163, 29), 'multiprocessing.Queue', 'Queue', ({}, {}), '()', False, 'from multiprocessing import Process, Queue, Lock\n'), ((164, 23, 164, 30), 'multiprocessing.Queue', 'Queue', ({}, {}), '()', False, 'from multiprocessing import Process, Queue, Lock\n'), ((165, 22, 165, 28), 'multiprocessing.Lock', 'Lock', ({}, {}), '()', False, 'from multiprocessing import Process, Queue, Lock\n'), ((166, 19, 166, 45), 'PyMaSC.utils.progress.MultiLineProgressManager', 'MultiLineProgressManager', ({}, {}), '()', False, 'from PyMaSC.utils.progress import ProgressHook, MultiLineProgressManager\n'), ((200, 26, 200, 80), 'PyMaSC.core.mappability.MappableLengthCalculator', 'MappableLengthCalculator', ({(200, 51, 200, 55): 'path', (200, 57, 200, 66): 'max_shift', (200, 68, 200, 79): 'logger_lock'}, {}), '(path, max_shift, logger_lock)', False, 'from PyMaSC.core.mappability import MappableLengthCalculator\n'), ((205, 36, 205, 62), 'PyMaSC.utils.progress.ProgressHook', 'ProgressHook', ({(205, 49, 205, 61): 'report_queue'}, {}), '(report_queue)', False, 'from PyMaSC.utils.progress import ProgressHook, MultiLineProgressManager\n'), ((48, 15, 48, 39), 'os.access', 'os.access', ({(48, 25, 48, 29): 'path', (48, 31, 48, 38): 'os.R_OK'}, {}), '(path, os.R_OK)', False, 'import os\n'), ((63, 15, 63, 44), 'os.path.exists', 'os.path.exists', ({(63, 30, 63, 43): 'self.map_path'}, {}), '(self.map_path)', False, 'import os\n'), ((82, 15, 82, 46), 'PyMaSC.utils.output.prepare_outdir', 'prepare_outdir', ({(82, 30, 82, 37): 'dirname', (82, 39, 82, 45): 'logger'}, {}), '(dirname, logger)', False, 'from PyMaSC.utils.output import prepare_outdir\n'), ((100, 20, 100, 32), 'json.load', 'json.load', ({(100, 30, 100, 31): 'f'}, {}), '(f)', False, 'import json\n'), ((134, 15, 134, 48), 'os.access', 'os.access', ({(134, 25, 134, 38): 'self.map_path', (134, 40, 134, 47): 'os.W_OK'}, {}), '(self.map_path, os.W_OK)', False, 'import os\n'), ((159, 25, 159, 33), 'PyMaSC.utils.compatible.tostr', 'tostr', ({(159, 31, 159, 32): 'c'}, {}), '(c)', False, 'from PyMaSC.utils.compatible import tostr, xrange\n'), ((171, 13, 171, 66), 'PyMaSC.utils.calc.exec_worker_pool', 'exec_worker_pool', ({(171, 30, 171, 37): 'workers', (171, 39, 171, 52): 'target_chroms', (171, 54, 171, 65): 'order_queue'}, {}), '(workers, target_chroms, order_queue)', False, 'from PyMaSC.utils.calc import exec_worker_pool\n'), ((192, 21, 192, 47), 'PyMaSC.utils.compatible.xrange', 'xrange', ({(192, 28, 192, 46): '(self.max_shift + 1)'}, {}), '(self.max_shift + 1)', False, 'from PyMaSC.utils.compatible import tostr, xrange\n'), ((49, 46, 49, 66), 'os.path.isfile', 'os.path.isfile', ({(49, 61, 49, 65): 'path'}, {}), '(path)', False, 'import os\n'), ((66, 17, 66, 46), 'os.path.isfile', 'os.path.isfile', ({(66, 32, 66, 45): 'self.map_path'}, {}), '(self.map_path)', False, 'import os\n'), ((147, 16, 151, 65), 'json.dump', 'json.dump', (), '', False, 'import json\n'), ((61, 28, 61, 50), 'os.path.splitext', 'os.path.splitext', ({(61, 45, 61, 49): 'path'}, {}), '(path)', False, 'import os\n'), ((69, 17, 69, 50), 'os.access', 'os.access', ({(69, 27, 69, 40): 'self.map_path', (69, 42, 69, 49): 'os.R_OK'}, {}), '(self.map_path, os.R_OK)', False, 'import os\n'), ((209, 70, 209, 81), 'os.getpid', 'os.getpid', ({}, {}), '()', False, 'import os\n')]