classdef LayerGraph < deep.internal.graph.EditableLayerGraph
    % LayerGraph   A graph consisting of network layers
    %
    %   A LayerGraph object is used to describe the layout of a Directed
    %   Acyclic Graph (DAG) network. You can add layers to the graph and
    %   connect them together.    
    %
    %   LayerGraph properties:
    %       Layers              - The layers of the network
    %       Connections         - The connections between the layers
    %       InputNames          - The names of the input layers
    %       OutputNames         - The names of the output layers
    %
    %   LayerGraph methods:
    %       addLayers           - Add layers to the LayerGraph
    %       connectLayers       - Connect layers in the LayerGraph
    %       removeLayers        - Remove layers from the LayerGraph
    %       disconnectLayers    - Disconnect layers in the LayerGraph
    %       replaceLayer        - Replace a layer in the LayerGraph
    %       plot                - Plot a diagram of the LayerGraph
    %
    %   Example:
    %       Create a layer graph to describe a network with a skip layer
    %       connection:
    %
    %       layers = [
    %           imageInputLayer([28 28 1])
    %           convolution2dLayer(5, 20)
    %           reluLayer('Name', 'relu_1')
    %           convolution2dLayer(3, 20, 'Padding', 1)
    %           reluLayer()
    %           convolution2dLayer(3, 20, 'Padding', 1)
    %           reluLayer()
    %           additionLayer(2, 'Name', 'add')
    %           fullyConnectedLayer(10)
    %           softmaxLayer()
    %           classificationLayer()];
    %
    %       lgraph = layerGraph(layers);
    %
    %       lgraph = connectLayers(lgraph, 'relu_1', 'add/in2');
    %
    %   See also Layer.
    
    %   Copyright 2016-2022 The MathWorks, Inc.
    
    properties(Dependent, Hidden)
        % HiddenConnections
        HiddenConnections
    end

    properties(Dependent, SetAccess = private)
        % InputNames   Names of the input layers of the network
        InputNames
        
        % OutputNames   Names of the output layers of the network
        OutputNames
    end
    
    % Methods from EditableLayerGraph.
    methods(Access=protected)
        function nodeNames = getNodeNamesForPlot(this)
            % Helper method to generate the names of the nodes in the
            % plotted graph

            % Get the names from the NodeModel
            nodeNames = names(this.FlowGraph.NodeModel);
        end

        function layers = getLayers(this)
            % Helper method to retrieve the external layers
            
            % LayerGraph contains directly the external layers as its nodes
            layers = getnodes(this.FlowGraph.NodeModel);
        end

        function internalLayer = getInternalLayerAtIndex(~, flowgraph, layerIndex)
            % Helper method to retrieve the internal layer of the layer at
            % index layerIndex
            nodes = getnodes(flowgraph.NodeModel);
            internalLayer = nnet.internal.cnn.layer.util.ExternalInternalConverter.getInternalLayers(...
                nodes{layerIndex});
            internalLayer = internalLayer{1};
        end
    end

    methods
        function val = get.HiddenConnections(this)
            internalConnections = this.FlowGraph.Connections;
            internalConnections = sortrows(internalConnections);
            val = nnet.internal.cnn.util.internalToHiddenConnections(internalConnections);
        end
    end

    properties(Access = private)        
        % TopologicalOrder   Original indices of layers. This property
        % is empty initially but after calling toposort, this is filled
        % with indices of original layers when the layers are arranged in a
        % topologically sorted order.
        TopologicalOrder
    end

    methods(Hidden)
        function this = LayerGraph(layersOrFlowgraph, externalConnections, options)
            % Construct a LayerGraph object.
            arguments
                layersOrFlowgraph (:,1)
                externalConnections = []
                options.ValidateLayers = false
            end
            % LayerGraph   Constructor for layer graph object

            if isa(layersOrFlowgraph, "deep.internal.graph.FlowGraph")
                flowgraph = layersOrFlowgraph;
            else
                % Construct a flowgraph

                % Create the NodeModel for LayerGraph
                nodeModel = deep.internal.graph.flowgraph.LayerGraphNodes(layersOrFlowgraph, ValidateLayers=options.ValidateLayers);

                if isempty(externalConnections)
                    numLayers = height(layersOrFlowgraph);
                    internalConnections = deep.internal.graph.generateSeriesConnections(numLayers);
                else
                    internalConnections = nnet.internal.cnn.util.hiddenToInternalConnections(...
                        externalConnections );
                end
                flowgraph = deep.internal.graph.FlowGraph(nodeModel, internalConnections);
            end

            % Assign the flowgraph
            this.FlowGraph = flowgraph;
        end

        function directedGraph = extractPrivateDirectedGraph(this)
            % extractPrivateDirectedGraph   

            % Reconstruct the legacy representation of the graph, with
            % tables for nodes and edges
            layerTable = table(this.Layers, 'VariableNames', {'Layers'});
            internalConnections = sortrows(this.FlowGraph.Connections);
            layerConnections = internalConnections(:,[1 3]);

            % Multiple connections between the same two nodes must be
            % representated as a single edge with several rows for the
            % EndPort. So we need to remove the duplicate layer connections
            % and merge the EndPorts together for those values
            [uniqueLayerConnections,~,uniqueIndex] = unique(layerConnections,'rows');
            if isempty(uniqueLayerConnections)
                uniquePortConnections = [];
            else
                cnt = histcounts(uniqueIndex, max(uniqueIndex));
                uniquePortConnections = mat2cell(internalConnections(:,[2 4]), cnt, 2);
            end

            edgeTable = table(uniqueLayerConnections, uniquePortConnections, 'VariableNames', {'EndNodes','EndPorts'});
            directedGraph = digraph(edgeTable, layerTable);
        end

        function this = setTopologicalOrder(this, order)
            % Set the TopologicalOrder property.
            this.TopologicalOrder = order;
        end

        function this = toposort(this)
            [sortedIndices, sortedGraph] = toposort(this.FlowGraph, 'Order','stable');
            this.FlowGraph = sortedGraph;
            this.TopologicalOrder = sortedIndices;
        end

        function this = inferParameters(this)
            [~, analysis] = nnet.internal.cnn.layer.util.inferParameters(this);
            this = analysis.LayerGraph;
        end

        function topologicalOrder = extractTopologicalOrder(this)
            topologicalOrder = this.TopologicalOrder;
        end
    end

    methods
        function this = addLayers(this, layerArray)
            % addLayers   Add layers to layer graph
            %
            %   LGRAPH = addLayers(LGRAPH, LAYERS) adds the specified layer 
            %   array to the layer graph LGRAPH. The layer names must be
            %   unique or empty. The function connects layers in the layer
            %   array sequentially and automatically assigns unique names
            %   to any unnamed layers.
            try
                this = addLayers@deep.internal.graph.EditableLayerGraph(this, layerArray);
            catch e
                throw(e)
            end
            iWarnIfThereIsNoReturnArgument(nargout)
        end

        function this = removeLayers(this, layerNames)
            % removeLayers   Remove layers from layer graph
            %
            %   LGRAPH = removeLayers(LGRAPH, layerNames) removes the
            %   layers specified by layerNames from LGRAPH. layerNames is
            %   either the name of a layer or a cell array where each entry
            %   is a layer name. The function also removes any connections
            %   to the removed layers.
            
            %   Note that removing a layer will remove all ingoing and 
            %   outgoing connections for that layer.
            arguments
                this
                layerNames {iValidateLayerNameForRemoveLayers(layerNames)}
            end
            try
                this = removeLayers@deep.internal.graph.EditableLayerGraph(this, string(layerNames));
            catch e
                throw(e)
            end
            iWarnIfThereIsNoReturnArgument(nargout)
        end

        function this = connectLayers(this, sourceName, destinationName)
            % connectLayers   Connect layers in layer graph
            %
            %   LGRAPH = connectLayers(LGRAPH, S, D) connects a source
            %   layer S to a destination layer D in the layer graph LGRAPH,
            %   where S and D are character arrays or string scalars.
            %     - If the source layer has a single output, then S is the
            %       name of the layer. If the source layer has multiple 
            %       outputs, then S is the name of the layer followed by 
            %       the '/' character, followed by the name of the layer 
            %       output.
            %     - If the destination layer has a single input, then D is
            %       the name of the layer. If the destination layer has
            %       multiple inputs, then D is the name of the layer
            %       followed by the '/' character, followed by the name of
            %       the layer input.
            try
                this = connectLayers@deep.internal.graph.EditableLayerGraph(this, sourceName, destinationName);
            catch e
                throw(e)
            end
            iWarnIfThereIsNoReturnArgument(nargout)
        end

        function this = disconnectLayers(this, sourceName, destinationName)
            % disconnectLayers   Disconnect layers in layer graph
            %
            %   LGRAPH = disconnectLayers(LGRAPH, S, D) will disconnect a
            %   source layer S from a destination layer D in the layer
            %   graph LGRAPH, where S and D are character arrays or strings
            %   described below:
            %     - If the source layer has a single output, then S is the
            %       name of the layer. If the source layer has multiple 
            %       outputs, then S is the name of the layer followed by 
            %       the '/' character, followed by the name of the layer 
            %       output.
            %     - If the destination layer has a single input, then D is
            %       the name of the layer. If the destination layer has
            %       multiple inputs, then D is the name of the layer
            %       followed by the '/' character, followed by the name of
            %       the layer input.
            try
                this = disconnectLayers@deep.internal.graph.EditableLayerGraph(this, sourceName, destinationName);
            catch e
                throw(e)
            end
            iWarnIfThereIsNoReturnArgument(nargout)
        end

        function this = replaceLayer(this, layerName, layerArray, args)
            % replaceLayer   Replace layer in layer graph
            %
            %   LGRAPH = replaceLayer(LGRAPH, layerName, LAYERS) replaces
            %   the layer with name layerName in the layer graph LGRAPH
            %   with the specified layer array LAYERS. The layers names in 
            %   LAYERS must be unique or empty. The function connects
            %   the LAYERS to the graph and the layers in LAYERS
            %   sequentially and automatically assigns unique names
            %   to unnamed layers.
            %
            %   LGRAPH = replaceLayer(LGRAPH, layerName, LAYERS, 'ReconnectBy', MODE) 
            %   specifies the mode to use to reconnect the layers to their 
            %   neighboring layers. To connect neighboring layers to inputs 
            %   and outputs with the same names as the layer to replace, 
            %   specify MODE as 'name'. To connect neighboring layers to 
            %   inputs and outputs in the same order as they are connected 
            %   in the layer to replace, specify MODE as 'order'. The order 
            %   of the inputs and outputs is the same order as they appear 
            %   in the InputNames and OutputNames properties of the layers, 
            %   respectively. The default value is 'name'.
            arguments
                this
                layerName {iValidateLayerNameForReplaceLayers(layerName)}
                layerArray
                args.ReconnectBy (1,1) string {mustBeMember(args.ReconnectBy, {'name', 'order'})} = 'name';
            end
            try
                this = replaceLayer@deep.internal.graph.EditableLayerGraph(this, string(layerName), layerArray, ReconnectBy=args.ReconnectBy);
            catch e
                throw(e)
            end
            iWarnIfThereIsNoReturnArgument(nargout)
        end
        
        function val = get.InputNames(this)
            val = iGetInputLayerNames(this.Layers);
        end
        
        function val = get.OutputNames(this)
            val = iGetOutputLayerNames(this.Layers);
        end
    end

    methods(Static, Hidden)
        function obj = createFromLayersAndInternalConnections(externalLayers, internalConnections, topologicalOrder)
            % Construct a layer graph from the external layers and internal
            % connections
            nodeModel = deep.internal.graph.flowgraph.LayerGraphNodes(externalLayers);
            flowGraph = deep.internal.graph.FlowGraph(nodeModel, internalConnections);
            obj = nnet.cnn.LayerGraph(flowGraph);
            if nargin>2
                obj = obj.setTopologicalOrder(topologicalOrder);
            end
        end

        function this = loadobj(in)
            internalConnections = nnet.internal.cnn.util.externalToInternalConnections(in.Connections, in.Layers);
            this = nnet.cnn.LayerGraph.createFromLayersAndInternalConnections(in.Layers, internalConnections);
        end
    end
    
    methods(Hidden)
        function out = saveobj(this)
            out.Version = 1.0;
            out.Layers = this.Layers;
            out.Connections = this.Connections;
        end
    end
