import logging
import numpy as np

from gunpowder.caffe.net_io_wrapper import NetIoWrapper
from gunpowder.ext import caffe
from gunpowder.nodes.generic_train import GenericTrain
from gunpowder.array import Array, ArrayKey

logger = logging.getLogger(__name__)

class Train(GenericTrain):
    '''Caffe implementation of :class:`gunpowder.nodes.Train`.

    Args:

        solver_parameters (:class:`SolverParameters`):

            Parameters of the solver to use for training, contains the network
            description as well.

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

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

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

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

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

            Dictionary from the names of output layers in the network to array
            keys. New arrays containing the gradient of an output with respect
            to the loss will be generated by this node for each entry (if
            requested downstream).

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

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

        use_gpu (``int`` or ``None``, optional):

            Which GPU to use. Set to ``None`` for CPU mode.
    '''

    def __init__(
            self,
            solver_parameters,
            inputs,
            outputs,
            gradients,
            array_specs=None,
            use_gpu=None):

        super(Train, self).__init__(
            inputs,
            outputs,
            gradients,
            array_specs,
            spawn_subprocess=True)
        self.solver_parameters = solver_parameters
        self.use_gpu = use_gpu
        self.solver = None
        self.net_io = None

    def start(self):

        logger.info("Initializing solver...")

        if self.use_gpu is not None:

            logger.debug("Train process: using GPU %d", self.use_gpu)
            caffe.enumerate_devices(False)
            caffe.set_devices((self.use_gpu,))
            caffe.set_mode_gpu()
            caffe.select_device(self.use_gpu, False)

        self.solver = caffe.get_solver(self.solver_parameters)
        if self.solver_parameters.resume_from is not None:
            logger.debug(
                "Train process: restoring solver state from %s",
                self.solver_parameters.resume_from)
            self.solver.restore(self.solver_parameters.resume_from)

        names_net_outputs = self.outputs.keys() + self.gradients.keys()
        self.net_io = NetIoWrapper(self.solver.net, names_net_outputs)

    def train_step(self, batch, request):

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

        loss = self.solver.step(1)
        # self.__consistency_check()

        requested_outputs = {
            name: array_key
            for name, array_key in self.outputs.items()
            if array_key in request.array_specs }

        if requested_outputs:

            output = self.net_io.get_outputs()

            for output_name, array_key in requested_outputs.items():

                spec = self.spec[array_key].copy()
                spec.roi = request[array_key].roi
                batch.arrays[array_key] = Array(
                    output[output_name][0], # strip #batch dimension
                    spec)

        requested_gradients = {
            name: array_key
            for name, array_key in self.gradients.items()
            if array_key in request.array_specs }

        if requested_gradients:

            diffs = self.net_io.get_output_diffs()

            for output_name, array_key in requested_gradients.items():

                spec = self.spec[array_key].copy()
                spec.roi = request[array_key].roi
                batch.arrays[array_key] = Array(
                    diffs[output_name][0], # strip #batch dimension
                    spec)

        batch.loss = loss
        batch.iteration = self.solver.iter

    def __consistency_check(self):

        diffs = self.net_io.get_output_diffs()
        for k in diffs:
            assert not np.isnan(diffs[k]).any(), "Detected NaN in output diff " + k
