/* @file
 * @copyright This code is licensed under the 3-clause BSD license.
 *   Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.
 *   See LICENSE.txt for details.
 */

#include <boost/test/unit_test.hpp>
#include <Eigen/Geometry>

#include "Molassembler/Temple/Adaptors/AllPairs.h"
#include "Molassembler/Temple/Adaptors/Zip.h"
#include "Molassembler/Temple/Functional.h"
#include "Molassembler/Temple/Functor.h"
#include "Molassembler/Temple/SetAlgorithms.h"
#include "Molassembler/Temple/Stringify.h"
#include "Molassembler/Temple/Stl17.h"
#include "Molassembler/Temple/constexpr/Numeric.h"
#include "Molassembler/Temple/constexpr/ToStl.h"
#include "Molassembler/Temple/constexpr/TupleTypePairs.h"

#include "Molassembler/Shapes/PropertyCaching.h"
#include "Molassembler/Shapes/Data.h"

#include <set>
#include <numeric>
#include <iostream>
#include <iomanip>

using namespace Scine::Molassembler;
using namespace Shapes;

template<typename EnumType>
constexpr inline auto underlying(const EnumType e) {
  return static_cast<std::underlying_type_t<EnumType>>(e);
}

template<typename ShapeClass>
struct LockstepTest {
  static bool value() {
    return Shapes::nameIndex(ShapeClass::shape) == underlying(ShapeClass::shape);
  }
};

BOOST_AUTO_TEST_CASE(SymmetryTypeAndPositionInEnumLockstep, *boost::unit_test::label("Shapes")) {
  BOOST_CHECK_MESSAGE(
    Temple::all_of(
      Temple::Tuples::map<
        Shapes::Data::allShapeDataTypes,
        LockstepTest
      >()
    ),
    "Not all symmetries have the same order in Name and allShapeDataTypes"
  );
}

BOOST_AUTO_TEST_CASE(AngleFunctionInputSymmetry, *boost::unit_test::label("Shapes")) {
  // every angle function must be symmetrical on input of valid unsigned indices
  for(const Shape shape: allShapes) {
    bool passesAll = true;

    for(Vertex i(0); i < size(shape) && passesAll; ++i) {
      for(Vertex j(i + 1); j < size(shape); ++j) {
        if(angleFunction(shape)(i, j) != angleFunction(shape)(j, i)) {
          passesAll = false;
          std::cout << name(shape)
            << " is not symmetrical w.r.t. input indices: falsified by ("
            << i << ", " << j <<") -> (" << angleFunction(shape)(i, j)
            << ", " << angleFunction(shape)(j, i) << ")." << std::endl;
          break;
        }
      }
    }

    BOOST_CHECK(passesAll);
  }

}

BOOST_AUTO_TEST_CASE(AngleFunctionZeroForIdenticalInput, *boost::unit_test::label("Shapes")) {
  // every angle function must return 0 for identical indices
  for(const Shape shape: allShapes) {
    bool passesAll = true;

    for(Vertex i(0); i < size(shape); i++) {
      if(angleFunction(shape)(i, i) != 0) {
        passesAll = false;
        std::cout << name(shape)
          << "'s angle function does not return zero for identical indices ("
          << i << ", " << i << ")." << std::endl;
        break;
      }
    }

    BOOST_CHECK(passesAll);
  }
}

BOOST_AUTO_TEST_CASE(AnglesWithinRadiansBounds, *boost::unit_test::label("Shapes")) {
  for(const Shape shape : allShapes) {
    bool passesAll = true;

    for(Vertex i(0); i < size(shape); i++) {
      for(Vertex j(0); j < size(shape); j++) {
        if(
          !(
            0 <= angleFunction(shape)(i, j)
          ) || !(
            angleFunction(shape)(i, j) <= M_PI
          )
        ) {
          passesAll = false;
          std::cout << name(shape)
            << "'s angle function is not within radians bounds for indices ("
            << i << ", " << j << ") -> " << angleFunction(shape)(i, j)
            << std::endl;
          break;
        }
      }
    }

    BOOST_CHECK(passesAll);
  }
}

BOOST_AUTO_TEST_CASE(AnglesMatchCoordinates, *boost::unit_test::label("Shapes")) {

  /* The results of the angle functions ought to match the geometries specified
   * by the coordinates
   */

  for(const auto& shape: allShapes) {
    auto getCoordinates =  [shape](const Vertex index) -> Eigen::Vector3d {
      return coordinates(shape).col(index);
    };

    bool all_pass = true;

    for(Vertex i {0}; i < size(shape); i++) {
      for(Vertex j {i + 1}; j < size(shape); j++) {
        const double angleInCoordinates = std::acos(
          Temple::Stl17::clamp(
            getCoordinates(i).dot(getCoordinates(j)) / (
              getCoordinates(i).norm() * getCoordinates(j).norm()
            ),
            -1.0,
            1.0
          )
        );

        const double angleDifference = angleInCoordinates - angleFunction(shape)(i, j);

        // Tolerate only one degree difference
        if(std::fabs(angleDifference) > 1) {
          all_pass = false;

          std::cout << name(shape)
            << ": angleFunction != angles from coordinates ("
            << i << ", " << j << "): " << angleDifference
            << ", angleFunction = " << angleFunction(shape)(i, j)
            << ", angle from coordinates = " << angleInCoordinates << std::endl;
        }
      }
    }

    BOOST_CHECK(all_pass);
  }
}

BOOST_AUTO_TEST_CASE(AllTetrahedraPositive, *boost::unit_test::label("Shapes")) {
  /* Checks if sequence that tetrahedra are defined in leads to a positive
   * volume when calculated via
   *
   *  (1 - 4) dot [ (2 - 4) x (3 - 4) ]
   *
   */
  for(const auto& shape: allShapes) {
    auto getCoordinates = [shape](const boost::optional<Vertex>& indexOption) -> Eigen::Vector3d {
      if(indexOption) {
        return coordinates(shape).col(indexOption.value());
      }

      return {0, 0, 0};
    };

    bool all_pass = true;

    for(const auto& tetrahedron: tetrahedra(shape)) {

      double tetrahedronVolume = (
        getCoordinates(tetrahedron[0]) - getCoordinates(tetrahedron[3])
      ).dot(
        (
          getCoordinates(tetrahedron[1]) - getCoordinates(tetrahedron[3])
        ).cross(
          getCoordinates(tetrahedron[2]) - getCoordinates(tetrahedron[3])
        )
      );

      if(tetrahedronVolume < 0) {
        all_pass = false;
        std::cout << name(shape) << ": Tetrahedron {";

        for(unsigned i = 0; i < 4; i++) {
          if(tetrahedron[i]) {
            std::cout << tetrahedron[i].value();
          } else {
            std::cout << "C";
          }

          if(i != 3) {
            std::cout << ", ";
          }
        }

        std::cout << "} has negative volume (" << tetrahedronVolume << ")."
          << std::endl;
      }
    }

    BOOST_CHECK(all_pass);
  }
}

BOOST_AUTO_TEST_CASE(TetrahedraDefinitionIndicesUnique, *boost::unit_test::label("Shapes")) {
  for(const auto& shape : allShapes) {
    for(const auto& tetrahedron : tetrahedra(shape)) {
      bool containsAnEmptyOption = false;

      for(const auto& edgeOption : tetrahedron) {
        if(!edgeOption) {
          containsAnEmptyOption = true;
          break;
        }
      }

      std::set<unsigned> indices;

      for(const auto& edgeOption : tetrahedron) {
        if(edgeOption) {
          indices.insert(edgeOption.value());
        }
      }

      BOOST_CHECK(indices.size() + static_cast<unsigned>(containsAnEmptyOption) == 4);
    }
  }
}

BOOST_AUTO_TEST_CASE(SmallestAngleValueCorrect, *boost::unit_test::label("Shapes")) {
  auto shapeSmallestAngle = [](const Shape shape) -> double {
    return Temple::accumulate(
      Temple::Adaptors::allPairs(
        Temple::iota<Vertex>(Vertex(size(shape)))
      ),
      M_PI,
      [shape](const double currentSmallest, const auto& vertexPair) {
        return std::min(
          currentSmallest,
          Temple::invoke(angleFunction(shape), vertexPair)
        );
      }
    );
  };

  const double comparisonSmallestAngle = Temple::min(
    Temple::map(allShapes, shapeSmallestAngle)
  );

  BOOST_CHECK(0 < smallestAngle && smallestAngle < M_PI);
  BOOST_CHECK_MESSAGE(
    std::fabs(
      smallestAngle - comparisonSmallestAngle
    ) < 1e-4,
    "The constant smallest angle set by the library is NOT the smallest "
    << "returned angle within the library. Current value of smallestAngle: "
    << smallestAngle
    << ", true smallest angle:" << comparisonSmallestAngle
  );
}

#ifdef USE_CONSTEXPR_TRANSITION_MAPPINGS
/* NOTE: can refactor out doLigandGainTestIfAdjacent with a simple if-constexpr
 * in C++17
 */
template<class ShapeClassFrom, class ShapeClassTo>
std::enable_if_t<
  (
    ShapeClassFrom::size + 1 == ShapeClassTo::size
    || ShapeClassFrom::size == ShapeClassTo::size
  ),
  bool
> doLigandGainTestIfAdjacent() {
  /* Struct:
   * .mappings - dynamic array of fixed-size index mappings
   * .angularDistortion, .chiralDistortion - doubles
   */
  auto constexprMappings = allMappings.at(
    static_cast<unsigned>(ShapeClassFrom::shape),
    static_cast<unsigned>(ShapeClassTo::shape)
  ).value();
  /* Vector of structs:
   * .indexMapping - vector containing the index mapping
   * .totalDistortion, .chiralDistortion - doubles
   */
  auto dynamicMappings = Properties::selectBestTransitionMappings(
    Properties::shapeTransitionMappings(
      ShapeClassFrom::shape,
      ShapeClassTo::shape
    )
  );

  Temple::Floating::ExpandedRelativeEqualityComparator<double> comparator {
    Properties::floatingPointEqualityThreshold
  };

  if(
    comparator.isUnequal(
      dynamicMappings.angularDistortion,
      constexprMappings.angularDistortion
    ) || comparator.isUnequal(
      dynamicMappings.chiralDistortion,
      constexprMappings.chiralDistortion
    )
  ) {
    return false;
  }

  // Do a full set comparison
  auto convertedMappings = Temple::map_stl(
    Temple::toSTL(constexprMappings.mappings),
    [&](const auto& indexList) -> std::vector<Vertex> {
      std::vector<Vertex> v;
      for(unsigned i : indexList) {
        v.emplace_back(i);
      }
      return v;
    }
  );

  decltype(convertedMappings) dynamicResultSet {
    std::begin(dynamicMappings.indexMappings),
    std::end(dynamicMappings.indexMappings)
  };

  return Temple::set_symmetric_difference(
    convertedMappings,
    dynamicResultSet
  ).empty();
}

