import logging
import multiprocessing
import time

from gunpowder.nodes.batch_filter import BatchFilter
from gunpowder.producer_pool import ProducerPool, WorkersDied, NoResult
from gunpowder.array import ArrayKey
from gunpowder.array_spec import ArraySpec
from gunpowder.batch_request import BatchRequest
from queue import Full

logger = logging.getLogger(__name__)


class PredictProcessDied(Exception):
    pass


class GenericPredict(BatchFilter):
    """Generic predict node to add predictions of a trained network to each each
    batch that passes through. This node alone does nothing and should be
    subclassed for concrete implementations.

    Args:

        inputs (dict): Dictionary from names of input layers in the network to
            :class:``ArrayKey`` or batch attribute name as string.

        outputs (dict): Dictionary from the names of output layers in the
            network to :class:``ArrayKey``. New arrays will be generated by
            this node for each entry (if requested downstream).

        array_specs (dict, optional): An optional dictionary of
            :class:`ArrayKey` to :class:`ArraySpec` to set the array specs
            generated arrays (``outputs`` and ``gradients``). This is useful
            to set the ``voxel_size``, for example, if they differ from the
            voxel size of the input arrays. Only fields that are not ``None``
            in the given :class:`ArraySpec` will be used.

        spawn_subprocess (bool, optional): Whether to run ``predict`` in a
            separate process. Default is false.
    """

    def __init__(self, inputs, outputs, array_specs=None, spawn_subprocess=False):
        self.initialized = False
        self.inputs = inputs
        self.outputs = outputs
        self.array_specs = {} if array_specs is None else array_specs
        self.spawn_subprocess = spawn_subprocess
        self.timer_start = None

    def setup(self):
        # get common voxel size of inputs, or None if they differ
        common_voxel_size = None
        for key in self.inputs.values():
            if not isinstance(key, ArrayKey):
                continue

            voxel_size = self.spec[key].voxel_size

            if common_voxel_size is None:
                common_voxel_size = voxel_size
            elif common_voxel_size != voxel_size:
                common_voxel_size = None
                break

        # announce provided outputs
        for key in self.outputs.values():
            if key in self.array_specs:
                spec = self.array_specs[key].copy()
            else:
                spec = ArraySpec()

            if spec.voxel_size is None and not spec.nonspatial:
                assert common_voxel_size is not None, (
                    "There is no common voxel size of the inputs, and no "
                    "ArraySpec has been given for %s that defines "
                    "voxel_size." % key
                )

                spec.voxel_size = common_voxel_size

            if spec.interpolatable is None:
                # default for predictions
                spec.interpolatable = False

            self.provides(key, spec)

        if self.spawn_subprocess:
            # start prediction as a producer pool, so that we can gracefully
            # exit if anything goes wrong
            self.worker = ProducerPool([self._produce_predict_batch], queue_size=1)
            self.batch_in = multiprocessing.Queue(maxsize=1)
            self.batch_in_lock = multiprocessing.Lock()
            self.batch_out_lock = multiprocessing.Lock()
            self.worker.start()

    def teardown(self):
        if self.spawn_subprocess:
            # signal "stop"
            try:
                self.batch_in.put((None, None), timeout=2)
            except Full:
                # worker process might be stopped already
                pass
            try:
                self.worker.get(timeout=2)
            except NoResult:
                pass
            self.worker.stop()
        else:
            self.stop()

    def prepare(self, request):
        if not self.initialized and not self.spawn_subprocess:
            self.start()
            self.initialized = True

        deps = BatchRequest()
        for key in self.inputs.values():
            deps[key] = request[key]
        return deps

    def process(self, batch, request):
        if self.spawn_subprocess:
            start = time.time()
            self.batch_in_lock.acquire()
            logger.debug("waited for batch in lock for %.3fs", time.time() - start)
            start = time.time()
            self.batch_in.put((batch, request))
            logger.debug("queued batch for %.3fs", time.time() - start)

            start = time.time()
            with self.batch_out_lock:
                logger.debug("waited for batch out lock for %.3fs", time.time() - start)

                start = time.time()
                self.batch_in_lock.release()
                logger.debug("released batch in lock for %.3fs", time.time() - start)
                try:
                    start = time.time()
                    out = self.worker.get()
                    logger.debug("retreived batch for %.3fs", time.time() - start)
                except WorkersDied:
                    raise PredictProcessDied()

            for array_key in self.outputs.values():
                if array_key in request:
                    batch.arrays[array_key] = out.arrays[array_key]

        else:
            self.predict(batch, request)

    def start(self):
        """To be implemented in subclasses.

        This method will be called before the first call to :fun:`predict`,
        from the same process that :fun:`predict` will be called from. Use
        this to initialize your model and hardware.
        """
        pass

    def predict(self, batch, request):
        """To be implemented in subclasses.

        In this method, an implementation should predict arrays on the given
        batch. Output arrays should be created according to the given request
        and added to ``batch``."""
        raise NotImplementedError("Class %s does not implement 'predict'" % self.name())

    def stop(self):
        """To be implemented in subclasses.

        This method will be called after the last call to :fun:`predict`,
        from the same process that :fun:`predict` will be called from. Use
        this to tear down your model and free training hardware.
        """
        pass

    def _produce_predict_batch(self):
        """Process one batch."""

        if not self.initialized:
            self.start()
            self.initialized = True

        self.time_out = 0
        if self.timer_start is not None:
            self.time_out = time.time() - self.timer_start

        self.timer_start = time.time()
        batch, request = self.batch_in.get()
        self.time_in = time.time() - self.timer_start

        # stop signal
        if batch is None:
            self.stop()
            return None

        self.timer_start = time.time()
        self.predict(batch, request)
        self.time_predict = time.time() - self.timer_start
        self.timer_start = time.time()

        logger.debug(
            "batch in: %.3fs, predict: %.3fs, batch out: %.3fs",
            self.time_in,
            self.time_predict,
            self.time_out,
        )

        return batch
