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

  Program:   Visualization Toolkit
  Module:    vtkTreeFieldAggregator.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.

=========================================================================*/
/*-------------------------------------------------------------------------
  Copyright 2008 Sandia Corporation.
  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
  the U.S. Government retains certain rights in this software.
-------------------------------------------------------------------------*/

#include "vtkTreeFieldAggregator.h"

#include "vtkAdjacentVertexIterator.h"
#include "vtkCellData.h"
#include "vtkDataArray.h"
#include "vtkDoubleArray.h"
#include "vtkGraph.h"
#include "vtkIdList.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkIntArray.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkSmartPointer.h"
#include "vtkStringArray.h"
#include "vtkTree.h"
#include "vtkTreeDFSIterator.h"
#include "vtkVariantArray.h"

#include <cmath>

vtkStandardNewMacro(vtkTreeFieldAggregator);

vtkTreeFieldAggregator::vtkTreeFieldAggregator()
{
  this->MinValue = 0;
  this->Field = nullptr;
  this->LeafVertexUnitSize = true;
  this->LogScale = false;
}

vtkTreeFieldAggregator::~vtkTreeFieldAggregator()
{
  this->SetField(nullptr);
}

int vtkTreeFieldAggregator::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  // get the input and output
  vtkTree* input = vtkTree::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkTree* output = vtkTree::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Shallow copy the input
  output->ShallowCopy(input);

  // Check for the existence of the field to be aggregated
  if (!output->GetVertexData()->HasArray(this->Field))
  {
    this->LeafVertexUnitSize = true;
  }

  // Extract the field from the tree
  vtkAbstractArray* arr;
  if (this->LeafVertexUnitSize)
  {
    arr = vtkIntArray::New();
    arr->SetNumberOfTuples(output->GetNumberOfVertices());
    arr->SetName(this->Field);
    for (vtkIdType i = 0; i < arr->GetNumberOfTuples(); i++)
    {
      vtkArrayDownCast<vtkIntArray>(arr)->SetTuple1(i, 1);
    }
    output->GetVertexData()->AddArray(arr);
    arr->Delete();
  }
  else
  {
    vtkAbstractArray* oldArr = output->GetVertexData()->GetAbstractArray(this->Field);
    if (oldArr->GetNumberOfComponents() != 1)
    {
      vtkErrorMacro(<< "The field " << this->Field << " must have one component per tuple");
    }
    if (oldArr->IsA("vtkStringArray"))
    {
      vtkDoubleArray* doubleArr = vtkDoubleArray::New();
      doubleArr->Resize(oldArr->GetNumberOfTuples());
      for (vtkIdType i = 0; i < oldArr->GetNumberOfTuples(); i++)
      {
        doubleArr->InsertNextTuple1(vtkTreeFieldAggregator::GetDoubleValue(oldArr, i));
      }
      arr = doubleArr;
    }
    else
    {
      arr = vtkAbstractArray::CreateArray(oldArr->GetDataType());
      arr->DeepCopy(oldArr);
    }
    arr->SetName(this->Field);

    // We would like to do just perform output->GetVertexData()->RemoveArray(this->Field),
    // but because of a bug in vtkDataSetAttributes::RemoveArray(char*), we need to do it this way.
    vtkFieldData* data = vtkFieldData::SafeDownCast(output->GetVertexData());
    data->RemoveArray(this->Field);

    output->GetVertexData()->AddArray(arr);
    arr->Delete();
  }

  // Set up DFS iterator that traverses
  // children before the parent (i.e. bottom-up).
  vtkSmartPointer<vtkTreeDFSIterator> dfs = vtkSmartPointer<vtkTreeDFSIterator>::New();
  dfs->SetTree(output);
  dfs->SetMode(vtkTreeDFSIterator::FINISH);

  // Create a iterator for getting children.
  vtkSmartPointer<vtkAdjacentVertexIterator> it = vtkSmartPointer<vtkAdjacentVertexIterator>::New();

  // Iterator through the tree, aggregating child values into parent.
  while (dfs->HasNext())
  {
    vtkIdType vertex = dfs->Next();
    double value = 0;
    if (output->IsLeaf(vertex))
    {
      value = vtkTreeFieldAggregator::GetDoubleValue(arr, vertex);
      if (this->LogScale)
      {
        value = log10(value);
        if (value < this->MinValue)
        {
          value = this->MinValue;
        }
      }
    }
    else
    {
      output->GetChildren(vertex, it);
      while (it->HasNext())
      {
        value += vtkTreeFieldAggregator::GetDoubleValue(arr, it->Next());
      }
    }
    vtkTreeFieldAggregator::SetDoubleValue(arr, vertex, value);
  }

  return 1;
}

void vtkTreeFieldAggregator::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "Field: " << (this->Field ? this->Field : "(none)") << endl;
  os << indent << "LeafVertexUnitSize: " << (this->LeafVertexUnitSize ? "On" : "Off") << endl;
  os << indent << "MinValue: " << this->MinValue << endl;
  os << indent << "LogScale: " << (this->LogScale ? "On" : "Off") << endl;
}

double vtkTreeFieldAggregator::GetDoubleValue(vtkAbstractArray* arr, vtkIdType id)
{
  if (arr->IsA("vtkDataArray"))
  {
    double d = vtkArrayDownCast<vtkDataArray>(arr)->GetTuple1(id);
    if (d < this->MinValue)
    {
      return MinValue;
    }
    return d;
  }
  else if (arr->IsA("vtkVariantArray"))
  {
    vtkVariant v = vtkArrayDownCast<vtkVariantArray>(arr)->GetValue(id);
    if (!v.IsValid())
    {
      return this->MinValue;
    }
    bool ok;
    double d = v.ToDouble(&ok);
    if (!ok)
    {
      return this->MinValue;
    }
    if (d < this->MinValue)
    {
      return MinValue;
    }
    return d;
  }
  else if (arr->IsA("vtkStringArray"))
  {
    vtkVariant v(vtkArrayDownCast<vtkStringArray>(arr)->GetValue(id));
    bool ok;
    double d = v.ToDouble(&ok);
    if (!ok)
    {
      return this->MinValue;
    }
    if (d < this->MinValue)
    {
      return MinValue;
    }
    return d;
  }
  return this->MinValue;
}

void vtkTreeFieldAggregator::SetDoubleValue(vtkAbstractArray* arr, vtkIdType id, double value)
{
  if (arr->IsA("vtkDataArray"))
  {
    vtkArrayDownCast<vtkDataArray>(arr)->SetTuple1(id, value);
  }
  else if (arr->IsA("vtkVariantArray"))
  {
    vtkArrayDownCast<vtkVariantArray>(arr)->SetValue(id, vtkVariant(value));
  }
  else if (arr->IsA("vtkStringArray"))
  {
    vtkArrayDownCast<vtkStringArray>(arr)->SetValue(id, vtkVariant(value).ToString());
  }
}