// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <boost/functional/hash.hpp>
#include <cstddef>
#include <optional>
#include <ostream>
#include <pup.h>
#include <tuple>
#include <unordered_map>
#include <utility>

#include "DataStructures/DataBox/DataBox.hpp"
#include "DataStructures/Tensor/EagerMath/Magnitude.hpp"
#include "DataStructures/Variables.hpp"
#include "Domain/Domain.hpp"
#include "Domain/FaceNormal.hpp"
#include "Domain/Structure/Direction.hpp"
#include "Domain/Structure/DirectionMap.hpp"
#include "Domain/Structure/ElementId.hpp"
#include "Domain/Tags.hpp"
#include "Domain/Tags/FaceNormal.hpp"
#include "Domain/Tags/Faces.hpp"
#include "Domain/Tags/SurfaceJacobian.hpp"
#include "Elliptic/BoundaryConditions/ApplyBoundaryCondition.hpp"
#include "Elliptic/DiscontinuousGalerkin/DgOperator.hpp"
#include "Elliptic/DiscontinuousGalerkin/SubdomainOperator/Tags.hpp"
#include "Elliptic/DiscontinuousGalerkin/Tags.hpp"
#include "Elliptic/Systems/GetSourcesComputer.hpp"
#include "Elliptic/Utilities/ApplyAt.hpp"
#include "NumericalAlgorithms/DiscontinuousGalerkin/MortarHelpers.hpp"
#include "NumericalAlgorithms/DiscontinuousGalerkin/Tags.hpp"
#include "NumericalAlgorithms/Spectral/Mesh.hpp"
#include "Parallel/PupStlCpp17.hpp"
#include "ParallelAlgorithms/LinearSolver/Schwarz/ElementCenteredSubdomainData.hpp"
#include "ParallelAlgorithms/LinearSolver/Schwarz/OverlapHelpers.hpp"
#include "ParallelAlgorithms/LinearSolver/Schwarz/SubdomainOperator.hpp"
#include "ParallelAlgorithms/LinearSolver/Schwarz/Tags.hpp"
#include "Utilities/EqualWithinRoundoff.hpp"
#include "Utilities/ErrorHandling/Assert.hpp"
#include "Utilities/Gsl.hpp"
#include "Utilities/TMPL.hpp"

