// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkMoleculeToBondStickFilter.h"

#include "vtkCellArray.h"
#include "vtkCylinderSource.h"
#include "vtkMath.h"
#include "vtkMolecule.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPoints.h"
#include "vtkPolyData.h"
#include "vtkTransform.h"
#include "vtkUnsignedShortArray.h"

#include <vector>

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkMoleculeToBondStickFilter);

//------------------------------------------------------------------------------
vtkMoleculeToBondStickFilter::vtkMoleculeToBondStickFilter() = default;

//------------------------------------------------------------------------------
vtkMoleculeToBondStickFilter::~vtkMoleculeToBondStickFilter() = default;

//------------------------------------------------------------------------------
int vtkMoleculeToBondStickFilter::RequestData(
  vtkInformation*, vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  vtkMolecule* input = vtkMolecule::SafeDownCast(vtkDataObject::GetData(inputVector[0]));
  vtkPolyData* output = vtkPolyData::SafeDownCast(vtkDataObject::GetData(outputVector));

  // Extract data from input molecule
  vtkIdType numBonds = input->GetNumberOfBonds();

  // Prep the output
  output->Initialize();
  vtkCellArray* polys = vtkCellArray::New();
  vtkPoints* points = vtkPoints::New();
  vtkUnsignedShortArray* bondOrders = vtkUnsignedShortArray::New();
  bondOrders->SetName(input->GetBondOrdersArrayName());

  // Initialize a CylinderSource
  vtkCylinderSource* cylSource = vtkCylinderSource::New();
  cylSource->SetResolution(20);
  cylSource->SetHeight(1.0);
  cylSource->Update();

  // Preallocate memory
  points->Allocate(3 * numBonds * cylSource->GetOutput()->GetPoints()->GetNumberOfPoints());
  polys->AllocateEstimate(numBonds * cylSource->GetOutput()->GetPolys()->GetNumberOfCells(), 3);
  bondOrders->Allocate(3 * numBonds * cylSource->GetOutput()->GetPoints()->GetNumberOfPoints());

  // Create a transform object to map the cylinder source to the bond
  vtkTransform* xform = vtkTransform::New();
  xform->PostMultiply();

  // Declare some variables for later
  vtkIdType numCellPoints;
  const vtkIdType* cellPoints;
  unsigned short bondOrder;
  double bondLength;
  double radius;
  double delta[3];
  double initialDisp[3];
  double rotAngle;
  double rotAxis[3];
  double bondCenter[3];
  double pos1[3], pos2[3];
  // Normalized vector pointing along the cylinder (y axis);
  static const double cylVec[3] = { 0.0, 1.0, 0.0 };
  // Normalized vector pointing along bond
  double bondVec[3];
  // Unit z vector
  static const double unitZ[3] = { 0.0, 0.0, 1.0 };

  // Build a cylinder for each bond and append its data to the output arrays.
  for (vtkIdType bondInd = 0; bondInd < numBonds; ++bondInd)
  {
    // Extract bond info
    vtkBond bond = input->GetBond(bondInd);
    bondOrder = bond.GetOrder();
    bond.GetBeginAtom().GetPosition(pos1);
    bond.GetEndAtom().GetPosition(pos2);

    // Compute additional bond info
    // - Normalized vector in direction of bond
    vtkMath::Subtract(pos2, pos1, bondVec);
    bondLength = vtkMath::Normalize(bondVec);
    // - Axis for cylinder rotation, bondVec [cross] cylVec
    vtkMath::Cross(bondVec, cylVec, rotAxis);
    // Rotation angle
    rotAngle = -vtkMath::DegreesFromRadians(acos(vtkMath::Dot(bondVec, cylVec)));
    // - Center of bond for translation
    vtkMath::Add(pos2, pos1, bondCenter);
    vtkMath::MultiplyScalar(bondCenter, 0.5);

    // Set up delta step vector and bond radius from bond order:
    switch (bondOrder)
    {
      case 1:
      default:
        radius = 0.1;
        delta[0] = 0.0;
        delta[1] = 0.0;
        delta[2] = 0.0;
        initialDisp[0] = 0.0;
        initialDisp[1] = 0.0;
        initialDisp[2] = 0.0;
        break;
      case 2:
        radius = 0.1;
        vtkMath::Cross(bondVec, unitZ, delta);
        vtkMath::MultiplyScalar(delta, radius + radius);
        initialDisp[0] = delta[0] * -0.5;
        initialDisp[1] = delta[1] * -0.5;
        initialDisp[2] = delta[2] * -0.5;
        break;
      case 3:
        radius = 0.1;
        vtkMath::Cross(bondVec, unitZ, delta);
        vtkMath::MultiplyScalar(delta, radius + radius);
        initialDisp[0] = -delta[0];
        initialDisp[1] = -delta[1];
        initialDisp[2] = -delta[2];
        break;
    }

    // Construct transform
    xform->Identity();
    xform->Scale(radius, bondLength, radius);
    xform->RotateWXYZ(rotAngle, rotAxis);
    xform->Translate(bondCenter[0], bondCenter[1], bondCenter[2]);
    xform->Translate(initialDisp);

    // For each bond order, add a cylinder to output, translate by
    // delta, and repeat.
    for (unsigned short iter = 0; iter < bondOrder; ++iter)
    {
      vtkPolyData* cylinder = cylSource->GetOutput();
      vtkPoints* cylPoints = cylinder->GetPoints();
      vtkCellArray* cylPolys = cylinder->GetPolys();

      // Get offset for the new point IDs that will be added to points
      vtkIdType pointOffset = points->GetNumberOfPoints();
      // Total number of new points
      vtkIdType numPoints = cylPoints->GetNumberOfPoints();

      // Create points by transforming those of the cylinder
      xform->TransformPoints(cylPoints, points);

      // Use bond order for point scalar data.
      for (vtkIdType i = 0; i < numPoints; ++i)
      {
        bondOrders->InsertNextValue(bondOrder);
      }

      // Add new cells (polygons) that represent the cylinder
      cylPolys->InitTraversal();
      while (cylPolys->GetNextCell(numCellPoints, cellPoints) != 0)
      {
        std::vector<vtkIdType> newCellPoints(numCellPoints);
        for (vtkIdType i = 0; i < numCellPoints; ++i)
        {
          // The new point ids should be offset by the pointOffset above
          newCellPoints[i] = cellPoints[i] + pointOffset;
        }
        polys->InsertNextCell(numCellPoints, newCellPoints.data());
      }

      // Setup for the next cylinder in a multi-bond
      xform->Translate(delta);
    }
  }

  // Release extra memory
  points->Squeeze();
  bondOrders->Squeeze();
  polys->Squeeze();

  // update output
  output->SetPoints(points);
  output->GetPointData()->SetScalars(bondOrders);
  output->SetPolys(polys);

  // Clean up
  xform->Delete();
  polys->Delete();
  points->Delete();
  bondOrders->Delete();
  cylSource->Delete();

  return 1;
}

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