using IndexAndMappingsPairType = std::pair<
  unsigned,
  Shapes::ConstexprProperties::MappingsReturnType
>;

constexpr bool pairEqualityComparator(
  const IndexAndMappingsPairType& a,
  const IndexAndMappingsPairType& b
) {
  Temple::Floating::ExpandedRelativeEqualityComparator<double> comparator {
    Shapes::Properties::floatingPointEqualityThreshold
  };

  return (
    comparator.isEqual(a.second.angularDistortion, b.second.angularDistortion)
    && comparator.isEqual(a.second.chiralDistortion, b.second.chiralDistortion)
  );
}

template<class ShapeClassFrom, class ShapeClassTo>
std::enable_if_t<
  ShapeClassFrom::size == ShapeClassTo::size + 1,
  bool
> doLigandGainTestIfAdjacent() {
  // Ligand loss situation

  /* Constexpr part */
  Temple::Array<
    std::pair<
      unsigned,
      Shapes::ConstexprProperties::MappingsReturnType
    >,
    ShapeClassFrom::size
  > constexprMappings;

  for(unsigned i = 0; i < ShapeClassFrom::size; ++i) {
    constexprMappings.at(i) = std::make_pair(
      i,
      Shapes::ConstexprProperties::ligandLossMappings<
        ShapeClassFrom,
        ShapeClassTo
      >(i)
    );
  }

  // Group the results
  auto constexprGroups = Temple::groupByEquality(
    constexprMappings,
    pairEqualityComparator // C++17 constexpr lambda
  );

  /* Dynamic part */
  std::vector<
    std::pair<unsigned, Shapes::Properties::ShapeTransitionGroup>
  > dynamicMappings;

  for(unsigned i = 0; i < ShapeClassFrom::size; ++i) {
    dynamicMappings.emplace_back(
      i,
      selectBestTransitionMappings(
        Properties::ligandLossTransitionMappings(
          ShapeClassFrom::shape,
          ShapeClassTo::shape,
          i
        )
      )
    );
  }

  // Analyze all mappings - which indices have "identical" target mappings?
  auto dynamicGroups = Temple::groupByEquality(
    dynamicMappings,
    [&](const auto& firstMappingPair, const auto& secondMappingPair) -> bool {
      return (
        Temple::Floating::isCloseRelative(
          firstMappingPair.second.angularDistortion,
          secondMappingPair.second.angularDistortion,
          Shapes::Properties::floatingPointEqualityThreshold
        ) && Temple::Floating::isCloseRelative(
          firstMappingPair.second.chiralDistortion,
          secondMappingPair.second.chiralDistortion,
          Shapes::Properties::floatingPointEqualityThreshold
        )
      );
    }
  );

  /* Comparison */
  // Quick check
  if(dynamicGroups.size() != constexprGroups.size()) {
    return false;
  }

  // Compare an unsigned set of sub-group sizes from each
  std::multiset<unsigned> dynamicGroupSizes, constexprGroupSizes;

  for(const auto& dynamicGroup : dynamicGroups) {
    dynamicGroupSizes.insert(dynamicGroup.size());
  }

  for(const auto& constexprGroup : constexprGroups) {
    constexprGroupSizes.insert(constexprGroup.size());
  }

  return dynamicGroupSizes == constexprGroupSizes;
}

// Base case in which source and target symmetries are non-adjacent
template<class ShapeClassFrom, class ShapeClassTo>
std::enable_if_t<
  (
    ShapeClassFrom::size != ShapeClassTo::size + 1
    && ShapeClassFrom::size + 1 != ShapeClassTo::size
    && ShapeClassFrom::size != ShapeClassTo::size
  ),
  bool
> doLigandGainTestIfAdjacent() {
  return true;
}

template<class ShapeClassFrom, class ShapeClassTo>
struct LigandGainTest {
  static bool value() {
    return doLigandGainTestIfAdjacent<ShapeClassFrom, ShapeClassTo>();
  }
};
#endif

template<class ShapeClass>
struct RotationGenerationTest {
  static bool value() {

    // This is a DynamicSet of ShapeClass-sized Arrays
    auto constexprRotations = ConstexprProperties::generateAllRotations<ShapeClass>(
      ConstexprProperties::startingIndexSequence<ShapeClass>()
    );

    // This is a std::set of ShapeClass-sized std::vectors
    auto dynamicRotations = Properties::generateAllRotations(
      ShapeClass::shape,
      Temple::iota<Vertex>(Vertex(ShapeClass::size))
    );

    auto convertedRotations = Temple::map_stl(
      Temple::toSTL(constexprRotations),
      [&](const auto& indexList) -> std::vector<Vertex> {
        std::vector<Vertex> v;
        for(unsigned i : indexList) {
          v.emplace_back(i);
        }
        return v;
      }
    );

    if(convertedRotations.size() != constexprRotations.size()) {
      std::cout << "In symmetry " << ShapeClass::stringName << ", "
        << "constexpr rotations set reports " << constexprRotations.size()
        << " elements but the STL mapped variant has only "
        << convertedRotations.size() << " elements!" << std::endl;
    }

    bool pass = true;

    // Size mismatch
    if(convertedRotations.size() != dynamicRotations.size()) {
      pass = false;
    } else {
      pass = (
        Temple::set_symmetric_difference(
          convertedRotations,
          dynamicRotations
        ).empty()
      );
    }

    if(!pass) {
      std::cout << "Rotation generation differs for "
        << ShapeClass::stringName
        << " shape: Sizes of generated sets are different. "
        << "constexpr - " << convertedRotations.size() << " != "
        << dynamicRotations.size() << " - dynamic" << std::endl;
      std::cout << " Maximum #rotations: " << ConstexprProperties::maxRotations<ShapeClass>()
        << std::endl;

      std::cout << " Converted constexpr:" << std::endl;
      for(const auto& element : convertedRotations) {
        std::cout << " {" << Temple::condense(element)
          << "}\n";
      }

      std::cout << " Dynamic:" << std::endl;
      for(const auto& element : dynamicRotations) {
        std::cout << " {" << Temple::condense(element)
          << "}\n";
      }
    }

    return pass;
  }
/* Previously, when the interface with which this is used (unpackToFunction) was
 * unable to cope with both value data members and value function members
 * equally, this was of the form::
 *
 *   static bool initialize() {
 *     ...
 *     return value;
 *   }
 *   static bool value = initialize();
 *
 * Although this seems equivalent, it really isn't, since initialize() is called
 * at static initialization time instead of at first use as when value is
 * a function. Since, in this case, the value function depends on another static
 * value (generateAllRotations -> shapeData), the value-initialize variant
 * leads to a static initialization fiasco, where it is unclear whether the
 * value data member or shapeData is initialized first.
 *
 * Only in the case of static constexpr is the value-initialize variant
 * semantically equivalent to the data member variant.
 */
};

std::string getGraphvizNodeName(const Shapes::Shape shape) {
  auto stringName = Shapes::name(shape);

  stringName.erase(
    std::remove_if(
      stringName.begin(),
      stringName.end(),
      [](const char& singleChar) -> bool {
        return (
          singleChar == ' '
          || singleChar == '-'
        );
      }
    ),
    stringName.end()
  );
  return stringName;
}

BOOST_AUTO_TEST_CASE(ConstexprRotationsMatch, *boost::unit_test::label("Shapes")) {
  // Full test of rotation algorithm equivalency for all symmetries
  BOOST_CHECK_MESSAGE(
    Temple::all_of(
      Temple::Tuples::map<
        Shapes::Data::allShapeDataTypes,
        RotationGenerationTest
      >()
    ),
    "There is a discrepancy between constexpr and dynamic rotation generation"
  );

#ifdef USE_CONSTEXPR_TRANSITION_MAPPINGS
  // Test transitions generation/evaluation algorithm equivalency for all
  BOOST_CHECK_MESSAGE(
    Temple::all_of(
      Temple::Tuples::mapAllPairs<
        Shapes::Data::allShapeDataTypes,
        LigandGainTest
      >()
    ),
    "There is a discrepancy between constexpr and dynamic ligand gain mapping"
    << " generation!"
  );
#endif
}

template<typename ShapeClass>
struct NumUnlinkedTestFunctor {
  static bool value() {
    for(unsigned i = 1; i < ShapeClass::size; ++i) {
      unsigned constexprResult = ConstexprProperties::numUnlinkedStereopermutations<ShapeClass>(i);

      unsigned dynamicResult = Properties::numUnlinkedStereopermutations(ShapeClass::shape, i);

      if(constexprResult != dynamicResult) {
        std::cout << "Mismatch for " << Shapes::name(ShapeClass::shape) << " and " << i << " identical ligands between constexpr and dynamic number of unlinked: " << constexprResult << " vs. " << dynamicResult << "\n";
        return false;
      }

      // Cross-check with constexpr hasMultiple
      bool constexprHasMultiple = ConstexprProperties::hasMultipleUnlinkedStereopermutations<ShapeClass>(i);
      if((constexprResult > 1) != constexprHasMultiple) {
        std::cout << "Mismatch between constexpr count and constexpr "
          << "hasMultiple unlinked ligands for "
          << Shapes::name(ShapeClass::shape) << " and "
          << i << " identical ligands: " << constexprResult << " and "
          << std::boolalpha << constexprHasMultiple << "\n";
        return false;
      }

      // Cross-check with dynamic hasMultiple
      bool dynamicHasMultiple = Properties::hasMultipleUnlinkedStereopermutations(ShapeClass::shape, i);
      if((constexprResult > 1) != dynamicHasMultiple) {
        std::cout << "Mismatch between constexpr count and dynamic "
          << "hasMultiple unlinked ligands for "
          << Shapes::name(ShapeClass::shape) << " and "
          << i << " identical ligands: " << constexprResult << " and "
          << std::boolalpha << dynamicHasMultiple << "\n";
        return false;
      }
    }

    return true;
  }
};