/// Items related to the restriction of the DG operator to an element-centered
/// subdomain
namespace elliptic::dg::subdomain_operator {

namespace detail {
// Wrap the `Tag` in `LinearSolver::Schwarz::Tags::Overlaps`, except if it is
// included in `TakeFromCenterTags`.
template <typename Tag, typename Dim, typename OptionsGroup,
          typename TakeFromCenterTags>
struct make_overlap_tag_impl {
  using type = tmpl::conditional_t<
      tmpl::list_contains_v<TakeFromCenterTags, Tag>, Tag,
      LinearSolver::Schwarz::Tags::Overlaps<Tag, Dim::value, OptionsGroup>>;
};

// Wrap the `Tag` in `Tags::NeighborMortars`
template <typename Tag, typename Dim>
struct make_neighbor_mortars_tag_impl {
  using type = Tags::NeighborMortars<Tag, Dim::value>;
};
}  // namespace detail

/*!
 * \brief The elliptic DG operator on an element-centered subdomain
 *
 * This operator is a restriction of the full (linearized) DG-operator to an
 * element-centered subdomain with a few points overlap into neighboring
 * elements. It is a `LinearSolver::Schwarz::SubdomainOperator` to be used with
 * the Schwarz linear solver when it solves the elliptic DG operator.
 *
 * This operator requires the following tags are available on overlap regions
 * with neighboring elements:
 *
 * - Geometric quantities provided by
 *   `elliptic::dg::subdomain_operator::InitializeSubdomain`.
 * - All `System::fluxes_computer::argument_tags` and
 *   `System::sources_computer::argument_tags` (or
 *   `System::sources_computer_linearized::argument_tags` for nonlinear
 *   systems), except those listed in `ArgsTagsFromCenter`. The latter will be
 *   taken from the central element's DataBox, so they don't need to be made
 *   available on overlaps.
 * - The `System::fluxes_computer::argument_tags` on internal and external
 *   interfaces, except those listed in `System::fluxes_computer::volume_tags`.
 *
 * Some of these tags may require communication between elements. For example,
 * nonlinear system fields are constant background fields for the linearized DG
 * operator, but are updated in every nonlinear solver iteration. Therefore, the
 * updated nonlinear fields must be communicated across overlaps between
 * nonlinear solver iterations. To perform the communication you can use
 * `LinearSolver::Schwarz::Actions::SendOverlapFields` and
 * `LinearSolver::Schwarz::Actions::ReceiveOverlapFields`, setting
 * `RestrictToOverlap` to `false`. See
 * `LinearSolver::Schwarz::SubdomainOperator` for details.
 *
 * \par Overriding boundary conditions
 * Sometimes the subdomain operator should not use the boundary conditions that
 * have been selected when setting up the domain. For example, when the
 * subdomain operator is cached between non-linear solver iterations but the
 * boundary conditions depend on the non-linear fields, the preconditioning can
 * become ineffective (see
 * `LinearSolver::Schwarz::Actions::ResetSubdomainSolver`). Another example is
 * `elliptic::subdomain_preconditioners::MinusLaplacian`, where an auxiliary
 * Poisson system is used for preconditioning that doesn't have boundary
 * conditions set up in the domain. In these cases, the boundary conditions used
 * for the subdomain operator can be overridden with the optional
 * `override_boundary_conditions` argument to the `operator()`. If the
 * overriding boundary conditions are different from those listed in
 * `Metavariables::factory_creation`, you can supply the list of
 * boundary-condition classes to the `BoundaryConditionClasses` template
 * parameter. Note that the subdomain operator always applies the _linearized_
 * boundary conditions.
 *
 * \warning The subdomain operator hasn't been tested with periodic boundary
 * conditions so far.
 */
template <typename System, typename OptionsGroup,
          typename ArgsTagsFromCenter = tmpl::list<>,
          typename BoundaryConditionClasses = tmpl::list<>>
struct SubdomainOperator
    : LinearSolver::Schwarz::SubdomainOperator<System::volume_dim> {
 public:
  using system = System;
  using options_group = OptionsGroup;

 private:
  static constexpr size_t Dim = System::volume_dim;

  // Operator applications happen sequentially so we don't have to keep track of
  // the temporal id
  static constexpr size_t temporal_id = 0;

  // The subdomain operator always applies the linearized DG operator
  static constexpr bool linearized = true;

  using BoundaryConditionsBase = typename System::boundary_conditions_base;

  // These are the arguments that we need to retrieve from the DataBox and pass
  // to the functions in `elliptic::dg`, both on the central element and on
  // neighbors
  using prepare_args_tags =
      tmpl::list<domain::Tags::Element<Dim>, domain::Tags::Mesh<Dim>,
                 domain::Tags::InverseJacobian<Dim, Frame::ElementLogical,
                                               Frame::Inertial>,
                 domain::Tags::Faces<Dim, domain::Tags::FaceNormal<Dim>>,
                 domain::Tags::Faces<
                     Dim, domain::Tags::UnnormalizedFaceNormalMagnitude<Dim>>,
                 ::Tags::Mortars<domain::Tags::Mesh<Dim - 1>, Dim>,
                 ::Tags::Mortars<::Tags::MortarSize<Dim - 1>, Dim>>;
  using apply_args_tags = tmpl::list<
      domain::Tags::Element<Dim>, domain::Tags::Mesh<Dim>,
      domain::Tags::InverseJacobian<Dim, Frame::ElementLogical,
                                    Frame::Inertial>,
      domain::Tags::DetInvJacobian<Frame::ElementLogical, Frame::Inertial>,
      domain::Tags::Faces<Dim,
                          domain::Tags::UnnormalizedFaceNormalMagnitude<Dim>>,
      domain::Tags::Faces<Dim, domain::Tags::DetSurfaceJacobian<
                                   Frame::ElementLogical, Frame::Inertial>>,
      ::Tags::Mortars<domain::Tags::Mesh<Dim - 1>, Dim>,
      ::Tags::Mortars<::Tags::MortarSize<Dim - 1>, Dim>,
      ::Tags::Mortars<domain::Tags::DetSurfaceJacobian<Frame::ElementLogical,
                                                       Frame::Inertial>,
                      Dim>,
      elliptic::dg::Tags::PenaltyParameter, elliptic::dg::Tags::Massive>;
  using fluxes_args_tags = typename System::fluxes_computer::argument_tags;
  using sources_args_tags =
      typename elliptic::get_sources_computer<System,
                                              linearized>::argument_tags;

  // We need the fluxes args also on interfaces (internal and external). The
  // volume tags are the subset that don't have to be taken from interfaces.
  using fluxes_args_volume_tags = typename System::fluxes_computer::volume_tags;

  // These tags can be taken directly from the central element's DataBox, even
  // when evaluating neighbors
  using args_tags_from_center = tmpl::remove_duplicates<
      tmpl::push_back<ArgsTagsFromCenter, elliptic::dg::Tags::PenaltyParameter,
                      elliptic::dg::Tags::Massive>>;

  // Data on neighbors is stored in the central element's DataBox in
  // `LinearSolver::Schwarz::Tags::Overlaps` maps, so we wrap the argument tags
  // with this prefix
  using make_overlap_tag =
      detail::make_overlap_tag_impl<tmpl::_1, tmpl::pin<tmpl::size_t<Dim>>,
                                    tmpl::pin<OptionsGroup>,
                                    tmpl::pin<args_tags_from_center>>;
  using prepare_args_tags_overlap =
      tmpl::transform<prepare_args_tags, make_overlap_tag>;
  using apply_args_tags_overlap =
      tmpl::transform<apply_args_tags, make_overlap_tag>;
  using fluxes_args_tags_overlap =
      tmpl::transform<fluxes_args_tags, make_overlap_tag>;
  using sources_args_tags_overlap =
      tmpl::transform<sources_args_tags, make_overlap_tag>;
  using fluxes_args_tags_overlap_faces = tmpl::transform<
      domain::make_faces_tags<Dim, fluxes_args_tags, fluxes_args_volume_tags>,
      make_overlap_tag>;

  // We also need some data on the remote side of all neighbors' mortars. Such
  // data is stored in the central element's DataBox in `Tags::NeighborMortars`
  // maps
  using make_neighbor_mortars_tag =
      detail::make_neighbor_mortars_tag_impl<tmpl::_1,
                                             tmpl::pin<tmpl::size_t<Dim>>>;

 public:
  /// \warning This function is not thread-safe because it accesses mutable
  /// memory buffers.
  template <typename ResultTags, typename OperandTags, typename DbTagsList>
  void operator()(
      const gsl::not_null<
          LinearSolver::Schwarz::ElementCenteredSubdomainData<Dim, ResultTags>*>
          result,
      const LinearSolver::Schwarz::ElementCenteredSubdomainData<
          Dim, OperandTags>& operand,
      const db::DataBox<DbTagsList>& box,
      const std::unordered_map<std::pair<size_t, Direction<Dim>>,
                               const BoundaryConditionsBase&,
                               boost::hash<std::pair<size_t, Direction<Dim>>>>&
          override_boundary_conditions = {}) const {
    // Used to retrieve items out of the DataBox to forward to functions. This
    // replaces a long series of db::get calls.
    const auto get_items = [](const auto&... args) {
      return std::forward_as_tuple(args...);
    };

    // Retrieve data out of the DataBox
    using tags_to_retrieve = tmpl::flatten<tmpl::list<
        domain::Tags::Domain<Dim>, domain::Tags::Element<Dim>,
        ::Tags::Mortars<domain::Tags::Mesh<Dim - 1>, Dim>,
        // Data on overlaps with neighbors
        tmpl::transform<
            tmpl::flatten<tmpl::list<
                Tags::ExtrudingExtent, domain::Tags::Element<Dim>,
                domain::Tags::Mesh<Dim>,
                domain::Tags::Faces<
                    Dim, domain::Tags::UnnormalizedFaceNormalMagnitude<Dim>>,
                ::Tags::Mortars<domain::Tags::Mesh<Dim - 1>, Dim>,
                ::Tags::Mortars<::Tags::MortarSize<Dim - 1>, Dim>,
                // Data on the remote side of the neighbor's mortars
                tmpl::transform<
                    tmpl::list<
                        domain::Tags::Mesh<Dim>,
                        domain::Tags::UnnormalizedFaceNormalMagnitude<Dim>,
                        domain::Tags::Mesh<Dim - 1>,
                        ::Tags::MortarSize<Dim - 1>>,
                    make_neighbor_mortars_tag>>>,
            make_overlap_tag>>>;
    const auto& [domain, central_element, central_mortar_meshes,
                 all_overlap_extents, all_neighbor_elements,
                 all_neighbor_meshes,
                 all_neighbor_face_normal_magnitudes,
                 all_neighbor_mortar_meshes, all_neighbor_mortar_sizes,
                 all_neighbors_neighbor_meshes,
                 all_neighbors_neighbor_face_normal_magnitudes,
                 all_neighbors_neighbor_mortar_meshes,
                 all_neighbors_neighbor_mortar_sizes] =
        db::apply<tags_to_retrieve>(get_items, box);
    const auto fluxes_args = db::apply<fluxes_args_tags>(get_items, box);
    const auto sources_args = db::apply<sources_args_tags>(get_items, box);
    using FluxesArgs = std::decay_t<decltype(fluxes_args)>;
    DirectionMap<Dim, FluxesArgs> fluxes_args_on_faces{};
    for (const auto& direction : Direction<Dim>::all_directions()) {
      fluxes_args_on_faces.emplace(
          direction, elliptic::util::apply_at<
                         domain::make_faces_tags<Dim, fluxes_args_tags,
                                                 fluxes_args_volume_tags>,
                         fluxes_args_volume_tags>(get_items, box, direction));
    }

    // Setup boundary conditions
    const auto apply_boundary_condition =
        [&box, &local_domain = domain, &override_boundary_conditions](
            const ElementId<Dim>& local_element_id,
            const Direction<Dim>& local_direction, auto is_overlap,
            const auto& map_keys, const auto... fields_and_fluxes) {
          constexpr bool is_overlap_v =
              std::decay_t<decltype(is_overlap)>::value;
          // Get boundary conditions from domain, or use overridden boundary
          // conditions
          const auto& boundary_condition = [&local_domain, &local_element_id,
                                            &local_direction,
                                            &override_boundary_conditions]()
              -> const BoundaryConditionsBase& {
            if (not override_boundary_conditions.empty()) {
              const auto found_overridden_boundary_conditions =
                  override_boundary_conditions.find(
                      {local_element_id.block_id(), local_direction});
              ASSERT(found_overridden_boundary_conditions !=
                         override_boundary_conditions.end(),
                     "Overriding boundary conditions in subdomain operator, "
                     "but none is available for block "
                         << local_element_id.block_id() << " in direction "
                         << local_direction
                         << ". Make sure you have considered this external "
                            "boundary of the subdomain. If this is "
                            "intentional, add support to "
                            "elliptic::dg::SubdomainOperator.");
              return found_overridden_boundary_conditions->second;
            }
            const auto& boundary_conditions =
                local_domain.blocks()
                    .at(local_element_id.block_id())
                    .external_boundary_conditions();
            ASSERT(boundary_conditions.contains(local_direction),
                   "No boundary condition is available in block "
                       << local_element_id.block_id() << " in direction "
                       << local_direction
                       << ". Make sure you are setting up boundary conditions "
                          "when creating the domain.");
            ASSERT(
                dynamic_cast<const BoundaryConditionsBase*>(
                    boundary_conditions.at(local_direction).get()) != nullptr,
                "The boundary condition in block "
                    << local_element_id.block_id() << " in direction "
                    << local_direction
                    << " has an unexpected type. Make sure it derives off the "
                       "'boundary_conditions_base' class set in the system.");
            return dynamic_cast<const BoundaryConditionsBase&>(
                *boundary_conditions.at(local_direction));
          }();
          elliptic::apply_boundary_condition<
              linearized,
              tmpl::conditional_t<is_overlap_v, make_overlap_tag, void>,
              BoundaryConditionClasses>(boundary_condition, box, map_keys,
                                        fields_and_fluxes...);
        };

    // Check if the subdomain data is sparse, i.e. if some elements have zero
    // data. If they are, the operator is a lot cheaper to apply due to its
    // linearity.
    std::unordered_set<ElementId<Dim>> elements_in_subdomain{
        central_element.id()};
    std::unordered_set<ElementId<Dim>> elements_with_zero_data{};
    if (equal_within_roundoff(operand.element_data, 0.)) {
      elements_with_zero_data.insert(central_element.id());
    }
    for (const auto& [overlap_id, overlap_data] : operand.overlap_data) {
      elements_in_subdomain.insert(overlap_id.second);
      if (equal_within_roundoff(overlap_data, 0.)) {
        elements_with_zero_data.insert(overlap_id.second);
      }
    }
    const auto is_in_subdomain =
        [&elements_in_subdomain](const ElementId<Dim>& element_id) {
          return elements_in_subdomain.find(element_id) !=
                 elements_in_subdomain.end();
        };
    const auto data_is_zero = [&elements_with_zero_data, &is_in_subdomain](
                                  const ElementId<Dim>& element_id) {
      return elements_with_zero_data.find(element_id) !=
                 elements_with_zero_data.end() or
             // Data outside the subdomain is zero by definition
             not is_in_subdomain(element_id);
    };
    const bool central_data_is_zero = data_is_zero(central_element.id());

    // The subdomain operator essentially does two sweeps over all elements in
    // the subdomain: In the first sweep it prepares the mortar data and stores
    // them on both sides of all mortars, and in the second sweep it consumes
    // the mortar data to apply the operator. This implementation is relatively
    // simple because it can re-use the implementation for the parallel DG
    // operator. However, it is also possible to apply the subdomain operator in
    // a single sweep over all elements, incrementally building up the mortar
    // data and applying boundary corrections immediately to both adjacent
    // elements once the data is available. That approach is possibly a
    // performance optimization but requires re-implementing a lot of logic for
    // the DG operator here. It should be considered once the subdomain operator
    // has been identified as the performance bottleneck. An alternative to
    // optimizing the subdomain operator performance is to precondition the
    // subdomain solve with a _much_ simpler subdomain operator, such as a
    // finite-difference Laplacian, so fewer applications of the more expensive
    // DG subdomain operator are necessary.

    // 1. Prepare mortar data on all elements in the subdomain and store them on
    //    mortars, reorienting if needed
    //
    // Prepare central element
    const auto apply_boundary_condition_center =
        [&apply_boundary_condition, &local_central_element = central_element](
            const Direction<Dim>& local_direction,
            const auto... fields_and_fluxes) {
          apply_boundary_condition(local_central_element.id(), local_direction,
                                   std::false_type{}, local_direction,
                                   fields_and_fluxes...);
        };
    db::apply<prepare_args_tags>(
        [this, &operand](const auto&... args) {
          elliptic::dg::prepare_mortar_data<System, linearized>(
              make_not_null(&central_auxiliary_vars_),
              make_not_null(&central_auxiliary_fluxes_),
              make_not_null(&central_primal_fluxes_),
              make_not_null(&central_mortar_data_), operand.element_data,
              args...);
        },
        box, temporal_id, apply_boundary_condition_center, fluxes_args,
        sources_args, fluxes_args_on_faces, data_is_zero);
    // Prepare neighbors
    for (const auto& [direction, neighbors] : central_element.neighbors()) {
      const auto& orientation = neighbors.orientation();
      const auto direction_from_neighbor = orientation(direction.opposite());
      for (const auto& neighbor_id : neighbors) {
        const LinearSolver::Schwarz::OverlapId<Dim> overlap_id{direction,
                                                               neighbor_id};
        const auto& overlap_extent = all_overlap_extents.at(overlap_id);
        const auto& neighbor = all_neighbor_elements.at(overlap_id);
        const auto& neighbor_mesh = all_neighbor_meshes.at(overlap_id);
        const auto& mortar_id = overlap_id;
        const auto& mortar_mesh = central_mortar_meshes.at(mortar_id);
        const ::dg::MortarId<Dim> mortar_id_from_neighbor{
            direction_from_neighbor, central_element.id()};
        const bool neighbor_data_is_zero = data_is_zero(neighbor_id);

        // Intercept empty overlaps. In the unlikely case that overlaps have
        // zero extent, meaning no point of the neighbor is part of the
        // subdomain (which is fairly useless, except for testing), the
        // subdomain is identical to the central element and no communication
        // with neighbors is necessary. We can just handle the mortar between
        // central element and neighbor and continue.
        if (UNLIKELY(overlap_extent == 0)) {
          const auto& mortar_mesh_from_neighbor =
              all_neighbor_mortar_meshes.at(overlap_id)
                  .at(mortar_id_from_neighbor);
          const auto& mortar_size_from_neighbor =
              all_neighbor_mortar_sizes.at(overlap_id)
                  .at(mortar_id_from_neighbor);
          auto remote_boundary_data =
              elliptic::dg::zero_boundary_data_on_mortar<
                  typename System::primal_fields,
                  typename System::primal_fluxes>(
                  direction_from_neighbor, neighbor_mesh,
                  all_neighbor_face_normal_magnitudes.at(overlap_id)
                      .at(direction_from_neighbor),
                  mortar_mesh_from_neighbor, mortar_size_from_neighbor);
          if (not orientation.is_aligned()) {
            remote_boundary_data.orient_on_slice(
                mortar_mesh_from_neighbor.extents(),
                direction_from_neighbor.dimension(), orientation.inverse_map());
          }
          central_mortar_data_.at(mortar_id).remote_insert(
              temporal_id, std::move(remote_boundary_data));
          continue;
        }

        // Copy the central element's mortar data to the neighbor
        if (not(central_data_is_zero and neighbor_data_is_zero)) {
          auto oriented_mortar_data =
              central_mortar_data_.at(mortar_id).local_data(temporal_id);
          if (not orientation.is_aligned()) {
            oriented_mortar_data.orient_on_slice(
                mortar_mesh.extents(), direction.dimension(), orientation);
          }
          neighbors_mortar_data_[overlap_id][::dg::MortarId<Dim>{
                                                 direction_from_neighbor,
                                                 central_element.id()}]
              .remote_insert(temporal_id, std::move(oriented_mortar_data));
        }

        // Now we switch perspective to the neighbor. First, we extend the
        // overlap data to the full neighbor mesh by padding it with zeros. This
        // is necessary because spectral operators such as derivatives require
        // data on the full mesh.
        if (not neighbor_data_is_zero) {
          LinearSolver::Schwarz::extended_overlap_data(
              make_not_null(&extended_operand_vars_[overlap_id]),
              operand.overlap_data.at(overlap_id), neighbor_mesh.extents(),
              overlap_extent, direction_from_neighbor);
        }

        const auto apply_boundary_condition_neighbor =
            [&apply_boundary_condition, &local_neighbor_id = neighbor_id,
             &overlap_id](const Direction<Dim>& local_direction,
                          const auto... fields_and_fluxes) {
              apply_boundary_condition(
                  local_neighbor_id, local_direction, std::true_type{},
                  std::forward_as_tuple(overlap_id, local_direction),
                  fields_and_fluxes...);
            };

        const auto fluxes_args_on_overlap =
            elliptic::util::apply_at<fluxes_args_tags_overlap,
                                     args_tags_from_center>(get_items, box,
                                                            overlap_id);
        const auto sources_args_on_overlap =
            elliptic::util::apply_at<sources_args_tags_overlap,
                                     args_tags_from_center>(get_items, box,
                                                            overlap_id);
        DirectionMap<Dim, FluxesArgs> fluxes_args_on_overlap_faces{};
        for (const auto& neighbor_direction :
             Direction<Dim>::all_directions()) {
          fluxes_args_on_overlap_faces.emplace(
              neighbor_direction,
              elliptic::util::apply_at<fluxes_args_tags_overlap_faces,
                                       args_tags_from_center>(
                  get_items, box,
                  std::forward_as_tuple(overlap_id, neighbor_direction)));
        }

        elliptic::util::apply_at<prepare_args_tags_overlap,
                                 args_tags_from_center>(
            [this, &overlap_id](const auto&... args) {
              elliptic::dg::prepare_mortar_data<System, linearized>(
                  make_not_null(&neighbors_auxiliary_vars_[overlap_id]),
                  make_not_null(&neighbors_auxiliary_fluxes_[overlap_id]),
                  make_not_null(&neighbors_primal_fluxes_[overlap_id]),
                  make_not_null(&neighbors_mortar_data_[overlap_id]),
                  extended_operand_vars_[overlap_id], args...);
            },
            box, overlap_id, temporal_id, apply_boundary_condition_neighbor,
            fluxes_args_on_overlap, sources_args_on_overlap,
            fluxes_args_on_overlap_faces, data_is_zero);

        // Copy this neighbor's mortar data to the other side of the mortars. On
        // the other side we either have the central element, or another element
        // that may or may not be part of the subdomain.
        const auto& neighbor_mortar_meshes =
            all_neighbor_mortar_meshes.at(overlap_id);
        for (const auto& neighbor_mortar_id_and_data :
             neighbors_mortar_data_.at(overlap_id)) {
          // No structured bindings because capturing these in lambdas doesn't
          // work until C++20
          const auto& neighbor_mortar_id = neighbor_mortar_id_and_data.first;
          const auto& neighbor_mortar_data = neighbor_mortar_id_and_data.second;
          const auto& neighbor_direction = neighbor_mortar_id.first;
          const auto& neighbors_neighbor_id = neighbor_mortar_id.second;
          // No need to do anything on external boundaries
          if (neighbors_neighbor_id == ElementId<Dim>::external_boundary_id()) {
            continue;
          }
          const auto& neighbor_orientation =
              neighbor.neighbors().at(neighbor_direction).orientation();
          const auto neighbors_neighbor_direction =
              neighbor_orientation(neighbor_direction.opposite());
          const ::dg::MortarId<Dim> mortar_id_from_neighbors_neighbor{
              neighbors_neighbor_direction, neighbor_id};
          const auto send_mortar_data =
              [&neighbor_orientation, &neighbor_mortar_meshes,
               &neighbor_mortar_data, &neighbor_mortar_id, &neighbor_direction,
               &neighbor_data_is_zero,
               &data_is_zero](auto& remote_mortar_data,
                              const ElementId<Dim>& remote_element_id) {
                if (neighbor_data_is_zero and data_is_zero(remote_element_id)) {
                  return;
                }
                const auto& neighbor_mortar_mesh =
                    neighbor_mortar_meshes.at(neighbor_mortar_id);
                auto oriented_neighbor_mortar_data =
                    neighbor_mortar_data.local_data(temporal_id);
                if (not neighbor_orientation.is_aligned()) {
                  oriented_neighbor_mortar_data.orient_on_slice(
                      neighbor_mortar_mesh.extents(),
                      neighbor_direction.dimension(), neighbor_orientation);
                }
                remote_mortar_data.remote_insert(
                    temporal_id, std::move(oriented_neighbor_mortar_data));
              };
          if (neighbors_neighbor_id == central_element.id() and
              mortar_id_from_neighbors_neighbor == mortar_id) {
            send_mortar_data(central_mortar_data_.at(mortar_id),
                             central_element.id());
            continue;
          }
          // Determine whether the neighbor's neighbor overlaps with the
          // subdomain and find its overlap ID if it does.
          const auto neighbors_neighbor_overlap_id =
              [&local_all_neighbor_mortar_meshes = all_neighbor_mortar_meshes,
               &neighbors_neighbor_id, &mortar_id_from_neighbors_neighbor,
               &is_in_subdomain]()
              -> std::optional<LinearSolver::Schwarz::OverlapId<Dim>> {
            if (not is_in_subdomain(neighbors_neighbor_id)) {
              return std::nullopt;
            }
            for (const auto& [local_overlap_id, local_mortar_meshes] :
                 local_all_neighbor_mortar_meshes) {
              if (local_overlap_id.second != neighbors_neighbor_id) {
                continue;
              }
              for (const auto& local_mortar_id_and_mesh : local_mortar_meshes) {
                if (local_mortar_id_and_mesh.first ==
                    mortar_id_from_neighbors_neighbor) {
                  return local_overlap_id;
                }
              }
            }
            ERROR("The neighbor's neighbor "
                  << neighbors_neighbor_id
                  << " is part of the subdomain, but we didn't find its "
                     "overlap ID. This is a bug, so please file an issue.");
          }();
          if (neighbors_neighbor_overlap_id.has_value()) {
            // The neighbor's neighbor is part of the subdomain so we copy the
            // mortar data over. Once the loop is complete we will also have
            // received mortar data back. At that point, both neighbors have a
            // copy of each other's mortar data, which is the subject of the
            // possible optimizations mentioned above. Note that the data may
            // differ by orientations.
            send_mortar_data(
                neighbors_mortar_data_[*neighbors_neighbor_overlap_id]
                                      [mortar_id_from_neighbors_neighbor],
                neighbors_neighbor_overlap_id->second);
          } else if (not neighbor_data_is_zero) {
            // The neighbor's neighbor does not overlap with the subdomain, so
            // we don't copy mortar data and also don't expect to receive any.
            // Instead, we assume the data on it is zero and manufacture
            // appropriate remote boundary data.
            const auto& neighbors_neighbor_mortar_mesh =
                all_neighbors_neighbor_mortar_meshes.at(overlap_id)
                    .at(neighbor_mortar_id);
            auto zero_mortar_data = elliptic::dg::zero_boundary_data_on_mortar<
                typename System::primal_fields, typename System::primal_fluxes>(
                neighbors_neighbor_direction,
                all_neighbors_neighbor_meshes.at(overlap_id)
                    .at(neighbor_mortar_id),
                all_neighbors_neighbor_face_normal_magnitudes.at(overlap_id)
                    .at(neighbor_mortar_id),
                neighbors_neighbor_mortar_mesh,
                all_neighbors_neighbor_mortar_sizes.at(overlap_id)
                    .at(neighbor_mortar_id));
            // The data is zero, but auxiliary quantities such as the face
            // normal magnitude may need re-orientation
            if (not neighbor_orientation.is_aligned()) {
              zero_mortar_data.orient_on_slice(
                  neighbors_neighbor_mortar_mesh.extents(),
                  neighbors_neighbor_direction.dimension(),
                  neighbor_orientation.inverse_map());
            }
            neighbors_mortar_data_.at(overlap_id)
                .at(neighbor_mortar_id)
                .remote_insert(temporal_id, std::move(zero_mortar_data));
          }
        }  // loop over neighbor's mortars
      }    // loop over neighbors
    }      // loop over directions

    // 2. Apply the operator on all elements in the subdomain
    //
    // Apply on central element
    db::apply<apply_args_tags>(
        [this, &result, &operand](const auto&... args) {
          elliptic::dg::apply_operator<System, linearized>(
              make_not_null(&result->element_data),
              make_not_null(&central_mortar_data_), operand.element_data,
              central_primal_fluxes_, args...);
        },
        box, temporal_id, sources_args, data_is_zero);
    // Apply on neighbors
    for (const auto& [direction, neighbors] : central_element.neighbors()) {
      const auto& orientation = neighbors.orientation();
      const auto direction_from_neighbor = orientation(direction.opposite());
      for (const auto& neighbor_id : neighbors) {
        const LinearSolver::Schwarz::OverlapId<Dim> overlap_id{direction,
                                                               neighbor_id};
        const auto& overlap_extent = all_overlap_extents.at(overlap_id);
        const auto& neighbor_mesh = all_neighbor_meshes.at(overlap_id);

        if (UNLIKELY(overlap_extent == 0)) {
          continue;
        }

        elliptic::util::apply_at<apply_args_tags_overlap,
                                 args_tags_from_center>(
            [this, &overlap_id](const auto&... args) {
              elliptic::dg::apply_operator<System, linearized>(
                  make_not_null(&extended_results_[overlap_id]),
                  make_not_null(&neighbors_mortar_data_.at(overlap_id)),
                  extended_operand_vars_.at(overlap_id),
                  neighbors_primal_fluxes_.at(overlap_id), args...);
            },
            box, overlap_id, temporal_id,
            elliptic::util::apply_at<sources_args_tags_overlap,
                                     args_tags_from_center>(get_items, box,
                                                            overlap_id),
            data_is_zero);

        // Restrict the extended operator data back to the subdomain, assuming
        // we can discard any data outside the overlaps. WARNING: This
        // assumption may break with changes to the DG operator that affect its
        // sparsity. For example, multiplying the DG operator with the _full_
        // inverse mass-matrix ("massless" scheme with no "mass-lumping"
        // approximation) means that lifted boundary corrections bleed into the
        // volume.
        if (UNLIKELY(
                result->overlap_data[overlap_id].number_of_grid_points() !=
                operand.overlap_data.at(overlap_id).number_of_grid_points())) {
          result->overlap_data[overlap_id].initialize(
              operand.overlap_data.at(overlap_id).number_of_grid_points());
        }
        LinearSolver::Schwarz::data_on_overlap(
            make_not_null(&result->overlap_data[overlap_id]),
            extended_results_.at(overlap_id), neighbor_mesh.extents(),
            overlap_extent, direction_from_neighbor);
      }  // loop over neighbors
    }    // loop over directions
  }

  // NOLINTNEXTLINE(google-runtime-references)
  void pup(PUP::er& /*p*/) {}

 private:
  // Memory buffers for repeated operator applications
  // NOLINTNEXTLINE(spectre-mutable)
  mutable Variables<typename System::auxiliary_fields>
      central_auxiliary_vars_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable Variables<typename System::primal_fluxes> central_primal_fluxes_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable Variables<typename System::auxiliary_fluxes>
      central_auxiliary_fluxes_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, Variables<typename System::auxiliary_fields>>
      neighbors_auxiliary_vars_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, Variables<typename System::primal_fluxes>>
      neighbors_primal_fluxes_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, Variables<typename System::auxiliary_fluxes>>
      neighbors_auxiliary_fluxes_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, Variables<typename System::primal_fields>>
      extended_operand_vars_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, Variables<typename System::primal_fields>>
      extended_results_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable ::dg::MortarMap<
      Dim, elliptic::dg::MortarData<size_t, typename System::primal_fields,
                                    typename System::primal_fluxes>>
      central_mortar_data_{};
  // NOLINTNEXTLINE(spectre-mutable)
  mutable LinearSolver::Schwarz::OverlapMap<
      Dim, ::dg::MortarMap<Dim, elliptic::dg::MortarData<
                                    size_t, typename System::primal_fields,
                                    typename System::primal_fluxes>>>
      neighbors_mortar_data_{};
};

}  // namespace elliptic::dg::subdomain_operator


