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

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

#include "vtkBoundingBox.h"
#include "vtkCellData.h"
#include "vtkCellType.h"
#include "vtkDataSetAttributes.h"
#include "vtkFloatArray.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMolecule.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPointLocator.h"
#include "vtkPoints.h"
#include "vtkStreamingDemandDrivenPipeline.h"
#include "vtkUnsignedShortArray.h"
#include "vtkUnstructuredGrid.h"
#include "vtkVector.h"
#include "vtkVectorOperators.h"

#include "vtksys/FStream.hxx"
#include <vtksys/RegularExpression.hxx>

#include <algorithm>
#include <cassert>
#include <cmath>
#include <fstream>
#include <set>
#include <sstream>
#include <vector>

vtkStandardNewMacro(vtkVASPTessellationReader);

typedef vtksys::RegularExpression RegEx;
typedef vtkStreamingDemandDrivenPipeline vtkSDDP;

namespace
{

template <typename T>
bool parse(const std::string& str, T& result)
{
  if (!str.empty())
  {
    std::istringstream tmp(str);
    tmp >> result;
    return !tmp.fail();
  }
  return false;
}

template <typename T>
bool parseCommaSepList(const std::string& input, std::vector<T>& data)
{
  std::istringstream in(input);
  for (std::string valStr; std::getline(in, valStr, ',');)
  {
    T tmp;
    if (!parse(valStr, tmp))
    {
      return false;
    }
    data.push_back(tmp);
  }
  return true;
}

// This parses the voronoi points/faces list. The input is expected to be:
// [number of commaSepLists], ([commaSepList]) ([commaSepList]) ...
template <typename T>
bool parseVariableLists(
  const std::string& input, std::vector<std::vector<T>>& data, RegEx* parenExtract)
{
  std::istringstream in(input);
  size_t nLists;
  in >> nLists;
  if (in.fail())
  {
    return false;
  }

  data.resize(nLists);

  std::string parseBuffer = input;
  for (size_t i = 0; i < nLists; ++i)
  {
    if (!parenExtract->find(parseBuffer))
    {
      return false;
    }
    if (!parseCommaSepList(parenExtract->match(1), data[i]))
    {
      return false;
    }
    // Chop the current match off of the buffer to prepare for the next iter.
    parseBuffer = parseBuffer.substr(parenExtract->end());
  }

  return true;
}

} // end anon namespace

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

//------------------------------------------------------------------------------
vtkVASPTessellationReader::vtkVASPTessellationReader()
  : FileName(nullptr)
  , TimeParser(new RegEx("^ *time *= *([0-9EeDd.+-]+) *$"))
  ,                                                       // time = (timeVal)
  LatticeParser(new RegEx("^ *Rx1 *= *([0-9EeDd.+-]+) *," // Rx1
                          " *Rx2 *= *([0-9EeDd.+-]+) *,"  // Rx2
                          " *Rx3 *= *([0-9EeDd.+-]+) *,"  // Rx3
                          " *Ry2 *= *([0-9EeDd.+-]+) *,"  // Ry2
                          " *Ry3 *= *([0-9EeDd.+-]+) *,"  // Ry3
                          " *Rz3 *= *([0-9EeDd.+-]+) *$"  // Rz3
    ))
  , AtomCountParser(new RegEx("^ *Natoms *= *([0-9]+) *$"))
  ,                                           // Natoms = (int)),
  AtomParser(new RegEx("^ *([0-9]+) *,"       // Atom index
                       " *\\("                // Open paren
                       " *([0-9EeDd.+-]+) *," // X coord
                       " *([0-9EeDd.+-]+) *," // Y coord
                       " *([0-9EeDd.+-]+)"    // Z coord
                       " *\\) *,"             // Close paren
                       " *([0-9EeDd.+-]+) *$" // Radius
    ))
  , ParenExtract(new RegEx("\\(([^(]+)\\)")) // Extract contents of (...)
{
  this->SetNumberOfInputPorts(0);
  this->SetNumberOfOutputPorts(2);
}

