// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause

#include "vtkProteinRibbonFilter.h"

#include "vtkCellArray.h"
#include "vtkCellData.h"
#include "vtkDoubleArray.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkMath.h"
#include "vtkMolecule.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPeriodicTable.h"
#include "vtkPointData.h"
#include "vtkPolyDataNormals.h"
#include "vtkSphereSource.h"
#include "vtkSplineFilter.h"
#include "vtkStringArray.h"
#include "vtkTubeFilter.h"
#include "vtkUnsignedCharArray.h"
#include "vtkVector.h"

#include <map>
#include <vector>

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkProteinRibbonFilter);

namespace
{
// Shamelessly copied from vtkColorSeries vtkColor3ubFromHex3, should be added
// to the vtkColor code.
vtkColor3ub ToColor3ubFromHex3(vtkTypeUInt32 hex)
{
  unsigned char b = hex & 0xff;
  hex >>= 8;
  unsigned char g = hex & 0xff;
  hex >>= 8;
  unsigned char r = hex & 0xff;
  return vtkColor3ub(r, g, b);
}
// Again, should be added to the vtkColor code.
vtkColor3ub ToColor3ubFromColor3f(const vtkColor3f& color)
{
  return vtkColor3ub(static_cast<unsigned char>(color[0] * 255.0f),
    static_cast<unsigned char>(color[1] * 255.0f), static_cast<unsigned char>(color[2] * 255.0f));
}
} // End of anonymous namespace.

vtkProteinRibbonFilter::vtkProteinRibbonFilter()
{
  this->CoilWidth = 0.3f;
  this->HelixWidth = 1.3f;
  this->SphereResolution = 20;
  this->SubdivideFactor = 20;
  this->DrawSmallMoleculesAsSpheres = true;

  this->ElementColors["H"] = ToColor3ubFromHex3(0xCCCCCC);
  this->ElementColors["C"] = ToColor3ubFromHex3(0xAAAAAA);
  this->ElementColors["O"] = ToColor3ubFromHex3(0xCC0000);
  this->ElementColors["N"] = ToColor3ubFromHex3(0x0000CC);
  this->ElementColors["S"] = ToColor3ubFromHex3(0xCCCC00);
  this->ElementColors["P"] = ToColor3ubFromHex3(0x6622CC);
  this->ElementColors["F"] = ToColor3ubFromHex3(0x00CC00);
  this->ElementColors["CL"] = ToColor3ubFromHex3(0x00CC00);
  this->ElementColors["BR"] = ToColor3ubFromHex3(0x882200);
  this->ElementColors["I"] = ToColor3ubFromHex3(0x6600AA);
  this->ElementColors["FE"] = ToColor3ubFromHex3(0xCC6600);
  this->ElementColors["CA"] = ToColor3ubFromHex3(0xDDDDDD);
}

vtkProteinRibbonFilter::~vtkProteinRibbonFilter() = default;

int vtkProteinRibbonFilter::FillInputPortInformation(int port, vtkInformation* info)
{
  if (!this->Superclass::FillInputPortInformation(port, info))
  {
    return 0;
  }

  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkPolyData");
  return 1;
}

