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

#pragma once

#include <array>
#include <boost/functional/hash.hpp>  // IWYU pragma: keep
#include <cstddef>
#include <functional>
#include <limits>
#include <unordered_map>
#include <utility>
#include <vector>

#include "DataStructures/ApplyMatrices.hpp"
#include "DataStructures/DataVector.hpp"
#include "DataStructures/Index.hpp"
#include "DataStructures/Matrix.hpp"
#include "DataStructures/Tags.hpp"       // IWYU pragma: keep
#include "DataStructures/Variables.hpp"  // IWYU pragma: keep
#include "Domain/Structure/Direction.hpp"
#include "Domain/Structure/DirectionMap.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoGridHelpers.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoHelpers.hpp"
#include "Evolution/DiscontinuousGalerkin/Limiters/WenoOscillationIndicator.hpp"
#include "NumericalAlgorithms/LinearOperators/MeanValue.hpp"
#include "Utilities/Algorithm.hpp"
#include "Utilities/EqualWithinRoundoff.hpp"
#include "Utilities/ErrorHandling/Assert.hpp"
#include "Utilities/Gsl.hpp"

/// \cond
template <size_t VolumeDim>
class Element;
template <size_t VolumeDim>
class ElementId;
template <size_t>
class Mesh;
/// \endcond

namespace Limiters::Weno_detail {

// Caching class that holds various precomputed terms used in the constrained-
// fit algebra on each element.
//
// The terms to precompute and cache depend on the configuration of neighbors to
// the element (how many neighbors, in which directions, of what resolution).
// Each instance of the caching class represents one particular neighbor
// configuration, and typical simulations will create many instances of
// the caching class: one instance for each neighbor-element configuration
// present in the computational domain.
//
// Because the current implementation of the HWENO fitting makes the simplifying
// restrictions of no h/p-refinement, the structure of the caching is also
// simplified. In particular, with no h/p-refinement, the allowable neighbor
// configurations satisfy,
// 1. the element has at most one neighbor per dimension, AND
// 2. the mesh on every neighbor is the same as on the element.
// With these restrictions, the number of independent configurations to be
// cached is greatly reduced. Because an element has 2*VolumeDim boundaries,
// it has 2^(2*VolumeDim) possible configurations of internal/external
// boundaries, and therefore there are 2^(2*VolumeDim) configurations to cache.
// Different elements with the same configuration of internal/external
// boundaries vs. direction can share the same caching-class instance.
//
// Each instance of the caching class holds several terms, some of which also
// depend on the neighbor configuration. The restriction of no h/p-refinement
// therefore simplifies each cache instance, as well as reducing the necessary
// number of instances.
//
// The most complicated term to cache is the A^{-1} matrix. The complexity
// arises because the matrix can take many values depending on (runtime) choices
// made for each individual HWNEO fit: which element is the primary neighbor,
// and which element(s) are the excluded neighbors.
// Fits with one excluded neighbor are overwhelmingly the most likely, and the
// cache is designed for this particular case. Fits with no excluded neighbors
// can arise in elements with only one neighbor (always the primary neighbor);
// these cases are also handled. However, the very rare case in which more than
// one neighbor is excluded is not handled by the cache; the caller must compute
// A^{-1} from scratch if this scenario arises.
template <size_t VolumeDim>
class ConstrainedFitCache {
 public:
  ConstrainedFitCache(const Element<VolumeDim>& element,
                      const Mesh<VolumeDim>& mesh) noexcept;

  // Valid calls must satisfy these constraints on directions_to_exclude:
  // - the vector is empty, OR
  // - the vector contains one element, and this element is a direction
  //   different from the direction to the primary neighbor.
  // The very rare case where more than one neighbor is excluded from the fit is
  // not cached. In this case, A^{-1} must be computed from scratch.
  const Matrix& retrieve_inverse_a_matrix(
      const Direction<VolumeDim>& primary_direction,
      const std::vector<Direction<VolumeDim>>& directions_to_exclude) const
      noexcept;

