// SpaFEDTe, a Template based C++ library for creating 
// Discontinuous Finite Element Spaces,
// Copyright (C) 2012 Lorenzo Alessio Botti

/* This library is free software; you can redistribute it and/or */
/* modify it under the terms of the GNU Lesser General Public */ 
/* License as published by the Free Software Foundation either */ 
/* version 3.0 of the License, or (at your option) any later version. */

/* This software is distributed in the hope that it will be useful, */
/* but WITHOUT ANY WARRANTY; without even the implied warranty of */
/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU */
/* Lesser General Public License for more details. */

/* You should have received a copy of the GNU Lesser General Public */
/* License along with this software; if not, a copy of the full */
/* GNU Lesser General Public License can be found at */
/* http://www.gnu.org/licenses/. */

// This implementation is mostly borrowed from MOAB writers

#ifndef __write_mesh_vtu__h
#define __write_mesh_vtu__h

#include "FELib.hpp"
#include "VtuMeshUtil.hpp"
#include "Mesh.hpp"
#include <string>
#include <map>
#include <list>

using namespace SpaMesh;
using namespace Connectivity;

class WriteMeshVtu 
{
public:
  template<unsigned int dim>
  void write_file(const char *filename, const Mesh<dim>& mesh);

protected:
  template<unsigned int dim>
  void vtu_write_piece(const Mesh<dim>& mesh);
  
  template<unsigned int dim>
  void init_vtu_piece(const Mesh<dim>& mesh);
  
  void close_vtu_piece();

  template<unsigned int dim>
  unsigned long computeNodesSize(const Mesh<dim>& mesh)
  {
    return mesh.nodesCoords.cols();
  }

  template<unsigned int dim>
  unsigned long computeCellsSize(const Mesh<dim>& mesh)
  {
    unsigned long size = 0;
    for (unsigned int i = 0; i < mesh.elemMesh.elemNodesConnectivity.size(); i++)
      size += mesh.elemMesh.elemNodesConnectivity[i].cols();
    for (unsigned int i = 0; i < mesh.faceMesh.elemNodesConnectivity.size(); i++)
      size += mesh.faceMesh.elemNodesConnectivity[i].cols();
    return size;
  }
private:  
  
  std::string mPartitionTagName;

  FILE* pFile;
};

#define VERBOSE false

template<unsigned int dim>
void WriteMeshVtu::write_file(const char *filename, const Mesh<dim>& mesh)
{
  std::string fileNameString = std::string(filename);
  int dotindex = fileNameString.find_last_of("."); 
  std::string procInput = fileNameString.substr(0, dotindex); 
  std::string extension = fileNameString.substr(dotindex+1); 
  
  if (FELib::n_processors() > 1)
  {
    if (strcmp(extension.c_str(),"pvtu") != 0)
      std::cerr<<"VTUWriter error: pvtu format expected (parallel vtu)."
               <<" Got "<<extension;
    char tail[1024];
    sprintf(tail,"_%02d.vtu",FELib::processor_id());
    pFile=fopen((procInput+std::string(tail)).c_str(),"w");
  }
  else if (FELib::n_processors() == 1)
  {
    if (strcmp(extension.c_str(),"vtu") != 0)
      std::cerr<<"VTUWriter error: vtu format expected (serial vtu)."
                      <<" Got "<<extension<<", switching to vtu"<<std::endl;
    pFile=fopen((std::string(procInput)+".vtu").c_str(),"w");
  }
  this->init_vtu_piece(mesh);
  this->vtu_write_piece(mesh);
  this->close_vtu_piece();

#if VERBOSE
    std::cout<<"  Number of nodes = "<<nodeSet.range.size()<<std::endl;
    std::cout<<"  Number of element types "<<exportList.size()<<std::endl;
    for (std::list<ExportUnit>::iterator iter = exportList.begin(); iter !=  exportList.end(); iter++ )
      std::cout<<"  Element type "<<iter->type<<", number of elements "<<iter->range.size()<<std::endl;
    std::cout<<"  Tags on nodes "<<nodeTags.size()<<std::endl;
    std::cout<<"  Tags on elems "<<elemTags.size()<<std::endl;
#endif
}

