/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/*
 * @file testSymbolicBayesTree.cpp
 * @date sept 15, 2012
 * @author  Frank Dellaert
 * @author  Michael Kaess
 * @author  Viorela Ila
 */

#include <gtsam/symbolic/SymbolicBayesTree.h>
#include <gtsam/symbolic/SymbolicBayesNet.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam/symbolic/tests/symbolicExampleGraphs.h>

#include <boost/assign/list_of.hpp>
#include <boost/assign/std/vector.hpp>
#include <boost/assign/std/list.hpp>
#include <boost/range/adaptor/indirected.hpp>
using namespace boost::assign;
using boost::adaptors::indirected;

#include <CppUnitLite/TestHarness.h>
#include <gtsam/base/TestableAssertions.h>

using namespace std;
using namespace gtsam;
using namespace gtsam::symbol_shorthand;

static bool debug = false;

namespace {

  /* ************************************************************************* */
  // Helper functions for below
  template<typename KEYS>
  SymbolicBayesTreeClique::shared_ptr MakeClique(const KEYS& keys, DenseIndex nrFrontals)
  {
    return boost::make_shared<SymbolicBayesTreeClique>(
      boost::make_shared<SymbolicConditional>(
      SymbolicConditional::FromKeys(keys, nrFrontals)));
  }

  template<typename KEYS, typename CHILDREN>
  SymbolicBayesTreeClique::shared_ptr MakeClique(
    const KEYS& keys, DenseIndex nrFrontals, const CHILDREN& children)
  {
    SymbolicBayesTreeClique::shared_ptr clique =
      boost::make_shared<SymbolicBayesTreeClique>(
      boost::make_shared<SymbolicConditional>(
      SymbolicConditional::FromKeys(keys, nrFrontals)));
    clique->children.assign(children.begin(), children.end());
    for(typename CHILDREN::const_iterator child = children.begin(); child != children.end(); ++child)
      (*child)->parent_ = clique;
    return clique;
  }

}

/* ************************************************************************* */
TEST(SymbolicBayesTree, clear)
{
  SymbolicBayesTree bayesTree = asiaBayesTree;
  bayesTree.clear();

  SymbolicBayesTree expected;

  // Check whether cleared BayesTree is equal to a new BayesTree
  CHECK(assert_equal(expected, bayesTree));
}

