//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt 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 "vtkCMBStreamTracer.h"

#include "vtkCellArray.h"
#include "vtkCellData.h"
#include "vtkCompositeDataIterator.h"
#include "vtkCompositeDataPipeline.h"
#include "vtkDataSetAttributes.h"
#include "vtkDoubleArray.h"
#include "vtkExecutive.h"
#include "vtkGenericCell.h"
#include "vtkIdList.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkIntArray.h"
#include "vtkFloatArray.h"
#include "vtkDoubleArray.h"
#include "vtkCharArray.h"
#include "vtkMath.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPointSet.h"
#include "vtkPolyData.h"
#include "vtkPolyLine.h"
#include "vtkSmartPointer.h"
#include "vtkOutputWindow.h"
#include "vtkCMBInitialValueProblemSolver.h"
#include "vtkAbstractInterpolatedVelocityField.h"

#include "vtkCMBInitialValueProblemSolver.h"

#include <functional>
#include <algorithm>

//---------------------------------------------------------------------------
vtkStandardNewMacro(vtkCMBStreamTracer);

//---------------------------------------------------------------------------
vtkCMBStreamTracer::vtkCMBStreamTracer()
{
  this->SetInterpolatorTypeToCellLocator();
  this->SetComputeVorticity(false);

  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::New();
  this->SetIntegrator(solver);
  solver->Delete();

  //this->SetIntegratorTypeToRungeKutta45();
}

//---------------------------------------------------------------------------
vtkCMBStreamTracer::~vtkCMBStreamTracer()
{
}

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

  if (!this->SetupOutput(inInfo, outInfo))
    {
    return 0;
    }

  vtkInformation *sourceInfo = inputVector[1]->GetInformationObject(0);
  vtkDataSet *source = 0;
  if (sourceInfo)
    {
    source = vtkDataSet::SafeDownCast(
      sourceInfo->Get(vtkDataObject::DATA_OBJECT()));
    }
  vtkPolyData *output = vtkPolyData::SafeDownCast(
    outInfo->Get(vtkDataObject::DATA_OBJECT()));

  vtkDataArray* seeds = 0;
  vtkIdList* seedIds = 0;
  vtkIntArray* integrationDirections = 0;
  this->InitializeSeeds(seeds, seedIds, integrationDirections, source);

  if (seeds)
    {
    vtkAbstractInterpolatedVelocityField * func;
    int maxCellSize = 0;
    if (this->CheckInputs(func, &maxCellSize) != VTK_OK)
      {
      vtkDebugMacro("No appropriate inputs have been found. Can not execute.");
      func->Delete();
      seeds->Delete();
      integrationDirections->Delete();
      seedIds->Delete();
      this->InputData->UnRegister(this);
      return 1;
      }

  // initialize test locations

    vtkCompositeDataIterator* iter = this->InputData->NewIterator();
    vtkSmartPointer<vtkCompositeDataIterator> iterP(iter);
    iter->Delete();

    iterP->GoToFirstItem();
    vtkDataSet* input0 = 0;
    if (!iterP->IsDoneWithTraversal())
      {
      input0 = vtkDataSet::SafeDownCast(iterP->GetCurrentDataObject());
      }
    vtkDataArray *vectors = this->GetInputArrayToProcess(0,input0);
    if (vectors)
      {
      const char *vecName = vectors->GetName();
      double propagation = 0;
      vtkIdType numSteps = 0;
      double lastPoint[3];
      this->Integrate(input0, output,
        seeds, seedIds,
        integrationDirections,
        lastPoint, func,
        maxCellSize, vecName,
        propagation, numSteps);
      }
    func->Delete();
    seeds->Delete();
    }

  integrationDirections->Delete();
  seedIds->Delete();

  this->InputData->UnRegister(this);
  return 1;
}

