/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */


#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <vtkPolyData.h>
#include <vtkCellArray.h>
#include <vtkFloatArray.h>
#include <vtkPointData.h>
#include <vtkSmartPointer.h>
#include <vtkPolyDataReader.h>
#include <vtkPolyDataWriter.h>


#include <vtk/vtkmesh.hh>
#include <mia/mesh/triangulate.hh>


NS_BEGIN(vtkmia)
using namespace mia;
using namespace std;

CVtkMeshIO::CVtkMeshIO():
       CMeshIOPlugin("vtk")
{
       add_suffix(".vtk");
       add_suffix(".VTK");
       add_suffix(".vtkmesh");
       add_suffix(".VTKMESH");
}

const char *const  CVtkMeshIO::s_scale_array = "scale";
const char *const  CVtkMeshIO::s_normal_array = "normals";
const char *const  CVtkMeshIO::s_color_array = "colors";


static CVtkMeshIO::PVertexfield read_vertices( /*const*/ vtkPolyData& mesh)
{
       auto n_vertices = mesh.GetNumberOfPoints();
       cvdebug() << "Got " << n_vertices << " vertices \n";
       auto vertices = CVtkMeshIO::PVertexfield(new CVtkMeshIO::CVertexfield(n_vertices));
       // copy all vertices
       auto iv = vertices->begin();

       for (auto i = 0; i < n_vertices; i++, ++iv) {
              double p[3];
              mesh.GetPoint(i, p);
              iv->x = p[0];
              iv->y = p[1];
              iv->z = p[2];
       }

       return vertices;
}

static CVtkMeshIO::PTrianglefield read_triangles(const CVtkMeshIO::CVertexfield& /*vertices*/, /*const*/ vtkPolyData& mesh)
{
       // read all cells, if a cell is formed of more than 3 corners, then triangulate,
       // if it hes less then 3 corners, ignore it (no wireframes supported here
       auto triangles = CVtkMeshIO::PTrianglefield(new CVtkMeshIO::CTrianglefield ());
       vtkIdType npts;
       vtkIdType const *pts;
       auto strips = mesh.GetStrips();

       while (strips->GetNextCell(npts, pts)) {
              cvdebug() << "Read cell of " << npts << " vertices\n";

              if (npts == 3) {
                     CVtkMeshIO::CTrianglefield::value_type t(pts[0], pts[1], pts[2]);
                     triangles->push_back(t);
              } else { // this is a real triangle strip
                     if (npts < 3) {
                            cvinfo() << "ignoring a" << ((npts == 2) ? "n edge " : " point") << "\n";
                            continue;
                     }

                     for (int i = 2; i < npts; ++i)  {
                            CVtkMeshIO::CTrianglefield::value_type t;
                            t.x = pts[i - 2];

                            // in triangle strips, every other triangle is winded in the other direction
                            // misses test to see if we got it right
                            if (i & 1) {
                                   t.z = pts[i - 1];
                                   t.y = pts[i];
                            } else {
                                   t.y = pts[i - 1];
                                   t.z = pts[i];
                            }

                            triangles->push_back(t);
                     }
              }
       }

       // we may also read polygons using the triangulator
       // TPolyTriangulator<CVtkMeshIO::CVertexfield, vector<vtkIdType>> triangulator(vertices);
       cvdebug() << "Got " << triangles->size() << " triangles \n";
       return triangles;
}

static void read_scalars(CTriangleMesh& mesh, /*const*/ vtkPointData& point_data)
{
       auto abstract_scalars = point_data.GetScalars();

       if (!abstract_scalars)
              return;

       auto scales = dynamic_cast<vtkFloatArray *>(abstract_scalars);

       if (!scales) {
              cvinfo() << " got scales field, but is not of type 'vtkFloatArray'\n";
              return;
       }

       if (scales->GetNumberOfComponents() != 1)  {
              cvinfo() << " got scales field, but it has " << scales->GetNumberOfComponents() << " "
                       << " instead of one.\n";
              return;
       }

       auto n_scale = scales->GetNumberOfTuples();

       if (n_scale != mesh.vertices_size()) {
              cverr() << "Have " << mesh.vertices_size() << " but got " << n_scale
                      << " scale values, ignoring the scales\n";
              return;
       }

       auto is = mesh.scale_begin();

       for (auto i = 0; i < n_scale; ++i, ++is)
              *is = scales->GetValue(i);
}

static void read_normals(CTriangleMesh& mesh, /*const*/ vtkPointData& point_data)
{
       auto abstract_normals = point_data.GetNormals();

       if (!abstract_normals)
              return;

       auto normals = dynamic_cast<vtkFloatArray *>(abstract_normals);

       if (!normals) {
              cvinfo() << " got normals field, but is not of type 'vtkFloatArray'\n";
              return;
       }

       if (normals->GetNumberOfComponents() != 3)  {
              cvinfo() << " got normals field, but it has " << normals->GetNumberOfComponents() << " "
                       << " instead of three.\n";
              return;
       }

       auto n_normals = normals->GetNumberOfTuples();

       if (n_normals != mesh.vertices_size()) {
              cverr() << "Have " << mesh.vertices_size() << " but got " << n_normals
                      << " normals values, ignoring the normals\n";
              return;
       }

       auto is = mesh.normals_begin();

       for (auto i = 0; i < n_normals; ++i, ++is) {
              normals->GetTypedTuple(i, &is->x);
              cvdebug() << i << ": read normal " << *is << "\n";
       }
}