/* ************************************************************************* */
TEST(SymbolicBayesTree, clique_structure)
{
  //        l1                  l2
  //    /    |                /  |
  // x1 --- x2 --- x3 --- x4 --- x5
  //                          \  |
  //                            l3
  SymbolicFactorGraph graph;
  graph += SymbolicFactor(X(1), L(1));
  graph += SymbolicFactor(X(1), X(2));
  graph += SymbolicFactor(X(2), L(1));
  graph += SymbolicFactor(X(2), X(3));
  graph += SymbolicFactor(X(3), X(4));
  graph += SymbolicFactor(X(4), L(2));
  graph += SymbolicFactor(X(4), X(5));
  graph += SymbolicFactor(L(2), X(5));
  graph += SymbolicFactor(X(4), L(3));
  graph += SymbolicFactor(X(5), L(3));

  SymbolicBayesTree expected;
  expected.insertRoot(
    MakeClique(list_of(X(2)) (X(3)), 2, list_of
      (MakeClique(list_of(X(4)) (X(3)), 1, list_of
        (MakeClique(list_of(X(5)) (L(2)) (X(4)), 2, list_of
          (MakeClique(list_of(L(3)) (X(4)) (X(5)), 1))))))
      (MakeClique(list_of(X(1)) (L(1)) (X(2)), 2))));

  Ordering order = list_of(X(1)) (L(3)) (L(1)) (X(5)) (X(2)) (L(2)) (X(4)) (X(3));

  SymbolicBayesTree actual = *graph.eliminateMultifrontal(order);

  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* *
Bayes Tree for testing conversion to a forest of orphans needed for incremental.
       A,B
   C|A    E|B
   D|C    F|E
   */
/* ************************************************************************* */
TEST( BayesTree, removePath )
{
  const Key _A_=A(0), _B_=B(0), _C_=C(0), _D_=D(0), _E_=E(0), _F_=F(0);

  SymbolicBayesTree bayesTreeOrig;
  bayesTreeOrig.insertRoot(
    MakeClique(list_of(_A_)(_B_), 2, list_of
      (MakeClique(list_of(_C_)(_A_), 1, list_of
        (MakeClique(list_of(_D_)(_C_), 1))))
      (MakeClique(list_of(_E_)(_B_), 1, list_of
        (MakeClique(list_of(_F_)(_E_), 1))))));

  SymbolicBayesTree bayesTree = bayesTreeOrig;

  // remove C, expected outcome: factor graph with ABC,
  // Bayes Tree now contains two orphan trees: D|C and E|B,F|E
  SymbolicFactorGraph expected;
  expected += SymbolicFactor(_A_,_B_);
  expected += SymbolicFactor(_C_,_A_);
  SymbolicBayesTree::Cliques expectedOrphans;
  expectedOrphans += bayesTree[_D_], bayesTree[_E_];

  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removePath(bayesTree[_C_], &bn, &orphans);
  SymbolicFactorGraph factors(bn);
  CHECK(assert_equal(expected, factors));
  CHECK(assert_container_equal(expectedOrphans|indirected, orphans|indirected));

  bayesTree = bayesTreeOrig;

  // remove E: factor graph with EB; E|B removed from second orphan tree
  SymbolicFactorGraph expected2;
  expected2 += SymbolicFactor(_A_,_B_);
  expected2 += SymbolicFactor(_E_,_B_);
  SymbolicBayesTree::Cliques expectedOrphans2;
  expectedOrphans2 += bayesTree[_F_];
  expectedOrphans2 += bayesTree[_C_];

  SymbolicBayesNet bn2;
  SymbolicBayesTree::Cliques orphans2;
  bayesTree.removePath(bayesTree[_E_], &bn2, &orphans2);
  SymbolicFactorGraph factors2(bn2);
  CHECK(assert_equal(expected2, factors2));
  CHECK(assert_container_equal(expectedOrphans2|indirected, orphans2|indirected));
}

/* ************************************************************************* */
TEST( BayesTree, removePath2 )
{
  SymbolicBayesTree bayesTree = asiaBayesTree;

  // Call remove-path with clique B
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removePath(bayesTree[_B_], &bn, &orphans);
  SymbolicFactorGraph factors(bn);

  // Check expected outcome
  SymbolicFactorGraph expected;
  expected += SymbolicFactor(_E_,_L_,_B_);
  CHECK(assert_equal(expected, factors));
  SymbolicBayesTree::Cliques expectedOrphans;
  expectedOrphans += bayesTree[_S_], bayesTree[_T_], bayesTree[_X_];
  CHECK(assert_container_equal(expectedOrphans|indirected, orphans|indirected));
}

/* ************************************************************************* */
TEST(BayesTree, removePath3)
{
  SymbolicBayesTree bayesTree = asiaBayesTree;

  // Call remove-path with clique T
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removePath(bayesTree[_T_], &bn, &orphans);
  SymbolicFactorGraph factors(bn);

  // Check expected outcome
  SymbolicFactorGraph expected;
  expected += SymbolicFactor(_E_, _L_, _B_);
  expected += SymbolicFactor(_T_, _E_, _L_);
  CHECK(assert_equal(expected, factors));
  SymbolicBayesTree::Cliques expectedOrphans;
  expectedOrphans += bayesTree[_S_], bayesTree[_X_];
  CHECK(assert_container_equal(expectedOrphans|indirected, orphans|indirected));
}

void getAllCliques(const SymbolicBayesTree::sharedClique& subtree, SymbolicBayesTree::Cliques& cliques)  {
  // Check if subtree exists
  if (subtree) {
    cliques.push_back(subtree);
    // Recursive call over all child cliques
    for(SymbolicBayesTree::sharedClique& childClique: subtree->children) {
      getAllCliques(childClique,cliques);
    }
  }
}

/* ************************************************************************* */
TEST( BayesTree, shortcutCheck )
{
  const Key _A_=6, _B_=5, _C_=4, _D_=3, _E_=2, _F_=1, _G_=0;
  SymbolicFactorGraph chain = list_of
    (SymbolicFactor(_A_))
    (SymbolicFactor(_B_, _A_))
    (SymbolicFactor(_C_, _A_))
    (SymbolicFactor(_D_, _C_))
    (SymbolicFactor(_E_, _B_))
    (SymbolicFactor(_F_, _E_))
    (SymbolicFactor(_G_, _F_));
  Ordering ordering(list_of(_G_)(_F_)(_E_)(_D_)(_C_)(_B_)(_A_));
  SymbolicBayesTree bayesTree = *chain.eliminateMultifrontal(ordering);

  //bayesTree.saveGraph("BT1.dot");

  SymbolicBayesTree::sharedClique rootClique = bayesTree.roots().front();
  //rootClique->printTree();
  SymbolicBayesTree::Cliques allCliques;
  getAllCliques(rootClique,allCliques);

  for(SymbolicBayesTree::sharedClique& clique: allCliques) {
    //clique->print("Clique#");
    SymbolicBayesNet bn = clique->shortcut(rootClique);
    //bn.print("Shortcut:\n");
    //cout << endl;
  }

  // Check if all the cached shortcuts are cleared
  rootClique->deleteCachedShortcuts();
  for(SymbolicBayesTree::sharedClique& clique: allCliques) {
    bool notCleared = clique->cachedSeparatorMarginal().is_initialized();
    CHECK( notCleared == false);
  }
  EXPECT_LONGS_EQUAL(0, (long)rootClique->numCachedSeparatorMarginals());

//  for(SymbolicBayesTree::sharedClique& clique: allCliques) {
//    clique->print("Clique#");
//    if(clique->cachedShortcut()){
//      bn = clique->cachedShortcut().get();
//      bn.print("Shortcut:\n");
//    }
//    else
//      cout << "Not Initialized" << endl;
//    cout << endl;
//  }
}

/* ************************************************************************* */
TEST( BayesTree, removeTop )
{
  SymbolicBayesTree bayesTree = asiaBayesTree;

  // create a new factor to be inserted
  //boost::shared_ptr<IndexFactor> newFactor(new IndexFactor(_S_,_B_));

  // Remove the contaminated part of the Bayes tree
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removeTop(list_of(_B_)(_S_), &bn, &orphans);

  // Check expected outcome
  SymbolicBayesNet expected;
  expected += SymbolicConditional::FromKeys(list_of(_E_)(_L_)(_B_), 3);
  expected += SymbolicConditional::FromKeys(list_of(_S_)(_B_)(_L_), 1);
  CHECK(assert_equal(expected, bn));

  SymbolicBayesTree::Cliques expectedOrphans;
  expectedOrphans += bayesTree[_T_], bayesTree[_X_];
  CHECK(assert_container_equal(expectedOrphans|indirected, orphans|indirected));

  // Try removeTop again with a factor that should not change a thing
  //boost::shared_ptr<IndexFactor> newFactor2(new IndexFactor(_B_));
  SymbolicBayesNet bn2;
  SymbolicBayesTree::Cliques orphans2;
  bayesTree.removeTop(list_of(_B_), &bn2, &orphans2);
  SymbolicFactorGraph factors2(bn2);
  SymbolicFactorGraph expected2;
  CHECK(assert_equal(expected2, factors2));
  SymbolicBayesTree::Cliques expectedOrphans2;
  CHECK(assert_container_equal(expectedOrphans2|indirected, orphans2|indirected));
}

/* ************************************************************************* */
TEST( BayesTree, removeTop2 )
{
  SymbolicBayesTree bayesTree = asiaBayesTree;

  // create two factors to be inserted
  //SymbolicFactorGraph newFactors;
  //newFactors.push_factor(_B_);
  //newFactors.push_factor(_S_);

  // Remove the contaminated part of the Bayes tree
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removeTop(list_of(_T_), &bn, &orphans);

  // Check expected outcome
  SymbolicBayesNet expected = list_of
    (SymbolicConditional::FromKeys(list_of(_E_)(_L_)(_B_), 3))
    (SymbolicConditional::FromKeys(list_of(_T_)(_E_)(_L_), 1));
  CHECK(assert_equal(expected, bn));

  SymbolicBayesTree::Cliques expectedOrphans;
  expectedOrphans += bayesTree[_S_], bayesTree[_X_];
  CHECK(assert_container_equal(expectedOrphans|indirected, orphans|indirected));
}

/* ************************************************************************* */
TEST( BayesTree, removeTop3 )
{
  SymbolicFactorGraph graph = list_of
    (SymbolicFactor(L(5)))
    (SymbolicFactor(X(4), L(5)))
    (SymbolicFactor(X(2), X(4)))
    (SymbolicFactor(X(3), X(2)));
  Ordering ordering(list_of (X(3)) (X(2)) (X(4)) (L(5)) );
  SymbolicBayesTree bayesTree = *graph.eliminateMultifrontal(ordering);

  // remove all
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removeTop(list_of(L(5))(X(4))(X(2))(X(3)), &bn, &orphans);

  SymbolicBayesNet expectedBn = list_of
    (SymbolicConditional::FromKeys(list_of(X(4))(L(5)), 2))
    (SymbolicConditional(X(2), X(4)))
    (SymbolicConditional(X(3), X(2)));
  EXPECT(assert_equal(expectedBn, bn));
  EXPECT(orphans.empty());
}

/* ************************************************************************* */
TEST( BayesTree, removeTop4 )
{
  SymbolicFactorGraph graph = list_of
    (SymbolicFactor(L(5)))
    (SymbolicFactor(X(4), L(5)))
    (SymbolicFactor(X(2), X(4)))
    (SymbolicFactor(X(3), X(2)));
  Ordering ordering(list_of (X(3)) (X(2)) (X(4)) (L(5)) );
  SymbolicBayesTree bayesTree = *graph.eliminateMultifrontal(ordering);

  // remove all
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removeTop(list_of(X(2))(L(5))(X(4))(X(3)), &bn, &orphans);

  SymbolicBayesNet expectedBn = list_of
    (SymbolicConditional::FromKeys(list_of(X(4))(L(5)), 2))
    (SymbolicConditional(X(2), X(4)))
    (SymbolicConditional(X(3), X(2)));
  EXPECT(assert_equal(expectedBn, bn));
  EXPECT(orphans.empty());
}

/* ************************************************************************* */
TEST( BayesTree, removeTop5 )
{
  // Remove top called with variables that are not in the Bayes tree
  SymbolicFactorGraph graph = list_of
    (SymbolicFactor(L(5)))
    (SymbolicFactor(X(4), L(5)))
    (SymbolicFactor(X(2), X(4)))
    (SymbolicFactor(X(3), X(2)));
  Ordering ordering(list_of (X(3)) (X(2)) (X(4)) (L(5)) );
  SymbolicBayesTree bayesTree = *graph.eliminateMultifrontal(ordering);

  // Remove nonexistant
  SymbolicBayesNet bn;
  SymbolicBayesTree::Cliques orphans;
  bayesTree.removeTop(list_of(X(10)), &bn, &orphans);

  SymbolicBayesNet expectedBn;
  EXPECT(assert_equal(expectedBn, bn));
  EXPECT(orphans.empty());
}

/* ************************************************************************* */
TEST( SymbolicBayesTree, thinTree ) {

  // create a thin-tree Bayesnet, a la Jean-Guillaume
  SymbolicBayesNet bayesNet;
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(14));

  bayesNet.push_back(boost::make_shared<SymbolicConditional>(13, 14));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(12, 14));

  bayesNet.push_back(boost::make_shared<SymbolicConditional>(11, 13, 14));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(10, 13, 14));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(9, 12, 14));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(8, 12, 14));

  bayesNet.push_back(boost::make_shared<SymbolicConditional>(7, 11, 13));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(6, 11, 13));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(5, 10, 13));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(4, 10, 13));

  bayesNet.push_back(boost::make_shared<SymbolicConditional>(3, 9, 12));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(2, 9, 12));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(1, 8, 12));
  bayesNet.push_back(boost::make_shared<SymbolicConditional>(0, 8, 12));

  if (debug) {
    GTSAM_PRINT(bayesNet);
    bayesNet.saveGraph("/tmp/symbolicBayesNet.dot");
  }

  // create a BayesTree out of a Bayes net
  SymbolicBayesTree bayesTree = *SymbolicFactorGraph(bayesNet).eliminateMultifrontal();
  if (debug) {
    GTSAM_PRINT(bayesTree);
    bayesTree.saveGraph("/tmp/SymbolicBayesTree.dot");
  }

  SymbolicBayesTree::Clique::shared_ptr R = bayesTree.roots().front();

  {
    // check shortcut P(S9||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[9];
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(14, 11, 13));
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S8||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[8];
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of
      (SymbolicConditional(12, 14))
      (SymbolicConditional(14, 11, 13));
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S4||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[4];
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(10, 11, 13));
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S2||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[2];
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(9, 11, 12, 13))
                                       (SymbolicConditional(12, 11, 13));
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S0||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[0];
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(8, 11, 12, 13))
                                       (SymbolicConditional(12, 11, 13));
    EXPECT(assert_equal(expected, shortcut));
  }

  SymbolicBayesNet::shared_ptr actualJoint;

  // Check joint P(8,2)
  if (false) { // TODO, not disjoint
    actualJoint = bayesTree.jointBayesNet(8, 2);
    SymbolicBayesNet expected;
    expected.push_back(boost::make_shared<SymbolicConditional>(8));
    expected.push_back(boost::make_shared<SymbolicConditional>(2, 8));
    EXPECT(assert_equal(expected, *actualJoint));
  }

  // Check joint P(1,2)
  if (false) { // TODO, not disjoint
    actualJoint = bayesTree.jointBayesNet(1, 2);
    SymbolicBayesNet expected;
    expected.push_back(boost::make_shared<SymbolicConditional>(2));
    expected.push_back(boost::make_shared<SymbolicConditional>(1, 2));
    EXPECT(assert_equal(expected, *actualJoint));
  }

  // Check joint P(2,6)
  if (true) {
    actualJoint = bayesTree.jointBayesNet(2, 6);
    SymbolicBayesNet expected;
    expected.push_back(boost::make_shared<SymbolicConditional>(2, 6));
    expected.push_back(boost::make_shared<SymbolicConditional>(6));
    EXPECT(assert_equal(expected, *actualJoint));
  }

  // Check joint P(4,6)
  if (false) { // TODO, not disjoint
    actualJoint = bayesTree.jointBayesNet(4, 6);
    SymbolicBayesNet expected;
    expected.push_back(boost::make_shared<SymbolicConditional>(6));
    expected.push_back(boost::make_shared<SymbolicConditional>(4, 6));
    EXPECT(assert_equal(expected, *actualJoint));
  }
}

/* ************************************************************************* */
TEST(SymbolicBayesTree, forest_joint)
{
  // Create forest
  SymbolicBayesTreeClique::shared_ptr root1 = MakeClique(list_of(1), 1);
  SymbolicBayesTreeClique::shared_ptr root2 = MakeClique(list_of(2), 1);
  SymbolicBayesTree bayesTree;
  bayesTree.insertRoot(root1);
  bayesTree.insertRoot(root2);

  // Check joint
  SymbolicBayesNet expected = list_of
    (SymbolicConditional(1))
    (SymbolicConditional(2));
  SymbolicBayesNet actual = *bayesTree.jointBayesNet(1, 2);

  EXPECT(assert_equal(expected, actual));
}

