#include "MinimalFEM/VTPWriter.h"

namespace minimalfem
{

VTPWriter::VTPWriter(std::string outfile)
  : m_outStream(outfile)
  , m_xmlWriter(m_outStream)
  , m_numPoints(0)
  , m_numPolys(0)
{
}

VTPWriter::~VTPWriter()
{
  if (m_outStream.is_open())
  {
    m_outStream << std::endl;
    m_outStream.close();
  }
}

void VTPWriter::beginPolyData(int numberPoints, int numberPolygons)
{
  m_numPoints = numberPoints;
  m_numPolys = numberPolygons;

  m_xmlWriter.openElt("VTKFile").attr("type", "PolyData").attr("version", "0.1");
  m_xmlWriter.openElt("PolyData");
  m_xmlWriter.openElt("Piece")
    .attr("NumberOfPoints", std::to_string(numberPoints))
    .attr("NumberOfPolys", std::to_string(numberPolygons));
}

void VTPWriter::writeNodes(Eigen::VectorXf& nodesX, Eigen::VectorXf& nodesY)
{
  m_xmlWriter.openElt("Points");
  m_xmlWriter.openElt("DataArray")
    .attr("type", "Float32")
    .attr("Name", "Points")
    .attr("NumberOfComponents", "3")
    .attr("format", "ascii")
    .content("");
  m_xmlWriter.writeContentLine("");

  std::stringstream ss;
  for (std::ptrdiff_t i = 0; i < nodesX.size(); ++i)
  {
    ss.str("");
    ss << nodesX[i] << " " << nodesY[i] << " 0";
    m_xmlWriter.writeContentLine(ss.str().c_str());
  }

  m_xmlWriter.closeElt().closeElt();
}

// void writePolys(int conn[])
void VTPWriter::writePolys(Eigen::VectorXi& conn)
{
  std::stringstream ss;
  m_xmlWriter.openElt("Polys");

  // Write connectivity
  m_xmlWriter.openElt("DataArray")
    .attr("type", "Int64")
    .attr("Name", "connectivity")
    .attr("format", "ascii")
    .content("");
  m_xmlWriter.writeContentLine("");

  std::size_t index = 0;
  for (std::size_t i = 0, index = 0; i < m_numPolys; ++i, index += 3)
  {
    ss.str("");
    ss << conn[index] << " " << conn[index + 1] << " " << conn[index + 2];
    m_xmlWriter.writeContentLine(ss.str().c_str());
  }
  m_xmlWriter.closeElt(); // DataArray

  // Write offsets
  m_xmlWriter.openElt("DataArray")
    .attr("type", "Int64")
    .attr("Name", "offsets")
    .attr("format", "ascii")
    .content("");
  m_xmlWriter.writeContentLine("");

  for (std::size_t i = 0, offset = 3; i < m_numPolys; ++i, offset += 3)
  {
    m_xmlWriter.writeContentLine(std::to_string(offset).c_str());
  }
  m_xmlWriter.closeElt(); // DataArray
  m_xmlWriter.closeElt(); // Polys
}

void VTPWriter::writeDisplacementData(Eigen::VectorXf& displacement)
{
  std::stringstream ss;
  // Displacement is point data (vector)
  m_xmlWriter.openElt("PointData").attr("Vectors", "displacement");

  m_xmlWriter.openElt("DataArray")
    .attr("type", "Float32")
    .attr("Name", "displacement")
    .attr("NumberOfComponents", "3")
    .attr("format", "ascii")
    .content("");
  m_xmlWriter.writeContentLine("");

  std::size_t index = 0;
  for (std::size_t i = 0, index = 0; i < m_numPoints; ++i, index += 2)
  {
    ss.str("");
    ss << displacement[index] << " " << displacement[index + 1] << " 0";
    m_xmlWriter.writeContentLine(ss.str().c_str());
  }
  m_xmlWriter.closeElt(); // DataArray
  m_xmlWriter.closeElt(); // PointData
}

void VTPWriter::writeStressData(Eigen::VectorXf& stress)
{
  // Stress is cell data (scalar)
  m_xmlWriter.openElt("CellData").attr("Scalars", "stress");
  m_xmlWriter.openElt("DataArray")
    .attr("type", "Float32")
    .attr("Name", "stress")
    .attr("format", "ascii")
    .content("");
  m_xmlWriter.writeContentLine("");

  for (std::size_t i = 0; i < m_numPolys; ++i)
  {
    m_xmlWriter.writeContentLine(std::to_string(stress[i]).c_str());
  }
  m_xmlWriter.closeElt(); // DataArray
  m_xmlWriter.closeElt(); // CellData
}

void VTPWriter::endPolyData()
{
  m_xmlWriter.closeAll();
  m_outStream << std::endl;
  m_outStream.close();
}

} // namespace minimalfem
