/* Copyright (c) 2008-2025 the MRtrix3 contributors.
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Covered Software is provided under this License on an "as is"
 * basis, without warranty of any kind, either expressed, implied, or
 * statutory, including, without limitation, warranties that the
 * Covered Software is free of defects, merchantable, fit for a
 * particular purpose or non-infringing.
 * See the Mozilla Public License v. 2.0 for more details.
 *
 * For more details, see http://www.mrtrix.org/.
 */

#include <mutex>

#include "command.h"
#include "image.h"
#include "progressbar.h"
#include "thread_queue.h"
#include "types.h"

#include "algo/loop.h"
#include "adapter/subset.h"

#include "connectome/connectome.h"

#include "surface/mesh.h"
#include "surface/mesh_multi.h"
#include "surface/algo/image2mesh.h"


using namespace MR;
using namespace App;
using namespace MR::Surface;


void usage ()
{

  AUTHOR = "Robert E. Smith (robert.smith@florey.edu.au)";

  SYNOPSIS = "Generate meshes from a label image";

  ARGUMENTS
  + Argument ("nodes_in", "the input node parcellation image").type_image_in()
  + Argument ("mesh_out", "the output mesh file").type_file_out();

  OPTIONS
  + Option ("blocky", "generate 'blocky' meshes with precise delineation of voxel edges, "
                      "rather than the default Marching Cubes approach");

}





void run ()
{

  Header labels_header = Header::open (argument[0]);
  Connectome::check (labels_header);
  check_3D_nonunity (labels_header);
  auto labels = labels_header.get_image<uint32_t>();

  using voxel_corner_t = Eigen::Array<int, 3, 1>;

  vector<voxel_corner_t> lower_corners, upper_corners;
  {
    for (auto i = Loop ("Importing label image", labels) (labels); i; ++i) {
      const uint32_t index = labels.value();
      if (index) {

        if (index >= lower_corners.size()) {
          lower_corners.resize (index+1, voxel_corner_t (labels.size(0), labels.size(1), labels.size(2)));
          upper_corners.resize (index+1, voxel_corner_t (-1, -1, -1));
        }

        for (size_t axis = 0; axis != 3; ++axis) {
          lower_corners[index][axis] = std::min (lower_corners[index][axis], int(labels.index (axis)));
          upper_corners[index][axis] = std::max (upper_corners[index][axis], int(labels.index (axis)));
        }

      }
    }
  }

  MeshMulti meshes (lower_corners.size(), MR::Surface::Mesh());
  meshes[0].set_name ("none");
  const bool blocky = get_options ("blocky").size();

  vector<uint32_t> missing_nodes;
  for (uint32_t i = 1; i != upper_corners.size(); ++i) {
    if (upper_corners[i][0] < 0)
      missing_nodes.push_back (i);
  }
  if (missing_nodes.size()) {
    WARN ("The following labels are absent from the parcellation image "
          "and so will have an empty mesh in the output file: "
          + join(missing_nodes, ", "));
  }

  {
    std::mutex mutex;
    ProgressBar progress ("Generating meshes from labels", lower_corners.size() - 1);
    auto loader = [&] (size_t& out) { static size_t i = 1; out = i++; return (out != lower_corners.size()); };

    auto worker = [&] (const size_t& in)
    {
      meshes[in].set_name (str(in));
      vector<int> from, dimensions;
      for (size_t axis = 0; axis != 3; ++axis) {
        from.push_back (lower_corners[in][axis]);
        dimensions.push_back (upper_corners[in][axis] - lower_corners[in][axis] + 1);
        if (dimensions.back() < 1)
          return true;
      }

      Adapter::Subset<Image<uint32_t>> subset (labels, from, dimensions);

      auto scratch = Image<bool>::scratch (subset, "Node " + str(in) + " mask");
      for (auto i = Loop (subset) (subset, scratch); i; ++i)
        scratch.value() = (subset.value() == in);

      if (blocky)
        MR::Surface::Algo::image2mesh_blocky (scratch, meshes[in]);
      else
        MR::Surface::Algo::image2mesh_mc (scratch, meshes[in], 0.5);
      std::lock_guard<std::mutex> lock (mutex);
      ++progress;
      return true;
    };

    Thread::run_queue (loader, size_t(), Thread::multi (worker));
  }

  meshes.save (argument[1]);
}



