--- Implements a concecpt that stores a collection of layers.

--- The class for storing a collection of layers.
-- @type nerv.LayerRepo

local LayerRepo = nerv.class("nerv.LayerRepo")

--- The constructor.
-- @param layer_spec the *layer specification*, a declarative way of creating layers in the collection. The layer specification is structured as follow:
--
--    {
--        [<layer_typename1>] =
--        {
--            <layer_id1> = <layer_conf1>,
--            <layer_id2> = <layer_conf2>,
--            <layer_id3> = <layer_conf3>,
--            ...
--        },
--        [<layer_typename2>] =
--        {
--            ...
--        },
--        ...
--    }
-- To be short, the specification is a table containing pairs of a layer type
-- name string (such as `"nerv.AffineLayer"`) and a table which maps layer
-- identifiers to `layer_conf`. For `layer_conf`, see `nerv.Layer.__init` and
-- the `__init` doc for an individual layer type.
--
-- Here is an example:
--
--    {
--        ["nerv.AffineLayer"] =
--        {
--            affine0 = {dim_in = {429}, dim_out = {2048},
--                        params = {ltp = "affine0_ltp", bp = "affine0_bp"}},
--            affine1 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine1_ltp", bp = "affine1_bp"}},
--            affine2 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine2_ltp", bp = "affine2_bp"}},
--            affine3 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine3_ltp", bp = "affine3_bp"}},
--            affine4 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine4_ltp", bp = "affine4_bp"}},
--            affine5 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine5_ltp", bp = "affine5_bp"}},
--            affine6 = {dim_in = {2048}, dim_out = {2048},
--                        params = {ltp = "affine6_ltp", bp = "affine6_bp"}},
--            affine7 = {dim_in = {2048}, dim_out = {3001},
--                        params = {ltp = "affine7_ltp", bp = "affine7_bp"}}
--        },
--        ["nerv.SigmoidLayer"] =
--        {
--            sigmoid0 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid1 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid2 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid3 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid4 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid5 = {dim_in = {2048}, dim_out = {2048}},
--            sigmoid6 = {dim_in = {2048}, dim_out = {2048}}
--        },
--        ["nerv.SoftmaxCELayer"] = -- softmax + ce criterion layer for finetune output
--        {
--            ce_crit = {dim_in = {3001, 1}, dim_out = {1}, compressed = true}
--        },
--        ["nerv.SoftmaxLayer"] = -- softmax for decode output
--        {
--            softmax = {dim_in = {3001}, dim_out = {3001}}
--        }
--    }
-- @param param_repo the default parameter repo to be used for binding parameters, if one layer
-- does not specify `pr` in its layer config `layer_conf`
-- @param global_conf a table describing the computation state and providing
-- with some global settings

function LayerRepo:__init(layer_spec, param_repo, global_conf)
    self.layers = {}
    self:add_layers(layer_spec, param_repo, global_conf)
end

--- Add more layers to the collection (repo).
-- @param layer_spec the *layer specification*, a declarative way of creating layers in the collection.
-- @param param_repo the default parameter repo to be used for binding parameters, if one layer
-- does not specify `pr` in its layer config `layer_conf`
-- @param global_conf a table describing the computation state and providing
-- with some global settings
--
-- Here is an example for adding graph layers based on the previous example:
--    layer_repo:add_layers(
--    {
--        ["nerv.GraphLayer"] =
--        {
--            global_transf = {
--                dim_in = {429}, dim_out = {429},
--                layer_repo = layer_repo,
--                connections = {
--                    {"<input>[1]", "blayer1[1]", 0},
--                    {"blayer1[1]", "wlayer1[1]", 0},
--                    {"wlayer1[1]", "blayer2[1]", 0},
--                    {"blayer2[1]", "wlayer2[1]", 0},
--                    {"wlayer2[1]", "<output>[1]", 0}
--                }
--            },
--            main = {
--                dim_in = {429}, dim_out = {3001},
--                layer_repo = layer_repo,
--                connections = {
--                    {"<input>[1]", "affine0[1]", 0},
--                    {"affine0[1]", "sigmoid0[1]", 0},
--                    {"sigmoid0[1]", "affine1[1]", 0},
--                    {"affine1[1]", "sigmoid1[1]", 0},
--                    {"sigmoid1[1]", "affine2[1]", 0},
--                    {"affine2[1]", "sigmoid2[1]", 0},
--                    {"sigmoid2[1]", "affine3[1]", 0},
--                    {"affine3[1]", "sigmoid3[1]", 0},
--                    {"sigmoid3[1]", "affine4[1]", 0},
--                    {"affine4[1]", "sigmoid4[1]", 0},
--                    {"sigmoid4[1]", "affine5[1]", 0},
--                    {"affine5[1]", "sigmoid5[1]", 0},
--                    {"sigmoid5[1]", "affine6[1]", 0},
--                    {"affine6[1]", "sigmoid6[1]", 0},
--                    {"sigmoid6[1]", "affine7[1]", 0},
--                    {"affine7[1]", "<output>[1]", 0}
--                }
--            }
--        }
--    }, param_repo, gconf)
--
-- To fully understand the example, please check the doc for `nerv.GraphLayer`,
-- and notice that `layer_repo` itself is passed to the graph layer config because
-- primitive layers such as `"affine0"` have been created by the layer
-- specification during the construction (see the example in `__init`).

function LayerRepo:add_layers(layer_spec, param_repo, global_conf)
    local layers = self.layers
    for ltype, llist in pairs(layer_spec) do
        local layer_type = nerv.get_type(ltype)
        if layer_type == nil then
            nerv.error('layer type `%s` not found', ltype)
        end
        for id, lconf in pairs(llist) do
            if type(lconf) ~= "table" then
                nerv.error("layer config table is need")
            end
            if lconf.pr == nil then
                lconf.pr = param_repo
            end
            if layers[id] ~= nil then
                nerv.error("a layer with id %s already exists", id)
            end
            nerv.info("create layer: %s", id)
            layers[id] = layer_type(id, global_conf, lconf)
        end
    end
end

--- Rebind the parameters.
-- @param param_repo the new parameter repo used for parameter rebinding

function LayerRepo:rebind(param_repo)
    if self.__rebinding then
        return
    end
    self.__rebinding = true
    for _, layer in pairs(self.layers) do
        if not layer.__already_rebound then
            layer.__already_rebound = true
            layer.lconf.pr = param_repo
            layer:bind_params()
        end
    end
    for _, layer in pairs(self.layers) do
        layer.__already_rebound = false
    end
    self.__rebinding = false
end

--- Get a layer from the collection (repo) by its identifier.
-- @param lid the layer id
function LayerRepo:get_layer(lid)
    local layer = self.layers[lid]
    if layer == nil then
        nerv.error("layer with id %s not found", lid)
    end
    return layer
end