import ctypes
import logging
import multiprocessing as mp
from functools import reduce
from operator import mul

import numpy as np

from gunpowder.array import Array, ArrayKey
from gunpowder.ext import tensorflow as tf
from gunpowder.nodes.generic_predict import GenericPredict
from gunpowder.tensorflow.local_server import LocalServer

logger = logging.getLogger(__name__)


class Predict(GenericPredict):
    """Tensorflow implementation of :class:`gunpowder.nodes.Predict`.

    Args:

        checkpoint (``string``):

            Basename of a tensorflow checkpoint storing the tensorflow graph
            and associated tensor values and metadata, as created by
            :class:`gunpowder.nodes.Train`, for example.

        inputs (``dict``, ``string`` -> :class:`ArrayKey`):

            Dictionary from the names of input tensors in the network to
            array keys.

        outputs (``dict``, ``string`` -> :class:`ArrayKey`):

            Dictionary from the names of output tensors in the network to array
            keys. New arrays will be generated by this node for each entry (if
            requested downstream).

        array_specs (``dict``, :class:`ArrayKey` -> :class:`ArraySpec`, optional):

            Used to set the specs of generated arrays (``outputs``). 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.

        graph: (``string``, optional):

            An optional path to a tensorflow computation graph that should be
            used for prediction. The checkpoint is used to restore the values
            of matching variable names in the graph. Note that the graph
            specified here can differ from the one associated to the
            checkpoint.

        skip_empty (``bool``, optional):

            Skip prediction, if all inputs are empty (contain only 0). In this
            case, outputs are simply set to 0.

        max_shared_memory (``int``, optional):

            The maximal amount of shared memory in bytes to allocate to send
            batches to the GPU processes. Defaults to 1GB.
    """

    def __init__(
        self,
        checkpoint,
        inputs,
        outputs,
        array_specs=None,
        graph=None,
        skip_empty=False,
        max_shared_memory=1024 * 1024 * 1024,
    ):
        super(Predict, self).__init__(inputs, outputs, array_specs)

        self.checkpoint = checkpoint
        self.meta_graph = graph
        self.session = None
        self.graph = None
        self.skip_empty = skip_empty

        self.manager = mp.Manager()
        self.max_shared_memory = max_shared_memory
        self.shared_input_array_config = self.manager.dict()
        self.shared_output_array_config = self.manager.dict()
        self.shared_input_arrays = {}
        self.shared_output_arrays = {}
        self.shared_input_memory = mp.RawArray(ctypes.c_float, self.max_shared_memory)
        self.shared_output_memory = mp.RawArray(ctypes.c_float, self.max_shared_memory)

        self.send_lock = mp.Lock()
        self.receive_lock = mp.Lock()
        self.predict_process_initialized = mp.Event()
        self.worker_sent_inputs = mp.Event()
        self.predict_received_inputs = mp.Event()
        self.predict_sent_outputs = mp.Event()

        self.predict_process = mp.Process(target=self.__predict)
        self.predict_process_crashed = mp.Value("i", False)
        self.predict_process.start()
        self.predict_process_initialized.wait()

    def predict(self, batch, request):
        if not self.shared_output_arrays:
            self.__init_shared_output_arrays()

        if self.skip_empty:
            can_skip = True
            for array_key in self.inputs.values():
                if batch[array_key].data.any():
                    can_skip = False
                    break

            if can_skip:
                logger.info(f"Skipping batch for request: {request} (all inputs are 0)")

                for name, array_key in self.outputs.items():
                    shape = self.shared_output_arrays[name].shape
                    dtype = self.shared_output_arrays[name].dtype

                    spec = self.spec[array_key].copy()
                    spec.roi = request[array_key].roi.copy()
                    batch.arrays[array_key] = Array(np.zeros(shape, dtype=dtype), spec)

                return

        logger.debug(f"predicting for request: {request}")

        output_tensors = self.__collect_outputs(request)
        input_data = self.__collect_provided_inputs(batch)

        self.send_lock.acquire()

        if not self.shared_input_arrays:
            if not self.shared_input_array_config:
                self.__create_shared_input_array_config(batch, request)
            self.__init_shared_input_arrays()

        self.__write_inputs_to_shared(input_data)
        self.worker_sent_inputs.set()

        self.receive_lock.acquire()

        self.predict_received_inputs.wait()
        self.__check_background_process([self.receive_lock, self.send_lock])

        self.predict_received_inputs.clear()
        self.send_lock.release()

        self.predict_sent_outputs.wait()

        self.predict_sent_outputs.clear()

        output_data = self.__read_outputs_from_shared(output_tensors)

        self.receive_lock.release()

        for array_key in output_tensors:
            spec = self.spec[array_key].copy()
            spec.roi = request[array_key].roi
            batch.arrays[array_key] = Array(output_data[array_key], spec)

        logger.debug("predicted")

    def __predict(self):
        """The background predict process."""

        try:
            # TODO: is the server still needed?
            target = LocalServer.get_target()
            logger.info("Initializing tf session, connecting to %s...", target)

            self.graph = tf.Graph()
            self.session = tf.Session(target=target, graph=self.graph)

            with self.graph.as_default():
                self.__read_checkpoint()

            if not self.shared_output_arrays:
                if not self.shared_output_array_config:
                    self.__create_shared_output_array_config()
                self.__init_shared_output_arrays()

            # from now on it is save to access the shared array configuration
            self.predict_process_initialized.set()

            # loop predict
            while True:
                # wait for inputs
                self.worker_sent_inputs.wait()
                self.worker_sent_inputs.clear()

                if not self.shared_input_arrays:
                    self.__init_shared_input_arrays()

                # read inputs
                input_data = self.__read_inputs_from_shared()
                self.predict_received_inputs.set()

                # compute outputs
                output_data = self.session.run(
                    {t: t for t in self.outputs.keys()}, feed_dict=input_data
                )

                # write outputs
                self.__write_outputs_to_shared(output_data)
                self.predict_sent_outputs.set()

        except Exception as e:
            self.predict_process_crashed.value = True

            # release locks and events
            self.predict_process_initialized.set()
            self.worker_sent_inputs.clear()
            self.predict_received_inputs.set()
            self.predict_sent_outputs.set()
            raise e

    def teardown(self):
        self.predict_process.terminate()
        self.predict_process.join()

    def __check_background_process(self, locks=[]):
        if self.predict_process_crashed.value:
            # release all locks before raising exception
            for l in locks:
                l.release()
            raise RuntimeError("Background process died.")

    def __read_checkpoint(self):
        # read the graph associated to the checkpoint
        if self.meta_graph is None:
            meta_graph_file = self.checkpoint + ".meta"
        # read alternative, custom graph
        else:
            meta_graph_file = self.meta_graph

        logger.info(
            "Reading graph from %s and weights from %s...",
            meta_graph_file,
            self.checkpoint,
        )

        saver = tf.train.import_meta_graph(meta_graph_file, clear_devices=True)

        # restore variables from checkpoint
        saver.restore(self.session, self.checkpoint)

    def __collect_outputs(self, request=None):
        """Get a dict:

            array key: tensor name

        If request is not None, return only outputs that are in request.
        """

        array_outputs = {}

        for tensor_name, array_key in self.outputs.items():
            if request is None or array_key in request:
                array_outputs[array_key] = tensor_name

        return array_outputs

    def __collect_provided_inputs(self, batch):
        """Get a dict:

        tensor name: ndarray
        """

        inputs = {}

        for input_name, input_key in self.inputs.items():
            if isinstance(input_key, ArrayKey):
                if input_key in batch.arrays:
                    inputs[input_name] = batch.arrays[input_key].data
                else:
                    logger.warn(
                        "batch does not contain %s, input %s will not be set",
                        input_key,
                        input_name,
                    )
            elif isinstance(input_key, np.ndarray):
                inputs[input_name] = input_key
            elif isinstance(input_key, str):
                inputs[input_name] = getattr(batch, input_key)
            else:
                raise Exception(
                    "Unknown network input key {}, can't be given to network".format(
                        input_key
                    )
                )

        return inputs

    def __create_shared_input_array_config(self, batch, request):
        """Store the shared array config in a shared dictionary. Should be run
        once by the first worker to submit a batch."""

        begin = 0
        for name, array_key in self.inputs.items():
            shape = batch[array_key].data.shape
            size = reduce(mul, shape, 1)
            dtype = batch[array_key].data.dtype

            self.shared_input_array_config[name] = (begin, size, shape, dtype)

            begin += size * np.dtype(dtype).itemsize
            assert (
                begin <= self.max_shared_memory
            ), "The input arrays exceed the max_shared_memory"

    def __create_shared_output_array_config(self):
        """To be called by predict process."""

        begin = 0
        for name, array_key in self.outputs.items():
            tensor = self.graph.get_tensor_by_name(name)
            shape = tensor.get_shape().as_list()
            size = reduce(mul, shape, 1)
            dtype = tensor.dtype.as_numpy_dtype

            self.shared_output_array_config[name] = (begin, size, tuple(shape), dtype)

            begin += size * np.dtype(dtype).itemsize
            assert (
                begin <= self.max_shared_memory
            ), "The output arrays exceed the max_shared_memory"

    def __init_shared_input_arrays(self):
        """Assign the shared memory to numpy arrays."""

        for name, (begin, size, shape, dtype) in self.shared_input_array_config.items():
            self.shared_input_arrays[name] = np.frombuffer(
                self.shared_input_memory, dtype=dtype, offset=begin, count=size
            ).reshape(shape)

    def __init_shared_output_arrays(self):
        """Assign the shared memory to numpy arrays."""

        for name, (
            begin,
            size,
            shape,
            dtype,
        ) in self.shared_output_array_config.items():
            self.shared_output_arrays[name] = np.frombuffer(
                self.shared_output_memory, dtype=dtype, offset=begin, count=size
            ).reshape(shape)

    def __write_inputs_to_shared(self, input_data):
        for tensor_name, data in input_data.items():
            self.shared_input_arrays[tensor_name][:] = data

    def __read_inputs_from_shared(self):
        return {
            tensor_name: self.shared_input_arrays[tensor_name].copy()
            for tensor_name in self.inputs.keys()
        }

    def __write_outputs_to_shared(self, output_data):
        for tensor_name, data in output_data.items():
            self.shared_output_arrays[tensor_name][:] = data

    def __read_outputs_from_shared(self, output_tensors):
        return {
            array_key: self.shared_output_arrays[tensor_name].copy()
            for array_key, tensor_name in output_tensors.items()
        }