  DataVector quadrature_weights;
  DirectionMap<VolumeDim, Matrix> interpolation_matrices;
  DirectionMap<VolumeDim, DataVector>
      quadrature_weights_dot_interpolation_matrices;
  // The many possible values of A^{-1} are stored in a map of maps. The outer
  // map indexes over the primary neighbor, and the inner map indexes over the
  // excluded neighbor.
  // This data structure is not perfect: configurations with only one neighbor
  // (always the primary neighbor) lead to no excluded neighbors, and there is
  // no natural place to hold A^{-1} in the maps. For this case, we simply store
  // the data in the normally-nonsensical slot where
  // excluded_neighbor == primary_neighbor.
  DirectionMap<VolumeDim, DirectionMap<VolumeDim, Matrix>> inverse_a_matrices;
};

// Return the appropriate cache for the given element and mesh.
template <size_t VolumeDim>
const ConstrainedFitCache<VolumeDim>& constrained_fit_cache(
    const Element<VolumeDim>& element, const Mesh<VolumeDim>& mesh) noexcept;

// Compute the inverse of the matrix A_st for the constrained fit.
// See the documentation of `hweno_modified_neighbor_solution`.
template <size_t VolumeDim>
Matrix inverse_a_matrix(
    const Mesh<VolumeDim>& mesh, const Element<VolumeDim>& element,
    const DataVector& quadrature_weights,
    const DirectionMap<VolumeDim, Matrix>& interpolation_matrices,
    const DirectionMap<VolumeDim, DataVector>&
        quadrature_weights_dot_interpolation_matrices,
    const Direction<VolumeDim>& primary_direction,
    const std::vector<Direction<VolumeDim>>& directions_to_exclude) noexcept;

// Not all secondary neighbors (i.e., neighbors that are not the primary) are
// included in the HWENO constrained fit. In particular, for the tensor
// component specified by `Tag` and `tensor_index`, the secondary neighbor whose
// mean is most different from the troubled cell's mean is excluded.
//
// Zhu2016 indicate that when multiple secondary neighbors share the property of
// being "most different" from the troubled cell, then all of these secondary
// neighbors should be excluded from the minimization. The probability of this
// occuring is exceedingly low, but we handle it anyway. We return the excluded
// secondary neighbors in a vector.
//
// Note that if there is only one neighbor, it is the primary neighbor, and so
// there are no secondary neighbors to exclude. We return an empty vector. This
// scenario can arise in various test cases, but is unlikely to arise in science
// cases.
template <typename Tag, size_t VolumeDim, typename Package>
std::vector<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>
secondary_neighbors_to_exclude_from_fit(
    const double local_mean, const size_t tensor_index,
    const std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, Package,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
        neighbor_data,
    const std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>&
        primary_neighbor) noexcept {
  // Rare case: with only one neighbor, there is no secondary to exclude
  if (UNLIKELY(neighbor_data.size() == 1)) {
    return std::vector<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>{};
  }

  // Identify element with maximum mean difference
  const auto mean_difference =
      [&tensor_index, &local_mean](
          const std::pair<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>,
                          Package>& neighbor_and_data) noexcept {
        return fabs(get<::Tags::Mean<Tag>>(
                        neighbor_and_data.second.means)[tensor_index] -
                    local_mean);
      };

  double max_difference = std::numeric_limits<double>::lowest();
  std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>
      neighbor_max_difference{};
  for (const auto& neighbor_and_data : neighbor_data) {
    const auto& neighbor = neighbor_and_data.first;
    if (neighbor == primary_neighbor) {
      continue;
    }
    const double difference = mean_difference(neighbor_and_data);
    if (difference > max_difference) {
      max_difference = difference;
      neighbor_max_difference = neighbor;
    }
  }

  std::vector<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>
      neighbors_to_exclude{{neighbor_max_difference}};

  // See if other elements share this maximum mean difference. This loop should
  // only rarely find other neighbors with the same maximal mean difference to
  // add to the vector, so it will usually not change the vector.
  for (const auto& neighbor_and_data : neighbor_data) {
    const auto& neighbor = neighbor_and_data.first;
    if (neighbor == primary_neighbor or neighbor == neighbor_max_difference) {
      continue;
    }
    const double difference = mean_difference(neighbor_and_data);
    if (UNLIKELY(equal_within_roundoff(difference, max_difference))) {
      neighbors_to_exclude.push_back(neighbor);
    }
  }

  ASSERT(not alg::found(neighbors_to_exclude, primary_neighbor),
         "Logical inconsistency: trying to exclude the primary neighbor.");

  return neighbors_to_exclude;
}

// Compute the vector b_s for the constrained fit. For details, see the
// documentation of `hweno_modified_neighbor_solution` below.
template <typename Tag, size_t VolumeDim, typename Package>
DataVector b_vector(
    const size_t tensor_index, const Mesh<VolumeDim>& mesh,
    const DataVector& quadrature_weights,
    const DirectionMap<VolumeDim, Matrix>& interpolation_matrices,
    const DirectionMap<VolumeDim, DataVector>&
        quadrature_weights_dot_interpolation_matrices,
    const std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, Package,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
        neighbor_data,
    const std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>&
        primary_neighbor,
    const std::vector<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>&
        neighbors_to_exclude) noexcept {
  const size_t number_of_grid_points = mesh.number_of_grid_points();
  DataVector b(number_of_grid_points, 0.);

  for (const auto& neighbor_and_data : neighbor_data) {
    // Generally, neighbors_to_exclude contains just one neighbor to exclude,
    // and the loop over neighbor_data will hit this (in 3D) roughly 1/6 times.
    // So the condition is somewhat, but not overwhelmingly, unlikely:
    if (UNLIKELY(alg::found(neighbors_to_exclude, neighbor_and_data.first))) {
      continue;
    }

    const auto& direction = neighbor_and_data.first.first;
    ASSERT(interpolation_matrices.contains(direction),
           "interpolation_matrices does not contain key: " << direction);
    ASSERT(
        quadrature_weights_dot_interpolation_matrices.contains(direction),
        "quadrature_weights_dot_interpolation_matrices does not contain key: "
            << direction);

    const auto& neighbor_mesh = mesh;
    const auto& neighbor_quadrature_weights = quadrature_weights;
    const auto& interpolation_matrix = interpolation_matrices.at(direction);
    const auto& quadrature_weights_dot_interpolation_matrix =
        quadrature_weights_dot_interpolation_matrices.at(direction);

    const auto& neighbor_tensor_component =
        get<Tag>(neighbor_and_data.second.volume_data)[tensor_index];

    // Add terms from the primary neighbor
    if (neighbor_and_data.first == primary_neighbor) {
      for (size_t r = 0; r < neighbor_mesh.number_of_grid_points(); ++r) {
        for (size_t s = 0; s < number_of_grid_points; ++s) {
          b[s] += neighbor_tensor_component[r] *
                  neighbor_quadrature_weights[r] * interpolation_matrix(r, s);
        }
      }
    }
    // Add terms from the secondary neighbors
    else {
      const double quadrature_weights_dot_u = [&]() noexcept {
        double result = 0.;
        for (size_t r = 0; r < neighbor_mesh.number_of_grid_points(); ++r) {
          result +=
              neighbor_tensor_component[r] * neighbor_quadrature_weights[r];
        }
        return result;
      }();
      b += quadrature_weights_dot_u *
           quadrature_weights_dot_interpolation_matrix;
    }
  }
  return b;
}

// Solve the constrained fit problem that gives the HWENO modified solution,
// for one particular tensor component. For details, see documentation of
// `hweno_modified_neighbor_solution` below.
template <typename Tag, size_t VolumeDim, typename Package>
void solve_constrained_fit(
    const gsl::not_null<DataVector*> constrained_fit_result,
    const DataVector& u, const size_t tensor_index, const Mesh<VolumeDim>& mesh,
    const Element<VolumeDim>& element,
    const std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, Package,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
        neighbor_data,
    const std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>&
        primary_neighbor,
    const std::vector<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>&
        neighbors_to_exclude) noexcept {
  ASSERT(not alg::found(neighbors_to_exclude, primary_neighbor),
         "Logical inconsistency: trying to exclude the primary neighbor.");
  ASSERT(not neighbors_to_exclude.empty() or neighbor_data.size() == 1,
         "The HWENO constrained fit algorithm expects at least one neighbor \n"
         "to exclude from the fit (unless if the element has a single \n"
         "neighbor, which would automatically be the primary neighbor).");

  // Get the cache of linear algebra quantities for this element
  const ConstrainedFitCache<VolumeDim>& cache =
      constrained_fit_cache(element, mesh);

  // Because we don't support h-refinement, the direction is the only piece
  // of the neighbor information that we actually need.
  const Direction<VolumeDim> primary_direction = primary_neighbor.first;
  const std::vector<Direction<VolumeDim>> directions_to_exclude =
      [&neighbors_to_exclude]() noexcept {
        std::vector<Direction<VolumeDim>> result(neighbors_to_exclude.size());
        for (size_t i = 0; i < result.size(); ++i) {
          result[i] = neighbors_to_exclude[i].first;
        }
        return result;
      }();

  const DataVector& w = cache.quadrature_weights;
  const DirectionMap<VolumeDim, Matrix>& interp_matrices =
      cache.interpolation_matrices;
  const DirectionMap<VolumeDim, DataVector>& w_dot_interp_matrices =
      cache.quadrature_weights_dot_interpolation_matrices;

  // Use cache if possible, or compute matrix if we are in edge case
  const Matrix& inverse_a =
      LIKELY(directions_to_exclude.size() < 2)
          ? cache.retrieve_inverse_a_matrix(primary_direction,
                                            directions_to_exclude)
          : inverse_a_matrix(mesh, element, w, interp_matrices,
                             w_dot_interp_matrices, primary_direction,
                             directions_to_exclude);

  const DataVector b = b_vector<Tag>(tensor_index, mesh, w, interp_matrices,
                                     w_dot_interp_matrices, neighbor_data,
                                     primary_neighbor, neighbors_to_exclude);

  const size_t number_of_points = b.size();
  const DataVector inverse_a_times_b = apply_matrices(
      std::array<std::reference_wrapper<const Matrix>, 1>{{inverse_a}}, b,
      Index<1>(number_of_points));
  const DataVector inverse_a_times_w = apply_matrices(
      std::array<std::reference_wrapper<const Matrix>, 1>{{inverse_a}}, w,
      Index<1>(number_of_points));

  // Compute Lagrange multiplier:
  // Note: we take w as an argument (instead of as a lambda capture), because
  //       some versions of Clang incorrectly warn about capturing w.
  const double lagrange_multiplier = [&number_of_points, &inverse_a_times_b,
                                      &inverse_a_times_w,
                                      &u](const DataVector& local_w) noexcept {
    double numerator = 0.;
    double denominator = 0.;
    for (size_t s = 0; s < number_of_points; ++s) {
      numerator += local_w[s] * (inverse_a_times_b[s] - u[s]);
      denominator += local_w[s] * inverse_a_times_w[s];
    }
    return -numerator / denominator;
  }(w);

  // Compute solution:
  *constrained_fit_result =
      inverse_a_times_b + lagrange_multiplier * inverse_a_times_w;
}

/*!
 * \ingroup LimitersGroup
 * \brief Compute the HWENO solution for one tensor
 *
 * The HWENO limiter reconstructs a new solution from the linear combination of
 * the local DG solution and a "modified" solution from each neighbor element.
 * This function computes the modified solution for a particular tensor and
 * neighbor, following Section 3 of \cite Zhu2016.
 *
 * The modified solution associated with a particular neighbor (the "primary"
 * neighbor) is obtained by solving a constrained fit over the local element,
 * the primary neighbor, and the other ("secondary") neighbors of the local
 * element. This fit seeks to minimize in a least-squared sense:
 * 1. The distance between the modified solution and the original solution on
 *    the primary neighbor.
 * 2. The distance between the cell average of the modified solution and the
 *    cell average of the original solution on each secondary neighbor. Note
 *    however that one secondary neighbor (or, rarely, several) is excluded from
 *    this minimization: the secondary neighbor(s) where the original solution
 *    has the most different cell average from the local element. This helps to
 *    prevent an outlier (e.g., near a shock) from biasing the fit.
 *
 * The constraint on the minimization is the following: the cell average of the
 * modified solution on the local element must equal the cell average of the
 * local element's original solution.
 *
 * Below we give the mathematical form of the constraints described above and
 * show how these are translated into a numerical algorithm.
 *
 * Consider an element \f$I_0\f$ with neighbors \f$I_1, I_2, ...\f$. For a
 * given tensor component \f$u\f$, the values on each of these elements are
 * \f$u^{(0)}, u^{(1)}, u^{(2)}, ...\f$. Taking for the sake of example the
 * primary neighbor to be \f$I_1\f$, the modified solution \f$\phi\f$ must
 * minimize
 *
 * \f[
 * \chi^2 = \int_{I_1} (\phi - u^{(1)})^2 dV
 *          + \sum_{\ell \in L} \left(
 *                              \int_{I_{\ell}} ( \phi - u^{(\ell)} ) dV
 *                              \right)^2,
 * \f]
 *
 * subject to the constaint
 *
 * \f[
 * C = \int_{I_0} ( \phi - u^{(0)} ) dV = 0.
 * \f]
 *
 * where \f$\ell\f$ ranges over a subset \f$L\f$ of the secondary neighbors.
 * \f$L\f$ excludes the one (or more) secondary neighbor(s) where the mean of
 * \f$u\f$ is the most different from the mean of \f$u^{(0)}\f$. Typically, only
 * one secondary neighbor is excluded, so \f$L\f$ contains two fewer neighbors
 * than the total number of neighbors to the element. Note that in 1D, this
 * implies that \f$L\f$ is the empty set; for each modified solution, one
 * neighbor is the primary neighbor and the other is the excluded neighbor.
 *
 * The integrals are evaluated by quadrature. We denote the quadrature weights
 * by \f$w_s\f$ and the values of some data \f$X\f$ at the quadrature
 * nodes by \f$X_s\f$. We use subscripts \f$r,s\f$ to denote quadrature nodes
 * on the neighbor and local elements, respectively. The minimization becomes
 *
 * \f[
 * \chi^2 = \sum_r w^{(1)}_r ( \phi_r - u^{(1)}_r )^2
 *          + \sum_{\ell \in L} \left(
 *                              \sum_r w^{(\ell)}_r ( \phi_r - u^{(\ell)}_r )
 *                              \right)^2,
 * \f]
 *
 * subject to the constraint
 *
 * \f[
 * C = \sum_s w^{(0)}_s ( \phi_s - u^{(0)}_s ) = 0.
 * \f]
 *
 * Note that \f$\phi\f$ is a function defined on the local element \f$I_0\f$,
 * and so is fully represented by its values \f$\phi_s\f$ at the quadrature
 * points on this element. When evaluating \f$\phi\f$ on element \f$I_{\ell}\f$,
 * we obtain the function values \f$\phi_r\f$ by polynomial extrapolation,
 * \f$\phi_r = \sum_s \mathcal{I}^{(\ell)}_{rs} \phi_s\f$, where
 * \f$\mathcal{I}^{(\ell)}_{rs}\f$ is the interpolation/extrapolation matrix
 * that interpolates data defined at grid points \f$x_s\f$ and evaluates it at
 * grid points \f$x_r\f$ of \f$I_{\ell}\f$. Thus,
 *
 * \f[
 * \chi^2 = \sum_r w^{(1)}_r
 *                 \left(
 *                 \sum_s \mathcal{I}^{(1)}_{rs} \phi_s - u^{(1)}_r
 *                 \right)^2
 *          + \sum_{\ell \in L} \left(
 *                              \sum_r w^{(\ell)}_r
 *                                     \left(
 *                                     \sum_s \mathcal{I}^{(\ell)}_{rs} \phi_s
 *                                             - u^{(\ell)}_r
 *                                     \right)
 *                      \right)^2.
 * \f]
 *
 * The solution to this optimization problem is found in the standard way,
 * using a Lagrange multiplier \f$\lambda\f$ to impose the constraint:
 *
 * \f[
 * 0 = \frac{d}{d \phi_s} \left( \chi^2 + \lambda C \right).
 * \f]
 *
 * Working out the differentiation with respect to \f$\phi_s\f$ leads to the
 * linear problem that must be inverted to obtain the solution,
 *
 * \f[
 * 0 = A_{st} \phi_t - b_s - \lambda w^{(0)}_s,
 * \f]
 *
 * where
 *
 * \f{align*}{
 * A_{st} &= \sum_r \left( w^{(1)}_r
 *                  \mathcal{I}^{(1)}_{rs} \mathcal{I}^{(1)}_{rt}
 *                  \right)
 *          + \sum_{\ell \in L} \left(
 *                              \sum_r \left(
 *                                     w^{(\ell)}_r \mathcal{I}^{(\ell)}_{rt}
 *                                     \right)
 *                              \cdot
 *                              \sum_r \left(
 *                                     w^{(\ell)}_r \mathcal{I}^{(\ell)}_{rs}
 *                                     \right)
 *                              \right)
 * \\
 * b_s &= \sum_r \left( w^{(1)}_r u^{(1)}_r \mathcal{I}^{(1)}_{rs} \right)
 *         + \sum_{\ell \in L} \left(
 *                             \sum_r \left( w^{(\ell)}_r u^{(\ell)}_r \right)
 *                             \cdot
 *                             \sum_r \left(
 *                                    w^{(\ell)}_r \mathcal{I}^{(\ell)}_{rs}
 *                                    \right)
 *                             \right).
 * \f}
 *
 * Finally, the solution to the constrained fit is
 *
 * \f{align*}{
 * \lambda &= - \frac{ \sum_s w^{(0)}_s
 *                            \left( (A^{-1})_{st} b_t - u^{(0)}_s \right)
 *              }{ \sum_s w^{(0)}_s (A^{-1})_{st} w^{(0)}_t }
 * \\
 * \phi_s &= (A^{-1})_{st} ( b_t + \lambda w^{(0)}_t ).
 * \f}
 *
 * Note that the matrix \f$A\f$ does not depend on the values of the tensor
 * \f$u\f$, so its inverse \f$A^{-1}\f$ can be precomputed and stored.
 *
 * \warning
 * Note also that the implementation currently does not support h- or
 * p-refinement; this is checked by some assertions. The implementation is
 * untested for grids where elements are curved, and it should not be expected
 * to work in these cases.
 *
 * When calling `hweno_impl`, `modified_neighbor_solution_buffer` should contain
 * one DataVector for each neighboring element (i.e. for each entry in
 * `neighbor_data`).
 */
template <typename Tag, size_t VolumeDim, typename PackagedData>
void hweno_impl(
    const gsl::not_null<std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, DataVector,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>*>
        modified_neighbor_solution_buffer,
    const gsl::not_null<typename Tag::type*> tensor,
    const double neighbor_linear_weight, const Mesh<VolumeDim>& mesh,
    const Element<VolumeDim>& element,
    const std::unordered_map<
        std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>, PackagedData,
        boost::hash<std::pair<Direction<VolumeDim>, ElementId<VolumeDim>>>>&
        neighbor_data) noexcept {
  ASSERT(
      modified_neighbor_solution_buffer->size() == neighbor_data.size(),
      "modified_neighbor_solution_buffer->size() = "
          << modified_neighbor_solution_buffer->size()
          << "\nneighbor_data.size() = " << neighbor_data.size()
          << "\nmodified_neighbor_solution_buffer was incorrectly initialized "
             "before calling hweno_impl.");

  alg::for_each(
      neighbor_data, [&element, &mesh](const auto& neighbor_and_data) noexcept {
        ASSERT(Weno_detail::check_element_has_one_similar_neighbor_in_direction(
                   element, neighbor_and_data.first.first),
               "Found some amount of h-refinement; this is not supported");
        ASSERT(neighbor_and_data.second.mesh == mesh,
               "Found some amount of p-refinement; this is not supported");
      });

  for (size_t tensor_index = 0; tensor_index < tensor->size(); ++tensor_index) {
    const auto& tensor_component = (*tensor)[tensor_index];
    for (const auto& neighbor_and_data : neighbor_data) {
      const auto& primary_neighbor = neighbor_and_data.first;
      const auto neighbors_to_exclude =
          secondary_neighbors_to_exclude_from_fit<Tag>(
              mean_value(tensor_component, mesh), tensor_index, neighbor_data,
              primary_neighbor);

      DataVector& buffer =
          modified_neighbor_solution_buffer->at(primary_neighbor);
      solve_constrained_fit<Tag>(make_not_null(&buffer), tensor_component,
                                 tensor_index, mesh, element, neighbor_data,
                                 primary_neighbor, neighbors_to_exclude);
    }

    // Sum local and modified neighbor polynomials for the WENO reconstruction
    Weno_detail::reconstruct_from_weighted_sum(
        make_not_null(&((*tensor)[tensor_index])), neighbor_linear_weight,
        Weno_detail::DerivativeWeight::PowTwoEllOverEllFactorial, mesh,
        *modified_neighbor_solution_buffer);
  }
}

}  // namespace Limiters::Weno_detail


