# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""
Nerual net class for constructing the nets using layers and providing access
functions for net info, e.g., parameters.
"""


from .proto.model_pb2 import kTrain, kEval
import tensor
import layer
import cPickle as pickle


class FeedForwardNet(object):

    def __init__(self, loss=None, metric=None):
        self.loss = loss
        self.metric = metric
        self.layers = []
        self.src_of_layer = {}
        self.dst_of_layer = None
        self.ordered_layers = None

    def to_device(self, dev):
        for lyr in self.layers:
            lyr.to_device(dev)

    def add(self, lyr, src=None):
        """Append a layer into the layer list.

        This function will get the sample shape from the last layer to setup
        the newly added layer. For the first layer, it is setup outside.
        The calling function should ensure the correctness of the layer order.

        Args:
            lyr (Layer): the layer to be added
        """
        if src is not None:
            if isinstance(src, layer.Layer):
                assert src.has_setup is True, 'the source layer must be set up'
                self.src_of_layer[lyr.name] = [src]
            else:
                assert type(src) == list, 'the src must be a list of layers'
                self.src_of_layer[lyr.name] = src
                # print 'merge------', len(src)
        else:
            assert len(self.layers) > 0 or lyr.has_setup, \
                'Source layers are needed to set up this layer'
            if len(self.layers) > 0:
                self.src_of_layer[lyr.name] = [self.layers[-1]]
            else:
                self.src_of_layer[lyr.name] = []
        if lyr.has_setup is False:
            # print shape
            in_shape = self.src_of_layer[lyr.name][0].get_output_sample_shape()
            lyr.setup(in_shape)
            print lyr.name, lyr.get_output_sample_shape()
        self.layers.append(lyr)
        return lyr

    def param_values(self):
        values = []
        layers = self.layers
        if self.ordered_layers is not None:
            layers = self.ordered_layers
        for lyr in layers:
            values.extend(lyr.param_values())
        return values

    def param_specs(self):
        specs = []
        layers = self.layers
        if self.ordered_layers is not None:
            layers = self.ordered_layers
        for lyr in layers:
            specs.extend(lyr.param_specs)
        return specs

    def param_names(self):
        return [spec.name for spec in self.param_specs()]

    def train(self, x, y):
        out = self.forward(kTrain, x)
        l = self.loss.forward(kTrain, out, y)
        if self.metric is not None:
            m = self.metric.evaluate(out, y)
        return self.backward(), (l.l1(), m)

    def evaluate(self, x, y):
        """Evaluate the loss and metric of the given data"""
        out = self.forward(kEval, x)
        l = None
        m = None
        assert self.loss is not None or self.metric is not None,\
            'Cannot do evaluation, as neither loss nor metic is set'
        if self.loss is not None:
            l = self.loss.evaluate(kEval, out, y)
        if self.metric is not None:
            m = self.metric.evaluate(out, y)
        return l, m

    def predict(self, x):
        xx = self.forward(kEval, x)
        return tensor.softmax(xx)

    def topo_sort(self, cur, src_of_layer, visited=None, order=None):
        if visited is None:
            visited = {}
            for name in src_of_layer.keys():
                visited[name] = False
            order = []
        srcs = src_of_layer[cur.name]
        for src in srcs:
            if visited[src.name] is False:
                visited[src.name] = True
                self.topo_sort(src, src_of_layer, visited, order)
        order.append(cur)
        visited[cur.name] = True
        return order

    def forward(self, flag, x):
        # print x.l1()
        if self.ordered_layers is None:
            self.ordered_layers = self.topo_sort(self.layers[-1],
                                                 self.src_of_layer)
        inputs = [x]
        output_of_layer = {}
        for cur in self.ordered_layers:
            srcs = self.src_of_layer[cur.name]
            disp_src = cur.name + '<--'
            for src in srcs:
                outs = output_of_layer[src.name]
                if type(outs) == list:
                    inputs.append(outs[0])
                else:
                    inputs.append(outs)
                disp_src += '+' + src.name
                # del output_of_layer[src.name]
            # print disp_src
            if len(inputs) == 1:
                inputs = inputs[0]
            output_of_layer[cur.name] = cur.forward(flag, inputs)
            inputs = []
            # print lyr.name, x.l1()
        # print output_of_layer
        return output_of_layer[self.ordered_layers[-1].name]

    def backward(self):
        if self.dst_of_layer is None:
            self.dst_of_layer = {}
            for cur in self.layers:
                self.dst_of_layer[cur.name] = []
            for cur in self.ordered_layers[1:]:
                srcs = self.src_of_layer[cur.name]
                for src in srcs:
                    self.dst_of_layer[src.name].append(cur)
        grad = self.loss.backward()
        if len(grad.shape) > 1:
            grad /= grad.shape[0]  # average across the batch
        # print 'grad', grad.l1()
        grads = [grad]
        output_of_layer = {}
        pgrads = []
        for cur in reversed(self.ordered_layers):
            for dst in self.dst_of_layer[cur.name]:
                outputs = output_of_layer[dst.name]
                if type(outputs) == list:
                    grads.append(outputs[0])
                else:
                    grads.append(outputs)
                # del output_of_layer[dst.name]
            if len(grads) == 1:
                grads = grads[0]
            outs, _pgrads = cur.backward(kTrain, grads)
            pgrads.append(_pgrads)
            output_of_layer[cur.name] = outs
            grads = []

        ret = []
        for pgrad in reversed(pgrads):
            ret.extend(pgrad)
        return ret

    def save(self, f):
        """Save model parameters using cpickle"""
        params = {}
        for (specs, val) in zip(self.param_specs(), self.param_values()):
            val.to_host()
            params[specs.name] = tensor.to_numpy(val)
        with open(f, 'wb') as fd:
            pickle.dump(params, fd)

    def load(self, f):
        """Load model parameters using cpickle"""
        with open(f, 'rb') as fd:
            params = pickle.load(fd)
        for (specs, val) in zip(self.param_specs(), self.param_values()):
            val.copy_from_numpy(params[specs.name])