BOOST_AUTO_TEST_CASE(NumUnlinkedStereopermsMatch, *boost::unit_test::label("Shapes")) {
  using TestTypes = std::tuple<
    Data::Line,
    Data::Bent,
    Data::EquilateralTriangle, // 3
    Data::VacantTetrahedron,
    Data::T,
    Data::Tetrahedron, // 4
    Data::Square,
    Data::Seesaw,
    Data::TrigonalPyramid,
    Data::SquarePyramid, // 5
    Data::TrigonalBipyramid,
    Data::Pentagon,
    Data::Octahedron, // 6
    Data::TrigonalPrism,
    Data::PentagonalPyramid,
    Data::Hexagon
  >;

  BOOST_CHECK_MESSAGE(
    Temple::all_of(
      Temple::Tuples::map<TestTypes, NumUnlinkedTestFunctor>()
    ),
    "Not all numbers of unlinked stereopermutations match across constexpr "
    " and dynamic algorithms"
  );

  BOOST_CHECK(Properties::numUnlinkedStereopermutations(Shapes::Shape::Line, 0) == 1);
  BOOST_CHECK(Properties::numUnlinkedStereopermutations(Shapes::Shape::Bent, 0) == 1);
  BOOST_CHECK(Properties::numUnlinkedStereopermutations(Shapes::Shape::EquilateralTriangle, 0) == 1);
  BOOST_CHECK(Properties::numUnlinkedStereopermutations(Shapes::Shape::Tetrahedron, 0) == 2);
  BOOST_CHECK(Properties::numUnlinkedStereopermutations(Shapes::Shape::Octahedron, 0) == 30);
}

static_assert(
  nShapes == std::tuple_size<Data::allShapeDataTypes>::value,
  "nShapes does not equal number of shape data class types in "
  "allShapeDataTypes"
);

#ifdef USE_CONSTEXPR_TRANSITION_MAPPINGS
BOOST_AUTO_TEST_CASE(ConstexprMappingsExist, *boost::unit_test::label("Shapes")) {
  /* In every case where allMappings has a value, getMapping must also return
   * a some optional
   */
  bool pass = true;
  for(const auto& fromShape : Shapes::allShapes) {
    auto i = static_cast<unsigned>(fromShape);
    for(const auto& toShape : Shapes::allShapes) {
      auto j = static_cast<unsigned>(toShape);
      if(
        i < j
        && allMappings.at(i, j).hasValue()
          != static_cast<bool>(
            Shapes::getMapping(fromShape, toShape)
          )
      ) {
        pass = false;
        break;
      }
    }
  }

  BOOST_CHECK_MESSAGE(
    pass,
    "Not all constexpr mappings from allMappings are available from getMapping!"
  );
}
#endif

BOOST_AUTO_TEST_CASE(AngleBoundsCorrect, *boost::unit_test::label("Shapes")) {
  BOOST_CHECK(Shapes::minimumAngle(Shapes::Shape::T) == M_PI / 2);
  BOOST_CHECK(Shapes::maximumAngle(Shapes::Shape::T) == M_PI);

  BOOST_CHECK(Shapes::minimumAngle(Shapes::Shape::Octahedron) == M_PI / 2);
  BOOST_CHECK(Shapes::maximumAngle(Shapes::Shape::Octahedron) == M_PI);

  BOOST_CHECK(Shapes::minimumAngle(Shapes::Shape::TrigonalBipyramid) == M_PI / 2);
  BOOST_CHECK(Shapes::maximumAngle(Shapes::Shape::TrigonalBipyramid) == M_PI);

  BOOST_CHECK(
    Shapes::minimumAngle(Shapes::Shape::Tetrahedron) == Shapes::maximumAngle(Shapes::Shape::Tetrahedron)
  );
}

std::ostream& operator << (std::ostream& os, const std::vector<char>& chars) {
  std::cout << "{";
  for(const char x : chars) {
    std::cout << x;
  }
  std::cout << "}";
  return os;
}

BOOST_AUTO_TEST_CASE(PositionGroups, *boost::unit_test::label("Shapes")) {
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::Line) == std::vector<char> (2, 'A'));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::Tetrahedron) == std::vector<char> (4, 'A'));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::TrigonalBipyramid) == (std::vector<char> {'A','A','A','B','B'}));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::CappedTrigonalPrism) == (std::vector<char> {'A','B','C','B','C','D','D'}));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::Octahedron) == std::vector<char> (6, 'A'));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::Cube) == std::vector<char> (8, 'A'));
  BOOST_CHECK(Properties::positionGroupCharacters(Shape::Icosahedron) == std::vector<char> (12, 'A'));
}

BOOST_AUTO_TEST_CASE(DimensionalityProperty, *boost::unit_test::label("Shapes")) {
  BOOST_CHECK(threeDimensional(Shape::Tetrahedron));
  BOOST_CHECK(!threeDimensional(Shape::Line));
  BOOST_CHECK(!threeDimensional(Shape::Bent));
  BOOST_CHECK(!threeDimensional(Shape::Pentagon));
  BOOST_CHECK(!threeDimensional(Shape::Hexagon));
  BOOST_CHECK(threeDimensional(Shape::Icosahedron));
}


#include "read_map.hpp"
#include "timer.hpp"
#include <mpl_planner/planner/map_planner.h>

#include <boost/geometry.hpp>
#include <boost/geometry/geometries/point_xy.hpp>
#include <boost/geometry/geometries/polygon.hpp>

int main(int argc, char **argv) {
  if (argc != 2) {
    printf(ANSI_COLOR_RED "Input yaml required!\n" ANSI_COLOR_RESET);
    return -1;
  }

  // Load the map
  MapReader<Vec2i, Vec2f> reader(argv[1]);
  if (!reader.exist()) {
    printf(ANSI_COLOR_RED "Cannot find input file [%s]!\n" ANSI_COLOR_RESET,
           argv[1]);
    return -1;
  }

  // Pass the data into a VoxelMapUtil class for collision checking
  std::shared_ptr<MPL::OccMapUtil> map_util = std::make_shared<MPL::OccMapUtil>();
  map_util->setMap(reader.origin(), reader.dim(), reader.data(),
                   reader.resolution());
  map_util->freeUnknown();

  // Initialize start and goal, using acc control
  Waypoint2D start, goal;
  start.pos = Vec2f(reader.start(0), reader.start(1));
  start.vel = Vec2f::Zero();
  start.acc = Vec2f::Zero();
  start.jrk = Vec2f::Zero();
  start.yaw = 0;
  start.use_pos = true;
  start.use_vel = true;
  start.use_acc = false;
  start.use_jrk = false;
  start.use_yaw = false;

  goal.pos = Vec2f(reader.goal(0), reader.goal(1));
  goal.vel = Vec2f::Zero();
  goal.acc = Vec2f::Zero();
  goal.jrk = Vec2f::Zero();
  goal.yaw = 0;
  goal.control = start.control;

  // Initialize control input
  decimal_t u = 0.5;
  decimal_t du = u;
  vec_E<VecDf> U;
  for (decimal_t dx = -u; dx <= u; dx += du)
    for (decimal_t dy = -u; dy <= u; dy += du)
      U.push_back(Vec2f(dx, dy));

  // Initialize planner
  std::unique_ptr<MPL::OccMapPlanner> planner(
      new MPL::OccMapPlanner(true));    // Declare a mp planner using voxel map
  planner->setMapUtil(map_util); // Set collision checking function
  planner->setVmax(1.0);         // Set max velocity
  planner->setAmax(1.0);         // Set max acceleration
  planner->setDt(1.0);           // Set dt for each primitive
  planner->setU(U);              // Set control input

  // Planning
  Timer time(true);
  bool valid = planner->plan(start, goal); // Plan from start to goal
  double dt = time.Elapsed().count();
  printf("MPL Planner takes: %f ms\n", dt);
	printf("MPL Planner expanded states: %zu\n", planner->getCloseSet().size());
	const Trajectory2D traj = planner->getTraj();

  // Create a path from planned traj
	const auto ws = traj.getWaypoints();
	vec_Vec2f path;
	for(const auto& w: ws)
		path.push_back(w.pos);
  // Initiaize planner as a distance map planner
	planner.reset(new MPL::OccMapPlanner(true));
  planner->setMapUtil(map_util); // Set collision checking function
  planner->setVmax(1.0);         // Set max velocity
  planner->setAmax(1.0);         // Set max acceleration
  planner->setDt(1.0);           // Set dt for each primitive
  planner->setU(U);              // Set control input
  planner->setEpsilon(1.0);           // Set heursitic to zero

	planner->setValidRegion(path, Vec2f(0.5, 0.5)); // Set search region around path
	planner->setPotentialRadius(Vec2f(1.0, 1.0)); // Set potential distance
  planner->setPotentialWeight(1); // Set potential weight
  planner->setGradientWeight(0); // Set gradient weight
  planner->updatePotentialMap(start.pos); // Update potential map

  // Planning
  Timer time_dist(true);
  bool valid_dist = planner->plan(start, goal); // Plan from start to goal
  double dt_dist = time_dist.Elapsed().count();
  printf("MPL Distance Planner takes: %f ms\n", dt_dist);
	printf("MPL Distance Planner expanded states: %zu\n", planner->getCloseSet().size());
	const Trajectory2D traj_dist = planner->getTraj();

  // Plot the result in svg image
  typedef boost::geometry::model::d2::point_xy<double> point_2d;
  // Declare a stream and an SVG mapper
  std::ofstream svg("output.svg");
  boost::geometry::svg_mapper<point_2d> mapper(svg, 1000, 1000);

  // Draw the canvas
  const Vec2i dim = reader.dim();
  boost::geometry::model::polygon<point_2d> bound;
  const double origin_x = reader.origin()(0);
  const double origin_y = reader.origin()(1);
  const double range_x = reader.dim()(0) * reader.resolution();
  const double range_y = reader.dim()(1) * reader.resolution();
  std::vector<point_2d> points;
  points.push_back(point_2d(origin_x, origin_y));
  points.push_back(point_2d(origin_x, origin_y + range_y));
  points.push_back(point_2d(origin_x + range_x, origin_y + range_y));
  points.push_back(point_2d(origin_x + range_x, origin_y));
  points.push_back(point_2d(origin_x, origin_y));
  boost::geometry::assign_points(bound, points);
  boost::geometry::correct(bound);

  mapper.add(bound);
  mapper.map(bound, "fill-opacity:1.0;fill:rgb(255,255,255);stroke:rgb(0,0,0);"
                    "stroke-width:2"); // White

  // Draw start and goal
  point_2d start_pt, goal_pt;
  boost::geometry::assign_values(start_pt, start.pos(0), start.pos(1));
  mapper.add(start_pt);
  mapper.map(start_pt, "fill-opacity:1.0;fill:rgb(255,0,0);", 10); // Red
  boost::geometry::assign_values(goal_pt, goal.pos(0), goal.pos(1));
  mapper.add(goal_pt);
  mapper.map(goal_pt, "fill-opacity:1.0;fill:rgb(255,0,0);", 10); // Red

  // Draw the obstacles
  const auto data = map_util->getMap();
  for(int x = 0; x < dim(0); x ++) {
    for(int y = 0; y < dim(1); y ++) {
        Vec2f pt = map_util->intToFloat(Vec2i(x, y));
        decimal_t occ = (decimal_t) data[map_util->getIndex(Vec2i(x,y))]/100;
        if(occ <= 0)
          continue;
        point_2d a;
        boost::geometry::assign_values(a, pt(0), pt(1));
        mapper.add(a);
        if(occ < 1)
          mapper.map(a, "fill-opacity:"+std::to_string(occ/2.0) +";fill:rgb(118,215,234);", 1);
        else
          mapper.map(a, "fill-opacity:1.0;fill:rgb(0,0,0);", 1);
    }
  }

  // Draw searched region
  for (const auto &pt : planner->getSearchRegion()) {
 // for (const auto &pt : planner->getCloseSet()) {
    point_2d a;
    boost::geometry::assign_values(a, pt(0), pt(1));
    mapper.add(a);
    mapper.map(a, "fill-opacity:0.2;fill:rgb(100,200,100);", 1); // Green
  }

  if (valid) {
    // Draw the trajectory
    double total_t = traj.getTotalTime();
    printf("Total time T: %f\n", total_t);
    printf("Total J:  J(VEL) = %f, J(ACC) = %f, J(JRK) = %f, J(SNP) = %f\n",
           traj.J(Control::VEL), traj.J(Control::ACC), traj.J(Control::JRK), traj.J(Control::SNP));
    int num = 200; // number of points on trajectory to draw
    const auto ws = traj.sample(num);
    boost::geometry::model::linestring<point_2d> line;
    for (const auto& w: ws)
        line.push_back(point_2d(w.pos(0), w.pos(1)));
    mapper.add(line);
    mapper.map(line,
        "opacity:0.4;fill:none;stroke:rgb(212,0,0);stroke-width:5"); // Red
	}

	if (valid_dist) {
		// Draw the trajectory
		double total_t = traj_dist.getTotalTime();
		printf("Total dist time T: %f\n", total_t);
		printf("Total dist J:  J(VEL) = %f, J(ACC) = %f, J(JRK) = %f, J(SNP) = %f\n",
					 traj_dist.J(Control::VEL), traj_dist.J(Control::ACC), traj_dist.J(Control::JRK), traj_dist.J(Control::SNP));
		int num = 200; // number of points on trajectory to draw
		const auto ws = traj_dist.sample(num);
		boost::geometry::model::linestring<point_2d> line;
		for (const auto& w: ws)
			line.push_back(point_2d(w.pos(0), w.pos(1)));
		mapper.add(line);
		mapper.map(line,
							 "opacity:0.8;fill:none;stroke:rgb(10,10,250);stroke-width:5"); // Blue
	}

	// Write title at the lower right corner on canvas
	mapper.text(point_2d(origin_x + range_x - 11, origin_y+2.4), "test_distance_map_planner_2d",
							"fill-opacity:1.0;fill:rgb(10,10,250);");

  mapper.text(point_2d(origin_x + range_x - 13, origin_y+1.8), "Green: ",
              "fill-opacity:1.0;fill:rgb(100,200,100);");
  mapper.text(point_2d(origin_x + range_x - 10.5, origin_y+1.8), "search region",
              "fill-opacity:1.0;fill:rgb(0,0,0);");

  mapper.text(point_2d(origin_x + range_x - 13, origin_y+1.2), "Red: ",
              "fill-opacity:1.0;fill:rgb(237,10,63);");
  mapper.text(point_2d(origin_x + range_x - 10.5, origin_y+1.2), "original trajectory",
              "fill-opacity:1.0;fill:rgb(0,0,0);");

  mapper.text(point_2d(origin_x + range_x - 13, origin_y+0.6), "Blue:",
              "fill-opacity:1.0;fill:rgb(10,10,250);");
  mapper.text(point_2d(origin_x + range_x - 10.5, origin_y+0.6), "perturbed trajectory",
              "fill-opacity:1.0;fill:rgb(0,0,0);");

  return 0;
}


