/* Copyright (c) 2008-2021 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 "command.h"
#include "image.h"
#include "filter/dwi_brain_mask.h"
#include "filter/mask_clean.h"

using namespace MR;
using namespace App;

#define DEFAULT_CLEAN_SCALE 2

void usage () {
  AUTHOR = "David Raffelt (david.raffelt@florey.edu.au), Thijs Dhollander (thijs.dhollander@gmail.com) and Ben Jeurissen (ben.jeurissen@uantwerpen.be)";

SYNOPSIS = "Generates a whole brain mask from a DWI image";

DESCRIPTION
  + "All diffusion weighted and b=0 volumes are used to "
    "obtain a mask that includes both brain tissue and CSF."

  + "In a second step peninsula-like extensions, where the "
    "peninsula itself is wider than the bridge connecting it "
    "to the mask, are removed. This may help removing "
    "artefacts and non-brain parts, e.g. eyes, from "
    "the mask.";

REFERENCES
  + "Dhollander T, Raffelt D, Connelly A. " // Internal
    "Unsupervised 3-tissue response function estimation from single-shell or multi-shell diffusion MR data without a co-registered T1 image. "
    "ISMRM Workshop on Breaking the Barriers of Diffusion MRI, 2016, 5";

ARGUMENTS
   + Argument ("input",
    "the input DWI image containing volumes that are both diffusion weighted and b=0")
    .type_image_in ()

   + Argument ("output",
    "the output whole-brain mask image")
    .type_image_out ();

OPTIONS

   + Option ("clean_scale", "the maximum scale used to cut bridges. A certain maximum scale cuts "
                            "bridges up to a width (in voxels) of 2x the provided scale. Setting "
                            "this to 0 disables the mask cleaning step. (Default: " + str(DEFAULT_CLEAN_SCALE, 2) + ")")
    + Argument ("value").type_integer (0, 1e6)

   + DWI::GradImportOptions();
}


void run () {

  auto input = Image<float>::open (argument[0]).with_direct_io (3);
  auto grad = DWI::get_DW_scheme (input);

  if (input.ndim() != 4)
    throw Exception ("input DWI image must be 4D");

  Filter::DWIBrainMask dwi_brain_mask_filter (input, grad);
  dwi_brain_mask_filter.set_message ("computing dwi brain mask");
  auto temp_mask = Image<bool>::scratch (dwi_brain_mask_filter, "brain mask");
  dwi_brain_mask_filter (input, temp_mask);

  Header H_out (temp_mask);
  auto output = Image<bool>::create (argument[1], H_out);

  unsigned int scale = get_option_value<unsigned int> ("clean_scale", DEFAULT_CLEAN_SCALE);

  if (scale > 0) {
    try {
      Filter::MaskClean clean_filter (temp_mask, std::string("applying mask cleaning filter"));
      clean_filter.set_scale (scale);
      clean_filter (temp_mask, temp_mask);
      for (auto l = Loop (0,3) (input, temp_mask, output); l; ++l) {
        if (temp_mask.value()) {
          bool all_zero = true;
          for (auto vl = Loop (3) (input); vl; ++vl) {
            if (std::isfinite (float(input.value())) && input.value()) {
              all_zero = false;
              break;
            }
          }
          output.value() = !all_zero;
        } else {
          output.value() = false;
        }
      }
    } catch (...) {
      WARN("Unable to run mask cleaning filter (image is not truly 3D); skipping");
      for (auto l = Loop (0,3) (temp_mask, output); l; ++l)
        output.value() = temp_mask.value();
    }
  }
  else {
    for (auto l = Loop (0,3) (temp_mask, output); l; ++l)
      output.value() = temp_mask.value();
  }

}
