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

  Program:   Visualization Toolkit
  Module:    vtkBoostBetweennessClustering.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 "vtkBoostBetweennessClustering.h"

#include "vtkBoostConnectedComponents.h"
#include "vtkBoostGraphAdapter.h"
#include "vtkDataSetAttributes.h"
#include "vtkDirectedGraph.h"
#include "vtkFloatArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMutableDirectedGraph.h"
#include "vtkMutableUndirectedGraph.h"
#include "vtkObjectFactory.h"
#include "vtkSmartPointer.h"
#include "vtkUndirectedGraph.h"

#include <boost/graph/bc_clustering.hpp>

// @note: This piece of code is modification of algorithm from boost graph
// library. This modified version allows the user to pass edge weight map
namespace boost
{
// Graph clustering based on edge betweenness centrality.
//
// This algorithm implements graph clustering based on edge
// betweenness centrality. It is an iterative algorithm, where in each
// step it compute the edge betweenness centrality (via @ref
// brandes_betweenness_centrality) and removes the edge with the
// maximum betweenness centrality. The @p done function object
// determines when the algorithm terminates (the edge found when the
// algorithm terminates will not be removed).
//
// @param g The graph on which clustering will be performed. The type
// of this parameter (@c MutableGraph) must be a model of the
// VertexListGraph, IncidenceGraph, EdgeListGraph, and Mutable Graph
// concepts.
//
// @param done The function object that indicates termination of the
// algorithm. It must be a ternary function object that accepts the
// maximum centrality, the descriptor of the edge that will be
// removed, and the graph @p g.
//
// @param edge_centrality (UTIL/out2) The property map that will store
// the betweenness centrality for each edge. When the algorithm
// terminates, it will contain the edge centralities for the
// graph. The type of this property map must model the
// ReadWritePropertyMap concept. Defaults to a @c
// iterator_property_map whose value type is
// @c Done::centrality_type and using @c get(edge_index, g) for the
// index map.
//
// @param vertex_index (IN) The property map that maps vertices to
// indices in the range @c [0, num_vertices(g)). This type of this
// property map must model the ReadablePropertyMap concept and its
// value type must be an integral type. Defaults to
// @c get(vertex_index, g).
template <typename MutableGraph, typename Done, typename EdgeCentralityMap, typename EdgeWeightMap,
  typename VertexIndexMap>
void betweenness_centrality_clustering(MutableGraph& g, Done done,
  EdgeCentralityMap edge_centrality, EdgeWeightMap edge_weight_map, VertexIndexMap vertex_index)
{
  typedef typename property_traits<EdgeCentralityMap>::value_type centrality_type;
  typedef typename graph_traits<MutableGraph>::edge_iterator edge_iterator;
  typedef typename graph_traits<MutableGraph>::edge_descriptor edge_descriptor;

  if (has_no_edges(g))
    return;

  // Function object that compares the centrality of edges
  indirect_cmp<EdgeCentralityMap, std::less<centrality_type>> cmp(edge_centrality);

  bool is_done;
  do
  {
    brandes_betweenness_centrality(g,
      edge_centrality_map(edge_centrality)
        .vertex_index_map(vertex_index)
        .weight_map(edge_weight_map));
    std::pair<edge_iterator, edge_iterator> edges_iters = edges(g);
    edge_descriptor e = *max_element(edges_iters.first, edges_iters.second, cmp);
    is_done = done(get(edge_centrality, e), e, g);
    if (!is_done)
      remove_edge(e, g);
  } while (!is_done && !has_no_edges(g));
}
}

vtkStandardNewMacro(vtkBoostBetweennessClustering);

//------------------------------------------------------------------------------
vtkBoostBetweennessClustering::vtkBoostBetweennessClustering()
  : Threshold(0)
  , UseEdgeWeightArray(false)
  , InvertEdgeWeightArray(false)
  , EdgeWeightArrayName(nullptr)
  , EdgeCentralityArrayName(nullptr)
{
  this->SetNumberOfOutputPorts(2);
}