#include "optimizers.hpp"
#include "modeling.hpp"
#include "utils/logging.hpp"
#include <boost/foreach.hpp>
#include "solver_interface.hpp"
#include "expr_ops.hpp"
#include <cmath>
#include <cstdio>
#include "sco_common.hpp"
#include "utils/stl_to_string.hpp"
#include "macros.h"
#include <boost/format.hpp>
using namespace std;
using namespace util;

namespace sco {

typedef vector<double> DblVec;

std::ostream& operator<<(std::ostream& o, const OptResults& r) {
  o << "Optimization results:" << endl
    << "status: " << statusToString(r.status) << endl
    << "cost values: " << Str(r.cost_vals) << endl
    << "constraint violations: " << Str(r.cnt_viols) << endl
    << "n func evals: " << r.n_func_evals << endl
    << "n qp solves: " << r.n_qp_solves << endl;
  return o;
}

//////////////////////////////////////////////////
////////// private utility functions for  sqp /////////
//////////////////////////////////////////////////

static DblVec evaluateCosts(vector<CostPtr>& costs, const DblVec& x) {
  DblVec out(costs.size());
  for (size_t i=0; i < costs.size(); ++i) {
    out[i] = costs[i]->value(x);
  }
  return out;
}
static DblVec evaluateConstraintViols(vector<ConstraintPtr>& constraints, const DblVec& x) {
  DblVec out(constraints.size());
  for (size_t i=0; i < constraints.size(); ++i) {
    out[i] = constraints[i]->violation(x);
  }
  return out;
}
static vector<ConvexObjectivePtr> convexifyCosts(vector<CostPtr>& costs, const DblVec& x, Model* model) {
  vector<ConvexObjectivePtr> out(costs.size());
  for (size_t i=0; i < costs.size(); ++i) {
    out[i] = costs[i]->convex(x,  model);
  }
  return out;
}
static vector<ConvexConstraintsPtr> convexifyConstraints(vector<ConstraintPtr>& cnts, const DblVec& x, Model* model) {
  vector<ConvexConstraintsPtr> out(cnts.size());
  for (size_t i=0; i < cnts.size(); ++i) {
    out[i] = cnts[i]->convex(x, model);
  }
  return out;
}

DblVec evaluateModelCosts(vector<ConvexObjectivePtr>& costs, const DblVec& x) {
  DblVec out(costs.size());
  for (size_t i=0; i < costs.size(); ++i) {
    out[i] = costs[i]->value(x);
  }
  return out;
}
DblVec evaluateModelCntViols(vector<ConvexConstraintsPtr>& cnts, const DblVec& x) {
  DblVec out(cnts.size());
  for (size_t i=0; i < cnts.size(); ++i) {
    out[i] = cnts[i]->violation(x);
  }
  return out;
}

static vector<string> getCostNames(const vector<CostPtr>& costs) {
  vector<string> out(costs.size());
  for (size_t i=0; i < costs.size(); ++i) out[i] = costs[i]->name();
  return out;
}
static vector<string> getCntNames(const vector<ConstraintPtr>& cnts) {
  vector<string> out(cnts.size());
  for (size_t i=0; i < cnts.size(); ++i) out[i] = cnts[i]->name();
  return out;
}

void printCostInfo(const vector<double>& old_cost_vals, const vector<double>& model_cost_vals, const vector<double>& new_cost_vals,
                  const vector<double>& old_cnt_vals, const vector<double>& model_cnt_vals, const vector<double>& new_cnt_vals,
    const vector<string>& cost_names, const vector<string>& cnt_names, double merit_coeff) {
    printf("%15s | %10s | %10s | %10s | %10s\n", "", "oldexact", "dapprox", "dexact", "ratio");
    printf("%15s | %10s---%10s---%10s---%10s\n", "COSTS", "----------", "----------", "----------", "----------");
    for (size_t i=0; i < old_cost_vals.size(); ++i) {
      double approx_improve = old_cost_vals[i] - model_cost_vals[i];
      double exact_improve = old_cost_vals[i] - new_cost_vals[i];
      if (fabs(approx_improve) > 1e-8)
        printf("%15s | %10.3e | %10.3e | %10.3e | %10.3e\n", cost_names[i].c_str(), old_cost_vals[i], approx_improve, exact_improve, exact_improve/approx_improve);
      else
        printf("%15s | %10.3e | %10.3e | %10.3e | %10s\n",cost_names[i].c_str(), old_cost_vals[i], approx_improve, exact_improve, "  ------  ");
    }
    if (cnt_names.size() == 0) return;
    printf("%15s | %10s---%10s---%10s---%10s\n", "CONSTRAINTS", "----------", "----------", "----------", "----------");
    for (size_t i=0; i < old_cnt_vals.size(); ++i) {
      double approx_improve = old_cnt_vals[i] - model_cnt_vals[i];
      double exact_improve = old_cnt_vals[i] - new_cnt_vals[i];
      if (fabs(approx_improve > 1e-8))
        printf("%15s | %10.3e | %10.3e | %10.3e | %10.3e\n", cnt_names[i].c_str(), merit_coeff*old_cnt_vals[i], merit_coeff*approx_improve, merit_coeff*exact_improve, exact_improve/approx_improve);
      else
        printf("%15s | %10.3e | %10.3e | %10.3e | %10s\n", cnt_names[i].c_str(), merit_coeff*old_cnt_vals[i], merit_coeff*approx_improve, merit_coeff*exact_improve, "  ------  ");
    }

}

// todo: use different coeffs for each constraint
vector<ConvexObjectivePtr> cntsToCosts(const vector<ConvexConstraintsPtr>& cnts, double err_coeff, Model* model) {
  vector<ConvexObjectivePtr> out;
  BOOST_FOREACH(const ConvexConstraintsPtr& cnt, cnts) {
    ConvexObjectivePtr obj(new ConvexObjective(model));
    BOOST_FOREACH(const AffExpr& aff, cnt->eqs_) {
      obj->addAbs(aff, err_coeff);
    }
    BOOST_FOREACH(const AffExpr& aff, cnt->ineqs_) {
      obj->addHinge(aff, err_coeff);
    }
    out.push_back(obj);
  }
  return out;
}

void Optimizer::addCallback(const Callback& cb) {
  callbacks_.push_back(cb);
}
void Optimizer::callCallbacks(DblVec& x) {
  for (int i=0; i < callbacks_.size(); ++i) {
    callbacks_[i](prob_.get(), x);
  }
}

void Optimizer::initialize(const vector<double>& x) {
  if (!prob_) PRINT_AND_THROW("need to set the problem before initializing");
  if (prob_->getVars().size() != x.size())
    PRINT_AND_THROW(boost::format("initialization vector has wrong length. expected %i got %i")%prob_->getVars().size()%x.size());
  results_.clear();
  results_.x = x;
}

BasicTrustRegionSQP::BasicTrustRegionSQP() {
  initParameters();
}
BasicTrustRegionSQP::BasicTrustRegionSQP(OptProbPtr prob) {
  initParameters();
  setProblem(prob);
}

void BasicTrustRegionSQP::initParameters() {

  improve_ratio_threshold_ = .25;
  min_trust_box_size_ = 1e-4;
  min_approx_improve_= 1e-4;
  min_approx_improve_frac_ = -INFINITY;
  max_iter_ = 50;
  trust_shrink_ratio_=.1;
  trust_expand_ratio_ = 1.5;
  cnt_tolerance_ = 1e-4;
  max_merit_coeff_increases_ = 5;
  merit_coeff_increase_ratio_ = 10;
  max_time_ = INFINITY;

  merit_error_coeff_ = 10;
  trust_box_size_ = 1e-1;

}

void BasicTrustRegionSQP::setProblem(OptProbPtr prob) {
  Optimizer::setProblem(prob);
  model_ = prob->getModel();
}

void BasicTrustRegionSQP::adjustTrustRegion(double ratio) {
  trust_box_size_ *= ratio;
}
void BasicTrustRegionSQP::setTrustBoxConstraints(const DblVec& x) {
  vector<Var>& vars = prob_->getVars();
  assert(vars.size() == x.size());
  DblVec& lb=prob_->getLowerBounds(), ub=prob_->getUpperBounds();
  DblVec lbtrust(x.size()), ubtrust(x.size());
  vector<bool> incmask = prob_->getIncrementMask();
  if (incmask.empty()) incmask = vector<bool>(x.size(), false);
  for (size_t i=0; i < x.size(); ++i) {
    lbtrust[i] = fmax((incmask[i] ? 0 : x[i]) - trust_box_size_, lb[i]);
    ubtrust[i] = fmin((incmask[i] ? 0 : x[i]) + trust_box_size_, ub[i]);
  }
  model_->setVarBounds(vars, lbtrust, ubtrust);
}

#if 0
struct MultiCritFilter {
  /**
   * Checks if you're making an improvement on a multidimensional objective
   * Given a set of past error vectors, the improvement is defined as
   * min_{olderrvec in past_err_vecs} | olderrvec - errvec |^+
   */
  vector<DblVec> errvecs;
  double improvement(const DblVec& errvec) {
    double leastImprovement=INFINITY;
    BOOST_FOREACH(const DblVec& olderrvec, errvecs) {
      double improvement=0;
      for (int i=0; i < errvec.size(); ++i) improvement += pospart(olderrvec[i] - errvec[i]);
      leastImprovement = fmin(leastImprovement, improvement);
    }
    return leastImprovement;
  }
  void insert(const DblVec& x) {errvecs.push_back(x);}
  bool empty() {return errvecs.size() > 0;}
};
#endif

OptStatus BasicTrustRegionSQP::optimize() {

  vector<string> cost_names = getCostNames(prob_->getCosts());
  vector<ConstraintPtr> constraints = prob_->getConstraints();
  vector<string> cnt_names = getCntNames(constraints);

  DblVec& x_ = results_.x; // just so I don't have to rewrite code
  if (x_.size() == 0) PRINT_AND_THROW("you forgot to initialize!");
  if (!prob_) PRINT_AND_THROW("you forgot to set the optimization problem");

  x_ = prob_->getClosestFeasiblePoint(x_);

  assert(x_.size() == prob_->getVars().size());
  assert(prob_->getCosts().size() > 0 || constraints.size() > 0);

  OptStatus retval = INVALID;

  for (int merit_increases=0; merit_increases < max_merit_coeff_increases_; ++merit_increases) { /* merit adjustment loop */
    for (int iter=1; ; ++iter) { /* sqp loop */
      callCallbacks(x_);

      LOG_DEBUG("current iterate: %s", CSTR(x_));
      LOG_INFO("iteration %i", iter);

      // speedup: if you just evaluated the cost when doing the line search, use that
      if (results_.cost_vals.empty()) { //only happens on the first iteration
        results_.cnt_viols = evaluateConstraintViols(constraints, x_);
        results_.cost_vals = evaluateCosts(prob_->getCosts(), x_);
        assert(results_.n_func_evals == 0);
        ++results_.n_func_evals;
      }

      vector<ConvexObjectivePtr> cost_models = convexifyCosts(prob_->getCosts(),x_, model_.get());
      vector<ConvexConstraintsPtr> cnt_models = convexifyConstraints(constraints, x_, model_.get());
      vector<ConvexObjectivePtr> cnt_cost_models = cntsToCosts(cnt_models, merit_error_coeff_, model_.get());
      model_->update();
      BOOST_FOREACH(ConvexObjectivePtr& cost, cost_models)cost->addConstraintsToModel();
      BOOST_FOREACH(ConvexObjectivePtr& cost, cnt_cost_models)cost->addConstraintsToModel();
      model_->update();
      QuadExpr objective;
      BOOST_FOREACH(ConvexObjectivePtr& co, cost_models)exprInc(objective, co->quad_);
      BOOST_FOREACH(ConvexObjectivePtr& co, cnt_cost_models){
      exprInc(objective, co->quad_);
    }
//    objective = cleanupExpr(objective);
      model_->setObjective(objective);

//    if (logging::filter() >= IPI_LEVEL_DEBUG) {
//      DblVec model_cost_vals;
//      BOOST_FOREACH(ConvexObjectivePtr& cost, cost_models) {
//        model_cost_vals.push_back(cost->value(x));
//      }
//      LOG_DEBUG("model costs %s should equalcosts  %s", printer(model_cost_vals), printer(cost_vals));
//    }

      while (trust_box_size_ >= min_trust_box_size_) {

        setTrustBoxConstraints(x_);
        CvxOptStatus status = model_->optimize();
        ++results_.n_qp_solves;
        if (status != CVX_SOLVED) {
          LOG_ERROR("convex solver failed! set LOG_DEBUG_LEVEL=DEBUG to see solver output. saving model to /tmp/fail.lp");
          model_->writeToFile("/tmp/fail.lp");
          retval = OPT_FAILED;
          goto cleanup;
        }
        DblVec model_var_vals = model_->getVarValues(model_->getVars());

        DblVec model_cost_vals = evaluateModelCosts(cost_models, model_var_vals);
        DblVec model_cnt_viols = evaluateModelCntViols(cnt_models, model_var_vals);

        // the n variables of the OptProb happen to be the first n variables in the Model
        DblVec new_x(model_var_vals.begin(), model_var_vals.begin() + x_.size());

        if (GetLogLevel() >= util::LevelDebug) {
          DblVec model_cnt_viols2 = evaluateModelCosts(cnt_cost_models, model_var_vals);
          LOG_DEBUG("SHOULD BE THE SAME: %.2f*%s ?= %s", merit_error_coeff_, CSTR(model_cnt_viols), CSTR(model_cnt_viols2));
        }

        DblVec new_cost_vals = evaluateCosts(prob_->getCosts(), new_x);
        DblVec new_cnt_viols = evaluateConstraintViols(constraints, new_x);
        ++results_.n_func_evals;

        double old_merit = vecSum(results_.cost_vals) + merit_error_coeff_ * vecSum(results_.cnt_viols);
        double model_merit = vecSum(model_cost_vals) + merit_error_coeff_ * vecSum(model_cnt_viols);
        double new_merit = vecSum(new_cost_vals) + merit_error_coeff_ * vecSum(new_cnt_viols);
        double approx_merit_improve = old_merit - model_merit;
        double exact_merit_improve = old_merit - new_merit;
        double merit_improve_ratio = exact_merit_improve / approx_merit_improve;

        if (util::GetLogLevel() >= util::LevelInfo) {
          LOG_INFO("");
          printCostInfo(results_.cost_vals, model_cost_vals, new_cost_vals,
                        results_.cnt_viols, model_cnt_viols, new_cnt_viols, cost_names,
                        cnt_names, merit_error_coeff_);
          printf("%15s | %10.3e | %10.3e | %10.3e | %10.3e\n", "TOTAL", old_merit, approx_merit_improve, exact_merit_improve, merit_improve_ratio);
        }

        //cout << "approx merit improve: " << approx_merit_improve << endl;

        if (approx_merit_improve < -1e-4) {
          LOG_ERROR("approximate merit function got worse (%.3e). (convexification is probably wrong to zeroth order)", approx_merit_improve);
        }
        if (approx_merit_improve < min_approx_improve_) {
          LOG_INFO("converged because improvement was small (%.3e < %.3e)", approx_merit_improve, min_approx_improve_);
          retval = OPT_CONVERGED;
          goto penaltyadjustment;
        }
        if (approx_merit_improve / old_merit < min_approx_improve_frac_) {
          LOG_INFO(
              "converged because improvement ratio was small (%.3e < %.3e)",
              approx_merit_improve/old_merit, min_approx_improve_frac_);
          retval = OPT_CONVERGED;
          goto penaltyadjustment;
        }
        else if (exact_merit_improve < 0 || merit_improve_ratio < improve_ratio_threshold_) {
          adjustTrustRegion(trust_shrink_ratio_);
          LOG_INFO("shrunk trust region. new box size: %.4f",
              trust_box_size_);
        } else {
          x_ = new_x;
          results_.cost_vals = new_cost_vals;
          results_.cnt_viols = new_cnt_viols;
          adjustTrustRegion(trust_expand_ratio_);
          LOG_INFO("expanded trust region. new box size: %.4f",trust_box_size_);
          break;
        }
      }

      if (trust_box_size_ < min_trust_box_size_) {
        LOG_INFO("converged because trust region is tiny");
        retval = OPT_CONVERGED;
        goto penaltyadjustment;
      } else if (iter >= max_iter_) {
        LOG_INFO("iteration limit");
        retval = OPT_ITERATION_LIMIT;
        goto cleanup;
      }
    }

    penaltyadjustment:
    if (results_.cnt_viols.empty() || vecMax(results_.cnt_viols) < cnt_tolerance_) {
      if (results_.cnt_viols.size() > 0) LOG_INFO("woo-hoo! constraints are satisfied (to tolerance %.2e)", cnt_tolerance_);
      goto cleanup;
    }
    else {
      merit_error_coeff_ *= merit_coeff_increase_ratio_;
      LOG_INFO("not all constraints are satisfied. increasing penalties to %.4f", merit_error_coeff_);
      trust_box_size_ = fmax(trust_box_size_, min_trust_box_size_ * 5);
    }

  }
  retval = OPT_ITERATION_LIMIT;
  LOG_INFO("optimization couldn't satisfy all constraints");

  cleanup:
  assert(retval != INVALID && "should never happen");
  results_.status = retval;
  results_.total_cost = vecSum(results_.cost_vals);
  LOG_INFO("\n==================\n%s==================", CSTR(results_));
  callCallbacks(x_);

  return retval;

}

}