#include <svo/common/frame.h>
#include <svo/common/point.h>
#include <svo/common/imu_calibration.h>
#include <svo/backend/graph_manager.h>
#include <vikit/math_utils.h>

// boost
#include <boost/make_shared.hpp>
#include <boost/foreach.hpp>
#include <boost/shared_ptr.hpp>

// gtsam
#include <gtsam/geometry/Point2.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/slam/ProjectionFactor.h>
#include <gtsam/slam/SmartProjectionPoseFactor.h>
#include <gtsam/navigation/ImuFactor.h>
#include <gtsam/navigation/CombinedImuFactor.h>
#include <gtsam/navigation/ImuBias.h>
#include <gtsam/linear/linearExceptions.h>

namespace svo {

GraphManager::GraphManager(const GraphManagerOptions& options)
  : options_(options)
{}

GraphManager::~GraphManager()
{}

void GraphManager::initialize()
{
  // camera
  cam_calib_.reset(new gtsam::Cal3_S2(1.0, 1.0, 0.0, 0.0, 0.0));

  // reprojection error noise models
  const size_t kMaxPyrLevels = 6;
  if(options_.use_robust_px_noise)
  {
    const double kRobustKurtosis = 3.0;
    px_noise_pyr_.resize(kMaxPyrLevels);
    for(size_t level=0u; level<kMaxPyrLevels; ++level)
    {
      const double px_noise = options_.reproj_error_noise * (1<<level);
      px_noise_pyr_[level] =
          gtsam::noiseModel::Robust::Create(
            gtsam::noiseModel::mEstimator::Cauchy::Create(kRobustKurtosis),
            gtsam::noiseModel::Isotropic::Sigma(2, px_noise));
    }
  }
  else
  {
    px_noise_pyr_.resize(kMaxPyrLevels);
    for(size_t level=0u; level<kMaxPyrLevels; ++level)
    {
      const double px_noise = options_.reproj_error_noise; // (1<<level);
      px_noise_pyr_[level] = gtsam::noiseModel::Isotropic::Sigma(2, px_noise);
    }
  }
  smart_noise_ = gtsam::noiseModel::Isotropic::Sigma(2, options_.reproj_error_noise);
  smart_factor_params_.reset(new gtsam::SmartProjectionParams(
                               gtsam::HESSIAN,
                               gtsam::IGNORE_DEGENERACY, // creates rotation-only constraint
                               true, true));

  // TODO: retriangulation threshold is a const variable and cant be set in constructor.
  smart_factor_params_->setRankTolerance(1e-3);
//  smart_factor_params_->setEnableEPI(false); // nonlinear refinement after triangulation
//  smart_factor_params_->setLandmarkDistanceThreshold(100);
  smart_factor_params_->setDynamicOutlierRejectionThreshold(options_.smart_reproj_threshold);

  // position prior
  pose_prior_covariance_.setZero();
  pose_prior_covariance_.diagonal()[0] = options_.init_roll_pitch_sigma * options_.init_roll_pitch_sigma;
  pose_prior_covariance_.diagonal()[1] = options_.init_roll_pitch_sigma * options_.init_roll_pitch_sigma;
  pose_prior_covariance_.diagonal()[2] = options_.init_yaw_sigma * options_.init_yaw_sigma;
  pose_prior_covariance_.diagonal()[3] = options_.init_pos_sigma * options_.init_pos_sigma;
  pose_prior_covariance_.diagonal()[4] = options_.init_pos_sigma * options_.init_pos_sigma;
  pose_prior_covariance_.diagonal()[5] = options_.init_pos_sigma * options_.init_pos_sigma;
}

void GraphManager::initializeImuNoiseModels(
      const ImuCalibration& imu_calib,
      const ImuInitialization& imu_init)
{
  // prior bias
  Eigen::Matrix<double, 6, 1> prior_bias_sigmas;
  prior_bias_sigmas.head<3>().setConstant(imu_init.acc_bias_sigma);
  prior_bias_sigmas.tail<3>().setConstant(imu_init.omega_bias_sigma);
  imu_bias_prior_noise_ =
      gtsam::noiseModel::Diagonal::Sigmas(prior_bias_sigmas);

  // velocity prior
  velocity_prior_noise_ =
      gtsam::noiseModel::Isotropic::Sigma(3, imu_init.velocity_sigma);

  // Create the constant zero velocity noise model.
  static const double kZeroVelocitySigma = 1.0e-3;  // [m/s]
  zero_velocity_prior_noise_ = gtsam::noiseModel::Isotropic::Sigma(3, kZeroVelocitySigma);

  // Calculate the measurement and integration noise covariances.
  // Note that the covariances are specified in continous-time.
  preintegration_params_ = PreintegratedImuMeasurements::Params::MakeSharedD(
      -imu_calib.gravity_magnitude);
  preintegration_params_->gyroscopeCovariance =
      std::pow(imu_calib.gyro_noise_density, 2.0) * Eigen::Matrix3d::Identity();
  preintegration_params_->accelerometerCovariance =
      std::pow(imu_calib.acc_noise_density, 2.0) * Eigen::Matrix3d::Identity();
  preintegration_params_->integrationCovariance =
      std::pow(imu_calib.imu_integration_sigma, 2.0) * Eigen::Matrix3d::Identity();
  preintegration_params_->biasAccCovariance =
      std::pow(imu_calib.acc_bias_random_walk_sigma, 2.0) * Eigen::Matrix3d::Identity();
  preintegration_params_->biasOmegaCovariance =
      std::pow(imu_calib.gyro_bias_random_walk_sigma, 2.0) * Eigen::Matrix3d::Identity();
  preintegration_params_->biasAccOmegaInt.setZero();
  preintegration_params_->omegaCoriolis = imu_calib.omega_coriolis;
  preintegration_params_->use2ndOrderCoriolis = false;
}

void GraphManager::reset()
{
  VLOG(3) << "Backend: Graph Manager reset";
  new_factors_.resize(0);
  new_values_.clear();
  new_smart_factors_.clear();
  smart_factors_map_.clear();
  last_added_state_index_ = -1;
}

// -----------------------------------------------------------------------------
void GraphManager::addLandmark(svo::Point& point)
{
  CHECK(!point.in_ba_graph_);

  // Add the initial landmark position expressed in the world frame.
  new_values_.insert(gtsam::Symbol('l', point.id()), gtsam::Point3(point.pos()));

  // add observations to smart-factor
  for(const KeypointIdentifier& obs : point.obs_)
  {
    if(const FramePtr frame = obs.frame.lock())
    {
      CHECK_GE(frame->in_ba_graph_vec_.size(), obs.keypoint_index_);
      CHECK(frame->landmark_vec_[obs.keypoint_index_]->id() == point.id()) << "Feature observation inconsistent!";
      CHECK(!frame->in_ba_graph_vec_[obs.keypoint_index_]) << "Feature is already in ba graph!";
      const Eigen::Vector2d uv = vk::project2(frame->f_vec_.col(obs.keypoint_index_));
      const size_t level = frame->level_vec_(obs.keypoint_index_);
      new_factors_.push_back(
            boost::make_shared<ProjectionFactor>(
              uv, px_noise_pyr_.at(level),
              gtsam::Symbol('x', frame->bundleId()),
              gtsam::Symbol('l', point.id()),
              cam_calib_,
              gtsam::Pose3(frame->T_imu_cam().getTransformationMatrix())));
      frame->in_ba_graph_vec_[obs.keypoint_index_] = true;
      VLOG(100) << "Reproj. error = "
                << dynamic_cast<ProjectionFactor*>(new_factors_.back().get())->unwhitenedError(new_values_).norm();

    }
  }

  VLOG(40) << "Added new smart landmark for point " << point.id();
}

void GraphManager::addObservationToLandmark(
    svo::Frame& frame,
    const size_t keypoint_index)
{
  CHECK(!frame.in_ba_graph_vec_[keypoint_index]);
  const svo::Point& point = *CHECK_NOTNULL(frame.landmark_vec_.at(keypoint_index).get());
  CHECK(point.in_ba_graph_);

  const Eigen::Vector2d uv = vk::project2(frame.f_vec_.col(keypoint_index));
  const size_t level = frame.level_vec_(keypoint_index);
  new_factors_.push_back(
        boost::make_shared<ProjectionFactor>(
          uv, px_noise_pyr_.at(level),
          gtsam::Symbol('x', frame.bundleId()),
          gtsam::Symbol('l', point.id()),
          cam_calib_,
          gtsam::Pose3(frame.T_imu_cam().getTransformationMatrix())));
  frame.in_ba_graph_vec_[keypoint_index] = true;

  VLOG(40) << "Added observation to point " << point.id();
}

// -----------------------------------------------------------------------------
void GraphManager::addSmartLandmark(svo::Point& point)
{
  CHECK(!point.in_ba_graph_);

  // TODO: We need to modify the Smart factor to take a different Camera-Body
  // transformation for every observation to generalize to camera-rig.
  gtsam::Pose3 T_imu_cam(point.getSeedFrame()->T_imu_cam().getTransformationMatrix());
  SmartFactor::shared_ptr new_factor(
        new SmartFactor(smart_noise_, cam_calib_, T_imu_cam, *smart_factor_params_));

  // add observations to smart-factor
  for(const KeypointIdentifier& obs : point.obs_)
  {
    if(const FramePtr frame = obs.frame.lock())
    {
      CHECK(!frame->in_ba_graph_vec_[obs.keypoint_index_]) << "Point is already in ba graph!";
      const Eigen::Vector2d uv = vk::project2(frame->f_vec_.col(obs.keypoint_index_));
      new_factor->add(uv, frame->bundleId());
      frame->in_ba_graph_vec_[obs.keypoint_index_] = true;
    }
  }

  // add new factor: Todo, add to optimization problem.
  new_smart_factors_.insert(std::make_pair(point.id(), new_factor));
  smart_factors_map_.insert(std::make_pair(point.id(), std::make_pair(new_factor, -1)));
  point.in_ba_graph_ = true;

  VLOG(40) << "Added new smart landmark for point " << point.id();
}

void GraphManager::addObservationToSmartLandmark(
    svo::Frame& frame,
    const size_t keypoint_index)
{
  CHECK(!frame.in_ba_graph_vec_[keypoint_index]);
  const svo::Point& point = *CHECK_NOTNULL(frame.landmark_vec_.at(keypoint_index).get());
  CHECK(point.in_ba_graph_);

  // Update existing smart-factor.
  auto smart_factors_map_it = smart_factors_map_.find(point.id());
  CHECK(smart_factors_map_it != smart_factors_map_.end())
      << "Tried to add a new observation to an existing landmark that is not in graph";

  const Eigen::Vector2d uv = vk::project2(frame.f_vec_.col(keypoint_index));

  SmartFactor::shared_ptr old_factor = smart_factors_map_it->second.first;
  SmartFactor::shared_ptr new_factor = boost::make_shared<SmartFactor>(*old_factor);
  new_factor->add(uv, gtsam::Symbol('x', frame.bundleId()));

  // update the factor
  if(smart_factors_map_it->second.second != -1)
  {
    // if slot is still -1, it means that the factor is not inserted yet in the graph
    new_smart_factors_.insert(std::make_pair(point.id(), new_factor));
  }
  smart_factors_map_it->second.first = new_factor;
  frame.in_ba_graph_vec_[keypoint_index] = true;

  VLOG(40) << "Added smart observation to point " << point.id();
}

// -----------------------------------------------------------------------------
void GraphManager::addCombinedImuFactor(
    const NFramesId& from_id,
    const NFramesId& to_id,
    const Eigen::Vector3d& acc_bias,
    const Eigen::Vector3d& gyr_bias,
    const double to_timestamp_sec,
    const ImuMeasurements& imu_measurements)
{
  // Create the preintegrated measurement.
  PreintegratedImuMeasurements pim(
        preintegration_params_,
        gtsam::imuBias::ConstantBias(acc_bias, gyr_bias));

  auto it = imu_measurements.rbegin();
  auto it_plus = imu_measurements.rbegin();
  for(++it_plus; it != imu_measurements.rend(); ++it, ++it_plus)
  {
    double dt = 0.0;
    if(it_plus == imu_measurements.rend()) // only for newest measurement
      dt = to_timestamp_sec - it->timestamp_;
    else
      dt = it_plus->timestamp_ - it->timestamp_;
    pim.integrateMeasurement(it->linear_acceleration_, it->angular_velocity_, dt);
  }

  new_factors_.push_back(
        boost::make_shared<gtsam::CombinedImuFactor>(
          gtsam::Symbol('x',from_id), gtsam::Symbol('v',from_id),
          gtsam::Symbol('x',to_id), gtsam::Symbol('v',to_id),
          gtsam::Symbol('b',from_id), gtsam::Symbol('b',to_id),
          pim));

  VLOG(40) << "Integrated " << imu_measurements.size() << " measurements"
           << " between frame " << from_id << " and " << to_id;
}

void GraphManager::addPosePriorFactor(
    const NFramesId& state_index,
    const Transformation& T_W_B)
{
  // Rotate initial pose uncertainty in world coordinates to body coordinates.
  Eigen::Matrix3d R_W_B = T_W_B.getRotationMatrix();
  Eigen::Matrix3d R_B_W = R_W_B.transpose();
  pose_prior_covariance_.topLeftCorner(3,3) =
      R_B_W * pose_prior_covariance_.topLeftCorner(3,3) * R_W_B;

  // Add pose prior.
  gtsam::SharedNoiseModel noise_init_pose =
      gtsam::noiseModel::Gaussian::Covariance(pose_prior_covariance_);
  new_factors_.push_back(
        boost::make_shared<gtsam::PriorFactor<gtsam::Pose3> >(
          gtsam::Symbol('x', state_index),
          gtsam::Pose3(T_W_B.getTransformationMatrix()),
          noise_init_pose));

  VLOG(40) << "Priors - Rotation Cov: \n" << pose_prior_covariance_.topLeftCorner(3,3);
  VLOG(40) << "Priors - Position Cov: \n" << pose_prior_covariance_.bottomRightCorner(3,3);

}

void GraphManager::addBiasPriorFactor(
    const NFramesId& state_index,
    const Eigen::Vector3d& acc_bias,
    const Eigen::Vector3d& gyr_bias)
{
  new_factors_.push_back(
        boost::make_shared<gtsam::PriorFactor<gtsam::imuBias::ConstantBias>>(
          gtsam::Symbol('b', state_index),
          gtsam::imuBias::ConstantBias(acc_bias, gyr_bias),
          imu_bias_prior_noise_));

  VLOG(40) << "Priors - Bias Mean: Acc (" << acc_bias.transpose()
           << "), Gyro ( " << gyr_bias.transpose() << ")";
  VLOG(40) << "Priors - Bias Sigmas: " << imu_bias_prior_noise_->sigmas().transpose();
}

void GraphManager::addVelocityPriorFactor(
    const NFramesId& state_index,
    const Eigen::Vector3d& W_v)
{
  new_factors_.push_back(
        boost::make_shared<gtsam::PriorFactor<gtsam::Vector3>>(
          gtsam::Symbol('v', state_index),
          W_v,
          velocity_prior_noise_));

  VLOG(40) << "Priors - Velocity";
}

void GraphManager::addZeroVelocityPriorFactor(
    const NFramesId& state_index)
{
  new_factors_.push_back(
        boost::make_shared<gtsam::PriorFactor<gtsam::Vector3>>(
          gtsam::Symbol('v', state_index),
          Eigen::Vector3d::Zero(),
          zero_velocity_prior_noise_));
}

void GraphManager::augmentStateWithPose(
    const NFramesId& state_index,
    const Transformation& T_W_B)
{
  VLOG(4) << "Backend: Add new pose: " << state_index;
  new_values_.insert(
        gtsam::Symbol('x', state_index), gtsam::Pose3(T_W_B.getTransformationMatrix()));
  last_added_state_index_ = state_index;
}

void GraphManager::augmentStateWithVelocityAndBias(
    const NFramesId& state_index,
    const Eigen::Vector3d& W_v,
    const Eigen::Vector3d& acc_bias,
    const Eigen::Vector3d& gyr_bias)
{
  new_values_.insert(
        gtsam::Symbol('v', state_index), W_v);
  new_values_.insert(
        gtsam::Symbol('b', state_index), gtsam::imuBias::ConstantBias(acc_bias, gyr_bias));
}

void GraphManager::getUpdatesCopy(
    gtsam::NonlinearFactorGraph& new_factors,
    gtsam::Values& new_states,
    std::vector<size_t>& delete_indices,
    std::vector<int>& smart_factor_point_ids)
{
  for(auto& s : new_smart_factors_)
  {
    new_factors.push_back(s.second);
    smart_factor_point_ids.push_back(s.first);
    const auto& it = smart_factors_map_.find(s.first);
    if(it->second.second != -1) // get current slot
      delete_indices.push_back(it->second.second);
  }
  new_factors.push_back(new_factors_.begin(), new_factors_.end());
  new_states = new_values_;
  new_smart_factors_.clear();
  new_factors_.resize(0);
  new_values_.clear();
}

void GraphManager::updateFactorSlots(
    const gtsam::FastVector<size_t>& new_slots,
    const std::vector<int>& smart_factor_point_ids)
{
  // update slots of new inserted indices:
  for(size_t i=0; i<smart_factor_point_ids.size(); ++i)
  {
    const auto& it = smart_factors_map_.find(smart_factor_point_ids.at(i));
    it->second.second = new_slots.at(i);
  }
}

} // namespace svo