//------------------------------------------------------------------------------
vtkBoostBetweennessClustering::~vtkBoostBetweennessClustering()
{
  this->SetEdgeWeightArrayName(nullptr);
  this->SetEdgeCentralityArrayName(nullptr);
}

//------------------------------------------------------------------------------
void vtkBoostBetweennessClustering::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);

  os << indent << "Threshold: " << this->Threshold << endl;
  os << indent << "UseEdgeWeightArray: " << this->UseEdgeWeightArray << endl;
  os << indent << "InvertEdgeWeightArray: " << this->InvertEdgeWeightArray << endl;
  (EdgeWeightArrayName)
    ? os << indent << "EdgeWeightArrayName: " << this->EdgeWeightArrayName << endl
    : os << indent << "EdgeWeightArrayName: nullptr" << endl;

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

//------------------------------------------------------------------------------
int vtkBoostBetweennessClustering::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // Helpful vars.
  bool isDirectedGraph(false);

  // Get the info objects
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  if (!inInfo)
  {
    vtkErrorMacro("Failed to get input information.");
    return 1;
  }

  vtkInformation* outInfo1 = outputVector->GetInformationObject(0);
  if (!outInfo1)
  {
    vtkErrorMacro("Failed get output1 on information first port.");
  }

  vtkInformation* outInfo2 = outputVector->GetInformationObject(1);
  if (!outInfo2)
  {
    vtkErrorMacro("Failed to get output2 information on second port.");
    return 1;
  }

  // Get the input, output1 and output2.
  vtkGraph* input = vtkGraph::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  if (!input)
  {
    vtkErrorMacro("Failed to get input graph.");
    return 1;
  }

  if (vtkDirectedGraph::SafeDownCast(input))
  {
    isDirectedGraph = true;
  }

  vtkGraph* output1 = vtkGraph::SafeDownCast(outInfo1->Get(vtkDataObject::DATA_OBJECT()));
  if (!output1)
  {
    vtkErrorMacro("Failed to get output1 graph.");
    return 1;
  }

  vtkGraph* output2 = vtkGraph::SafeDownCast(outInfo2->Get(vtkDataObject::DATA_OBJECT()));
  if (!output2)
  {
    vtkErrorMacro("Failed to get output2 graph.");
    return 1;
  }

  vtkSmartPointer<vtkFloatArray> edgeCM = vtkSmartPointer<vtkFloatArray>::New();
  if (this->EdgeCentralityArrayName)
  {
    edgeCM->SetName(this->EdgeCentralityArrayName);
  }
  else
  {
    edgeCM->SetName("edge_centrality");
  }

  boost::vtkGraphEdgePropertyMapHelper<vtkFloatArray*> helper(edgeCM);

  vtkSmartPointer<vtkDataArray> edgeWeight(nullptr);
  if (this->UseEdgeWeightArray && this->EdgeWeightArrayName)
  {
    if (!this->InvertEdgeWeightArray)
    {
      edgeWeight = input->GetEdgeData()->GetArray(this->EdgeWeightArrayName);
    }
    else
    {
      vtkDataArray* weights = input->GetEdgeData()->GetArray(this->EdgeWeightArrayName);

      if (!weights)
      {
        vtkErrorMacro(<< "Error: Edge weight array " << this->EdgeWeightArrayName
                      << " is set but not found or not a data array.\n");
        return 1;
      }

      edgeWeight.TakeReference(vtkDataArray::CreateDataArray(weights->GetDataType()));

      double range[2];
      weights->GetRange(range);

      if (weights->GetNumberOfComponents() > 1)
      {
        vtkErrorMacro("Expecting single component array.");
        return 1;
      }

      for (int i = 0; i < weights->GetDataSize(); ++i)
      {
        edgeWeight->InsertNextTuple1(range[1] - weights->GetTuple1(i));
      }
    }

    if (!edgeWeight)
    {
      vtkErrorMacro(<< "Error: Edge weight array " << this->EdgeWeightArrayName
                    << " is set but not found or not a data array.\n");
      return 1;
    }
  }

  // First compute the second output and the result will be used
  // as input for the first output.
  if (isDirectedGraph)
  {
    vtkMutableDirectedGraph* out2 = vtkMutableDirectedGraph::New();

    // Copy the data to the second output (as this algorithm most likely
    // going to removed edges (and hence modifies the graph).
    out2->DeepCopy(input);

    if (edgeWeight)
    {
      boost::vtkGraphEdgePropertyMapHelper<vtkDataArray*> helper2(edgeWeight);
      boost::betweenness_centrality_clustering(out2,
        boost::bc_clustering_threshold<double>(this->Threshold, out2, false), helper, helper2,
        boost::get(boost::vertex_index, out2));
    }
    else
    {
      boost::betweenness_centrality_clustering(
        out2, boost::bc_clustering_threshold<double>(this->Threshold, out2, false), helper);
    }
    out2->GetEdgeData()->AddArray(edgeCM);

    // Finally copy the results to the output.
    output2->ShallowCopy(out2);

    out2->Delete();
  }
  else
  {
    vtkMutableUndirectedGraph* out2 = vtkMutableUndirectedGraph::New();

    // Send the data to output2.
    out2->DeepCopy(input);

    if (edgeWeight)
    {
      boost::vtkGraphEdgePropertyMapHelper<vtkDataArray*> helper2(edgeWeight);
      boost::betweenness_centrality_clustering(out2,
        boost::bc_clustering_threshold<double>(this->Threshold, out2, false), helper, helper2,
        boost::get(boost::vertex_index, out2));
    }
    else
    {
      boost::betweenness_centrality_clustering(
        out2, boost::bc_clustering_threshold<double>(this->Threshold, out2, false), helper);
    }
    out2->GetEdgeData()->AddArray(edgeCM);

    // Finally copy the results to the output.
    output2->ShallowCopy(out2);

    out2->Delete();
  }

  // Now take care of the first output.
  vtkSmartPointer<vtkBoostConnectedComponents> bcc(
    vtkSmartPointer<vtkBoostConnectedComponents>::New());

  vtkSmartPointer<vtkGraph> output2Copy(nullptr);

  if (isDirectedGraph)
  {
    output2Copy = vtkSmartPointer<vtkDirectedGraph>::New();
  }
  else
  {
    output2Copy = vtkSmartPointer<vtkUndirectedGraph>::New();
  }

  output2Copy->ShallowCopy(output2);

  bcc->SetInputData(0, output2Copy);
  bcc->Update();

  vtkSmartPointer<vtkGraph> bccOut = bcc->GetOutput(0);

  vtkSmartPointer<vtkAbstractArray> compArray(nullptr);
  if (isDirectedGraph)
  {
    vtkSmartPointer<vtkDirectedGraph> out1(vtkSmartPointer<vtkDirectedGraph>::New());
    out1->ShallowCopy(input);

    compArray = bccOut->GetVertexData()->GetAbstractArray("component");

    if (!compArray)
    {
      vtkErrorMacro("Unable to get component array.");
      return 1;
    }

    out1->GetVertexData()->AddArray(compArray);

    // Finally copy the output to the algorithm output.
    output1->ShallowCopy(out1);
  }
  else
  {
    vtkSmartPointer<vtkUndirectedGraph> out1(vtkSmartPointer<vtkUndirectedGraph>::New());
    out1->ShallowCopy(input);

    compArray = bccOut->GetVertexData()->GetAbstractArray("component");

    if (!compArray)
    {
      vtkErrorMacro("Unable to get component array.");
      return 1;
    }

    out1->GetVertexData()->AddArray(compArray);

    // Finally copy the output to the algorithm output.
    output1->ShallowCopy(out1);
  }

  // Also add the components array to the second output.
  output2->GetVertexData()->AddArray(compArray);

  return 1;
}

//------------------------------------------------------------------------------
int vtkBoostBetweennessClustering::FillOutputPortInformation(int port, vtkInformation* info)
{
  if (port == 0 || port == 1)
  {
    info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkGraph");
  }
  return 1;
}