/* CirKit: A circuit toolkit
 * Copyright (C) 2009-2015  University of Bremen
 * Copyright (C) 2015-2017  EPFL
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

/**
 * @file xmg_exact_decomp.hpp
 *
 * @brief Exact MAJ decomposition
 *
 * @author Mathias Soeken
 * @author Eleonora Testa
 * @since  2.4
 */

#ifndef XMG_EXACT_DECOMP_HPP
#define XMG_EXACT_DECOMP_HPP

#include <array>
#include <chrono>
#include <cstdint>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

#include <boost/optional.hpp>

#include <classical/abc/abc_api.hpp>
#include <classical/xmg/xmg.hpp>
#include <fmt/format.h>
#include <sat/glucose/AbcGlucose.h>

namespace cirkit
{

namespace detail
{

template<uint32_t NumVars, uint32_t NumGates>
struct xmg_exact_decomposition_impl
{
public:
  xmg_exact_decomposition_impl()
  {
    for ( auto& a : select )
    {
      for ( auto& a2 : a )
      {
        a2.fill( 0 );
      }
    }

    add_structure_variables();

    solver = abc::bmcg_sat_solver_start();
    abc::bmcg_sat_solver_set_nvars( solver, var_index );
  }

  ~xmg_exact_decomposition_impl()
  {
    abc::bmcg_sat_solver_stop( solver );
  }

  void write_dimacs(std::string filename)
  {
  	file.open(filename, std::ofstream::out);
	file << fmt::format( "P cnf {} {}",
                       abc::bmcg_sat_solver_varnum( solver ),
					   abc::bmcg_sat_solver_clausenum( solver ))
		 << std::endl;
    file << str;
  }