/* Copyright 2017 The sfcpp Authors. All Rights Reserved.

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

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

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

#pragma once

#include <latex/tikz/TikzElement.hpp>

#include <Eigen/Dense>

#include <string>

namespace sfcpp {
namespace latex {
namespace tikz {

class TikzCoordinatePlot : public TikzElement {
  Eigen::MatrixXd coordinates;  // a 2 x n matrix with x and y coordinates
  std::string legend;

 public:
  TikzCoordinatePlot(Eigen::MatrixXd coordinates, std::string legend);
  TikzCoordinatePlot(TikzCoordinatePlot const &other);

  virtual std::string getCode() const;
  virtual std::shared_ptr<TikzElement> clone() const;
  virtual void setCoordinateSystem(
      std::shared_ptr<sfcpp::geo::CoordinateSystem> const &coordinateSystem);
};

} /* namespace tikz */
} /* namespace latex */
} /* namespace sfcpp */


#include "astar.cpp"
#include <Eigen/Dense>
#include "forward_kinematics.h"
#include <vector>
#include <stdio.h>
int main(int argc, char** argv )
{
  // std::cout << "Start Program\n";

  double x;
  double y;
  double z;
  double alpha;
  for(int i = 1; i < argc; i++){
    char* c = argv[i];

    if(i == 1){
      x = atof(c);
    }
    if(i == 2){
      y = atof(c);
    }
    if(i == 3){
      z = atof(c);
    }
    if(i == 4){
      alpha = atof(c);
    }
  }

  Astar planner;

  int start[3] = {49,49,49};

  //  fk(&x, &z, &alpha, 99,99,49,planner.numticks);
  double target[3] = {x,z,alpha};
  double targett0 = atan2(y,x);
   // printf("Target: %f, %f, %f\n", x, z, alpha);
  std::vector<PState*> path = planner.run(start, target, targett0);
  int n = path.size();
  for(int i = 0; i < n; i++){
    PState* s = path.at(i);
    double t0 = targett0;
    double t1 = tick_to_radians(s->state[0],planner.numticks);
    double t2 = tick_to_radians(s->state[1],planner.numticks);
    double t3 = tick_to_radians(s->state[2],planner.numticks);
    double t4 = 0;

    printf("%f %f %f %f %f\n", t0, t1,t2,t3,t4);
  }

  return 0;
}