/* Copyright (C) 2016-2018 Alibaba Group Holding Limited

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <glog/logging.h>
#include <Eigen/Dense>

#include "ps-plus/server/slice.h"
#include "ps-plus/common/hashmap.h"
#include "ps-plus/server/streaming_model_utils.h"
#include "ps-plus/server/udf/simple_udf.h"
#include "ps-plus/common/initializer/constant_initializer.h"

namespace ps {
namespace server {
namespace udf {

class HashFeatureScoreFilter : public SimpleUdf<float, float, float, float, float ,int64_t> {
 public:
  virtual Status SimpleRun(UdfContext* ctx, const float& decay_rate,
                           const float& nonclk_weight, const float& clk_weight,
                           const float& train_threshold,
                           const float& export_threshold,
                           const int64_t& cur_step) const {
    Variable* variable = GetVariable(ctx);
    if (variable == nullptr) {
      return Status::ArgumentError("HashFeatureScoreFilter: Variable should not be empty");
    }

    if (variable->GetData()->Shape().IsScalar()) {
      return Status::ArgumentError("HashFeatureScoreFilter: Variable should not be Scalar");
    }

    WrapperData<HashMap>* hashmap = dynamic_cast<WrapperData<HashMap>*>(variable->GetSlicer());
    if (hashmap == nullptr) {
      return Status::ArgumentError("HashFeatureScoreFilter: Variable Should be a Hash Variable");
    }

    HashMap::HashMapStruct map;
    if (hashmap->Internal().GetHashKeys(&map) != 0) {
      return Status::Unknown("HashFeatureScoreFilter Get Hash Keys Error");
    }

    auto& items = map.items;
    if (!items.size()) {
      return Status::Ok();
    }

    variable->SetFeaExportThreshold(export_threshold);

    //1. decay fea stats
    auto stats_vec = variable->GetStatsVec();
    LOG_ASSERT(stats_vec.size()) << "fea stats num eq 0.";

    Tensor* tensor = nullptr;
    PS_CHECK_STATUS(variable->GetExistSlot("fea_stats", &tensor));

    size_t stats_num = stats_vec.size();
    auto& shape = tensor->Shape();

    LOG_ASSERT(shape.Size() == 2) << "fea_stats should have 2 dims.";
    LOG_ASSERT(shape[1] == stats_num) << "fea_stats dims 1 doesn't eq fea stats num.";

    CASES(tensor->Type(), {
        T* pstats = tensor->Raw<T>();
        T* dst;
        for (size_t i = 0; i < items.size(); ++i) {
          dst = pstats + items[i].id * stats_num;
          for (size_t j = 0; j < stats_num; ++j) {
            *(dst + j) *= decay_rate;
          }
        }
    });

    //2. build show and click vector
    Eigen::VectorXd show_vector(items.size());
    Eigen::VectorXd clk_vector(items.size());

    auto show_iter = std::find(stats_vec.begin(), stats_vec.end(), "show");
    auto click_iter = std::find(stats_vec.begin(), stats_vec.end(), "click");
    assert(show_iter != stats_vec.end() && click_iter != stats_vec.end());

    size_t show_idx = show_iter - stats_vec.begin();
    size_t click_idx = click_iter - stats_vec.begin();

    CASES(tensor->Type(), {
      T* pstats = tensor->Raw<T>();
      T* dst;
      for (size_t i = 0; i < items.size(); ++i) {
        auto& item = items[i];
        dst = pstats + item.id * stats_num;
        show_vector(i) = *(dst + show_idx);
        clk_vector(i) = *(dst + click_idx);
      }
    });

    //3. compute fea score
    auto score = (show_vector - clk_vector) * nonclk_weight + clk_vector * clk_weight;
    std::string var_name = ctx->GetVariableName();
    printf("HashFeatureScoreFilter for %s fea score min %f max %f, cur_step:%ld\n",
           var_name.c_str(), score.minCoeff(), score.maxCoeff(), cur_step);

    //4. select keys and store fea score
    std::vector<int64_t> keys;
    Tensor* fea_scores = variable->GetVariableLikeSlot("fea_score", DataType::kFloat, TensorShape(),
                                                       []{ return new initializer::ConstantInitializer(0); });
    float* pscores = fea_scores->Raw<float>();
    for (size_t i = 0; i < items.size(); ++i) {
      pscores[items[i].id] = score(i);
      if (score(i) < train_threshold) {
        keys.push_back(items[i].x);
        keys.push_back(items[i].y);
      }
    }

    //5. delete
    hashmap->Internal().Del(&(keys[0]), keys.size() / 2, 2);

    if (!ctx->GetStreamingModelArgs()->streaming_hash_model_addr.empty()) {
      PS_CHECK_STATUS(StreamingModelUtils::DelHash(ctx->GetVariableName(), keys));
    }

    printf("HashFeatureScoreFilter for %s origin= %ld, clear= %ld\n",
           var_name.c_str(), items.size(), keys.size() / 2);

    return Status::Ok();
  }
};

SIMPLE_UDF_REGISTER(HashFeatureScoreFilter, HashFeatureScoreFilter);

}
}
}


//
// Created by Kanishka Ganguly on 6/25/19.
// Copyright (c) 2019 Amazon Robotics. All rights reserved.
//

#pragma once

/** C++ Imports */
#include <algorithm>
#include <chrono>
#include <ctime>
#include <cmath>
#include <functional>
#include <future>
#include <iomanip>
#include <iostream>
#include <locale>
#include <math.h>
#include <mutex>
#include <queue>
#include <sstream>
#include <string>
#include <thread>
#include <tinyxml2.h>
#include <type_traits>
#include <vector>
#include "../include/logger.hpp"
#include "../include/prettyprint.hpp"

