import copy
import logging
import os

from gunpowder.caffe.net_io_wrapper import NetIoWrapper
from gunpowder.ext import caffe
from gunpowder.nodes.generic_predict import GenericPredict
from gunpowder.volume import VolumeType, Volume


logger = logging.getLogger(__name__)

class StupidPredict(object):
    '''Augments a batch with network predictions.

    Args:

        prototxt (string): Filename of the network prototxt.

        weights (string): Filename of the network weights.

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

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

        volume_specs (dict, optional): An optional dictionary of
            :class:`VolumeType` to :class:`VolumeSpec` to set the volume specs
            generated volumes (``outputs``). This is useful to set the
            ``voxel_size``, for example, if they differ from the voxel size of
            the input volumes. Only fields that are not ``None`` in the given
            :class:`VolumeSpec` will be used.

        use_gpu (int): Which GPU to use. Set to ``None`` for CPU mode.
    '''

    def __init__(
            self,
            prototxt,
            weights,
            inputs,
            outputs,
            volume_specs=None,
            use_gpu=None):

        for f in [prototxt, weights]:
            if not os.path.isfile(f):
                raise RuntimeError("%s does not exist"%f)
        self.prototxt = prototxt
        self.weights = weights
        self.inputs = inputs
        self.outputs = outputs

        if use_gpu is not None:

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

        self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST)
        self.net_io = NetIoWrapper(self.net, self.outputs.values())

    def __call__(self, input_data):
        assert isinstance(input_data, dict)
        self.net_io.set_inputs({
            input_name: data
            for input_name, data in input_data.items()
        })

        self.net.forward()
        output = self.net_io.get_outputs()
        return output
