"""
nodes specific to dealing with hyperparameters
"""

# FIXME add HyperparameterNode to this file

from __future__ import division, absolute_import
from __future__ import print_function, unicode_literals

import theano
import theano.tensor as T

from .. import utils
from .. import core

fX = theano.config.floatX


@core.register_node("variable_hyperparameter")
class VariableHyperparameterNode(core.Wrapper1NodeImpl):

    """
    provides a theano variable for the given hyperparameter to its children
    """

    hyperparameter_names = ("hyperparameter", "dtype", "shape")

    def init_state(self, network):
        # perform default
        super(VariableHyperparameterNode, self).init_state(network)
        # also create variable
        dtype = network.find_hyperparameter(["dtype"], fX)
        shape = network.find_hyperparameter(["shape"], ())
        network.create_variable(
            "hyperparameter",
            shape=shape,
            dtype=dtype,
            is_shared=False,
            tags={"hyperparameter", "monitor"},
        )

    def get_hyperparameter(self, network, name):
        if name == "hyperparameter":
            # prevent infinite loop
            return super(VariableHyperparameterNode, self).get_hyperparameter(
                network, name)

        rel_network = network[self.name]
        hyperparameter = rel_network.find_hyperparameter(["hyperparameter"])
        if name == hyperparameter:
            return rel_network.get_variable("hyperparameter").variable
        else:
            return super(VariableHyperparameterNode, self).get_hyperparameter(
                network, name)


@core.register_node("shared_hyperparameter")
class SharedHyperparameterNode(core.Wrapper1NodeImpl):

    """
    provides a theano variable for the given hyperparameter to its children,
    but also keeps the last value as a shared variable

    use case:
    - loading a network that was saved while the hyperparameter was changing
    """
    # FIXME a lot copy pasted from VariableHyperparameterNode

    hyperparameter_names = ("hyperparameter", "dtype", "shape")

    def init_state(self, network):
        # perform default
        super(SharedHyperparameterNode, self).init_state(network)
        # also create variable
        dtype = network.find_hyperparameter(["dtype"], fX)
        shape = network.find_hyperparameter(["shape"], ())
        # TODO take in optional initial value instead of dtype/shape
        raw_vw = network.create_variable(
            "raw_hyperparameter",
            shape=shape,
            dtype=dtype,
            is_shared=True,
            tags={"state"},
        )
        # create a copy of the shared variable, so we can use this to
        # update the variable
        network.copy_variable(
            name="hyperparameter",
            previous_variable=raw_vw,
            tags={"hyperparameter", "monitor"},
        )

    def new_update_deltas(self, network):
        raw = network.get_variable("raw_hyperparameter").variable
        actual = network.get_variable("hyperparameter").variable
        return core.UpdateDeltas.from_updates({raw: actual})

    def get_hyperparameter(self, network, name):
        if name == "hyperparameter":
            # prevent infinite loop
            return super(SharedHyperparameterNode, self).get_hyperparameter(
                network, name)

        rel_network = network[self.name]
        hyperparameter = rel_network.find_hyperparameter(["hyperparameter"])
        if name == hyperparameter:
            return rel_network.get_variable("hyperparameter").variable
        else:
            return super(SharedHyperparameterNode, self).get_hyperparameter(
                network, name)


@core.register_node("output_hyperparameter")
class OutputHyperparameterNode(core.NodeImpl):

    """
    outputs a hyperparameter of the node
    """

    hyperparameter_names = ("hyperparameter", "shape")
    input_keys = ()

    def compute_output(self, network):
        hyperparameter_name = network.find_hyperparameter(["hyperparameter"])
        # TODO add default hyperparameter
        res = network.find_hyperparameter([hyperparameter_name])
        if utils.is_number(res):
            var = T.constant(res)
            shape = ()
        elif utils.is_ndarray(res):
            var = T.constant(res)
            shape = res.shape
        elif utils.is_shared_variable(res):
            var = res
            shape = res.get_value().shape
        elif utils.is_nonshared_variable(res):
            var = res
            if res.ndim == 0:
                shape = ()
            else:
                shape = network.find_hyperparameter(["shape"])
        else:
            raise ValueError("Unknown hyperparameter type of %s" % res)

        network.create_variable(
            "default",
            variable=var,
            shape=shape,
            tags={"output"},
        )

@@ -200,16 +200,29 @@
 terals%0A%0A
+import toolz%0A
 import t
@@ -2107,34 +2107,43 @@
 %22dtype%22, %22shape%22
+, %22inits%22
 )%0A
-
 %0A    def init_st
@@ -2395,32 +2395,137 @@
 (%5B%22shape%22%5D, ())%0A
+        inits = list(toolz.concat(network.find_hyperparameters(%0A            %5B%22inits%22%5D,%0A            %5B%5D)))%0A
         # TODO t
@@ -2695,32 +2695,32 @@
    dtype=dtype,%0A
-
             is_s
@@ -2731,16 +2731,41 @@
 d=True,%0A
+            inits=inits,%0A
         

add inits to shared hyperparameter