/** Boost imports */
#include <boost/bind.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/thread.hpp>

/** Gazebo Imports */
#include <gazebo/common/common.hh>
#include <gazebo/common/Plugin.hh>
#include <gazebo/gazebo.hh>
#include <gazebo/gazebo_client.hh>
#include <gazebo/gazebo_config.h>
#include <gazebo/msgs/msgs.hh>
#include <gazebo/physics/physics.hh>
#include <gazebo/sensors/sensors.hh>
#include <gazebo/transport/transport.hh>
#include <gazebo_plugins/gazebo_ros_utils.h>
#include <sdf/sdf.hh>

/** Eigen Imports */
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/StdVector>
#include <eigen3/unsupported/Eigen/CXX11/Tensor>

/** ROS imports */
#ifdef USE_ROS
#include <actionlib/client/simple_action_client.h>
#include <actionlib/client/terminal_state.h>
#include <actionlib/server/simple_action_server.h>
#include <controller_manager_msgs/ControllerState.h>
#include <controller_manager_msgs/ListControllers.h>
#include <control_msgs/FollowJointTrajectoryAction.h>
#include <ros/ros.h>
#include <ros/package.h>
#include <sensor_msgs/JointState.h>
#include <trajectory_msgs/JointTrajectory.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>
#include <voxel_grid_plugin/ModelControlAction.h>
#include <voxel_grid_plugin/RobotControlAction.h>
#include <voxel_grid_plugin/SimControlAction.h>
#include <voxel_grid_plugin/VoxelGridAction.h>
#include <voxel_grid_plugin/VoxelData.h>
#endif

/** Visualization */
#ifdef SHOW_VIZ
#include <cilantro/common_renderables.hpp>
#include <cilantro/point_cloud.hpp>
#include <cilantro/visualizer.hpp>
#endif

typedef Eigen::TensorFixedSize<int, Eigen::Sizes<127, 127, 127>> Tensor;    /**< @typedef Eigen::TensorFixedSizefor grid management */
typedef Eigen::Vector3f Color;                                              /**< @typedef Eigen::Vector3f for managing cilantro::PointCloud colors */
typedef Eigen::Vector3f v3f;                                                /**< @typedef Eigen::Vector3f shorthand */
typedef Eigen::Vector3i v3i;                                                /**< @typedef Eigen::Vector3i shorthand */
typedef std::vector<std::tuple<std::string, int, Color>> CustomCellState;   /**< @typedef Vector of data for voxel grid with multiple items */

/**
 * @namespace math_helper
 * @brief Small helper functions for math
 */
namespace math_helper {

/**
 * @fn inline static double DEG2RAD(const T &DEG)
 * @brief Convert degrees to radians
 * @tparam T
 * @param DEG - Degrees
 * @return Radians
 */
	template<typename T>
	inline static double DEG2RAD(const T &DEG) {
		return DEG * 0.017453293;
	}

/**
 * @fn inline static double RAD2DEG(const T &RAD)
 * @brief Convert radians to degrees
 * @tparam T
 * @param RAD - Radians
 * @return Degrees
 */
	template<typename T>
	inline static double RAD2DEG(const T &RAD) {
		return RAD * 57.29577951;
	}

