/*
 * Copyright (c) 2008-2018 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/
 *
 * MRtrix3 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.
 *
 * For more details, see http://www.mrtrix.org/
 */


#include "command.h"
#include "image.h"
#include "algo/threaded_loop.h"
#include "registration/warp/helpers.h"


using namespace MR;
using namespace App;


void usage ()
{
  AUTHOR = "David Raffelt (david.raffelt@florey.edu.au)";

  SYNOPSIS = "Replaces voxels in a deformation field that point to 0,0,0 with nan,nan,nan";

  DESCRIPTION
  + "This can be used in conjunction with the warpinit command to compute a MRtrix "
    "compatible deformation field from non-linear transformations generated by any other registration package.";

  ARGUMENTS
  + Argument ("in", "the input warp image.").type_image_in ()
  + Argument ("out", "the output warp image.").type_image_out ();
}


using value_type = float;


void run ()
{
  auto input = Image<value_type>::open (argument[0]).with_direct_io (3);
  Registration::Warp::check_warp (input);

  auto output = Image<value_type>::create (argument[1], input);

  auto func = [&](Image<value_type>& in, Image<value_type>& out) {
    if (Eigen::Matrix<value_type, 3, 1>(in.row(3)).norm() == 0.0) {
      for (auto l = Loop (3) (out); l; ++l)
        out.value() = NaN;
    } else {
      for (auto l = Loop (3) (in, out); l; ++l)
        out.value() = in.value();
    }
  };

  ThreadedLoop ("correcting warp", input, 0, 3)
    .run (func, input, output);
}