#pragma once

#include <Eigen/Core>
#include <array>
#include <doubleCCD/Rational.hpp>
#include <igl/Timer.h>
#include <igl/predicates/predicates.h>
#include <vector>
namespace doubleccd {

typedef Eigen::Matrix<Rational, 3, 1> Vector3r;
typedef Eigen::Matrix<double, 3, 1> Vector3d;
typedef Eigen::Matrix<double, 2, 1> Vector2d;
static const int COPLANAR = -1;
static const int INTERSECTED = 1;
static const int NOT_INTERSECTED1 = 2;
static const int NOT_INTERSECTED2 = 3;
static const Vector3d ORIGIN = Vector3d(0, 0, 0);

static const int BI_DEGE_PLANE = 1;
static const int BI_DEGE_XOR_02 = 2;
static const int BI_DEGE_XOR_13 = 3;

class bilinear {
public:
    // v0, v1 are vertices of one triangle, v2, v3 are the vertices of another
    // one.
    bilinear(
        const Vector3d& v0,
        const Vector3d& v1,
        const Vector3d& v2,
        const Vector3d& v3);
    bool is_degenerated;
    std::vector<std::array<int, 3>> facets;
    std::array<int, 2> phi_f = { { 2, 2 } };
    std::array<Vector3d, 4> v;
};
template <typename V1, typename V2> Vector3r cross(const V1& v1, const V2& v2)
{
    Vector3r res;
    res[0] = v1[1] * v2[2] - v1[2] * v2[1];
    res[1] = v1[2] * v2[0] - v1[0] * v2[2];
    res[2] = v1[0] * v2[1] - v1[1] * v2[0];

    return res;
}

inline int orient_3d(
    const Vector3d& p, const Vector3d& q, const Vector3d& r, const Vector3d& s)
{
    igl::predicates::exactinit();
    return (int)igl::predicates::orient3d(p, q, r, s);
}
inline int orient_2d(const Vector2d& p, const Vector2d& q, const Vector2d& r)
{
    igl::predicates::exactinit();
    return (int)igl::predicates::orient2d(p, q, r);
}

Rational func_g(
    const Vector3r& x,
    const std::array<Vector3d, 4>& corners,
    const std::array<int, 3>& indices);

Rational phi(const Vector3d x, const std::array<Vector3d, 4>& corners);
Rational phi(const Vector3r x, const std::array<Vector3d, 4>& corners);

void get_tet_phi(bilinear& bl);
// bool XOR(const bool a, const bool b)
//{
//    if (a && b)
//        return false;
//    if (!a && !b)
//        return false;
//    return true;
//}

// accept 0,1,2,3 as inputs
bool int_seg_XOR(const int a, const int b);

// accept -1,0,1,2,3 as inputs
int int_ray_XOR(const int a, const int b);
template <typename V> void print(const V& v)
{
    std::cout << v[0] << " " << v[1] << " " << v[2] << std::endl;
}

void write(const Vector3d& v, std::ostream& out);
Vector3d read(std::istream& in);

bool segment_segment_intersection_2d(
    const Vector2d& s0,
    const Vector2d& e0,
    const Vector2d& s1,
    const Vector2d& e1);

// 0 not intersected, 1 intersected, 2 s0 on segment
// can deal with degenerated cases
int ray_segment_intersection(
    const Vector3d& s0,
    const Vector3d& e0,
    const Vector3d& dir0,
    const Vector3d& s1,
    const Vector3d& e1);

// this function can also tell us if they are parallel and overlapped
// and also tell us if the parallel case has seg-seg overlapping:

bool is_triangle_degenerated(
    const Vector3d& t1, const Vector3d& t2, const Vector3d& t3);

bool same_point(const Vector3d& p1, const Vector3d& p2);

template <typename T>
static bool orient3D_LPI_prefilter_multiprecision(
    const T& px,
    const T& py,
    const T& pz,
    const T& qx,
    const T& qy,
    const T& qz,
    const T& rx,
    const T& ry,
    const T& rz,
    const T& sx,
    const T& sy,
    const T& sz,
    const T& tx,
    const T& ty,
    const T& tz,
    T& a11,
    T& a12,
    T& a13,
    T& d,
    T& n,
    const std::function<int(T)>& checker)
{

    a11 = (px - qx);
    a12 = (py - qy);
    a13 = (pz - qz);
    T a21(sx - rx);
    T a22(sy - ry);
    T a23(sz - rz);
    T a31(tx - rx);
    T a32(ty - ry);
    T a33(tz - rz);
    T a2233((a22 * a33) - (a23 * a32));
    T a2133((a21 * a33) - (a23 * a31));
    T a2132((a21 * a32) - (a22 * a31));
    d = (((a11 * a2233) - (a12 * a2133)) + (a13 * a2132));
    int flag1 = checker(d);
    if (flag1 == -2 || flag1 == 0) {
        return false; // not enough precision
    }
    T px_rx(px - rx);
    T py_ry(py - ry);
    T pz_rz(pz - rz);

    n
        = ((((py_ry)*a2133) - ((px_rx)*a2233))
           - ((pz_rz)*a2132)); // caution, this is -n actually

    a11 = a11 * n;
    a12 = a12 * n;
    a13 = a13 * n;
    return true;
}
template <typename T>
static int orient3D_LPI_postfilter_multiprecision(
    const T& a11,
    const T& a12,
    const T& a13,
    const T& d,
    const T& px,
    const T& py,
    const T& pz,
    const T& ax,
    const T& ay,
    const T& az,
    const T& bx,
    const T& by,
    const T& bz,
    const T& cx,
    const T& cy,
    const T& cz,
    const std::function<int(T)>& checker)
{

    T px_cx(px - cx);
    T py_cy(py - cy);
    T pz_cz(pz - cz);

    T d11((d * px_cx) + (a11));
    T d21(ax - cx);
    T d31(bx - cx);
    T d12((d * py_cy) + (a12));
    T d22(ay - cy);
    T d32(by - cy);
    T d13((d * pz_cz) + (a13));
    T d23(az - cz);
    T d33(bz - cz);

    T d2233(d22 * d33);
    T d2332(d23 * d32);
    T d2133(d21 * d33);
    T d2331(d23 * d31);
    T d2132(d21 * d32);
    T d2231(d22 * d31);

    T det(
        d11 * (d2233 - d2332) - d12 * (d2133 - d2331) + d13 * (d2132 - d2231));

    int flag2 = checker(det);
    if (flag2 == -2) {
        return 100; // not enough precision, only happens when using floating
                    // points
    }
    if (flag2 == 1) {
        if (d > 0) {
            return 1;
        }
        if (d < 0) {
            return -1;
        }
    }
    if (flag2 == -1) {
        if (d > 0) {
            return -1;
        }
        if (d < 0) {
            return 1;
        }
    }
    return 0;
}

static const std::function<int(Rational)> check_rational = [](Rational v) {
    if (v.get_sign() > 0)
        return 1;
    if (v.get_sign() < 0)
        return -1;
    return 0;
};

int point_inter_triangle(
    const Vector3d& pt,
    const Vector3d& t1,
    const Vector3d& t2,
    const Vector3d& t3,
    const bool& dege,
    const bool halfopen);

// already know lpi exist;
// 0 not intersected, 1 intersect open triangle, 2 shoot on edge, 3 shoot on
// edge t2-t3
int is_line_cut_triangle(
    const Vector3d& e0,
    const Vector3d& e1,
    const Vector3d& t1,
    const Vector3d& t2,
    const Vector3d& t3,
    const bool halfopen);
int seg_triangle_inter_return_t(
    const Vector3d& e0,
    const Vector3d& e1,
    const Vector3d& t1,
    const Vector3d& t2,
    const Vector3d& t3,
    Rational& t);
// if a line (going across pt, pt+dir) intersects triangle
// triangle is not degenerated
// 0 not intersected, 1 intersected, 3 intersected t2-t3 edge

// we check if triangle intersect segment,
// this function is used in cube edge--prism tri and cube edge--bilinear tri
// if halfopen= true, can tell us if intersect the edge t2-t3
// 0 not intersected, 1 intersected, 2 intersect edge, 3 intersect t2-t3 edge
int segment_triangle_intersection(
    const Vector3d& e0,
    const Vector3d& e1,
    const Vector3d& t1,
    const Vector3d& t2,
    const Vector3d& t3,
    const bool halfopen);

int segment_triangle_intersection(
    const Vector3r& e0,
    const Vector3r& e1,
    const Vector3r& t1,
    const Vector3r& t2,
    const Vector3r& t3,
    const bool halfopen);
// 0 no intersection, 1 intersect, 2 point on triangle, 3 point or ray go to on
// t2-t3 edge, -1 shoot on border
int ray_triangle_intersection(
    const Vector3d& pt,
    const Vector3d& pt1,
    const Vector3d& dir,
    const Vector3d& t1,
    const Vector3d& t2,
    const Vector3d& t3,
    const bool halfopen);
void tri_bilinear(
    const bilinear& bl, int n, std::vector<std::array<Vector3r, 3>>& patch);
void save_obj(
    const std::string& name, const std::vector<std::array<Vector3r, 3>>& tris);
bool seg_discrete_bilinear_intersection(
    const bilinear& bl, int n, const Vector3d& s0, const Vector3d& s1);
int lpi_in_triangle(
    const Vector3d& p,
    const Vector3d& q,
    const Vector3d& r,
    const Vector3d& s,
    const Vector3d& t,
    bool halfopen);
bool lpi_rational(
    const Vector3d& p,
    const Vector3d& q,
    const Vector3d& r,
    const Vector3d& s,
    const Vector3d& t,
    Rational& a11,
    Rational& a12,
    Rational& a13,
    Rational& d);
void compare_lpi_results();
double print_phi_time();
} // namespace doubleccd