//------------------------------------------------------------------------------
vtkVASPTessellationReader::~vtkVASPTessellationReader()
{
  this->SetFileName(nullptr);
  delete this->TimeParser;
  delete this->LatticeParser;
  delete this->AtomCountParser;
  delete this->AtomParser;
  delete this->ParenExtract;
}

//------------------------------------------------------------------------------
int vtkVASPTessellationReader::RequestData(
  vtkInformation*, vtkInformationVector**, vtkInformationVector* outInfos)
{
  vtkInformation* outInfo0 = outInfos->GetInformationObject(0);
  vtkInformation* outInfo1 = outInfos->GetInformationObject(1);

  vtkMolecule* molecule = vtkMolecule::SafeDownCast(outInfo0->Get(vtkDataObject::DATA_OBJECT()));
  assert(molecule);

  vtkUnstructuredGrid* voronoi =
    vtkUnstructuredGrid::SafeDownCast(outInfo1->Get(vtkDataObject::DATA_OBJECT()));
  assert(voronoi);

  vtksys::ifstream in(this->FileName);
  if (!in)
  {
    vtkErrorMacro("Could not open file for reading: " << (this->FileName ? this->FileName : ""));
    return 1;
  }

  // Advance to the selected timestep:
  size_t stepIdx = this->SelectTimeStepIndex(outInfo0);
  double time = 0.;
  for (size_t i = 0; i <= stepIdx; ++i) // <= to read the "time=" line
  {
    if (!this->NextTimeStep(in, time))
    {
      vtkErrorMacro("Error -- attempting to read timestep #"
        << (stepIdx + 1) << " but encountered a parsing error at timestep #" << (i + 1) << ".");
      return 1;
    }
  }

  if (this->ReadTimeStep(in, molecule, voronoi))
  {
    molecule->GetInformation()->Set(vtkDataObject::DATA_TIME_STEP(), time);
    voronoi->GetInformation()->Set(vtkDataObject::DATA_TIME_STEP(), time);
  }
  else
  {
    molecule->Initialize();
    voronoi->Initialize();
  }

  return 1;
}

//------------------------------------------------------------------------------
int vtkVASPTessellationReader::RequestInformation(
  vtkInformation*, vtkInformationVector**, vtkInformationVector* outInfos)
{
  vtksys::ifstream in(this->FileName);
  if (!in)
  {
    vtkErrorMacro("Could not open file for reading: " << (this->FileName ? this->FileName : ""));
    return 1;
  }

  // Scan the file for timesteps:
  double time;
  std::vector<double> times;
  double timeRange[2] = { VTK_DOUBLE_MAX, VTK_DOUBLE_MIN };
  while (this->NextTimeStep(in, time))
  {
    times.push_back(time);
    timeRange[0] = std::min(timeRange[0], time);
    timeRange[1] = std::max(timeRange[1], time);
  }

  if (!times.empty())
  {
    for (int port = 0; port < 2; ++port)
    {
      vtkInformation* outInfo = outInfos->GetInformationObject(port);
      outInfo->Set(vtkSDDP::TIME_RANGE(), timeRange, 2);
      outInfo->Set(vtkSDDP::TIME_STEPS(), &times[0], static_cast<int>(times.size()));
    }
  }

  return 1;
}

//------------------------------------------------------------------------------
int vtkVASPTessellationReader::FillOutputPortInformation(int port, vtkInformation* info)
{
  switch (port)
  {
    case 0:
      info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkMolecule");
      break;
    case 1:
      info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkUnstructuredGrid");
      break;
    default:
      return 0;
  }

  return 1;
}

//------------------------------------------------------------------------------
bool vtkVASPTessellationReader::NextTimeStep(std::istream& in, double& time)
{
  std::string line;
  while (std::getline(in, line))
  {
    if (this->TimeParser->find(line))
    {
      // Parse timestamp:
      if (!parse(this->TimeParser->match(1), time))
      {
        vtkErrorMacro("Error parsing time information from line: " << line);
        return false;
      }
      return true;
    }
  }

  return false;
}