/* ************************************************************************* *
 Bayes tree for smoother with "natural" ordering:
 C1 5 6
 C2   4 : 5
 C3     3 : 4
 C4       2 : 3
 C5         1 : 2
 C6           0 : 1
 **************************************************************************** */

TEST( SymbolicBayesTree, linear_smoother_shortcuts ) {
  // Create smoother with 7 nodes
  SymbolicFactorGraph smoother;
  smoother.push_factor(0);
  smoother.push_factor(0, 1);
  smoother.push_factor(1, 2);
  smoother.push_factor(2, 3);
  smoother.push_factor(3, 4);
  smoother.push_factor(4, 5);
  smoother.push_factor(5, 6);

  // Eliminate in numerical order 0..6
  Ordering ordering(smoother.keys());
  SymbolicBayesNet bayesNet = *smoother.eliminateSequential(ordering);

  if (debug) {
    GTSAM_PRINT(bayesNet);
    bayesNet.saveGraph("/tmp/symbolicBayesNet.dot");
  }

  // create a BayesTree
  SymbolicBayesTree bayesTree = *smoother.eliminateMultifrontal(ordering);
  if (debug) {
    GTSAM_PRINT(bayesTree);
    bayesTree.saveGraph("/tmp/SymbolicBayesTree.dot");
  }

  SymbolicBayesTree::Clique::shared_ptr R = bayesTree.roots().front();

  {
    // check shortcut P(S2||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[4]; // 4 is frontal in C2
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected;
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S3||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[3]; // 3 is frontal in C3
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(4, 5));
    EXPECT(assert_equal(expected, shortcut));
  }

  {
    // check shortcut P(S4||R) to root
    SymbolicBayesTree::Clique::shared_ptr c = bayesTree[2]; // 2 is frontal in C4
    SymbolicBayesNet shortcut = c->shortcut(R);
    SymbolicBayesNet expected = list_of(SymbolicConditional(3, 5));
    EXPECT(assert_equal(expected, shortcut));
  }
}

/* ************************************************************************* */
// from testSymbolicJunctionTree, which failed at one point
TEST(SymbolicBayesTree, complicatedMarginal)
{
  // Create the conditionals to go in the BayesTree
  SymbolicBayesTreeClique::shared_ptr cur;
  SymbolicBayesTreeClique::shared_ptr root = MakeClique(list_of(11)(12), 2);
  cur = root;

  root->children += MakeClique(list_of(9)(10)(11)(12), 2);
  root->children.back()->parent_ = root;

  root->children += MakeClique(list_of(7)(8)(11), 2);
  root->children.back()->parent_ = root;
  cur = root->children.back();

  cur->children += MakeClique(list_of(5)(6)(7)(8), 2);
  cur->children.back()->parent_ = cur;
  cur = cur->children.back();

  cur->children += MakeClique(list_of(3)(4)(6), 2);
  cur->children.back()->parent_ = cur;

  cur->children += MakeClique(list_of(1)(2)(5), 2);
  cur->children.back()->parent_ = cur;

  // Create Bayes Tree
  SymbolicBayesTree bt;
  bt.insertRoot(root);
  if (debug) {
    GTSAM_PRINT(bt);
    bt.saveGraph("/tmp/SymbolicBayesTree.dot");
  }

  // Shortcut on 9
  {
    SymbolicBayesTree::Clique::shared_ptr c = bt[9];
    SymbolicBayesNet shortcut = c->shortcut(root);
    EXPECT(assert_equal(SymbolicBayesNet(), shortcut));
  }

  // Shortcut on 7
  {
    SymbolicBayesTree::Clique::shared_ptr c = bt[7];
    SymbolicBayesNet shortcut = c->shortcut(root);
    EXPECT(assert_equal(SymbolicBayesNet(), shortcut));
  }

  // Shortcut on 5
  {
    SymbolicBayesTree::Clique::shared_ptr c = bt[5];
    SymbolicBayesNet shortcut = c->shortcut(root);
    SymbolicBayesNet expected = list_of
      (SymbolicConditional(7, 8, 11))
      (SymbolicConditional(8, 11));
    EXPECT(assert_equal(expected, shortcut));
  }

  // Shortcut on 3
  {
    SymbolicBayesTree::Clique::shared_ptr c = bt[3];
    SymbolicBayesNet shortcut = c->shortcut(root);
    SymbolicBayesNet expected = list_of(SymbolicConditional(6, 11));
    EXPECT(assert_equal(expected, shortcut));
  }

  // Shortcut on 1
  {
    SymbolicBayesTree::Clique::shared_ptr c = bt[1];
    SymbolicBayesNet shortcut = c->shortcut(root);
    SymbolicBayesNet expected = list_of(SymbolicConditional(5, 11));
    EXPECT(assert_equal(expected, shortcut));
  }

  // Marginal on 5
  {
    SymbolicFactor::shared_ptr actual = bt.marginalFactor(5);
    EXPECT(assert_equal(SymbolicFactor(5), *actual, 1e-1));
  }

  // Shortcut on 6
  {
    SymbolicFactor::shared_ptr actual = bt.marginalFactor(6);
    EXPECT(assert_equal(SymbolicFactor(6), *actual, 1e-1));
  }

}

/* ************************************************************************* */
TEST(SymbolicBayesTree, COLAMDvsMETIS) {

  // create circular graph
  SymbolicFactorGraph sfg;
  sfg.push_factor(0, 1);
  sfg.push_factor(1, 2);
  sfg.push_factor(2, 3);
  sfg.push_factor(3, 4);
  sfg.push_factor(4, 5);
  sfg.push_factor(0, 5);

  // COLAMD
  {
    Ordering ordering = Ordering::Create(Ordering::COLAMD, sfg);
    EXPECT(assert_equal(Ordering(list_of(0)(5)(1)(4)(2)(3)), ordering));

    //  - P( 4 2 3)
    //  | - P( 1 | 2 4)
    //  | | - P( 5 | 1 4)
    //  | | | - P( 0 | 1 5)
    SymbolicBayesTree expected;
    expected.insertRoot(
        MakeClique(list_of(4)(2)(3), 3,
            list_of(
                MakeClique(list_of(1)(2)(4), 1,
                    list_of(
                        MakeClique(list_of(5)(1)(4), 1,
                            list_of(MakeClique(list_of(0)(1)(5), 1))))))));

    SymbolicBayesTree actual = *sfg.eliminateMultifrontal(ordering);
    EXPECT(assert_equal(expected, actual));
  }

#ifdef GTSAM_SUPPORT_NESTED_DISSECTION
  // METIS
  {
    Ordering ordering = Ordering::Create(Ordering::METIS, sfg);
// Linux and Mac split differently when using mettis
#if !defined(__APPLE__)
    EXPECT(assert_equal(Ordering(list_of(3)(2)(5)(0)(4)(1)), ordering));
#else
    EXPECT(assert_equal(Ordering(list_of(5)(4)(2)(1)(0)(3)), ordering));
#endif

    //  - P( 1 0 3)
    //  | - P( 4 | 0 3)
    //  | | - P( 5 | 0 4)
    //  | - P( 2 | 1 3)
    SymbolicBayesTree expected;
#if !defined(__APPLE__)
    expected.insertRoot(
        MakeClique(list_of(2)(4)(1), 3,
            list_of(
                MakeClique(list_of(0)(1)(4), 1,
                    list_of(MakeClique(list_of(5)(0)(4), 1))))(
                MakeClique(list_of(3)(2)(4), 1))));
#else
    expected.insertRoot(
            MakeClique(list_of(1)(0)(3), 3,
                list_of(
                    MakeClique(list_of(4)(0)(3), 1,
                        list_of(MakeClique(list_of(5)(0)(4), 1))))(
                    MakeClique(list_of(2)(1)(3), 1))));
#endif
    SymbolicBayesTree actual = *sfg.eliminateMultifrontal(ordering);
    EXPECT(assert_equal(expected, actual));
  }
#endif
}

/* ************************************************************************* */
int main() {
  TestResult tr;
  return TestRegistry::runAllTests(tr);
}
/* ************************************************************************* */