int vtkProteinRibbonFilter::RequestData(
  vtkInformation*, vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  vtkPolyData* input = vtkPolyData::GetData(inputVector[0]);
  vtkPolyData* output = vtkPolyData::GetData(outputVector);

  vtkPointData* pointData = input->GetPointData();

  // Extract alpha-carbon backbone from input poly data
  vtkStringArray* atomTypes =
    vtkArrayDownCast<vtkStringArray>(pointData->GetAbstractArray("atom_types"));
  vtkIdTypeArray* atomType =
    vtkArrayDownCast<vtkIdTypeArray>(pointData->GetAbstractArray("atom_type"));

  if (!atomTypes || !atomType)
  {
    vtkErrorMacro(<< "Atom Type String & Ids Arrays Required");
    return 0;
  }

  // Extract secondary structures information from input poly data
  vtkIdTypeArray* resi = vtkArrayDownCast<vtkIdTypeArray>(pointData->GetAbstractArray("residue"));
  vtkUnsignedCharArray* chain =
    vtkArrayDownCast<vtkUnsignedCharArray>(pointData->GetAbstractArray("chain"));
  vtkUnsignedCharArray* atom_ss =
    vtkArrayDownCast<vtkUnsignedCharArray>(pointData->GetAbstractArray("secondary_structures"));
  vtkUnsignedCharArray* atom_ssbegin = vtkArrayDownCast<vtkUnsignedCharArray>(
    pointData->GetAbstractArray("secondary_structures_begin"));
  vtkUnsignedCharArray* atom_ssend =
    vtkArrayDownCast<vtkUnsignedCharArray>(pointData->GetAbstractArray("secondary_structures_end"));
  vtkUnsignedCharArray* ishetatm =
    vtkArrayDownCast<vtkUnsignedCharArray>(pointData->GetAbstractArray("ishetatm"));

  if (!resi || !chain || !atom_ss || !atom_ssbegin || !atom_ssend || !ishetatm)
  {
    vtkErrorMacro(<< "Atom Secondary Structures Arrays Required");
    return 0;
  }

  unsigned char currentChain = 0;
  unsigned char ss = 0;
  vtkIdType currentResi = 0;
  vtkVector3f currentCA(0.f);
  vtkVector3f prevCO(0.f);
  bool hasPrevCO = false;

  vtkNew<vtkPoints> strandPoints;
  vtkNew<vtkPolyData> strand;
  strand->AllocateExact(1024, 1024);
  strand->SetPoints(strandPoints);

  vtkNew<vtkUnsignedCharArray> pointsColors;
  pointsColors->SetName("RGB");
  pointsColors->SetNumberOfComponents(3);

  // Initialize colors per point/atom
  std::vector<vtkColor3ub> atomsColors;
  this->SetColorByAtom(atomsColors, atomTypes);
  this->SetColorByStructure(
    atomsColors, atomTypes, atom_ss, ToColor3ubFromHex3(0xFF0080), ToColor3ubFromHex3(0xFFC800));

  std::vector<vtkColor3ub> colors;
  std::vector<std::pair<vtkVector3f, bool>> borderPoints[2];

  // Need this for radius / color lookups
  vtkNew<vtkPeriodicTable> pTab;

  for (int i = 0; i < input->GetNumberOfPoints(); i++)
  {
    std::string type = atomTypes->GetValue(i);
    unsigned short atomicNum = static_cast<unsigned short>(atomType->GetValue(i));

    if (ishetatm->GetValue(i) && this->DrawSmallMoleculesAsSpheres)
    {
      if (type != "O")
      {
        CreateAtomAsSphere(strand, pointsColors, input->GetPoint(i),
          ToColor3ubFromColor3f(pTab->GetDefaultRGBTuple(atomicNum)), pTab->GetVDWRadius(atomicNum),
          1.f);
      }
    }
    else if (type == "CA")
    {
      // Create a ribbon between 2 CA atoms passing through each O atoms found in-between
      double* xyz = input->GetPoint(i);
      unsigned char atomChain = chain->GetValue(i);
      vtkIdType atomResi = resi->GetValue(i);

      if (currentChain != atomChain || currentResi + 1 != atomResi)
      {
        this->CreateThinStrip(
          strand, pointsColors, strandPoints, borderPoints[0], borderPoints[1], colors);
        borderPoints[0].clear();
        borderPoints[1].clear();
        colors.clear();
        hasPrevCO = false;
      }
      currentCA.Set(xyz[0], xyz[1], xyz[2]);
      currentChain = atomChain;
      currentResi = atomResi;
      ss = atom_ss->GetValue(i);
      colors.push_back(atomsColors[i]);
    }
    else if (type == "O")
    {
      // Insert a new step in the next ribbon
      double* xyz = input->GetPoint(i);
      vtkVector3f p(xyz[0], xyz[1], xyz[2]);
      p = (p - currentCA).Normalized() * ((ss == 'c') ? this->CoilWidth : this->HelixWidth);
      if (hasPrevCO && p.Dot(prevCO) < 0)
      {
        p = p * -1.f;
      }
      hasPrevCO = true;
      prevCO = p;
      bool isSheet = (ss == 's');
      borderPoints[0].emplace_back(currentCA - prevCO, isSheet);
      borderPoints[1].emplace_back(currentCA + prevCO, isSheet);
    }
  }

  // Create the last ribbon strip if needed
  this->CreateThinStrip(
    strand, pointsColors, strandPoints, borderPoints[0], borderPoints[1], colors);

  strand->GetPointData()->SetScalars(pointsColors);

  // Compute the model normals
  vtkNew<vtkPolyDataNormals> pdnormals;
  pdnormals->SetInputData(strand);
  pdnormals->SetFeatureAngle(150.0);
  pdnormals->Update();

  output->ShallowCopy(pdnormals->GetOutput());

  return 1;
}

