/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkCollapseVerticesByArray.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/

#include "vtkCollapseVerticesByArray.h"

#include "vtkAbstractArray.h"
#include "vtkDataArray.h"
#include "vtkDataObject.h"
#include "vtkDataSetAttributes.h"
#include "vtkDoubleArray.h"
#include "vtkEdgeListIterator.h"
#include "vtkGraphEdge.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkIntArray.h"
#include "vtkMutableDirectedGraph.h"
#include "vtkObjectFactory.h"
#include "vtkOutEdgeIterator.h"
#include "vtkSmartPointer.h"
#include "vtkStringArray.h"
#include "vtkVertexListIterator.h"

#include <map>    // Using STL.
#include <vector> // Using STL.
#include <string> // Using STL.

vtkStandardNewMacro(vtkCollapseVerticesByArray);

//------------------------------------------------------------------------------
class vtkCollapseVerticesByArrayInternal
{
public:
  std::vector<std::string> AggregateEdgeArrays;
};

//------------------------------------------------------------------------------
vtkCollapseVerticesByArray::vtkCollapseVerticesByArray() :
  vtkGraphAlgorithm(),
  AllowSelfLoops(false),
  VertexArray(nullptr),
  CountEdgesCollapsed(0),
  EdgesCollapsedArray(nullptr),
  CountVerticesCollapsed(0),
  VerticesCollapsedArray(nullptr)
{
  // Setting default names.
  this->SetVerticesCollapsedArray("VerticesCollapsedCountArray");
  this->SetEdgesCollapsedArray("EdgesCollapsedCountArray");

  this->Internal = new vtkCollapseVerticesByArrayInternal();
}

// ----------------------------------------------------------------------------
vtkCollapseVerticesByArray::~vtkCollapseVerticesByArray()
{
  delete this->Internal;
  delete [] this->VertexArray;
  delete [] this->VerticesCollapsedArray;
  delete [] this->EdgesCollapsedArray;
}

// ----------------------------------------------------------------------------
void vtkCollapseVerticesByArray::PrintSelf(ostream &os, vtkIndent indent)
{
  // Base class print.
  vtkGraphAlgorithm::PrintSelf(os, indent);

  os << indent << "AllowSelfLoops: "  << this->AllowSelfLoops << endl;
  os << indent << "VertexArray: "     <<
    (this->VertexArray ? this->VertexArray : "nullptr") << endl;

  os << indent << "CountEdgesCollapsed: "  << this->CountEdgesCollapsed << endl;
  os << indent << "EdgesCollapsedArray: "
    << (this->EdgesCollapsedArray ? this->EdgesCollapsedArray : "nullptr") << endl;

  os << indent << "CountVerticesCollapsed: "  << this->CountVerticesCollapsed << endl;
  os << indent << "VerticesCollapsedArray: "
    << (this->VerticesCollapsedArray ? this->VerticesCollapsedArray : "nullptr") << endl;
}

//------------------------------------------------------------------------------
void vtkCollapseVerticesByArray::AddAggregateEdgeArray(
  const char* arrName)
{
  this->Internal->AggregateEdgeArrays.push_back(std::string(arrName));
}

//------------------------------------------------------------------------------
void vtkCollapseVerticesByArray::ClearAggregateEdgeArray()
{
  this->Internal->AggregateEdgeArrays.clear();
}

//------------------------------------------------------------------------------
int vtkCollapseVerticesByArray::RequestData(vtkInformation* vtkNotUsed(request),
                                            vtkInformationVector** inputVector,
                                            vtkInformationVector* outputVector)
{
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);

  if(!inInfo)
  {
    vtkErrorMacro("Error: nullptr input vtkInformation");
    return 0;
  }

  vtkDataObject* inObj = inInfo->Get(vtkDataObject::DATA_OBJECT());

  if(!inObj)
  {
    vtkErrorMacro(<< "Error: nullptr vtkDataObject");
    return 0;
  }

  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  if(!outInfo)
  {
    vtkErrorMacro("Error: nullptr output vtkInformation");
    return 0;
  }

  vtkDataObject* outObj = outInfo->Get(vtkDataObject::DATA_OBJECT());
  if(!outObj)
  {
    vtkErrorMacro("Error: nullptr output vtkDataObject");
    return 0;
  }

  vtkGraph* outGraph =
    this->Create(vtkGraph::SafeDownCast(inObj));
  if(outGraph)
  {
    vtkDirectedGraph::SafeDownCast(outObj)->ShallowCopy(outGraph);
    outGraph->Delete();
  }
  else
  {
    return 0;
  }

  return 1;
}