/*
@copyright Louis Dionne 2014
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana/detail/assert.hpp>
#include <boost/hana/detail/constexpr.hpp>
#include <boost/hana/integral.hpp>
#include <boost/hana/list/instance.hpp>
using namespace boost::hana;

int main() {
    //! [main]
    BOOST_HANA_CONSTEXPR_LAMBDA auto xs = list(0, '1', 2.0);
    BOOST_HANA_CONSTEXPR_ASSERT(drop(int_<0>, xs) == xs);
    BOOST_HANA_CONSTEXPR_ASSERT(drop(int_<1>, xs) == list('1', 2.0));
    BOOST_HANA_CONSTEXPR_ASSERT(drop(int_<2>, xs) == list(2.0));
    BOOST_HANA_CONSTANT_ASSERT(drop(int_<3>, xs) == list());
    BOOST_HANA_CONSTANT_ASSERT(drop(int_<4>, xs) == list());
    //! [main]
}


//==================================================================================================
/*!
  @file

  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#include <boost/simd/pack.hpp>
#include <boost/simd/function/dec.hpp>
#include <boost/simd/function/bits.hpp>
#include <boost/simd/meta/cardinal_of.hpp>
#include <simd_test.hpp>
#include <boost/simd/detail/dispatch/meta/as_integer.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/minf.hpp>
#include <boost/simd/constant/mone.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/constant/valmin.hpp>
#include <boost/simd/constant/two.hpp>
#include <boost/simd/constant/mtwo.hpp>

template <typename T, std::size_t N, typename Env>
void test(Env& runtime)
{
  namespace bs = boost::simd;
  using p_t = bs::pack<T, N>;

  T a1[N], b[N];
  for(std::size_t i = 0; i < N; ++i)
  {
    a1[i] = (i%2) ? T(i) : T(-i);
    b[i] = bs::saturated_(bs::dec)(a1[i]) ;
  }
  p_t aa1(&a1[0], &a1[0]+N);
  p_t bb (&b[0], &b[0]+N);
  STF_IEEE_EQUAL(bs::saturated_(bs::dec)(aa1), bb);
}

STF_CASE_TPL("Check dec_s on pack" , STF_NUMERIC_TYPES)
{
  namespace bs = boost::simd;
  using p_t = bs::pack<T>;
  static const std::size_t N = bs::cardinal_of<p_t>::value;
  test<T, N>(runtime);
  test<T, N/2>(runtime);
  test<T, N*2>(runtime);
}

STF_CASE_TPL (" bs::saturated_(dec) signed_int",  STF_SIGNED_INTEGRAL_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::dec;
  using p_t = bs::pack<T>;

  using r_t = decltype(bs::saturated_(dec)(p_t()));

  // return type conformity test
  STF_TYPE_IS(r_t, p_t);

  // specific values tests
  STF_EQUAL(bs::saturated_(dec)(bs::Mone<p_t>()), bs::Mtwo<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::One<p_t>()), bs::Zero<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Valmin<p_t>()), bs::Valmin<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Zero<p_t>()), bs::Mone<p_t>());
}

STF_CASE_TPL (" bs::saturated_(dec) unsigned_uint",  STF_UNSIGNED_INTEGRAL_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::dec;
  using p_t = bs::pack<T>;

  using r_t = decltype(bs::saturated_(dec)(p_t()));

  // return type conformity test
  STF_TYPE_IS(r_t, p_t);

  // specific values tests
  STF_EQUAL(bs::saturated_(dec)(bs::One<p_t>()), bs::Zero<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Two<p_t>()), bs::One<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Zero<p_t>()), bs::Zero<p_t>());
}

STF_CASE_TPL(" bs::saturated_(dec) floating", STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::dec;
  using p_t = bs::pack<T>;

  using r_t = decltype(bs::saturated_(dec)(p_t()));

  STF_TYPE_IS(r_t, p_t);

  // specific values tests
#ifndef BOOST_SIMD_NO_INVALIDS
  STF_EQUAL(bs::saturated_(dec)(bs::Inf<p_t>()), bs::Inf<p_t>());
  STF_IEEE_EQUAL(bs::saturated_(dec)(bs::Nan<p_t>()), bs::Nan<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Minf<p_t>()), bs::Minf<p_t>());
#endif
  STF_EQUAL(bs::saturated_(dec)(bs::One<p_t>()), bs::Zero<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Two<p_t>()), bs::One<p_t>());
  STF_EQUAL(bs::saturated_(dec)(bs::Zero<p_t>()), bs::Mone<p_t>());
}


//==================================================================================================
/*!

  Copyright 2015 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#include <boost/simd/function/scalar/fmod.hpp>
#include <boost/simd/function/std.hpp>
#include <boost/simd/function/fast.hpp>
#include <simd_test.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/minf.hpp>
#include <boost/simd/constant/mone.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/function/is_positive.hpp>
#include <boost/simd/function/is_negative.hpp>
#include <cmath>

STF_CASE_TPL (" fmod real",  STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
#ifndef STF_NO_INVALIDS
  STF_IEEE_EQUAL(fmod(bs::Inf<T>(), bs::Inf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(fmod(bs::Minf<T>(), bs::Minf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(fmod(bs::Nan<T>(), bs::Nan<T>()), bs::Nan<T>());
#endif
  STF_EQUAL(fmod(bs::Mone<T>(), bs::Mone<T>()), bs::Zero<T>());
  STF_EQUAL(fmod(bs::One<T>(), bs::One<T>()), bs::Zero<T>());
  STF_IEEE_EQUAL(fmod(bs::Zero<T>(), bs::Zero<T>()), bs::Nan<T>());
} // end of test for floating_

STF_CASE_TPL (" fmod real fast",  STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
#ifndef STF_NO_INVALIDS
  STF_IEEE_EQUAL(bs::fast_(fmod)(bs::Inf<T>(), bs::Inf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(bs::fast_(fmod)(bs::Minf<T>(), bs::Minf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(bs::fast_(fmod)(bs::Nan<T>(), bs::Nan<T>()), bs::Nan<T>());
#endif
  STF_EQUAL(bs::fast_(fmod)(bs::Mone<T>(), bs::Mone<T>()), bs::Zero<T>());
  STF_EQUAL(bs::fast_(fmod)(bs::One<T>(), bs::One<T>()), bs::Zero<T>());
} // end of test for floating_

STF_CASE_TPL (" fmod unsigned_int",  STF_UNSIGNED_INTEGRAL_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
  STF_EQUAL(fmod(bs::One<T>(), bs::One<T>()), bs::Zero<T>());
  STF_EQUAL(fmod(bs::Zero<T>(), bs::Zero<T>()), bs::Zero<T>());
} // end of test for unsigned_int_

STF_CASE_TPL (" fmod signed_int",  STF_SIGNED_INTEGRAL_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
  STF_EQUAL(fmod(bs::Mone<T>(), bs::Mone<T>()), bs::Zero<T>());
  STF_EQUAL(fmod(bs::One<T>(), bs::One<T>()), bs::Zero<T>());
  STF_EQUAL(fmod(bs::Zero<T>(), bs::Zero<T>()), bs::Zero<T>());
} // end of test for signed_int_

STF_CASE_TPL (" fmod limits",  STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
  STF_EQUAL(fmod(T(0), T(1)),  T(0));
  STF_EQUAL(fmod(-T(0), T(1)),  T(-0));
  STF_EXPECT(bs::is_negative(fmod(-T(0), T(1))));
  STF_EXPECT(bs::is_positive(fmod(T(0), T(1))));
  STF_IEEE_EQUAL(fmod(bs::Inf<T>(), T(1)),  bs::Nan<T>());
  STF_IEEE_EQUAL(fmod(bs::Minf<T>(), T(1)), bs::Nan<T>());
  STF_IEEE_EQUAL(fmod(T(1), T(0)),  bs::Nan<T>());
  STF_IEEE_EQUAL(fmod(T(1), T(-0)), bs::Nan<T>());
  STF_EQUAL(fmod(T(2), bs::Inf<T>()), T(2));
  STF_EQUAL(fmod(T(2), bs::Minf<T>()), T(2));

}

STF_CASE_TPL (" fmod std",  STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::fmod;
  using r_t = decltype(fmod(T(), T()));

  // return type conformity test
  STF_TYPE_IS(r_t, T);

  // specific values tests
#ifndef STF_NO_INVALIDS
  STF_IEEE_EQUAL(bs::std_(fmod)(bs::Inf<T>(), bs::Inf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(bs::std_(fmod)(bs::Minf<T>(), bs::Minf<T>()), bs::Nan<T>());
  STF_IEEE_EQUAL(bs::std_(fmod)(bs::Nan<T>(), bs::Nan<T>()), bs::Nan<T>());
#endif
  STF_EQUAL(bs::std_(fmod)(bs::Mone<T>(), bs::Mone<T>()), bs::Zero<T>());
  STF_EQUAL(bs::std_(fmod)(bs::One<T>(), bs::One<T>()), bs::Zero<T>());
  STF_IEEE_EQUAL(bs::std_(fmod)(bs::Zero<T>(), bs::Zero<T>()), bs::Nan<T>());
} // end of


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_DETAIL_CONSTANT_MAXFLINT_HPP_INCLUDED
#define BOOST_SIMD_DETAIL_CONSTANT_MAXFLINT_HPP_INCLUDED

#include <boost/simd/config.hpp>
#include <boost/simd/detail/brigand.hpp>
#include <boost/simd/detail/dispatch.hpp>
#include <boost/simd/detail/constant_traits.hpp>
#include <boost/simd/detail/dispatch/function/make_callable.hpp>
#include <boost/simd/detail/dispatch/hierarchy/functions.hpp>
#include <boost/simd/detail/dispatch/as.hpp>

/*

    @ingroup group-constant

    Generate the least integer value which is exactly
    representable in floating point numbers and
    equal to its integral successor.

    All floating numbers greater than Maxflint are integral.

    @par Semantic:

    @code
    T r = Maxflint<T>();
    @endcode

    is similar to:

    @code
    if T is double
      r = 9007199254740992.0
    else if T is float
      r = 16777216.0f
    @endcode

    @return The Maxflint constant for the proper type
  */

namespace boost { namespace simd
{
  namespace tag
  {
    struct maxflint_ : boost::dispatch::constant_value_<maxflint_>
    {
      BOOST_DISPATCH_MAKE_CALLABLE(ext,maxflint_,boost::dispatch::constant_value_<maxflint_>);
      BOOST_SIMD_REGISTER_CONSTANT(1, 0x4b800000, 0x4340000000000000ll);
    };
  }

  namespace ext
  {
    BOOST_DISPATCH_FUNCTION_DECLARATION(tag, maxflint_)
  }

  namespace detail
  {
    BOOST_DISPATCH_CALLABLE_DEFINITION(tag::maxflint_,maxflint);
  }

  template<typename T> BOOST_FORCEINLINE auto Maxflint()
  BOOST_NOEXCEPT_DECLTYPE(detail::maxflint( boost::dispatch::as_<T>{}))
  {
    return detail::maxflint( boost::dispatch::as_<T>{} );
  }
} }

#include <boost/simd/arch/common/scalar/constant/constant_value.hpp>
#include <boost/simd/arch/common/simd/constant/constant_value.hpp>

#endif


/*
 * This benchmark causes some versions of gcc to crash.
 */

#include <blitz/vector2.h>
#include <blitz/array.h>
#include <random/uniform.h>
#include <blitz/benchext.h>

BZ_NAMESPACE(blitz)
extern void sink();
BZ_NAMESPACE_END