//---------------------------------------------------------------------------
void vtkCMBStreamTracer::Integrate(vtkDataSet *input0,
                                vtkPolyData* output,
                                vtkDataArray* seedSource,
                                vtkIdList* seedIds,
                                vtkIntArray* integrationDirections,
                                double lastPoint[3],
                                vtkAbstractInterpolatedVelocityField* func,
                                int maxCellSize,
                                const char *vecName,
                                double& inPropagation,
                                vtkIdType& inNumSteps)
{
  int i;
  vtkIdType numLines = seedIds->GetNumberOfIds();
  double propagation = inPropagation;
  vtkIdType numSteps = inNumSteps;

  // Useful pointers
  vtkDataSetAttributes* outputPD = output->GetPointData();
  vtkDataSetAttributes* outputCD = output->GetCellData();
  vtkPointData* inputPD;
  vtkDataSet* input;
  vtkDataArray* inVectors;

  int direction=1;

  double* weights = 0;
  if ( maxCellSize > 0 )
    {
    weights = new double[maxCellSize];
    }

  if (this->GetIntegrator() == 0)
    {
    vtkErrorMacro("No integrator is specified.");
    return;
    }

  // Used in GetCell()
  vtkGenericCell* cell = vtkGenericCell::New();

  // Create a new integrator, the type is the same as Integrator
  vtkInitialValueProblemSolver* integrator =
    vtkInitialValueProblemSolver::SafeDownCast(
    this->GetIntegrator())->NewInstance();
  integrator->SetFunctionSet(func);

  // Since we do not know what the total number of points
  // will be, we do not allocate any. This is important for
  // cases where a lot of streamers are used at once. If we
  // were to allocate any points here, potentially, we can
  // waste a lot of memory if a lot of streamers are used.
  // Always insert the first point
  vtkPoints* outputPoints = vtkPoints::New();
  vtkCellArray* outputLines = vtkCellArray::New();

  // We will keep track of integration time in this array
  vtkDoubleArray* time = vtkDoubleArray::New();
  time->SetName("IntegrationTime");

  // This array explains why the integration stopped
  vtkIntArray* retVals = vtkIntArray::New();
  retVals->SetName("ReasonForTermination");

  // We will interpolate all point attributes of the input on each point of
  // the output (unless they are turned off). Note that we are using only
  // the first input, if there are more than one, the attributes have to match.
  //
  // Note: We have to use a specific value (safe to employ the maximum number
  //       of steps) as the size of the initial memory allocation here. The
  //       use of the default argument might incur a crash problem (due to
  //       "insufficient memory") in the parallel mode. This is the case when
  //       a streamline intensely shuttles between two processes in an exactly
  //       interleaving fashion --- only one point is produced on each process
  //       (and actually two points, after point duplication, are saved to a
  //       vtkPolyData in vtkDistributedStreamTracer::NoBlockProcessTask) and
  //       as a consequence a large number of such small vtkPolyData objects
  //       are needed to represent a streamline, consuming up the memory before
  //       the intermediate memory is timely released.
  outputPD->InterpolateAllocate( input0->GetPointData(),
    this->MaximumNumberOfSteps );

  vtkIdType numPtsTotal=0;
  double velocity[3];

  int shouldAbort = 0;

  for(int currentLine = 0; currentLine < numLines; currentLine++)
    {
    double progress = static_cast<double>(currentLine)/numLines;
    this->UpdateProgress(progress);

    switch (integrationDirections->GetValue(currentLine))
    {
    case FORWARD:
      direction = 1;
      break;
    case BACKWARD:
      direction = -1;
      break;
    }

    // temporary variables used in the integration
    double point1[3], point2[3];
    vtkIdType numPts=0;

    // Clear the last cell to avoid starting a search from
    // the last point in the streamline
    func->ClearLastCellId();

    // Initial point
    seedSource->GetTuple(seedIds->GetId(currentLine), point1);
    memcpy(point2, point1, 3*sizeof(double));
    if (!func->FunctionValues(point1, velocity))
      {
      continue;
      }

    if ( propagation >= this->MaximumPropagation ||
      numSteps    >  this->MaximumNumberOfSteps)
      {
      continue;
      }

    numPts++;
    numPtsTotal++;
    vtkIdType nextPoint = outputPoints->InsertNextPoint(point1);
    time->InsertNextValue(0.0);

    // We will always pass an arc-length step size to the integrator.
    // If the user specifies a step size in cell length unit, we will
    // have to convert it to arc length.
    IntervalInformation stepSize;  // either positive or negative
    stepSize.Unit  = LENGTH_UNIT;
    stepSize.Interval = 0;
    IntervalInformation aStep; // always positive
    aStep.Unit = LENGTH_UNIT;
    double step, minStep=0, maxStep=0;
    double stepTaken, accumTime=0;
    double speed;
    double cellLength;
    int retVal=OUT_OF_LENGTH, tmp;

    // Make sure we use the dataset found by the vtkAbstractInterpolatedVelocityField
    input = func->GetLastDataSet();
    inputPD = input->GetPointData();
    inVectors = inputPD->GetVectors(vecName);

    // Convert intervals to arc-length unit
    input->GetCell(func->GetLastCellId(), cell);
    cellLength = sqrt(static_cast<double>(cell->GetLength2()));
    speed = vtkMath::Norm(velocity);
    // Never call conversion methods if speed == 0
    if ( speed != 0.0 )
      {
      this->ConvertIntervals( stepSize.Interval, minStep, maxStep,
        direction, cellLength );
      }

    // Interpolate all point attributes on first point
    func->GetLastWeights(weights);
    outputPD->InterpolatePoint(inputPD, nextPoint, cell->PointIds, weights);

    double error = 0;
    // Integrate until the maximum propagation length is reached,
    // maximum number of steps is reached or until a boundary is encountered.
    // Begin Integration
    while ( propagation < this->MaximumPropagation )
      {

      if (numSteps > this->MaximumNumberOfSteps)
        {
        retVal = OUT_OF_STEPS;
        break;
        }

      if ( numSteps++ % 1000 == 1 )
        {
        progress =
          ( currentLine + propagation / this->MaximumPropagation ) / numLines;
        this->UpdateProgress(progress);

        if (this->GetAbortExecute())
          {
          shouldAbort = 1;
          break;
          }
        }

      // Never call conversion methods if speed == 0
      if ( (speed == 0) || (speed <= this->TerminalSpeed) )
        {
        retVal = STAGNATION;
        break;
        }

      // If, with the next step, propagation will be larger than
      // max, reduce it so that it is (approximately) equal to max.
      aStep.Interval = fabs( stepSize.Interval );

      if ( ( propagation + aStep.Interval ) > this->MaximumPropagation )
        {
        aStep.Interval = this->MaximumPropagation - propagation;
        if ( stepSize.Interval >= 0 )
          {
          stepSize.Interval = this->ConvertToLength( aStep, cellLength );
          }
        else
          {
          stepSize.Interval = this->ConvertToLength( aStep, cellLength ) * ( -1.0 );
          }
        maxStep = stepSize.Interval;
        }
      this->LastUsedStepSize = stepSize.Interval;

      // Calculate the next step using the integrator provided
      // Break if the next point is out of bounds.
      func->SetNormalizeVector( true );
      tmp = integrator->ComputeNextStep( point1, point2, 0, stepSize.Interval,
        stepTaken, minStep, maxStep,
        this->MaximumError, error );
      func->SetNormalizeVector( false );
      if ( tmp != 0 )
        {
        retVal = tmp;
        memcpy(lastPoint, point2, 3*sizeof(double));
        break;
        }

      // It is not enough to use the starting point for stagnation calculation
      // Use delX/stepSize to calculate speed and check if it is below
      // stagnation threshold
      double disp[3];
      for (i=0; i<3; i++)
        {
        disp[i] = point2[i] - point1[i];
        }
      if ( (stepSize.Interval == 0) ||
        (vtkMath::Norm(disp) / fabs(stepSize.Interval) <= this->TerminalSpeed) )
        {
        retVal = STAGNATION;
        break;
        }

      accumTime += stepTaken / speed;
      // Calculate propagation (using the same units as MaximumPropagation
      propagation += fabs( stepSize.Interval );

      // This is the next starting point
      for(i=0; i<3; i++)
        {
        point1[i] = point2[i];
        }

      // Interpolate the velocity at the next point
      if ( !func->FunctionValues(point2, velocity) )
        {
        retVal = OUT_OF_DOMAIN;
        memcpy(lastPoint, point2, 3*sizeof(double));
        break;
        }
      // Make sure we use the dataset found by the vtkAbstractInterpolatedVelocityField
      input = func->GetLastDataSet();
      inputPD = input->GetPointData();
      inVectors = inputPD->GetVectors(vecName);

      // Point is valid. Insert it.
      numPts++;
      numPtsTotal++;
      nextPoint = outputPoints->InsertNextPoint(point1);
      time->InsertNextValue(accumTime);

      // Calculate cell length and speed to be used in unit conversions
      input->GetCell(func->GetLastCellId(), cell);
      cellLength = sqrt(static_cast<double>(cell->GetLength2()));
      speed = vtkMath::Norm(velocity);
      // Interpolate all point attributes on current point
      func->GetLastWeights(weights);
      outputPD->InterpolatePoint(inputPD, nextPoint, cell->PointIds, weights);

      // Never call conversion methods if speed == 0
      if ( (speed == 0) || (speed <= this->TerminalSpeed) )
        {
        retVal = STAGNATION;
        break;
        }

      // Convert all intervals to arc length
      this->ConvertIntervals( step, minStep, maxStep, direction, cellLength );


      // If the solver is adaptive and the next step size (stepSize.Interval)
      // that the solver wants to use is smaller than minStep or larger
      // than maxStep, re-adjust it. This has to be done every step
      // because minStep and maxStep can change depending on the cell
      // size (unless it is specified in arc-length unit)
      if (integrator->IsAdaptive())
        {
        if (fabs(stepSize.Interval) < fabs(minStep))
          {
          stepSize.Interval = fabs( minStep ) *
            stepSize.Interval / fabs( stepSize.Interval );
          }
        else if (fabs(stepSize.Interval) > fabs(maxStep))
          {
          stepSize.Interval = fabs( maxStep ) *
            stepSize.Interval / fabs( stepSize.Interval );
          }
        }
      else
        {
        stepSize.Interval = step;
        }

      // End Integration
      }

    if (shouldAbort)
      {
      break;
      }

    if (numPts > 1)
      {
      outputLines->InsertNextCell(numPts);
      for (i=numPtsTotal-numPts; i<numPtsTotal; i++)
        {
        outputLines->InsertCellPoint(i);
        }
      retVals->InsertNextValue(retVal);
      }

    // Initialize these to 0 before starting the next line.
    // The values passed in the function call are only used
    // for the first line.
    inPropagation = propagation;
    inNumSteps = numSteps;

    propagation = 0;
    numSteps = 0;
    }

  if (!shouldAbort)
    {
    // Create the output polyline
    output->SetPoints(outputPoints);
    outputPD->AddArray(time);

    vtkIdType numPts = outputPoints->GetNumberOfPoints();
    if ( numPts > 1 )
      {
      // Assign geometry and attributes
      output->SetLines(outputLines);
      if (this->GenerateNormalsInIntegrate)
        {
        this->GenerateNormals(output, 0, vecName);
        }

      outputCD->AddArray(retVals);
      }
    }

  retVals->Delete();

  outputPoints->Delete();
  outputLines->Delete();

  time->Delete();

  integrator->Delete();
  cell->Delete();

  delete[] weights;

  output->Squeeze();
  return;
}