  bool run(std::string filename)
  {

    if ( !add_main_cnf() )
    {
      return false;
    }

    // add value CNF for each interesting minterm
    for ( uint32_t i = 0; i < ( uint32_t( 1 ) << NumVars ); ++i )
    {
      auto cnt = __builtin_popcount( i );

      switch ( cnt )
      {
      default:
        break;
      case NumVars / 2 - 1:
        if ( !add_minterm( i, 0 ) )
        {
          return false;
        }
        break;
      case NumVars / 2 + 1:
        if ( !add_minterm( i, 1 ) )
        {
          return false;
        }
        break;
      case NumVars / 2:
        if ( !add_minterm( i, 2 ) )
        {
          return false;
        }
        break;
      }
    }

    /* write dimacs file */
	write_dimacs(filename);

    /* solving time */
    const auto start = std::chrono::high_resolution_clock::now();
    const auto result = abc::bmcg_sat_solver_solve( solver, nullptr, 0 );
    const auto stop = std::chrono::high_resolution_clock::now();

    solving_time = std::chrono::duration<double>{stop - start}.count();

    if ( result == GLUCOSE_SAT )
    {
      return true;
    }
    return false;
  }

  void print_statistics( std::ostream& os = std::cout )
  {
    os << fmt::format( "[i] vars = {}   cls = {}   learnt = {}   conf = {}   time = {:.2f}",
                       abc::bmcg_sat_solver_varnum( solver ),
                       abc::bmcg_sat_solver_clausenum( solver ),
                       abc::bmcg_sat_solver_learntnum( solver ),
                       abc::bmcg_sat_solver_conflictnum( solver ),
                       solving_time )
       << std::endl;
  }

  void print_solution( std::ostream& os = std::cout )
  {
    for ( auto i = 0; i < NumGates; ++i )
    {
      os << fmt::format( "n{} = <", i + NumVars + 1 );

      for ( auto k = 0; k < 3; ++k )
      {
        for ( auto j = 0; j < NumVars + i; ++j )
        {
          if ( select[i][k][j] && abc::bmcg_sat_solver_read_cex_varvalue( solver, select[i][k][j] ) )
          {
            os << fmt::format( "{}{}", j < NumVars ? 'x' : 'n', j + 1 );
          }
        }
      }

      os << ">" << std::endl;
    }
  }

  xmg_graph extract_circuit()
  {
    xmg_graph g;

    std::vector<xmg_function> functions;

    for ( auto i = 0; i < NumVars; ++i )
    {
      functions.push_back( g.create_pi( fmt::format( "x{}", i + 1 ) ) );
    }

    for ( auto i = 0; i < NumGates; ++i )
    {
      std::vector<xmg_function> children;

      for ( auto k = 0; k < 3; ++k )
      {
        for ( auto j = 0; j < NumVars + i; ++j )
        {
          if ( select[i][k][j] && abc::bmcg_sat_solver_read_cex_varvalue( solver, select[i][k][j] ) )
          {
            children.push_back( functions[j] );
          }
        }
      }

      assert( children.size() == 3u );

      functions.push_back( g.create_maj( children[0], children[1], children[2] ) );
    }

    const auto out = g.create_maj( g.create_pi( fmt::format( "x{}", NumVars ) ), functions[functions.size() - 2], functions[functions.size() - 1] );
    g.create_po( out, "f" );

    return g;
  }

private:
  bool add_clause(abc::bmcg_sat_solver* solver, std::vector<int> lits, int size)
  {
	  for (auto& l : lits)
	  {
		  auto divs = l/2;
		  if (l%2 == 0)
			  //str << divs << " ";
		      str += fmt::format( "{} ", divs);
		  else
			  str += fmt::format( "-{} ", divs);

	  }
	  str += "0\n";

	  if (!abc::bmcg_sat_solver_addclause( solver, &lits[0], size))
		  {
			  return false;
		  }
	  return true;
  }

  void add_structure_variables()
  {
    /* first gate is <x1x2x3> */
    for ( auto k = 0; k < 3; ++k )
    {
      select[0][k][k] = var_index++;
    }

    for ( auto i = 1; i < NumGates; ++i )
    {
      for ( auto k = 0; k < 3; ++k )
      {
        const auto max = ( k == 0 ) ? NumVars : ( ( i == NumGates - 1 ) ? NumVars + i - 1 : NumVars + i );
        for ( auto j = 0; j < max; ++j )
        {
          select[i][k][j] = var_index;
          if ( j >= NumVars && ( j - NumVars ) < NumGates - 2 )
          {
            output_lits[j - NumVars].push_back( abc::Abc_Var2Lit( var_index, 0 ) );
          }
          var_index++;
        }
      }
    }
  }

  bool add_main_cnf()
  {
    /* every node has three children */
    for ( auto i = 0; i < NumGates; ++i )
    {
      for ( auto k = 0; k < 3; ++k )
      {
        std::vector<int> lits;

        for ( auto j = 0; j < NumVars + i; ++j )
        {
          if ( !select[i][k][j] )
            continue;

          lits.push_back( abc::Abc_Var2Lit( select[i][k][j], 0 ) );
        }

        if ( !add_clause( solver, lits, lits.size() ) )
        {
          return false;
        }

        for ( auto l = 1; l < lits.size(); ++l )
        {
          for ( auto m = 0; m < l; ++m )
          {
            std::vector<int> plits = {abc::Abc_LitNot( lits[m] ), abc::Abc_LitNot( lits[l] )};
            if ( !add_clause( solver, plits, 2 ) )
            {
              return false;
            }
          }
        }

        /* symmetry breaking */
        if ( k == 2 )
          break;

        for ( auto l = 1; l < NumVars + i; ++l )
        {
          for ( auto m = 0; m < l; ++m )
          {
            // select[i][k][l] -> !select[i][k + 1][m]
            if ( !select[i][k][l] || !select[i][k + 1][m] )
              continue;

            std::vector<int> plits = {abc::Abc_Var2Lit( select[i][k][l], 1 ), abc::Abc_Var2Lit( select[i][k + 1][m], 1 )};
            if ( !add_clause( solver, plits, 2 ) )
            {
              return false;
            }
          }
        }
      }

      /* no gate appears twice */
      for ( auto ii = 0; ii < i; ++ii )
      {
        for ( auto l = 2; l < NumVars + ii; ++l )
        {
          for ( auto m = 1; m < l; ++m )
          {
            for ( auto n = 0; n < m; ++n )
            {
              if ( !select[i][0][n] || !select[i][1][m] || !select[i][2][l] || !select[ii][0][n] || !select[ii][1][m] || !select[ii][2][l] )
                continue;

              std::vector<int> plits = {abc::Abc_Var2Lit( select[i][0][n], 1 ), abc::Abc_Var2Lit( select[i][1][m], 1 ), abc::Abc_Var2Lit( select[i][2][l], 1 ),
                             abc::Abc_Var2Lit( select[ii][0][n], 1 ), abc::Abc_Var2Lit( select[ii][1][m], 1 ), abc::Abc_Var2Lit( select[ii][2][l], 1 )};
              if ( !add_clause( solver, plits, 6 ) )
              {
                return false;
              }
            }
          }
        }
      }
    }

    /* every gate needs to be used */
    for ( auto& lits : output_lits )
    {
      if ( !add_clause( solver, lits, lits.size() ) )
      {
        return false;
      }
    }

    return true;
  }

  bool add_minterm( uint32_t minterm, int expected_value )
  {
    std::array<bool, NumVars> assignment;

    for ( auto i = 0; i < NumVars; ++i )
    {
      assignment[i] = ( minterm >> i ) & 1;
    }

    abc::bmcg_sat_solver_set_nvars( solver, var_index + 4 * NumGates );

    /* structure propagation */
    for ( auto i = 0; i < NumGates; ++i )
    {
      const auto base_var_i = var_index + 4 * i;

      for ( auto k = 0; k < 3; ++k )
      {

        for ( auto j = 0; j < NumVars + i; ++j )
        {
          if ( !select[i][k][j] )
            continue;

          std::vector<int> lits;
          lits.push_back( abc::Abc_Var2Lit( select[i][k][j], 1 ) );

          if ( j < NumVars ) /* node i'th k'th input should have value of PI j */
          {
            lits.push_back( abc::Abc_Var2Lit( base_var_i + k, assignment[j] ? 0 : 1 ) );
            if ( !add_clause( solver, lits, lits.size() ) )
            {
              return false;
            }
          }
          else /* node i'th k'th input should have value of node j */
          {
            const auto base_var_j = var_index + 4 * ( j - NumVars );
            lits.push_back( abc::Abc_Var2Lit( base_var_j + 3, 0 ) );
            lits.push_back( abc::Abc_Var2Lit( base_var_i + k, 1 ) );
            if ( !add_clause( solver, lits, lits.size() ) )
            {
              return false;
            }
            lits[1] = abc::Abc_LitNot( lits[1] );
            lits[2] = abc::Abc_LitNot( lits[2] );
            if ( !add_clause( solver, lits, lits.size() ) )
            {
              return false;
            }
          }
        }
      }

      for ( auto n = 0; n < 2; ++n )
      {
        if ( i >= NumGates - 2 && expected_value < 2 && n == expected_value )
        {
          continue;
        }

        for ( auto k = 0; k < 3; k++ )
        {
          std::vector<int> pLits(3);
		  int nLits = 0;
          if ( k != 0 )
          {
            pLits[nLits++] = abc::Abc_Var2Lit( base_var_i + 0, n );
          }
          if ( k != 1 )
          {
            pLits[nLits++] = abc::Abc_Var2Lit( base_var_i + 1, n );
          }
          if ( k != 2 )
          {
            pLits[nLits++] = abc::Abc_Var2Lit( base_var_i + 2, n );
          }
          if ( i < NumGates - 2 || expected_value == 2 )
          {
            pLits[nLits++] = abc::Abc_Var2Lit( base_var_i + 3, !n );
          }
          assert( nLits <= 3 );
          if ( !add_clause( solver, pLits, nLits ) )
          {
            return false;
          }
        }
      }
    }

    /* special case for expected_value == 2 */
    if ( expected_value == 2 )
    {
      std::vector<int> pLits(2);
      pLits[0] = abc::Abc_Var2Lit( var_index + 4 * ( NumGates - 2 ) + 3, 0 );
      pLits[1] = abc::Abc_Var2Lit( var_index + 4 * ( NumGates - 1 ) + 3, 0 );
      if ( !add_clause( solver, pLits, 2 ) )
      {
        return false;
      }
      pLits[0] = abc::Abc_LitNot( pLits[0] );
      pLits[1] = abc::Abc_LitNot( pLits[1] );
      if ( !add_clause( solver, pLits, 2 ) )
      {
        return false;
      }
    }

    var_index += 4 * NumGates;

    return true;
  }

private:
  std::array<std::array<std::array<int, NumGates + NumVars - 1>, 3>, NumGates> select;
  std::array<std::vector<int>, NumGates - 2> output_lits;

  abc::bmcg_sat_solver* solver;

  int var_index = 1;

  std::ofstream file;
  std::string str;

  double solving_time = 0.0;
};
}

template<uint32_t NumVars, uint32_t NumGates>
boost::optional<xmg_graph> xmg_exact_decomposition (std::string filename)
{
  detail::xmg_exact_decomposition_impl<NumVars, NumGates> impl;
  const auto result = impl.run(filename);
  impl.print_statistics();
  if ( result )
  {
    impl.print_solution();

    return impl.extract_circuit();
  }
  else
  {
    std::cout << "[i] no circuit found" << std::endl;
    return boost::none;
  }
}
}

#endif

