#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include <Eigen/Geometry>
#include <arc_utilities/arc_helpers.hpp>
#include <arc_utilities/eigen_typedefs.hpp>
#include <map>
#include <vector>

#ifndef IIWA_14_FK_FAST_HPP
#define IIWA_14_FK_FAST_HPP

namespace IIWA_14_FK_FAST {
const size_t IIWA_14_NUM_ACTIVE_JOINTS = 7;
const size_t IIWA_14_NUM_LINKS = 8;

const std::string IIWA_14_ACTIVE_JOINT_1_NAME = "iiwa_joint_1";
const std::string IIWA_14_ACTIVE_JOINT_2_NAME = "iiwa_joint_2";
const std::string IIWA_14_ACTIVE_JOINT_3_NAME = "iiwa_joint_3";
const std::string IIWA_14_ACTIVE_JOINT_4_NAME = "iiwa_joint_4";
const std::string IIWA_14_ACTIVE_JOINT_5_NAME = "iiwa_joint_5";
const std::string IIWA_14_ACTIVE_JOINT_6_NAME = "iiwa_joint_6";
const std::string IIWA_14_ACTIVE_JOINT_7_NAME = "iiwa_joint_7";

const std::string IIWA_14_LINK_1_NAME = "iiwa_link_0";
const std::string IIWA_14_LINK_2_NAME = "iiwa_link_1";
const std::string IIWA_14_LINK_3_NAME = "iiwa_link_2";
const std::string IIWA_14_LINK_4_NAME = "iiwa_link_3";
const std::string IIWA_14_LINK_5_NAME = "iiwa_link_4";
const std::string IIWA_14_LINK_6_NAME = "iiwa_link_5";
const std::string IIWA_14_LINK_7_NAME = "iiwa_link_6";
const std::string IIWA_14_LINK_8_NAME = "iiwa_link_7";

inline Eigen::Isometry3d Get_link_0_joint_1_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.1575);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(0.0, 0.0, 0.0);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_1_joint_2_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.2025);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(M_PI_2, 0.0, M_PI);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_2_joint_3_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.2045, 0.0);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(M_PI_2, 0.0, M_PI);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_3_joint_4_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.2155);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(M_PI_2, 0.0, 0.0);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_4_joint_5_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.1845, 0.0);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(-M_PI_2, M_PI, 0.0);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_5_joint_6_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.0, 0.2155);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(M_PI_2, 0.0, 0.0);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline Eigen::Isometry3d Get_link_6_joint_7_LinkJointTransform(const double joint_val) {
  const Eigen::Translation3d pre_joint_translation(0.0, 0.081, 0.0);
  const Eigen::Quaterniond pre_joint_rotation = EigenHelpers::QuaternionFromUrdfRPY(-M_PI_2, M_PI, 0.0);
  const Eigen::Isometry3d pre_joint_transform = pre_joint_translation * pre_joint_rotation;
  const Eigen::Translation3d joint_translation(0.0, 0.0, 0.0);
  const Eigen::Quaterniond joint_rotation(Eigen::AngleAxisd(joint_val, Eigen::Vector3d::UnitZ()));
  const Eigen::Isometry3d joint_transform = joint_translation * joint_rotation;
  return (pre_joint_transform * joint_transform);
}

inline EigenHelpers::VectorIsometry3d GetLinkTransforms(
    const std::vector<double>& configuration, const Eigen::Isometry3d& base_transform = Eigen::Isometry3d::Identity()) {
  assert(configuration.size() == IIWA_14_NUM_ACTIVE_JOINTS);
  EigenHelpers::VectorIsometry3d link_transforms(IIWA_14_NUM_LINKS);
  link_transforms[0] = base_transform;
  link_transforms[1] = link_transforms[0] * Get_link_0_joint_1_LinkJointTransform(configuration[0]);
  link_transforms[2] = link_transforms[1] * Get_link_1_joint_2_LinkJointTransform(configuration[1]);
  link_transforms[3] = link_transforms[2] * Get_link_2_joint_3_LinkJointTransform(configuration[2]);
  link_transforms[4] = link_transforms[3] * Get_link_3_joint_4_LinkJointTransform(configuration[3]);
  link_transforms[5] = link_transforms[4] * Get_link_4_joint_5_LinkJointTransform(configuration[4]);
  link_transforms[6] = link_transforms[5] * Get_link_5_joint_6_LinkJointTransform(configuration[5]);
  link_transforms[7] = link_transforms[6] * Get_link_6_joint_7_LinkJointTransform(configuration[6]);
  return link_transforms;
}

inline EigenHelpers::VectorIsometry3d GetLinkTransforms(
    const std::map<std::string, double>& configuration,
    const Eigen::Isometry3d& base_transform = Eigen::Isometry3d::Identity()) {
  std::vector<double> configuration_vector(IIWA_14_NUM_ACTIVE_JOINTS);
  configuration_vector[0] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_1_NAME, 0.0);
  configuration_vector[1] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_2_NAME, 0.0);
  configuration_vector[2] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_3_NAME, 0.0);
  configuration_vector[3] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_4_NAME, 0.0);
  configuration_vector[4] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_5_NAME, 0.0);
  configuration_vector[5] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_6_NAME, 0.0);
  configuration_vector[6] = arc_helpers::RetrieveOrDefault(configuration, IIWA_14_ACTIVE_JOINT_7_NAME, 0.0);
  return GetLinkTransforms(configuration_vector, base_transform);
}

inline EigenHelpers::MapStringIsometry3d GetLinkTransformsMap(
    const std::vector<double>& configuration, const Eigen::Isometry3d& base_transform = Eigen::Isometry3d::Identity()) {
  const EigenHelpers::VectorIsometry3d link_transforms = GetLinkTransforms(configuration, base_transform);
  EigenHelpers::MapStringIsometry3d link_transforms_map;
  link_transforms_map[IIWA_14_LINK_1_NAME] = link_transforms[0];
  link_transforms_map[IIWA_14_LINK_2_NAME] = link_transforms[1];
  link_transforms_map[IIWA_14_LINK_3_NAME] = link_transforms[2];
  link_transforms_map[IIWA_14_LINK_4_NAME] = link_transforms[3];
  link_transforms_map[IIWA_14_LINK_5_NAME] = link_transforms[4];
  link_transforms_map[IIWA_14_LINK_6_NAME] = link_transforms[5];
  link_transforms_map[IIWA_14_LINK_7_NAME] = link_transforms[6];
  link_transforms_map[IIWA_14_LINK_8_NAME] = link_transforms[7];
  return link_transforms_map;
}

inline EigenHelpers::MapStringIsometry3d GetLinkTransformsMap(
    const std::map<std::string, double>& configuration,
    const Eigen::Isometry3d& base_transform = Eigen::Isometry3d::Identity()) {
  const EigenHelpers::VectorIsometry3d link_transforms = GetLinkTransforms(configuration, base_transform);
  EigenHelpers::MapStringIsometry3d link_transforms_map;
  link_transforms_map[IIWA_14_LINK_1_NAME] = link_transforms[0];
  link_transforms_map[IIWA_14_LINK_2_NAME] = link_transforms[1];
  link_transforms_map[IIWA_14_LINK_3_NAME] = link_transforms[2];
  link_transforms_map[IIWA_14_LINK_4_NAME] = link_transforms[3];
  link_transforms_map[IIWA_14_LINK_5_NAME] = link_transforms[4];
  link_transforms_map[IIWA_14_LINK_6_NAME] = link_transforms[5];
  link_transforms_map[IIWA_14_LINK_7_NAME] = link_transforms[6];
  link_transforms_map[IIWA_14_LINK_8_NAME] = link_transforms[7];
  return link_transforms_map;
}
}  // namespace IIWA_14_FK_FAST

#endif  // IIWA_14_FK_FAST_HPP


#ifndef INCLUDED_STDDEFX
#include "stddefx.h"
#define INCLUDED_STDDEFX
#endif

#ifndef INCLUDED_CALC
#include "calc.h"
#define INCLUDED_CALC
#endif

// Library headers.
#ifndef INCLUDED_IOSTREAM
#include <iostream>
#define INCLUDED_IOSTREAM
#endif
#ifndef INCLUDED_BOOST_NONCOPYABLE
#include <boost/noncopyable.hpp>
#define INCLUDED_BOOST_NONCOPYABLE
#endif

// PCRaster library headers.
#ifndef INCLUDED_FIELDAPI_INTERFACE
#include "fieldapi_interface.h"
#define INCLUDED_FIELDAPI_INTERFACE
#endif

#ifndef INCLUDED_FIELDAPI_SCALARDOMAINCHECK
#include "fieldapi_scalardomaincheck.h"
#define INCLUDED_FIELDAPI_SCALARDOMAINCHECK
#endif

#ifndef INCLUDED_GEO_CELLLOCVISITOR
#include "geo_celllocvisitor.h"
#define INCLUDED_GEO_CELLLOCVISITOR
#endif

#ifndef INCLUDED_GEO_SIMPLERASTER
#include "geo_simpleraster.h"
#define INCLUDED_GEO_SIMPLERASTER
#endif

#ifndef INCLUDED_MISC
#include "misc.h"
#define INCLUDED_MISC
#endif

#ifndef INCLUDED_COM_INTERVALTYPES
#include "com_intervaltypes.h"
#define INCLUDED_COM_INTERVALTYPES
#endif

// Module headers.

//------------------------------------------------------------------------------
// Implicit finite difference according to:
// Introduction to groundwater modeling
// Wang & Anderson
// Crank-Nicolson method, Gauss-Seidel iteration
//------------------------------------------------------------------------------

namespace calc {

struct CalculationArgs {

  double d_area;     // Cell area, constant.
  double d_alpha;    // Constant.
  double d_nrCells;  // Number of surrounding cells contributing to current cel.
  double d_interval;
  double d_currentElevation;
  double d_recharge;
  double d_transmissivity;   // Average t for all directions.
  double d_storageCoefficient;
  double d_h1;       // h * t for all directions, timestep n
  double d_h2;       // h * t for all directions, timestep n + 1.

};

double calcElevation(const CalculationArgs& args)
{
  double f1 = (args.d_area * args.d_storageCoefficient) / args.d_interval;
  double f2 = 1.0 / (f1 + args.d_nrCells * args.d_transmissivity *
                   args.d_alpha);

  return f2 * (args.d_alpha * args.d_h2 + f1 * args.d_currentElevation +
                   (1.0 - args.d_alpha) *
                   (args.d_h1 - args.d_currentElevation * args.d_nrCells *
                   args.d_transmissivity) + args.d_area * args.d_recharge);
}

class DiffuseAlgorithm: public boost::noncopyable
{

public:

  enum Direction {
    NoDirection = 0x000,
    Left = 0x0001,
    Top = 0x0002,
    Right = 0x0004,
    Bottom = 0x0008,
    UpperLeft = Top | Left,
    UpperRight = Top | Right,
    LowerLeft = Bottom | Left,
    LowerRight = Bottom | Right,
    Enclosed = Left | Top | Right | Bottom
  };

  enum FlowCondition  {
    NOFLOW = 0,
    CALCULATE = 1,
    FIXED = 2
  };

private:

  fieldapi::ReadWriteReal8&  d_resultElevation;
  const fieldapi::ReadOnlyReal8& d_elevation;
  const fieldapi::ReadOnlyReal8& d_recharge;
  const fieldapi::ReadOnlyReal8& d_transmissivity;
  const fieldapi::ReadOnlyInt4& d_flowCondition;
  const fieldapi::ReadOnlyReal8& d_storageCoefficient;

  CalculationArgs d_args;
  geo::CellLoc d_location;    // A cell location.

  geo::SimpleRaster<double> d_tmRight; // Average transm. values, right.
  geo::SimpleRaster<double> d_tmBottom; // Average transm. values, bottom.

  geo::SimpleRaster<unsigned int> d_noFlowBoundaries;

  double calcElevationWithNoFlowBoundary(size_t r, size_t c, unsigned int direction)
  {
    if(direction == Enclosed) {
      // Current cell is surrounded by no flow boundaries. Current elevation
      // remains unchanged.
      return d_elevation.value(r, c);
    }

    // Calculate new elevation.
    d_args.d_h1 = 0.0;
    d_args.d_h2 = 0.0;
    d_args.d_nrCells = 0;

    // Harmonic means of transmissivities.
    double tmLeft(0), tmTop(0), tmRight(0), tmBottom(0);

    // Cell to the left.
    if(!(direction & Left)) {
      tmLeft = d_tmRight.cell(r, c - 1);
      d_args.d_h1 += tmLeft * d_elevation.value(r, c - 1);
      d_args.d_h2 += tmLeft * d_resultElevation.value(r, c - 1);
      ++d_args.d_nrCells;
    }
    else if(!(direction & Right)) {
      // Copy value from Right cell.
      tmLeft = d_tmRight.cell(r, c);
      d_args.d_h1 += tmLeft * d_elevation.value(r, c + 1);
      d_args.d_h2 += tmLeft * d_resultElevation.value(r, c + 1);
      ++d_args.d_nrCells;
    }

    // Cell to the top.
    if(!(direction & Top)) {
      tmTop = d_tmBottom.cell(r - 1, c);
      d_args.d_h1 += tmTop * d_elevation.value(r - 1, c);
      d_args.d_h2 += tmTop * d_resultElevation.value(r - 1, c);
      ++d_args.d_nrCells;
    }
    else if(!(direction & Bottom)) {
      // Copy value from Bottom cell.
      tmTop = d_tmBottom.cell(r, c);
      d_args.d_h1 += tmTop * d_elevation.value(r + 1, c);
      d_args.d_h2 += tmTop * d_resultElevation.value(r + 1, c);
      ++d_args.d_nrCells;
    }

    // Cell to the right.
    if(!(direction & Right)) {
      tmRight = d_tmRight.cell(r, c);
      d_args.d_h1 += tmRight * d_elevation.value(r, c + 1);
      d_args.d_h2 += tmRight * d_resultElevation.value(r, c + 1);
      ++d_args.d_nrCells;
    }
    else if(!(direction & Left)) {
      // Copy value from Left cell.
      tmRight = d_tmRight.cell(r, c - 1);
      d_args.d_h1 += tmRight * d_elevation.value(r, c - 1);
      d_args.d_h2 += tmRight * d_resultElevation.value(r, c - 1);
      ++d_args.d_nrCells;
    }

    // Cell to the bottom.
    if(!(direction & Bottom)) {
      tmBottom = d_tmBottom.cell(r, c);
      d_args.d_h1 += tmBottom * d_elevation.value(r + 1, c);
      d_args.d_h2 += tmBottom * d_resultElevation.value(r + 1, c);
      ++d_args.d_nrCells;
    }
    else if(!(direction & Top)) {
      // Copy value from Top cell.
      tmBottom = d_tmBottom.cell(r - 1, c);
      d_args.d_h1 += tmBottom * d_elevation.value(r - 1, c);
      d_args.d_h2 += tmBottom * d_resultElevation.value(r - 1, c);
      ++d_args.d_nrCells;
    }

    PRECOND(d_args.d_nrCells > 0);

    d_args.d_currentElevation = d_elevation.value(r, c);
    d_args.d_recharge = d_recharge.value(r, c);
    // Arithmic mean.
    d_args.d_transmissivity = (tmLeft + tmTop + tmRight + tmBottom) /
                   d_args.d_nrCells;
    d_args.d_storageCoefficient = d_storageCoefficient.value(r, c);

    return calc::calcElevation(d_args);
  }

  double calcElevationWithoutNoFlowBoundary(size_t r, size_t c)
  {
    d_args.d_h1 = 0.0;       // Elevation at t = n, * transmissivities.
    d_args.d_h2 = 0.0;       // Elevation at t = n + 1, * transmissivities.
    d_args.d_nrCells = 0;    // Nr of cells contributing to the current value.
    static double tmLeft, tmTop, tmRight, tmBottom; // Harm. means of transm.

    // The current cell has no no flow boundaries. All neighbours have
    // valid values.

    // Cell to the left.
    tmLeft = d_tmRight.cell(r, c - 1);
    d_args.d_h1 += tmLeft * d_elevation.value(r, c - 1);
    d_args.d_h2 += tmLeft * d_resultElevation.value(r, c - 1);
    ++d_args.d_nrCells;

    // Cell to the top.
    tmTop = d_tmBottom.cell(r - 1, c);
    d_args.d_h1 += tmTop * d_elevation.value(r - 1, c);
    d_args.d_h2 += tmTop * d_resultElevation.value(r - 1, c);
    ++d_args.d_nrCells;

    // Cell to the right.
    tmRight = d_tmRight.cell(r, c);
    d_args.d_h1 += tmRight * d_elevation.value(r, c + 1);
    d_args.d_h2 += tmRight * d_resultElevation.value(r, c + 1);
    ++d_args.d_nrCells;

    // Cell to the bottom.
    tmBottom = d_tmBottom.cell(r, c);
    d_args.d_h1 += tmBottom * d_elevation.value(r + 1, c);
    d_args.d_h2 += tmBottom * d_resultElevation.value(r + 1, c);
    ++d_args.d_nrCells;

    PRECOND(d_args.d_nrCells > 0);
    PRECOND(d_args.d_nrCells == 4);

    d_args.d_currentElevation = d_elevation.value(r, c);
    d_args.d_recharge = d_recharge.value(r, c);
    // Arithmic mean.
    d_args.d_transmissivity = (tmLeft + tmTop + tmRight + tmBottom) /
                   d_args.d_nrCells;
    d_args.d_storageCoefficient = d_storageCoefficient.value(r, c);

    return calc::calcElevation(d_args);
  }

public:

  DiffuseAlgorithm(fieldapi::ReadWriteReal8& resultElevation,
         const fieldapi::ReadOnlyReal8& elevation,
         const fieldapi::ReadOnlyReal8& recharge,
         const fieldapi::ReadOnlyReal8& transmissivity,
         const fieldapi::ReadOnlyInt4& flowCondition,
         const fieldapi::ReadOnlyReal8& storageCoefficient,
         double area, double alpha, double interval)

    : d_resultElevation(resultElevation), d_elevation(elevation),
      d_recharge(recharge), d_transmissivity(transmissivity),
      d_flowCondition(flowCondition), d_storageCoefficient(storageCoefficient),
      d_tmRight(resultElevation.nrRows(), resultElevation.nrCols()),
      d_tmBottom(resultElevation.nrRows(), resultElevation.nrCols()),
      d_noFlowBoundaries(resultElevation.nrRows(), resultElevation.nrCols(),
      static_cast<const unsigned int>(NoDirection))

  {
    d_args.d_area = area;
    d_args.d_alpha = alpha;
    d_args.d_interval = interval;

    geo::CellLoc loc;
    size_t r, c; // comes from C ya know

    //--------------------------------------------------------------------------
    // Calculate transmissivities. These stay constant within a call to the
    // function and are needed for each iteration so we calculate them here
    // once. We trade memory space for speed.

    // Harmonic mean transmissivities current and right cell.
    // For every row.
    for( r = 0; r < d_tmRight.nrRows(); ++r) {

      loc.setRow(r);

      // For every but the last col.
      for( c = 0; c < d_tmRight.nrCols() - 1; ++c) {

        loc.setCol(c);

        if(!d_transmissivity.isMV(loc) &&
                   !d_transmissivity.isMV(geo::CellLoc(r, c + 1))) {
          d_tmRight.cell(r, c) =
              2 * d_transmissivity.value(r, c + 1) *
              d_transmissivity.value(r, c) /
              (d_transmissivity.value(r, c + 1) + d_transmissivity.value(r, c));
        }
      }
    }

    // Harmonic mean transmissivities current and bottom cell.
    // For every but the last row.
    for( r = 0; r < d_tmBottom.nrRows() - 1; ++r) {

      loc.setRow(r);

      // For every col.
      for( c = 0; c < d_tmBottom.nrCols(); ++c) {

        loc.setCol(c);

        if(!d_transmissivity.isMV(loc) &&
                   !d_transmissivity.isMV(geo::CellLoc(r + 1, c))) {
          d_tmBottom.cell(r, c) =
              2 * d_transmissivity.value(r + 1, c) *
              d_transmissivity.value(r, c) /
              (d_transmissivity.value(r + 1, c) + d_transmissivity.value(r, c));
        }
      }
    }

    //--------------------------------------------------------------------------
    // Calculate no flow boundaries. These stay constant within a call to the
    // function and are needed for each iteration so we calculate them here
    // once. We trade memory space for speed.
    // d_noFlowBoundaries is already initialized with NoDirection.
    // A bounding cell is regarded as a now flow boundary if:
    // - it is missing (border cell of the raster)
    // - it contains a missing value
    // - if its flow condition is NOFLOW

    PRECOND(d_noFlowBoundaries.nrRows() >= 1);
    PRECOND(d_noFlowBoundaries.nrCols() >= 1);

    // For every but the first and last row.
    for( r = 1; r < d_noFlowBoundaries.nrRows() - 1; ++r) {

      // For every but the first and last col.
      for( c = 1; c < d_noFlowBoundaries.nrCols() - 1; ++c) {

        // Cell to the left.
        loc.setRow(r);
        loc.setCol(c - 1);
        if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
          d_noFlowBoundaries.cell(r, c) |= Left;
        }

        // Cell to the top.
        loc.setRow(r - 1);
        loc.setCol(c);
        if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
          d_noFlowBoundaries.cell(r, c) |= Top;
        }

        // Cell to the right.
        loc.setRow(r);
        loc.setCol(c + 1);
        if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
          d_noFlowBoundaries.cell(r, c) |= Right;
        }