#include "cnn/nodes.h"
#include "cnn/cnn.h"
#include "cnn/training.h"
#include "cnn/timing.h"
#include "cnn/rnn.h"
#include "cnn/gru.h"
#include "cnn/lstm.h"
#include "cnn/dglstm.h"
#include "cnn/dict.h"
#include "cnn/expr.h"
#include "cnn/cnn-helper.h"
#include "cnn/expr-xtra.h"
#include <iostream>
#include <fstream>
#include <sstream>

#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>

using namespace std;
using namespace cnn;

unsigned int LAYERS = 2;
unsigned int INPUT_DIM = 8;  //256
unsigned int HIDDEN_DIM = 24;  // 1024
unsigned int VOCAB_SIZE = 0;

int verbose = 0;

cnn::Dict d;
int kSOS;
int kEOS;

template <class Builder>
struct RNNLanguageModel {
  LookupParameters* p_c;
  vector<Parameters*> p_R;
  vector<Parameters*> p_bias;
  Parameters* p_cls, *p_cls_bias;
  unsigned int ncls;
  Builder builder;
  vector<int> clssize;
  vector<long> word2cls;
  vector<long> acc_cls2size;
  vector<long> dict_wrd_id2within_class_id;
  explicit RNNLanguageModel(const vector<int>& cls2nbrwords, /// #words for each class, class starts from 0
      const vector<long> & acc_cls2size, /// the accumulated class size
      const vector<long>& word2cls,
      const vector<long>& dict_wrd_id2within_class_id,
      Model& model) : builder(LAYERS, vector<unsigned>{INPUT_DIM, HIDDEN_DIM}, &model),
      clssize(cls2nbrwords), word2cls(word2cls), acc_cls2size(acc_cls2size), dict_wrd_id2within_class_id(dict_wrd_id2within_class_id)
  {
      unsigned int n_cls = clssize.size();
      p_c = model.add_lookup_parameters(VOCAB_SIZE, {INPUT_DIM});
      p_cls = model.add_parameters({ n_cls, HIDDEN_DIM });
      p_cls_bias = model.add_parameters({ n_cls});
      for (size_t id = 0; id < n_cls; id++)
      {
          unsigned int  clssize = cls2nbrwords[id];
          p_R.push_back(model.add_parameters({ clssize, HIDDEN_DIM }));
          p_bias.push_back(model.add_parameters({ clssize }));
      }
  }

  // return Expression of total loss
  Expression BuildLMGraph(const vector<int>& sent, ComputationGraph& cg) {
    const unsigned slen = sent.size() - 1;
    builder.new_graph(cg);  // reset RNN builder for new graph
    builder.start_new_sequence();

    Expression i_cls, i_cls_bias;
    i_cls = parameter(cg, p_cls);
    i_cls_bias = parameter(cg, p_cls_bias);

    vector<Expression> i_R, i_bias;
    for (auto& p : p_R)
        i_R.push_back(parameter(cg, p)); // hidden -> word rep parameter
    for (auto& p: p_bias)
        i_bias.push_back(parameter(cg, p));  // word bias

    vector<Expression> errs;
    for (unsigned t = 0; t < slen; ++t) {
      Expression i_x_t = lookup(cg, p_c, sent[t]);
      // y_t = RNN(x_t)
      Expression i_y_t = builder.add_input(i_x_t);
      int cls_id = word2cls[sent[t + 1]];
      Expression i_r_t =  i_bias[cls_id] + i_R[cls_id] * i_y_t;
      Expression i_c_t = i_cls_bias + i_cls * i_y_t;
      if (verbose)
		  display_value(i_r_t, cg, "response at " + t);

#if 1
      Expression i_err_cls = pick(log_softmax(i_c_t), cls_id);
      Expression i_err_prb = pick(log_softmax(i_r_t), dict_wrd_id2within_class_id[sent[t + 1]]);
      errs.push_back(i_err_cls + i_err_prb);
#if 0
      Expression i_ydist = softmax(i_r_t);
      i_ydist = log(i_ydist)
      errs.push_back(pick(i_ydist, sent[t+1]));
#endif
#else
      ppl turns to be 0 at epoch 4, so there is a bug
      Expression i_err_cls = pickneglogsoftmax(i_c_t, cls_id);
      Expression i_err_prb = pickneglogsoftmax(i_r_t, dict_wrd_id2within_class_id[sent[t + 1]]);
      errs.push_back(i_err_cls + i_err_prb);
#endif
    }
    Expression i_nerr = sum(errs);
#if 1
    return -i_nerr;
#else
    return i_nerr;
#endif
  }

};

template <class LM_t>
void train(Model &model, LM_t &lm,
    const vector<vector<int>>& training,
    const vector<vector<int>>& dev,
    Trainer *sgd, const string& fname,
    bool randomSample)
{
    cnn::real best = 9e+99;
    unsigned report_every_i = 50;
    unsigned dev_every_i_reports = 500;
    unsigned si = training.size();
    vector<unsigned> order(training.size());
    for (unsigned i = 0; i < order.size(); ++i) order[i] = i;
    bool first = true;
    int report = 0;
    unsigned lines = 0;
    unsigned total_epoch = 40;

    ofstream out(fname, ofstream::out);
    boost::archive::text_oarchive oa(out);
    oa << model;
    out.close();

    size_t i_epoch = 0;
    while (sgd->epoch < total_epoch) {
        Timer iteration("completed in");
        cnn::real loss = 0;
        unsigned chars = 0;
        for (unsigned i = 0; i < report_every_i; ++i) {
            if (si == training.size()) {
                si = 0;
                if (first) { first = false; }
                else { sgd->update_epoch(); }
                cerr << "**SHUFFLE\n";
                shuffle(order.begin(), order.end(), *rndeng);
            }

            // build graph for this instance
            ComputationGraph cg;
            auto& sent = training[order[si]];
            chars += sent.size() - 1;
            ++si;
            lm.BuildLMGraph(sent, cg);
            loss += as_scalar(cg.forward());
            cg.backward();
            sgd->update();
            ++lines;
        }
        sgd->status();
        cerr << " report = " << report << " E = " << (loss / chars) << " ppl=" << exp(loss / chars) << ' ';
        iteration.WordsPerSecond(chars);
        // show score on dev data?
        report++;
        if (report % dev_every_i_reports == 0) {
            cnn::real dloss = 0;
            int dchars = 0;
            for (auto& sent : dev) {
                ComputationGraph cg;
                lm.BuildLMGraph(sent, cg);
                dloss += as_scalar(cg.forward());
                dchars += sent.size() - 1;
            }
            if (dloss < best) {
                best = dloss;
                ofstream out(fname, ofstream::out);
                boost::archive::text_oarchive oa(out);
                oa << model;
                out.close();
            }
            else{
                sgd->eta *= 0.5;
            }
            cerr << "\n***TEST E = " << (dloss / dchars) << " ppl=" << exp(dloss / dchars) << ' ';
        }

        i_epoch++;
    }
}

template <class LM_t>
void testcorpus(Model &model, LM_t &lm,
    const vector<vector<int>>& dev)
{
    unsigned lines = 0;
    cnn::real dloss = 0;
    int dchars = 0;
    for (auto& sent : dev) {
        ComputationGraph cg;
        lm.BuildLMGraph(sent, cg);
        dloss += as_scalar(cg.forward());
        dchars += sent.size() - 1;
    }

    cerr << "\n***DEV [epoch=" << (lines / (cnn::real)dev.size()) << "] E = " << (dloss / dchars) << " ppl=" << exp(dloss / dchars) << ' ';
}