	namespace gazebo_eigen_conversions {
		/**
		* @brief Converts gazebo::math::Matrix3 to Eigen::Matrix3f
		* @param input_mat - Gazebo Matrix3 to be converted
		* @return output_mat - Eigen Matrix3f after conversion
		*/
		inline static Eigen::Matrix3f GazeboMat3ToEigenMat3(const gazebo::math::Matrix3 &input_mat) {
			Eigen::Matrix3f output_mat;
			for (int i = 0; i < 3; ++i) {
				for (int j = 0; j < 3; ++j) {
					output_mat(i, j) = input_mat[i][j];
				}
			}
			return output_mat;
		}

		/**
		 * @brief Converts Eigen::Matrix3f to gazebo::math::Matrix3
		 * @param input_mat - Eigen Matrix3f to be converted
		 * @return output_mat - Gazebo Matrix3 after conversion
		 */
		inline static gazebo::math::Matrix3 EigenMat3ToGazeboMat3(const Eigen::Matrix3f &input_mat) {
			gazebo::math::Matrix3 output_mat;
			for (int i = 0; i < 3; ++i) {
				for (int j = 0; j < 3; ++j) {
					output_mat[i][j] = input_mat(i, j);
				}
			}
			return output_mat;
		}

		/**
		 * @brief Converts gazebo::math::Vector3 to Eigen::Vector3f
		 * @param input_vec - Gazebo Vector3 to be converted
		 * @return output_vec - Eigen Vector3f after conversion
		 */
		inline static Eigen::Vector3f GazeboVec3ToEigenVec3(const gazebo::math::Vector3 &input_vec) {
			Eigen::Vector3f output_vec;
			for (int i = 0; i < 3; ++i) {
				output_vec(i) = input_vec[i];
			}
			return output_vec;
		}

		/**
		 * @brief Converts Eigen::Vector3f to gazebo::math::Vector3
		 * @param input_vec - Eigen Vector3f to be converted
		 * @return output_vec - Gazebo Vector3 after conversion
		 */
		inline static gazebo::math::Vector3 EigenVec3ToGazeboVec3(const Eigen::Vector3f &input_vec) {
			gazebo::math::Vector3 output_vec;
			output_vec.x = input_vec(0);
			output_vec.y = input_vec(1);
			output_vec.z = input_vec(2);
			return output_vec;
		}

		/**
		 * @brief Converts Eigen::Quaternionf to gazebo::math::Quaternion
		 * @param input_quat - Eigen Quaternion to be converted
		 * @return output_quat - Gazebo Quaternion after conversion
		 */
		inline static gazebo::math::Quaternion EigenQuaternionToGazeboQuaternion(const Eigen::Quaternionf &input_quat) {
			gazebo::math::Quaternion output_quat;
			output_quat.Set(input_quat.w(), input_quat.x(), input_quat.y(), input_quat.z());
			return output_quat;
		}

		/**
		 * @brief Converts gazebo::math::Quaternion to Eigen::Quaternionf
		 * @param input_quat - Gazebo Quaternion to be converted
		 * @return output_quat - Eigen Quaternion after conversion
		 */
		inline static Eigen::Quaternionf GazeboQuaternionToEigenQuaternion(const gazebo::math::Quaternion &input_quat) {
			Eigen::Quaternionf output_quat;
			output_quat.w() = input_quat.w;
			output_quat.x() = input_quat.x;
			output_quat.y() = input_quat.y;
			output_quat.z() = input_quat.z;
			return output_quat;
		}
	}

	namespace eigen_conversions {
		/**
		 * @brief Convert roll, pitch, yaw to Quaternion
		 * @param r - Roll
		 * @param p - Pitch
		 * @param y - Yaw
		 * @return q - Quaternion
		 */
		inline static Eigen::Quaternionf EulerXYZToEigenQuaternion(const float &r, const float &p, const float &y) {
			Eigen::Quaternionf q;
			q = Eigen::AngleAxisf(r, Eigen::Vector3f::UnitX())
				* Eigen::AngleAxisf(p, Eigen::Vector3f::UnitY())
				* Eigen::AngleAxisf(y, Eigen::Vector3f::UnitZ());
			return q;
		}

		/**
		 * @brief Convert Quaternion to roll, pitch, yaw
		 * @param q - Quaternion
		 * @return - Vector3 of roll, pitch, yaw
		 */
		inline static Eigen::Vector3f EigenQuaternionToEulerXYZ(const Eigen::Quaternionf &q) {
			return q.toRotationMatrix().eulerAngles(0, 1, 2);
		};

		/**
		 * @brief Get 3x3 rotation and 3x1 translation components from 4x4 transformation matrix
		 * @param matrix4 Transformation matrix to decompose
		 * @param rot 3x3 rotation matrix
		 * @param pos 3x1 position vector
		 */
		inline void GetRotationTranslationEigen(const Eigen::Matrix4f &matrix4, Eigen::Matrix3f &rot, Eigen::Vector3f &pos) {
			// Rightmost 3x1 column of 4x4 matrix is translation
			pos = matrix4.block<3, 1>(0, 3);
			// Top left 3x3 matrix of 4x4 matrix is rotation
			rot = matrix4.block<3, 3>(0, 0);
		}

		/**
		 * @brief Get 3x1 translation components from 4x4 transformation matrix
		 * @param matrix4 Transformation matrix to decompose
		 * @param pos 3x1 position vector
		 */
		inline void GetTranslationEigen(const Eigen::Matrix4f &matrix4, Eigen::Vector3f &pos) {
			// Rightmost 3x1 column of 4x4 matrix is translation
			pos = matrix4.block<3, 1>(0, 3);
		}

		/**
		 * @brief Get 3x3 rotation components from 4x4 transformation matrix
		 * @param matrix4 Transformation matrix to decompose
		 * @param rot 3x3 rotation matrix
		 */
		inline void GetRotationEigen(const Eigen::Matrix4f &matrix4, Eigen::Matrix3f &rot) {
			// Top left 3x3 matrix of 4x4 matrix is rotation
			rot = matrix4.block<3, 3>(0, 0);
		}

		/**
		 * @brief Converts Eigen::Vector3f and Eigen::Matrix3f to Eigen::Matrix4f
		 * @param pos - Translation Vector3f to convert
		 * @param rot - Matrix3f to convert
		 * @param matrix4 - Converted Eigen::Matrix4f output
		 */
		inline void PoseToMatrix4(const Eigen::Vector3f &pos, const Eigen::Matrix3f &rot, Eigen::Matrix4f &matrix4) {
			matrix4.topLeftCorner(3, 3) = rot;
			matrix4.topRightCorner(3, 1) = pos;
			matrix4(3, 3) = 1;
		}
	}

	namespace gazebo_conversions {
		/**
		 * @brief Convert roll, pitch, yaw to Quaternion
		 * @param r - Roll
		 * @param p - Pitch
		 * @param y - Yaw
		 * @return q - Quaternion
		 */
		inline static gazebo::math::Quaternion EulerXYZToGazeboQuaternion(const double &r, const double &p, const double &y) {
			gazebo::math::Quaternion q;
			q.SetFromEuler(r, p, y);
			return q;
		}

		/**
		 * @brief Convert Quaternion to roll, pitch, yaw
		 * @return rpy - roll, pitch, yaw
		 */
		inline static gazebo::math::Vector3 GazeboQuaternionToEulerXYZ(const gazebo::math::Quaternion &q) {
			gazebo::math::Vector3 EulerZYX = q.GetAsEuler();
			gazebo::math::Vector3 EulerXYZ;
			EulerXYZ.Set(EulerZYX.z, EulerZYX.y, EulerZYX.x);
			return EulerXYZ;
		}

		/**
		 * @brief Converts math::Pose to math::Matrix4
		 * @param pose - Pose to convert
		 * @param matrix4 - Converted math::Matrix4 output
		 */
		inline void PoseToMatrix4(const gazebo::math::Pose &pose, gazebo::math::Matrix4 &matrix4) {
			/** Initialize to ZERO */
			matrix4 = gazebo::math::Matrix4::ZERO;
			/** Set last element to 1 for valid transformation matrix */
			matrix4[3][3] = 1;
			/** Set rotation data in transformation matrix */
			gazebo::math::Matrix3 rot_as_matrix3 = pose.rot.GetAsMatrix3();
			for (int i = 0; i < 3; ++i) {
				for (int j = 0; j < 3; ++j) {
					matrix4[i][j] = rot_as_matrix3[i][j];
				}
			}
			/** Set translation data in transformation matrix */
			for (int j = 0; j < 3; ++j) {
				matrix4[j][3] = pose.pos[j];
			}
		}

		/**
		* @brief Converts math::Vector3 and math::Quaternion to math::Matrix4
		* @param pos - Translation vector to convert
		* @param rot - Quaternion to convert
		* @param matrix4 - Converted math::Matrix4 output
		*/
		inline void PoseToMatrix4(const gazebo::math::Vector3 &pos, const gazebo::math::Quaternion &rot, gazebo::math::Matrix4 &matrix4) {
			/** Initialize to ZERO */
			matrix4 = gazebo::math::Matrix4::ZERO;
			/** Set last element to 1 for valid transformation matrix */
			matrix4[3][3] = 1;
			/** Set rotation data in transformation matrix */
			gazebo::math::Matrix3 rot_as_matrix3 = rot.GetAsMatrix3();
			for (int i = 0; i < 3; ++i) {
				for (int j = 0; j < 3; ++j) {
					matrix4[i][j] = rot_as_matrix3[i][j];
				}
			}
			/** Set translation data in transformation matrix */
			for (int j = 0; j < 3; ++j) {
				matrix4[j][3] = pos[j];
			}
		}

		/**
		 * @brief Converts math::Vector3 and math::Matrix3 to math::Matrix4
		 * @param pos - Translation vector to convert
		 * @param rot - Matrix3 to convert
		 * @param matrix4 - Converted math::Matrix4 output
		 */
		inline void PoseToMatrix4(const gazebo::math::Vector3 &pos, const gazebo::math::Matrix3 &rot, gazebo::math::Matrix4 &matrix4) {
			/** Initialize to ZERO */
			matrix4 = gazebo::math::Matrix4::ZERO;
			/** Set last element to 1 for valid transformation matrix */
			matrix4[3][3] = 1;
			/** Set rotation data in transformation matrix */
			for (int i = 0; i < 3; ++i) {
				for (int j = 0; j < 3; ++j) {
					matrix4[i][j] = rot[i][j];
				}
			}
			/** Set translation data in transformation matrix */
			for (int j = 0; j < 3; ++j) {
				matrix4[j][3] = pos[j];
			}
		}

	}

/**
* @fn inline void UpdateElement(const std::string &elem, const double &value, math::Vector3 &vec)
* @brief Updates value of math::Vector3 at index given by axis specifier "x", "y" or "z"
* @param elem - Axis to modify
* @param value - Value to set
* @param vec - Vector that will be updated
*/
	inline void UpdateElement(const std::string &elem, const double &value, gazebo::math::Vector3 &vec) {
		if (elem == "x") {
			vec.x = value;
		} else if (elem == "y") {
			vec.y = value;
		} else if (elem == "z") {
			vec.z = value;
		}
	}

/**
 * @overload
 */
	inline void UpdateElement(const int &idx, const double &value, gazebo::math::Vector3 &vec) {
		assert(idx >= 0 && idx <= 3);
		switch (idx) {
			case 0:
				vec.x = value;
				break;
			case 1:
				vec.y = value;
				break;
			case 2:
				vec.z = value;
				break;
		}
	}

}