end

function iWarnIfThereIsNoReturnArgument(numArguments)
if numArguments == 0
    warning(message('nnet_cnn:nnet:cnn:LayerGraph:NoReturnArgumentLayerGraph'));
end
end

function iValidateLayerNameForReplaceLayers(layerName)
if ~iIsValidStringOrCharArray(layerName)
    error(message('nnet_cnn:nnet:cnn:LayerGraph:ReplaceLayerBadArgType2'));
end
end

function tf = iIsValidStringOrCharArray(x)
tf = nnet.internal.cnn.layer.paramvalidation.isValidStringOrCharArray(x);
end

function iValidateLayerNameForRemoveLayers(layerNames)

if iIsValidStringOrCharArray(layerNames) || iIsValidVectorCellOfStringOrCharArray(layerNames) 
else
    error(message('nnet_cnn:nnet:cnn:LayerGraph:InvalidInputForRemoveLayers'));
end
end

function tf = iIsValidVectorCellOfStringOrCharArray(value)
tf = false;
if isvector(value)
    if iscell(value)
        tf = all(cellfun(@iIsValidStringOrCharArray, value));
    elseif isstring(value)
        tf = ~any(ismissing(value));
    end
end
end

function inputNames = iGetInputLayerNames(layers)
internalLayers = nnet.internal.cnn.layer.util.ExternalInternalConverter.getInternalLayers(layers);
inputIndices = iGetInputLayerIndices(internalLayers);
inputNames = arrayfun(@(l)l.Name, layers(inputIndices), 'UniformOutput', false);
% Return row vector
inputNames = inputNames(:)';
end

function outputNames = iGetOutputLayerNames(layers)
internalLayers = nnet.internal.cnn.layer.util.ExternalInternalConverter.getInternalLayers(layers);
outputIndices = iGetOutputLayerIndices(internalLayers);
outputNames = arrayfun(@(l)l.Name, layers(outputIndices), 'UniformOutput', false);
% Return row vector
outputNames = outputNames(:)';
end

function inputLayerIndices = iGetInputLayerIndices(internalLayers)
isInputLayer = cellfun( @(l)isa(l, 'nnet.internal.cnn.layer.InputLayer'), internalLayers );
inputLayerIndices = find( isInputLayer );
end

function outputLayerIndices = iGetOutputLayerIndices(internalLayers)
isOutputLayer = cellfun( @(l)isa(l, 'nnet.internal.cnn.layer.OutputLayer'), internalLayers );
outputLayerIndices = find( isOutputLayer );
end