//---------------------------------------------------------------------------
void vtkCMBStreamTracer::SetNumberOfTestLocations(int NumTests)
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    solver->SetNumberOfTestLocations(NumTests);
    }
}
//---------------------------------------------------------------------------
int vtkCMBStreamTracer::GetNumberOfTestLocations()
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    return solver->GetNumberOfTestLocations();
    }
  return 0;
}
//---------------------------------------------------------------------------
double vtkCMBStreamTracer::GetSensorDefaultRelativeOffset()
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    return solver->GetDefaultRelativeOffset();
    }
  return 0.0;
}
//---------------------------------------------------------------------------
void vtkCMBStreamTracer::SetSensorDefaultRelativeOffset(double offset)
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    solver->SetDefaultRelativeOffset(offset);
    }
}

//---------------------------------------------------------------------------
void vtkCMBStreamTracer::SetRelativeOffsetOfTestLocation(
  int index, double offset[3])
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    solver->SetRelativeOffsetOfTestLocation(index, offset);
    }
}
//---------------------------------------------------------------------------
void vtkCMBStreamTracer::GetRelativeOffsetOfTestLocation(
  int index, double offset[3])
{
  vtkCMBInitialValueProblemSolver* solver =
    vtkCMBInitialValueProblemSolver::SafeDownCast(this->GetIntegrator());
  if(solver)
    {
    solver->GetRelativeOffsetOfTestLocation(index, offset);
    }
}

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