#include "GenericReconCartesianNonLinearSpirit2DTGadget.h"
#include "hoSPIRIT2DTOperator.h"
#include "hoSPIRIT2DTDataFidelityOperator.h"
#include "hoWavelet2DTOperator.h"
#include "mri_core_spirit.h"
#include "mri_core_grappa.h"
#include "hoNDArray_reductions.h"
#include "hoGdSolver.h"
#include <boost/make_shared.hpp>

namespace Gadgetron {

    GenericReconCartesianNonLinearSpirit2DTGadget::GenericReconCartesianNonLinearSpirit2DTGadget() : BaseClass()
    {
    }

    GenericReconCartesianNonLinearSpirit2DTGadget::~GenericReconCartesianNonLinearSpirit2DTGadget()
    {
    }

    int GenericReconCartesianNonLinearSpirit2DTGadget::process_config(ACE_Message_Block* mb)
    {
        GADGET_CHECK_RETURN(BaseClass::process_config(mb) == GADGET_OK, GADGET_FAIL);

        // -------------------------------------------------

        ISMRMRD::IsmrmrdHeader h;
        try
        {
            deserialize(mb->rd_ptr(), h);
        }
        catch (...)
        {
            GDEBUG("Error parsing ISMRMRD Header");
        }

        // -------------------------------------------------
        // check the parameters
        if(this->spirit_nl_iter_max.value()==0)
        {
            this->spirit_nl_iter_max.value(15);
            GDEBUG_STREAM("spirit_iter_max: " << this->spirit_nl_iter_max.value());
        }

        if (this->spirit_nl_iter_thres.value()<FLT_EPSILON)
        {
            this->spirit_nl_iter_thres.value(0.004);
            GDEBUG_STREAM("spirit_nl_iter_thres: " << this->spirit_nl_iter_thres.value());
        }

        if (this->spirit_image_reg_lamda.value() < FLT_EPSILON)
        {
            if(this->spirit_reg_proximity_across_cha.value())
            {
                if(spirit_reg_estimate_noise_floor.value())
                {
                    this->spirit_image_reg_lamda.value(0.001);
                }
                else
                {
                    this->spirit_image_reg_lamda.value(0.0002);
                }
            }
            else
            {
                if(spirit_reg_estimate_noise_floor.value())
                {
                    this->spirit_image_reg_lamda.value(0.002);
                }
                else
                {
                    this->spirit_image_reg_lamda.value(0.00005);
                }
            }

            GDEBUG_STREAM("spirit_image_reg_lamda: " << this->spirit_image_reg_lamda.value());
        }

        if (this->spirit_reg_N_weighting_ratio.value() < FLT_EPSILON)
        {
            if(acceFactorE1_[0]<=5)
            {
                this->spirit_reg_N_weighting_ratio.value(10.0);
            }
            else
            {
                this->spirit_reg_N_weighting_ratio.value(20.0);
            }

            GDEBUG_STREAM("spirit_reg_N_weighting_ratio: " << this->spirit_reg_N_weighting_ratio.value());
        }

        return GADGET_OK;
    }