        // Cell to the bottom.
        loc.setRow(r + 1);
        loc.setCol(c);
        if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r + 1, c) == NOFLOW) {
          d_noFlowBoundaries.cell(r, c) |= Bottom;
        }
      }
    }

    // Handle no-flow conditions on the borders of the area. The value of head
    // along the fictious column reflects across the border.

    // Process borders of raster.

    // First col.
    c = 0;

    // For every but the first and last row.
    for(r = 1; r < d_noFlowBoundaries.nrRows() - 1; ++r) {

      // Cell to the left.
      d_noFlowBoundaries.cell(r, c) |= Left;

      // Cell to the top.
      loc.setRow(r - 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Top;
      }

      // Cell to the right.
      loc.setRow(r);
      loc.setCol(c + 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Right;
      }

      // Cell to the bottom.
      loc.setRow(r + 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                 d_flowCondition.value(r + 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Bottom;
      }
    }

    // Last col.
    c = d_noFlowBoundaries.nrCols() - 1;

    // For every but the first and last row.
    for(r = 1; r < d_noFlowBoundaries.nrRows() - 1; ++r) {

      // Cell to the left.
      loc.setRow(r);
      loc.setCol(c - 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Left;
      }

      // Cell to the top.
      loc.setRow(r - 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Top;
      }

      // Cell to the right.
      d_noFlowBoundaries.cell(r, c) |= Right;

      // Cell to the bottom.
      loc.setRow(r + 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                 d_flowCondition.value(r + 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Bottom;
      }
    }

    // First row.
    r = 0;

    // For every but the first and last col.
    for(c = 1; c < d_noFlowBoundaries.nrCols() - 1; ++c) {

      // Cell to the left.
      loc.setRow(r);
      loc.setCol(c - 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Left;
      }

      // Cell to the top.
      d_noFlowBoundaries.cell(r, c) |= Top;

      // Cell to the right.
      loc.setRow(r);
      loc.setCol(c + 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Right;
      }

      // Cell to the bottom.
      loc.setRow(r + 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                 d_flowCondition.value(r + 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Bottom;
      }
    }

    // Last row.
    r = d_noFlowBoundaries.nrRows() - 1;

    // For every but the first and last col.
    for(c = 1; c < d_noFlowBoundaries.nrCols() - 1; ++c) {

      // Cell to the left.
      loc.setRow(r);
      loc.setCol(c - 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Left;
      }

      // Cell to the top.
      loc.setRow(r - 1);
      loc.setCol(c);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Top;
      }

      // Cell to the right.
      loc.setRow(r);
      loc.setCol(c + 1);
      if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
        d_noFlowBoundaries.cell(r, c) |= Right;
      }

      // Cell to the bottom.
      d_noFlowBoundaries.cell(r, c) |= Bottom;
    }

    // Upper left corner.
    r = 0;
    c = 0;

    d_noFlowBoundaries.cell(r, c) |= UpperLeft;

    // Cell to the right.
    loc.setRow(r);
    loc.setCol(c + 1);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Right;
    }

    // Cell to the bottom.
    loc.setRow(r + 1);
    loc.setCol(c);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                 d_flowCondition.value(r + 1, c) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Bottom;
    }

    // Upper right corner.
    r = 0;
    c = d_noFlowBoundaries.nrCols() - 1;

    d_noFlowBoundaries.cell(r, c) |= UpperRight;

    // Cell to the left.
    loc.setRow(r);
    loc.setCol(c - 1);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Left;
    }

    // Cell to the bottom.
    loc.setRow(r + 1);
    loc.setCol(c);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                 d_flowCondition.value(r + 1, c) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Bottom;
    }

    // Lower right corner.
    r = d_noFlowBoundaries.nrRows() - 1;
    c = d_noFlowBoundaries.nrCols() - 1;

    d_noFlowBoundaries.cell(r, c) |= LowerRight;

    // Cell to the left.
    loc.setRow(r);
    loc.setCol(c - 1);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c - 1) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Left;
    }

    // Cell to the top.
    loc.setRow(r - 1);
    loc.setCol(c);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Top;
    }

    // Lower left corner.
    r = d_noFlowBoundaries.nrRows() - 1;
    c = 0;

    d_noFlowBoundaries.cell(r, c) |= LowerLeft;

    // Cell to the top.
    loc.setRow(r - 1);
    loc.setCol(c);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r - 1, c) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Top;
    }

    // Cell to the right.
    loc.setRow(r);
    loc.setCol(c + 1);
    if(d_elevation.isMV(loc) || d_flowCondition.isMV(loc) ||
                   d_flowCondition.value(r, c + 1) == NOFLOW) {
      d_noFlowBoundaries.cell(r, c) |= Right;
    }
  }

//! Calculates the new elevation at a cell position.
/*!
  \param     r Row of cell.
  \param     c Column of cell.
  \return    Elevation.
  \warning   All input variables and d_resultElevation must have a valid value
             at \a r, \a c. This is forced by checking d_resultElevation.
*/
  double calcElevation(size_t r, size_t c)
  {
    PRECOND(!d_resultElevation.isMV(geo::CellLoc(r, c)));

    if(d_noFlowBoundaries.cell(r, c) != NoDirection) {
      return calcElevationWithNoFlowBoundary(r, c, d_noFlowBoundaries.cell(r, c));
    }
    else {
      return calcElevationWithoutNoFlowBoundary(r, c);
    }
  }

};

} // namespace calc.

//!
/*!
  \param     .
  \return    .
  \exception .
  \warning   .
  \sa        .

  Elevation                  scalar [ L       ]
  Recharge                   scalar [ L , T-1 ]
  Transmissivity             scalar [ L2, T-1 ]
  Storage coefficient        scalar [ L3, L-3 ]
  Interval                   scalar [ T       ]       > 0.0

  Issues: cell length, error messages and handling of return values
*/
extern "C" int  Transient(void** out, const void** in, int nrArgs)
{
  using namespace calc; // hack hack
  PRECOND(nrArgs == 7);
  (void)nrArgs; // shut up compiler

  ReadWriteReal8_ref(resultElevation, static_cast<MAP_REAL8*>(out[0]));
  ReadOnlyReal8_ref(elevation, static_cast<const MAP_REAL8*>(in[0]));
  ReadOnlyReal8_ref(recharge, static_cast<const MAP_REAL8*>(in[1]));
  ReadOnlyReal8_ref(transmissivity, static_cast<const MAP_REAL8*>(in[2]));
  ReadOnlyInt4_ref(flowCondition, static_cast<const MAP_INT4*>(in[3]));
  ReadOnlyReal8_ref(storageCoefficient, static_cast<const MAP_REAL8*>(in[4]));
  ReadOnlyReal8_ref(intervalInterface, static_cast<const MAP_REAL8*>(in[5]));
  ReadOnlyReal8_ref(toleranceInterface, static_cast<const MAP_REAL8*>(in[6]));

  PRECOND(!intervalInterface.spatial());
  PRECOND(!toleranceInterface.spatial());

  std::vector<const fieldapi::Common*> inputs;
  inputs.push_back(&elevation);
  inputs.push_back(&recharge);
  inputs.push_back(&transmissivity);
  inputs.push_back(&flowCondition);
  inputs.push_back(&storageCoefficient);
  inputs.push_back(&intervalInterface);
  inputs.push_back(&toleranceInterface);
  PRECOND(inputs.size() == static_cast<size_t>(nrArgs));

  size_t nrRows = elevation.nrRows();
  size_t nrCols = elevation.nrCols();
  size_t r,c;
  PRECOND(nrRows >= 2 && nrCols >= 2);

  // Domain checks...
  std::vector<fieldapi::ScalarDomainCheck> scalarDomains, nsDomains;
  scalarDomains.push_back(fieldapi::ScalarDomainCheck(transmissivity,
                   "transmissivity", com::GreaterThan<double>(0)));
  scalarDomains.push_back(fieldapi::ScalarDomainCheck(storageCoefficient,
                   "storage coefficient", com::GreaterThan<double>(0)));
  nsDomains.push_back(fieldapi::ScalarDomainCheck(intervalInterface,
                   "interval", com::GreaterThan<double>(0)));
  nsDomains.push_back(fieldapi::ScalarDomainCheck(toleranceInterface,
                   "tolerance", com::GreaterThan<double>(0)));

  // Checks...
  int check;

  // Check non spatials.
  check = fieldapi::checkScalarDomains(nsDomains, geo::CellLoc(0, 0));
  if(check != -1) {
    return RetError(1, nsDomains[check].msg().c_str());
  }

  // Check spatials.
  for(geo::CellLocVisitor loc(elevation); loc.valid(); ++loc) {
    check = fieldapi::checkScalarDomains(scalarDomains, *loc);
    if(check != -1) {
      return RetError(1, scalarDomains[check].msg().c_str());
    }
  }

  // Inititalize...
  double oldValue;
  double tolerance = toleranceInterface.value(0, 0);
  double difference, maxDifference;    // Current and max difference.
  size_t nrIterations = 0;

  // yepyep: fieldapi::Common::cellLength();
  MAP_REAL8* map = static_cast<MAP_REAL8*>(out[0]);
  double cellLength = map->CellLength(map);
  PRECOND(cellLength > 0.0);

  // Result elevation is missing value if any of the inputs is MV.
  for(geo::CellLocVisitor visitor(elevation); visitor.valid(); ++visitor) {

    if(fieldapi::nonMV(inputs, *visitor)) {
      resultElevation.copy(elevation, *visitor);
    }
    else {
      resultElevation.putMV(*visitor);
    }
  }

  DiffuseAlgorithm algorithm(resultElevation, elevation, recharge,
                   transmissivity, flowCondition, storageCoefficient,
                   cellLength * cellLength, 0.5, intervalInterface.value(0, 0));
  // std::cout << std::endl;

  // Algorithm...
  // Loop over all cells in the raster untill the maximum difference between
  // the current elevation and the new elevation is smaller than the tolerance.
  do {
    maxDifference = 0.0;

    for( r = 0; r < nrRows; ++r) {
      for( c = 0; c < nrCols; ++c) {

        // Skip missing values.
        if(resultElevation.isMV(geo::CellLoc(r, c))) {
          continue;
        }

        // Remember current elevation.
        oldValue = resultElevation.value(r, c);

        // Select calculation method for current cell, based on flow condition.
        if(flowCondition.value(r, c) == DiffuseAlgorithm::CALCULATE) {
          resultElevation.put(algorithm.calcElevation(r, c), r, c);
        }
        // (else FIXED, NOFLOW)

        // Determine difference.
        difference = ABS(resultElevation.value(r, c) - oldValue);
        if(difference > maxDifference) {
          maxDifference = difference;
        }
      }
    }

    ++nrIterations;

    // yepyep: test op aantal iteraties?
    // PRECOND(nrIterations < 30);
    // std::cout << '.' << std::flush;

  } while (maxDifference > tolerance);

  // std::cout << ": " << nrIterations << std::endl;

  return 0;
}


#include <iostream>
#include <fstream>
#include <armadillo>
#include <math.h>
#include <omp.h>
#include <mpi.h>
#include <unordered_set>

#include "class_partition.hpp"
#include "constants.hpp"

/////////////////// Below are private member functions /////////////////

/*
Creats knots in the current region
Input:
    xMin, xMax, yMin, yMax: the boudary of the current region
    nKnotsX, nKnotsY: the number of knots in the current region in each direction
Output:
    currentKnotsX: the array of the coordinates in longitude for each knot
    currentKnotsY: the array of the coordinates in latitude for each knot
*/
void Partition::create_knots(double *&currentKnotsX, double *&currentKnotsY, const double &xMin, const double &xMax, const int &nKnotsX, const double &yMin, const double &yMax, const int &nKnotsY)
{
    double offsetX = (xMax-xMin)*OFFSET;
    double offsetY = (yMax-yMin)*OFFSET;

    double xStart = xMin+offsetX;
    double xEnd = xMax-offsetX;
    double yStart = yMin+offsetY;
    double yEnd = yMax-offsetY;

    double xIncrement = ( nKnotsX==1 ? 0 : (xEnd-xStart)/(nKnotsX-1));
    double yIncrement = ( nKnotsY==1 ? 0 : (yEnd-yStart)/(nKnotsY-1));

    #pragma omp parallel for
    for(int iKnotX = 0; iKnotX < nKnotsX; iKnotX++)
        for(int jKnotY = 0; jKnotY < nKnotsY; jKnotY++)
        {
            currentKnotsX[iKnotX*nKnotsY+jKnotY] = xEnd-(nKnotsX-iKnotX-1)*xIncrement;//xStart+iKnotX*xIncrement;
            currentKnotsY[iKnotX*nKnotsY+jKnotY] = yEnd-(nKnotsY-jKnotY-1)*yIncrement;//yStart+jKnotY*yIncrement;
        }
}

