#include <vtkMutableDirectedGraph.h>
#include <vtkNew.h>
#include <vtkTree.h>

#include <vector>

#include "vtkTreeDFSIterator.h"

int TestTreeDFSIterator(int, char*[])
{
  vtkNew<vtkMutableDirectedGraph> g;

  // Create vertices:
  vtkIdType v0 = g->AddVertex(); // Level 0
  vtkIdType v1 = g->AddVertex(); // Level 1
  vtkIdType v2 = g->AddVertex(); // Level 2
  vtkIdType v3 = g->AddVertex(); // Level 2
  vtkIdType v4 = g->AddVertex(); // Level 1
  vtkIdType v5 = g->AddVertex(); // Level 2
  vtkIdType v6 = g->AddVertex(); // Level 1
  vtkIdType v7 = g->AddVertex(); // Level 2
  vtkIdType v8 = g->AddVertex(); // Level 3

  // create a fully connected graph
  g->AddEdge(v0, v1);
  g->AddEdge(v1, v2);
  g->AddEdge(v1, v3);
  g->AddEdge(v0, v4);
  g->AddEdge(v4, v5);
  g->AddEdge(v0, v6);
  g->AddEdge(v6, v7);
  g->AddEdge(v7, v8);

  vtkNew<vtkTree> tree;
  tree->CheckedShallowCopy(g);

  std::vector<int> correctSequence;
  for (int i = 0; i <= 8; i++)
  {
    correctSequence.push_back(i);
  }

  vtkNew<vtkTreeDFSIterator> dfsIterator;
  dfsIterator->SetTree(tree);

  if (dfsIterator->GetStartVertex() != tree->GetRoot())
  {
    cout << "StartVertex is not defaulting to root" << endl;
    return EXIT_FAILURE;
  }

  // traverse the tree in a depth first fashion
  for (size_t i = 0; i < correctSequence.size(); i++)
  {
    if (!dfsIterator->HasNext())
    {
      cout << "HasNext() returned false before the end of the tree" << endl;
      return EXIT_FAILURE;
    }

    vtkIdType nextVertex = dfsIterator->Next();
    if (nextVertex != correctSequence[i])
    {
      cout << "Next vertex should be " << correctSequence[i] << " but it is " << nextVertex << endl;
      return EXIT_FAILURE;
    }
  }

  return EXIT_SUCCESS;
}