    void GenericReconCartesianNonLinearSpirit2DTGadget::perform_unwrapping(IsmrmrdReconBit& recon_bit, ReconObjType& recon_obj, size_t e)
    {
        try
        {
            size_t RO = recon_bit.data_.data_.get_size(0);
            size_t E1 = recon_bit.data_.data_.get_size(1);
            size_t E2 = recon_bit.data_.data_.get_size(2);
            size_t dstCHA = recon_bit.data_.data_.get_size(3);
            size_t N = recon_bit.data_.data_.get_size(4);
            size_t S = recon_bit.data_.data_.get_size(5);
            size_t SLC = recon_bit.data_.data_.get_size(6);

            hoNDArray< std::complex<float> >& src = recon_obj.ref_calib_;

            size_t ref_RO = src.get_size(0);
            size_t ref_E1 = src.get_size(1);
            size_t ref_E2 = src.get_size(2);
            size_t srcCHA = src.get_size(3);
            size_t ref_N = src.get_size(4);
            size_t ref_S = src.get_size(5);
            size_t ref_SLC = src.get_size(6);

            size_t convkRO = recon_obj.kernel_.get_size(0);
            size_t convkE1 = recon_obj.kernel_.get_size(1);
            size_t convkE2 = recon_obj.kernel_.get_size(2);

            recon_obj.recon_res_.data_.create(RO, E1, E2, 1, N, S, SLC);
            Gadgetron::clear(recon_obj.recon_res_.data_);
            recon_obj.full_kspace_ = recon_bit.data_.data_;
            Gadgetron::clear(recon_obj.full_kspace_);

            std::stringstream os;
            os << "encoding_" << e;
            std::string suffix = os.str();

            if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(recon_bit.data_.data_, debug_folder_full_path_ + "data_src_" + suffix); }

            // ------------------------------------------------------------------
            // compute effective acceleration factor
            // ------------------------------------------------------------------
            float effective_acce_factor(1), snr_scaling_ratio(1);
            this->compute_snr_scaling_factor(recon_bit, effective_acce_factor, snr_scaling_ratio);
            if (effective_acce_factor > 1)
            {
                Gadgetron::scal(snr_scaling_ratio, recon_bit.data_.data_);
            }

            Gadgetron::GadgetronTimer timer(false);

            // ------------------------------------------------------------------
            // compute the reconstruction
            // ------------------------------------------------------------------
            if(this->acceFactorE1_[e]<=1 && this->acceFactorE2_[e]<=1)
            {
                recon_obj.full_kspace_ = recon_bit.data_.data_;
            }
            else
            {
                hoNDArray< std::complex<float> >& kspace = recon_bit.data_.data_;
                hoNDArray< std::complex<float> >& res = recon_obj.full_kspace_;
                hoNDArray< std::complex<float> >& ref = recon_obj.ref_calib_;

                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_parallel_imaging_lamda             : " << this->spirit_parallel_imaging_lamda.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_image_reg_lamda                    : " << this->spirit_image_reg_lamda.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_data_fidelity_lamda                : " << this->spirit_data_fidelity_lamda.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_nl_iter_max                        : " << this->spirit_nl_iter_max.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_nl_iter_thres                      : " << this->spirit_nl_iter_thres.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_name                           : " << this->spirit_reg_name.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_level                          : " << this->spirit_reg_level.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_keep_approx_coeff              : " << this->spirit_reg_keep_approx_coeff.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_keep_redundant_dimension_coeff : " << this->spirit_reg_keep_redundant_dimension_coeff.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_proximity_across_cha           : " << this->spirit_reg_proximity_across_cha.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_use_coil_sen_map               : " << this->spirit_reg_use_coil_sen_map.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_RO_weighting_ratio             : " << this->spirit_reg_RO_weighting_ratio.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_E1_weighting_ratio             : " << this->spirit_reg_E1_weighting_ratio.value());
                GDEBUG_CONDITION_STREAM(this->verbose.value(), "spirit_reg_N_weighting_ratio              : " << this->spirit_reg_N_weighting_ratio.value());

                size_t slc, s;

                for (slc = 0; slc < SLC; slc++)
                {
                    for (s = 0; s < S; s++)
                    {
                        std::stringstream os;
                        os << "encoding_" << e << "_s" << s << "_slc" << slc;
                        std::string suffix_2DT = os.str();

                        // ------------------------------

                        std::complex<float>* pKspace = &kspace(0, 0, 0, 0, 0, s, slc);
                        hoNDArray< std::complex<float> > kspace2DT(RO, E1, E2, dstCHA, N, 1, 1, pKspace);

                        // ------------------------------

                        long long kernelS = s;
                        if (kernelS >= (long long)ref_S) kernelS = (long long)ref_S - 1;

                        std::complex<float>* pKIm = &recon_obj.kernelIm2D_(0, 0, 0, 0, 0, kernelS, slc);
                        hoNDArray< std::complex<float> > kIm2DT(RO, E1, srcCHA, dstCHA, ref_N, 1, 1, pKIm);

                        // ------------------------------

                        std::complex<float>* pRef = &ref(0, 0, 0, 0, 0, kernelS, slc);
                        hoNDArray< std::complex<float> > ref2DT(ref.get_size(0), ref.get_size(1), ref.get_size(2), dstCHA, ref_N, 1, 1, pRef);

                        // ------------------------------

                        hoNDArray< std::complex<float> > coilMap2DT;
                        if (recon_obj.coil_map_.get_size(6) == SLC)
                        {
                            size_t coil_S = recon_obj.coil_map_.get_size(5);
                            std::complex<float>* pCoilMap = &recon_obj.coil_map_(0, 0, 0, 0, 0, ((s>=coil_S) ? coil_S-1 : s), slc);
                            coilMap2DT.create(RO, E1, E2, dstCHA, ref_N, 1, 1, pCoilMap);
                        }

                        // ------------------------------

                        std::complex<float>* pRes = &res(0, 0, 0, 0, 0, s, slc);
                        hoNDArray< std::complex<float> > res2DT(RO, E1, E2, dstCHA, N, 1, 1, pRes);

                        // ------------------------------

                        if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(kspace2DT, debug_folder_full_path_ + "kspace2DT_nl_spirit_" + suffix_2DT); }
                        if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(kIm2DT, debug_folder_full_path_ + "kIm2DT_nl_spirit_" + suffix_2DT); }
                        if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(ref2DT, debug_folder_full_path_ + "ref2DT_nl_spirit_" + suffix_2DT); }

                        // ------------------------------

                        std::string timing_str = "SPIRIT, Non-linear unwrapping, 2DT_" + suffix_2DT;
                        if (this->perform_timing.value()) timer.start(timing_str.c_str());
                        this->perform_nonlinear_spirit_unwrapping(kspace2DT, kIm2DT, ref2DT, coilMap2DT, res2DT, e);
                        if (this->perform_timing.value()) timer.stop();

                        if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "res_nl_spirit_2DT_" + suffix_2DT); }
                    }
                }
            }

            // ---------------------------------------------------------------------
            // compute coil combined images
            // ---------------------------------------------------------------------
            if (this->perform_timing.value()) timer.start("SPIRIT Non linear, coil combination ... ");
            this->perform_spirit_coil_combine(recon_obj);
            if (this->perform_timing.value()) timer.stop();

            if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(recon_obj.recon_res_.data_, debug_folder_full_path_ + "unwrappedIm_" + suffix); }
        }
        catch (...)
        {
            GADGET_THROW("Errors happened in GenericReconCartesianNonLinearSpirit2DTGadget::perform_unwrapping(...) ... ");
        }
    }

    class solverCallBack : public hoGdSolverCallBack< hoNDArray< std::complex<float> >, hoWavelet2DTOperator< std::complex<float> > >
    {
        public:
        typedef hoGdSolverCallBack< hoNDArray< std::complex<float> >, hoWavelet2DTOperator< std::complex<float> > > BaseClass;

        solverCallBack() : BaseClass() {}
        virtual ~solverCallBack() {}

        void execute(const hoNDArray< std::complex<float> >& b, hoNDArray< std::complex<float> >& x)
        {
            typedef hoSPIRIT2DTDataFidelityOperator< std::complex<float> > SpiritOperType;
            SpiritOperType* pOper = dynamic_cast<SpiritOperType*> (this->solver_->oper_system_);
            pOper->restore_acquired_kspace(x);
        }
    };

    void GenericReconCartesianNonLinearSpirit2DTGadget::perform_nonlinear_spirit_unwrapping(hoNDArray< std::complex<float> >& kspace,
        hoNDArray< std::complex<float> >& kerIm, hoNDArray< std::complex<float> >& ref2DT, hoNDArray< std::complex<float> >& coilMap2DT, hoNDArray< std::complex<float> >& res, size_t e)
    {
        try
        {
            bool print_iter = this->spirit_print_iter.value();

            size_t RO = kspace.get_size(0);
            size_t E1 = kspace.get_size(1);
            size_t E2 = kspace.get_size(2);
            size_t CHA = kspace.get_size(3);
            size_t N = kspace.get_size(4);
            size_t S = kspace.get_size(5);
            size_t SLC = kspace.get_size(6);

            size_t ref_N = kerIm.get_size(4);
            size_t ref_S = kerIm.get_size(5);

            hoNDArray< std::complex<float> > kspaceLinear(kspace);
            res = kspace;

            // detect whether random sampling is used
            bool use_random_sampling = false;
            std::vector<long long> sampled_step_size;
            long long n, e1;
            for (n=0; n<(long long)N; n++)
            {
                long long prev_sampled_line = -1;
                for (e1=0; e1<(long long)E1; e1++)
                {
                    if(std::abs(kspace(RO/2, e1, 0, 0, 0, 0, 0))>0 && std::abs(kspace(RO/2, e1, 0, CHA-1, 0, 0, 0))>0)
                    {
                        if(prev_sampled_line>0)
                        {
                            sampled_step_size.push_back(e1 - prev_sampled_line);
                        }

                        prev_sampled_line = e1;
                    }
                }
            }

            if(sampled_step_size.size()>4)
            {
                size_t s;
                for (s=2; s<sampled_step_size.size()-1; s++)
                {
                    if(sampled_step_size[s]!=sampled_step_size[s-1])
                    {
                        use_random_sampling = true;
                        break;
                    }
                }
            }

            if(use_random_sampling)
            {
                GDEBUG_STREAM("SPIRIT Non linear, random sampling is detected ... ");
            }

            Gadgetron::GadgetronTimer timer(false);

            boost::shared_ptr< hoNDArray< std::complex<float> > > coilMap;

            bool hasCoilMap = false;
            if (coilMap2DT.get_size(0) == RO && coilMap2DT.get_size(1) == E1 && coilMap2DT.get_size(3)==CHA)
            {
                if (ref_N < N)
                {
                    coilMap = boost::shared_ptr< hoNDArray< std::complex<float> > >(new hoNDArray< std::complex<float> >(RO, E1, CHA, coilMap2DT.begin()));
                }
                else
                {
                    coilMap = boost::shared_ptr< hoNDArray< std::complex<float> > >(new hoNDArray< std::complex<float> >(RO, E1, CHA, ref_N, coilMap2DT.begin()));
                }

                hasCoilMap = true;
            }

            hoNDArray<float> gFactor;
            float gfactorMedian = 0;

            float smallest_eigen_value(0);

            // -----------------------------------------------------
            // estimate gfactor
            // -----------------------------------------------------

            // mean over N
            hoNDArray< std::complex<float> > meanKSpace;

            if(calib_mode_[e]==ISMRMRD_interleaved)
            {
                Gadgetron::compute_averaged_data_N_S(kspace, true, true, true, meanKSpace);
            }
            else
            {
                Gadgetron::compute_averaged_data_N_S(ref2DT, true, true, true, meanKSpace);
            }

            if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(meanKSpace, debug_folder_full_path_ + "spirit_nl_2DT_meanKSpace"); }

            hoNDArray< std::complex<float> > acsSrc(meanKSpace.get_size(0), meanKSpace.get_size(1), CHA, meanKSpace.begin());
            hoNDArray< std::complex<float> > acsDst(meanKSpace.get_size(0), meanKSpace.get_size(1), CHA, meanKSpace.begin());

            double grappa_reg_lamda = 0.0005;
            size_t kRO = 5;
            size_t kE1 = 4;

            hoNDArray< std::complex<float> > convKer;
            hoNDArray< std::complex<float> > kIm(RO, E1, CHA, CHA);

            Gadgetron::grappa2d_calib_convolution_kernel(acsSrc, acsDst, (size_t)this->acceFactorE1_[e], grappa_reg_lamda, kRO, kE1, convKer);
            Gadgetron::grappa2d_image_domain_kernel(convKer, RO, E1, kIm);

            hoNDArray< std::complex<float> > unmixC;

            if(hasCoilMap)
            {
                Gadgetron::grappa2d_unmixing_coeff(kIm, *coilMap, (size_t)acceFactorE1_[e], unmixC, gFactor);

                if (!debug_folder_full_path_.empty()) gt_exporter_.export_array(gFactor, debug_folder_full_path_ + "spirit_nl_2DT_gFactor");

                hoNDArray<float> gfactorSorted(gFactor);
                std::sort(gfactorSorted.begin(), gfactorSorted.begin()+RO*E1);
                gfactorMedian = gFactor((RO*E1 / 2));

                GDEBUG_STREAM("SPIRIT Non linear, the median gfactor is found to be : " << gfactorMedian);
            }

            if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(kIm, debug_folder_full_path_ + "spirit_nl_2DT_kIm");

            hoNDArray< std::complex<float> > complexIm;

            // compute linear solution as the initialization
            if(use_random_sampling)
            {
                if (this->perform_timing.value()) timer.start("SPIRIT Non linear, perform linear spirit recon ... ");
                this->perform_spirit_unwrapping(kspace, kerIm, kspaceLinear);
                if (this->perform_timing.value()) timer.stop();
            }
            else
            {
                if (this->perform_timing.value()) timer.start("SPIRIT Non linear, perform linear recon ... ");

                //size_t ref2DT_RO = ref2DT.get_size(0);
                //size_t ref2DT_E1 = ref2DT.get_size(1);

                //// mean over N
                //hoNDArray< std::complex<float> > meanKSpace;
                //Gadgetron::sum_over_dimension(ref2DT, meanKSpace, 4);

                //if (!debug_folder_full_path_.empty()) { gt_exporter_.export_array_complex(meanKSpace, debug_folder_full_path_ + "spirit_nl_2DT_meanKSpace"); }

                //hoNDArray< std::complex<float> > acsSrc(ref2DT_RO, ref2DT_E1, CHA, meanKSpace.begin());
                //hoNDArray< std::complex<float> > acsDst(ref2DT_RO, ref2DT_E1, CHA, meanKSpace.begin());

                //double grappa_reg_lamda = 0.0005;
                //size_t kRO = 5;
                //size_t kE1 = 4;

                //hoNDArray< std::complex<float> > convKer;
                //hoNDArray< std::complex<float> > kIm(RO, E1, CHA, CHA);

                //Gadgetron::grappa2d_calib_convolution_kernel(acsSrc, acsDst, (size_t)this->acceFactorE1_[e], grappa_reg_lamda, kRO, kE1, convKer);
                //Gadgetron::grappa2d_image_domain_kernel(convKer, RO, E1, kIm);

                //if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(kIm, debug_folder_full_path_ + "spirit_nl_2DT_kIm");

                Gadgetron::hoNDFFT<float>::instance()->ifft2c(kspace, complex_im_recon_buf_);
                if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(complex_im_recon_buf_, debug_folder_full_path_ + "spirit_nl_2DT_aliasedImage");

                hoNDArray< std::complex<float> > resKSpace(RO, E1, CHA, N);
                hoNDArray< std::complex<float> > aliasedImage(RO, E1, CHA, N, complex_im_recon_buf_.begin());
                Gadgetron::grappa2d_image_domain_unwrapping_aliased_image(aliasedImage, kIm, resKSpace);

                if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(resKSpace, debug_folder_full_path_ + "spirit_nl_2DT_linearImage");

                Gadgetron::hoNDFFT<float>::instance()->fft2c(resKSpace);

                memcpy(kspaceLinear.begin(), resKSpace.begin(), resKSpace.get_number_of_bytes());

                Gadgetron::apply_unmix_coeff_aliased_image(aliasedImage, unmixC, complexIm);

                if (this->perform_timing.value()) timer.stop();
            }

            if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(kspaceLinear, debug_folder_full_path_ + "spirit_nl_2DT_kspaceLinear");

            if(hasCoilMap)
            {
                if(N>=spirit_reg_minimal_num_images_for_noise_floor.value())
                {
                    // estimate the noise level

                    if(use_random_sampling)
                    {
                        Gadgetron::hoNDFFT<float>::instance()->ifft2c(kspaceLinear, complex_im_recon_buf_);

                        hoNDArray< std::complex<float> > complexLinearImage(RO, E1, CHA, N, complex_im_recon_buf_.begin());

                        Gadgetron::coil_combine(complexLinearImage, *coilMap, 2, complexIm);
                    }

                    if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(complexIm, debug_folder_full_path_ + "spirit_nl_2DT_linearImage_complexIm");

                    // if N is sufficiently large, we can estimate the noise floor by the smallest eigen value
                    hoMatrix< std::complex<float> > data;
                    data.createMatrix(RO*E1, N, complexIm.begin(), false);

                    hoNDArray< std::complex<float> > eigenVectors, eigenValues, eigenVectorsPruned;

                    // compute eigen
                    hoNDKLT< std::complex<float> > klt;
                    klt.prepare(data, (size_t)1, (size_t)0);
                    klt.eigen_value(eigenValues);

                    if (this->verbose.value())
                    {
                        GDEBUG_STREAM("SPIRIT Non linear, computes eigen values for all 2D kspaces ... ");
                        eigenValues.print(std::cout);

                        for (size_t i = 0; i<eigenValues.get_size(0); i++)
                        {
                            GDEBUG_STREAM(i << " = " << eigenValues(i));
                        }
                    }

                    smallest_eigen_value = std::sqrt( std::abs(eigenValues(N - 1).real()) / (RO*E1) );
                    GDEBUG_STREAM("SPIRIT Non linear, the smallest eigen value is : " << smallest_eigen_value);
                }
            }

            // perform nonlinear reconstruction
            {
                boost::shared_ptr<hoNDArray< std::complex<float> > > ker(new hoNDArray< std::complex<float> >(RO, E1, CHA, CHA, ref_N, kerIm.begin()));
                boost::shared_ptr<hoNDArray< std::complex<float> > > acq(new hoNDArray< std::complex<float> >(RO, E1, CHA, N, kspace.begin()));
                hoNDArray< std::complex<float> > kspaceInitial(RO, E1, CHA, N, kspaceLinear.begin());
                hoNDArray< std::complex<float> > res2DT(RO, E1, CHA, N, res.begin());

                if (this->spirit_data_fidelity_lamda.value() > 0)
                {
                    GDEBUG_STREAM("Start the NL SPIRIT data fidelity iteration - regularization strength : " << this->spirit_image_reg_lamda.value()
                                    << " - number of iteration : "                      << this->spirit_nl_iter_max.value()
                                    << " - proximity across cha : "                     << this->spirit_reg_proximity_across_cha.value()
                                    << " - redundant dimension weighting ratio : "      << this->spirit_reg_N_weighting_ratio.value()
                                    << " - using coil sen map : "                       << this->spirit_reg_use_coil_sen_map.value()
                                    << " - iter thres : "                               << this->spirit_nl_iter_thres.value()
                                    << " - wavelet name : "                             << this->spirit_reg_name.value()
                                    );

                    typedef hoGdSolver< hoNDArray< std::complex<float> >, hoWavelet2DTOperator< std::complex<float> > > SolverType;
                    SolverType solver;
                    solver.iterations_ = this->spirit_nl_iter_max.value();
                    solver.set_output_mode(this->spirit_print_iter.value() ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                    solver.grad_thres_ = this->spirit_nl_iter_thres.value();

                    if(spirit_reg_estimate_noise_floor.value() && std::abs(smallest_eigen_value)>0)
                    {
                        solver.scale_factor_ = smallest_eigen_value;
                        solver.proximal_strength_ratio_ = this->spirit_image_reg_lamda.value() * gfactorMedian;

                        GDEBUG_STREAM("SPIRIT Non linear, eigen value is used to derive the regularization strength : " << solver.proximal_strength_ratio_ << " - smallest eigen value : " << solver.scale_factor_);
                    }
                    else
                    {
                        solver.proximal_strength_ratio_ = this->spirit_image_reg_lamda.value();
                    }

                    boost::shared_ptr< hoNDArray< std::complex<float> > > x0 = boost::make_shared< hoNDArray< std::complex<float> > >(kspaceInitial);
                    solver.set_x0(x0);

                    // parallel imaging term
                    std::vector<size_t> dims;
                    acq->get_dimensions(dims);
                    hoSPIRIT2DTDataFidelityOperator< std::complex<float> > spirit(&dims);
                    spirit.set_forward_kernel(*ker, false);
                    spirit.set_acquired_points(*acq);

                    // image reg term
                    hoWavelet2DTOperator< std::complex<float> > wav3DOperator(&dims);
                    wav3DOperator.set_acquired_points(*acq);
                    wav3DOperator.scale_factor_first_dimension_ = this->spirit_reg_RO_weighting_ratio.value();
                    wav3DOperator.scale_factor_second_dimension_ = this->spirit_reg_E1_weighting_ratio.value();
                    wav3DOperator.scale_factor_third_dimension_ = this->spirit_reg_N_weighting_ratio.value();
                    wav3DOperator.with_approx_coeff_ = !this->spirit_reg_keep_approx_coeff.value();
                    wav3DOperator.change_coeffcients_third_dimension_boundary_ = !this->spirit_reg_keep_redundant_dimension_coeff.value();
                    wav3DOperator.proximity_across_cha_ = this->spirit_reg_proximity_across_cha.value();
                    wav3DOperator.no_null_space_ = true;
                    wav3DOperator.input_in_kspace_ = true;
                    wav3DOperator.select_wavelet(this->spirit_reg_name.value());

                    if (this->spirit_reg_use_coil_sen_map.value() && hasCoilMap)
                    {
                        wav3DOperator.coil_map_ = *coilMap;
                    }

                    // set operators

                    solver.oper_system_ = &spirit;
                    solver.oper_reg_ = &wav3DOperator;

                    if (this->perform_timing.value()) timer.start("NonLinear SPIRIT solver for 2DT with data fidelity ... ");
                    solver.solve(*acq, res2DT);
                    if (this->perform_timing.value()) timer.stop();

                    if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_data_fidelity_res");
                }
                else
                {
                    GDEBUG_STREAM("Start the NL SPIRIT iteration with regularization strength : "<< this->spirit_image_reg_lamda.value()
                                    << " - number of iteration : " << this->spirit_nl_iter_max.value()
                                    << " - proximity across cha : " << this->spirit_reg_proximity_across_cha.value()
                                    << " - redundant dimension weighting ratio : " << this->spirit_reg_N_weighting_ratio.value()
                                    << " - using coil sen map : " << this->spirit_reg_use_coil_sen_map.value()
                                    << " - iter thres : " << this->spirit_nl_iter_thres.value()
                                    << " - wavelet name : " << this->spirit_reg_name.value()
                                    );

                    typedef hoGdSolver< hoNDArray< std::complex<float> >, hoWavelet2DTOperator< std::complex<float> > > SolverType;
                    SolverType solver;
                    solver.iterations_ = this->spirit_nl_iter_max.value();
                    solver.set_output_mode(this->spirit_print_iter.value() ? SolverType::OUTPUT_VERBOSE : SolverType::OUTPUT_SILENT);
                    solver.grad_thres_ = this->spirit_nl_iter_thres.value();

                    if(spirit_reg_estimate_noise_floor.value() && std::abs(smallest_eigen_value)>0)
                    {
                        solver.scale_factor_ = smallest_eigen_value;
                        solver.proximal_strength_ratio_ = this->spirit_image_reg_lamda.value() * gfactorMedian;

                        GDEBUG_STREAM("SPIRIT Non linear, eigen value is used to derive the regularization strength : " << solver.proximal_strength_ratio_ << " - smallest eigen value : " << solver.scale_factor_);
                    }
                    else
                    {
                        solver.proximal_strength_ratio_ = this->spirit_image_reg_lamda.value();
                    }

                    boost::shared_ptr< hoNDArray< std::complex<float> > > x0 = boost::make_shared< hoNDArray< std::complex<float> > >(kspaceInitial);
                    solver.set_x0(x0);

                    // parallel imaging term
                    std::vector<size_t> dims;
                    acq->get_dimensions(dims);

                    hoSPIRIT2DTOperator< std::complex<float> > spirit(&dims);
                    spirit.set_forward_kernel(*ker, false);
                    spirit.set_acquired_points(*acq);
                    spirit.no_null_space_ = true;
                    spirit.use_non_centered_fft_ = false;

                    // image reg term
                    std::vector<size_t> dim;
                    acq->get_dimensions(dim);

                    hoWavelet2DTOperator< std::complex<float> > wav3DOperator(&dim);
                    wav3DOperator.set_acquired_points(*acq);
                    wav3DOperator.scale_factor_first_dimension_ = this->spirit_reg_RO_weighting_ratio.value();
                    wav3DOperator.scale_factor_second_dimension_ = this->spirit_reg_E1_weighting_ratio.value();
                    wav3DOperator.scale_factor_third_dimension_ = this->spirit_reg_N_weighting_ratio.value();
                    wav3DOperator.with_approx_coeff_ = !this->spirit_reg_keep_approx_coeff.value();
                    wav3DOperator.change_coeffcients_third_dimension_boundary_ = !this->spirit_reg_keep_redundant_dimension_coeff.value();
                    wav3DOperator.proximity_across_cha_ = this->spirit_reg_proximity_across_cha.value();
                    wav3DOperator.no_null_space_ = true;
                    wav3DOperator.input_in_kspace_ = true;
                    wav3DOperator.select_wavelet(this->spirit_reg_name.value());

                    if (this->spirit_reg_use_coil_sen_map.value() && hasCoilMap)
                    {
                        wav3DOperator.coil_map_ = *coilMap;
                    }

                    // set operators
                    solver.oper_system_ = &spirit;
                    solver.oper_reg_ = &wav3DOperator;

                    // set call back
                    solverCallBack cb;
                    cb.solver_ = &solver;
                    solver.call_back_ = &cb;

                    hoNDArray< std::complex<float> > b(kspaceInitial);
                    Gadgetron::clear(b);

                    if (this->perform_timing.value()) timer.start("NonLinear SPIRIT solver for 2DT ... ");
                    solver.solve(b, res2DT);
                    if (this->perform_timing.value()) timer.stop();

                    if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res");

                    spirit.restore_acquired_kspace(kspace, res2DT);

                    if (!debug_folder_full_path_.empty()) gt_exporter_.export_array_complex(res2DT, debug_folder_full_path_ + "spirit_nl_2DT_res_restored");
                }
            }
        }
        catch (...)
        {
            GADGET_THROW("Errors happened in GenericReconCartesianNonLinearSpirit2DTGadget::perform_nonlinear_spirit_unwrapping(...) ... ");
        }
    }

    GADGET_FACTORY_DECLARE(GenericReconCartesianNonLinearSpirit2DTGadget)
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_PREDICATES_FUNCTIONS_SIMD_SSE_SSE2_IS_LTZ_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_PREDICATES_FUNCTIONS_SIMD_SSE_SSE2_IS_LTZ_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_SSE2_SUPPORT
#include <boost/simd/toolbox/predicates/functions/is_ltz.hpp>
#include <boost/simd/sdk/meta/as_logical.hpp>
#include <boost/simd/sdk/meta/make_dependent.hpp>
#include <boost/simd/toolbox/swar/functions/details/shuffle.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION ( boost::simd::tag::is_ltz_
                                    , boost::simd::tag::sse2_
                                    , (A0)
                                    , ((simd_<int64_<A0>,boost::simd::tag::sse_>))
                                    )
  {
    typedef typename meta::as_logical<A0>::type                     result_type;
    typedef typename meta::make_dependent<int32_t,A0>::type         int32_type;
    typedef boost::simd::native<int32_type,boost::simd::tag::sse_>  type;

    BOOST_SIMD_FUNCTOR_CALL(1)
    {
      return  bitwise_cast<result_type>
              (
                details::shuffle<1,1,3,3>
                ( is_ltz(bitwise_cast<type>(a0))() )
              );
    }
  };
} } }

#endif
#endif