/*
Creats partitions in the current region
*/
void Partition::create_partitions(unsigned long region, double *partitionXMin, double *partitionXMax, double *partitionYMin, double *partitionYMax)
{
    double xMin=partitionXMin[region];
    double xMax=partitionXMax[region];
    double yMin=partitionYMin[region];
    double yMax=partitionYMax[region];

    if(NUM_PARTITIONS_J==2)
    {
        if((xMax-xMin)>=(yMax-yMin))
        {
            double xMid = (xMax+xMin)/2;
            partitionXMin[region*2+1]=xMin; partitionXMin[region*2+2]=xMid;
            partitionXMax[region*2+1]=xMid; partitionXMax[region*2+2]=xMax;

            partitionYMin[region*2+1]=yMin; partitionYMin[region*2+2]=yMin;
            partitionYMax[region*2+1]=yMax; partitionYMax[region*2+2]=yMax;
        }else
        {
            double yMid = (yMax+yMin)/2;
            partitionXMin[region*2+1]=xMin; partitionXMin[region*2+2]=xMin;
            partitionXMax[region*2+1]=xMax; partitionXMax[region*2+2]=xMax;

            partitionYMin[region*2+1]=yMin; partitionYMin[region*2+2]=yMid;
            partitionYMax[region*2+1]=yMid; partitionYMax[region*2+2]=yMax;
        }
    }else
    {
        double xMid = (xMax+xMin)/2;
        double yMid = (yMax+yMin)/2;

        partitionXMin[region*4+1]=xMin; partitionXMin[region*4+2]=xMid; partitionXMin[region*4+3]=xMin; partitionXMin[region*4+4]=xMid;

        partitionXMax[region*4+1]=xMid; partitionXMax[region*4+2]=xMax; partitionXMax[region*4+3]=xMid; partitionXMax[region*4+4]=xMax;

        partitionYMin[region*4+1]=yMin; partitionYMin[region*4+2]=yMin; partitionYMin[region*4+3]=yMid; partitionYMin[region*4+4]=yMid;

        partitionYMax[region*4+1]=yMid; partitionYMax[region*4+2]=yMid; partitionYMax[region*4+3]=yMax; partitionYMax[region*4+4]=yMax;
    }
}

/////////////////// Above are private member functions /////////////////

/////////////////// Below are public member functions /////////////////

/*Implement the constructor of Partition. Assign values to nRegionsAtEachLevel for the array of the number of regions in each level, nRegionsInTotal for the total number of regions, and MPI quantities.
*/
Partition::Partition()
{
    //If user did not predetermine the number of levels, find the number of levels such that the average number of observations per region similar to the number of knots
    if(NUM_LEVELS_M == -99)//-99 stands for the default method to automatically determine NUM_LEVELS_M
    {
        //Find the number of regions required at finest level, nRegionsAtFinestLevel
        unsigned long nRegionsAtFinestLevel = ceil((double)NUM_OBSERVATIONS/(double)NUM_KNOTS_r);

        //Find the number of levels, NUM_LEVELS_M, which satisfies NUM_PARTITIONS_J^(NUM_LEVELS_M-1)>=nRegionsAtFinestLevel
        NUM_LEVELS_M = ceil(log(nRegionsAtFinestLevel)/log(NUM_PARTITIONS_J))+1;

        if(NUM_LEVELS_M < 2)
	    {
            if(WORKER == 0) cout<<"Program exits with an error: the NUM_LEVELS_M calculated by default is "<<NUM_LEVELS_M<<", which is required to be larger than 1. Please decrease NUM_KNOTS_r to get a larger NUM_LEVELS_M.\n";
            MPI_Barrier(MPI_COMM_WORLD);
            MPI_Finalize();
            exit(EXIT_FAILURE);
	    }
    }

    //Find the array of the number of regions in each level, nRgnsVec, and the total number of regions, nRegionsInTotal
    nRegionsAtEachLevel = new unsigned long [NUM_LEVELS_M];
    nRegionsAtEachLevel[0] = 1;
    nRegionsInTotal = 1;
    for(int iLevel = 1; iLevel < NUM_LEVELS_M; iLevel++)
    {
        nRegionsAtEachLevel[iLevel] = nRegionsAtEachLevel[iLevel-1]*NUM_PARTITIONS_J;
        nRegionsInTotal += nRegionsAtEachLevel[iLevel];
    }

    //Calculate MPI information
        WORKING_REGION_FLAG = new bool [nRegionsInTotal]();

        unsigned long indexRegionAtFinestLevelStartThisWorker, indexRegionAtFinestLevelEndThisWorker;
        WORKERS_FOR_EACH_REGION = new std::set<unsigned long>[nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]];

        //Assign regions in the finest level to the worker
        unsigned long nRegionsPerWorker = nRegionsAtEachLevel[NUM_LEVELS_M-1] / MPI_SIZE;
        unsigned long nWorkersWithAdditionalRegion = nRegionsAtEachLevel[NUM_LEVELS_M-1] - nRegionsPerWorker * MPI_SIZE;

        if(WORKER < nWorkersWithAdditionalRegion)
        {
            indexRegionAtFinestLevelStartThisWorker =  WORKER * (nRegionsPerWorker+1);
            indexRegionAtFinestLevelEndThisWorker =  indexRegionAtFinestLevelStartThisWorker + nRegionsPerWorker;
        }
        else
        {
            indexRegionAtFinestLevelStartThisWorker =  nWorkersWithAdditionalRegion * (nRegionsPerWorker+1) + (WORKER -  nWorkersWithAdditionalRegion)* nRegionsPerWorker;
            indexRegionAtFinestLevelEndThisWorker =  indexRegionAtFinestLevelStartThisWorker + nRegionsPerWorker - 1;
        }

        indexRegionAtFinestLevelStartThisWorker += nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1];
        indexRegionAtFinestLevelEndThisWorker += nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1];

        for(unsigned long i = indexRegionAtFinestLevelStartThisWorker; i < indexRegionAtFinestLevelEndThisWorker+1; i++)
            WORKING_REGION_FLAG[i] = true;

        //Find the regions dealt with by this worker at the second finest level
        for(unsigned long index = indexRegionAtFinestLevelStartThisWorker; index <= indexRegionAtFinestLevelEndThisWorker; index++)
        {
            unsigned long ancestor = index;
            while(ancestor > 0)
            {
                ancestor = (ancestor-1)/NUM_PARTITIONS_J;
                //INDICES_REGIONS.insert(ancestor);
                WORKING_REGION_FLAG[ancestor] = true;
            }
            INDICES_REGIONS_AT_CURRENT_LEVEL.insert((index-1)/NUM_PARTITIONS_J);
        }

        //Assign WORKERS_FOR_EACH_REGION at the second finest level
        unsigned long indexStartFinestLevel = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1];
        unsigned long tmpIndexStart = indexStartFinestLevel, tmpIndexEnd;
        if(DYNAMIC_SCHEDULE_FLAG) origninalWorker = new int [nRegionsInTotal];

        for(int iWorker = 0; iWorker < nWorkersWithAdditionalRegion; iWorker++)
        {
            tmpIndexEnd =  tmpIndexStart + nRegionsPerWorker + 1;
            for(unsigned long index = tmpIndexStart; index < tmpIndexEnd; index++)
            {
                unsigned long ancestor = (index-1)/NUM_PARTITIONS_J;
                WORKERS_FOR_EACH_REGION[ancestor].insert(iWorker);
                if(DYNAMIC_SCHEDULE_FLAG)
                {
                    origninalWorker[index] = iWorker;
                    origninalWorker[ancestor] = iWorker;
                    while(ancestor > 0)
                    {
                        ancestor = (ancestor-1)/NUM_PARTITIONS_J;
                        origninalWorker[ancestor] = iWorker;
                    }
                }
            }
            tmpIndexStart += nRegionsPerWorker+1;
            tmpIndexEnd += nRegionsPerWorker+1;
        }
        for(int iWorker = nWorkersWithAdditionalRegion; iWorker < MPI_SIZE; iWorker++)
        {
            tmpIndexEnd =  tmpIndexStart + nRegionsPerWorker;
            for(unsigned long index = tmpIndexStart; index < tmpIndexEnd; index++)
            {
                unsigned long ancestor = (index-1)/NUM_PARTITIONS_J;
                WORKERS_FOR_EACH_REGION[ancestor].insert(iWorker);
                if(DYNAMIC_SCHEDULE_FLAG)
                {
                    origninalWorker[index] = iWorker;
                    origninalWorker[ancestor] = iWorker;
                    while(ancestor > 0)
                    {
                        ancestor = (ancestor-1)/NUM_PARTITIONS_J;
                        origninalWorker[ancestor] = iWorker;
                    }
                }
            }
            tmpIndexStart += nRegionsPerWorker;
            tmpIndexEnd += nRegionsPerWorker;
        }

        //Assign regionStart and regionEnd for each level
        REGION_START = new unsigned long [NUM_LEVELS_M];
        REGION_END = new unsigned long [NUM_LEVELS_M];
        REGION_START[NUM_LEVELS_M-1]=indexRegionAtFinestLevelStartThisWorker;
        REGION_END[NUM_LEVELS_M-1]=indexRegionAtFinestLevelEndThisWorker;
        for(int iLevel = NUM_LEVELS_M-2; iLevel > -1; iLevel--)
        {
            REGION_START[iLevel] = (REGION_START[iLevel+1]-1)/NUM_PARTITIONS_J;
            REGION_END[iLevel] = (REGION_END[iLevel+1]-1)/NUM_PARTITIONS_J;
        }
}

/*
Implement the member function of Partition, build_partition, which builds the hierarchical partition
*/