BZ_USING_NAMESPACE(blitz)

void VectorVersion(BenchmarkExt<int>& bench);
void ArrayVersion(BenchmarkExt<int>& bench);

int dontActuallyRunBenchmark()
{
    return 1;
}

int main()
{
    if (dontActuallyRunBenchmark())
        return 0;

    int numBenchmarks = 2;

    BenchmarkExt<int> bench("loop1: $x=sqrt($y)", numBenchmarks);

    const int numSizes = 26;
    bench.setNumParameters(numSizes);
    //bench.setRateDescription("Mflops/s"); removed

    Vector<int> parameters(numSizes);
    Vector<long> iters(numSizes);
    Vector<double> flops(numSizes);

    for (int i=0; i < numSizes; ++i)
    {
      parameters(i) = (int)pow(10.0, (i+1)/4.0);
      iters[i] = 50000000L / parameters(i);
        if (iters(i) < 2)
	  iters(i) = 2;
        flops(i) = 1 * parameters(i);
    }

    bench.setParameterVector(parameters);
    bench.setIterations(iters);
    bench.setOpsPerIteration(flops);

    bench.beginBenchmarking();

    VectorVersion(bench);
    ArrayVersion(bench);

    bench.endBenchmarking();

    bench.saveMatlabGraph("loop1.m");

    return 0;
}

template<typename T>
void initializeRandomDouble(T* data, int numElements, int stride = 1)
{
  ranlib::Uniform<T> rnd;

    for (int i=0; i < numElements; ++i)
        data[size_t(i*stride)] = rnd.random();
}

void VectorVersion(BenchmarkExt<int>& bench)
{
    bench.beginImplementation("Vector<T>");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();
        cout << "Vector<T>: N = " << N << endl;
        cout.flush();

        long iters = bench.getIterations();

        Vector<double> x(N);
        initializeRandomDouble(x.data(), N);
        Vector<double> y(N);
        initializeRandomDouble(y.data(), N);

        bench.start();
        long i;
        for (i=0; i < iters; ++i)
        {
            x=sqrt(y);
            sink();
        }
        bench.stop();

        bench.startOverhead();
        for (i=0; i < iters; ++i)
            sink();
        bench.stopOverhead();
    }

    bench.endImplementation();
}

void ArrayVersion(BenchmarkExt<int>& bench)
{
    bench.beginImplementation("Array<T,1>");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();
        cout << "Array<T,1>: N = " << N << endl;
        cout.flush();

        long iters = bench.getIterations();

        Array<double, 1> x(N);
        initializeRandomDouble(x.dataFirst(), N);
        Array<double, 1> y(N);
        initializeRandomDouble(y.dataFirst(), N);

        bench.start();
        long i;
        for (i=0; i < iters; ++i)
        {
            x=sqrt(y);
            sink();
        }
        bench.stop();

        bench.startOverhead();
        for (i=0; i < iters; ++i)
            sink();
        bench.stopOverhead();
    }

    bench.endImplementation();
}


