#ifndef VTK_WRITER_HPP
#define VTK_WRITER_HPP

#include <string>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <TinyVector.hpp>

class VTKWriter
{
 private:
  const std::string m_base_filename;
  unsigned int m_file_number;
  double m_last_time;
  const double m_time_period;

 public:
  template <typename MeshType>
  void write(const MeshType& mesh,
             const double& time,
             const bool& forced_output = false)
  {
    if (time == m_last_time) return; // output already performed
    if ((time - m_last_time >= m_time_period) or forced_output) {
      m_last_time = time;
    } else {
      return;
    }
    std::ostringstream sout;
    sout << m_base_filename << '.' << std::setfill('0') << std::setw(4) << m_file_number << ".vtu" << std::ends;
    std::ofstream fout(sout.str());
    fout << "<?xml version=\"1.0\"?>\n";
    fout << "<VTKFile type=\"UnstructuredGrid\">\n";
    fout << "<UnstructuredGrid>\n";
    fout << "<Piece NumberOfPoints=\""<< mesh.numberOfNodes()
         << "\" NumberOfCells=\"" << mesh.numberOfCells() << "\">\n";

    fout << "<Points>\n";
    fout << "<DataArray Name=\"Positions\" NumberOfComponents=\"3\" type=\"Float64\" format=\"ascii\">\n";
    if constexpr(MeshType::dimension ==1) {
      const Kokkos::View<const TinyVector<1>*> xr = mesh.xr();
      for (unsigned int r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<1; ++i) {
          fout << xr[r][i] << ' ';
        }
        fout << "0 0 "; // VTK requires 3 components
      }
    } else if constexpr (MeshType::dimension ==2) {
      const Kokkos::View<const TinyVector<2>*> xr = mesh.xr();
      for (unsigned int r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<2; ++i) {
          fout << xr[r][i] << ' ';
        }
        fout << "0 "; // VTK requires 3 components
      }
    } else {
      const Kokkos::View<const TinyVector<3>*> xr = mesh.xr();
      for (unsigned int r=0; r<mesh.numberOfNodes(); ++r) {
        for (unsigned short i=0; i<3; ++i) {
          fout << xr[r][i] << ' ';
        }
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";
    fout << "</Points>\n";

    fout << "<Cells>\n";
    const Kokkos::View<const unsigned int**> cell_nodes = mesh.connectivity().cellNodes();
    const Kokkos::View<const unsigned short*> cell_nb_nodes = mesh.connectivity().cellNbNodes();

    fout << "<DataArray type=\"Int32\" Name=\"connectivity\" NumberOfComponents=\"1\" format=\"ascii\">\n";
    for (unsigned int j=0; j<mesh.numberOfCells(); ++j) {
      for (unsigned short r=0; r<cell_nb_nodes[j]; ++r) {
        fout << cell_nodes(j,r) << ' ';
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "<DataArray type=\"UInt32\" Name=\"offsets\" NumberOfComponents=\"1\" format=\"ascii\">\n";
    {
      unsigned int offset=0;
      for (unsigned int j=0; j<mesh.numberOfCells(); ++j) {
        offset += cell_nb_nodes[j];
        fout << offset << ' ';
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "<DataArray type=\"Int8\" Name=\"types\" NumberOfComponents=\"1\" format=\"ascii\">\n";
    for (unsigned int j=0; j<mesh.numberOfCells(); ++j) {
      switch (cell_nb_nodes[j]) {
        case 2: {
          fout << "3 ";
          break;
        }
        case 3: {
          fout << "5 ";
          break;
        }
        case 4: {
          fout << "9 ";
          break;
        }
        case 8: {
          fout << "12 ";
          break;
        }
        default: {
          fout << "7 ";
          break;
        }
      }
    }
    fout << '\n';
    fout << "</DataArray>\n";

    fout << "</Cells>\n";
    fout << "</Piece>\n";
    fout << "</UnstructuredGrid>\n";
    fout << "</VTKFile>\n";

    m_file_number++;
  }
  VTKWriter(const std::string& base_filename,
            const double time_period)
      : m_base_filename(base_filename),
        m_file_number  (0),
        m_last_time    (-std::numeric_limits<double>::max()),
        m_time_period  (time_period)
  {}

  ~VTKWriter() = default;
};

#endif // VTK_WRITER_HPP