void Partition::build_partition()
{
    if(WORKER == 0)
	{
		gettimeofday(&timeNow, NULL);
		cout<<"===========================> Processor 1: building hierarchical grid starts. Elapsed time: "<<(double)timeNow.tv_sec-(double)timeBegin.tv_sec+((double)timeNow.tv_usec-(double)timeBegin.tv_usec)/1000000.0<<" seconds.\n";
	}

    //The number of knots in each direction
    int nKnotsX = ceil(sqrt(NUM_KNOTS_r));
    int nKnotsY = (int)(NUM_KNOTS_r/nKnotsX);
    nKnots = nKnotsX*nKnotsY;

    //Allocate memory for the coordinates of knots
    knotsX = new double* [nRegionsInTotal];
    knotsY = new double* [nRegionsInTotal];

    //Array of coordinates for the partition in each region with dimension [nRegionsInTotal]
    double *partitionXMin, *partitionYMin, *partitionXMax, *partitionYMax;
    partitionXMin = new double [nRegionsInTotal];
    partitionXMax = new double [nRegionsInTotal];
    partitionYMin = new double [nRegionsInTotal];
    partitionYMax = new double [nRegionsInTotal];

    partitionXMin[0] = data->domainBoundaries[0];
    partitionXMax[0] = data->domainBoundaries[1];
    partitionYMin[0] = data->domainBoundaries[2];
    partitionYMax[0] = data->domainBoundaries[3];

    for(int iLevel = 0; iLevel < NUM_LEVELS_M-1; iLevel++)
    {
        #pragma omp parallel for
        for(unsigned long jRegion = REGION_START[iLevel]; jRegion < REGION_END[iLevel] + 1; jRegion++)
            create_partitions(jRegion, partitionXMin, partitionXMax, partitionYMin, partitionYMax);
    }

    #pragma omp parallel for schedule(dynamic,1)
    for(unsigned long iRegion = 0; iRegion < nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
    {
        if(WORKING_REGION_FLAG[iRegion])
        {
            knotsX[iRegion] = new double [nKnots];
            knotsY[iRegion] = new double [nKnots];
            create_knots(knotsX[iRegion], knotsY[iRegion], partitionXMin[iRegion], partitionXMax[iRegion], nKnotsX, partitionYMin[iRegion], partitionYMax[iRegion], nKnotsY);
        }
    }

    //Allocate memory for nKnotsAtFinestLevel, nPredictionsAtFinestLevel
    nKnotsAtFinestLevel = new unsigned long [nRegionsAtEachLevel[NUM_LEVELS_M-1]];
    nPredictionsAtFinestLevel = new unsigned long [nRegionsAtEachLevel[NUM_LEVELS_M-1]];

    //Allocate memory for nKnotsAtFinestLevel, nPredictionsAtFinestLevel
    knotsResidual = new double* [nRegionsAtEachLevel[NUM_LEVELS_M-1]];

    //Allocate memory for the coordinates of predictions
    predictionX = new double* [nRegionsAtEachLevel[NUM_LEVELS_M-1]];
    predictionY = new double* [nRegionsAtEachLevel[NUM_LEVELS_M-1]];

    //Get coordinates of the knots and predictions at the regions in the finest level
    int maxOpenMPThreads=omp_get_max_threads();

    #pragma omp parallel for num_threads(maxOpenMPThreads) schedule(dynamic,1)
    for(unsigned long iRegion = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion < nRegionsInTotal; iRegion++)
    {
        std::set<unsigned long> *indexObservationsInThisRegion = new std::set<unsigned long>;
        std::set<unsigned long> *indexPredictionsInThisRegion = new std::set<unsigned long>;

        unsigned long indexStartThisLevel = iRegion+nRegionsAtEachLevel[NUM_LEVELS_M-1]-nRegionsInTotal;
        if(WORKING_REGION_FLAG[iRegion])
        {
            double xMin=partitionXMin[iRegion];
            double xMax=partitionXMax[iRegion];
            double yMin=partitionYMin[iRegion];
            double yMax=partitionYMax[iRegion];

            indexObservationsInThisRegion->clear();

            for(unsigned long jObservation = 0; jObservation < NUM_OBSERVATIONS; jObservation++)
            {
                if(data->observationLon[jObservation] >= xMin && data->observationLon[jObservation] < xMax && data->observationLat[jObservation] >= yMin && data->observationLat[jObservation] < yMax)
                {

                    bool skip = false;
                    //Skip if the observation location coincides with knots in levels above
                    unsigned long ancestor = iRegion;
                    for(int kLevel = 0; kLevel < NUM_LEVELS_M-1; kLevel++)
                    {
                        ancestor = (ancestor-1)/NUM_PARTITIONS_J;
                        for(int pKnot = 0; pKnot < nKnots; pKnot++)
                        {
                            if(knotsX[ancestor][pKnot] == data->observationLon[jObservation] && knotsY[ancestor][pKnot] == data->observationLat[jObservation])
                            {
                                skip = true;
                                break;
                            }
                        }
                        if(skip) break;
                    }

                    //Add this observation
                    if(!skip) indexObservationsInThisRegion->insert(jObservation);
                }
            }

            nKnotsAtFinestLevel[indexStartThisLevel] = indexObservationsInThisRegion->size();
            knotsX[iRegion] = new double [nKnotsAtFinestLevel[indexStartThisLevel]];
            knotsY[iRegion] = new double [nKnotsAtFinestLevel[indexStartThisLevel]];
            knotsResidual[indexStartThisLevel] = new double [nKnotsAtFinestLevel[indexStartThisLevel]];

            unsigned long tmpIndex=0;
            for(std::set<unsigned long>::iterator jObservation = indexObservationsInThisRegion->begin(); jObservation != indexObservationsInThisRegion->end(); jObservation++)
            {
                knotsX[iRegion][tmpIndex] = data->observationLon[*jObservation];
                knotsY[iRegion][tmpIndex] = data->observationLat[*jObservation];
                knotsResidual[indexStartThisLevel][tmpIndex++] = data->observationResiduals[*jObservation];
            }

            if(CALCULATION_MODE=="prediction")
            {
                indexPredictionsInThisRegion->clear();
                for(unsigned long jPrediction = 0; jPrediction < NUM_PREDICTIONS; jPrediction++)
                {
                    if(data->predictionLon[jPrediction] >= xMin && data->predictionLon[jPrediction] < xMax && data->predictionLat[jPrediction] >= yMin && data->predictionLat[jPrediction] < yMax)
                        indexPredictionsInThisRegion->insert(jPrediction);
                }

                nPredictionsAtFinestLevel[indexStartThisLevel] = indexPredictionsInThisRegion->size();
                predictionX[indexStartThisLevel] = new double [nPredictionsAtFinestLevel[indexStartThisLevel]];
                predictionY[indexStartThisLevel] = new double [nPredictionsAtFinestLevel[indexStartThisLevel]];

                unsigned long tmpIndex=0;
                for(std::set<unsigned long>::iterator jPrediction = indexPredictionsInThisRegion->begin(); jPrediction != indexPredictionsInThisRegion->end(); jPrediction++)
                {
                    predictionX[indexStartThisLevel][tmpIndex] = data->predictionLon[*jPrediction];
                    predictionY[indexStartThisLevel][tmpIndex++] = data->predictionLat[*jPrediction];
                }
            }
        }
        indexObservationsInThisRegion->clear();
        indexPredictionsInThisRegion->clear();
    }

    //Release the memory for the temporary arrays
    delete[] partitionXMin;
    delete[] partitionYMin;
    delete[] partitionXMax;
    delete[] partitionYMax;

    //Delete variables that are not used anymore
    delete[] data->observationLon;
    delete[] data->observationLat;
    delete[] data->observationResiduals;

    if(DYNAMIC_SCHEDULE_FLAG)
		dynamic_schedule();
	else
		WORLD = MPI_COMM_WORLD;

    if(WORKER == 0)
	{
		gettimeofday(&timeNow, NULL);
		cout<<"===========================> Processor 1: building hierarchical grid is complete. Elapsed time: "<<(double)timeNow.tv_sec-(double)timeBegin.tv_sec+((double)timeNow.tv_usec-(double)timeBegin.tv_usec)/1000000.0<<" seconds.\n\n";
		if(PRINT_DETAIL_FLAG) print_partition_summary();
	}
}

//Implement the member function of Partition, dynamic_schedule, which assign the work load for each worker by dynamic scheduling
void Partition::dynamic_schedule()
{
    //Synchronize the knots of regions before the finest level
    for(unsigned long iRegion = 0; iRegion < nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
    {
        if(!WORKING_REGION_FLAG[iRegion])
        {
            knotsX[iRegion] = new double [nKnots];
            knotsY[iRegion] = new double [nKnots];
        }
    }

    for(unsigned long iRegion = 0; iRegion < nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
    {
        MPI_Bcast(knotsX[iRegion],nKnots,MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
        MPI_Bcast(knotsY[iRegion],nKnots,MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
    }

    //Synchronize the knots of regions at the finest level and predictions if predicting
    unsigned long indexStartFinestLevel = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1];
    for(unsigned long iRegion = indexStartFinestLevel; iRegion < nRegionsInTotal; iRegion++)
    {

        MPI_Bcast(&nKnotsAtFinestLevel[iRegion-indexStartFinestLevel],1,MPI_UNSIGNED_LONG,origninalWorker[iRegion],MPI_COMM_WORLD);
        if(!WORKING_REGION_FLAG[iRegion])
        {
            knotsX[iRegion] = new double [nKnotsAtFinestLevel[iRegion-indexStartFinestLevel]];
            knotsY[iRegion] = new double [nKnotsAtFinestLevel[iRegion-indexStartFinestLevel]];
            knotsResidual[iRegion-indexStartFinestLevel] = new double [nKnotsAtFinestLevel[iRegion-indexStartFinestLevel]];
        }

        if(CALCULATION_MODE=="prediction")
        {
            MPI_Bcast(&nPredictionsAtFinestLevel[iRegion-indexStartFinestLevel],1,MPI_UNSIGNED_LONG,origninalWorker[iRegion],MPI_COMM_WORLD);
            if(!WORKING_REGION_FLAG[iRegion])
            {
                predictionX[iRegion-indexStartFinestLevel] = new double [nPredictionsAtFinestLevel[iRegion-indexStartFinestLevel]];
                predictionY[iRegion-indexStartFinestLevel] = new double [nPredictionsAtFinestLevel[iRegion-indexStartFinestLevel]];
            }
        }

        MPI_Bcast(knotsX[iRegion],nKnotsAtFinestLevel[iRegion-indexStartFinestLevel],MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
        MPI_Bcast(knotsY[iRegion],nKnotsAtFinestLevel[iRegion-indexStartFinestLevel],MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
        MPI_Bcast(knotsResidual[iRegion-indexStartFinestLevel],nKnotsAtFinestLevel[iRegion-indexStartFinestLevel],MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);

        if(CALCULATION_MODE=="prediction")
        {
            MPI_Bcast(predictionX[iRegion-indexStartFinestLevel],nPredictionsAtFinestLevel[iRegion-indexStartFinestLevel],MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
            MPI_Bcast(predictionY[iRegion-indexStartFinestLevel],nPredictionsAtFinestLevel[iRegion-indexStartFinestLevel],MPI_DOUBLE,origninalWorker[iRegion],MPI_COMM_WORLD);
        }
    }

    //Dynamic schedule work load for each worker

    ///Clear WORKERS_FOR_EACH_REGION at the second finest level
    for(unsigned long iRegion = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1] - nRegionsAtEachLevel[NUM_LEVELS_M-2]; iRegion < nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
        WORKERS_FOR_EACH_REGION[iRegion].clear();

    ///Clear WORKING_REGION_FLAG for all regions
    memset(WORKING_REGION_FLAG,false,nRegionsInTotal);

    ///Clear INDICES_REGIONS_AT_CURRENT_LEVEL
    INDICES_REGIONS_AT_CURRENT_LEVEL.clear();

    ///Calculate work load
    double totalComplexity = 0;
    for(unsigned long index = 0; index < nRegionsAtEachLevel[NUM_LEVELS_M-1]; index++)
    {
        double tmp = nKnotsAtFinestLevel[index];
        totalComplexity += tmp*tmp;//*tmp;
    }
    double complexityEachWorker = totalComplexity/MPI_SIZE;

    ///Assign working regions for each worker
    REGION_START[NUM_LEVELS_M-1]=nRegionsInTotal;
    REGION_END[NUM_LEVELS_M-1]=0;
    int thisWorker=-1;
    double complexityThisWorker=0;
    unsigned long indexRegionAtFinestLevelStartThisWorker = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1], indexRegionAtFinestLevelEndThisWorker;
    for(unsigned long index = 0; index < nRegionsAtEachLevel[NUM_LEVELS_M-1]; index++)
    {
        double tmp = nKnotsAtFinestLevel[index];
        double complexityThisRegion = tmp*tmp;//*tmp;
        complexityThisWorker += complexityThisRegion;
        //if(WORKER==0) cout<<complexityThisRegion<<" "<<complexityEachWorker<<endl;
        if(complexityThisWorker >= complexityEachWorker || index == nRegionsAtEachLevel[NUM_LEVELS_M-1]-1)
        {
            thisWorker++;
            if(thisWorker == MPI_SIZE) thisWorker = MPI_SIZE-1;

            indexRegionAtFinestLevelEndThisWorker = indexStartFinestLevel+index;

            //Check whether this region should be handled by this worker or next worker
            if(complexityEachWorker - complexityThisWorker + complexityThisRegion < complexityThisWorker - complexityEachWorker
                && index != nRegionsAtEachLevel[NUM_LEVELS_M-1]-1
                && indexRegionAtFinestLevelEndThisWorker > indexRegionAtFinestLevelStartThisWorker)
            {
                //Next worker handle this region
                indexRegionAtFinestLevelEndThisWorker--;
                complexityThisWorker = complexityThisRegion;
            }
            else
            {
                //This worker handle this region
                complexityThisWorker = 0;
            }

            //if(WORKER==0) cout<<thisWorker<<" "<<indexRegionAtFinestLevelStartThisWorker<<" "<<indexRegionAtFinestLevelEndThisWorker<<endl;

            for(unsigned long jRegion = indexRegionAtFinestLevelStartThisWorker; jRegion < indexRegionAtFinestLevelEndThisWorker+1; jRegion++)
            {
                unsigned long ancestor = (jRegion-1)/NUM_PARTITIONS_J;
                WORKERS_FOR_EACH_REGION[ancestor].insert(thisWorker);
            }

            if(thisWorker == WORKER)
            {
                for(unsigned long jRegion = indexRegionAtFinestLevelStartThisWorker; jRegion < indexRegionAtFinestLevelEndThisWorker+1; jRegion++)
                {
                    unsigned long ancestor = jRegion;
                    WORKING_REGION_FLAG[ancestor] = true;
                    while(ancestor > 0)
                    {
                        ancestor = (ancestor-1)/NUM_PARTITIONS_J;
                        WORKING_REGION_FLAG[ancestor] = true;
                    }
                    INDICES_REGIONS_AT_CURRENT_LEVEL.insert((jRegion-1)/NUM_PARTITIONS_J);
                }
                if(indexRegionAtFinestLevelStartThisWorker<REGION_START[NUM_LEVELS_M-1]) REGION_START[NUM_LEVELS_M-1]=indexRegionAtFinestLevelStartThisWorker;
                if(indexRegionAtFinestLevelEndThisWorker>REGION_END[NUM_LEVELS_M-1]) REGION_END[NUM_LEVELS_M-1]=indexRegionAtFinestLevelEndThisWorker;
            }

            indexRegionAtFinestLevelStartThisWorker = indexRegionAtFinestLevelEndThisWorker+1;
        }
    }

    //Deactivate MPI processes that are not assigned any regions, which rarely happens
    if(thisWorker<MPI_SIZE-1)
    {
        MPI_Group worldGroup;
        MPI_Comm_group(MPI_COMM_WORLD,&worldGroup);

        MPI_Group newWorldGroup;
        int ranges[1][3]={{thisWorker+1,MPI_SIZE-1,1}};
        MPI_Group_range_excl(worldGroup,1,ranges,&newWorldGroup);

        MPI_Comm_create(MPI_COMM_WORLD, newWorldGroup, &WORLD);

        MPI_SIZE=thisWorker+1;

        if(WORKER>thisWorker)
        {
            MPI_Finalize();
            exit(EXIT_SUCCESS);
        }
    }
    else
        WORLD = MPI_COMM_WORLD;

    for(int iLevel = NUM_LEVELS_M-2; iLevel > -1; iLevel--)
    {
        REGION_START[iLevel] = (REGION_START[iLevel+1]-1)/NUM_PARTITIONS_J;
        REGION_END[iLevel] = (REGION_END[iLevel+1]-1)/NUM_PARTITIONS_J;
    }

    //Delete regions that are not handled by this WORKER
    for(unsigned long iRegion = 0; iRegion < nRegionsInTotal; iRegion++)
        if(!WORKING_REGION_FLAG[iRegion])
        {
            delete[] knotsX[iRegion]; knotsX[iRegion]=NULL;
            delete[] knotsY[iRegion]; knotsY[iRegion]=NULL;
        }

    for(unsigned long iRegion = nRegionsInTotal - nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion < nRegionsInTotal; iRegion++)
        if(!WORKING_REGION_FLAG[iRegion])
        {
            delete[] knotsResidual[iRegion-indexStartFinestLevel]; knotsResidual[iRegion-indexStartFinestLevel]=NULL;

            if(CALCULATION_MODE=="prediction")
            {
                delete[] predictionX[iRegion-indexStartFinestLevel]; predictionX[iRegion-indexStartFinestLevel]=NULL;
                delete[] predictionY[iRegion-indexStartFinestLevel]; predictionY[iRegion-indexStartFinestLevel]=NULL;
            }
        }

}

//Implement the member function of Partition, print_partition_summary, which shows a brief summary of the partition
void Partition::print_partition_summary()
{
    cout<<">>The number of levels: "<<NUM_LEVELS_M<<endl<<endl;
    cout<<">>The total number of regions: "<<nRegionsInTotal<<endl<<endl;
    cout<<">>The number of regions in each level: \n    ";
    for(int iLevel = 0; iLevel < NUM_LEVELS_M; iLevel++) cout<<nRegionsAtEachLevel[iLevel]<<" ";
    cout<<endl<<endl;

    cout<<">>Knot coordinates in the coarsest region, up to ten knots are shown: \n    longitude: ";
    for(int iKnot = 0; iKnot < std::min(nKnots,10); iKnot++)
        printf("%8.2lf ",knotsX[0][iKnot]);
    cout<<"\n    latitude:  ";
    for(int iKnot = 0; iKnot < std::min(nKnots,10); iKnot++)
        printf("%8.2lf ",knotsY[0][iKnot]);
    cout<<endl<<endl;

    if(CALCULATION_MODE=="predict")
        cout<<">>The number of predictions after eliminating duplicates: "<<NUM_PREDICTIONS<<endl<<endl;

}

void Partition::dump_structure_information()
{
    //Open the file "structure_information.txt" to store structure information
    std::ofstream file;
    file.open("structure_information.txt");

    //Dump the number of levels
    file<<"The number of levels: "<<NUM_LEVELS_M<<endl<<endl;

    //Dump the total number of regions
    file<<"The total number of regions: "<<nRegionsInTotal<<endl<<endl;

    //Dump the number of knots in regions before the finest level
    file<<"The actual number of knots in regions before the finest level: "<<nKnots<<endl<<endl;

    //Dump some statistics of the number of knots in each region at the finest level
    unsigned long nKnotsMin = nKnotsAtFinestLevel[0];
    unsigned long nKnotsMax = nKnotsAtFinestLevel[0];
    unsigned long nRegionsHaveZeroKnots = 0;

    for(unsigned long iRegion = 0; iRegion < nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
    {
        if(nKnotsMin>nKnotsAtFinestLevel[iRegion]) nKnotsMin=nKnotsAtFinestLevel[iRegion];
        if(nKnotsMax<nKnotsAtFinestLevel[iRegion]) nKnotsMax=nKnotsAtFinestLevel[iRegion];
        if(nKnotsAtFinestLevel[iRegion] == 0) nRegionsHaveZeroKnots++;
    }

    file<<"The minimal number of knots in regions at the finest level: "<<nKnotsMin<<endl<<endl;
    file<<"The maximal number of knots in regions at the finest level: "<<nKnotsMax<<endl<<endl;
    file<<"The number of in regions at the finest level that have zero knots: "<<nRegionsHaveZeroKnots<<endl<<endl;

    unsigned long thresholdLeft, thresholdRight;
    unsigned long nRegionsBetweenThresholds;
    double interval = (nKnotsMax-nKnotsMin)/10;

    thresholdLeft=nKnotsMin;
    file<<"The number of regions at the finest level with number of knots in the following intervals:"<<endl;

    for(int i = 0; i < 10; i++)
    {
        nRegionsBetweenThresholds = 0;
        thresholdRight=thresholdLeft+interval;
        if(i == 9) thresholdRight = nKnotsMax;
        for(unsigned long iRegion = 0; iRegion < nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
           if(thresholdLeft<=nKnotsAtFinestLevel[iRegion] && nKnotsAtFinestLevel[iRegion]<=thresholdRight) nRegionsBetweenThresholds++;
        file<<"[ "<<thresholdLeft<<" , "<<thresholdRight<<" ]: "<<nRegionsBetweenThresholds<<endl;
        thresholdLeft = thresholdRight+1;
    }
    file<<endl;

    //Dump the number of regions in each level
    file<<"The number of regions in each level: \n";
    for(int iLevel = 0; iLevel < NUM_LEVELS_M; iLevel++)
        file<<"Level "<<iLevel+1<<": "<<nRegionsAtEachLevel[iLevel]<<endl;
    file<<endl;

    //Dump the number of knots in each region at the finest level
    file<<"The number of knots in each region at the finest level: \n";
    for(unsigned long iRegion = 0; iRegion < nRegionsAtEachLevel[NUM_LEVELS_M-1]; iRegion++)
        file<<"Region "<<nRegionsInTotal-nRegionsAtEachLevel[NUM_LEVELS_M-1]+iRegion<<": "<<nKnotsAtFinestLevel[iRegion]<<endl;

    //Close the file
    file.close();
}

Partition::~Partition()
{
    delete[] nRegionsAtEachLevel;

    delete[] knotsX;
    delete[] knotsY;
    delete[] knotsResidual;

    delete[] predictionX,
    delete[] predictionY;

    delete[] nKnotsAtFinestLevel,
    delete[] nPredictionsAtFinestLevel;

    if(DYNAMIC_SCHEDULE_FLAG) delete[] origninalWorker;
    //cout<<"Partition is deleted\n";
}


/**
 * @file methods/ann/rnn_impl.hpp
 * @author Marcus Edel
 *
 * Definition of the RNN class, which implements recurrent neural networks.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_METHODS_ANN_RNN_IMPL_HPP
#define MLPACK_METHODS_ANN_RNN_IMPL_HPP

// In case it hasn't been included yet.
#include "rnn.hpp"

#include "visitor/load_output_parameter_visitor.hpp"
#include "visitor/save_output_parameter_visitor.hpp"
#include "visitor/forward_visitor.hpp"
#include "visitor/backward_visitor.hpp"
#include "visitor/reset_cell_visitor.hpp"
#include "visitor/deterministic_set_visitor.hpp"
#include "visitor/gradient_set_visitor.hpp"
#include "visitor/gradient_visitor.hpp"
#include "visitor/weight_set_visitor.hpp"

#include <boost/serialization/variant.hpp>

namespace mlpack {
namespace ann /** Artificial Neural Network. */ {

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::RNN(
    const size_t rho,
    const bool single,
    OutputLayerType outputLayer,
    InitializationRuleType initializeRule) :
    rho(rho),
    outputLayer(std::move(outputLayer)),
    initializeRule(std::move(initializeRule)),
    inputSize(0),
    outputSize(0),
    targetSize(0),
    reset(false),
    single(single),
    numFunctions(0),
    deterministic(true)
{
  /* Nothing to do here */
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::~RNN()
{
  for (LayerTypes<CustomLayers...>& layer : network)
  {
    boost::apply_visitor(deleteVisitor, layer);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename OptimizerType>
typename std::enable_if<
      HasMaxIterations<OptimizerType, size_t&(OptimizerType::*)()>
      ::value, void>::type
RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::
WarnMessageMaxIterations(OptimizerType& optimizer, size_t samples) const
{
  if (optimizer.MaxIterations() < samples &&
      optimizer.MaxIterations() != 0)
  {
    Log::Warn << "The optimizer's maximum number of iterations "
              << "is less than the size of the dataset; the "
              << "optimizer will not pass over the entire "
              << "dataset. To fix this, modify the maximum "
              << "number of iterations to be at least equal "
              << "to the number of points of your dataset "
              << "(" << samples << ")." << std::endl;
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename OptimizerType>
typename std::enable_if<
      !HasMaxIterations<OptimizerType, size_t&(OptimizerType::*)()>
      ::value, void>::type
RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::
WarnMessageMaxIterations(OptimizerType& /* optimizer */,
                         size_t /* samples */) const
{
  return;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename OptimizerType, typename... CallbackTypes>
double RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Train(
    arma::cube predictors,
    arma::cube responses,
    OptimizerType& optimizer,
    CallbackTypes&&... callbacks)
{
  numFunctions = responses.n_cols;

  this->predictors = std::move(predictors);
  this->responses = std::move(responses);

  this->deterministic = true;
  ResetDeterministic();

  if (!reset)
  {
    ResetParameters();
  }

  WarnMessageMaxIterations<OptimizerType>(optimizer, this->predictors.n_cols);

  // Train the model.
  Timer::Start("rnn_optimization");
  const double out = optimizer.Optimize(*this, parameter, callbacks...);
  Timer::Stop("rnn_optimization");

  Log::Info << "RNN::RNN(): final objective of trained model is " << out
      << "." << std::endl;
  return out;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::ResetCells()
{
  for (size_t i = 1; i < network.size(); ++i)
  {
    boost::apply_visitor(ResetCellVisitor(rho), network[i]);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename OptimizerType, typename... CallbackTypes>
double RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Train(
    arma::cube predictors,
    arma::cube responses,
    CallbackTypes&&... callbacks)
{
  numFunctions = responses.n_cols;

  this->predictors = std::move(predictors);
  this->responses = std::move(responses);

  this->deterministic = true;
  ResetDeterministic();

  if (!reset)
  {
    ResetParameters();
  }

  OptimizerType optimizer;

  WarnMessageMaxIterations<OptimizerType>(optimizer, this->predictors.n_cols);

  // Train the model.
  Timer::Start("rnn_optimization");
  const double out = optimizer.Optimize(*this, parameter, callbacks...);
  Timer::Stop("rnn_optimization");

  Log::Info << "RNN::RNN(): final objective of trained model is " << out
      << "." << std::endl;
  return out;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Predict(
    arma::cube predictors, arma::cube& results, const size_t batchSize)
{
  ResetCells();

  if (parameter.is_empty())
  {
    ResetParameters();
  }

  if (!deterministic)
  {
    deterministic = true;
    ResetDeterministic();
  }

  const size_t effectiveBatchSize = std::min(batchSize,
      size_t(predictors.n_cols));

  Forward(arma::mat(predictors.slice(0).colptr(0), predictors.n_rows,
      effectiveBatchSize, false, true));
  arma::mat resultsTemp = boost::apply_visitor(outputParameterVisitor,
      network.back());

  outputSize = resultsTemp.n_rows;
  results = arma::zeros<arma::cube>(outputSize, predictors.n_cols, rho);
  results.slice(0).submat(0, 0, results.n_rows - 1,
      effectiveBatchSize - 1) = resultsTemp;

  // Process in accordance with the given batch size.
  for (size_t begin = 0; begin < predictors.n_cols; begin += batchSize)
  {
    const size_t effectiveBatchSize = std::min(batchSize,
        size_t(predictors.n_cols - begin));
    for (size_t seqNum = !begin; seqNum < rho; ++seqNum)
    {
      Forward(arma::mat(predictors.slice(seqNum).colptr(begin),
          predictors.n_rows, effectiveBatchSize, false, true));

      results.slice(seqNum).submat(0, begin, results.n_rows - 1, begin +
          effectiveBatchSize - 1) = boost::apply_visitor(outputParameterVisitor,
          network.back());
    }
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
double RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Evaluate(
    const arma::mat& /* parameters */,
    const size_t begin,
    const size_t batchSize,
    const bool deterministic)
{
  if (parameter.is_empty())
  {
    ResetParameters();
  }

  if (deterministic != this->deterministic)
  {
    this->deterministic = deterministic;
    ResetDeterministic();
  }

  if (!inputSize)
  {
    inputSize = predictors.n_rows;
    targetSize = responses.n_rows;
  }
  else if (targetSize == 0)
  {
    targetSize = responses.n_rows;
  }

  ResetCells();

  double performance = 0;
  size_t responseSeq = 0;

  for (size_t seqNum = 0; seqNum < rho; ++seqNum)
  {
    // Wrap a matrix around our data to avoid a copy.
    arma::mat stepData(predictors.slice(seqNum).colptr(begin),
        predictors.n_rows, batchSize, false, true);
    Forward(stepData);
    if (!single)
    {
      responseSeq = seqNum;
    }

    performance += outputLayer.Forward(boost::apply_visitor(
        outputParameterVisitor, network.back()),
        arma::mat(responses.slice(responseSeq).colptr(begin),
            responses.n_rows, batchSize, false, true));
  }

  if (outputSize == 0)
  {
    outputSize = boost::apply_visitor(outputParameterVisitor,
        network.back()).n_elem / batchSize;
  }

  return performance;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
double RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Evaluate(
    const arma::mat& parameters,
    const size_t begin,
    const size_t batchSize)
{
  return Evaluate(parameters, begin, batchSize, true);
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename GradType>
double RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::
EvaluateWithGradient(const arma::mat& /* parameters */,
                     const size_t begin,
                     GradType& gradient,
                     const size_t batchSize)
{
  // Initialize passed gradient.
  if (gradient.is_empty())
  {
    if (parameter.is_empty())
    {
      ResetParameters();
    }

    gradient = arma::zeros<arma::mat>(parameter.n_rows, parameter.n_cols);
  }
  else
  {
    gradient.zeros();
  }

  if (this->deterministic)
  {
    this->deterministic = false;
    ResetDeterministic();
  }

  if (!inputSize)
  {
    inputSize = predictors.n_rows;
    targetSize = responses.n_rows;
  }
  else if (targetSize == 0)
  {
    targetSize = responses.n_rows;
  }

  ResetCells();

  double performance = 0;
  size_t responseSeq = 0;
  const size_t effectiveRho = std::min(rho, size_t(responses.size()));

  for (size_t seqNum = 0; seqNum < effectiveRho; ++seqNum)
  {
    // Wrap a matrix around our data to avoid a copy.
    arma::mat stepData(predictors.slice(seqNum).colptr(begin),
        predictors.n_rows, batchSize, false, true);
    Forward(stepData);
    if (!single)
    {
      responseSeq = seqNum;
    }

    for (size_t l = 0; l < network.size(); ++l)
    {
      boost::apply_visitor(SaveOutputParameterVisitor(moduleOutputParameter),
          network[l]);
    }

    performance += outputLayer.Forward(boost::apply_visitor(
        outputParameterVisitor, network.back()),
        arma::mat(responses.slice(responseSeq).colptr(begin),
            responses.n_rows, batchSize, false, true));
  }

  if (outputSize == 0)
  {
    outputSize = boost::apply_visitor(outputParameterVisitor,
        network.back()).n_elem / batchSize;
  }

  // Initialize current/working gradient.
  if (currentGradient.is_empty())
  {
    currentGradient = arma::zeros<arma::mat>(parameter.n_rows,
        parameter.n_cols);
  }

  ResetGradients(currentGradient);

  for (size_t seqNum = 0; seqNum < effectiveRho; ++seqNum)
  {
    currentGradient.zeros();
    for (size_t l = 0; l < network.size(); ++l)
    {
      boost::apply_visitor(LoadOutputParameterVisitor(moduleOutputParameter),
          network[network.size() - 1 - l]);
    }

    if (single && seqNum > 0)
    {
      error.zeros();
    }
    else if (single && seqNum == 0)
    {
      outputLayer.Backward(boost::apply_visitor(
          outputParameterVisitor, network.back()),
          arma::mat(responses.slice(0).colptr(begin),
          responses.n_rows, batchSize, false, true), error);
    }
    else
    {
      outputLayer.Backward(boost::apply_visitor(
          outputParameterVisitor, network.back()),
          arma::mat(responses.slice(effectiveRho - seqNum - 1).colptr(begin),
          responses.n_rows, batchSize, false, true), error);
    }

    Backward();
    Gradient(
        arma::mat(predictors.slice(effectiveRho - seqNum - 1).colptr(begin),
        predictors.n_rows, batchSize, false, true));
    gradient += currentGradient;
  }

  return performance;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Gradient(
    const arma::mat& parameters,
    const size_t begin,
    arma::mat& gradient,
    const size_t batchSize)
{
  this->EvaluateWithGradient(parameters, begin, gradient, batchSize);
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Shuffle()
{
  arma::cube newPredictors, newResponses;
  math::ShuffleData(predictors, responses, newPredictors, newResponses);

  predictors = std::move(newPredictors);
  responses = std::move(newResponses);
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::ResetParameters()
{
  ResetDeterministic();

  // Reset the network parameter with the given initialization rule.
  NetworkInitialization<InitializationRuleType,
                        CustomLayers...> networkInit(initializeRule);
  networkInit.Initialize(network, parameter);

  reset = true;
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Reset()
{
  ResetParameters();
  ResetCells();
  currentGradient.zeros();
  ResetGradients(currentGradient);
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::ResetDeterministic()
{
  DeterministicSetVisitor deterministicSetVisitor(deterministic);
  std::for_each(network.begin(), network.end(),
      boost::apply_visitor(deterministicSetVisitor));
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::ResetGradients(
    arma::mat& gradient)
{
  size_t offset = 0;
  for (LayerTypes<CustomLayers...>& layer : network)
  {
    offset += boost::apply_visitor(GradientSetVisitor(gradient, offset), layer);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename InputType>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::Forward(const InputType& input)
{
  boost::apply_visitor(ForwardVisitor(input,
      boost::apply_visitor(outputParameterVisitor, network.front())),
      network.front());

  for (size_t i = 1; i < network.size(); ++i)
  {
    boost::apply_visitor(ForwardVisitor(
        boost::apply_visitor(outputParameterVisitor, network[i - 1]),
        boost::apply_visitor(outputParameterVisitor, network[i])),
        network[i]);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::Backward()
{
  boost::apply_visitor(BackwardVisitor(
        boost::apply_visitor(outputParameterVisitor, network.back()),
        error, boost::apply_visitor(deltaVisitor,
        network.back())), network.back());

  for (size_t i = 2; i < network.size(); ++i)
  {
    boost::apply_visitor(BackwardVisitor(
        boost::apply_visitor(outputParameterVisitor,
        network[network.size() - i]), boost::apply_visitor(
        deltaVisitor, network[network.size() - i + 1]),
        boost::apply_visitor(deltaVisitor, network[network.size() - i])),
        network[network.size() - i]);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename InputType>
void RNN<OutputLayerType, InitializationRuleType,
         CustomLayers...>::Gradient(const InputType& input)
{
  boost::apply_visitor(GradientVisitor(input,
      boost::apply_visitor(deltaVisitor, network[1])), network.front());

  for (size_t i = 1; i < network.size() - 1; ++i)
  {
    boost::apply_visitor(GradientVisitor(
        boost::apply_visitor(outputParameterVisitor, network[i - 1]),
        boost::apply_visitor(deltaVisitor, network[i + 1])),
        network[i]);
  }
}

template<typename OutputLayerType, typename InitializationRuleType,
         typename... CustomLayers>
template<typename Archive>
void RNN<OutputLayerType, InitializationRuleType, CustomLayers...>::serialize(
    Archive& ar, const unsigned int version)
{
  ar & BOOST_SERIALIZATION_NVP(parameter);
  ar & BOOST_SERIALIZATION_NVP(rho);
  ar & BOOST_SERIALIZATION_NVP(single);
  ar & BOOST_SERIALIZATION_NVP(inputSize);
  ar & BOOST_SERIALIZATION_NVP(outputSize);
  ar & BOOST_SERIALIZATION_NVP(targetSize);

  // Earlier versions of the RNN code did not serialize the 'reset' variable.
  if (version > 0)
  {
    ar & BOOST_SERIALIZATION_NVP(reset);
  }

  if (Archive::is_loading::value)
  {
    std::for_each(network.begin(), network.end(),
        boost::apply_visitor(deleteVisitor));
    network.clear();
  }

  ar & BOOST_SERIALIZATION_NVP(network);

  // If we are loading, we need to initialize the weights.
  if (Archive::is_loading::value)
  {
    // Earlier versions of the RNN code assumed that the weights needed to be
    // reset on load.
    if (version == 0)
      reset = false;

    size_t offset = 0;
    for (LayerTypes<CustomLayers...>& layer : network)
    {
      offset += boost::apply_visitor(WeightSetVisitor(parameter, offset),
          layer);

      boost::apply_visitor(resetVisitor, layer);
    }

    deterministic = true;
    ResetDeterministic();
  }
}

} // namespace ann
} // namespace mlpack

#endif


// This file is part of libigl, a simple c++ geometry processing library.
//
// Copyright (C) 2013 Alec Jacobson <alecjacobson@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla Public License
// v. 2.0. If a copy of the MPL was not distributed with this file, You can
// obtain one at http://mozilla.org/MPL/2.0/.
#include "arap_linear_block.h"
#include "verbose.h"
#include "cotmatrix_entries.h"
#include <Eigen/Dense>

template <typename MatV, typename MatF, typename Scalar>
IGL_INLINE void igl::arap_linear_block(
  const MatV & V,
  const MatF & F,
  const int d,
  const igl::ARAPEnergyType energy,
  Eigen::SparseMatrix<Scalar> & Kd)
{
  switch(energy)
  {
    case ARAP_ENERGY_TYPE_SPOKES:
      return igl::arap_linear_block_spokes(V,F,d,Kd);
      break;
    case ARAP_ENERGY_TYPE_SPOKES_AND_RIMS:
      return igl::arap_linear_block_spokes_and_rims(V,F,d,Kd);
      break;
    case ARAP_ENERGY_TYPE_ELEMENTS:
      return igl::arap_linear_block_elements(V,F,d,Kd);
      break;
    default:
      verbose("Unsupported energy type: %d\n",energy);
      assert(false);
  }
}

template <typename MatV, typename MatF, typename Scalar>
IGL_INLINE void igl::arap_linear_block_spokes(
  const MatV & V,
  const MatF & F,
  const int d,
  Eigen::SparseMatrix<Scalar> & Kd)
{
  using namespace std;
  using namespace Eigen;
  // simplex size (3: triangles, 4: tetrahedra)
  int simplex_size = F.cols();
  // Number of elements
  int m = F.rows();
  // Temporary output
  Matrix<int,Dynamic,2> edges;
  Kd.resize(V.rows(), V.rows());
  vector<Triplet<Scalar> > Kd_IJV;
  if(simplex_size == 3)
  {
    // triangles
    Kd.reserve(7*V.rows());
    Kd_IJV.reserve(7*V.rows());
    edges.resize(3,2);
    edges <<
      1,2,
      2,0,
      0,1;
  }else if(simplex_size == 4)
  {
    // tets
    Kd.reserve(17*V.rows());
    Kd_IJV.reserve(17*V.rows());
    edges.resize(6,2);
    edges <<
      1,2,
      2,0,
      0,1,
      3,0,
      3,1,
      3,2;
  }
  // gather cotangent weights
  Matrix<Scalar,Dynamic,Dynamic> C;
  cotmatrix_entries(V,F,C);
  // should have weights for each edge
  assert(C.cols() == edges.rows());
  // loop over elements
  for(int i = 0;i<m;i++)
  {
    // loop over edges of element
    for(int e = 0;e<edges.rows();e++)
    {
      int source = F(i,edges(e,0));
      int dest = F(i,edges(e,1));
      double v = 0.5*C(i,e)*(V(source,d)-V(dest,d));
      Kd_IJV.push_back(Triplet<Scalar>(source,dest,v));
      Kd_IJV.push_back(Triplet<Scalar>(dest,source,-v));
      Kd_IJV.push_back(Triplet<Scalar>(source,source,v));
      Kd_IJV.push_back(Triplet<Scalar>(dest,dest,-v));
    }
  }
  Kd.setFromTriplets(Kd_IJV.begin(),Kd_IJV.end());
  Kd.makeCompressed();
}

template <typename MatV, typename MatF, typename Scalar>
IGL_INLINE void igl::arap_linear_block_spokes_and_rims(
  const MatV & V,
  const MatF & F,
  const int d,
  Eigen::SparseMatrix<Scalar> & Kd)
{
  using namespace std;
  using namespace Eigen;
  // simplex size (3: triangles, 4: tetrahedra)
  int simplex_size = F.cols();
  // Number of elements
  int m = F.rows();
  // Temporary output
  Kd.resize(V.rows(), V.rows());
  vector<Triplet<Scalar> > Kd_IJV;
  Matrix<int,Dynamic,2> edges;
  if(simplex_size == 3)
  {
    // triangles
    Kd.reserve(7*V.rows());
    Kd_IJV.reserve(7*V.rows());
    edges.resize(3,2);
    edges <<
      1,2,
      2,0,
      0,1;
  }else if(simplex_size == 4)
  {
    // tets
    Kd.reserve(17*V.rows());
    Kd_IJV.reserve(17*V.rows());
    edges.resize(6,2);
    edges <<
      1,2,
      2,0,
      0,1,
      3,0,
      3,1,
      3,2;
    // Not implemented yet for tets
    assert(false);
  }
  // gather cotangent weights
  Matrix<Scalar,Dynamic,Dynamic> C;
  cotmatrix_entries(V,F,C);
  // should have weights for each edge
  assert(C.cols() == edges.rows());
  // loop over elements
  for(int i = 0;i<m;i++)
  {
    // loop over edges of element
    for(int e = 0;e<edges.rows();e++)
    {
      int source = F(i,edges(e,0));
      int dest = F(i,edges(e,1));
      double v = C(i,e)*(V(source,d)-V(dest,d))/3.0;
      // loop over edges again
      for(int f = 0;f<edges.rows();f++)
      {
        int Rs = F(i,edges(f,0));
        int Rd = F(i,edges(f,1));
        if(Rs == source && Rd == dest)
        {
          Kd_IJV.push_back(Triplet<Scalar>(Rs,Rd,v));
          Kd_IJV.push_back(Triplet<Scalar>(Rd,Rs,-v));
        }else if(Rd == source)
        {
          Kd_IJV.push_back(Triplet<Scalar>(Rd,Rs,v));
        }else if(Rs == dest)
        {
          Kd_IJV.push_back(Triplet<Scalar>(Rs,Rd,-v));
        }
      }
      Kd_IJV.push_back(Triplet<Scalar>(source,source,v));
      Kd_IJV.push_back(Triplet<Scalar>(dest,dest,-v));
    }
  }
  Kd.setFromTriplets(Kd_IJV.begin(),Kd_IJV.end());
  Kd.makeCompressed();
}

template <typename MatV, typename MatF, typename Scalar>
IGL_INLINE void igl::arap_linear_block_elements(
  const MatV & V,
  const MatF & F,
  const int d,
  Eigen::SparseMatrix<Scalar> & Kd)
{
  using namespace std;
  using namespace Eigen;
  // simplex size (3: triangles, 4: tetrahedra)
  int simplex_size = F.cols();
  // Number of elements
  int m = F.rows();
  // Temporary output
  Kd.resize(V.rows(), F.rows());
  vector<Triplet<Scalar> > Kd_IJV;
  Matrix<int,Dynamic,2> edges;
  if(simplex_size == 3)
  {
    // triangles
    Kd.reserve(7*V.rows());
    Kd_IJV.reserve(7*V.rows());
    edges.resize(3,2);
    edges <<
      1,2,
      2,0,
      0,1;
  }else if(simplex_size == 4)
  {
    // tets
    Kd.reserve(17*V.rows());
    Kd_IJV.reserve(17*V.rows());
    edges.resize(6,2);
    edges <<
      1,2,
      2,0,
      0,1,
      3,0,
      3,1,
      3,2;
  }
  // gather cotangent weights
  Matrix<Scalar,Dynamic,Dynamic> C;
  cotmatrix_entries(V,F,C);
  // should have weights for each edge
  assert(C.cols() == edges.rows());
  // loop over elements
  for(int i = 0;i<m;i++)
  {
    // loop over edges of element
    for(int e = 0;e<edges.rows();e++)
    {
      int source = F(i,edges(e,0));
      int dest = F(i,edges(e,1));
      double v = C(i,e)*(V(source,d)-V(dest,d));
      Kd_IJV.push_back(Triplet<Scalar>(source,i,v));
      Kd_IJV.push_back(Triplet<Scalar>(dest,i,-v));
    }
  }
  Kd.setFromTriplets(Kd_IJV.begin(),Kd_IJV.end());
  Kd.makeCompressed();
}

#ifdef IGL_STATIC_LIBRARY
// Explicit template instantiation
template IGL_INLINE void igl::arap_linear_block<Eigen::Matrix<double, -1, -1, 0, -1, -1>, Eigen::Matrix<int, -1, -1, 0, -1, -1>, double>(Eigen::Matrix<double, -1, -1, 0, -1, -1> const&, Eigen::Matrix<int, -1, -1, 0, -1, -1> const&, int, igl::ARAPEnergyType, Eigen::SparseMatrix<double, 0, int>&);
#endif