/**
 * GeoDa TM, Copyright (C) 2011-2015 by Luc Anselin - all rights reserved
 *
 * This file is part of GeoDa.
 *
 * GeoDa is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * GeoDa is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <cmath>
#include <limits>
#include <math.h>
#include <boost/math/distributions/students_t.hpp>
#include <boost/math/distributions/fisher_f.hpp>
#include "../logger.h"
#include "LineChartStats.h"

LineChartStats::LineChartStats(const vec_vec_dbl_type& Y_,
                               const std::vector<bool>& Y_undef_,
                               const wxString& Yname_,
                               const std::vector<bool>& tms_subset0_,
                               const std::vector<bool>& tms_subset1_,
                               const bool& compare_regimes_,
                               const bool& compare_time_periods_,
                               const bool& compare_r_and_t_)
: Y(Y_),
Y_undef(Y_undef_),
Yname(Yname_),
Y_avg_min(0),
Y_avg_max(0),
Y_sel_avg_min(0),
Y_excl_avg_min(0),
Y_sel_avg_max(0),
Y_excl_avg_max(0),
tms_subset0(tms_subset0_),
tms_subset1(tms_subset1_),
compare_regimes(compare_regimes_),
compare_time_periods(compare_time_periods_),
compare_r_and_t(compare_r_and_t_),
test_stat_valid(false),
time_invariant(tms_subset0_.size() > 1)
{
	test_stat_c.resize(6);
	test_stat_valid_c.resize(6);
	deg_free_c.resize(6);
	p_val_c.resize(6);
	for (size_t i=0; i<6; ++i) {
		test_stat_c[i] = 0;
		test_stat_valid_c[i] = false;
		deg_free_c[i] = 0;
		p_val_c[i] = 1;
	}

	ss_ptrs.push_back(&s0);
	ss_ptrs.push_back(&s1);
	ss_ptrs.push_back(&s2);
	ss_ptrs.push_back(&s3);
}

LineChartStats::SampStats::SampStats()
{
	mean = 0; // sample mean
	mean_v = false; // is sample mean valid
	var = 0; // sample variance
	sd = 0; // sample standard deviation
	sem = 0; // standard error of the mean
	var_v = false; // is sample variance valid
	sz_i = 0; // sample size as integer
	sz_d = 0; // sample size as double
}

void LineChartStats::UpdateNonRegimesNonTmsStats()
{
	Y_avg_valid = false;
	size_t tms=Y.size();

	size_t valid_num_obs = 0;

    for (size_t i=0; i<Y_undef.size(); i++) {
        if ( !Y_undef[i]) {
            valid_num_obs += 1;
        }
    }

	double num_obs_d = (double) valid_num_obs;
	obs_sz_i = valid_num_obs;
	obs_sz_d = obs_sz_i;
	Y_avg.resize(tms);
	Y_ss.resize(tms);

	for (size_t t=0; t<tms; ++t) {
		Y_avg[t] = 0;
		Y_ss[t] = 0;

        if (Y[t].size() > 0)  {
    		for (size_t i=0; i<Y_undef.size(); ++i) {
                if (Y_undef[i])
                    continue;
    			Y_avg[t] += Y[t][i];
    			Y_ss[t] += Y[t][i] * Y[t][i];
    		}
    		if (valid_num_obs > 0) {
    			Y_avg[t] /= num_obs_d;
    			Y_avg_valid = true;
    			if (t==0) {
    				Y_avg_min = Y_avg[0];
    				Y_avg_max = Y_avg[0];
    			} else {
    				if (Y_avg[t] < Y_avg_min) {
    					Y_avg_min = Y_avg[t];
    				} else if (Y_avg[t] > Y_avg_max) {
    					Y_avg_max = Y_avg[t];
    				}
    			}
    		}
        }
	}
}

void LineChartStats::UpdateRegimesStats(const std::vector<bool>& hs,
                                        int default_Y_sel_avg_valid,
                                        int default_Y_excl_avg_valid)
{
    Y_sel_avg_min = 0;
    Y_sel_avg_max = 0;
    Y_excl_avg_min = 0;
    Y_excl_avg_max = 0;

	Y_sel_avg_valid = false;
	Y_excl_avg_valid = false;
	size_t tms=Y.size();
    size_t valid_num_obs = 0;

    for (size_t i=0; i<Y_undef.size(); i++) {
        if ( !Y_undef[i]) {
            valid_num_obs += 1;
        }
    }

	double num_obs_d = (double) valid_num_obs;
	Y_sel_avg.resize(tms);
	Y_sel_ss.resize(tms);
	Y_excl_avg.resize(tms);
	Y_excl_ss.resize(tms);

	double num_sel = 0;
    double num_excl = 0;
    bool num_sel_valid = false;
    bool num_unsel_valid = false;

    for (size_t i=0; i<hs.size(); ++i) {
        if ( hs[i] ) {
            if ( Y_undef[i] == false ) {
                num_sel_valid = true;
			    num_sel += 1.0;
            }
        } else {
            if ( Y_undef[i] == false ) {
                num_unsel_valid = true;
                num_excl += 1.0;
            }
        }
    }

	sel_sz_i = (int) num_sel;
	sel_sz_d = num_sel;
	excl_sz_i = (int) num_excl;
	excl_sz_d = num_excl;

	for (size_t t=0; t<tms; ++t) {
		Y_sel_avg[t] = 0;
		Y_sel_ss[t] = 0;
		Y_excl_avg[t] = 0;
		Y_excl_ss[t] = 0;
        if (Y[t].size() == 0)
            continue;

		for (size_t i=0; i<Y_undef.size(); ++i) {
            if (Y_undef[i])
                continue;
			if (hs[i]) {
                if (num_sel_valid) {
                    Y_sel_avg[t] += Y[t][i];
                    Y_sel_ss[t] += Y[t][i] * Y[t][i];
                }
			} else {
                if (num_unsel_valid) {
                    Y_excl_avg[t] += Y[t][i];
                    Y_excl_ss[t] += Y[t][i] * Y[t][i];
                }
			}
		}
		if (num_sel > 0 && num_sel_valid) {
            if (num_sel > 0)
                Y_sel_avg[t] /= num_sel;
            else
                Y_sel_avg[t] = 0.0;

			Y_sel_avg_valid = true;
			if (t==0) {
				Y_sel_avg_min = Y_sel_avg[0];
				Y_sel_avg_max = Y_sel_avg[0];
			}
			if (Y_sel_avg[t] < Y_sel_avg_min) {
				Y_sel_avg_min = Y_sel_avg[t];
			} else if (Y_sel_avg[t] > Y_sel_avg_max) {
				Y_sel_avg_max = Y_sel_avg[t];
			}
		}
		if (num_excl > 0 && num_unsel_valid) {
            if (num_excl > 0)
                Y_excl_avg[t] /= num_excl;
            else
                Y_excl_avg[t] = 0.0;

			Y_excl_avg_valid = true;
			if (t==0) {
				Y_excl_avg_min = Y_excl_avg[0];
				Y_excl_avg_max = Y_excl_avg[0];
			}
			if (Y_excl_avg[t] < Y_excl_avg_min) {
				Y_excl_avg_min = Y_excl_avg[t];
			} else if (Y_excl_avg[t] > Y_excl_avg_max) {
				Y_excl_avg_max = Y_excl_avg[t];
			}
		}

	}

    if (!num_sel_valid)
        Y_sel_avg_valid = num_sel_valid;
    if (!num_unsel_valid)
        Y_excl_avg_valid = num_unsel_valid;

    // override Y_sel_avg_valid if user selected in UI
    if (default_Y_excl_avg_valid >=0) {
        Y_excl_avg_valid = default_Y_excl_avg_valid;
    }
    if (default_Y_sel_avg_valid >= 0 ) {
        Y_sel_avg_valid = default_Y_sel_avg_valid;
    }
}

void LineChartStats::UpdateOtherStats()
{
	if (compare_time_periods)
        UpdateCompareTmsStats();

	//if (compare_regimes) UpdateCompareRegimesStats();

	if (compare_regimes || compare_r_and_t)
        UpdateCompareRegAndTmStats();

	UpdateTtest();
}

void LineChartStats::UpdateCompareTmsStats()
{
	// Instead of comparing highlight/non-highlight regimes, time subsets
	// 0 and 1 are used for regimes.

	Y_avg_tm0_valid = false;
	Y_avg_tm1_valid = false;
	Y_avg_tm0 = 0;
	Y_avg_tm1 = 0;

	size_t tms = tms_subset0.size();
	size_t tsub0_sz = 0;
	for (size_t t=0; t<tms; ++t)
		if (tms_subset0[t])
			tsub0_sz++;
	size_t tsub1_sz = 0;
	for (size_t t=0; t<tms; ++t)
		if (tms_subset1[t])
			tsub1_sz++;

	if (tsub0_sz > 0 && Y_avg_valid) {
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				Y_avg_tm0 += Y_avg[t];
		Y_avg_tm0 /= (double) tsub0_sz;
		Y_avg_tm0_valid = true;
	}

	if (tsub1_sz > 0 && Y_avg_valid) {
		for (size_t t=0; t<tms; ++t)
			if (tms_subset1[t])
				Y_avg_tm1 += Y_avg[t];
		Y_avg_tm1 /= (double) tsub1_sz;
		Y_avg_tm1_valid = true;
	}

    size_t valid_num_obs = 0;

    for (size_t i=0; i<Y_undef.size(); i++) {
        if ( !Y_undef[i]) {
            valid_num_obs += 1;
        }
    }

    double obs_d = (double) valid_num_obs;

	SampStats blank;
	s0 = blank;
	s1 = blank;
	s2 = blank;
	s3 = blank;
	s0.sz_i = tsub0_sz * valid_num_obs;
	s1.sz_i = tsub1_sz * valid_num_obs;
	s0.sz_d = (double) s0.sz_i;
	s1.sz_d = (double) s1.sz_i;
	s0.mean = Y_avg_tm0;
	s0.mean_v = Y_avg_tm0_valid;
	s1.mean = Y_avg_tm1;
	s1.mean_v = Y_avg_tm1_valid;
	if (s0.mean_v && s0.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				ss += Y_ss[t];
		s0.var = (ss - s0.sz_d*s0.mean*s0.mean)/(s0.sz_d-1);
		if (s0.var > 16*DBL_MIN) {
			s0.var_v = true;
			s0.sd = sqrt(s0.var);
			s0.sem = s0.sd/sqrt(s0.sz_d);
		}
	}
	if (s1.mean_v && s1.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset1[t])
				ss += Y_ss[t];
		s1.var = (ss - s1.sz_d*s1.mean*s1.mean)/(s1.sz_d-1);
		if (s1.var > 16*DBL_MIN) {
			s1.var_v = true;
			s1.sd = sqrt(s1.var);
			s1.sem = s1.sd/sqrt(s1.sz_d);
		}
	}
}

void LineChartStats::UpdateCompareRegimesStats()
{
	Y_sel_tm0_avg_valid = false;
	Y_excl_tm0_avg_valid = false;
	Y_sel_tm0_avg = 0;
	Y_excl_tm0_avg = 0;

	size_t tms = tms_subset0.size();
	size_t tsub0_sz = 0;
	if (Y.size() > 1) {
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				tsub0_sz++;
	} else {
		tsub0_sz = 1;
	}

	if (tsub0_sz > 0 && Y_sel_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset0[t])
				Y_sel_tm0_avg += Y_sel_avg[t];
		}
		Y_sel_tm0_avg /= (double) tsub0_sz;
		Y_sel_tm0_avg_valid = true;
	}

	if (tsub0_sz > 0 && Y_excl_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset0[t])
				Y_excl_tm0_avg += Y_excl_avg[t];
		}
		Y_excl_tm0_avg /= (double) tsub0_sz;
		Y_excl_tm0_avg_valid = true;
	}

    size_t valid_num_obs = 0;

    for (size_t i=0; i<Y_undef.size(); i++) {
        if ( !Y_undef[i]) {
            valid_num_obs += 1;
        }
    }

	double obs_d = (double) valid_num_obs;
	SampStats blank;
	s0 = blank;
	s1 = blank;
	s2 = blank;
	s3 = blank;
	s0.sz_i = tsub0_sz * sel_sz_i;
	s1.sz_i = tsub0_sz * excl_sz_i;
	s0.sz_d = (double) s0.sz_i;
	s1.sz_d = (double) s1.sz_i;
	s0.mean = Y_sel_tm0_avg;
	s0.mean_v = Y_sel_tm0_avg_valid;
	s1.mean = Y_excl_tm0_avg;
	s1.mean_v = Y_excl_tm0_avg_valid;
	if (s0.mean_v && s0.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				ss += Y_sel_ss[t];
		s0.var = (ss - s0.sz_d*s0.mean*s0.mean)/(s0.sz_d-1);
		if (s0.var > 16*DBL_MIN) {
			s0.var_v = true;
			s0.sd = sqrt(s0.var);
			s0.sem = s0.sd/sqrt(s0.sz_d);
		}
	}
	if (s1.mean_v && s1.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				ss += Y_excl_ss[t];
		s1.var = (ss - s1.sz_d*s1.mean*s1.mean)/(s1.sz_d-1);
		if (s1.var > 16*DBL_MIN) {
			s1.var_v = true;
			s1.sd = sqrt(s1.var);
			s1.sem = s1.sd/sqrt(s1.sz_d);
		}
	}
}

void LineChartStats::UpdateCompareRegAndTmStats()
{
	Y_sel_tm0_avg_valid = false;
	Y_excl_tm0_avg_valid = false;
	Y_sel_tm1_avg_valid = false;
	Y_excl_tm1_avg_valid = false;
	Y_sel_tm0_avg = 0;
	Y_excl_tm0_avg = 0;
	Y_sel_tm1_avg = 0;
	Y_excl_tm1_avg = 0;

	size_t tms = tms_subset0.size();
	size_t tsub0_sz = 0;
	if (Y.size() > 1) {
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				tsub0_sz++;
	} else {
		tsub0_sz = 1;
	}
	size_t tsub1_sz = 0;
	if (Y.size() > 1) {
		for (size_t t=0; t<tms; ++t)
			if (tms_subset1[t])
				tsub1_sz++;
	} else {
		tsub1_sz = 1;
	}

	if (tsub0_sz > 0 && Y_sel_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset0[t])
				Y_sel_tm0_avg += Y_sel_avg[t];
		}
		Y_sel_tm0_avg /= (double) tsub0_sz;
        if (Y_sel_avg_valid == 1) {
            Y_sel_tm0_avg_valid = true;
        }
	}

	if (tsub0_sz > 0 && Y_excl_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset0[t])
				Y_excl_tm0_avg += Y_excl_avg[t];
		}
		Y_excl_tm0_avg /= (double) tsub0_sz;
        if (Y_excl_avg_valid == 1) {
            Y_excl_tm0_avg_valid = true;
        }
	}

	if (tsub1_sz > 0 && Y_sel_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset1[t])
				Y_sel_tm1_avg += Y_sel_avg[t];
		}
		Y_sel_tm1_avg /= (double) tsub1_sz;
        if (Y_sel_avg_valid == 2 ||
            (Y_sel_avg_valid == 1 && Y_excl_avg_valid == 0)) {
            Y_sel_tm1_avg_valid = true;
        }
	}

	if (tsub1_sz > 0 && Y_excl_avg_valid) {
		for (size_t t=0; t<tms; ++t) {
			if (tms_subset1[t])
				Y_excl_tm1_avg += Y_excl_avg[t];
		}
		Y_excl_tm1_avg /= (double) tsub1_sz;
        if (Y_excl_avg_valid == 2 ||
            (Y_excl_avg_valid == 1 && Y_sel_avg_valid == 0)) {
            Y_excl_tm1_avg_valid = true;
        }
	}

    size_t valid_num_obs = 0;

    for (size_t i=0; i<Y_undef.size(); i++) {
        if ( !Y_undef[i]) {
            valid_num_obs += 1;
        }
    }
	double obs_d = (double) valid_num_obs;
	SampStats blank;
	s0 = blank;
	s1 = blank;
	s2 = blank;
	s3 = blank;
	s0.sz_i = tsub0_sz * sel_sz_i;
	s1.sz_i = tsub0_sz * excl_sz_i;
	s0.sz_d = (double) s0.sz_i;
	s1.sz_d = (double) s1.sz_i;
	s0.mean = Y_sel_tm0_avg;
	s0.mean_v = Y_sel_tm0_avg_valid;
	s1.mean = Y_excl_tm0_avg;
	s1.mean_v = Y_excl_tm0_avg_valid;

	s2.sz_i = tsub1_sz * sel_sz_i;
	s3.sz_i = tsub1_sz * excl_sz_i;
	s2.sz_d = (double) s2.sz_i;
	s3.sz_d = (double) s3.sz_i;
	s2.mean = Y_sel_tm1_avg;
	s2.mean_v = Y_sel_tm1_avg_valid;
	s3.mean = Y_excl_tm1_avg;
	s3.mean_v = Y_excl_tm1_avg_valid;

	if (s0.mean_v && s0.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				ss += Y_sel_ss[t];
		s0.var = (ss - s0.sz_d*s0.mean*s0.mean)/(s0.sz_d-1);
		if (s0.var > 16*DBL_MIN) {
			s0.var_v = true;
			s0.sd = sqrt(s0.var);
			s0.sem = s0.sd/sqrt(s0.sz_d);
		}
	}
	if (s1.mean_v && s1.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset0[t])
				ss += Y_excl_ss[t];
		s1.var = (ss - s1.sz_d*s1.mean*s1.mean)/(s1.sz_d-1);
		if (s1.var > 16*DBL_MIN) {
			s1.var_v = true;
			s1.sd = sqrt(s1.var);
			s1.sem = s1.sd/sqrt(s1.sz_d);
		}
	}

	if (s2.mean_v && s2.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset1[t])
				ss += Y_sel_ss[t];
		s2.var = (ss - s2.sz_d*s2.mean*s2.mean)/(s2.sz_d-1);
		if (s2.var > 16*DBL_MIN) {
			s2.var_v = true;
			s2.sd = sqrt(s2.var);
			s2.sem = s2.sd/sqrt(s2.sz_d);
		}
	}
	if (s3.mean_v && s3.sz_i >= 2) {
		double ss = 0;
		for (size_t t=0; t<tms; ++t)
			if (tms_subset1[t])
				ss += Y_excl_ss[t];
		s3.var = (ss - s3.sz_d*s3.mean*s3.mean)/(s3.sz_d-1);
		if (s3.var > 16*DBL_MIN) {
			s3.var_v = true;
			s3.sd = sqrt(s3.var);
			s3.sem = s3.sd/sqrt(s3.sz_d);
		}
	}
}

void LineChartStats::UpdateTtest()
{
	using namespace boost::math;
	test_stat_valid = false;
	if (s0.var_v && s1.var_v) {

        double mu = (s0.mean * s0.sz_d + s1.mean * s1.sz_d ) /
        (s0.sz_d + s1.sz_d);

        double est_eff_i = s0.mean - mu;
        double est_eff_j = s1.mean - mu;

        double ss_treat = est_eff_i * est_eff_i * s0.sz_d +
        est_eff_j * est_eff_j * s1.sz_d;

        double ss_res = s0.var * (s0.sz_d -1) + s1.var * (s1.sz_d -1);

        double ss_tot = ss_treat + ss_res;

        double df_treat = 2 - 1;
        double df_tot = s0.sz_d + s1.sz_d -1;
        double df_res = df_tot - df_treat;

        double MS_treat = ss_treat / df_treat;

        double MS_res = ss_res / df_res;

        double f_val = MS_treat / MS_res;

        fisher_f dist(df_treat, df_res);
        double q = cdf(complement(dist, fabs(f_val)));

        deg_free = obs_sz_i-1;
        test_stat = f_val;
        p_val = q;

        /*
		test_stat = (s0.mean-s1.mean) / sqrt(s0.var/s0.sz_d + s1.var/s1.sz_d);

		// Estimating the combined degrees of freedom we using the
		// Welch-Satterthwaite approximation:
		deg_free = pow(s0.var/s0.sz_d + s1.var/s1.sz_d, 2) /
		(pow(s0.var/s0.sz_d,2)/(s0.sz_d-1) +
		 pow(s1.var/s1.sz_d,2)/(s1.sz_d-1) );

		// T-score to 2-sided p-val
		students_t dist(deg_free);
		double q = cdf(complement(dist, fabs(test_stat)));
		p_val = 2 * q;
		wxString msg;
		msg << "Probability that difference is due to chance" << " =  ";
		msg << p_val;
		LOG_MSG(msg);
         */
		test_stat_valid = true;
	}

	size_t c = 0;
	for (size_t i=0; i<ss_ptrs.size(); ++i) {
		const SampStats& si = *ss_ptrs[i];
		for (size_t j=i+1; j<ss_ptrs.size(); ++j) {
			const SampStats& sj = *ss_ptrs[j];
			test_stat_valid_c[c] = false;
			if (si.var_v && sj.var_v) {

                double mu = (si.mean * si.sz_d + sj.mean * sj.sz_d ) /
                            (si.sz_d + sj.sz_d);

                double est_eff_i = si.mean - mu;
                double est_eff_j = sj.mean - mu;

                double ss_treat = est_eff_i * est_eff_i * si.sz_d +
                                    est_eff_j * est_eff_j * sj.sz_d;

                double ss_res = si.var * (si.sz_d -1) + sj.var * (sj.sz_d -1);

                double ss_tot = ss_treat + ss_res;

                double df_treat = 2 - 1;
                double df_tot = si.sz_d + sj.sz_d -1;
                double df_res = df_tot - df_treat;

                double MS_treat = ss_treat / df_treat;

                double MS_res = ss_res / df_res;

                double f_val = MS_treat / MS_res;

                fisher_f dist(df_treat, df_res);
                double q = cdf(complement(dist, fabs(f_val)));
                if (si.sz_d == sj.sz_d) {
                    if (i == 0 || i == 2)
                        deg_free_c[c] = (sel_sz_i - 1 );
                    else
                        deg_free_c[c] = (excl_sz_i -1);

                } else {
                    deg_free_c[c] = (obs_sz_i - 1);
                }
                //deg_free_c[c] = (obs_sz_i - 1);
                test_stat_c[c] = f_val;
                p_val_c[c] = q;
                /*
                // following is t-test
				test_stat_c[c] = ((si.mean-sj.mean) /
								  sqrt(si.var/si.sz_d + sj.var/sj.sz_d));

				// Estimating the combined degrees of freedom we using the
				// Welch-Satterthwaite approximation:
				deg_free_c[c] = (pow(si.var/si.sz_d + sj.var/sj.sz_d, 2) /
								 (pow(si.var/si.sz_d,2)/(si.sz_d-1) +
								  pow(sj.var/sj.sz_d,2)/(sj.sz_d-1)) );

				// T-score to 2-sided p-val
				students_t dist(deg_free_c[c]);
				double q = cdf(complement(dist, fabs(test_stat_c[c])));
				p_val_c[c] = 2 * q;
				wxString msg;
				msg << "Probability avg s"<<i<<", s"<<j<<" diff ";
				msg << "due to chance:  " << p_val_c[c];
				LOG_MSG(msg);
                 */
				test_stat_valid_c[c] = true;

			}
			++c;
		}
	}
}


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