template<unsigned int dim>
void WriteMeshVtu::init_vtu_piece(const Mesh<dim>& mesh)
{
  // write header
  fprintf (pFile,"%s","<VTKFile type=\"UnstructuredGrid\"");
  fprintf (pFile,"%s"," version=\"0.1\"");
  fprintf (pFile,"%s\n"," byte_order=\"LittleEndian\">");
  fprintf (pFile,"%s\n","   <UnstructuredGrid>");
  // init piece
  fprintf (pFile,"%s","      <Piece  ");
  fprintf (pFile,"%s%lu%s","NumberOfPoints=\"",this->computeNodesSize(mesh),"\"  ");
  fprintf (pFile,"%s%lu%s\n","NumberOfCells=\"",this->computeCellsSize(mesh),"\">");
}

template<unsigned int dim>
void WriteMeshVtu::vtu_write_piece(const Mesh<dim>& mesh)
{
  // write mesh nodes  
  fprintf (pFile,"%s\n","         <Points>");
  fprintf (pFile,"%s%lu%s%s%s\n","            <DataArray type=\"Float",sizeof(double)*8,
                                "\" NumberOfComponents=\"","3","\" format=\"ascii\">");
  double coords[3];
  for (unsigned long i = 0; i != mesh.nodesCoords.cols(); ++i) {
    fprintf (pFile,"%f %f %f ",mesh.nodesCoords(0,i),mesh.nodesCoords(1,i),mesh.nodesCoords(2,i));
  }
  fprintf (pFile,"%s\n","");
  fprintf (pFile,"%s\n","            </DataArray>");
  fprintf (pFile,"%s\n","         </Points>");
  // write elements connectivity
  fprintf (pFile,"%s\n","         <Cells>");
  fprintf (pFile,"%s%lu%s\n","            <DataArray type=\"Int",sizeof(int)*8,
                            "\" Name=\"connectivity\" format=\"ascii\">");
  unsigned long numberOfCells = this->computeCellsSize(mesh);
  std::vector<unsigned> vtu_types( numberOfCells);
  std::vector<unsigned> vtu_offsets( numberOfCells+1 );
  std::vector<unsigned>::iterator t = vtu_types.begin();
  std::vector<unsigned>::iterator o = vtu_offsets.begin();
  for (unsigned int type = 1; type !=  ElemTypesMesh::NumberOfElemTypesMesh; type++)
  {
    bool typeExists = mesh.trueIfTypeExists(static_cast<ElemTypesMesh>(type));
    if (!typeExists)
      continue;
    const ElemNodesConnMatrixDynamic& connectivity = mesh.getElemNodesConnForType(static_cast<ElemTypesMesh>(type));
    for (unsigned long i = 0; i != connectivity.cols(); ++i)
    {
      int conn_len = connectivity.rows();
        // Get VTK type
      const MeshVtuElemType* vtk_type = VtuMeshUtil::get_vtu_type( static_cast<ElemTypesMesh>(type), conn_len );
      if (!vtk_type) {
          // try connectivity with 1 fewer node
        vtk_type = VtuMeshUtil::get_vtu_type( static_cast<ElemTypesMesh>(type), conn_len-1 );
        if (vtk_type) conn_len--;
        else {
          std::cerr<<"Vtk file format does not support elements of type "
                     "( "<<(int)type<<") with "<<conn_len<<" nodes";
        }
      }

      // Save VTK type index for later
      *t = vtk_type->vtu_type;
      ++t;
      *o += conn_len;
      *(o+1) = *o;
      ++o;
      
        // Write connectivity list
      if (vtk_type->node_order)
        for (int k = 0; k < conn_len; ++k)
          fprintf (pFile, "%d ", connectivity(vtk_type->node_order[k],i));
      else
        for (int k = 0; k < conn_len; ++k)
          fprintf (pFile, "%d ", connectivity(k,i));;
    }
  }
  fprintf (pFile,"%s\n","");
  fprintf (pFile,"%s\n","            </DataArray>");
  // write connectivity offsets 
  fprintf (pFile,"%s%lu%s\n","            <DataArray type=\"Int",sizeof(int)*8,
                            "\" Name=\"offsets\" format=\"ascii\">");
  for (unsigned long i = 0; i < numberOfCells; i ++)
    fprintf (pFile, "%d ", vtu_offsets[i]);
  fprintf (pFile,"%s\n","");
  fprintf (pFile,"%s\n","            </DataArray>");

  // write elements types 
  fprintf (pFile,"%s%lu%s\n","            <DataArray type=\"Int",sizeof(int)*8,
                            "\" Name=\"types\" format=\"ascii\">");
  for (unsigned long i = 0; i < numberOfCells; i ++)
    fprintf (pFile, "%d ", vtu_types[i]);
  fprintf (pFile,"%s\n","");
  fprintf (pFile,"%s\n","            </DataArray>");
  fprintf (pFile,"%s\n","         </Cells>");
  
}
#endif