//------------------------------------------------------------------------------
int vtkCollapseVerticesByArray::FillOutputPortInformation(
    int vtkNotUsed(port), vtkInformation* info)
{
  info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkDirectedGraph");
  return 1;
}

//------------------------------------------------------------------------------
vtkGraph* vtkCollapseVerticesByArray::Create(vtkGraph* inGraph)
{
  if(!inGraph)
  {
    return nullptr;
  }

  if(!this->VertexArray)
  {
    return nullptr;
  }

  typedef vtkSmartPointer<vtkMutableDirectedGraph>
                                              vtkMutableDirectedGraphRefPtr;
  typedef vtkSmartPointer<vtkEdgeListIterator>
                                              vtkEdgeListIteratorRefPtr;
  typedef vtkSmartPointer<vtkVertexListIterator>
                                              vtkVertexListIteratorRefPtr;
  typedef vtkSmartPointer<vtkIntArray>        vtkIntArrayRefPtr;
  typedef std::pair<vtkVariant, vtkIdType> NameIdPair;


  // Create a new merged graph.
  vtkMutableDirectedGraphRefPtr outGraph(vtkMutableDirectedGraphRefPtr::New());

  vtkVertexListIteratorRefPtr itr(vtkVertexListIteratorRefPtr::New());
  itr->SetGraph(inGraph);

  // Copy the input vertex data and edge data arrays to the
  // output graph vertex and edge data.
  outGraph->GetVertexData()->CopyAllocate(inGraph->GetVertexData());
  outGraph->GetEdgeData()->CopyAllocate(inGraph->GetEdgeData());

  vtkDataSetAttributes* inVtxDsAttrs = inGraph->GetVertexData();
  vtkDataSetAttributes* inEgeDsAttrs = inGraph->GetEdgeData();

  if(!inVtxDsAttrs)
  {
    vtkErrorMacro("Error: No vertex data found on the graph.")
    return nullptr;
  }

  // Find the vertex array of interest.
  vtkAbstractArray* inVertexAOI =
    inVtxDsAttrs->GetAbstractArray(this->VertexArray);

  // Cannot proceed if vertex array of interest is not found.
  if(!inVertexAOI)
  {
    vtkErrorMacro("Error: Could not find the key vertex array.")
    return nullptr;
  }

  // Optional.
  vtkIntArrayRefPtr countEdgesCollapsedArray(nullptr);
  if(this->CountEdgesCollapsed)
  {
    countEdgesCollapsedArray = vtkIntArrayRefPtr::New();
    countEdgesCollapsedArray->SetName(this->EdgesCollapsedArray);
    countEdgesCollapsedArray->SetNumberOfComponents(1);
    outGraph->GetEdgeData()->AddArray(countEdgesCollapsedArray);
  }

  // Optional.
  vtkIntArrayRefPtr countVerticesCollapsedArray(nullptr);
  if(this->CountVerticesCollapsed)
  {
    countVerticesCollapsedArray = vtkIntArrayRefPtr::New();
    countVerticesCollapsedArray->SetName(this->VerticesCollapsedArray);
    countVerticesCollapsedArray->SetNumberOfComponents(1);
    outGraph->GetVertexData()->AddArray(countVerticesCollapsedArray);
  }


  // Arrays of interest.
  std::vector<vtkDataArray*>     inEdgeDataArraysOI;
  std::vector<vtkDataArray*>     outEdgeDataArraysOI;

  // All other arrays.
  std::vector<vtkAbstractArray*> inEdgeDataArraysAO;
  std::vector<vtkAbstractArray*> outEdgeDataArraysAO;

  std::vector<vtkAbstractArray*> inVertexDataArraysAO;
  std::vector<vtkAbstractArray*> outVertexDataArraysAO;

  //++
  // Find all the input vertex data arrays except the one set as key.
  for(int i=0; i < inVtxDsAttrs->GetNumberOfArrays(); ++i)
  {
    vtkAbstractArray* absArray = inVtxDsAttrs->GetAbstractArray(i);
    if(strcmp(absArray->GetName(), inVertexAOI->GetName()) == 0)
    {
      continue;
    }

    inVertexDataArraysAO.push_back(absArray);
  }


  for(size_t i=0; i < inVertexDataArraysAO.size(); ++i)
  {
    if(!inVertexDataArraysAO[i]->GetName())
    {
      vtkErrorMacro("Error: Name on the array is nullptr or not set.")
      return nullptr;
    }

    outVertexDataArraysAO.push_back(outGraph->GetVertexData()
      ->GetAbstractArray(inVertexDataArraysAO[i]->GetName()));
    outVertexDataArraysAO.back()->SetNumberOfTuples(inVertexDataArraysAO[i]->
                                                    GetNumberOfTuples());
  }
  //--

  //++
  // Find all the input edge data arrays of interest or not.
  for(int i=0; i < inEgeDsAttrs->GetNumberOfArrays(); ++i)
  {
    vtkAbstractArray* absArray = inEgeDsAttrs->GetAbstractArray(i);

    bool alreadyAdded(false);
    for(size_t j=0; j < this->Internal->AggregateEdgeArrays.size(); ++j)
    {
      if(strcmp(absArray->GetName(),
                this->Internal->AggregateEdgeArrays[j].c_str()) == 0)
      {
        vtkDataArray* inDataArray =
          vtkArrayDownCast<vtkDataArray>(absArray);
        if(inDataArray)
        {
          inEdgeDataArraysOI.push_back(inDataArray);
        }
        else
        {
          inEdgeDataArraysAO.push_back(absArray);
        }
        alreadyAdded = true;
        break;
      }
      else
      {
        continue;
      }
    } // End inner for.

    if(!alreadyAdded)
    {
       inEdgeDataArraysAO.push_back(absArray);
    }
  }

  // Find the corresponding empty arrays in output graph.
  vtkAbstractArray* outVertexAOI
    (outGraph->GetVertexData()->GetAbstractArray(this->VertexArray));

  // Arrays of interest.
  if(!inEdgeDataArraysOI.empty())
  {
    for(size_t i=0; i < inEdgeDataArraysOI.size(); ++i)
    {
      if(!inEdgeDataArraysOI[i]->GetName())
      {
        vtkErrorMacro("Error: Name on the array is nullptr or not set.")
        return nullptr;
      }

      vtkDataArray* outDataArray = vtkArrayDownCast<vtkDataArray>(
                                   outGraph->GetEdgeData()->GetAbstractArray(
                                   inEdgeDataArraysOI[i]->GetName()));

      outEdgeDataArraysOI.push_back(outDataArray);
      outDataArray->SetNumberOfTuples
        (inEdgeDataArraysOI[i]->GetNumberOfTuples());
    }
  }

  // All others.
  if(!inEdgeDataArraysAO.empty())
  {
    for(size_t i=0; i < inEdgeDataArraysAO.size(); ++i)
    {
      if(!inEdgeDataArraysAO[i]->GetName())
      {
        vtkErrorMacro("Error: Name on the array is nullptr or not set.")
        return nullptr;
      }

      vtkAbstractArray* outAbsArray = outGraph->GetEdgeData()->GetAbstractArray(
                                      inEdgeDataArraysAO[i]->GetName());

      outEdgeDataArraysAO.push_back(outAbsArray);
      outAbsArray->SetNumberOfTuples
        (inEdgeDataArraysAO[i]->GetNumberOfTuples());
    }
  }
  //--

  std::map<vtkVariant, vtkIdType>            myMap;
  std::map<vtkVariant, vtkIdType>::iterator  myItr;

  vtkIdType inSourceId;
  vtkIdType inTargetId;
  vtkIdType outSourceId;
  vtkIdType outTargetId;
  vtkIdType outEdgeId;

  // Iterate over all the vertices.
  while(itr->HasNext())
  {
    inSourceId = itr->Next();
    vtkVariant source = inVertexAOI->GetVariantValue(inSourceId);

    myItr = myMap.find(source);

    if(myItr != myMap.end())
    {
      // If we already have a vertex for this "source" get its id.
      outSourceId = myItr->second;
      if(this->CountVerticesCollapsed)
      {
        countVerticesCollapsedArray->SetValue(outSourceId, countVerticesCollapsedArray->GetValue(outSourceId) + 1);
      }
    }
    else
    {
      // If not then add a new vertex to the output graph.
      outSourceId = outGraph->AddVertex();
      outVertexAOI->InsertVariantValue(outSourceId, source);
      myMap.insert(NameIdPair(source, outSourceId));

      if(this->CountVerticesCollapsed)
      {
        countVerticesCollapsedArray->InsertValue(outSourceId, 1);
      }
    }


    for(size_t i=0; i < inVertexDataArraysAO.size(); ++i)
    {
      outVertexDataArraysAO[i]->SetTuple(outSourceId, inSourceId,
                                         inVertexDataArraysAO[i]);
    }
  }

  // Now itereate over all the edges in the graph.
  // Result vary dependeing on whether the input graph is
  // directed or not.
  vtkEdgeListIteratorRefPtr elItr (vtkEdgeListIteratorRefPtr::New());
  inGraph->GetEdges(elItr);

  while(elItr->HasNext())
  {
    vtkGraphEdge* edge = elItr->NextGraphEdge();
    inSourceId = edge->GetSource();
    inTargetId = edge->GetTarget();

    vtkVariant source = inVertexAOI->GetVariantValue(inSourceId);
    vtkVariant target = inVertexAOI->GetVariantValue(inTargetId);

    // Again find if there is associated vertex with this name.
    myItr = myMap.find(source);
    outSourceId = myItr->second;

    myItr = myMap.find(target);
    outTargetId = myItr->second;


    // Find if there is an edge between the out source and target.
    FindEdge(outGraph, outSourceId, outTargetId, outEdgeId);

    if((outSourceId == outTargetId) && !this->AllowSelfLoops)
    {
      continue;
    }

    //++
    if(outEdgeId == -1)
    {
      outEdgeId = outGraph->AddEdge(outSourceId, outTargetId).Id;

      // Edge does not exist. Add a new one.
      if(inEdgeDataArraysOI.empty() && inEdgeDataArraysAO.empty())
      {
        continue;
      }

      // Arrays of interest.
      for(size_t i=0; i < inEdgeDataArraysOI.size(); ++i)
      {

        outEdgeDataArraysOI[i]->SetTuple(
          outEdgeId, edge->GetId(), inEdgeDataArraysOI[i]);
      }

      // All others. Last entered will override previous ones.
      for(size_t i=0; i < inEdgeDataArraysAO.size(); ++i)
      {
        outEdgeDataArraysAO[i]->SetTuple(outEdgeId, edge->GetId(),
                                         inEdgeDataArraysAO[i]);
      }

      if(this->CountEdgesCollapsed)
      {
        countEdgesCollapsedArray->InsertValue(outEdgeId, 1);
      }
    }
    else
    {
      if(inEdgeDataArraysOI.empty() && inEdgeDataArraysAO.empty())
      {
        continue;
      }

      // Find the data on the out edge and add the data fron in edge
      // and set it on the out edge.
      for(size_t i=0; i < inEdgeDataArraysOI.size(); ++i)
      {
        double* outEdgeData = outEdgeDataArraysOI[i]->GetTuple(outEdgeId);
        double* inEdgeData  = inEdgeDataArraysOI[i]->GetTuple(edge->GetId());

        if(!outEdgeData && !inEdgeData)
        {
        continue;
        }
        for(int j=0; j < inEdgeDataArraysOI[i]->GetNumberOfComponents(); ++j)
        {
          outEdgeData[j] = outEdgeData[j] + inEdgeData[j];
        }

        outEdgeDataArraysOI[i]->SetTuple(outEdgeId, outEdgeData);
      }

      // All others. Last entered will override previous ones.
      for(size_t i=0; i < inEdgeDataArraysAO.size(); ++i)
      {
        outEdgeDataArraysAO[i]->SetTuple(outEdgeId, edge->GetId(),
                                         inEdgeDataArraysAO[i]);
      }

      if(this->CountEdgesCollapsed)
      {
        countEdgesCollapsedArray->SetValue(
          outEdgeId,
          static_cast<int>(countEdgesCollapsedArray->GetValue(outEdgeId) + 1));
      }
    }
    //--

  } // while(elItr->HasNext())

  outGraph->Register(nullptr);
  return outGraph;
}

//------------------------------------------------------------------------------
void vtkCollapseVerticesByArray::FindEdge(vtkGraph* outGraph, vtkIdType source,
                                           vtkIdType target, vtkIdType& edgeId)
{
  typedef vtkSmartPointer<vtkOutEdgeIterator> vtkOutEdgeIteratorRefPtr;
  edgeId = -1;

  if(!outGraph)
  {
    return;
  }

 vtkOutEdgeIteratorRefPtr itr (vtkOutEdgeIteratorRefPtr::New());

 outGraph->GetOutEdges(source, itr);
 while(itr->HasNext())
 {
   vtkGraphEdge* edge = itr->NextGraphEdge();
     if(edge->GetTarget() == target)
     {
     edgeId = edge->GetId();
     break;
     }
 }
}