void initialise(Model &model, const string &filename)
{
    cerr << "Initialising model parameters from file: " << filename << endl;
    ifstream in(filename, ifstream::in);
    boost::archive::text_iarchive ia(in);
    ia >> model;
}

/**
dict_wrd_id2within_class_id : the dictionary word id to the id inside a class
*/
void load_word2cls_fn(string word2clsfn, Dict& sd, std::vector<long>& wrd2cls, std::vector<long>& dict_wrd_id2within_class_id)
{
    ifstream in(word2clsfn);
    string line;

    wrd2cls.resize(sd.size());
    dict_wrd_id2within_class_id.resize(sd.size());
    map<int, int> cls2acccnt; /// the count for each class so far
    while (getline(in, line)) {

        std::istringstream in(line);
        std::string word;
        string cls;

        in >> word;
        in >> cls;

        int icls = boost::lexical_cast<int>(cls) - 1;
        int wridx = sd.Convert(word);

        wrd2cls[wridx] = icls;
        if (cls2acccnt.find(icls) == cls2acccnt.end())
            cls2acccnt[icls] = 1;
        else
            cls2acccnt[icls] += 1;
        dict_wrd_id2within_class_id[wridx] = cls2acccnt[icls] - 1;
    }
    in.close();
}

/**
acc_cls2size : the accumulated class size
*/
void load_clssize_fn(string clsszefn, std::vector<int> & cls2size, std::vector<long>& acc_cls2size)
{
    ifstream in(clsszefn);
    string line;

    cls2size.clear();
    acc_cls2size.clear();
    int idx = 1;

    while (getline(in, line)) {

        std::istringstream in(line);
        std::string cls;
        string sze;

        in >> cls;
        in >> sze;

        int icls = boost::lexical_cast<int>(cls) - 1;
        if (icls != idx - 1)
            throw("class id should start from 1 and then consecuitively increasing with step 1");

        cls2size.push_back(boost::lexical_cast<int>(sze));
        if (acc_cls2size.size() == 0)
            acc_cls2size.push_back(cls2size.back());
        else
            acc_cls2size.push_back(acc_cls2size.back() + cls2size.back());

        idx++;
    }
    in.close();
}

bool check_info_correct(Dict& sd, const std::vector<long>& wrd2cls, const std::vector<long>& dict_wrd_id2within_class_id, const std::vector<int> & cls2size, const std::vector<long>& acc_cls2size)
{
    for (auto&p : sd.GetWordList())
    {
        long wd = sd.Convert(p);
        int  cls = wrd2cls[wd];
        int clssize = cls2size[cls];
        int pos = dict_wrd_id2within_class_id[wd];
        if (pos < 0 || pos >= clssize)
        {
            cerr << "word " << p << " id " << wd << " cls " << cls << " clssize " << clssize << " pos in cls " << pos << " wrong" << endl;
            return false;
        }
    }
    return true;
}

int main(int argc, char** argv) {
  cnn::Initialize(argc, argv);

  // command line processing
  using namespace boost::program_options;
  variables_map vm;
  options_description opts("Allowed options");
  opts.add_options()
      ("help", "print help message")
      ("seed,s", value<int>()->default_value(217), "random seed number")
      ("train,t", value<string>(), "file containing training sentences")
      ("devel,d", value<string>(), "file containing development sentences.")
      ("test,T", value<string>(), "file containing testing source sentences")
      ("word2cls", value<string>(), "word2class info file")
      ("cls2size", value<string>(), "class size information file")
      ("initialise,i", value<string>(), "load initial parameters from file")
      ("parameters,p", value<string>(), "save best parameters to this file")
      ("layers,l", value<int>()->default_value(LAYERS), "use <num> layers for RNN components")
      ("hidden,h", value<int>()->default_value(HIDDEN_DIM), "use <num> dimensions for recurrent hidden states")
      ("gru", "use Gated Recurrent Unit (GRU) for recurrent structure; default RNN")
      ("lstm", "use Long Short Term Memory (GRU) for recurrent structure; default RNN")
      ("dglstm", "use depth-gated LSTM for recurrent structure; default RNN")
      ("verbose,v", "be extremely chatty")
      ("generate,g", value<bool>()->default_value(false), "generate random samples")
      ;
  store(parse_command_line(argc, argv, opts), vm);

  string flavour;
  if (vm.count("gru"))	flavour = "gru";
  else if (vm.count("lstm"))	flavour = "lstm";
  else if (vm.count("rnnem"))	flavour = "rnnem";
  else if (vm.count("dglstm")) flavour = "dglstm";
  else if (vm.count("nmn")) flavour = "nmn";
  else			flavour = "rnn";

  LAYERS = vm["layers"].as<int>();
  HIDDEN_DIM = vm["hidden"].as<int>();

  bool generateSample = false;
  generateSample = vm["generate"].as<bool>();

  string fname;
  if (vm.count("parameters")) {
      fname = vm["parameters"].as<string>();
  }
  else {
      ostringstream os;
      os << "lm"
          << '_' << LAYERS
          << '_' << HIDDEN_DIM
          << '_' << flavour
          << "-pid" << getpid() << ".params";
      fname = os.str();
  }

  cerr << "Parameters will be written to: " << fname << endl;

  if (vm.count("help") || vm.count("train") != 1 || (vm.count("devel") != 1 && vm.count("test") != 1)) {
      cout << opts << "\n";
      return 1;
  }

  kSOS = d.Convert("<s>");
  kEOS = d.Convert("</s>");
  vector<vector<int>> training, dev, test;
  string line;
  int tlc = 0;
  int ttoks = 0;

  string infile = vm["train"].as<string>();
  cerr << "Reading training data from " << infile << "...\n";

  {
    ifstream in(infile);
    assert(in);
    while(getline(in, line)) {
      ++tlc;
      training.push_back(ReadSentence(line, &d));
      ttoks += training.back().size();
      if (training.back().front() != kSOS && training.back().back() != kEOS) {
		  throw("Training sentence in %s : %d didnt start or end with <s>, </s>", infile.c_str(), tlc );
      }
    }
    cerr << tlc << " lines, " << ttoks << " tokens, " << d.size() << " types\n";
  }
  d.Freeze(); // no new word types allowed
  VOCAB_SIZE = d.size();

  if (vm.count("devel") > 0)
  {
      int dlc = 0;
      int dtoks = 0;
      string devfile = vm["devel"].as<string>();
      cerr << "Reading training data from " << devfile << "...\n";
      {
          ifstream in(devfile);
          assert(in);
          while (getline(in, line)) {
              ++dlc;
              dev.push_back(ReadSentence(line, &d));
              dtoks += dev.back().size();
			  if (dev.back().front() != kSOS && dev.back().back() != kEOS) {
				  throw("Dev sentence in %s : %d didn't start or end with <s>, </s> ", devfile.c_str(), tlc);
			  }
          }
          cerr << dlc << " lines, " << dtoks << " tokens\n";
      }
  }

  Model model;
  bool use_momentum = false;
  Trainer* sgd = nullptr;
  if (use_momentum)
    sgd = new MomentumSGDTrainer(&model);
  else
    sgd = new SimpleSGDTrainer(&model);

  std::vector<long> wrd2cls;
  std::vector<int> cls2size;
  std::vector<long> acc_cls2size;
  std::vector<long> dict_wrd_id2within_class_id;
  if (vm["word2cls"].as<string>() != "")
  {
      load_word2cls_fn(vm["word2cls"].as<string>(), d, wrd2cls, dict_wrd_id2within_class_id);
      load_clssize_fn(vm["cls2size"].as<string>(), cls2size, acc_cls2size);
      check_info_correct(d, wrd2cls, dict_wrd_id2within_class_id, cls2size, acc_cls2size);
  }
  else{
      throw std::invalid_argument("need to specify word2cls and cls2size files for word clustering information.");
  }

  if (vm.count("test") == 0)
  {
      if (vm.count("lstm")) {
          cerr << "%% Using LSTM recurrent units" << endl;
          RNNLanguageModel<LSTMBuilder> lm(cls2size, acc_cls2size, wrd2cls, dict_wrd_id2within_class_id, model);
          train(model, lm, training, dev, sgd, fname, generateSample);
      }
      else if (vm.count("dglstm")) {
          cerr << "%% Using DGLSTM recurrent units" << endl;
          RNNLanguageModel<DGLSTMBuilder> lm(cls2size, acc_cls2size, wrd2cls, dict_wrd_id2within_class_id, model);
          train(model, lm, training, dev, sgd, fname, generateSample);
      }
  }
  else
  {
      string testfile = vm["test"].as<string>();
      int dlc = 0;
      int dtoks = 0;
      cerr << "Reading training data from " << testfile << "...\n";
      {
          ifstream in(testfile);
          assert(in);
          while (getline(in, line)) {
              ++dlc;
              test.push_back(ReadSentence(line, &d));
              dtoks += test.back().size();
			  if (test.back().front() != kSOS && test.back().back() != kEOS) {
				  throw("Dev sentence in %s : %d didnt start or end with <s>, </s> ", testfile.c_str(), tlc);
			  }
          }
          cerr << dlc << " lines, " << dtoks << " tokens\n";
      }

      if (vm.count("test"))
      {
          if (vm.count("lstm")){
              cerr << "%% using LSTM recurrent units" << endl;
              RNNLanguageModel<LSTMBuilder> lm(cls2size, acc_cls2size, wrd2cls, dict_wrd_id2within_class_id, model);
              if (vm.count("initialise"))
                  initialise(model, vm["initialise"].as<string>());
              testcorpus(model, lm, test);
          }
          if (vm.count("dglstm")){
              cerr << "%% using DGLSTM recurrent units" << endl;
              RNNLanguageModel<DGLSTMBuilder> lm(cls2size, acc_cls2size, wrd2cls, dict_wrd_id2within_class_id, model);
              if (vm.count("initialise"))
                  initialise(model, vm["initialise"].as<string>());
              testcorpus(model, lm, test);
          }
      }
  }

  //RNNLanguageModel<SimpleRNNBuilder> lm(model);
  if (argc == 4) {
    string fname = argv[3];
    ifstream in(fname);
    boost::archive::text_iarchive ia(in);
    ia >> model;
  }

  delete sgd;
}