//------------------------------------------------------------------------------
size_t vtkVASPTessellationReader::SelectTimeStepIndex(vtkInformation* info)
{
  if (!info->Has(vtkSDDP::TIME_STEPS()) || !info->Has(vtkSDDP::UPDATE_TIME_STEP()))
  {
    return 0;
  }

  double* times = info->Get(vtkSDDP::TIME_STEPS());
  int nTimes = info->Length(vtkSDDP::TIME_STEPS());
  double t = info->Get(vtkSDDP::UPDATE_TIME_STEP());

  double resultDiff = VTK_DOUBLE_MAX;
  size_t result = 0;
  for (int i = 0; i < nTimes; ++i)
  {
    double diff = std::fabs(times[i] - t);
    if (diff < resultDiff)
    {
      resultDiff = diff;
      result = static_cast<size_t>(i);
    }
  }

  return result;
}

//------------------------------------------------------------------------------
bool vtkVASPTessellationReader::ReadTimeStep(
  std::istream& in, vtkMolecule* molecule, vtkUnstructuredGrid* voronoi)
{
  // Assume the 'time = ...' line has already been read.
  std::string line;

  // Read the lattice info:
  if (!std::getline(in, line))
  {
    vtkErrorMacro("Unexpected EOF while reading lattice info.");
    return false;
  }
  if (!this->LatticeParser->find(line))
  {
    vtkErrorMacro("Error parsing lattice info from line: " << line);
    return false;
  }

  vtkVector3d latA(0.);
  vtkVector3d latB(0.);
  vtkVector3d latC(0.);
  vtkVector3d latO(0.);

  // Rxx
  if (!parse(this->LatticeParser->match(1), latA[0]))
  {
    vtkErrorMacro(
      "Error parsing Rxx component '" << this->LatticeParser->match(1) << "' from line: " << line);
    return false;
  }
  // Rxy
  if (!parse(this->LatticeParser->match(2), latB[0]))
  {
    vtkErrorMacro(
      "Error parsing Rxy component '" << this->LatticeParser->match(2) << "' from line: " << line);
    return false;
  }
  // Rxz
  if (!parse(this->LatticeParser->match(3), latC[0]))
  {
    vtkErrorMacro(
      "Error parsing Rxz component '" << this->LatticeParser->match(3) << "' from line: " << line);
    return false;
  }
  // Ryy
  if (!parse(this->LatticeParser->match(4), latB[1]))
  {
    vtkErrorMacro(
      "Error parsing Ryy component '" << this->LatticeParser->match(4) << "' from line: " << line);
    return false;
  }
  // Ryz
  if (!parse(this->LatticeParser->match(5), latC[1]))
  {
    vtkErrorMacro(
      "Error parsing Ryz component '" << this->LatticeParser->match(5) << "' from line: " << line);
    return false;
  }
  // Rzz
  if (!parse(this->LatticeParser->match(6), latC[2]))
  {
    vtkErrorMacro(
      "Error parsing Rzz component '" << this->LatticeParser->match(6) << "' from line: " << line);
    return false;
  }

  molecule->SetLattice(latA, latB, latC);
  molecule->SetLatticeOrigin(latO);

  // Number of atoms:
  if (!std::getline(in, line))
  {
    vtkErrorMacro("Unexpected EOF while parsing number of atoms.");
    return false;
  }
  if (!this->AtomCountParser->find(line))
  {
    vtkErrorMacro("Error parsing number of atoms from line: " << line);
    return false;
  }
  vtkIdType nAtoms;
  if (!parse(this->AtomCountParser->match(1), nAtoms))
  {
    vtkErrorMacro("Error parsing number atoms '" << this->AtomCountParser->match(1)
                                                 << "'' from line: " << line);
    return false;
  }

  // Skip 'Atomic_Numbers =' line:
  if (!std::getline(in, line))
  {
    vtkErrorMacro("Unexpected EOF.");
    return false;
  }

  // Atomic numbers:
  std::vector<unsigned short> atomicNumbers;
  atomicNumbers.reserve(static_cast<size_t>(nAtoms));
  if (!std::getline(in, line))
  {
    vtkErrorMacro("Unexpected EOF while reading atomic number list.");
    return false;
  }
  if (!parseCommaSepList(line, atomicNumbers))
  {
    vtkErrorMacro("Error while parsing atomic number list: " << line);
    return false;
  }

  // Initialize the molecule with atoms, setting just the atomic number.
  // We'll add positions as we parse them later.
  if (static_cast<size_t>(nAtoms) != atomicNumbers.size())
  {
    vtkErrorMacro("Error: expected " << nAtoms
                                     << " atomic numbers, but only "
                                        "parsed "
                                     << atomicNumbers.size());
    return false;
  }
  vtkVector3f pos(0.f);
  for (size_t i = 0; i < atomicNumbers.size(); ++i)
  {
    molecule->AppendAtom(atomicNumbers[i], pos);
  }

  // Now for the actual list of atoms/tessellations
  vtkNew<vtkFloatArray> radii;
  radii->SetNumberOfTuples(nAtoms);

  // Compute unit cell bounds to initialize point merging:
  vtkBoundingBox bbox;
  bbox.AddPoint(latO.GetData());
  bbox.AddPoint((latO + latA).GetData());
  bbox.AddPoint((latO + latB).GetData());
  bbox.AddPoint((latO + latC).GetData());
  bbox.AddPoint((latO + latA + latB).GetData());
  bbox.AddPoint((latO + latA + latC).GetData());
  bbox.AddPoint((latO + latB + latC).GetData());
  bbox.AddPoint((latO + latA + latB + latC).GetData());
  double bounds[6];
  bbox.GetBounds(bounds);

  // Merge the tessellation points using a locator:
  vtkNew<vtkPointLocator> locator;
  vtkNew<vtkPoints> tessPoints;
  tessPoints->SetDataTypeToFloat();
  voronoi->SetPoints(tessPoints);
  voronoi->Allocate(nAtoms);

  // Cell attributes for the voronoi tessellation:
  vtkNew<vtkUnsignedShortArray> tessAtomicNumbers;
  tessAtomicNumbers->SetName("Atomic Numbers");
  tessAtomicNumbers->Allocate(nAtoms);
  vtkNew<vtkIdTypeArray> tessAtomIds;
  tessAtomIds->SetName("Atom Ids");
  tessAtomIds->Allocate(nAtoms);

  // Estimate 10 unique points per atom:
  locator->InitPointInsertion(tessPoints, bounds, nAtoms * 10);

  // Storage for parsing the tessellation points/faces info
  std::vector<vtkIdType> faceStream;
  std::vector<vtkIdType> pointIds;
  std::set<vtkIdType> uniquePointIds;
  // parse as doubles for locator API, but store as floats:
  std::vector<std::vector<double>> pointData;
  std::vector<std::vector<vtkIdType>> faceData;

  for (vtkIdType atomEntry = 0; atomEntry < nAtoms; ++atomEntry)
  {
    // Skip any blank lines:
    line.clear();
    while (line.empty() || line.find_first_not_of(" \t") == std::string::npos)
    {
      if (!std::getline(in, line))
      {
        vtkErrorMacro("Unexpected EOF while reading atom entry " << atomEntry);
        return false;
      }
    }

    if (!this->AtomParser->find(line))
    {
      vtkErrorMacro("Error parsing atom position/radius specification: " << line);
      return false;
    }
    vtkIdType atomId;
    if (!parse(this->AtomParser->match(1), atomId))
    {
      vtkErrorMacro(
        "Error parsing atomId '" << this->AtomParser->match(1) << "' from line: " << line);
      return false;
    }
    if (!parse(this->AtomParser->match(2), pos[0]))
    {
      vtkErrorMacro(
        "Error parsing x coordinate '" << this->AtomParser->match(2) << "' from line: " << line);
      return false;
    }
    if (!parse(this->AtomParser->match(3), pos[1]))
    {
      vtkErrorMacro(
        "Error parsing y coordinate '" << this->AtomParser->match(3) << "' from line: " << line);
      return false;
    }
    if (!parse(this->AtomParser->match(4), pos[2]))
    {
      vtkErrorMacro(
        "Error parsing z coordinate '" << this->AtomParser->match(4) << "' from line: " << line);
      return false;
    }
    float radius;
    if (!parse(this->AtomParser->match(5), radius))
    {
      vtkErrorMacro(
        "Error parsing radius '" << this->AtomParser->match(5) << "' from line: " << line);
      return false;
    }

    if (atomId >= nAtoms)
    {
      vtkErrorMacro("Found entry for atom with id " << atomId
                                                    << ", but "
                                                       "only "
                                                    << nAtoms << " atoms exist.");
      return false;
    }
    vtkAtom atom = molecule->GetAtom(atomId);
    atom.SetPosition(pos);
    radii->SetTypedComponent(atomId, 0, radius);

    // Extract tessellation points:
    pointData.clear();
    if (!std::getline(in, line))
    {
      vtkErrorMacro("Unexpected EOF while reading voronoi points for atom " << atomId);
      return false;
    }
    if (!parseVariableLists(line, pointData, this->ParenExtract))
    {
      vtkErrorMacro(
        "Error while parsing voronoi point data for atom " << atomId << ". Input: " << line);
      return false;
    }

    // Extract tessellation faces:
    faceData.clear();
    if (!std::getline(in, line))
    {
      vtkErrorMacro("Unexpected EOF while reading voronoi faces for atom " << atomId);
      return false;
    }
    if (!parseVariableLists(line, faceData, this->ParenExtract))
    {
      vtkErrorMacro(
        "Error while parsing voronoi face data for atom " << atomId << ". Input: " << line);
      return false;
    }

    // Add points to locator:
    pointIds.resize(pointData.size());
    uniquePointIds.clear();
    for (size_t i = 0; i < pointData.size(); ++i)
    {
      const std::vector<double>& p = pointData[i];
      if (p.size() != 3)
      {
        vtkErrorMacro("Error: Tessellation point " << i << " for atom " << atomId << " has "
                                                   << p.size()
                                                   << " components. "
                                                      "Expected a 3D coordinate.");
        return false;
      }
      locator->InsertUniquePoint(&p[0], pointIds[i]);
      uniquePointIds.insert(pointIds[i]);
    }

    // Create face stream:
    faceStream.clear();
    for (size_t faceId = 0; faceId < faceData.size(); ++faceId)
    {
      const std::vector<vtkIdType>& face = faceData[faceId];
      faceStream.push_back(static_cast<vtkIdType>(face.size()));
      for (std::vector<vtkIdType>::const_iterator it = face.begin(), itEnd = face.end();
           it != itEnd; ++it)
      {
        // Convert the local point id into the dataset point id:
        vtkIdType datasetId = pointIds[*it];
        faceStream.push_back(datasetId);
      }
    }

    // Reuse pointIds to prepare a contiguous buffer of the unique pointId set
    pointIds.resize(uniquePointIds.size());
    std::copy(uniquePointIds.begin(), uniquePointIds.end(), pointIds.begin());

    // Add cell to tessellation dataset:
    voronoi->InsertNextCell(VTK_POLYHEDRON, static_cast<vtkIdType>(pointIds.size()),
      pointIds.empty() ? nullptr : &pointIds[0], static_cast<vtkIdType>(faceData.size()),
      faceStream.empty() ? nullptr : &faceStream[0]);
    tessAtomicNumbers->InsertNextValue(atom.GetAtomicNumber());
    tessAtomIds->InsertNextValue(atom.GetId());
  }

  molecule->GetVertexData()->AddArray(radii);
  voronoi->GetCellData()->SetScalars(tessAtomicNumbers);
  voronoi->GetCellData()->AddArray(tessAtomIds);

  return true;
}