// nnet3/nnet-graph.cc

// Copyright      2015  Johns Hopkins University (author: Daniel Povey)
//                2015  Guoguo Chen

// See ../../COPYING for clarification regarding multiple authors
//
// Licensed 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include <iterator>
#include <sstream>
#include "nnet3/nnet-graph.h"

namespace kaldi {
namespace nnet3 {



void NnetToDirectedGraph(const Nnet &nnet,
                         std::vector<std::vector<int32> > *graph) {
  graph->clear();
  int32 num_nodes = nnet.NumNodes();
  graph->resize(num_nodes);
  for (int32 n = 0; n < num_nodes; n++) {
    const NetworkNode &node = nnet.GetNode(n);
    // handle dependencies of this node.
    std::vector<int32> node_dependencies;
    switch (node.node_type) {
      case kInput:
        break;  // no node dependencies.
      case kDescriptor:
        node.descriptor.GetNodeDependencies(&node_dependencies);
        break;
      case kComponent:
        node_dependencies.push_back(n - 1);
        break;
      case kDimRange:
        node_dependencies.push_back(node.u.node_index);
        break;
      default:
        KALDI_ERR << "Invalid node type";
    }
    SortAndUniq(&node_dependencies);
    for (size_t i = 0; i < node_dependencies.size(); i++) {
      int32 dep_n = node_dependencies[i];
      KALDI_ASSERT(dep_n >= 0 && dep_n < num_nodes);
      (*graph)[dep_n].push_back(n);
    }
  }
}

void ComputeGraphTranspose(const std::vector<std::vector<int32> > &graph,
                           std::vector<std::vector<int32> > *graph_transpose) {
  int32 size = graph.size();
  graph_transpose->clear();
  graph_transpose->resize(size);
  for (int32 n = 0; n < size; n++) {
    const std::vector<int32> &nodes = graph[n];
    std::vector<int32>::const_iterator iter = nodes.begin(), end = nodes.end();
    for (; iter != end; ++iter) {
      int32 dest = *iter;
      (*graph_transpose)[dest].push_back(n);
    }
  }
}

struct TarjanNode {
  int32 index;
  int32 lowlink;
  bool on_stack;
  TarjanNode() : index(-1), lowlink(-1), on_stack(false) {}
};

void TarjanSccRecursive(int32 node,
                        const std::vector<std::vector<int32> > &graph,
                        int32 *global_index,
                        std::vector<TarjanNode> *tarjan_nodes,
                        std::vector<int32> *tarjan_stack,
                        std::vector<std::vector<int32> > *sccs) {
  KALDI_ASSERT(sccs != NULL);
  KALDI_ASSERT(tarjan_nodes != NULL);
  KALDI_ASSERT(tarjan_stack != NULL);
  KALDI_ASSERT(global_index != NULL);
  KALDI_ASSERT(node >= 0 && node < graph.size());

  // Initializes the current Tarjan node.
  (*tarjan_nodes)[node].index = *global_index;
  (*tarjan_nodes)[node].lowlink = *global_index;
  *global_index += 1;
  (*tarjan_nodes)[node].on_stack = true;
  tarjan_stack->push_back(node);

  // DFS from the current node.
  for (int32 i = 0; i < graph[node].size(); ++i) {
    int32 next = graph[node][i];

    if ((*tarjan_nodes)[next].index == -1) {
      // First time we see this node.
      TarjanSccRecursive(next, graph,
                         global_index, tarjan_nodes, tarjan_stack, sccs);
      (*tarjan_nodes)[node].lowlink = std::min((*tarjan_nodes)[node].lowlink,
                                               (*tarjan_nodes)[next].lowlink);
    } else if ((*tarjan_nodes)[next].on_stack) {
      // Next node is on the stack -- back edge. We can't use the lowlink of
      // next node, because that may point to the index of the root, while the
      // current node can't be the root.
      (*tarjan_nodes)[node].lowlink = std::min((*tarjan_nodes)[node].lowlink,
                                               (*tarjan_nodes)[next].index);
    }
  }

  // Output SCC.
  if ((*tarjan_nodes)[node].index == (*tarjan_nodes)[node].lowlink) {
    std::vector<int32> scc;
    int32 pop_node;
    do {
      pop_node = tarjan_stack->back();
      tarjan_stack->pop_back();
      (*tarjan_nodes)[pop_node].on_stack = false;
      scc.push_back(pop_node);
    } while (pop_node != node);
    KALDI_ASSERT(pop_node == node);
    sccs->push_back(scc);
  }
}

void FindSccsTarjan(const std::vector<std::vector<int32> > &graph,
                    std::vector<std::vector<int32> > *sccs) {
  KALDI_ASSERT(sccs != NULL);

  // Initialization.
  std::vector<TarjanNode> tarjan_nodes(graph.size());
  std::vector<int32> tarjan_stack;
  int32 global_index = 0;

  // Calls the recursive function.
  for (int32 n = 0; n < graph.size(); ++n) {
    if (tarjan_nodes[n].index == -1) {
      TarjanSccRecursive(n, graph,
                         &global_index, &tarjan_nodes, &tarjan_stack, sccs);
    }
  }
}

void FindSccs(const std::vector<std::vector<int32> > &graph,
              std::vector<std::vector<int32> > *sccs) {
  // Internally we call Tarjan's SCC algorithm, as it only requires one DFS. We
  // can change this to other methods later on if necessary.
  KALDI_ASSERT(sccs != NULL);
  FindSccsTarjan(graph, sccs);
}

void MakeSccGraph(const std::vector<std::vector<int32> > &graph,
                  const std::vector<std::vector<int32> > &sccs,
                  std::vector<std::vector<int32> > *scc_graph) {
  KALDI_ASSERT(scc_graph != NULL);
  scc_graph->clear();
  scc_graph->resize(sccs.size());

  // Hash map from node to SCC index.
  std::vector<int32> node_to_scc_index(graph.size());
  for (int32 i = 0; i < sccs.size(); ++i) {
    for (int32 j = 0; j < sccs[i].size(); ++j) {
      KALDI_ASSERT(sccs[i][j] >= 0 && sccs[i][j] < graph.size());
      node_to_scc_index[sccs[i][j]] = i;
    }
  }

  // Builds graph.
  for (int32 i = 0; i < sccs.size(); ++i) {
    for (int32 j = 0; j < sccs[i].size(); ++j) {
      int32 node = sccs[i][j];
      KALDI_ASSERT(node >= 0 && node < graph.size());
      for (int32 k = 0; k < graph[node].size(); ++k) {
        if (node_to_scc_index[graph[node][k]] != i) { // Exclucding self.
          (*scc_graph)[i].push_back(node_to_scc_index[graph[node][k]]);
        }
      }
    }
    // If necessary, we can use a hash maps to avoid this sorting.
    SortAndUniq(&((*scc_graph)[i]));
  }
}

void ComputeTopSortOrderRecursive(int32 node,
                                  const std::vector<std::vector<int32> > &graph,
                                  std::vector<bool> *cycle_detector,
                                  std::vector<bool> *is_visited,
                                  std::vector<int32> *reversed_orders) {
  KALDI_ASSERT(node >= 0 && node < graph.size());
  KALDI_ASSERT(cycle_detector != NULL);
  KALDI_ASSERT(is_visited != NULL);
  KALDI_ASSERT(reversed_orders != NULL);
  if ((*cycle_detector)[node]) {
    KALDI_ERR << "Cycle detected when computing the topological sorting order";
  }

  if (!(*is_visited)[node]) {
    (*cycle_detector)[node] = true;
    for (int32 i = 0; i < graph[node].size(); ++i) {
      ComputeTopSortOrderRecursive(graph[node][i], graph,
                                   cycle_detector, is_visited, reversed_orders);
    }
    (*cycle_detector)[node] = false;
    (*is_visited)[node] = true;
    // At this point we have added all the children to <reversed_orders>, so we
    // can add the current now.
    reversed_orders->push_back(node);
  }
}

void ComputeTopSortOrder(const std::vector<std::vector<int32> > &graph,
                         std::vector<int32> *node_to_order) {
  // Internally we use DFS, but we only put the node to <node_to_order> when all
  // its parents have been visited.
  KALDI_ASSERT(node_to_order != NULL);
  node_to_order->resize(graph.size());

  std::vector<bool> cycle_detector(graph.size(), false);
  std::vector<bool> is_visited(graph.size(), false);

  std::vector<int32> reversed_orders;
  for(int32 i = 0; i < graph.size(); ++i) {
    if (!is_visited[i]) {
      ComputeTopSortOrderRecursive(i, graph, &cycle_detector,
                                   &is_visited, &reversed_orders);
    }
  }

  KALDI_ASSERT(node_to_order->size() == reversed_orders.size());
  for (int32 i = 0; i < reversed_orders.size(); ++i) {
    KALDI_ASSERT(reversed_orders[i] >= 0 && reversed_orders[i] < graph.size());
    (*node_to_order)[reversed_orders[i]] = graph.size() - i - 1;
  }
}

std::string PrintGraphToString(const std::vector<std::vector<int32> > &graph) {
  std::ostringstream os;
  int32 num_nodes = graph.size();
  for (int32 i = 0; i < num_nodes; i++) {
    os << i << " -> (";
    const std::vector<int32> &vec = graph[i];
    int32 size = vec.size();
    for (int32 j = 0; j < size; j++) {
      os << vec[j];
      if (j + 1 < size) os << ",";
    }
    os << ")";
    if (i + 1 < num_nodes) os << "; ";
  }
  return os.str();
}

void ComputeNnetComputationEpochs(const Nnet &nnet,
                                  std::vector<int32> *node_to_epoch) {
  KALDI_ASSERT(node_to_epoch != NULL);

  std::vector<std::vector<int32> > graph;
  NnetToDirectedGraph(nnet, &graph);
  KALDI_VLOG(6) << "graph is: " << PrintGraphToString(graph);

  std::vector<std::vector<int32> > sccs;
  FindSccs(graph, &sccs);

  std::vector<std::vector<int32> > scc_graph;
  MakeSccGraph(graph, sccs, &scc_graph);
  KALDI_VLOG(6) << "scc graph is: " << PrintGraphToString(scc_graph);

  std::vector<int32> scc_node_to_epoch;
  ComputeTopSortOrder(scc_graph, &scc_node_to_epoch);
  if (GetVerboseLevel() >= 6) {
    std::ostringstream os;
    for (int32 i = 0; i < scc_node_to_epoch.size(); i++)
      os << scc_node_to_epoch[i] << ", ";
    KALDI_VLOG(6) << "scc_node_to_epoch is: " << os.str();
  }

  node_to_epoch->clear();
  node_to_epoch->resize(graph.size());
  for (int32 i = 0; i < sccs.size(); ++i) {
    for (int32 j = 0; j < sccs[i].size(); ++j) {
      int32 node = sccs[i][j];
      KALDI_ASSERT(node >= 0 && node < graph.size());
      (*node_to_epoch)[node] = scc_node_to_epoch[i];
    }
  }
}

bool GraphHasCycles(const std::vector<std::vector<int32> > &graph) {
  std::vector<std::vector<int32> > sccs;
  FindSccs(graph, &sccs);
  for (size_t i = 0; i < sccs.size(); i++) {
    if (sccs[i].size() > 1)
      return true;
  }
  // the next code checks for links from a state to itself.
  int32 num_nodes = graph.size();
  for (size_t i = 0; i < num_nodes; i++)
    for (std::vector<int32>::const_iterator iter = graph[i].begin(),
             end = graph[i].end(); iter != end; ++iter)
      if (*iter == i) return true;
  return false;
}

} // namespace nnet3
} // namespace kaldi