#include <cassie_common_toolbox/motion_transition.hpp>
#include <Eigen/Dense>
#include <iostream>

using namespace Eigen;

namespace motion_transition {

void startTransition(PhaseVariable &phase, Eigen::VectorXd &ya, Eigen::VectorXd &dya, Eigen::MatrixXd &param) {
    int nParams = param.cols();
    int nOutputs = param.rows();

    VectorXi reserveParams(nParams - 2);
    VectorXi updateParam(2);
    updateParam << 0,1;
    for (int i=0; i<reserveParams.size(); ++i)
        reserveParams(i) = 2 + i;

    double tau_0 = 0.0;
    double dtau_0 = phase.dtau;

    VectorXd ya0(ya.size()), Lfya0(dya.size());
    ya0 << ya;
    Lfya0 << dya;

    MatrixXd Y(ya.size(), 2);
    Y << ya0, Lfya0;

    MatrixXd Phi_rem(2,2), a_bias(2,1), a_rem(2,1);
    MatrixXd a_temp(2,nParams), a_rest(1,nParams);
    a_temp.setZero();
    a_rest.setZero();

    a_temp(0,updateParam(0)) = 1;
    a_temp(1,updateParam(1)) = 1;

    VectorXd yd0(2), Lfyd0(2), yd0r(1), Lfyd0r(1);

    for (int i=0; i<nOutputs; ++i) {
        a_rest.setZero();

        for (int j=0; j<reserveParams.size(); ++j)
            a_rest(0,reserveParams(j)) = param(i,reserveParams(j));

        bezier_tools::bezier(a_rest, tau_0, yd0r);
        bezier_tools::dbezier(a_rest, tau_0, Lfyd0r);
        Lfyd0r = Lfyd0r * dtau_0;

        bezier_tools::bezier(a_temp, tau_0, yd0);
        bezier_tools::dbezier(a_temp, tau_0, Lfyd0);
        Lfyd0 = Lfyd0 * dtau_0;

        a_bias << yd0r, Lfyd0r;
        Phi_rem << yd0.transpose(), Lfyd0.transpose();

        a_rem = Phi_rem.inverse() * (Y.row(i).transpose() - a_bias);

        for (int j=0; j<updateParam.size(); ++j) {
            param(i,updateParam(j)) = a_rem(j);
        }
    }
}

void endTransition(PhaseVariable &phase, Eigen::VectorXd &ya, Eigen::MatrixXd &param) {
    int nParams = param.cols();
    int nOutputs = param.rows();

    VectorXi reserveParams(nParams - 2);
    VectorXi updateParam(2);
    updateParam << nParams-2, nParams-1;
    for (int i=0; i<reserveParams.size(); ++i)
        reserveParams(i) = i;

    double tau_0 = 1.0;
    double dtau_0 = phase.dtau;

    VectorXd dya(nOutputs);
    bezier_tools::dbezier(param, tau_0, dya);
    dya = dya * phase.dtau;

    VectorXd ya0(ya.size()), Lfya0(dya.size());
    ya0 << ya;
    Lfya0 << dya;

    MatrixXd Y(ya.size(), 2);
    Y << ya0, Lfya0;

    MatrixXd Phi_rem(2,2), a_bias(2,1), a_rem(2,1);
    MatrixXd a_temp(2,nParams), a_rest(1,nParams);
    a_temp.setZero();
    a_rest.setZero();

    a_temp(0,updateParam(0)) = 1;
    a_temp(1,updateParam(1)) = 1;

    VectorXd yd0(2), Lfyd0(2), yd0r(1), Lfyd0r(1);

    for (int i=0; i<nOutputs; ++i) {
        a_rest.setZero();

        for (int j=0; j<reserveParams.size(); ++j)
            a_rest(0,reserveParams(j)) = param(i,reserveParams(j));

        bezier_tools::bezier(a_rest, tau_0, yd0r);
        bezier_tools::dbezier(a_rest, tau_0, Lfyd0r);
        Lfyd0r = Lfyd0r * dtau_0;

        bezier_tools::bezier(a_temp, tau_0, yd0);
        bezier_tools::dbezier(a_temp, tau_0, Lfyd0);
        Lfyd0 = Lfyd0 * dtau_0;

        a_bias << yd0r, Lfyd0r;
        Phi_rem << yd0.transpose(), Lfyd0.transpose();

        a_rem = Phi_rem.inverse() * (Y.row(i).transpose() - a_bias);

        for (int j=0; j<updateParam.size(); ++j) {
            param(i,updateParam(j)) = a_rem(j);
        }
    }
}

}