// Local Variables:
// c-basic-offset: 2
// eval: (c-set-offset 'substatement-open 0)
// eval: (c-set-offset 'innamespace 0)
// End:


//
// Created by bevan on 29/11/2020.
//

#ifndef LIBAPR_APRDENOISE_H
#define LIBAPR_APRDENOISE_H

#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <unsupported/Eigen/IterativeSolvers>

#include <random>
#include <unordered_set>

#include "data_structures/Mesh/PixelData.hpp"
#include "io/hdf5functions_blosc.h"
#include "data_structures/APR/APR.hpp"
#include "data_structures/APR/particles/ParticleData.hpp"
#include "numerics/APRReconstruction.hpp"

//Helper classes

template<typename stencilType>
class Stencil {
public:
    std::vector<stencilType> linear_coeffs;

    std::vector<int> stencil_dims = {0, 0, 0};

};

class StencilSetUp {

public:

    StencilSetUp(APR& apr) {
      auto it = apr.iterator();
      dim = it.number_dimensions();
    }

    size_t dim;
    std::vector<int64_t> index;
    std::vector<std::vector<int>> stencil_index_base;

    PixelData<double> stencil_index_1;
    PixelData<double> stencil_index_2;

    int stencil_span;

    std::vector<int> stencil_dims = {0, 0, 0};

    size_t center_index;

    std::vector<uint64_t> l_index_1;

    std::vector<uint64_t> nl_index_1;
    std::vector<uint64_t> nl_index_2;

    void stencil_l_index() {

        l_index_1.resize(stencil_index_base.size() - 1, 0);

        size_t counter = 0;

        for (size_t i = 0; i < l_index_1.size(); ++i) {
            if (i != center_index) {
                l_index_1[counter] = i;
                counter++;
            }
        }

    }

    void setup_standard(std::vector<int> &stencil_dims_in) {

        stencil_dims[0] = stencil_dims_in[0];
        stencil_dims[1] = stencil_dims_in[1];
        stencil_dims[2] = stencil_dims_in[2];

        std::vector<int> temp;
        temp.resize(3);

        uint64_t counter = 0;

        for (int i = -stencil_dims[2]; i < (stencil_dims[2] + 1); ++i) {
            for (int j = -stencil_dims[1]; j < (stencil_dims[1] + 1); ++j) {
                for (int k = -stencil_dims[0]; k < (stencil_dims[0] + 1); ++k) {

                    temp[0] = k;
                    temp[1] = j;
                    temp[2] = i;
                    stencil_index_base.push_back(temp);

                    if ((i == 0) && (j == 0) && (k == 0)) {
                        center_index = counter;
                    } else {
                        l_index_1.push_back(counter);
                    }

                    counter++;

                }
            }
        }

        stencil_span = std::max(stencil_dims[0], std::max(stencil_dims[1], stencil_dims[2]));

    }

    template<typename T>
    void calculate_global_index(PixelData<T> &img) {

        uint64_t x_num = img.x_num + 2 * stencil_span;
        uint64_t y_num = img.y_num + 2 * stencil_span;

        index.resize(stencil_index_base.size());

        for (size_t i = 0; i < stencil_index_base.size(); ++i) {
            index[i] = stencil_index_base[i][0] + stencil_index_base[i][1] * y_num +
                       stencil_index_base[i][2] * y_num * x_num;
        }

    }

};

class APRStencils {

public:
    std::vector<Stencil<float>> stencils;
    int dim;
    int number_levels;

    void init(int _number_levels,int _dim)  {
      number_levels = _number_levels;
      dim = _dim;
      stencils.resize(number_levels+1);
    }

  void read_stencil(const std::string& file_name) {

    hid_t fileId = H5Fopen(file_name.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);
    hid_t base = H5Gopen2(fileId, "/", H5P_DEFAULT);

    int number_levels;
    int dim;
    //stored relative to maximum level to allow the read and write functions to be used with different sized APRs. (inherent assumption of only the image size changing not resolution).
    readAttr(H5T_NATIVE_INT, "number_levels", base, &number_levels);
    readAttr(H5T_NATIVE_INT, "dim", base, &dim);

    this->dim = dim;
    this->number_levels = number_levels;

    stencils.resize(number_levels + 1);

    for (int d_level = 0; d_level <= number_levels; ++d_level) {

      auto &stencil = stencils[d_level];

      std::string level_name = "_level_" + std::to_string(d_level);

      int num_pts_1;
      int num_pts_2;

      //get the number of points
      readAttr(H5T_NATIVE_INT, "num_pts_l" + level_name, base, &num_pts_1);

      //read in the full stencil
      std::vector<double> coeff_full;
      coeff_full.resize(num_pts_1 + num_pts_2);

      std::string coeff_n = "coeff" + level_name;

      hdf5_load_data_blosc(base, H5T_NATIVE_DOUBLE, coeff_full.data(), coeff_n.c_str());

      //now compute the linear stencil

      stencil.linear_coeffs.resize(num_pts_1, 0); //need to include the 0 center
      auto offset = 0;

      for (size_t k1 = 0; k1 < stencil.linear_coeffs.size(); ++k1) {

        stencil.linear_coeffs[k1] = coeff_full[k1 + offset];

      }

    }

    H5Fclose(fileId);

  }

  void write_stencil(std::string &file_name) {

    hid_t fileId = hdf5_create_file_blosc(file_name);

    //hid_t base = H5Gcreate2(fileId, "/a",  H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);

    if(this->number_levels==0){
      std::cerr << "Attempting to write non-intialized APR Stencils" << std::endl;
    }

    int _number_levels = this->number_levels;
    //stored relative to maximum level to allow the read and write functions to be used with different sized APRs. (inherent assumption of only the image size changing not resolution).
    writeAttr(H5T_NATIVE_INT, "number_levels", fileId, &_number_levels);

    writeAttr(H5T_NATIVE_INT, "dim", fileId, &(this->dim));

    for (int d_level = 0; d_level <= number_levels; ++d_level) {

      auto &stencil = stencils[d_level];

      std::string level_name = "_level_" + std::to_string(d_level);

      int num_pts_linear = stencil.linear_coeffs.size();

      //get the number of points
      writeAttr(H5T_NATIVE_INT, "num_pts_l" + level_name, fileId, &num_pts_linear);

      int dim1 = stencil.stencil_dims[0];
      int dim2 = stencil.stencil_dims[1];
      int dim3 = stencil.stencil_dims[2];

      writeAttr(H5T_NATIVE_INT, "dim_1" + level_name, fileId, &dim1);
      writeAttr(H5T_NATIVE_INT, "dim_2" + level_name, fileId, &dim2);
      writeAttr(H5T_NATIVE_INT, "dim_3" + level_name, fileId, &dim3);

      //read in the full stencil
      std::vector<double> coeff_full;
      coeff_full.resize(num_pts_linear);

      auto offset = 0;

      for (size_t k1 = 0; k1 < stencil.linear_coeffs.size(); ++k1) {

        coeff_full[k1 + offset] = stencil.linear_coeffs[k1];

      }

      writeData(H5T_NATIVE_DOUBLE, "coeff" + level_name, coeff_full, fileId);

    }

    H5Fclose(fileId);

  }

private:

    static void readAttr(hid_t type, std::string name, hid_t aGroupId, void *aDest) {
        hid_t attr_id = H5Aopen(aGroupId, name.c_str(), H5P_DEFAULT);
        H5Aread(attr_id, type, aDest);
        H5Aclose(attr_id);
    }

    static void writeAttr(hid_t type, std::string name, hid_t aGroupId, const void *const aSrc) {
        hsize_t dims[] = {1};
        hdf5_write_attribute_blosc(aGroupId, type, name.c_str(), 1, dims, aSrc);
    }
    template<typename T>
    static void writeData(hid_t type, std::string name, T aContainer, hid_t location) {
        hsize_t dims[] = {aContainer.size()};
        const hsize_t rank = 1;
        hdf5_write_data_blosc(location, type, name.c_str(), rank, dims, aContainer.data(), BLOSC_ZSTD, 1l, 0);
    }

};

std::vector<int> sample_without_replacement(int k, int N, std::default_random_engine &gen) {
    // Sample k elements from the range [1, N] without replacement
    // k should be <= N

    // Create an unordered set to store the samples
    std::unordered_set<int> samples;

    // Sample and insert values into samples
    for (int r = N - k; r < N; ++r) {
        int v = std::uniform_int_distribution<>(1, r)(gen);
        if (!samples.insert(v).second) samples.insert(r);
    }

    // Copy samples into vector
    std::vector<int> result(samples.begin(), samples.end());

    // Shuffle vector
    std::shuffle(result.begin(), result.end(), gen);

    return result;
}

class APRDenoise {

public:

    bool verbose = true;

    int max_level = 4;
    int others_level = 1;

    int number_levels = 4;

    int iteration_max = 1000;
    int iteration_others = 600;

    int N_max = 1000;
    int N_ = 1000;

    int level_min = 0;

    bool estimate_center_flag = true;

    int train_factor = 3; // minimum number of points above the stencil size required for the stencil to train.

    float tolerance = 0.05;

    // Train APR

    template<typename S>
    void train_denoise(APR &apr, ParticleData <S> &parts, APRStencils &aprStencils) {

        train_denoise(apr, parts, parts, aprStencils);
    }

    template<typename S>
    void train_denoise(APR &apr, ParticleData <S> &parts, ParticleData <S> &parts_gt, APRStencils &aprStencils) {
        //
        //  Trains a level dependent de-noising model
        //

        auto it = apr.iterator();
        int viable_levels = 0;

        for(int level = it.level_max(); level >= it.level_min(); level--){
            uint64_t total_parts = it.particles_level_end(level) - it.particles_level_begin(level);

            //check if enough particles to train a kernel; (K*size of kernel?)

            int stencil_sz;
            if (level == apr.level_max()) {
              stencil_sz = max_level;

            } else {
              stencil_sz = others_level;

            }

            uint64_t pts = std::pow(2*stencil_sz+1,it.number_dimensions())*train_factor;

            if(total_parts >= pts){
                viable_levels++;
            }

        }

        if(viable_levels < this->number_levels) {
          std::cout << "Not enough particles at levels to train kernel of desired number of levels" << std::endl;
          std::cout << "Setting number of kernel levels to: " << viable_levels << std::endl;
          this->number_levels = viable_levels;
        }

        aprStencils.init(this->number_levels,it.number_dimensions());

        int stencil_level = aprStencils.number_levels;

        float tol_ = this->tolerance;

        APRTimer timer(this->verbose);

        timer.start_timer("train");

        for (int level = apr.level_max(); level >= apr.level_min(); --level) {

            StencilSetUp setUp(apr);

            int stencil_sz;

            int it;
            int N;

            if (level == apr.level_max()) {
                stencil_sz = max_level;
                it = iteration_max;
                N = N_max;
            } else {
                it = iteration_others;
                N = N_;
                stencil_sz = others_level;

            }

            //set the dimension of the stencils learned.
            std::vector<int> stencil_dim;
            stencil_dim.resize(3);

            if (aprStencils.dim == 3) {
                stencil_dim[0] = stencil_sz;
                stencil_dim[1] = stencil_sz;
                stencil_dim[2] = stencil_sz;

            } else if (aprStencils.dim == 2) {

                stencil_dim[0] = stencil_sz;
                stencil_dim[1] = stencil_sz;
                stencil_dim[2] = 0;

            } else if (aprStencils.dim == 1) {

                stencil_dim[0] = stencil_sz;
                stencil_dim[1] = 0;
                stencil_dim[2] = 0;

            }

            setUp.setup_standard(stencil_dim);

            //change the stencil size, the stencil is way to big for the lower levels.
            assemble_system_guided(apr, parts, parts_gt, setUp, aprStencils.stencils[stencil_level], N, it, level,
                                   tol_,verbose);

            tol_ = tol_ / 8; // check this is this necessary, and I don't think the decay makes much sense beyond the first two layers. (This was based on a noise reduction per level).

            if(stencil_level > 1){
              stencil_level--;
            } else {
              break;
            }

        }

        timer.stop_timer();

    }

// Apply model APR