static void read_colors(CTriangleMesh& mesh, /*const*/ vtkPointData& point_data)
{
       if (!point_data.HasArray(CVtkMeshIO::s_color_array))
              return;

       auto colors = dynamic_cast<vtkFloatArray *>(point_data.GetArray(CVtkMeshIO::s_color_array));

       if (!colors) {
              cvinfo() << " got colors field, but is not of type 'vtkFloatArray'\n";
              return;
       }

       if (colors->GetNumberOfComponents() != 3)  {
              cvinfo() << " got colors field, but it has " << colors->GetNumberOfComponents() << " "
                       << " instead of three.\n";
              return;
       }

       auto n_colors = colors->GetNumberOfTuples();

       if (n_colors != mesh.vertices_size()) {
              cverr() << "Have " << mesh.vertices_size() << " but got " << n_colors
                      << " colors values, ignoring the colors field\n";
              return;
       }

       auto is = mesh.color_begin();

       for (auto i = 0; i < n_colors; ++i, ++is)
              colors->GetTypedTuple(i, &is->x);
}

PTriangleMesh CVtkMeshIO::do_load(string const&   filename) const
{
       TRACE_FUNCTION;
       auto reader = vtkSmartPointer<vtkPolyDataReader>::New();
       reader->SetFileName(filename.c_str());
       reader->Update();
       auto mesh = reader->GetOutput();

       if (!mesh)
              return PTriangleMesh();

       auto vertices = read_vertices(*mesh);
       auto triangles = read_triangles(*vertices, *mesh);
       auto out_mesh = PTriangleMesh(new CTriangleMesh(triangles, vertices));
       // now check for scales, normals, and colors
       auto point_data = mesh->GetPointData();
       read_scalars(*out_mesh, *point_data);
       read_normals(*out_mesh, *point_data);
       read_colors(*out_mesh, *point_data);
       return out_mesh;
}

bool CVtkMeshIO::do_save(string const&   filename, const CTriangleMesh& mesh) const
{
       // construct VTK mesh
       auto points = vtkSmartPointer<vtkPoints>::New();
       for_each(mesh.vertices_begin(), mesh.vertices_end(),
       [&points](const CTriangleMesh::vertex_type & v) {
              points->InsertNextPoint(v.x, v.y, v.z);
       });
       auto triangles = vtkSmartPointer<vtkCellArray>::New();
       for_each(mesh.triangles_begin(), mesh.triangles_end(),
       [&triangles](const CTriangleMesh::triangle_type & x)->void {
              vtkIdType p[] = {
                     static_cast<int>(x.x),
                     static_cast<int>(x.y),
                     static_cast<int>(x.z)
              };
              triangles->InsertNextCell(3, p);
       });
       auto data = vtkSmartPointer<vtkPolyData>::New();
       data->SetPoints(points);
       data->SetStrips(triangles);
       auto point_data = data->GetPointData();

       // now convert the optional data
       if (mesh.get_available_data() & CTriangleMesh::ed_scale) {
              auto scales = vtkSmartPointer<vtkFloatArray>::New();
              scales->SetName(s_scale_array);
              for_each(mesh.scale_begin(), mesh.scale_end(),
              [&scales](CTriangleMesh::scale_type s) {
                     scales->InsertNextValue(s);
              });
              point_data->SetScalars(scales);
       }

       if (mesh.get_available_data() & CTriangleMesh::ed_color) {
              auto colors = vtkSmartPointer<vtkFloatArray>::New();
              colors->SetName(s_color_array);
              colors->SetNumberOfComponents (3);
              for_each(mesh.color_begin(), mesh.color_end(),
              [&colors](CTriangleMesh::color_type c) -> void {
                     colors->InsertNextTuple(&c.x);
              });
              point_data->AddArray(colors);
       }

       if (mesh.get_available_data() & CTriangleMesh::ed_normal) {
              auto normals = vtkSmartPointer<vtkFloatArray>::New();
              normals->SetName(s_normal_array);
              normals->SetNumberOfComponents (3);
              for_each(mesh.normals_begin(), mesh.normals_end(),
              [&normals](CTriangleMesh::normal_type n) -> void {
                     normals->InsertNextTuple(&n.x);
              });
              point_data->SetNormals(normals);
       }

       // write it
       auto writer = vtkSmartPointer<vtkPolyDataWriter>::New();
       writer->SetFileName(filename.c_str());
       writer->SetFileTypeToBinary();
#if  VTK_MAJOR_VERSION < 6
       writer->SetInput(data);
#else
       writer->SetInputData(data);
#endif
       return writer->Write();
}

const string  CVtkMeshIO::do_get_descr() const
{
       return "A subset of VTK mesh in-and output: Triangle meshes are written, and triangle "
              "meshes and triangle strips are read. Additional per-vertex attributes are supported: 'normals', "
              "'colors' for three component colors, and 'scale' for a scalar value attached to each vertex. "
              "The data is written by the vtkPolyDataWriter in binary format.";
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
       //vista::prepare_vista_logging();
       return new CVtkMeshIO;
}

NS_END