//
// Created by Sandra Alfaro on 24/05/18.
//

#include <vector>
#include <chrono>
#include <atomic>
#include <Eigen/Geometry>                 // MatrixBase.homogeneous()
#include <Eigen/SVD>
#include <Eigen/Core>                     // Transform.computeRotationScaling()

#ifdef OpenGR_USE_OPENMP
#include <omp.h>
#endif

#include "gr/utils/sampling.h"
#include "gr/utils/shared.h"
#include "gr/utils/logger.h"
#include "gr/accelerators/kdtree.h"
#include "match4pcsBase.h"

#ifdef TEST_GLOBAL_TIMINGS
#   include "gr/utils/timer.h"
#endif

namespace gr {
    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::Match4pcsBase (const OptionsType& options
            , const Utils::Logger& logger)
            : MatchBaseType(options,logger)
            , fun_(MatchBaseType::sampled_Q_3D_,MatchBaseType::base_3D_,MatchBaseType::options_)
    {
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::~Match4pcsBase() {}

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    bool Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::TryQuadrilateral(
        typename PointType::Scalar &invariant1,
        typename PointType::Scalar &invariant2,
        int &id1, int &id2, int &id3, int &id4) {

        Scalar min_distance = std::numeric_limits<Scalar>::max();
        int best1, best2, best3, best4;
        best1 = best2 = best3 = best4 = -1;
        for (int i = 0; i < 4; ++i) {
            for (int j = 0; j < 4; ++j) {
                if (i == j) continue;
                int k = 0;
                while (k == i || k == j) k++;
                int l = 0;
                while (l == i || l == j || l == k) l++;
                Scalar local_invariant1;
                Scalar local_invariant2;
                // Compute the closest points on both segments, the corresponding
                // invariants and the distance between the closest points.
                Scalar segment_distance = distSegmentToSegment(
                        MatchBaseType::base_3D_[i]->pos(), MatchBaseType::base_3D_[j]->pos(),
                        MatchBaseType::base_3D_[k]->pos(), MatchBaseType::base_3D_[l]->pos(),
                        local_invariant1, local_invariant2);
                // Retail the smallest distance and the best order so far.
                if (segment_distance < min_distance) {
                    min_distance = segment_distance;
                    best1 = i;
                    best2 = j;
                    best3 = k;
                    best4 = l;
                    invariant1 = local_invariant1;
                    invariant2 = local_invariant2;
                }
            }
        }

        if(best1 < 0 || best2 < 0 || best3 < 0 || best4 < 0 ) return false;

        Coordinates tmp = MatchBaseType::base_3D_;
        MatchBaseType::base_3D_[0] = tmp[best1];
        MatchBaseType::base_3D_[1] = tmp[best2];
        MatchBaseType::base_3D_[2] = tmp[best3];
        MatchBaseType::base_3D_[3] = tmp[best4];

        CongruentBaseType tmpId = {id1, id2, id3, id4};
        id1 = tmpId[best1];
        id2 = tmpId[best2];
        id3 = tmpId[best3];
        id4 = tmpId[best4];

        return true;
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    bool Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::SelectQuadrilateral(
        Scalar &invariant1,
        Scalar &invariant2,
        int& base1, int& base2, int& base3, int& base4)  {

        const Scalar kBaseTooSmall (0.2);
        int current_trial = 0;

        // Try fix number of times.
        while (current_trial < MatchBaseType::kNumberOfDiameterTrials) {
            // Select a triangle if possible. otherwise fail.
            if (!MatchBaseType::SelectRandomTriangle(MatchBaseType::max_base_diameter_, base1, base2, base3)){
                return false;
            }

            const auto& b0 = *(MatchBaseType::base_3D_[0] = &MatchBaseType::sampled_P_3D_[base1]);
            const auto& b1 = *(MatchBaseType::base_3D_[1] = &MatchBaseType::sampled_P_3D_[base2]);
            const auto& b2 = *(MatchBaseType::base_3D_[2] = &MatchBaseType::sampled_P_3D_[base3]);

            // The 4th point will be a one that is close to be planar to the other 3
            // while still not too close to them.
            const double x1 = b0.pos()(0);
            const double y1 = b0.pos()(1);
            const double z1 = b0.pos()(2);
            const double x2 = b1.pos()(0);
            const double y2 = b1.pos()(1);
            const double z2 = b1.pos()(2);
            const double x3 = b2.pos()(0);
            const double y3 = b2.pos()(1);
            const double z3 = b2.pos()(2);

            // Fit a plan.
            Scalar denom = (-x3 * y2 * z1 + x2 * y3 * z1 + x3 * y1 * z2 - x1 * y3 * z2 -
                            x2 * y1 * z3 + x1 * y2 * z3);

            if (denom != 0) {
                Scalar A =
                        (-y2 * z1 + y3 * z1 + y1 * z2 - y3 * z2 - y1 * z3 + y2 * z3) / denom;
                Scalar B =
                        (x2 * z1 - x3 * z1 - x1 * z2 + x3 * z2 + x1 * z3 - x2 * z3) / denom;
                Scalar C =
                        (-x2 * y1 + x3 * y1 + x1 * y2 - x3 * y2 - x1 * y3 + x2 * y3) / denom;
                base4 = -1;
                Scalar best_distance = std::numeric_limits<Scalar>::max();
                // Go over all points in P.
                const Scalar too_small = std::pow(MatchBaseType::max_base_diameter_ * kBaseTooSmall, 2);
                for (unsigned int i = 0; i < MatchBaseType::sampled_P_3D_.size(); ++i) {
                    const auto &p = MatchBaseType::sampled_P_3D_[i];
                    if ((p.pos() - b0.pos()).squaredNorm() >= too_small &&
                        (p.pos() - b1.pos()).squaredNorm() >= too_small &&
                        (p.pos() - b2.pos()).squaredNorm() >= too_small) {
                        // Not too close to any of the first 3.
                        const Scalar distance =
                                std::abs(A * p.pos()(0) + B * p.pos()(1) + C * p.pos()(2) - 1.0);
                        // Search for the most planar.
                        if (distance < best_distance) {
                            best_distance = distance;
                            base4 = int(i);
                        }
                    }
                }
                // If we have a good one we can quit.
                if (base4 != -1) {
                    MatchBaseType::base_3D_[3] = &MatchBaseType::sampled_P_3D_[base4];
                    if(TryQuadrilateral(invariant1, invariant2, base1, base2, base3, base4))
                        return true;
                }
            }
            current_trial++;
        }

        // We failed to find good enough base..
        return false;
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    // Initialize all internal data structures and data members.
    void Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::Initialize() {
        fun_.Initialize();
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    bool Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::generateCongruents (
        CongruentBaseType &base, Set& congruent_quads) {
//      std::cout << "------------------" << std::endl;
        Scalar invariant1, invariant2;

        if(!initBase(base, invariant1, invariant2)) return false;

//        std::cout << "Found a new base !" << std::endl;
        const auto& b0 = *MatchBaseType::base_3D_[0];
        const auto& b1 = *MatchBaseType::base_3D_[1];
        const auto& b2 = *MatchBaseType::base_3D_[2];
        const auto& b3 = *MatchBaseType::base_3D_[3];

        // Computes distance between pairs.
        const Scalar distance1 = (b0.pos()- b1.pos()).norm();
        const Scalar distance2 = (b2.pos()- b3.pos()).norm();

        std::vector<std::pair<int, int>> pairs1, pairs2;

        // Compute normal angles.
        const Scalar normal_angle1 = (b0.normal() - b1.normal()).norm();
        const Scalar normal_angle2 = (b2.normal() - b3.normal()).norm();

        fun_.ExtractPairs(distance1, normal_angle1, MatchBaseType::distance_factor * MatchBaseType::options_.delta, 0, 1, &pairs1);
        fun_.ExtractPairs(distance2, normal_angle2, MatchBaseType::distance_factor * MatchBaseType::options_.delta, 2, 3, &pairs2);

//        std::cout << "Pair set 1 has " << pairs1.size() << " elements" << std::endl;
//        std::cout << "Pair set 2 has " << pairs2.size() << " elements" << std::endl;

//  Log<LogLevel::Verbose>( "Pair creation ouput: ", pairs1.size(), " - ", pairs2.size());

        if (pairs1.size() == 0 || pairs2.size() == 0) {
            return false;
        }

        if (!fun_.FindCongruentQuadrilaterals(invariant1, invariant2,
                                         MatchBaseType::distance_factor * MatchBaseType::options_.delta,
                                         MatchBaseType::distance_factor * MatchBaseType::options_.delta,
                                         pairs1,
                                         pairs2,
                                         &congruent_quads)) {
            return false;
        }

        return true;
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    bool Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::initBase (CongruentBaseType &base)
    {
        Scalar invariant1, invariant2; // dummy

        return initBase(base, invariant1, invariant2);
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    bool Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::initBase (CongruentBaseType &base, Scalar& invariant1, Scalar& invariant2)
    {
        #ifdef STATIC_BASE
        static bool first_time = true;

        if (first_time){
            std::cerr << "Warning: Running with static base" << std::endl;
            base[0] = 0;
            base[1] = 3;
            base[2] = 1;
            base[3] = 4;

            MatchBaseType::base_3D_[0] = &MatchBaseType::sampled_P_3D_ [base[0]];
            MatchBaseType::base_3D_[1] = &MatchBaseType::sampled_P_3D_ [base[1]];
            MatchBaseType::base_3D_[2] = &MatchBaseType::sampled_P_3D_ [base[2]];
            MatchBaseType::base_3D_[3] = &MatchBaseType::sampled_P_3D_ [base[3]];
            TryQuadrilateral(invariant1, invariant2, base[0], base[1], base[2], base[3]);

            first_time = false;
        }
        else
            return false;

        #else
            if (!SelectQuadrilateral(invariant1, invariant2, base[0], base[1], base[2], base[3])) {
                // std::cout << "Skipping wrong base" << std::endl;
                return false;
            }
        #endif

        return true;
    }

    template <template <typename, typename, typename> class _Functor,
              typename PointType,
              typename TransformVisitor,
              typename PairFilteringFunctor,
              template < class, class > class PFO>
    typename Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::Scalar
    Match4pcsBase<_Functor, PointType, TransformVisitor, PairFilteringFunctor, PFO>::distSegmentToSegment(
        const VectorType& p1, const VectorType& p2,
        const VectorType& q1, const VectorType& q2,
        Scalar& invariant1, Scalar& invariant2) {

        static const Scalar kSmallNumber = 0.0001;
        VectorType u = p2 - p1;
        VectorType v = q2 - q1;
        VectorType w = p1 - q1;
        Scalar a = u.dot(u);
        Scalar b = u.dot(v);
        Scalar c = v.dot(v);
        Scalar d = u.dot(w);
        Scalar e = v.dot(w);
        Scalar f = a * c - b * b;
        // s1,s2 and t1,t2 are the parametric representation of the intersection.
        // they will be the invariants at the end of this simple computation.
        Scalar s1 = 0.0;
        Scalar s2 = f;
        Scalar t1 = 0.0;
        Scalar t2 = f;

        if (f < kSmallNumber) {
            s1 = 0.0;
            s2 = 1.0;
            t1 = e;
            t2 = c;
        } else {
            s1 = (b * e - c * d);
            t1 = (a * e - b * d);
            if (s1 < 0.0) {
                s1 = 0.0;
                t1 = e;
                t2 = c;
            } else if (s1 > s2) {
                s1 = s2;
                t1 = e + b;
                t2 = c;
            }
        }

        if (t1 < 0.0) {
            t1 = 0.0;
            if (-d < 0.0)
                s1 = 0.0;
            else if (-d > a)
                s1 = s2;
            else {
                s1 = -d;
                s2 = a;
            }
        } else if (t1 > t2) {
            t1 = t2;
            if ((-d + b) < 0.0)
                s1 = 0;
            else if ((-d + b) > a)
                s1 = s2;
            else {
                s1 = (-d + b);
                s2 = a;
            }
        }
        invariant1 = (std::abs(s1) < kSmallNumber ? 0.0 : s1 / s2);
        invariant2 = (std::abs(t1) < kSmallNumber ? 0.0 : t1 / t2);

        return ( w + (invariant1 * u) - (invariant2 * v)).norm();
    }
}


#include <Eigen/Dense>
#include "disc.hpp"
#include "evaluations.hpp"
#include "global_residual.hpp"
#include "local_residual.hpp"
#include "macros.hpp"
#include "nested.hpp"
#include "qoi.hpp"
#include "state.hpp"

namespace calibr8 {

void eval_forward_jacobian(RCP<State> state, RCP<Disc> disc, int step) {

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();

  // gather information from the state object
  RCP<LocalResidual<FADT>> local = state->d_residuals->local;
  RCP<GlobalResidual<FADT>> global = state->d_residuals->global;
  Array1D<RCP<VectorT>>& RHS = state->la->b[GHOST];
  Array2D<RCP<MatrixT>>& LHS = state->la->A[GHOST];
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  // modify the fields if we are doing a verification
  RCP<NestedDisc> nested;
  bool const is_verification = (disc->type() == VERIFICATION);
  if (is_verification) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    x = nested->primal_fine(step).global;
    xi = nested->primal_fine(step).local;
    x_prev = nested->primal_fine(step - 1).global;
    xi_prev = nested->primal_fine(step - 1).local;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      global->gather(x, x_prev);

      // loop over domain ip sets
      // ip_set = 0 -> coupled
      // ip_set > 0 -> global only
      Array1D<int> ip_sets = global->ip_sets();
      int const num_ip_sets = ip_sets.size();

      for (int ip_set = 0; ip_set < num_ip_sets; ++ip_set) {

        // get the quadrature order for the ip set
        int const q_order = ip_sets[ip_set];
        // loop over all integration points in the current element
        int const npts = apf::countIntPoints(me, q_order);

        for (int pt = 0; pt < npts; ++pt) {

          // get integration point specific information
          apf::Vector3 iota;
          apf::getIntPoint(me, q_order, pt, iota);
          double const w = apf::getIntWeight(me, q_order, pt);
          double const dv = apf::getDV(me, iota);

          if (ip_set == 0) {

            // solve the local constitutive equations at the integration point
            // and store the resultant local residual and its derivatives (dC_dxi)
            global->interpolate(iota);
            local->gather(pt, xi, xi_prev);
            local->seed_wrt_xi();
            int path = local->solve_nonlinear(global);
            if (is_verification) {
              nested->branch_paths()[step][es][elem] = path;
            }
            local->scatter(pt, xi);
            EMatrix const dC_dxi = local->eigen_jacobian();

            // re-evaluate the constitutive equations to obtain dC_dx
            local->unseed_wrt_xi();
            global->seed_wrt_x();
            global->interpolate(iota);
            local->evaluate(global);
            EMatrix const dC_dx = local->eigen_jacobian();

            // solve the forward sensitivty system to obtain dxi_dx
            EMatrix const dxi_dx = dC_dxi.fullPivLu().solve(-dC_dx);

            // evaluate and scatter point contributions to the global residual
            local->seed_wrt_x(dxi_dx);

          }

          else {

            global->seed_wrt_x();
            global->interpolate(iota);

          }

          global->zero_residual();
          global->evaluate(local, iota, w, dv, ip_set);
          EMatrix const dtotal = global->eigen_jacobian();
          EMatrix const elem_resid = global->eigen_residual();
          global->scatter_lhs(disc, dtotal, LHS);
          global->scatter_rhs(disc, elem_resid, RHS);
          global->unseed_wrt_x();

        }

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();

}

template<typename T>
void preprocess_qoi(RCP<QoI<T>> qoi,
    RCP<LocalResidual<T>> local,
    RCP<GlobalResidual<T>> global,
    RCP<State> state,
    RCP<Disc> disc,
    int step) {

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();
  int const q_order = disc->lv_shape()->getOrder();

  // gather information from the state object
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  RCP<NestedDisc> nested;
  if (disc->type() == VERIFICATION) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    x = nested->primal_fine(step).global;
    xi = nested->primal_fine(step).local;
    x_prev = nested->primal_fine(step - 1).global;
    xi_prev = nested->primal_fine(step - 1).local;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);
  qoi->before_elems(disc, step);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      qoi->set_elem(me);
      global->gather(x, x_prev);

      // loop over all integration points in the current element
      int const npts = apf::countIntPoints(me, q_order);
      for (int pt = 0; pt < npts; ++pt) {

        // get integration point specific information
        apf::Vector3 iota;
        apf::getIntPoint(me, q_order, pt, iota);
        double const w = apf::getIntWeight(me, q_order, pt);
        double const dv = apf::getDV(me, iota);

        // preprocess the quantities needed for QoI evaluation
        global->interpolate(iota);
        local->gather(pt, xi, xi_prev);
        qoi->preprocess(es, elem, global, local, iota, w, dv);

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();
      qoi->unset_elem();

    }

  }

  qoi->preprocess_finalize(step);

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();
  qoi->after_elems();

}

void eval_adjoint_jacobian(
    RCP<State> state,
    RCP<Disc> disc,
    Array3D<EVector>& g,
    Array3D<EVector>& f,
    int step) {

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();

  // preprocess the QoI
  RCP<LocalResidual<FADT>> local = state->d_residuals->local;
  RCP<GlobalResidual<FADT>> global = state->d_residuals->global;
  RCP<QoI<FADT>> qoi = state->d_qoi;
  preprocess_qoi(qoi, local, global, state, disc, step);

  // gather information from the state object
  Array2D<RCP<MatrixT>>& LHS = state->la->A[GHOST];
  Array1D<RCP<VectorT>>& RHS = state->la->b[GHOST];
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  // determine if we are doing verification
  RCP<NestedDisc> nested;
  bool force_path = false;
  int path = 0;
  if (disc->type() == VERIFICATION) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    force_path = true;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);
  qoi->before_elems(disc, step);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      qoi->set_elem(me);
      global->gather(x, x_prev);

      // grab the forced path if required
      if (force_path) {
        path = nested->branch_paths()[step][es][elem];
      }

      // loop over domain ip sets
      // ip_set = 0 -> coupled
      // ip_set > 0 -> global only
      Array1D<int> ip_sets = global->ip_sets();
      int const num_ip_sets = ip_sets.size();

      for (int ip_set = 0; ip_set < num_ip_sets; ++ip_set) {

        // get the quadrature order for the ip set
        int const q_order = ip_sets[ip_set];
        // loop over all integration points in the current element
        int const npts = apf::countIntPoints(me, q_order);

        for (int pt = 0; pt < npts; ++pt) {

          // get integration point specific information
          apf::Vector3 iota;
          apf::getIntPoint(me, q_order, pt, iota);
          double const w = apf::getIntWeight(me, q_order, pt);
          double const dv = apf::getDV(me, iota);

          if (ip_set == 0) {

            // solve the local constitutive equations at the integration point
            // and store the resultant local residual and its derivatives (dC_dxi)
            global->interpolate(iota);
            local->gather(pt, xi, xi_prev);
            local->seed_wrt_xi();
            local->evaluate(global, force_path, path);
            EMatrix const dC_dxi = local->eigen_jacobian();

            // re-evaluate the constitutive equations to obtain dC_dx
            local->unseed_wrt_xi();
            global->seed_wrt_x();
            global->interpolate(iota);
            local->evaluate(global, force_path, path);
            EMatrix const dC_dx = local->eigen_jacobian();

            // solve the forward sensitivty system to obtain dxi_dx
            EMatrix const dxi_dx = dC_dxi.fullPivLu().solve(-dC_dx);

            // evaluate and scatter point contributions to the global LHS
            local->seed_wrt_x(dxi_dx);

            global->zero_residual();
            global->evaluate(local, iota, w, dv, ip_set);
            EMatrix const dtotal = global->eigen_jacobian();
            EMatrix const dtotalT = dtotal.transpose();
            global->scatter_lhs(disc, dtotalT, LHS);
            local->unseed_wrt_xi();

            // evaluate the QoI derivatives to obtain dJ_dx
            qoi->evaluate(es, elem, global, local, iota, w, dv);
            EVector const dJ_dx = qoi->eigen_dvector();
            global->unseed_wrt_x();

            // evaluate the QoI derivatives to obtain dJ_dxi
            local->seed_wrt_xi();
            global->interpolate(iota);
            qoi->evaluate(es, elem, global, local, iota, w, dv);
            EVector const dJ_dxi = qoi->eigen_dvector();
            local->unseed_wrt_xi();

            // update the local history variable
            g[es][elem][pt] -= dJ_dxi;
            EVector const g_pt = g[es][elem][pt];
            EVector const f_pt = f[es][elem][pt];

            // evaluate and scatter point contributions to the global RHS
            EMatrix const dxi_dxT = dxi_dx.transpose();
            EVector const rhs = -dJ_dx + f_pt + dxi_dxT * g_pt;
            global->scatter_rhs(disc, rhs, RHS);

          }

          else {

            global->seed_wrt_x();
            global->interpolate(iota);
            global->zero_residual();
            global->evaluate(local, iota, w, dv, ip_set);
            EMatrix const dtotal = global->eigen_jacobian();
            EMatrix const dtotalT = dtotal.transpose();
            global->scatter_lhs(disc, dtotalT, LHS);

          }

        }

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();
      qoi->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();
  qoi->after_elems();

}

void solve_adjoint_local(
    RCP<State> state,
    RCP<Disc> disc,
    Array3D<EVector>& g,
    Array3D<EVector>& f,
    int step) {

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();
  int const q_order = disc->lv_shape()->getOrder();

  // gather information from the state object
  RCP<LocalResidual<FADT>> local = state->d_residuals->local;
  RCP<GlobalResidual<FADT>> global = state->d_residuals->global;
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;
  Array1D<apf::Field*> z = disc->adjoint(step).global;
  Array1D<apf::Field*> phi = disc->adjoint(step).local;

  // determine if we are doing verification
  RCP<NestedDisc> nested;
  bool force_path = false;
  int path = 0;
  if (disc->type() == VERIFICATION) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    force_path = true;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      global->gather(x, x_prev);

      // grab the forced path if required
      if (force_path) {
        path = nested->branch_paths()[step][es][elem];
      }

      // grab the adjoint nodal solution at the element
      EVector const z_nodes = global->gather_adjoint(z);

      // loop over all integration points in the current element
      int const npts = apf::countIntPoints(me, q_order);

      for (int pt = 0; pt < npts; ++pt) {

        // get integration point specific information
        apf::Vector3 iota;
        apf::getIntPoint(me, q_order, pt, iota);
        double const w = apf::getIntWeight(me, q_order, pt);
        double const dv = apf::getDV(me, iota);

        // evaluate local/global residuals and their derivatives, store
        // their transpose Jacobians, and grab g at the integration point
        global->interpolate(iota);
        local->gather(pt, xi, xi_prev);
        local->seed_wrt_xi();
        global->zero_residual();
        global->evaluate(local, iota, w, dv, 0);
        local->evaluate(global, force_path, path);
        EMatrix const dC_dxiT = local->eigen_jacobian().transpose();
        EMatrix const dR_dxiT = global->eigen_jacobian().transpose();
        EVector const g_pt = g[es][elem][pt];

        // Solve for the local adjoint variables and scatter them into fields
        EVector const phi_pt = dC_dxiT.fullPivLu().solve(g_pt - dR_dxiT * z_nodes);
        local->scatter_adjoint(pt, phi_pt, phi);

        // Solve for the global history vector
        local->unseed_wrt_xi();
        global->seed_wrt_x_prev();
        global->interpolate(iota);
        local->evaluate(global, force_path, path);
        EMatrix const dC_dx_prevT = local->eigen_jacobian().transpose();
        f[es][elem][pt] = -dC_dx_prevT * phi_pt;

        // Solve for the local history vector
        global->unseed_wrt_x_prev();
        global->interpolate(iota);
        local->seed_wrt_xi_prev();
        local->evaluate(global, force_path, path);
        EMatrix const dC_dxi_prevT = local->eigen_jacobian().transpose();
        g[es][elem][pt] = -dC_dxi_prevT * phi_pt;
        local->unseed_wrt_xi_prev();

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();

}

double eval_qoi(RCP<State> state, RCP<Disc> disc, int step) {

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();
  int const q_order = disc->lv_shape()->getOrder();

  // preprocess the QoI
  RCP<LocalResidual<double>> local = state->residuals->local;
  RCP<GlobalResidual<double>> global = state->residuals->global;
  RCP<QoI<double>> qoi = state->qoi;
  preprocess_qoi(qoi, local, global, state, disc, step);

  // gather information from the state object
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  RCP<NestedDisc> nested;
  if (disc->type() == VERIFICATION) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    x = nested->primal_fine(step).global;
    xi = nested->primal_fine(step).local;
    x_prev = nested->primal_fine(step - 1).global;
    xi_prev = nested->primal_fine(step - 1).local;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);
  qoi->before_elems(disc, step);

  // initialize the QoI value at the step
  double J = 0.;

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      qoi->set_elem(me);
      global->gather(x, x_prev);

      // loop over all integration points in the current element
      int const npts = apf::countIntPoints(me, q_order);
      for (int pt = 0; pt < npts; ++pt) {

        // get integration point specific information
        apf::Vector3 iota;
        apf::getIntPoint(me, q_order, pt, iota);
        double const w = apf::getIntWeight(me, q_order, pt);
        double const dv = apf::getDV(me, iota);

        // solve the local constitutive equations at the integration point
        // and store the resultant local residual and its derivatives (dC_dxi)
        global->interpolate(iota);
        local->gather(pt, xi, xi_prev);
        qoi->evaluate(es, elem, global, local, iota, w, dv);
        qoi->scatter(J);

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();
      qoi->unset_elem();

    }

  }

  qoi->postprocess(J);

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();
  qoi->after_elems();

  return J;

}

Array1D<double> eval_qoi_gradient(RCP<State> state, int step) {

  int const num_active_params = state->residuals->local->num_active_params();
  Array1D<double> grad(num_active_params);
  EVector Egrad = EVector::Zero(num_active_params);

  // gather discretization information
  RCP<Disc> disc = state->disc;
  apf::Mesh* mesh = disc->apf_mesh();

  // preprocess the QoI
  RCP<LocalResidual<FADT>> local = state->d_residuals->local;
  RCP<GlobalResidual<FADT>> global = state->d_residuals->global;
  RCP<QoI<FADT>> qoi = state->d_qoi;
  preprocess_qoi(qoi, local, global, state, disc, step);

  // gather information from the state object
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;
  Array1D<apf::Field*> z = disc->adjoint(step).global;
  Array1D<apf::Field*> phi = disc->adjoint(step).local;

  // perform initializations of the residual objects
  global->before_elems(disc);
  qoi->before_elems(disc, step);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      qoi->set_elem(me);
      global->gather(x, x_prev);

      // grab the adjoint nodal solution at the element
      EVector const z_nodes = global->gather_adjoint(z);

      // loop over domain ip sets
      // ip_set = 0 -> coupled
      // ip_set > 0 -> global only
      Array1D<int> ip_sets = global->ip_sets();
      int const num_ip_sets = ip_sets.size();

      for (int ip_set = 0; ip_set < num_ip_sets; ++ip_set) {

        // get the quadrature order for the ip set
        int const q_order = ip_sets[ip_set];
        // loop over all integration points in the current element
        int const npts = apf::countIntPoints(me, q_order);

        for (int pt = 0; pt < npts; ++pt) {
          // get integration point specific information
          apf::Vector3 iota;
          apf::getIntPoint(me, q_order, pt, iota);
          double const w = apf::getIntWeight(me, q_order, pt);
          double const dv = apf::getDV(me, iota);

          // evaluate local/global residuals and their derivatives
          // and dot with the corresponding adjoint solutions to
          // compute gradient contributions
          global->interpolate(iota);
          global->zero_residual();
          local->seed_wrt_params(es);

          if (ip_set == 0) {
            local->gather(pt, xi, xi_prev);
            local->evaluate(global);
            EMatrix const dC_dpT = local->eigen_jacobian().transpose();
            EVector const phi_pt = local->gather_adjoint(pt, phi);
            Egrad += dC_dpT * phi_pt;

            // evaluate the QoI derivatives to obtain dJ_dp
            qoi->evaluate(es, elem, global, local, iota, w, dv);
            EVector const dJ_dp = qoi->eigen_dvector();
            Egrad += dJ_dp;

          }

          global->evaluate(local, iota, w, dv, ip_set);
          EMatrix const dR_dpT = global->eigen_jacobian().transpose();
          Egrad += dR_dpT * z_nodes;
          local->unseed_wrt_params(es);

        }

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();
      qoi->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();
  qoi->after_elems();

  EVector::Map(&grad[0], num_active_params) = Egrad;

  return grad;
}

// TODO: write this using global->interpolate_error
//void eval_error_contributions_nodal();

void eval_error_contributions(
    RCP<State> state,
    RCP<Disc> disc,
    apf::Field* R_error_field,
    apf::Field* C_error_field,
    int step) {

  // gather the residuals from the state object
  RCP<LocalResidual<double>> local = state->residuals->local;
  RCP<GlobalResidual<double>> global = state->residuals->global;
  Array1D<RCP<VectorT>>& resid_vec = state->la->b[GHOST];
  Array1D<RCP<VectorT>>& z_vec = state->la->x[GHOST];

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();

  // gather the prolonged forward state variables
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  // gather the enriched adjoint state variables
  Array1D<apf::Field*> z = disc->adjoint(step).global;
  Array1D<apf::Field*> phi = disc->adjoint(step).local;

  // determine if we are doing verification
  RCP<NestedDisc> nested;
  bool force_path = false;
  int path = 0;
  if (disc->type() == VERIFICATION) {
    nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
    ALWAYS_ASSERT(nested != Teuchos::null);
    force_path = true;
  }

  // perform initializations of the residual objects
  global->before_elems(disc);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      global->gather(x, x_prev);

      // grab the adjoint nodal solution at the element
      EVector const z_nodes = global->gather_adjoint(z);
      global->assign_rhs(disc, z_nodes, z_vec);

      // grab the forced path if required
      if (force_path) {
        path = nested->branch_paths()[step][es][elem];
      }

      // loop over domain ip sets
      // ip_set = 0 -> coupled
      // ip_set > 0 -> global only
      Array1D<int> ip_sets = global->ip_sets();
      int const num_ip_sets = ip_sets.size();

      for (int ip_set = 0; ip_set < num_ip_sets; ++ip_set) {

        // get the quadrature order for the ip set
        int const q_order = ip_sets[ip_set];
        // loop over all integration points in the current element
        int const npts = apf::countIntPoints(me, q_order);

        for (int pt = 0; pt < npts; ++pt) {

          // get integration point specific information
          apf::Vector3 iota;
          apf::getIntPoint(me, q_order, pt, iota);
          double const w = apf::getIntWeight(me, q_order, pt);
          double const dv = apf::getDV(me, iota);

          if (ip_set == 0) {

            // evaluate the global residual error contributions
            local->gather(pt, xi, xi_prev);
            global->zero_residual();
            global->interpolate(iota);
            global->evaluate(local, iota, w, dv, ip_set);
            EVector const R = global->eigen_residual();
            double const E_R_elem = z_nodes.dot(R);
            double E_R = apf::getScalar(R_error_field, e, 0);
            apf::setScalar(R_error_field, e, 0, E_R + E_R_elem);
            global->scatter_rhs(disc, R, resid_vec);

            // evaluate the local residual error contributions
            local->evaluate(global, force_path, path);
            EVector const C = local->eigen_residual();
            EVector const phi_pt = local->gather_adjoint(pt, phi);
            double const E_C_elem = phi_pt.dot(C);
            double E_C = apf::getScalar(C_error_field, e, 0);
            apf::setScalar(C_error_field, e, 0, E_C + E_C_elem);

          }

          else {

            // evaluate the global residual error contributions
            global->zero_residual();
            global->interpolate(iota);
            global->evaluate(local, iota, w, dv, ip_set);
            EVector const R = global->eigen_residual();
            double const E_R_elem = z_nodes.dot(R);
            double E_R = apf::getScalar(R_error_field, e, 0);
            apf::setScalar(R_error_field, e, 0, E_R + E_R_elem);
            global->scatter_rhs(disc, R, resid_vec);

          }

        }

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();

}

void eval_linearization_errors(
    RCP<State> state,
    RCP<Disc> disc,
    int step,
    double& E_lin_R,
    double& E_lin_C) {

  // we must be a verification mesh to do this evaluation
  ALWAYS_ASSERT(disc->type() == VERIFICATION);
  RCP<NestedDisc> nested = Teuchos::rcp_static_cast<NestedDisc>(disc);
  bool force_path = true;

  // gather the residuals from the state object
  RCP<LocalResidual<FADT>> local = state->d_residuals->local;
  RCP<GlobalResidual<FADT>> global = state->d_residuals->global;

  // gather discretization information
  apf::Mesh* mesh = disc->apf_mesh();

  // gather the prolonged forward state variables
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  // gather the enriched forward state variables
  Array1D<apf::Field*> x_fine = nested->primal_fine(step).global;
  Array1D<apf::Field*> xi_fine = nested->primal_fine(step).local;
  Array1D<apf::Field*> x_prev_fine = nested->primal_fine(step - 1).global;
  Array1D<apf::Field*> xi_prev_fine = nested->primal_fine(step - 1).local;

  // gather the enriched adjoint state variables
  Array1D<apf::Field*> z = disc->adjoint(step).global;
  Array1D<apf::Field*> phi = disc->adjoint(step).local;

  // perform initializations of the residual objects
  global->before_elems(disc);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      global->gather(x, x_prev);

      // grab some nodal solution information at the element
      EVector const z_nodes = global->gather_adjoint(z);
      EVector const x_diff = global->gather_difference(x_fine, x);
      EVector const x_prev_diff =
          global->gather_difference(x_prev_fine, x_prev);

      // initialize the element level global linearization error
      EVector ELR_e = EVector::Zero(x_diff.size());

      // grab the forced path
      int const path = nested->branch_paths()[step][es][elem];

      // loop over domain ip sets
      // ip_set = 0 -> coupled
      // ip_set > 0 -> global only
      Array1D<int> ip_sets = global->ip_sets();
      int const num_ip_sets = ip_sets.size();

      for (int ip_set = 0; ip_set < num_ip_sets; ++ip_set) {

        // get the quadrature order for the ip set
        int const q_order = ip_sets[ip_set];
        // loop over all integration points in the current element
        int const npts = apf::countIntPoints(me, q_order);

        for (int pt = 0; pt < npts; ++pt) {

          // get integration point specific information
          apf::Vector3 iota;
          apf::getIntPoint(me, q_order, pt, iota);
          double const w = apf::getIntWeight(me, q_order, pt);
          double const dv = apf::getDV(me, iota);

          if (ip_set == 0) {

            // grab local state variable data at the point
            EVector const phi_pt = local->gather_adjoint(pt, phi);
            EVector const xi_diff = local->gather_difference(pt, xi_fine, xi);
            EVector const xi_prev_diff =
                local->gather_difference(pt, xi_prev_fine, xi_prev);

            // evaluate derivatives wrt x
            global->zero_residual();
            global->seed_wrt_x();
            global->interpolate(iota);
            local->gather(pt, xi, xi_prev);
            global->evaluate(local, iota, w, dv, ip_set);
            local->evaluate(global, force_path, path);
            EVector const R = global->eigen_residual();
            EVector const C = local->eigen_residual();
            EMatrix const dR_dx = global->eigen_jacobian();
            EMatrix const dC_dx = local->eigen_jacobian();

            // evaluate derivatives wrt xi
            global->unseed_wrt_x();
            global->zero_residual();
            local->seed_wrt_xi();
            global->interpolate(iota);
            global->evaluate(local, iota, w, dv, ip_set);
            local->evaluate(global, force_path, path);
            EMatrix const dR_dxi = global->eigen_jacobian();
            EMatrix const dC_dxi = local->eigen_jacobian();

            // evaluate derivatives wrt x_prev
            local->unseed_wrt_xi();
            global->seed_wrt_x_prev();
            global->interpolate(iota);
            local->evaluate(global, force_path, path);
            EMatrix const dC_dx_prev = local->eigen_jacobian();

            // evaluate derivatives wrt xi_prev
            global->unseed_wrt_x_prev();
            global->interpolate(iota);
            local->seed_wrt_xi_prev();
            local->evaluate(global, force_path, path);
            EMatrix const dC_dxi_prev = local->eigen_jacobian();

            // evaluate the point level local linearization error
            EVector const ELC_e =
              -C - (dC_dx * x_diff) - (dC_dxi * xi_diff) -
              (dC_dx_prev * x_prev_diff) - (dC_dxi_prev * xi_prev_diff);
            E_lin_C += phi_pt.dot(ELC_e);

            // evaluate point contribs to the global linearization error
            EVector const ELR_e = -R - (dR_dx * x_diff) - (dR_dxi * xi_diff);
            E_lin_R += z_nodes.dot(ELR_e);

            // unseed on output
            local->unseed_wrt_xi_prev();

          }

          else {

            // evaluate the global residual linearization error contributions
            global->zero_residual();
            global->seed_wrt_x();
            global->interpolate(iota);
            global->evaluate(local, iota, w, dv, ip_set);
            EVector const R = global->eigen_residual();
            EMatrix const dR_dx = global->eigen_jacobian();

            // evaluate point contribs to the global linearization error
            EVector const ELR_e = -R - (dR_dx * x_diff);
            E_lin_R += z_nodes.dot(ELR_e);

            // unseed on output
            global->unseed_wrt_x();

          }

        }

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();

    }

  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();

}

// TODO: Cauchy stress can have contributions from fields
// with different polynomials orders (e.g. constant + linear).
// Generalize to handle such cases.
apf::Field* eval_cauchy(RCP<State> state, int step) {

  // an assumption about the pressure index
  int const pressure_idx = 1;

  // gather discretization information
  RCP<Disc> disc = state->disc;
  apf::Mesh* mesh = disc->apf_mesh();
  int const ndims = mesh->getDimension();
  int const q_order = disc->lv_shape()->getOrder();

  // gather information from the state object
  RCP<LocalResidual<double>> local = state->residuals->local;
  RCP<GlobalResidual<double>> global = state->residuals->global;
  Array1D<apf::Field*> x = disc->primal(step).global;
  Array1D<apf::Field*> xi = disc->primal(step).local;
  Array1D<apf::Field*> x_prev = disc->primal(step - 1).global;
  Array1D<apf::Field*> xi_prev = disc->primal(step - 1).local;

  // create the field to fill in
  apf::FieldShape* shape = state->disc->lv_shape();
  apf::Field* field = apf::createField(mesh, "sigma", apf::MATRIX, shape);
  apf::zeroField(field);
  if (step == 0) return field;

  // perform initializations of the residual objects
  global->before_elems(disc);

  // loop over all element sets in the discretization
  for (int es = 0; es < disc->num_elem_sets(); ++es) {

    local->before_elems(es, disc);

    // gather the elements in the current element set
    std::string const& es_name = disc->elem_set_name(es);
    ElemSet const& elems = disc->elems(es_name);

    // loop over all elements in the element set
    for (size_t elem = 0; elem < elems.size(); ++elem) {

      // get the current mesh element
      apf::MeshEntity* e = elems[elem];
      apf::MeshElement* me = apf::createMeshElement(mesh, e);

      // peform operations on element input
      global->set_elem(me);
      local->set_elem(me);
      global->gather(x, x_prev);

      // loop over all integration points in the current element
      int const npts = apf::countIntPoints(me, q_order);
      for (int pt = 0; pt < npts; ++pt) {

        // get integration point specific information
        apf::Vector3 iota;
        apf::getIntPoint(me, q_order, pt, iota);

        // evaluate the cauchy stress tensor at the point
        global->interpolate(iota);
        local->gather(pt, xi, xi_prev);
        double const p = global->scalar_x(pressure_idx);
        Tensor<double> const sigma = local->cauchy(global, p);

        // set the cauchy stress tensor to a field
        apf::Matrix3x3 apf_sigma(0, 0, 0, 0, 0, 0, 0, 0, 0);
        for (int i = 0; i < ndims; ++i) {
          for (int j = 0; j < ndims; ++j) {
            apf_sigma[i][j] = sigma(i, j);
          }
        }
        apf::setMatrix(field, e, pt, apf_sigma);

      }

      // perform operations on element output
      apf::destroyMeshElement(me);
      global->unset_elem();
      local->unset_elem();

    }
  }

  // perform clean-ups of the residual objects
  local->after_elems();
  global->after_elems();

  return field;

}

}


/* pcmsolver_copyright_start */
/*
 *     PCMSolver, an API for the Polarizable Continuum Model
 *     Copyright (C) 2013-2016 Roberto Di Remigio, Luca Frediani and contributors
 *
 *     This file is part of PCMSolver.
 *
 *     PCMSolver is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU Lesser General Public License as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 *
 *     PCMSolver is distributed in the hope that it will be useful,
 *     but WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *     GNU Lesser General Public License for more details.
 *
 *     You should have received a copy of the GNU Lesser General Public License
 *     along with PCMSolver.  If not, see <http://www.gnu.org/licenses/>.
 *
 *     For information on the complete list of contributors to the
 *     PCMSolver API, see: <http://pcmsolver.readthedocs.io/>
 */
/* pcmsolver_copyright_end */

#ifndef NUMERICALINTEGRATOR_HPP
#define NUMERICALINTEGRATOR_HPP

#include <cmath>
#include <iosfwd>
#include <vector>

#include "Config.hpp"

#include <Eigen/Core>

#include "IntegratorHelperFunctions.hpp"
#include "cavity/Element.hpp"
#include "green/AnisotropicLiquid.hpp"
#include "green/IonicLiquid.hpp"
#include "green/SphericalDiffuse.hpp"
#include "green/UniformDielectric.hpp"
#include "green/Vacuum.hpp"

/*! \file NumericalIntegrator.hpp
 *  \struct NumericalIntegrator
 *  \brief Implementation of the single and double layer operators matrix representation using one-point collocation
 *  \author Roberto Di Remigio
 *  \date 2015
 *
 *  Calculates the diagonal elements of S and D by collocation, using numerical integration.
 */

struct NumericalIntegrator
{
    NumericalIntegrator() : factor(1.07) {}
    NumericalIntegrator(double f) : factor(f) {}
    ~NumericalIntegrator() {}
    /**@{ Single and double layer potentials for a Vacuum Green's function by collocation: numerical integration of diagonal */
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd singleLayer(const Vacuum<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelS kernelS = gf.exportKernelS();
        integrator::Diagonal diagS = pcm::bind(&integrator::integrateS<32, 16>, kernelS, pcm::_1);
        return integrator::singleLayer(e, diagS, kernelS);
    }
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd doubleLayer(const Vacuum<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelD kernelD = gf.exportKernelD();
        integrator::Diagonal diagD = pcm::bind(&integrator::integrateD<32, 16>, kernelD, pcm::_1);
        return integrator::doubleLayer(e, diagD, kernelD);
    }
    /**@}*/

    /**@{ Single and double layer potentials for a UniformDielectric Green's function by collocation: numerical integration of diagonal */
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd singleLayer(const UniformDielectric<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelS kernelS = gf.exportKernelS();
        integrator::Diagonal diagS = pcm::bind(&integrator::integrateS<32, 16>, kernelS, pcm::_1);
        return integrator::singleLayer(e, diagS, kernelS);
    }
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd doubleLayer(const UniformDielectric<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelD kernelD = gf.exportKernelD();
        integrator::Diagonal diagD = pcm::bind(&integrator::integrateD<32, 16>, kernelD, pcm::_1);
        return integrator::doubleLayer(e, diagD, kernelD);
    }
    /**@}*/

    /**@{ Single and double layer potentials for a IonicLiquid Green's function by collocation: numerical integration of diagonal */
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd singleLayer(const IonicLiquid<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelS kernelS = gf.exportKernelS();
        integrator::Diagonal diagS = pcm::bind(&integrator::integrateS<32, 16>, kernelS, pcm::_1);
        return integrator::singleLayer(e, diagS, kernelS);
    }
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd doubleLayer(const IonicLiquid<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelD kernelD = gf.exportKernelD();
        integrator::Diagonal diagD = pcm::bind(&integrator::integrateD<32, 16>, kernelD, pcm::_1);
        return integrator::doubleLayer(e, diagD, kernelD);
    }
    /**@}*/

    /**@{ Single and double layer potentials for a AnisotropicLiquid Green's function by collocation: numerical integration of diagonal */
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd singleLayer(const AnisotropicLiquid<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelS kernelS = gf.exportKernelS();
        integrator::Diagonal diagS = pcm::bind(&integrator::integrateS<32, 16>, kernelS, pcm::_1);
        return integrator::singleLayer(e, diagS, kernelS);
    }
    /*! \tparam DerivativeTraits how the derivatives of the Greens's function are calculated
     *  \param[in] gf Green's function
     *  \param[in] e  list of finite elements
     */
    template <typename DerivativeTraits>
    Eigen::MatrixXd doubleLayer(const AnisotropicLiquid<DerivativeTraits, NumericalIntegrator> & gf, const std::vector<Element> & e) const {
        integrator::KernelD kernelD = gf.exportKernelD();
        integrator::Diagonal diagD = pcm::bind(&integrator::integrateD<32, 16>, kernelD, pcm::_1);
        return integrator::doubleLayer(e, diagD, kernelD);
    }
    /**@}*/

    /**@{ Single and double layer potentials for a SphericalDiffuse Green's function by collocation: numerical integration of diagonal */
    template <typename ProfilePolicy>
    Eigen::MatrixXd singleLayer(const SphericalDiffuse<NumericalIntegrator, ProfilePolicy> & /* gf */, const std::vector<Element> & e) const {
//      // The singular part is "integrated" as usual, while the nonsingular part is evaluated in full
//      double area = e.area();
//      // Diagonal of S inside the cavity
//      double Sii_I = factor_ * std::sqrt(4 * M_PI / area);
//      // "Diagonal" of Coulomb singularity separation coefficient
//      double coulomb_coeff = gf.coefficientCoulomb(e.center(), e.center());
//      // "Diagonal" of the image Green's function
//      double image = gf.imagePotential(e.center(), e.center());

//      return (Sii_I / coulomb_coeff + image);
        return Eigen::MatrixXd::Zero(e.size(), e.size());
    }
    template <typename ProfilePolicy>
    Eigen::MatrixXd doubleLayer(const SphericalDiffuse<NumericalIntegrator, ProfilePolicy> & /* gf */, const std::vector<Element> & e) const {
//      // The singular part is "integrated" as usual, while the nonsingular part is evaluated in full
//      double area = e.area();
//      double radius = e.sphere().radius();
//      // Diagonal of S inside the cavity
//      double Sii_I = factor_ * std::sqrt(4 * M_PI / area);
//      // Diagonal of D inside the cavity
//      double Dii_I = -factor_ * std::sqrt(M_PI/ area) * (1.0 / radius);
//      // "Diagonal" of Coulomb singularity separation coefficient
//      double coulomb_coeff = gf.coefficientCoulomb(e.center(), e.center());
//      // "Diagonal" of the directional derivative of the Coulomb singularity separation coefficient
//      double coeff_grad = gf.coefficientCoulombDerivative(e.normal(), e.center(), e.center()) / std::pow(coulomb_coeff, 2);
//      // "Diagonal" of the directional derivative of the image Green's function
//      double image_grad = gf.imagePotentialDerivative(e.normal(), e.center(), e.center());

//      double eps_r2 = 0.0;
//      pcm::tie(eps_r2, pcm::ignore) = gf.epsilon(e.center());

//      return eps_r2 * (Dii_I / coulomb_coeff - Sii_I * coeff_grad + image_grad);
        return Eigen::MatrixXd::Zero(e.size(), e.size());
    }
    /**@}*/
    /// Scaling factor for the collocation formulas (unused here)
    double factor;
};

#endif // NUMERICALINTEGRATOR_HPP


#ifndef TEST_ANASAZI_HPP
#define TEST_ANASAZI_HPP

#include <vector>
#include <set>
#include <sstream>

#include <boost/random.hpp>

#include "Teuchos_LAPACK.hpp"
#include "Teuchos_SerialDenseMatrix.hpp"
#include "Teuchos_SerialDenseVector.hpp"

#include "AnasaziConfigDefs.hpp"
#include "AnasaziBasicEigenproblem.hpp"
#include "AnasaziSimpleLOBPCGSolMgr.hpp"
#include "AnasaziBasicOutputManager.hpp"
#include "AnasaziEpetraAdapter.hpp"
#include "Epetra_CrsMatrix.h"
#include "Epetra_Operator.h"
#include "EpetraExt_VectorOut.h"
#include "EpetraExt_MultiVectorOut.h"
#include "EpetraExt_RowMatrixOut.h"
#include "EpetraExt_MatrixMatrix.h"
#include "Teuchos_CommandLineProcessor.hpp"

#include "AnasaziBlockKrylovSchurSolMgr.hpp"
#include "Epetra_LinearProblem.h"
#include "Epetra_InvOperator.h"

// Include header for AztecOO solver and solver interface for Epetra_Operator
#include "AztecOO.h"
#include "AztecOO_Operator.h"

// Include header for Belos solver and solver interface for Epetra_Operator
#include "BelosEpetraOperator.h"
#include "BelosEpetraAdapter.hpp"

// Ifpack
#include "Ifpack.h"
#include "Ifpack_AdditiveSchwarz.h"

// Pliris
#include "Pliris.h"

#ifdef HAVE_MPI
#include "Epetra_MpiComm.h"
#include <mpi.h>
#else
#include "Epetra_SerialComm.h"
#endif
#include "Epetra_Map.h"

// 2*PI
static double const TWOPI = 6.28318530717958647693;

class test_params {

  public :

    test_params ( );
    ~test_params ( ) { }

    int np;
    int myp;

    int n_bundle;
    int bundle_size;

    int n_trace;
    int n_flux_trace;
    int n_flux;

    int n_pix_margin;
    int n_pix_gap;

    int n_pix_bundle;
    int n_pix_lambda;
    int n_pix_trace;
    int n_pix;

    double peak_amp;
    double peak_obj;

    int peak_space;

    double background;

    double psf_fwhm;

    int psf_corr;

};

class gls_operator : public Epetra_Operator {

  public :

    gls_operator ( Teuchos::RCP < Epetra_CrsMatrix > AT, Teuchos::RCP < Epetra_CrsMatrix > invN ) {
      AT_ = AT;
      invN_ = invN;
      strcpy ( label_, "GLS" );
    }

    ~gls_operator () { }

    gls_operator ( gls_operator const & orig ) {
      AT_ = orig.AT_;
      invN_ = orig.invN_;
      strcpy ( label_, orig.label_ );
    }

    int SetUseTranspose ( bool UseTranspose ) { return -1; }

    int Apply ( const Epetra_MultiVector& X, Epetra_MultiVector& Y ) const {
      Teuchos::RCP < Epetra_MultiVector > temp1 = Teuchos::rcp( new Epetra_MultiVector ( AT_->DomainMap(), X.NumVectors() ) );
      Teuchos::RCP < Epetra_MultiVector > temp2 = Teuchos::rcp( new Epetra_MultiVector ( AT_->DomainMap(), X.NumVectors() ) );
      int info = AT_->Multiply ( true, X, (*temp1) );
      info = invN_->Multiply ( false, (*temp1), (*temp2) );
      info = AT_->Multiply ( false, (*temp2), Y );
      return info;
    }

    int ApplyInverse ( const Epetra_MultiVector& X, Epetra_MultiVector& Y ) const {
      return -1;
    }

    double NormInf() const {
      return 0.0;
    }

    const char * Label() const {
      return label_;
    }

    bool UseTranspose() const {
      return false;
    }

    bool HasNormInf() const {
      return false;
    }

    const Epetra_Comm & Comm() const {
      return AT_->Comm();
    }

    //! Returns the Epetra_Map object associated with the domain of this operator.
    const Epetra_Map & OperatorDomainMap() const {
      return AT_->RangeMap();
    }

    //! Returns the Epetra_Map object associated with the range of this operator.
    const Epetra_Map & OperatorRangeMap() const {
      return AT_->RangeMap();
    }

  private :

    Teuchos::RCP < Epetra_CrsMatrix > AT_;
    Teuchos::RCP < Epetra_CrsMatrix > invN_;
    char label_[20];

};

void block_dist ( int n, int myp, int np, std::vector < int > & my_elems );

void gauss_sample ( std::vector < double > & vals, std::vector < double > & xrel, std::vector < double > & yrel, double amp, double maj, double min, double ang );

Teuchos::RCP < Epetra_CrsMatrix > build_design_matrix ( test_params & par, Epetra_Map & fluxmap, Epetra_Map & pixmap, double & timing );

Teuchos::RCP < Epetra_Vector > simulate_signal ( test_params & par, Teuchos::RCP < Epetra_CrsMatrix > AT, double & timing );

Teuchos::RCP < Epetra_CrsMatrix > noise_covariance ( test_params & par, Epetra_Map const & pixmap, Epetra_Vector & signal, Epetra_Vector & noise, double & timing );

void eigenvals_lobpcg ( test_params & par, Teuchos::RCP < gls_operator > mat, std::vector < double > & vals, double & timing );

void eigenvals_lapack ( test_params & par, Teuchos::RCP < Epetra_CrsMatrix > mat, std::vector < double > & vals, double & timing );

int eigenvals_full ( test_params & par, Teuchos::RCP < Epetra_CrsMatrix > mat, std::vector < double > & vals, double & timing );

void write_frame ( test_params & par, Epetra_Vector & frame, std::string const & filename );

#endif


#define BOOST_TEST_MODULE SolutionTest

#include "solution.hpp"

#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>

struct less_than_key {
    inline bool operator() (const vector<int>& vec1, const vector<int>& vec2) {
        if (vec1.size() == vec2.size()) {
            for (int i = 0; i < vec1.size(); i++) {
                if (vec1[i] < vec2[i]) return true;
                else if (vec1[i] > vec2[i]) return false;
            }
        } else return false;
        return false;
    }
};

BOOST_AUTO_TEST_SUITE(SolutionSuite)

BOOST_AUTO_TEST_CASE(PlainTest1)
{
    vector<int> nums{1,2,3};
    vector<vector<int>> result = Solution().permute(nums);

    vector<vector<int>> expected{
        {1,2,3},
        {1,3,2},
        {2,1,3},
        {2,3,1},
        {3,1,2},
        {3,2,1}
    };
    BOOST_CHECK_EQUAL(result.size(), expected.size());
    sort(result.begin(), result.end(), less_than_key());
    sort(expected.begin(), expected.end(), less_than_key());
    for (int i = 0; i < result.size(); i++)
        BOOST_CHECK_EQUAL_COLLECTIONS(result[i].begin(), result[i].end(), expected[i].begin(), expected[i].end());
}

BOOST_AUTO_TEST_SUITE_END()


////////////////////////////////////////////////////////////////////////////////
//  Copyright (c) 2015 Bryce Adelstein Lelbach aka wash
//
//  Distributed under the Boost Software License, Version 1.0. (See accompanying
//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
////////////////////////////////////////////////////////////////////////////////

#include <mdspan>

#include <vector>
#include <tuple>

#include <boost/core/lightweight_test.hpp>

using std::vector;
using std::tuple;

using std::experimental::dyn;
using std::experimental::dimensions;
using std::experimental::layout_mapping_left;

//#warning Fix all tests to use size() to determine buffer allocation size instead of computing it by hand.

// FIXME FIXME FIXME
//#warning Test stride == 0.
// FIXME FIXME FIXME

template <std::size_t N, std::size_t X>
void test_1d_static()
{ // {{{
    static_assert(0 == (X % N), "X must be divisable by N");

    layout_mapping_left<
        dimensions<X>, dimensions<1>, dimensions<0>
    > const l{};

    BOOST_TEST_EQ((l.is_regular()), true);

    BOOST_TEST_EQ((l.stride(0)), 1);

    BOOST_TEST_EQ((l.size()), X);
    BOOST_TEST_EQ((l.span()), X);

    layout_mapping_left<
        dimensions<X / N>, dimensions<N>, dimensions<0>
    > const sub_l{};

    BOOST_TEST_EQ((sub_l.is_regular()), true);

    BOOST_TEST_EQ((sub_l.stride(0)), N);

    BOOST_TEST_EQ((sub_l.size()), X / N);
    BOOST_TEST_EQ((sub_l.span()), X);

    int dptr[X];

    // Set all elements to 42.
    for (auto i = 0; i < l[0]; ++i)
    {
        BOOST_TEST_EQ((l.index(i)), i);

        BOOST_TEST_EQ(&(dptr[l.index(i)]), &(dptr[i]));

        dptr[l.index(i)] = 42;

        BOOST_TEST_EQ((dptr[l.index(i)]), 42);
    }

    // Set every Nth element to 17.
    for (auto i = 0; i < sub_l[0]; ++i)
    {
        auto const s = sub_l.steps();
        auto const true_idx = (s[0] * i);

        BOOST_TEST_EQ((sub_l.index(i)), true_idx);

        BOOST_TEST_EQ(&(dptr[sub_l.index(i)]), &(dptr[true_idx]));

        dptr[sub_l.index(i)] = 17;

        BOOST_TEST_EQ((dptr[sub_l.index(i)]), 17);
    }

    // Check final structure.
    for (auto i = 0; i < l[0]; ++i)
    {
        BOOST_TEST_EQ((l.index(i)), i);

        // Element not in the strided sub-box.
        if (0 == (i % sub_l.steps()[0]))
        {
            BOOST_TEST_EQ((dptr[l.index(i)]), 17);
        }
        // Element not in the strided sub-box.
        else
        {
            BOOST_TEST_EQ((dptr[l.index(i)]), 42);
        }
    }
} // }}}

template <std::size_t N, std::size_t X>
void test_1d_dynamic()
{ // {{{
    static_assert(0 == (X % N), "X must be divisable by N");

    layout_mapping_left<
        dimensions<dyn>, dimensions<dyn>, dimensions<0>
    > const l{{X}, {1}, {}};

    BOOST_TEST_EQ((l.is_regular()), true);

    BOOST_TEST_EQ((l.stride(0)), 1);

    BOOST_TEST_EQ((l.size()), X);
    BOOST_TEST_EQ((l.span()), X);

    layout_mapping_left<
        dimensions<dyn>, dimensions<dyn>, dimensions<0>
    > const sub_l{{X / N}, {N}, {}};

    BOOST_TEST_EQ((sub_l.is_regular()), true);

    BOOST_TEST_EQ((sub_l.stride(0)), N);

    BOOST_TEST_EQ((sub_l.size()), X / N);
    BOOST_TEST_EQ((sub_l.span()), X);

    // Initialize all elements to 42.
    std::vector<int> data(l[0], 42);
    int* dptr = data.data();

    // Set every Nth element to 17.
    for (auto i = 0; i < sub_l[0]; ++i)
    {
        auto const s = sub_l.steps();
        auto const true_idx = (s[0] * i);

        BOOST_TEST_EQ((sub_l.index(i)), true_idx);

        BOOST_TEST_EQ(&(dptr[sub_l.index(i)]), &(dptr[true_idx]));

        dptr[sub_l.index(i)] = 17;

        BOOST_TEST_EQ((dptr[sub_l.index(i)]), 17);

        // Bounds-checking.
        BOOST_TEST_EQ((data.at(sub_l.index(i))), 17);
    }

    // Check final structure.
    for (auto i = 0; i < l[0]; ++i)
    {
        BOOST_TEST_EQ((l.index(i)), i);

        // Element in the strided sub-box.
        if (0 == (i % sub_l.steps()[0]))
        {
            BOOST_TEST_EQ((dptr[l.index(i)]), 17);

            // Bounds-checking.
            BOOST_TEST_EQ((data.at(l.index(i))), 17);
        }
        // Element not in the strided sub-box.
        else
        {
            BOOST_TEST_EQ((dptr[l.index(i)]), 42);

            // Bounds-checking.
            BOOST_TEST_EQ((data.at(l.index(i))), 42);
        }
    }
} // }}}

template <std::size_t N, std::size_t M, std::size_t X, std::size_t Y>
void test_2d_static()
{ // {{{
    static_assert(0 == (X % N), "X must be divisable by N");
    static_assert(0 == (Y % M), "Y must be divisable by M");

    layout_mapping_left<
        dimensions<X, Y>, dimensions<1, 1>, dimensions<0, 0>
    > const l{};

    BOOST_TEST_EQ((l.is_regular()), true);

    BOOST_TEST_EQ((l.stride(0)), 1);
    BOOST_TEST_EQ((l.stride(1)), 1);

    BOOST_TEST_EQ((l.size()), X * Y);
    BOOST_TEST_EQ((l.span()), X * Y);

    layout_mapping_left<
        dimensions<X / N, Y / M>, dimensions<N, M>, dimensions<0, 0>
    > const sub_l{};

    BOOST_TEST_EQ((sub_l.is_regular()), true);

    BOOST_TEST_EQ((sub_l.stride(0)), N);
    BOOST_TEST_EQ((sub_l.stride(1)), M);

    BOOST_TEST_EQ((sub_l.size()), (X / N) * (Y / M));
    BOOST_TEST_EQ((sub_l.span()), X * Y);

    int dptr[X * Y];

    // Set all elements to 42.
    for (auto j = 0; j < l[1]; ++j)
    for (auto i = 0; i < l[0]; ++i)
    {
        auto const true_idx = (i) + (l[0]) * (j);

        BOOST_TEST_EQ((l.index(i, j)), true_idx);

        BOOST_TEST_EQ(&(dptr[l.index(i, j)]), &(dptr[true_idx]));

        dptr[l.index(i, j)] = 42;

        BOOST_TEST_EQ((dptr[l.index(i, j)]), 42);
    }

    // Set every (Nth, Mth) element to 17.
    for (auto j = 0; j < sub_l[1]; ++j)
    for (auto i = 0; i < sub_l[0]; ++i)
    {
        auto const s = sub_l.steps();
        auto const true_idx = (s[0] * i) + (sub_l[0] * s[0]) * (s[1] * j);

        BOOST_TEST_EQ((sub_l.index(i, j)), true_idx);

        BOOST_TEST_EQ(&(dptr[sub_l.index(i, j)]), &(dptr[true_idx]));

        dptr[sub_l.index(i, j)] = 17;

        BOOST_TEST_EQ((dptr[sub_l.index(i, j)]), 17);
    }

    // Check final structure.
    for (auto j = 0; j < l[1]; ++j)
    for (auto i = 0; i < l[0]; ++i)
    {
        auto const true_idx = (i) + (l[0]) * (j);

        BOOST_TEST_EQ((l.index(i, j)), true_idx);

        // Element in the strided sub-box.
        if (  (0 == (i % sub_l.steps()[0]))
           && (0 == (j % sub_l.steps()[1]))
           )
        {
            BOOST_TEST_EQ((dptr[l.index(i, j)]), 17);
        }
        // Element not in the strided sub-box.
        else
        {
            BOOST_TEST_EQ((dptr[l.index(i, j)]), 42);
        }
    }
} // }}}

template <std::size_t N, std::size_t M, std::size_t X, std::size_t Y>
void test_2d_dynamic()
{ // {{{
    static_assert(0 == (X % N), "X must be divisable by N");
    static_assert(0 == (Y % M), "Y must be divisable by M");

    layout_mapping_left<
        dimensions<dyn, dyn>, dimensions<dyn, dyn>, dimensions<0, 0>
    > const l{{X, Y}, {1, 1}, {}};

    BOOST_TEST_EQ((l.is_regular()), true);

    BOOST_TEST_EQ((l.stride(0)), 1);
    BOOST_TEST_EQ((l.stride(1)), 1);

    BOOST_TEST_EQ((l.size()), X * Y);
    BOOST_TEST_EQ((l.span()), X * Y);

    layout_mapping_left<
        dimensions<dyn, dyn>, dimensions<dyn, dyn>, dimensions<0, 0>
    > const sub_l{{X / N, Y / M}, {N, M}, {}};

    BOOST_TEST_EQ((sub_l.is_regular()), true);

    BOOST_TEST_EQ((sub_l.stride(0)), N);
    BOOST_TEST_EQ((sub_l.stride(1)), M);

    BOOST_TEST_EQ((sub_l.size()), (X / N) * (Y / M));
    BOOST_TEST_EQ((sub_l.span()), X * Y);

    // Initialize all elements to 42.
    std::vector<int> data(l[0] * l[1], 42);
    int* dptr = data.data();

    // Set every (Nth, Mth) element to 17.
    for (auto j = 0; j < sub_l[1]; ++j)
    for (auto i = 0; i < sub_l[0]; ++i)
    {
        auto const s = sub_l.steps();
        auto const true_idx = (s[0] * i) + (sub_l[0] * s[0]) * (s[1] * j);

        BOOST_TEST_EQ((sub_l.index(i, j)), true_idx);

        BOOST_TEST_EQ(&(dptr[sub_l.index(i, j)]), &(dptr[true_idx]));

        dptr[sub_l.index(i, j)] = 17;

        BOOST_TEST_EQ((dptr[sub_l.index(i, j)]), 17);

        // Bounds-checking.
        BOOST_TEST_EQ((data.at(sub_l.index(i, j))), 17);
    }

    // Check final structure.
    for (auto j = 0; j < l[1]; ++j)
    for (auto i = 0; i < l[0]; ++i)
    {
        auto const true_idx = (i) + (l[0]) * (j);

        BOOST_TEST_EQ((l.index(i, j)), true_idx);

        // Element in the strided sub-box.
        if (  (0 == (i % sub_l.steps()[0]))
           && (0 == (j % sub_l.steps()[1]))
           )
        {
            BOOST_TEST_EQ((dptr[l.index(i, j)]), 17);

            // Bounds-checking.
            BOOST_TEST_EQ((data.at(l.index(i, j))), 17);
        }
        // Element not in the strided sub-box.
        else
        {
            BOOST_TEST_EQ((dptr[l.index(i, j)]), 42);

            // Bounds-checking.
            BOOST_TEST_EQ((data.at(l.index(i, j))), 42);
        }
    }
} // }}}

int main()
{
    ///////////////////////////////////////////////////////////////////////////
    // 1D Static

    test_1d_static<1, 1>();
    test_1d_static<1, 30>();

    test_1d_static<2,  30>();
    test_1d_static<5,  30>();
    test_1d_static<10, 30>();
    test_1d_static<30, 30>();

    ///////////////////////////////////////////////////////////////////////////
    // 1D Dynamic

    test_1d_dynamic<1, 1>();
    test_1d_dynamic<1, 30>();

    test_1d_dynamic<2,  30>();
    test_1d_dynamic<5,  30>();
    test_1d_dynamic<10, 30>();
    test_1d_dynamic<30, 30>();

    ///////////////////////////////////////////////////////////////////////////
    // 2D Static

    test_2d_static<1, 1, 1,  1 >();

    test_2d_static<1, 1, 30, 1 >();
    test_2d_static<1, 1, 1,  30>();
    test_2d_static<1, 1, 30, 30>();

    test_2d_static<2, 1, 30, 1 >();
    test_2d_static<1, 2, 1,  30>();
    test_2d_static<2, 2, 30, 30>();

    test_2d_static<5, 1, 30, 1 >();
    test_2d_static<1, 5, 1,  30>();
    test_2d_static<5, 5, 30, 30>();

    ///////////////////////////////////////////////////////////////////////////
    // 2D Dynamic

    test_2d_dynamic<1, 1, 1,  1 >();

    test_2d_dynamic<1, 1, 30, 1 >();
    test_2d_dynamic<1, 1, 1,  30>();
    test_2d_dynamic<1, 1, 30, 30>();

    test_2d_dynamic<2, 1, 30, 1 >();
    test_2d_dynamic<1, 2, 1,  30>();
    test_2d_dynamic<2, 2, 30, 30>();

    test_2d_dynamic<5, 1, 30, 1 >();
    test_2d_dynamic<1, 5, 1,  30>();
    test_2d_dynamic<5, 5, 30, 30>();

    return boost::report_errors();
}


#include <iostream>
#include "opencv2/opencv.hpp"
#include "opencv2/core/core.hpp"
#include <opencv2/imgproc/imgproc.hpp>
//#include <boost/python.hpp>
#include "pybind11/pybind11.h"
#include "pybind11/numpy.h"
#include "pybind11/stl.h"
#include "pybind11/stl_bind.h"

#include "pybind11_opencv.h"

#include "conversion.h"
#include "imgwarp_mls_rigid.h"
#include "imgwarp_mls_similarity.h"
#include "imgwarp_piecewiseaffine.h"

#include <stdlib.h>
#include <sys/timeb.h>

using namespace cv;
using namespace std;

namespace py = pybind11;

typedef unsigned char uchar_t;

py::array_t<uint8_t> cv_mat_uint8_3c_to_numpy(cv::Mat& input);

namespace MLSTransformer {

	py::array_t<uint8_t> Distort(py::array_t<uint8_t>& image, const int segment) {

		ImgWarp_MLS_Similarity trans1;
		// ImgWarp_MLS_Rigid trans2;
		// ImgWarp_PieceWiseAffine trans3;
		vector<Point_<int> > qsrc;
		vector<Point_<int> > qdst;
		qsrc.clear();
		qdst.clear();

		// numpy to Mat
		cv::Mat img_input = pyarray_to_mat(image, CV_8U);

		int cut = img_input.cols / segment;
		int threshold = cut * 1. / 3;

		struct timeb timeSeed;
		ftime(&timeSeed);
		srand(timeSeed.time * 1000 + timeSeed.millitm);

		qsrc.push_back(Point(0, 0));
		qsrc.push_back(Point(img_input.cols, 0));
		qsrc.push_back(Point(img_input.cols, img_input.rows));
		qsrc.push_back(Point(0, img_input.rows));

		qdst.push_back(Point(rand() % threshold, rand() % threshold));
		qdst.push_back(Point(img_input.cols - rand() % threshold, rand() % threshold));
		qdst.push_back(Point(img_input.cols - rand() % threshold, img_input.rows - rand() % threshold));
		qdst.push_back(Point(rand() % threshold, img_input.rows - rand() % threshold));

		for (int i = 1; i < segment; i++) {
			qsrc.push_back(Point(cut * i, 0));
			qsrc.push_back(Point(cut * i, img_input.rows));
			qdst.push_back(Point(cut * i + rand() % threshold - 0.5 * threshold, rand() % threshold - 0.5 * threshold));
			qdst.push_back(Point(cut * i + rand() % threshold - 0.5 * threshold, img_input.rows + rand() % threshold - 0.5 * threshold));
		}

		cv::Mat result = trans1.setAllAndGenerate(img_input, qsrc, qdst, img_input.cols, img_input.rows);
		// Mat to numpy
		py::array_t<uint8_t> res = cv_mat_uint8_3c_to_numpy(result);
		return res;
	}

	py::array_t<uint8_t> Stretch(py::array_t<uint8_t>& image, const int segment) {

		ImgWarp_MLS_Similarity trans1;
		// ImgWarp_MLS_Rigid trans2;
		// ImgWarp_PieceWiseAffine trans3;
		vector<Point_<int> > qsrc;
		vector<Point_<int> > qdst;
		qsrc.clear();
		qdst.clear();

		// numpy to Mat
		cv::Mat img_input = pyarray_to_mat(image, CV_8U);

		qsrc.push_back(Point(0, 0));
		qsrc.push_back(Point(img_input.cols, 0));
		qsrc.push_back(Point(img_input.cols, img_input.rows));
		qsrc.push_back(Point(0, img_input.rows));

		qdst.push_back(Point(0, 0));
		qdst.push_back(Point(img_input.cols, 0));
		qdst.push_back(Point(img_input.cols, img_input.rows));
		qdst.push_back(Point(0, img_input.rows));

		int cut = img_input.cols / segment;
		int threshold = cut * 4 / 5;

		struct timeb timeSeed;
		ftime(&timeSeed);
		srand(timeSeed.time * 1000 + timeSeed.millitm);
		int move = 0;
		for (int i = 1; i < segment; i++) {
			move = rand() % threshold - 0.5 * threshold;
			qsrc.push_back(Point(cut * i, 0));
			qsrc.push_back(Point(cut * i, img_input.rows));
			qdst.push_back(Point(cut * i + move, 0));
			qdst.push_back(Point(cut * i + move, img_input.rows));
		}

		cv::Mat result = trans1.setAllAndGenerate(img_input, qsrc, qdst, img_input.cols, img_input.rows);
		// Mat to numpy
		py::array_t<uint8_t> res = cv_mat_uint8_3c_to_numpy(result);
		return res;
	}

	py::array_t<uint8_t> Perspective(py::array_t<uint8_t>& image) {

		ImgWarp_MLS_Similarity trans1;
		// ImgWarp_MLS_Rigid trans2;
		// ImgWarp_PieceWiseAffine trans3;
		vector<Point_<int> > qsrc;
		vector<Point_<int> > qdst;
		qsrc.clear();
		qdst.clear();

		// numpy to Mat
		cv::Mat img_input = pyarray_to_mat(image, CV_8U);

		int threshold = img_input.rows * 0.5;

		struct timeb timeSeed;
		ftime(&timeSeed);
		srand(timeSeed.time * 1000 + timeSeed.millitm);

		qsrc.push_back(Point(0, 0));
		qsrc.push_back(Point(img_input.cols, 0));
		qsrc.push_back(Point(img_input.cols, img_input.rows));
		qsrc.push_back(Point(0, img_input.rows));

		qdst.push_back(Point(0, rand() % threshold));
		qdst.push_back(Point(img_input.cols, rand() % threshold));
		qdst.push_back(Point(img_input.cols, img_input.rows - rand() % threshold));
		qdst.push_back(Point(0, img_input.rows - rand() % threshold));

		cv::Mat result = trans1.setAllAndGenerate(img_input, qsrc, qdst, img_input.cols, img_input.rows);
		// Mat to numpy
		py::array_t<uint8_t> res = cv_mat_uint8_3c_to_numpy(result);
		return res;
	}
}

py::array_t<uint8_t> cv_mat_uint8_3c_to_numpy(cv::Mat& input)
{
	py::array_t<uint8_t> dst = py::array_t<uint8_t>({ input.rows,input.cols,3 }, input.data);
	return dst;
}

PYBIND11_MODULE(MLSTransformer, m) {
	m.def("Distort", &MLSTransformer::Distort, "Moving Least Squares transformation using c++.", py::arg("image"), py::arg("segment"));
	m.def("Stretch", &MLSTransformer::Stretch, "Moving Least Squares transformation using c++.", py::arg("image"), py::arg("segment"));
	m.def("Perspective", &MLSTransformer::Perspective, "Moving Least Squares transformation using c++.", py::arg("image"));
}

//PyObject*
//GenerateDistort(PyObject *img, int segment)
//{
//    NDArrayConverter cvt;
//    cv::Mat img_input;
//    img_input = cvt.toMat(img);
//    cv::Mat result = Distort(img_input, segment);
//
//    PyObject* ret = cvt.toNDArray(result);
//
//    return ret;
//}
//
//PyObject*
//GenerateStretch(PyObject *img, int segment)
//{
//    NDArrayConverter cvt;
//    cv::Mat img_input;
//    img_input = cvt.toMat(img);
//    cv::Mat result = Stretch(img_input, segment);
//
//    PyObject* ret = cvt.toNDArray(result);
//
//    return ret;
//}
//
//PyObject*
//GeneratePerspective(PyObject *img)
//{
//    NDArrayConverter cvt;
//    cv::Mat img_input;
//    img_input = cvt.toMat(img);
//    cv::Mat result = Perspective(img_input);
//
//    PyObject* ret = cvt.toNDArray(result);
//
//    return ret;
//}

//static void init()
//{
//    Py_Initialize();
//    import_array();
//}

//BOOST_PYTHON_MODULE(Augment)
//{
//    init();
//    py::def("GenerateDistort", GenerateDistort);
//    py::def("GenerateStretch", GenerateStretch);
//    py::def("GeneratePerspective", GeneratePerspective);
//}