    template<typename R,typename S>
    void apply_denoise(APR &apr, ParticleData <R> &parts_in, ParticleData <S> &parts_out, APRStencils &aprStencils) {

        APRTimer timer(this->verbose);

        int stencil_level = aprStencils.number_levels;

        timer.start_timer("apply");

        for (int level = apr.level_max(); level >= std::max((int) apr.level_min(), level_min); --level) {

            StencilSetUp setUp(apr);
            setUp.setup_standard(aprStencils.stencils[stencil_level].stencil_dims);

            apply_conv_guided(apr, parts_in, parts_out, setUp, aprStencils.stencils[stencil_level], level);

            if(stencil_level > 1){
                stencil_level--;
            }

        }
        timer.stop_timer();

    }

    template<typename T, typename R,typename S>
    float
    apply_conv_guided(APR &apr, ParticleData <T> &parts_in,ParticleData <R> &parts_out, StencilSetUp &stencilSetUp, Stencil<S> &stencil, int level) {

        APRTimer timer(this->verbose);

        PixelData<T> img; //change to be level dependent

        if(parts_out.size() != parts_in.size()) {
          parts_out.init(apr);
        }

        int delta = level - apr.level_max();

        APRReconstruction::interp_img_us_smooth(apr, img, parts_in, false, delta);

        stencilSetUp.calculate_global_index(img);

        PixelData<T> pad_img;

        timer.start_timer("pad_img");

        paddPixels(img, pad_img, stencilSetUp.stencil_span, stencilSetUp.stencil_span, stencilSetUp.stencil_span);

        timer.stop_timer();

        const uint64_t off_y = (uint64_t) stencilSetUp.stencil_span;
        const uint64_t off_x = (uint64_t) std::min((int) stencilSetUp.stencil_span, img.x_num - 1);
        const uint64_t off_z = (uint64_t) std::min((int) stencilSetUp.stencil_span, img.z_num - 1);

        const uint64_t x_num_p = img.x_num + 2 * off_x;
        const uint64_t y_num_p = img.y_num + 2 * off_y;

        std::vector<float> local_vec;

        local_vec.resize(stencilSetUp.index.size(), 0);

        const int stencil_size = stencilSetUp.index.size();
        const int unroll_size = 16;
        const int loop_sz = (stencil_size / unroll_size) * unroll_size;

        timer.start_timer("conv guided");

        auto it = apr.iterator();

        int z = 0;
#ifdef HAVE_OPENMP
#pragma omp parallel for schedule(dynamic) private(z) firstprivate(local_vec, it)
#endif
        for (z = 0; z < img.z_num; ++z) {
            for (int x = 0; x < img.x_num; ++x) {

                float temp_val = 0;

                const uint64_t global_off = off_y + (off_x + x) * y_num_p + (off_z + z) * x_num_p * y_num_p;

                for (it.begin(level, z, x); it < it.end(); it++) {

                    uint64_t y = it.y();

                    //Get the local stencil of points
                    const uint64_t global_off_l = y + global_off;

                    int i = 0;

                    for (; i < loop_sz; i += unroll_size) {
                        local_vec[i] = pad_img.mesh[global_off_l + stencilSetUp.index[i]];
                        local_vec[i + 1] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 1]];
                        local_vec[i + 2] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 2]];
                        local_vec[i + 3] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 3]];
                        local_vec[i + 4] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 4]];
                        local_vec[i + 5] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 5]];
                        local_vec[i + 6] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 6]];
                        local_vec[i + 7] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 7]];
                        local_vec[i + 8] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 8]];
                        local_vec[i + 9] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 9]];
                        local_vec[i + 10] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 10]];
                        local_vec[i + 11] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 11]];
                        local_vec[i + 12] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 12]];
                        local_vec[i + 13] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 13]];
                        local_vec[i + 14] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 14]];
                        local_vec[i + 15] = pad_img.mesh[global_off_l + stencilSetUp.index[i + 15]];
                    }

                    for (; i < stencil_size; i++) {
                        local_vec[i] = pad_img.mesh[global_off_l + stencilSetUp.index[i]];
                    }

                    parts_out[it] = std::inner_product(local_vec.begin(),
                                                 local_vec.end(),
                                                 stencil.linear_coeffs.begin(),
                                                 temp_val);

                }
            }
        }

        timer.stop_timer();

        float time = timer.timings.back();
        return time;

    }

    template<typename T, typename S>
    void
    assemble_system_guided(APR &apr, ParticleData <T> &parts, ParticleData <T> &parts_g, StencilSetUp &stencilSetUp,
                           Stencil<S> &stencil, uint64_t N, int num_rep, int level, float factor = 0.05,
                           bool verbose = false) {

        APRTimer timer(this->verbose);

        PixelData<T> img; //change to be level dependent
        // apr.interp_img(img,parts);
        int delta = (level - apr.level_max());

        auto level_ = level;

        auto apr_iterator = apr.iterator();

        uint64_t total_parts = apr_iterator.particles_level_end(level_) - apr_iterator.particles_level_begin(level_);

        APRReconstruction::interp_img_us_smooth(apr, img, parts, false, delta);
        stencilSetUp.calculate_global_index(img);

        if (total_parts < (uint64_t) N) {
            N =  stencilSetUp.l_index_1.size()*this->train_factor;
        }

        PixelData<T> pad_img;

        timer.start_timer("pad_img");

        paddPixels(img, pad_img, stencilSetUp.stencil_span, stencilSetUp.stencil_span, stencilSetUp.stencil_span);

        timer.stop_timer();

        const uint64_t off_y = (uint64_t) stencilSetUp.stencil_span;
        const uint64_t off_x = (uint64_t) stencilSetUp.stencil_span;
        const uint64_t off_z = (uint64_t) std::min((int) stencilSetUp.stencil_span, img.z_num - 1);

        const uint64_t x_num_p = img.x_num + 2 * off_x;
        const uint64_t y_num_p = img.y_num + 2 * off_y;

        std::vector<T> local_vec;

        local_vec.resize(stencilSetUp.index.size(), 0);

        std::vector<float> nl_local_vec;
        nl_local_vec.resize(stencilSetUp.nl_index_1.size(), 0);

        size_t n =  stencilSetUp.l_index_1.size() + stencilSetUp.nl_index_1.size();

        std::vector<uint64_t> random_index;

        random_index.resize(N);

        auto l_num = stencilSetUp.l_index_1.size();

        //PixelData<float> A_temp;
        //A_temp.init(n,N,1);

        std::vector<double> b_temp;
        b_temp.resize(N);

        std::vector<std::vector<double>> coeff_store;
        coeff_store.resize(num_rep);

        timer.start_timer("assemble");

        auto total_p = (img.x_num * img.z_num * img.y_num) / N;

        total_p = total_parts / N;

        std::vector<double> A_temp;
        A_temp.resize(n * N, 0);

        Eigen::Map<Eigen::MatrixXd> A(A_temp.data(), n, N);

        //Eigen::MatrixXd A(n,N);

        Eigen::VectorXd coeff_prev(n);

        std::vector<uint16_t> x_vec;
        std::vector<uint16_t> y_vec;
        std::vector<uint16_t> z_vec;

        std::vector<uint64_t> global_index;

        int z = 0;
        int x = 0;

        for (z = 0; z < apr_iterator.z_num(level_); z++) {
            for (x = 0; x < apr_iterator.x_num(level_); ++x) {

                for (apr_iterator.begin(level_, z, x); apr_iterator < apr_iterator.end(); apr_iterator++) {

                    x_vec.push_back(x);
                    z_vec.push_back(z);
                    y_vec.push_back(apr_iterator.y());

                    global_index.push_back(apr_iterator.global_index());

                }
            }
        }

        timer.stop_timer();

        timer.start_timer("solve loop");

        for (int l = 0; l < num_rep; ++l) {

            APRTimer timer_l(verbose);

            timer_l.start_timer("random");

            random_index[0] = std::rand() % total_p + 1;

            for (size_t j = 1; j < random_index.size(); ++j) {

                random_index[j] = random_index[j - 1] + std::rand() % total_p + 1;
            }

            timer_l.stop_timer();

            uint64_t k = 0;

            timer_l.start_timer("A");

#ifdef HAVE_OPENMP
#pragma omp parallel for schedule(static) private(k) firstprivate(local_vec)
#endif
            for (k = 0; k < N; ++k) {

                const uint64_t r_i = random_index[k];

                const uint64_t z = z_vec[r_i];
                const uint64_t x = x_vec[r_i];
                const uint64_t y = y_vec[r_i];

                const uint64_t global_off = off_y + (off_x + x) * y_num_p + (off_z + z) * x_num_p * y_num_p;

                //Get the local stencil of points
                const uint64_t global_off_l = y + global_off;

                for (size_t i = 0; i < stencilSetUp.index.size(); ++i) {
                    local_vec[i] = pad_img.mesh[global_off_l + stencilSetUp.index[i]];
                }

                //b_temp[k] = local_vec[stencilSetUp.center_index];
                b_temp[k] = parts_g[global_index[r_i]];

                for (size_t i = 0; i < stencilSetUp.l_index_1.size(); ++i) {
                    //A_temp(i, k, 0) = local_vec[stencilSetUp.l_index_1[i]];
                    A(i, k) = local_vec[stencilSetUp.l_index_1[i]];
                }

                //Apply the non-linear kernel
                for (size_t i = 0; i < stencilSetUp.nl_index_1.size(); ++i) {
                    //A_temp(i + l_num, k, 0) =
                    //      local_vec[stencilSetUp.nl_index_1[i]] * local_vec[stencilSetUp.nl_index_2[i]];
                    A(i + l_num, k) =
                            local_vec[stencilSetUp.nl_index_1[i]] * local_vec[stencilSetUp.nl_index_2[i]];
                }

            }

            timer_l.stop_timer();

            timer_l.start_timer("solve");

            std::vector<double> norm_c;
            norm_c.resize(n, 0);

            Eigen::Map<Eigen::VectorXd> norm_c_(norm_c.data(), n);

            std::vector<double> std;

            for (int i = 0; i < A.rows(); i++) {
                Eigen::ArrayXd vec = A.row(i);
                double std_dev = std::sqrt((vec - vec.mean()).square().sum() / (vec.size() - 1));
                //std_dev = 1;
                norm_c_(i) = std_dev;
                A.row(i) = A.row(i) / std_dev;
                std.push_back(std_dev);
            }

            Eigen::Map<Eigen::VectorXd> b(b_temp.data(), N);

            coeff_store[l].resize(n, 0);

            if (l == 0) {
                float eps = .01;
                for (size_t i = 0; i < n; ++i) {
                    coeff_store[l][i] = std::rand() * 2 * eps - eps;
                }

            }

            Eigen::Map<Eigen::VectorXd> coeff(coeff_store[l].data(), n);

            Eigen::LeastSquaresConjugateGradient<Eigen::MatrixXd, Eigen::IdentityPreconditioner> solver;

            //Need to compute desired tol
            float val = factor;
            Eigen::VectorXf ones_v = val * Eigen::VectorXf::Ones(N);

            float norm_b = b.norm();
            float norm_e = ones_v.norm();

            float tol = norm_e / norm_b;

            solver.setMaxIterations(800);
            solver.setTolerance(tol);
            solver.compute(A.transpose());

            for (int m = 0; m < coeff_prev.size(); ++m) {
                coeff_prev(m) = coeff_prev(m) * norm_c_(m);
            }

            if (l > 0) {
                solver.setMaxIterations(20);
                coeff = solver.solveWithGuess(b, coeff_prev);
            } else {
                tol *= 0.0001;
                solver.setTolerance(tol);
                coeff = solver.solve(b);

            }

            for (size_t i1 = 0; i1 < n; ++i1) {
                coeff(i1) = coeff(i1) / norm_c_(i1);
            }

            coeff_prev = coeff;

            if (verbose) {
                std::cout << "#iterations: " << solver.iterations() << std::endl;
            }

            timer_l.stop_timer();

        }

        timer.stop_timer();

        //now compute the linear stencil

        stencil.linear_coeffs.resize(stencilSetUp.index.size(), 0); //need to include the 0 center
        auto offset = 0;

        stencil.stencil_dims[0] = stencilSetUp.stencil_dims[0];
        stencil.stencil_dims[1] = stencilSetUp.stencil_dims[1];
        stencil.stencil_dims[2] = stencilSetUp.stencil_dims[2];

        int include = std::floor(0.5 * num_rep);

        for (size_t k1 = 0; k1 < stencil.linear_coeffs.size(); ++k1) {

            if (k1 == stencilSetUp.center_index) {
                offset = -1;
            } else {

                double sum = 0;
                double counter = 0;

                for (int i = include; i < num_rep; ++i) {
                    sum += coeff_store[i][k1 + offset];
                    counter++;
                }
                stencil.linear_coeffs[k1] = sum / (counter * 1.0);

            }

        }

        if(this->estimate_center_flag) {
          // set the center pixel to the largest weight in the kernel, and then re-normalise the kernel

          float factor = 0;
          float sum = 0;

          for (size_t j = 0; j < stencil.linear_coeffs.size(); ++j) {
            factor = std::max(factor, std::abs(stencil.linear_coeffs[j]));
            sum += stencil.linear_coeffs[j];
          }

          stencil.linear_coeffs[stencilSetUp.center_index] = factor;

          for (size_t j = 0; j < stencil.linear_coeffs.size(); ++j) {
            stencil.linear_coeffs[j] = stencil.linear_coeffs[j]*(sum)/(sum+factor);
          }

        }

        //NAN check, this can occur if you have noise free data, and the system becomes ill-posed. (likely other reasons as well, hence the warning).

        bool valid_check = true;

        double coeff_sum = 0;

        for (size_t l1 = 0; l1 < stencil.linear_coeffs.size(); ++l1) {
          if(std::isnan(stencil.linear_coeffs[l1])){
            valid_check = false;
          }
            coeff_sum += stencil.linear_coeffs[l1];
        }

        // Does it sum close to 1? very liberal to detect agian non-convergence.
        float error_threshold = 0.1;
        double val = coeff_sum - 1.0;

        if(std::abs(val) > error_threshold){
            valid_check = false;
        }

        std::cout << "kernel sum: " << coeff_sum << std::endl;

        bool diverged_val = std::isnan(coeff_sum);

        if(diverged_val){
            valid_check = false;
        }

        if(!valid_check){
            std::wcerr << "Inference hasn't converged to a solution, setting the kernel to identity. This could be due to std(signal) = 0, or the signal is noise free." << std::endl;

            std::fill(stencil.linear_coeffs.begin(),stencil.linear_coeffs.end(),0);
            stencil.linear_coeffs[stencilSetUp.center_index] = 1;
        }

    }

};

