/* 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 <unsupported/Eigen/FFT>

#include "axes.h"
#include "command.h"
#include "image.h"
#include "progressbar.h"
#include "types.h"
#include "algo/threaded_loop.h"
#include "metadata/bids.h"
#include <numeric>

using namespace MR;
using namespace App;

void usage ()
{
  AUTHOR = "Ben Jeurissen (ben.jeurissen@uantwerpen.be) & J-Donald Tournier (jdtournier@gmail.com)";

  SYNOPSIS = "Remove Gibbs Ringing Artifacts";

  DESCRIPTION
    + "This application attempts to remove Gibbs ringing artefacts from MRI images using the method "
    "of local subvoxel-shifts proposed by Kellner et al. (see reference below for details)."

    + "This command is designed to run on data directly after it has been reconstructed by the scanner, "
    "before any interpolation of any kind has taken place. You should not run this command after any "
    "form of motion correction (e.g. not after dwifslpreproc). Similarly, if you intend running dwidenoise, "
    "you should run denoising before this command to not alter the noise structure, "
    "which would impact on dwidenoise's performance."

    + "For best results, any form of filtering performed by the scanner should be disabled, whether "
    "performed in the image domain or k-space. This includes elliptic filtering and other filters that "
    "are often applied to reduce Gibbs ringing artifacts. While this method can still safely be applied "
    "to such data, some residual ringing artefacts may still be present in the output."

    + "Note that this method is designed to work on images acquired with full k-space coverage. "
    "If this method is executed on data acquired with partial Fourier (eg. \"half-scan\") acceleration, "
    "it may not fully remove all ringing artifacts, "
    "and you may observe residuals of the original artifact in the partial Fourier direction. "
    "Nonetheless, application of the method is still considered safe and worthwhile. "
    "Users are however encouraged to acquired full-Fourier data where possible.";


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


  OPTIONS
  + Option ("axes",
            "select the slice axes (default: 0,1 - i.e. x-y).")
  +   Argument ("list").type_sequence_int ()

  + Option ("nshifts", "discretization of subpixel spacing (default: 20).")
  +   Argument ("value").type_integer (8, 128)

  + Option ("minW", "left border of window used for TV computation (default: 1).")
  +   Argument ("value").type_integer (0, 10)

  + Option ("maxW", "right border of window used for TV computation (default: 3).")
  +   Argument ("value").type_integer (0, 128)

  + DataType::options();


  REFERENCES
    + "Kellner, E; Dhital, B; Kiselev, V.G & Reisert, M. "
    "Gibbs-ringing artifact removal based on local subvoxel-shifts. "
    "Magnetic Resonance in Medicine, 2016, 76, 1574–1581.";

}


typedef double value_type;




class ComputeSlice
{ MEMALIGN (ComputeSlice)
  public:
    ComputeSlice (const vector<size_t>& outer_axes, const vector<size_t>& slice_axes, const int& nsh, const int& minW, const int& maxW, Image<value_type>& in, Image<value_type>& out) :
      outer_axes (outer_axes),
      slice_axes (slice_axes),
      nsh (nsh),
      minW (minW),
      maxW (maxW),
      in (in),
      out (out),
      im1 (in.size(slice_axes[0]), in.size(slice_axes[1])),
      im2 (im1.rows(), im1.cols()) {
        prealloc_FFT();
      }

    ComputeSlice (const ComputeSlice& other) :
      outer_axes (other.outer_axes),
      slice_axes (other.slice_axes),
      nsh (other.nsh),
      minW (other.minW),
      maxW (other.maxW),
      in (other.in),
      out (other.out),
      fft (),
      im1 (in.size(slice_axes[0]), in.size(slice_axes[1])),
      im2 (im1.rows(), im1.cols()) {
        prealloc_FFT();
      }


    void operator() (const Iterator& pos)
    {
      const int X = slice_axes[0];
      const int Y = slice_axes[1];
      assign_pos_of (pos, outer_axes).to (in, out);

      for (auto l = Loop (slice_axes) (in); l; ++l)
        im1 (ssize_t(in.index(X)), ssize_t(in.index(Y))) = cdouble (in.value(), 0.0);

      unring_2d ();

      for (auto l = Loop (slice_axes) (out); l; ++l)
        out.value() = im1 (ssize_t(out.index(X)), ssize_t(out.index(Y))).real();
    }

  private:
    const vector<size_t>& outer_axes;
    const vector<size_t>& slice_axes;
    const int nsh, minW, maxW;
    Image<value_type> in, out;
    Eigen::FFT<double> fft;
    Eigen::MatrixXcd im1, im2, shifted;
    Eigen::VectorXcd v;

    void prealloc_FFT () {
      // needed to avoid within-thread allocations,
      // which aren't thread-safe in FFTW:
#ifdef EIGEN_FFTW_DEFAULT
      Eigen::VectorXcd tmp (im1.rows());
      FFT (tmp);
      iFFT (tmp);
      tmp.resize (im1.cols());
      FFT (tmp);
      iFFT (tmp);
#endif
    }

    template <typename Derived> FORCE_INLINE void FFT      (Eigen::MatrixBase<Derived>&& vec) { fft.fwd (v, vec); vec = v; }
    template <typename Derived> FORCE_INLINE void FFT      (Eigen::MatrixBase<Derived>& vec) { FFT (std::move (vec)); }
    template <typename Derived> FORCE_INLINE void iFFT     (Eigen::MatrixBase<Derived>&& vec) { fft.inv (v, vec); vec = v; }
    template <typename Derived> FORCE_INLINE void iFFT     (Eigen::MatrixBase<Derived>& vec) { iFFT (std::move (vec)); }
    template <typename Derived> FORCE_INLINE void row_FFT  (Eigen::MatrixBase<Derived>& mat) { for (auto n = 0; n < mat.rows(); ++n)  FFT (mat.row(n)); }
    template <typename Derived> FORCE_INLINE void row_iFFT (Eigen::MatrixBase<Derived>& mat) { for (auto n = 0; n < mat.rows(); ++n) iFFT (mat.row(n)); }
    template <typename Derived> FORCE_INLINE void col_FFT  (Eigen::MatrixBase<Derived>& mat) { for (auto n = 0; n < mat.cols(); ++n)  FFT (mat.col(n)); }
    template <typename Derived> FORCE_INLINE void col_iFFT (Eigen::MatrixBase<Derived>& mat) { for (auto n = 0; n < mat.cols(); ++n) iFFT (mat.col(n)); }



    FORCE_INLINE void unring_2d ()
    {
      row_FFT (im1);
      col_FFT (im1);

      for (int k = 0; k < im1.cols(); k++) {
        double ck = (1.0+cos(2.0*Math::pi*(double(k)/im1.cols())))*0.5;
        for (int j = 0 ; j < im1.rows(); j++) {
          double cj = (1.0+cos(2.0*Math::pi*(double(j)/im1.rows())))*0.5;

          if (ck+cj != 0.0) {
            im2(j,k) = im1(j,k) * cj / (ck+cj);
            im1(j,k) *= ck / (ck+cj);
          }
          else
            im1(j,k) = im2(j,k) = cdouble(0.0, 0.0);
        }
      }

      row_iFFT (im1);
      col_iFFT (im2);

      unring_1d (im1);
      unring_1d (im2.transpose());

      im1 += im2;
    }





    template <typename Derived>
      FORCE_INLINE void unring_1d (Eigen::MatrixBase<Derived>&& eig)
      {
        const int n = eig.rows();
        const int numlines = eig.cols();
        shifted.resize (n, 2*nsh+1);

        vector<int> shifts(2*nsh+1);
        shifts[0] = 0;
        for (int j = 0; j < nsh; j++) {
          shifts[j+1] = j+1;
          shifts[1+nsh+j] = -(j+1);
        }

        vector<double> TV1arr(2*nsh+1);
        vector<double> TV2arr(2*nsh+1);

        for (int k = 0; k < numlines; k++) {
          shifted.col(0) = eig.col(k);

          const int maxn = (n&1) ? (n-1)/2 : n/2-1;

          for (int j = 1; j < 2*nsh+1; j++) {
            double phi = Math::pi*double(shifts[j])/double(n*nsh);
            cdouble u (std::cos(phi), std::sin(phi));
            cdouble e (1.0, 0.0);
            shifted(0,j) = shifted(0,0);

            if (!(n&1))
              shifted(n/2,j) = cdouble(0.0, 0.0);

            for (int l = 0; l < maxn; l++) {
              e = u*e;
              int L = l+1; shifted(L,j) = e * shifted(L,0);
              L = n-1-l;   shifted(L,j) = std::conj(e) * shifted(L,0);
            }
          }


          col_iFFT (shifted);

          for (int j = 0; j < 2*nsh+1; ++j) {
            TV1arr[j] = 0.0;
            TV2arr[j] = 0.0;
            for (int t = minW; t <= maxW; t++) {
              TV1arr[j] += abs (shifted((n-t)%n,j).real() - shifted((n-t-1)%n,j).real());
              TV1arr[j] += abs (shifted((n-t)%n,j).imag() - shifted((n-t-1)%n,j).imag());
              TV2arr[j] += abs (shifted((n+t)%n,j).real() - shifted((n+t+1)%n,j).real());
              TV2arr[j] += abs (shifted((n+t)%n,j).imag() - shifted((n+t+1)%n,j).imag());
            }
          }

          for (int l = 0; l < n; ++l) {
            double minTV = std::numeric_limits<double>::max();
            int minidx = 0;
            for (int j = 0; j < 2*nsh+1; ++j) {

              if (TV1arr[j] < minTV) {
                minTV = TV1arr[j];
                minidx = j;
              }
              if (TV2arr[j] < minTV) {
                minTV = TV2arr[j];
                minidx = j;
              }

              TV1arr[j] += abs (shifted((l-minW+1+n)%n,j).real() - shifted((l-(minW  )+n)%n,j).real());
              TV1arr[j] -= abs (shifted((l-maxW  +n)%n,j).real() - shifted((l-(maxW+1)+n)%n,j).real());
              TV2arr[j] += abs (shifted((l+maxW+1+n)%n,j).real() - shifted((l+(maxW+2)+n)%n,j).real());
              TV2arr[j] -= abs (shifted((l+minW  +n)%n,j).real() - shifted((l+(minW+1)+n)%n,j).real());

              TV1arr[j] += abs (shifted((l-minW+1+n)%n,j).imag() - shifted((l-(minW  )+n)%n,j).imag());
              TV1arr[j] -= abs (shifted((l-maxW  +n)%n,j).imag() - shifted((l-(maxW+1)+n)%n,j).imag());
              TV2arr[j] += abs (shifted((l+maxW+1+n)%n,j).imag() - shifted((l+(maxW+2)+n)%n,j).imag());
              TV2arr[j] -= abs (shifted((l+minW  +n)%n,j).imag() - shifted((l+(minW+1)+n)%n,j).imag());
            }

            double a0r = shifted((l-1+n)%n,minidx).real();
            double a1r = shifted(l,minidx).real();
            double a2r = shifted((l+1+n)%n,minidx).real();
            double a0i = shifted((l-1+n)%n,minidx).imag();
            double a1i = shifted(l,minidx).imag();
            double a2i = shifted((l+1+n)%n,minidx).imag();
            double s = double(shifts[minidx])/(2.0*nsh);

            if (s > 0.0)
              eig(l,k) = cdouble (a1r*(1.0-s) + a0r*s, a1i*(1.0-s) + a0i*s);
            else
              eig(l,k) = cdouble (a1r*(1.0+s) - a2r*s, a1i*(1.0+s) - a2i*s);
          }
        }
      }

    template <typename Derived>
      FORCE_INLINE void unring_1d (Eigen::MatrixBase<Derived>& eig) { unring_1d (std::move (eig)); }

};







void run ()
{
  const int nshifts = App::get_option_value ("nshifts", 20);
  const int minW = App::get_option_value ("minW", 1);
  const int maxW = App::get_option_value ("maxW", 3);

  if (minW >= maxW)
    throw Exception ("minW must be smaller than maxW");

  auto in = Image<value_type>::open (argument[0]);
  Header header (in);

  header.datatype() = DataType::from_command_line (DataType::Float32);
  auto out = Image<value_type>::create (argument[1], header);

  vector<size_t> slice_axes = { 0, 1 };
  auto opt = get_options ("axes");
  const bool axes_set_manually = opt.size();
  if (opt.size()) {
    vector<uint32_t> axes = parse_ints<uint32_t> (opt[0][0]);
    if (axes.size() != 2)
      throw Exception ("slice axes must be specified as a comma-separated 2-vector");
    if (size_t(std::max (axes[0], axes[1])) >= header.ndim())
      throw Exception ("slice axes must be within the dimensionality of the image");
    if (axes[0] == axes[1])
      throw Exception ("two independent slice axes must be specified");
    slice_axes = { size_t(axes[0]), size_t(axes[1]) };
  }

  auto slice_encoding_it = header.keyval().find ("SliceEncodingDirection");
  if (slice_encoding_it != header.keyval().end()) {
    try {
      const Metadata::BIDS::axis_vector_type slice_encoding_axis_onehot = Metadata::BIDS::axisid2vector (slice_encoding_it->second);
      vector<size_t> auto_slice_axes = { 0, 0 };
      if (slice_encoding_axis_onehot[0])
        auto_slice_axes = { 1, 2 };
      else if (slice_encoding_axis_onehot[1])
        auto_slice_axes = { 0, 2 };
      else if (slice_encoding_axis_onehot[2])
        auto_slice_axes = { 0, 1 };
      else
        throw Exception ("Fatal error: Invalid slice axis one-hot encoding [ " + str(slice_encoding_axis_onehot.transpose()) + " ]");
      if (axes_set_manually) {
        if (slice_axes == auto_slice_axes) {
          INFO ("User's manual selection of within-slice axes consistent with \"SliceEncodingDirection\" field in image header");
        } else {
          WARN ("Within-slice axes set using -axes option will be used, but is inconsistent with SliceEncodingDirection field present in image header (" + slice_encoding_it->second + ")");
        }
      } else {
        if (slice_axes == auto_slice_axes) {
          INFO ("\"SliceEncodingDirection\" field in image header is consistent with default selection of first two axes as being within-slice");
        } else {
          slice_axes = auto_slice_axes;
          CONSOLE ("Using axes { " + str(slice_axes[0]) + ", " + str(slice_axes[1]) + " } for Gibbs ringing removal based on \"SliceEncodingDirection\" field in image header");
        }
      }
    } catch (...) {
      WARN ("Invalid value for field \"SliceEncodingDirection\" in image header (" + slice_encoding_it->second + "); ignoring");
    }
  }

  // build vector of outer axes:
  vector<size_t> outer_axes (header.ndim());
  std::iota (outer_axes.begin(), outer_axes.end(), 0);
  for (const auto axis : slice_axes) {
    auto it = std::find (outer_axes.begin(), outer_axes.end(), axis);
    if (it == outer_axes.end())
      throw Exception ("slice axis out of range!");
    outer_axes.erase (it);
  }

  ThreadedLoop ("performing Gibbs ringing removal", in, outer_axes, slice_axes)
    .run_outer (ComputeSlice (outer_axes, slice_axes, nshifts, minW, maxW, in, out));
}