void vtkProteinRibbonFilter::SetColorByAtom(
  std::vector<vtkColor3ub>& colors, vtkStringArray* atomTypes)
{
  vtkNew<vtkPeriodicTable> pTab;
  unsigned int len = atomTypes->GetNumberOfValues();
  colors.resize(len);
  for (unsigned int i = 0; i < len; i++)
  {
    if (this->ElementColors.find(atomTypes->GetValue(i)) != this->ElementColors.end())
    {
      colors[i] = this->ElementColors[atomTypes->GetValue(i)];
    }
    else
    {
      colors[i] = vtkColor3ub(0xFFFFFF);
    }
  }
}

void vtkProteinRibbonFilter::SetColorByStructure(std::vector<vtkColor3ub>& colors,
  vtkStringArray* atomTypes, vtkUnsignedCharArray* ss, const vtkColor3ub& helixColor,
  const vtkColor3ub& sheetColor)
{
  unsigned int len = atomTypes->GetNumberOfValues();
  colors.resize(len);
  for (unsigned int i = 0; i < len; i++)
  {
    if (ss->GetValue(i) == 's')
    {
      colors[i] = sheetColor;
    }
    else if (ss->GetValue(i) == 'h')
    {
      colors[i] = helixColor;
    }
  }
}

void vtkProteinRibbonFilter::CreateAtomAsSphere(vtkPolyData* poly,
  vtkUnsignedCharArray* pointsColors, double* pos, const vtkColor3ub& color, float radius,
  float scale)
{
  // Create the sphere source at the atom size & position
  vtkNew<vtkSphereSource> sphereSource;
  sphereSource->SetPhiResolution(this->SphereResolution);
  sphereSource->SetThetaResolution(this->SphereResolution);
  sphereSource->SetCenter(pos);
  sphereSource->SetRadius(radius * scale);
  sphereSource->Update();

  // Extract polydata from sphere
  vtkPolyData* sphere = sphereSource->GetOutput();
  vtkPoints* spherePoints = sphere->GetPoints();
  vtkCellArray* spherePolys = sphere->GetPolys();

  vtkPoints* points = poly->GetPoints();
  // Get offset for the new point IDs that will be added to points
  vtkIdType pointOffset = points->GetNumberOfPoints();
  // Total number of new points
  vtkIdType numPoints = spherePoints->GetNumberOfPoints();
  vtkIdType numCellPoints;
  const vtkIdType* cellPoints;
  // Add new points
  for (vtkIdType i = 0; i < numPoints; ++i)
  {
    points->InsertNextPoint(spherePoints->GetPoint(i));
    for (int j = 0; j < 3; ++j)
    {
      pointsColors->InsertNextValue(color[j]);
    }
  }

  // Add new cells (polygons) that represent the sphere
  spherePolys->InitTraversal();
  while (spherePolys->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;
    }
    poly->InsertNextCell(VTK_TRIANGLE_STRIP, numCellPoints, newCellPoints.data());
  }
}