#endif //LIBAPR_APRDENOISE_H


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG, www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also tools/license/license.mtl.txt in the distribution.

#ifndef MTL_IO_WRITE_AST_DISPATCH_INCLUDE
#define MTL_IO_WRITE_AST_DISPATCH_INCLUDE

#include <string>
#include <sstream>
#include <fstream>
#include <boost/utility/enable_if.hpp>
#include <boost/type_traits/is_integral.hpp>
#include <boost/type_traits/is_floating_point.hpp>
#include <boost/numeric/mtl/mtl_fwd.hpp>
#include <boost/numeric/mtl/io/functor_symbol.hpp>
#include <boost/numeric/mtl/vector/vec_vec_aop_expr.hpp>

namespace mtl { namespace io {

template <typename Value, typename Parameters>
void write_ast_dispatch(const mtl::dense_vector<Value, Parameters>& v, std::string s, std::ofstream& f);
template <typename E1, typename E2, typename SFunctor>
void write_ast_dispatch(const mtl::vec_vec_aop_expr<E1, E2, SFunctor>& expr, std::string s, std::ofstream& f);
template <class E1, class E2, typename SFunctor>
void write_ast_dispatch(const mtl::vec_vec_pmop_expr<E1, E2, SFunctor>& expr, std::string s, std::ofstream& f);
template <typename Functor, typename Vector>
void write_ast_dispatch(const mtl::map_view<Functor, Vector>& expr, std::string s, std::ofstream& f);
template <typename Scaling, typename Vector>
void write_ast_dispatch(const mtl::scaled_view<Scaling, Vector>& expr, std::string s, std::ofstream& f);
template <typename Matrix, typename Vector>
void write_ast_dispatch(const mtl::mat_cvec_times_expr<Matrix, Vector>& expr, std::string s, std::ofstream& f);
template <typename Expr>
void write_ast_dispatch(const mtl::operation::compute_summand<Expr>& expr, std::string s, std::ofstream& f);
template <typename Matrix, typename Vector>
void write_ast_dispatch(const mtl::operation::compute_summand<mtl::mat_cvec_times_expr<Matrix, Vector> >& expr, std::string s, std::ofstream& f);
template <typename Vector1, typename Vector2>
void write_ast_dispatch(const mtl::mat::outer_product_matrix<Vector1, Vector2>& expr, std::string s, std::ofstream& f);

template <typename Value>
typename boost::enable_if_c<boost::is_floating_point<Value>::value || boost::is_integral<Value>::value>::type
write_ast_dispatch(const Value& v, std::string s, std::ofstream& f)
{
    f << "  " << s << "[shape=box,label=\"scalar\\n" << v << "\"]\n";
}

template <typename Value, typename Parameters>
void write_ast_dispatch(const mtl::dense_vector<Value, Parameters>& v, std::string s, std::ofstream& f)
{
    f << "  " << s << "[shape=box,label=\"vector\\n" << &v << "\"]\n";
}

template <typename E1, typename E2, typename SFunctor>
void write_ast_dispatch(const mtl::vec_vec_aop_expr<E1, E2, SFunctor>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"" << functor_symbol(SFunctor()) << "\"]\n";
    std::string target= s + "t", source= s + "s";
    write_ast_dispatch(expr.first_argument(), target, f);
    write_ast_dispatch(expr.second_argument(), source, f);
    f << "  " << s << "->" << target << '\n';
    f << "  " << s << "->" << source << '\n';
}

template <class E1, class E2, typename SFunctor>
void write_ast_dispatch(const mtl::vec_vec_pmop_expr<E1, E2, SFunctor>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"" << functor_symbol(SFunctor()) << "\"]\n";
    std::string first= s + "f", second= s + "s";
    write_ast_dispatch(expr.first_argument(), first, f);
    write_ast_dispatch(expr.second_argument(), second, f);
    f << "  " << s << "->" << first << '\n';
    f << "  " << s << "->" << second << '\n';
}

template <typename Scaling, typename Vector>
void write_ast_dispatch(const mtl::scaled_view<Scaling, Vector>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"scaled_view\"]\n";
    std::string functor= s + "f", ref= s + "r";

    write_ast_dispatch(expr.functor.value, functor, f);
    write_ast_dispatch(expr.ref, ref, f);

    f << "  " << s << "->" << functor << '\n';
    f << "  " << s << "->" << ref << '\n';
}

template <typename Value1, typename Value2>
void write_ast_dispatch(const mtl::tfunctor::scale<Value1, Value2, mtl::tag::scalar>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"scale\"]\n";
    std::string factor= s + "f", ref= s + "r";
    write_ast_dispatch(expr.value(), factor, f);
    f << "  " << ref << "[label=\".\"]\n";

    f << "  " << s << "->" << factor << '\n';
    f << "  " << s << "->" << ref << '\n';
}

template <typename Functor, typename Vector>
void write_ast_dispatch(const mtl::map_view<Functor, Vector>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"map\"]\n";
    std::string functor= s + "f", ref= s + "r";

    write_ast_dispatch(expr.functor, functor, f);
    write_ast_dispatch(expr.ref, ref, f);

    f << "  " << s << "->" << functor << '\n';
    f << "  " << s << "->" << ref << '\n';
}

// template <typename Matrix, typename Vector>
// void write_ast_dispatch(const mtl::mat_cvec_times_expr<Matrix, Vector>& expr, std::string s, std::ofstream& f)
// {
// }

template <typename Expr>
void write_ast_dispatch(const mtl::operation::compute_summand<Expr>& expr, std::string s, std::ofstream& f)
{
    write_ast_dispatch(expr.value, s, f);
}

template <typename Matrix, typename Vector>
void write_ast_dispatch(const mtl::operation::compute_summand<mtl::mat_cvec_times_expr<Matrix, Vector> >& expr, std::string s, std::ofstream& f)
{
# ifdef NDEBUG
    write_ast_dispatch(expr.value, s, f);
# else
    f << "  " << s << "[label=\"*\"]\n";
    std::string first= s + "f", second= s + "s";
    write_ast_dispatch(expr.first, first, f);
    write_ast_dispatch(expr.second, second, f);
    f << "  " << s << "->" << first << '\n';
    f << "  " << s << "->" << second << '\n';
# endif
}

template <typename Vector1, typename Vector2>
void write_ast_dispatch(const mtl::mat::outer_product_matrix<Vector1, Vector2>& expr, std::string s, std::ofstream& f)
{
    f << "  " << s << "[label=\"outer product\"]\n";
    std::string first= s + "f", second= s + "s";
    write_ast_dispatch(expr.v1(), first, f);
    write_ast_dispatch(expr.v2(), second, f);
    f << "  " << s << "->" << first << '\n';
    f << "  " << s << "->" << second << '\n';
}

}} // namespace mtl::io

#endif // MTL_IO_WRITE_AST_DISPATCH_INCLUDE