//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include "VTKMeshCreation.hpp"
#include <vtkDiscreteMarchingCubes.h>
#include <vtkTriangleFilter.h>
#include <vtkCellData.h>
#include <vtkCell.h>
#include <vtkIdList.h>

#include <Forall.hpp>
#include <Mesh.hpp>
#include "VTKProgress.hpp"
#include <Information.hpp>

#include "UnorderedMap.hpp"
#include <queue>

namespace mgx {
namespace process {

namespace
{

struct Triangle {
  Triangle()
    : vs(3, vertex(0))
  {
  }

  Triangle(const Triangle& copy)
    : vs(copy.vs)
  {
  }

  const vertex& operator[](int i) const {
    return vs[i];
  }
  vertex& operator[](int i) {
    return vs[i];
  }

  std::vector<vertex> vs;
};

}

bool VTKDiscreteMarchingCube::operator()(const Store* store, Mesh* mesh)
{
  // Find the list of labels
  std::vector<int> used_labels(65536, 0);
  const HVecUS& data = store->data();
  for(size_t i = 0; i < data.size(); ++i)
    used_labels[data[i]] = 1;

  std::vector<double> values;
  values.reserve(65536);
  for(size_t i = 1; i < used_labels.size(); ++i)
    if(used_labels[i] != 0)
      values.push_back(double(i));

  SETSTATUS("Number of cells: " << values.size());

  if(values.size() == 0) {
    setErrorMessage("Error, no label found in the image");
    return false;
  }

  // Allocate objects
  VTKImageConverter* converter = VTKImageConverter::New();
  vtkDiscreteMarchingCubes* filter = vtkDiscreteMarchingCubes::New();
  vtkTriangleFilter* triangulate = vtkTriangleFilter::New();

  // Link everything together
  converter->SetStore(store);

  // Create the progress bar
  VTKProgress progress("Discrete Marching Cube");
  progress.setFilter(filter);

  // Define the list of cells to extract
  filter->SetNumberOfContours(values.size());
  for(size_t i = 0; i < values.size(); ++i)
    filter->SetValue(i, values[i]);

#if VTK_MAJOR_VERSION < 6
  filter->SetInput((vtkDataObject*)converter->GetOutput());
#else
  filter->SetInputConnection(converter->GetOutputPort());
#endif

// Force the triangulation
#if VTK_MAJOR_VERSION < 6
  triangulate->SetInput(filter->GetOutput());
#else
  triangulate->SetInputConnection(filter->GetOutputPort());
#endif
  triangulate->PassVertsOn();
  triangulate->PassLinesOn();

  // Force the calculation now
  triangulate->Update();
  vtkPolyData* output = triangulate->GetOutput();

  // Write the output to the mesh
  mesh->reset();
  mesh->updateAll();

  vvgraph& S = mesh->graph();

  typedef std::unordered_map<vtkIdType, vertex> point_map_t;
  typedef std::unordered_map<ushort, point_map_t> label_map_t;

  label_map_t label_map;

  vtkCellData* cells_data = output->GetCellData();

  SETSTATUS("Number of triangles: " << output->GetNumberOfCells());

  vtkDataArray* labels = cells_data->GetScalars();

  if(labels->GetNumberOfComponents() != 1) {
    triangulate->Delete();
    filter->Delete();
    converter->Delete();
    throw QString("Error, there is more than one attribute on the cells.");
  }

  std::queue<Triangle> to_process;

  vtkIdType nbCells = output->GetNumberOfCells();
  for(vtkIdType id = 0; id < nbCells; ++id) {
    vtkCell* cell = output->GetCell(id);
    vtkIdList* pts = cell->GetPointIds();
    vtkIdType nbPts = pts->GetNumberOfIds();
    if(nbPts != 3) {
      triangulate->Delete();
      filter->Delete();
      converter->Delete();
      throw QString("Error, the mesh is not triangular. One cell has: %1 vertices.").arg(nbPts);
    }
    ushort lab = (ushort)labels->GetComponent(id, 0);
    point_map_t& pm = label_map[lab];
    Triangle tr;
    for(vtkIdType i = 0; i < nbPts; ++i) {
      vtkIdType pid = pts->GetId(i);
      point_map_t::iterator found = pm.find(pid);
      if(found == pm.end()) {
        double p[3] = { 0, 0, 0 };
        output->GetPoint(pid, p);
        vertex v;
        v->pos = Point3f(p[0], p[1], p[2]);
        v->label = lab;
        S.insert(v);
        found = pm.insert(std::make_pair(pid, v)).first;
      }
      tr[i] = found->second;
    }
    to_process.push(tr);
  }

  for(int cnt = 0; cnt < 10 and !to_process.empty(); ++cnt) {
    std::queue<Triangle> new_to_process;

    while(!to_process.empty()) {
      Triangle tr = to_process.front();
      to_process.pop();

      int prev = 2;
      int next = 1;
      for(int i = 0; i < 3; ++i) {
        vertex v = tr[i];
        vertex p = tr[prev];
        vertex n = tr[next];
        if(S.empty(v)) {
          S.insertEdge(v, p);
          S.insertEdge(v, n);
        } else if(S.edge(v, p)) {
          if(!S.edge(v, n))
            S.spliceBefore(v, p, n);
        } else if(S.edge(v, n)) {
          if(!S.edge(v, n))
            S.spliceAfter(v, n, p);
        } else
          new_to_process.push(tr);
        next = prev;
        prev = i;
      }
    }

    std::swap(to_process, new_to_process);
  }

  if(!to_process.empty()) {
    triangulate->Delete();
    filter->Delete();
    converter->Delete();
    throw QString("Error, cannot construct the mesh");
  }

  mesh->setNormals();

  // Delete objects
  triangulate->Delete();
  filter->Delete();
  converter->Delete();
  return true;
}

REGISTER_MESH_PROCESS(VTKDiscreteMarchingCube);
} // namespace process
} // namespace mgx