// POWELL Method


/*************************************************************************\
License
    Copyright (c) 2017 Kavvadias Ioannis.

    This file is part of SPHSimulator.

    Licensed under the MIT License. See LICENSE file in the project root for
    full license information.

\************************************************************************/

#include <iostream>
#include <fstream>
#include <algorithm>

#include <glm/gtx/norm.hpp>

#include "HashTable.hpp"

#include "Statistics/Statistics.hpp"
#include "Util/ZValue.hpp"
#include "Kernels.hpp"

#include "Parallel/Parallel.hpp"

#include <boost/range/adaptor/indexed.hpp>
#include <boost/range/irange.hpp>
#include <boost/range/algorithm_ext.hpp>

using boost::adaptors::indexed;

//********************************************************************************
void HashTable::setHashTable(const glm::dvec2& minPos, const glm::dvec2& dx)
//********************************************************************************
{
  minPos_ = minPos-glm::dvec2(Kernel::SmoothingLength::h);
  gridSize_.x = unsigned(floor(dx.x*Kernel::SmoothingLength::dh))+2;
  gridSize_.y = unsigned(floor(dx.y*Kernel::SmoothingLength::dh))+2;

  particlesIn_.resize(gridSize_.x, gridSize_.y);
}

//********************************************************************************
void HashTable::writeGridRAW(std::string fileName) const
//********************************************************************************
{
  std::string fileNameGNU = fileName+".raw";
  std::cout<<"#writing file "<<fileNameGNU<<std::endl;
  std::ofstream outfile (fileNameGNU.c_str(),std::ofstream::binary);
  for (unsigned i=0;i<gridSize_.x;i++)
  {
    for (unsigned j=0;j<gridSize_.y;j++)
    {
      glm::dvec2 point1 = minPos_+glm::dvec2((i  )*Kernel::SmoothingLength::h,(j  )*Kernel::SmoothingLength::h);
      glm::dvec2 point2 = minPos_+glm::dvec2((i+1)*Kernel::SmoothingLength::h,(j  )*Kernel::SmoothingLength::h);
      glm::dvec2 point3 = minPos_+glm::dvec2((i+1)*Kernel::SmoothingLength::h,(j+1)*Kernel::SmoothingLength::h);
      glm::dvec2 point4 = minPos_+glm::dvec2((i  )*Kernel::SmoothingLength::h,(j+1)*Kernel::SmoothingLength::h);
      glm::dvec2 point5 = minPos_+glm::dvec2((i  )*Kernel::SmoothingLength::h,(j  )*Kernel::SmoothingLength::h);
      outfile<<point1.x<<"\t"<<point1.y<<"\n"
        <<point2.x<<"\t"<<point2.y<<"\n"
        <<point3.x<<"\t"<<point3.y<<"\n"
        <<point4.x<<"\t"<<point4.y<<"\n"
        <<point5.x<<"\t"<<point5.y<<"\n\n";
    }
  }
  outfile.close();
}

//********************************************************************************
void HashTable::write() const
//********************************************************************************
{
  for (unsigned j=0;j<gridSize_.y;j++)
  {
    for (unsigned i=0;i<gridSize_.x;i++)
    {
      std::cout<<" "<<particlesIn_(i,j).size();
    }
    std::cout<<std::endl;
  }
}

//********************************************************************************
void HashTable::clear()
//********************************************************************************
{
  for (unsigned i=0;i<gridSize_.x;i++)
  {
    for (unsigned j=0;j<gridSize_.y;j++)
    {
      particlesIn_(i,j).clear();
    }
  }
}

//********************************************************************************
glm::uvec2 HashTable::findGridPos(glm::dvec2 pos) const
//********************************************************************************
{
  static auto findGridPosTimerID= Statistics::createTimer("HashTable::findGridPos");
  Statistics::TimerGuard findGridPosTimerGuard(findGridPosTimerID);

  glm::dvec2 dgrdPos = (pos-minPos_)*Kernel::SmoothingLength::dh;
  if ( dgrdPos.x < 0. ) dgrdPos.x = 0.;
  if ( dgrdPos.y < 0. ) dgrdPos.y = 0.;
  glm::uvec2 gridPos = glm::ivec2(unsigned(floor(dgrdPos.x)),unsigned(floor(dgrdPos.y)));
  if ( gridPos.x >= gridSize_.x ) gridPos.x = gridSize_.x;
  if ( gridPos.y >= gridSize_.y ) gridPos.y = gridSize_.y;

  return gridPos;
}

//********************************************************************************
void HashTable::findNei(ParticleCloud& cloud)
//********************************************************************************
{
  static auto findNeiTimerID = Statistics::createTimer("HashTable::findNei");
  Statistics::TimerGuard findNeiTimerGuard(findNeiTimerID);
  //static auto findNeiTimerID_1 = Statistics::createTimer("HashTable::findNei_P1");
  //static auto findNeiTimerID_2 = Statistics::createTimer("HashTable::findNei_P2");
  //static auto findNeiTimerID_3 = Statistics::createTimer("HashTable::findNei_P3");

  const auto& particlePos = cloud.get<Attr::ePosition>();

  auto& particleNei = cloud.get<Attr::eNei>();

  particleGridPos_.resize(cloud.size());

  const size_t nPart = cloud.size();

  {
    //Statistics::TimerGuard findNeiTimerGuard_1(findNeiTimerID_1);
    const double dh = Kernel::SmoothingLength::dh;
    //find grid position of each particle
    Parallel::For (nPart, [this, &particlePos, dh](size_t iPart){
      const glm::dvec2 pos = particlePos[iPart];
      const glm::dvec2 dgrdPos = (pos-minPos_)*dh;
      const glm::uvec2 gridPos = glm::uvec2(unsigned(floor(dgrdPos.x)),unsigned(floor(dgrdPos.y)));

      particleGridPos_[iPart] = gridPos;
    });

    for (const auto& elem : particleGridPos_ | indexed())
    {
      particlesIn_(elem.value()).push_back(unsigned(elem.index()));
    }
  }

  {
    //Statistics::TimerGuard findNeiTimerGuard_2(findNeiTimerID_2);
    const double h2 = Kernel::SmoothingLength::h2;
    auto makeRange = [](unsigned val, unsigned max) {
      return boost::irange(
          (val > 0) ? val-1 : val,
          (val == max-1) ? max : val+2
      );
    };

    //find nei of each particle
    Parallel::For (nPart, [this, &particlePos, &particleNei, makeRange, h2](size_t iPart){
      const glm::dvec2 iPos = particlePos[iPart];
      const glm::uvec2 iGridPos = particleGridPos_[iPart];
      auto& iNei = particleNei[iPart];

      auto iRange = makeRange(iGridPos.x, gridSize_.x);
      auto jRange = makeRange(iGridPos.y, gridSize_.y);

      for (unsigned iGrid : iRange)
      for (unsigned jGrid : jRange)
      for (unsigned neiPos : particlesIn_(iGrid, jGrid)) {
        const double dist2 = glm::length2(iPos - particlePos[neiPos]);
        if (dist2 < h2)
        {
          iNei.push_back(Neigbhor(neiPos));
        }
      }
    });
  }

  {
    //Statistics::TimerGuard findNeiTimerGuard_3(findNeiTimerID_3);
    auto comp = [&particlePos](const Neigbhor& i, const Neigbhor& j){
      return particlePos[i.ID].x < particlePos[j.ID].x;
    };

    Parallel::For (nPart, [ &particleNei, comp ](size_t iPart){
      auto& nei = particleNei[iPart];
      std::sort(nei.begin(), nei.end(), comp);
    });
  }
}

//********************************************************************************
std::vector<unsigned> HashTable::neiParticlesFor(glm::uvec2 gridPos) const
//********************************************************************************
{
  std::vector<unsigned> ret;

  auto makeRange = [](unsigned val, unsigned max) {
    return boost::irange(
        (val > 0) ? val-1 : val,
        (val == max-1) ? max : val+2
    );
  };
  auto iRange = makeRange(gridPos.x, gridSize_.x);
  auto jRange = makeRange(gridPos.y, gridSize_.y);

  for (unsigned iGrid : iRange)
  for (unsigned jGrid : jRange)
  {
    boost::push_back(ret, particlesIn_(iGrid,jGrid));
  }
  return ret;
}