void vtkProteinRibbonFilter::CreateThinStrip(vtkPolyData* poly, vtkUnsignedCharArray* pointsColors,
  vtkPoints* p, std::vector<std::pair<vtkVector3f, bool>>& p1,
  std::vector<std::pair<vtkVector3f, bool>>& p2, std::vector<vtkColor3ub>& colors)
{
  if (p1.size() < 2 || p2.size() < 2)
  {
    return;
  }

  // Get offset for the new point IDs that will be added to points
  vtkIdType pointOffset = p->GetNumberOfPoints();

  // Subdivide (smooth) the 2 ribbon borders
  std::vector<vtkVector3f>* points1 = Subdivide(p1, this->SubdivideFactor);
  std::vector<vtkVector3f>* points2 = Subdivide(p2, this->SubdivideFactor);

  int len = static_cast<int>(points1->size());

  // Insert smoothed ribbon borders points into the polydata
  for (int i = 0; i < len; i++)
  {
    p->InsertNextPoint((*points1)[i].GetData());
    p->InsertNextPoint((*points2)[i].GetData());

    vtkColor3ub color =
      colors[static_cast<int>(floor(0.5f + i / static_cast<float>(this->SubdivideFactor)))];
    for (int k = 0; k < 2; ++k)
    {
      for (int ci = 0; ci < 3; ++ci)
      {
        pointsColors->InsertNextValue(color[ci]);
      }
    }
  }
  delete points1;
  delete points2;

  // Fill in between the 2 ribbons borders with triangle strips
  vtkIdType connectivity[4];
  for (int i = 0, offset = pointOffset; i < len - 1; i++, offset += 2)
  {
    for (int j = 0; j < 4; j++)
    {
      connectivity[j] = offset + j;
    }
    poly->InsertNextCell(VTK_TRIANGLE_STRIP, 4, connectivity);
  }
}

std::vector<vtkVector3f>* vtkProteinRibbonFilter::Subdivide(
  std::vector<std::pair<vtkVector3f, bool>>& p, int div)
{
  std::vector<vtkVector3f>* ret = new std::vector<vtkVector3f>;
  std::vector<vtkVector3f> points;

  // Smoothing test
  points.push_back(p[0].first);
  for (int i = 1, lim = static_cast<int>(p.size()) - 1; i < lim; i++)
  {
    vtkVector3f& p1 = p[i].first;
    vtkVector3f& p2 = p[i + 1].first;
    if (p[i].second)
    {
      points.push_back((p1 + p2) * 0.5f);
    }
    else
    {
      points.push_back(p1);
    }
  }
  points.push_back(p[p.size() - 1].first);

  // Catmull-Rom subdivision
  for (int i = -1, size = static_cast<int>(points.size()); i <= size - 3; i++)
  {
    vtkVector3f& p0 = points[(i == -1) ? 0 : i];
    vtkVector3f& p1 = points[i + 1];
    vtkVector3f& p2 = points[i + 2];
    vtkVector3f& p3 = points[(i == size - 3) ? size - 1 : i + 3];
    vtkVector3f v0 = (p2 - p0) * 0.5f;
    vtkVector3f v1 = (p3 - p1) * 0.5f;
    for (int j = 0; j < div; j++)
    {
      double t = 1.0 / div * j;
      double t2 = t * t;
      double x = p1.GetX() + t * v0.GetX() +
        t2 * (-3 * p1.GetX() + 3 * p2.GetX() - 2 * v0.GetX() - v1.GetX()) +
        t2 * t * (2 * p1.GetX() - 2 * p2.GetX() + v0.GetX() + v1.GetX());
      double y = p1.GetY() + t * v0.GetY() +
        t2 * (-3 * p1.GetY() + 3 * p2.GetY() - 2 * v0.GetY() - v1.GetY()) +
        t2 * t * (2 * p1.GetY() - 2 * p2.GetY() + v0.GetY() + v1.GetY());
      double z = p1.GetZ() + t * v0.GetZ() +
        t2 * (-3 * p1.GetZ() + 3 * p2.GetZ() - 2 * v0.GetZ() - v1.GetZ()) +
        t2 * t * (2 * p1.GetZ() - 2 * p2.GetZ() + v0.GetZ() + v1.GetZ());
      ret->emplace_back(x, y, z);
    }
  }
  ret->push_back(points[points.size() - 1]);
  return ret;
}

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