//
// Created by gregor on 31.07.19.
//

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE unit_tests

#include <boost/test/unit_test.hpp>

namespace utf = boost::unit_test;
namespace tt = boost::test_tools;

#include <vector>
#include "../src/datatypes.h"
#include "../src/neighbors.h"

#include "../src/clustering.h"
#include "../src/core.h"
#include "../src/cnn.h"
#include "../src/vs_cnn.h"
#include "../src/geometry.h"

struct dataFixture {
    dataFixture() {
        nns::neighbors(shrt_neighbor_lists,
                       shrt,
                       fixed_cut,
                       0);

        nns::neighbors(mdm_neighbor_lists,
                       mdm,
                       fixed_cut,
                       0);

        nns::neighbors(lng_neighbor_lists,
                       lng,
                       fixed_cut,
                       0);
    };

    ~dataFixture() {};

    const float fixed_cut = 2.0 * std::sqrt(3);
    const unsigned int fixed_sim = 2;
    Neighbors shrt_neighbor_lists;
    Neighbors mdm_neighbor_lists;
    Neighbors lng_neighbor_lists;

    std::vector<std::vector<float>> shrt = {{1.,  2.,  3.},
                                            {2.,  3.,  4.},
                                            {3.,  4.,  5.},
                                            {4.,  5.,  6.},
                                            {5.,  6.,  7.},
                                            {6.,  7.,  8.},
                                            {7.,  8.,  9.},
                                            {91., 92., 93.},
                                            {92., 93., 94.},
                                            {93., 94., 95.},
                                            {94., 95., 96.},
                                            {95., 96., 97.},
                                            {96., 97., 98.},
                                            {97., 98., 99.}};

    std::vector<std::vector<float>> mdm = {{0.,  1.,  2.},
                                           {1.,  2.,  3.},
                                           {2.,  3.,  4.},
                                           {3.,  4.,  5.},
                                           {4.,  5.,  6.},
                                           {5.,  6.,  7.},
                                           {6.,  7.,  8.},
                                           {7.,  8.,  9.},
                                           {10,  11., 12.},
                                           {91., 92., 93.},
                                           {92., 93., 94.},
                                           {93., 94., 95.},
                                           {94., 95., 96.},
                                           {95., 96., 97.},
                                           {96., 97., 98.},
                                           {97., 98., 99.}};

    std::vector<std::vector<float>> lng = {{0.,  1.,  2.},
                                           {1.,  2.,  3.},
                                           {2.,  3.,  4.},
                                           {3.,  4.,  5.},
                                           {4.,  5.,  6.},
                                           {5.,  6.,  7.},
                                           {6.,  7.,  8.},
                                           {7.,  8.,  9.},
                                           {8.,  9.,  10.},
                                           {9.,  10., 11.},
                                           {20., 21., 22.},
                                           {21., 22., 23.},
                                           {22., 23., 24.},
                                           {23., 24., 25.},
                                           {24., 25., 26.},
                                           {25., 26., 27.},
                                           {26., 27., 28.},
                                           {27., 28., 29.},
                                           {28., 29., 30.},
                                           {29., 30., 31.},
                                           {91., 92., 93.},
                                           {92., 93., 94.},
                                           {93., 94., 95.},
                                           {94., 95., 96.},
                                           {95., 96., 97.},
                                           {96., 97., 98.},
                                           {97., 98., 99.}};
};

BOOST_FIXTURE_TEST_SUITE(NeighborsTestSuite, dataFixture)

    BOOST_AUTO_TEST_CASE(neighbors) {

        const float cut = std::sqrt(3);
        const unsigned int sim = 0;

        Neighbors neighbor_lists;
        nns::neighbors(neighbor_lists,
                       shrt,
                       cut,
                       sim);

        for (auto list : neighbor_lists) {
            if (list.first == 0 || list.first == 7) {
                BOOST_CHECK_EQUAL(list.second[0], list.first + 1);
            } else if (list.first == 6 || list.first == 13) {
                BOOST_CHECK_EQUAL(list.second[0], list.first - 1);
            } else {
                BOOST_CHECK_EQUAL(list.second[0], list.first - 1);
                BOOST_CHECK_EQUAL(list.second[1], list.first + 1);
            }
        }
    }

    BOOST_AUTO_TEST_CASE(neighbor_pruning) {

        const float cut = std::sqrt(3);
        const unsigned int sim = 1;

        Neighbors neighbor_lists;
        nns::neighbors(neighbor_lists,
                       lng,
                       cut,
                       sim);

        for (auto list : neighbor_lists) {
            BOOST_CHECK_EQUAL(list.second[0], list.first - 1);
            BOOST_CHECK_EQUAL(list.second[1], list.first + 1);
        }
    }

BOOST_AUTO_TEST_SUITE_END()

BOOST_FIXTURE_TEST_SUITE(CNNTestSuite, dataFixture)

    BOOST_AUTO_TEST_CASE(CNNsimilarity) {

        const float cut = 2.0f * std::sqrt(3);
        unsigned int sim = 0;
        Neighbors neighbor_lists;
        nns::neighbors(neighbor_lists,
                       shrt,
                       cut,
                       sim);
        sim = 2;

        BOOST_CHECK(!Clustering::CommonNearestNeighbor::similarity(shrt,
                                                                   neighbor_lists,
                                                                   0, 1,
                                                                   cut, sim));
        BOOST_CHECK(Clustering::CommonNearestNeighbor::similarity(shrt,
                                                                  neighbor_lists,
                                                                  2, 1,
                                                                  cut, sim));
    }

BOOST_AUTO_TEST_SUITE_END()

BOOST_FIXTURE_TEST_SUITE(vsCNNTestSuite, dataFixture)

    BOOST_AUTO_TEST_CASE(vsCNNdistance, *utf::tolerance(0.00001)) {
        BOOST_TEST(Clustering::CommonDensity::calc_distance(shrt[0], shrt[2]) == std::sqrt(12));
    }

    BOOST_AUTO_TEST_CASE(regularized_intersection_volume, *utf::tolerance(0.00001)) {
        BOOST_TEST(Geometry::regularized_intersection_volume(0.0f, std::sqrt(3.0f), 3) == 1.0f);
        BOOST_TEST(Geometry::regularized_intersection_volume(std::sqrt(3.0f), std::sqrt(3.0f), 3) == 0.3125f);
        BOOST_TEST(Geometry::regularized_intersection_volume(2.0f * std::sqrt(3.0f), std::sqrt(3.0f), 3) == 0.0f);
    }

    BOOST_AUTO_TEST_CASE(vsCNNsimilarity) {

        const float cut = 2.0f * std::sqrt(3);
        unsigned int sim = 0;
        Neighbors neighbor_lists;
        nns::neighbors(neighbor_lists,
                       shrt,
                       cut,
                       sim);
        sim = 2;

        // This is a different result to the CNN similarity and the reason
        // why the vsCNN clustering is more 'crisp'.
        BOOST_CHECK(Clustering::CommonDensity::similarity(shrt,
                                                          neighbor_lists,
                                                          0, 1,
                                                          cut, sim));
        BOOST_CHECK(Clustering::CommonDensity::similarity(shrt,
                                                          neighbor_lists,
                                                          2, 1,
                                                          cut, sim));
    }

BOOST_AUTO_TEST_SUITE_END()

BOOST_FIXTURE_TEST_SUITE(CoreTestSuite, dataFixture)

    BOOST_AUTO_TEST_CASE(intersection) {
        const vector<unsigned int> list1 = {0, 1, 2, 3};
        const vector<unsigned int> list2 = {2, 3, 4, 5};
        const vector<unsigned int> list3 = {4, 5, 6, 7};

        vector<unsigned int> out1(0);
        Clustering::Core::intersection(out1, list1, list2);

        vector<unsigned int> out2(0);
        Clustering::Core::intersection(out2, list1, list3);

        BOOST_CHECK(out1[0] == 2);
        BOOST_CHECK(out1[1] == 3);
        BOOST_CHECK(out2.empty());
    }

    BOOST_AUTO_TEST_CASE(similarity) {

        Clustering::Core::simptr similarity = Clustering::CommonNearestNeighbor::similarity;

        map<unsigned int, int> clustered;
        vector<vector<unsigned int> > clusters;

        unsigned int refpoint = 2;
        Clustering::Core::similarity_unclustered(similarity,
                                                 shrt,
                                                 clustered,
                                                 clusters,
                                                 shrt_neighbor_lists,
                                                 shrt_neighbor_lists[refpoint],
                                                 refpoint,
                                                 fixed_cut,
                                                 fixed_sim);

        refpoint = 11;
        Clustering::Core::similarity_unclustered(similarity,
                                                 shrt,
                                                 clustered,
                                                 clusters,
                                                 shrt_neighbor_lists,
                                                 shrt_neighbor_lists[refpoint],
                                                 refpoint,
                                                 fixed_cut,
                                                 fixed_sim);

        // Sort because of parallel loop in similarity_unclustered
        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK(clusters.size() == 2);
        BOOST_CHECK(clusters[0].size() == 3);
        BOOST_CHECK(clusters[1].size() == 3);
        BOOST_CHECK(clusters[0][0] == 3);
        BOOST_CHECK(clusters[0][1] == 2);
        BOOST_CHECK(clusters[0][2] == 1);
        BOOST_CHECK(clusters[1][0] == 12);
        BOOST_CHECK(clusters[1][1] == 11);
        BOOST_CHECK(clusters[1][2] == 10);

        refpoint = 3;
        Clustering::Core::similarity_clustered(similarity,
                                               shrt,
                                               clustered,
                                               clusters,
                                               shrt_neighbor_lists,
                                               shrt_neighbor_lists[refpoint],
                                               refpoint,
                                               fixed_cut,
                                               fixed_sim);

        refpoint = 10;
        Clustering::Core::similarity_clustered(similarity,
                                               shrt,
                                               clustered,
                                               clusters,
                                               shrt_neighbor_lists,
                                               shrt_neighbor_lists[refpoint],
                                               refpoint,
                                               fixed_cut,
                                               fixed_sim);

        // Sort because of parallel loop in similarity_unclustered
        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK(clusters.size() == 2);
        BOOST_CHECK(clusters[0].size() == 4);
        BOOST_CHECK(clusters[1].size() == 4);
        BOOST_CHECK(clusters[0][0] == 4);
        BOOST_CHECK(clusters[0][1] == 3);
        BOOST_CHECK(clusters[0][2] == 2);
        BOOST_CHECK(clusters[0][3] == 1);
        BOOST_CHECK(clusters[1][0] == 12);
        BOOST_CHECK(clusters[1][1] == 11);
        BOOST_CHECK(clusters[1][2] == 10);
        BOOST_CHECK(clusters[1][3] == 9);
    }

    BOOST_AUTO_TEST_CASE(CNNalgorithm) {

        const float cut = 5.0;
        const unsigned int sim = 2;

        Neighbors dummy_neighbors;
        vector<vector<unsigned int> > clusters;
        clusters = Clustering::Core::algorithm(Clustering::CommonNearestNeighbor::similarity,
                                               mdm,
                                               mdm_neighbor_lists,
                                               dummy_neighbors,
                                               cut,
                                               sim,
                                               0,
                                               true);

        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK_EQUAL(clusters.size(), 2);
        BOOST_CHECK_EQUAL(clusters[0].size(), 6);
        BOOST_CHECK_EQUAL(clusters[1].size(), 5);
        BOOST_CHECK_EQUAL(clusters[0][0], 6);
        BOOST_CHECK_EQUAL(clusters[0][1], 5);
        BOOST_CHECK_EQUAL(clusters[0][2], 4);
        BOOST_CHECK_EQUAL(clusters[0][3], 3);
        BOOST_CHECK_EQUAL(clusters[0][4], 2);
        BOOST_CHECK_EQUAL(clusters[0][5], 1);
        BOOST_CHECK_EQUAL(clusters[1][0], 14);
        BOOST_CHECK_EQUAL(clusters[1][1], 13);
        BOOST_CHECK_EQUAL(clusters[1][2], 12);
        BOOST_CHECK_EQUAL(clusters[1][3], 11);
        BOOST_CHECK_EQUAL(clusters[1][4], 10);

    }

    BOOST_AUTO_TEST_CASE(vsCNNalgorithm) {

        const float cut = 5.0;
        const unsigned int sim = 2;

        Neighbors dummy_neighbors;
        vector<vector<unsigned int> > clusters;
        clusters = Clustering::Core::algorithm(Clustering::CommonDensity::similarity,
                                               mdm,
                                               mdm_neighbor_lists,
                                               dummy_neighbors,
                                               cut,
                                               sim,
                                               0,
                                               true);

        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK_EQUAL(clusters.size(), 2);
        BOOST_CHECK_EQUAL(clusters[0].size(), 8);
        BOOST_CHECK_EQUAL(clusters[1].size(), 7);
        BOOST_CHECK_EQUAL(clusters[0][0], 7);
        BOOST_CHECK_EQUAL(clusters[0][1], 6);
        BOOST_CHECK_EQUAL(clusters[0][2], 5);
        BOOST_CHECK_EQUAL(clusters[0][3], 4);
        BOOST_CHECK_EQUAL(clusters[0][4], 3);
        BOOST_CHECK_EQUAL(clusters[0][5], 2);
        BOOST_CHECK_EQUAL(clusters[0][6], 1);
        BOOST_CHECK_EQUAL(clusters[0][7], 0);
        BOOST_CHECK_EQUAL(clusters[1][0], 15);
        BOOST_CHECK_EQUAL(clusters[1][1], 14);
        BOOST_CHECK_EQUAL(clusters[1][2], 13);
        BOOST_CHECK_EQUAL(clusters[1][3], 12);
        BOOST_CHECK_EQUAL(clusters[1][4], 11);
        BOOST_CHECK_EQUAL(clusters[1][5], 10);
        BOOST_CHECK_EQUAL(clusters[1][6], 9);
    }

BOOST_AUTO_TEST_SUITE_END()

BOOST_FIXTURE_TEST_SUITE(ClusteringTestSuite, dataFixture)

    BOOST_AUTO_TEST_CASE(clustering_cnn) {

        const float cut = 5.0;
        const unsigned int sim = 2;

        vector<vector<unsigned int> > clusters;
        clusters = Clustering::clustering(Clustering::CommonNearestNeighbor::similarity,
                                          mdm,
                                          cut,
                                          sim,
                                          0,
                                          true);

        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK_EQUAL(clusters.size(), 2);
        BOOST_CHECK_EQUAL(clusters[0].size(), 6);
        BOOST_CHECK_EQUAL(clusters[1].size(), 5);
        BOOST_CHECK_EQUAL(clusters[0][0], 6);
        BOOST_CHECK_EQUAL(clusters[0][1], 5);
        BOOST_CHECK_EQUAL(clusters[0][2], 4);
        BOOST_CHECK_EQUAL(clusters[0][3], 3);
        BOOST_CHECK_EQUAL(clusters[0][4], 2);
        BOOST_CHECK_EQUAL(clusters[0][5], 1);
        BOOST_CHECK_EQUAL(clusters[1][0], 14);
        BOOST_CHECK_EQUAL(clusters[1][1], 13);
        BOOST_CHECK_EQUAL(clusters[1][2], 12);
        BOOST_CHECK_EQUAL(clusters[1][3], 11);
        BOOST_CHECK_EQUAL(clusters[1][4], 10);
    }

    BOOST_AUTO_TEST_CASE(clustering_vscnn) {

        const float cut = 5.0;
        const unsigned int sim = 2;

        vector<vector<unsigned int> > clusters;
        clusters = Clustering::clustering(Clustering::CommonDensity::similarity,
                                          mdm,
                                          cut,
                                          sim,
                                          0,
                                          true);

        for (auto &cluster : clusters) {
            std::sort(cluster.begin(), cluster.end(),
                      [](const unsigned int &a, const unsigned int &b) { return a > b; });
        }

        BOOST_CHECK_EQUAL(clusters.size(), 2);
        BOOST_CHECK_EQUAL(clusters[0].size(), 8);
        BOOST_CHECK_EQUAL(clusters[1].size(), 7);
        BOOST_CHECK_EQUAL(clusters[0][0], 7);
        BOOST_CHECK_EQUAL(clusters[0][1], 6);
        BOOST_CHECK_EQUAL(clusters[0][2], 5);
        BOOST_CHECK_EQUAL(clusters[0][3], 4);
        BOOST_CHECK_EQUAL(clusters[0][4], 3);
        BOOST_CHECK_EQUAL(clusters[0][5], 2);
        BOOST_CHECK_EQUAL(clusters[0][6], 1);
        BOOST_CHECK_EQUAL(clusters[0][7], 0);
        BOOST_CHECK_EQUAL(clusters[1][0], 15);
        BOOST_CHECK_EQUAL(clusters[1][1], 14);
        BOOST_CHECK_EQUAL(clusters[1][2], 13);
        BOOST_CHECK_EQUAL(clusters[1][3], 12);
        BOOST_CHECK_EQUAL(clusters[1][4], 11);
        BOOST_CHECK_EQUAL(clusters[1][5], 10);
        BOOST_CHECK_EQUAL(clusters[1][6], 9);
    }

BOOST_AUTO_TEST_SUITE_END()


#ifndef STAN_MATH_PRIM_MAT_PROB_LKJ_COV_LOG_HPP
#define STAN_MATH_PRIM_MAT_PROB_LKJ_COV_LOG_HPP

#include <stan/math/prim/mat/fun/Eigen.hpp>
#include <stan/math/prim/mat/prob/lkj_cov_lpdf.hpp>
#include <boost/math/tools/promotion.hpp>

namespace stan {
  namespace math {

    /**
     * @deprecated use lkj_cov_lpdf
     */
    template <bool propto,
              typename T_y, typename T_loc, typename T_scale, typename T_shape>
    typename
    boost::math::tools::promote_args<T_y, T_loc, T_scale, T_shape>::type
    lkj_cov_log(const Eigen::Matrix<T_y, Eigen::Dynamic, Eigen::Dynamic>& y,
                const Eigen::Matrix<T_loc, Eigen::Dynamic, 1>& mu,
                const Eigen::Matrix<T_scale, Eigen::Dynamic, 1>& sigma,
                const T_shape& eta) {
      return lkj_cov_lpdf<propto, T_y, T_loc,
                          T_scale, T_shape>(y, mu, sigma, eta);
    }

    /**
     * @deprecated use lkj_cov_lpdf
     */
    template <typename T_y, typename T_loc, typename T_scale, typename T_shape>
    inline
    typename
    boost::math::tools::promote_args<T_y, T_loc, T_scale, T_shape>::type
    lkj_cov_log(const Eigen::Matrix<T_y, Eigen::Dynamic, Eigen::Dynamic>& y,
                const Eigen::Matrix<T_loc, Eigen::Dynamic, 1>& mu,
                const Eigen::Matrix<T_scale, Eigen::Dynamic, 1>& sigma,
                const T_shape& eta) {
      return lkj_cov_lpdf<T_y, T_loc, T_scale, T_shape>(y, mu, sigma, eta);
    }

    /**
     * @deprecated use lkj_cov_lpdf
     */
    template <bool propto,
              typename T_y, typename T_loc, typename T_scale, typename T_shape>
    typename
    boost::math::tools::promote_args<T_y, T_loc, T_scale, T_shape>::type
    lkj_cov_log(const Eigen::Matrix<T_y, Eigen::Dynamic, Eigen::Dynamic>& y,
                const T_loc& mu,
                const T_scale& sigma,
                const T_shape& eta) {
      return lkj_cov_lpdf<propto, T_y, T_loc,
                          T_scale, T_shape>(y, mu, sigma, eta);
    }

    /**
     * @deprecated use lkj_cov_lpdf
     */
    template <typename T_y, typename T_loc, typename T_scale, typename T_shape>
    inline
    typename boost::math::tools::promote_args
    <T_y, T_loc, T_scale, T_shape>::type
    lkj_cov_log(const Eigen::Matrix<T_y, Eigen::Dynamic, Eigen::Dynamic>& y,
                const T_loc& mu,
                const T_scale& sigma,
                const T_shape& eta) {
      return lkj_cov_lpdf<T_y, T_loc,
                          T_scale, T_shape>(y, mu, sigma, eta);
    }

  }
}
#endif


//==================================================================================================
/*!
  @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_FUNCTION_TOFLOAT_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_TOFLOAT_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

  /*!

    @ingroup group-arithmetic
    This function object converts its parameter to floating point value occupying the same bit size.
    This implies that this functor is only defined for (smd vectors of) 32 or 64 bit integer
    entries to provide (smd vectors of) float and or double.

    @see toint, touint, cast

    @par Header <boost/simd/function/tofloat.hpp>

    @par Example:

      @snippet tofloat.cpp tofloat

    @par Possible output:

      @snippet tofloat.txt tofloat
  **/
  as_floating_t<Value> tofloat(Value const& x);
} }
#endif

#include <boost/simd/function/scalar/tofloat.hpp>
#include <boost/simd/function/simd/tofloat.hpp>

#endif


#ifndef ASLAM_BACKEND_VECTOR_EXPRESSION_HPP
#define ASLAM_BACKEND_VECTOR_EXPRESSION_HPP

#include <aslam/backend/JacobianContainer.hpp>
#include <boost/shared_ptr.hpp>
#include <sm/boost/null_deleter.hpp>
#include "VectorExpressionNode.hpp"
#include <aslam/backend/ScalarExpression.hpp>
#include <aslam/backend/ScalarExpressionNode.hpp>

namespace aslam {
  namespace backend {

    template<int D>
    class VectorExpression
    {
    public:
      typedef Eigen::Matrix<double,D,1> vector_t;
      typedef Eigen::Matrix<double,D,1> value_t;

      VectorExpression(boost::shared_ptr< VectorExpressionNode<D> > root);
      VectorExpression(VectorExpressionNode<D> * root);
      virtual ~VectorExpression();

      vector_t evaluate() const;
      vector_t toValue() const;

      void evaluateJacobians(JacobianContainer & outJacobians) const;
      void evaluateJacobians(JacobianContainer & outJacobians, const Eigen::MatrixXd & applyChainRule) const;

      void getDesignVariables(DesignVariable::set_t & designVariables) const;

      ScalarExpression toScalarExpression() const;

      boost::shared_ptr< VectorExpressionNode<D> > root() const { return _root; }

    private:
      boost::shared_ptr< VectorExpressionNode<D> > _root;
    };

  } // namespace backend
} // namespace aslam

#include "implementation/VectorExpression.hpp"

#endif /* ASLAM_BACKEND_VECTOR_EXPRESSION_HPP */


#include "dynet/training.h"
#include "dynet/timing.h"
#include "dynet/lstm.h"
#include "dynet/dict.h"
#include "dynet/expr.h"
#include "dynet/globals.h"

#include <algorithm>
#include <iostream>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <set>

#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>

#include <boost/algorithm/string.hpp>
#include <boost/program_options.hpp>

using namespace std;
using namespace dynet;

namespace po = boost::program_options;

struct PKey {
  PKey(int x1, int x2, unsigned x3)
    : t(x1,x2,x3) {}
  tuple<int, int, unsigned> t;
};

bool operator==(const PKey &a, const PKey &b)
{
  return std::get<0>(a.t) == std::get<0>(b.t) &&
         std::get<1>(a.t) == std::get<1>(b.t) &&
         std::get<2>(a.t) == std::get<2>(b.t);
}

namespace std {
  template <>
  struct hash<PKey>
  {
    std::size_t operator()(const PKey& e) const
    {
      return (std::get<0>(e.t)) | (std::get<1>(e.t) << 16) | (std::get<2>(e.t) << 24);
    }
  };

}

unsigned LAYERS = 1;
unsigned INPUT_DIM = 64;
unsigned XCRIBE_DIM = 32;
unsigned SEG_DIM = 16;
unsigned H1DIM = 32;
unsigned H2DIM = 32;
unsigned TAG_DIM = 32;
unsigned DURATION_DIM = 4;

unsigned int DATA_MAX_SEG_LEN = 0;

bool use_pretrained_embeding = false;
bool use_dropout = false;
float dropout_rate;
bool ner_tagging = false;
string pretrained_embeding = "";

// returns embeddings of labels
struct SymbolEmbedding {
  SymbolEmbedding(ParameterCollection& m, unsigned n, unsigned dim) {
    p_labels = m.add_lookup_parameters(n, {dim});
  }
  void load_embedding(dynet::Dict& d, string pretrain_path){
    ifstream fin(pretrain_path);
       string s;
       while( getline(fin,s) )
       {
        vector <string> fields;
        boost::algorithm::trim(s);
        boost::algorithm::split( fields, s, boost::algorithm::is_any_of( " " ) );
        string word = fields[0];
        vector<float> p_embeding;
        for (int ind = 1; ind < fields.size(); ++ind){
          p_embeding.push_back(std::stod(fields[ind]));
        }
        if (d.contains(word)){
          // cout << "init" << endl;
          p_labels.initialize(d.convert(word), p_embeding);
        }
      }
  }
  void new_graph(ComputationGraph& g) { cg = &g; }
  Expression embed(unsigned label_id) {
    return lookup(*cg, p_labels, label_id);
  }
  ComputationGraph* cg;
  LookupParameter p_labels;
};

struct DurationEmbedding {
  virtual ~DurationEmbedding();
  virtual void new_graph(ComputationGraph& g) = 0;
  virtual Expression embed(unsigned dur) = 0;
};

DurationEmbedding::~DurationEmbedding() {}

struct MLPDurationEmbedding : public DurationEmbedding {
  MLPDurationEmbedding(ParameterCollection& m, unsigned hidden, unsigned dim) {
    p_zero = m.add_parameters({dim});
    p_d2h = m.add_parameters({hidden, 2});
    p_hb = m.add_parameters({hidden});
    p_h2o = m.add_parameters({dim, hidden});
    p_ob = m.add_parameters({dim});
    dur_xs.resize(10000, vector<float>(2));
    for (unsigned i = 1; i < dur_xs.size(); ++i) {
      dur_xs[i][0] = i;
      dur_xs[i][1] = float(i) / logf(2);
    }
  }
  void new_graph(ComputationGraph& g) override {
    zero = parameter(g, p_zero);
    d2h = parameter(g, p_d2h);
    hb = parameter(g, p_hb);
    h2o = parameter(g, p_h2o);
    ob = parameter(g, p_ob);
    cg = &g;
  }
  Expression embed(unsigned dur) override {
    if (dur) {
      Expression x = input(*cg, {2}, &dur_xs[dur]);
      Expression h = rectify(affine_transform({hb, d2h, x}));
      return affine_transform({ob, h2o, h});
    } else {
      return zero;
    }
  }

  ComputationGraph* cg;
  vector<vector<float>> dur_xs;
  Expression zero, d2h, hb, h2o, ob;
  Parameter p_zero;
  Parameter p_d2h;
  Parameter p_hb;
  Parameter p_h2o;
  Parameter p_ob;
};

struct BinnedDurationEmbedding : public DurationEmbedding {
  BinnedDurationEmbedding(ParameterCollection& m, unsigned dim, unsigned num_bins = 8) : max_bin(num_bins - 1) {
    p_e = m.add_lookup_parameters(num_bins, {dim});
  }
  void new_graph(ComputationGraph& g) override {
    cg = &g;
  }
  Expression embed(unsigned dur) override {
    if (dur) dur = static_cast<unsigned>(log(dur) / log(1.6)) + 1;
    if (dur > max_bin) { dur = max_bin; }
    return lookup(*cg, p_e, dur);
  }
  ComputationGraph* cg;
  int max_bin;
  LookupParameter p_e;
};

template <class Builder>
struct BiTrans {
  Builder l2rbuilder;
  Builder r2lbuilder;
  Parameter p_f2c;
  Parameter p_r2c;
  Parameter p_cb;

  explicit BiTrans(ParameterCollection& model) :
      l2rbuilder(LAYERS, INPUT_DIM, XCRIBE_DIM, model),
      r2lbuilder(LAYERS, INPUT_DIM, XCRIBE_DIM, model) {
    p_f2c = model.add_parameters({XCRIBE_DIM, XCRIBE_DIM});
    p_r2c = model.add_parameters({XCRIBE_DIM, XCRIBE_DIM});
    p_cb = model.add_parameters({XCRIBE_DIM});
  }

  vector<Expression> transcribe(ComputationGraph& cg, const vector<Expression>& x) {
    l2rbuilder.new_graph(cg);
    if(use_dropout){
      l2rbuilder.set_dropout(dropout_rate);
    }else{
      l2rbuilder.disable_dropout();
    }
    l2rbuilder.start_new_sequence();
    r2lbuilder.new_graph(cg);
    if(use_dropout){
      r2lbuilder.set_dropout(dropout_rate);
    }else{
      r2lbuilder.disable_dropout();
    }
    r2lbuilder.start_new_sequence();
    Expression f2c = parameter(cg, p_f2c);
    Expression r2c = parameter(cg, p_r2c);
    Expression cb = parameter(cg, p_cb);

    const int len = x.size();
    vector<Expression> fwd(len), rev(len), res(len);
    for (int i = 0; i < len; ++i)
      fwd[i] = l2rbuilder.add_input(x[i]);
    for (int i = len - 1; i >= 0; --i)
      rev[i] = r2lbuilder.add_input(x[i]);
    for (int i = 0; i < len; ++i)
      res[i] = affine_transform({cb, f2c, fwd[i], r2c, rev[i]});
    return res;
  }
};

// uni-directional segment embeddings
// first call construct_chart(sequence), then access the embeddings of the spans with operator()(i,j)
template <class Builder>
struct SegEmbedUni {
  Parameter p_h0;
  int len;
  Builder builder;
  vector<vector<Expression>> h;  // h[i][length of segment - 1]
  explicit SegEmbedUni(ParameterCollection& m) :
      builder(LAYERS, XCRIBE_DIM, SEG_DIM, m) {
    p_h0 = m.add_parameters({XCRIBE_DIM});
  }
  void construct_chart(ComputationGraph& cg, const vector<Expression>& c, int max_seg_len = 0) {
    len = c.size();
    h.clear();
    h.resize(len);
    Expression h0 = parameter(cg, p_h0);
    builder.new_graph(cg);
    if(use_dropout){
      builder.set_dropout(dropout_rate);
    }else{
      builder.disable_dropout();
    }
    for (int i = 0; i < len; ++i) {
      int max_j = i + len;
      if (max_seg_len) max_j = i + max_seg_len;
      if (max_j > len) max_j = len;
      int seg_len = max_j - i;
      auto& hi = h[i];
      hi.resize(seg_len);

      builder.start_new_sequence();
      builder.add_input(h0);
      for (int k = 0; k < seg_len; ++k)
        hi[k] = builder.add_input(c[i+k]);
    }
  }
  // returns the embedding of segment [i,j]
  const Expression& operator()(int i, int j) const {
    assert(j < len);
    assert(j >= i);
    return h[i][j - i];
  }
};

// uni-directional segment embeddings
// first call construct_chart(sequence), then access the embeddings of the spans with operator()(i,j)
template <class Builder>
struct SegEmbedBi {
  int len;
  vector<vector<pair<Expression, Expression>>> h;
  SegEmbedUni<Builder> fwd, rev;
  explicit SegEmbedBi(ParameterCollection& m) : fwd(m), rev(m) {}
  void construct_chart(ComputationGraph& cg, const vector<Expression>& c, int max_seg_len = 0) {
    len = c.size();
    fwd.construct_chart(cg, c, max_seg_len);
    vector<Expression> rc(len);
    for (int i = 0; i < len; ++i) rc[i] = c[len - 1 - i];
    rev.construct_chart(cg, rc, max_seg_len);
    h.clear();
    h.resize(len);
    for (int i = 0; i < len; ++i) {
      int max_j = i + len;
      if (max_seg_len) max_j = i + max_seg_len;
      if (max_j > len) max_j = len;
      int seg_len = max_j - i;
      auto& hi = h[i];
      hi.resize(seg_len);
      for (int k = 0; k < seg_len; ++k) {
        int j = i + k;
        const Expression& fe = fwd(i, j);
        const Expression& re = rev(len-1-j, len-1-i);
        hi[k] = make_pair(fe,re);
      }
    }
  }
  const pair<Expression, Expression>& operator()(int i, int j) const {
    assert(j >= i);
    assert(j < len);
    return h[i][j-i];
  }
};

template <class Builder>
struct SegmentalRNN {
  SymbolEmbedding* xe;
  SymbolEmbedding* ye;
  DurationEmbedding* de;
  BiTrans<Builder> bt;
  SegEmbedBi<Builder> seb;
  dynet::Dict d;
  dynet::Dict td;
  explicit SegmentalRNN(ParameterCollection& model, dynet::Dict& d_, dynet::Dict& td_) :
      bt(model), seb(model) {
    d = d_;
    td = td_;
    xe = new SymbolEmbedding(model, d.size(), INPUT_DIM);
    if (use_pretrained_embeding) {
       xe->load_embedding(d, pretrained_embeding);
    }
    ye = new SymbolEmbedding(model, td.size(), TAG_DIM);
    de = new BinnedDurationEmbedding(model, DURATION_DIM);

    // potential lower layer
    p_d2h1 = model.add_parameters({H1DIM, DURATION_DIM});
    p_y2h1 = model.add_parameters({H1DIM, TAG_DIM});
    p_fwd2h1 = model.add_parameters({H1DIM, SEG_DIM});
    p_rev2h1 = model.add_parameters({H1DIM, SEG_DIM});
    p_h1b = model.add_parameters({H1DIM});

    // potential upper layer
    p_h12h2 = model.add_parameters({H2DIM, H1DIM});
    p_h2b = model.add_parameters({H2DIM});

    // potential output layer
    p_h22o = model.add_parameters({1, H2DIM});
    p_ob = model.add_parameters({1});

    // context aware
    p_c_start = model.add_parameters({XCRIBE_DIM});
    p_c_end = model.add_parameters({XCRIBE_DIM});
    p_cf2h1 = model.add_parameters({H1DIM, XCRIBE_DIM});
    p_ce2h1 = model.add_parameters({H1DIM, XCRIBE_DIM});
  }

  // Adapted the ConstructInput to apply the SegmentRNNs to different inputs
  vector<Expression> ConstructInput(const vector<int>& x,
                                  ComputationGraph& cg){
  xe->new_graph(cg);
  vector<Expression> xins(x.size());
  for (int i = 0; i < x.size(); ++i)
    xins[i] = xe->embed(x[i]);
    return xins;
  }

  unordered_map<PKey,Expression> ConstructSegmentMap(vector<Expression>& xins,
                                                     ComputationGraph& cg,
                                                     int max_seg_len = 0)
  {
    unordered_map<PKey, Expression> p_map;
    int len = xins.size();

    // context aware
    Expression c_start = parameter(cg, p_c_start);
    Expression c_end = parameter(cg, p_c_end);
    Expression cf2h1 = parameter(cg, p_cf2h1);
    Expression ce2h1 = parameter(cg, p_ce2h1);

    ye->new_graph(cg);
    de->new_graph(cg);
    Expression d2h1 = parameter(cg, p_d2h1);
    Expression y2h1 = parameter(cg, p_y2h1);
    Expression fwd2h1 = parameter(cg, p_fwd2h1);
    Expression rev2h1 = parameter(cg, p_rev2h1);
    Expression h1b = parameter(cg, p_h1b);
    Expression h12h2 = parameter(cg, p_h12h2);
    Expression h2b = parameter(cg, p_h2b);
    Expression h22o = parameter(cg, p_h22o);
    Expression ob = parameter(cg, p_ob);

    vector<Expression> c = bt.transcribe(cg, xins);
    seb.construct_chart(cg, c, max_seg_len);

    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      for (unsigned tag = 0; tag < td.size(); ++tag) {
        Expression y = ye->embed(tag);
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          PKey key_tuple(i,j,tag);
          auto seg_embedding_ij = seb(i, j-1); // pair<expr, expr>
          Expression d = de->embed(j - i);
          // To be context aware, add c_{i-1} and c_{j}
          Expression cf_embeding = (i == 0) ? c_start : c[i-1];
          Expression ce_embeding = (j == len) ? c_end : c[j];

          Expression h1 = rectify(affine_transform({h1b, d2h1, d, y2h1, y, fwd2h1, seg_embedding_ij.first, rev2h1, seg_embedding_ij.second, cf2h1, cf_embeding, ce2h1, ce_embeding}));
          Expression h2;
          if(use_dropout){
            h2 = dropout(tanh(affine_transform({h2b, h12h2, h1})), dropout_rate);
          }else{
            h2 = tanh(affine_transform({h2b, h12h2, h1}));
          }
          Expression p = affine_transform({ob, h22o, h2});
          p_map[key_tuple] = p;
        }
      }
    }
    return p_map;

  }

  // return Expression of total loss
  Expression SupervisedCRFLoss(vector<Expression>& xins,
                            const vector<pair<int,int>>& yz,  // .first = y, .second = duration (z)
                            ComputationGraph& cg,
                            int max_seg_len = 0) {
    int len = xins.size();
    unordered_map<PKey, Expression> p_map = ConstructSegmentMap(xins, cg, max_seg_len);
    // is_ref[i][j][y] returns true iff, in the reference, span(i,j) is labeled y
    vector<vector<vector<bool>>> is_ref(len, vector<vector<bool>>(len+1, vector<bool>(td.size(), false)));
    unsigned cur = 0;
    for (unsigned ri = 0; ri < yz.size(); ++ri) {
      assert(cur < len);
      int y = yz[ri].first;
      int dur = yz[ri].second;
      if (max_seg_len && dur > max_seg_len) {
        cerr << "max_seg_len=" << max_seg_len << " but reference duration is " << dur << endl;
        abort();
      }
      unsigned j = cur + dur;
      assert(j <= len);
      is_ref[cur][j][y] = true;
      // cerr << "Span[" << cur << "," << j << ")=" << td.convert(y) << endl;
      cur = j;
    }
    assert(cur == len);

    vector<Expression> fwd(len+1);  // fwd trellis for model
    vector<Expression> ref_fwd(len+1); // fwd trellis for reference
    vector<Expression> f, fr;

    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      // fwd[j] is the total unnoramlized probability for all segmentations / labels
      // ending (and including) the symbol at position j-1
      f.clear(); // f stores all additive contributions to item [j]
      fr.clear();
      for (unsigned tag = 0; tag < td.size(); ++tag) {
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          bool matches_ref = is_ref[i][j][tag];
          PKey key_tuple(i,j,tag);
          assert(p_map.find (key_tuple) != p_map.end());
          Expression p = p_map.at(key_tuple);
          if (i == 0) { // fwd[0] is the path up and including -1, so it's the empty set, i.e., its probability is 1
            f.push_back(p);
            if (matches_ref) fr.push_back(p);
          } else {
            f.push_back(p + fwd[i]);
            if (matches_ref) fr.push_back(p + ref_fwd[i]);
          }
        }
      }
      // cerr << "info\t" << j << "\tsize: " << f.size() << endl;
      // for(auto& e : f){
      //   cerr << "node\t" << as_scalar(cg.get_value(e.i)) << endl;
      // }
      // cerr << "size of vector f at j = " << j << " equals " << f.size() << endl;
      // cerr << "info end" << endl;
      fwd[j] = logsumexp(f);

      // cerr << "shock me\t" << as_scalar(cg.get_value((fwd[j]))) << endl;

      if (fr.size()) ref_fwd[j] = logsumexp(fr);
    }
    //return log(fwd.back()) - log(ref_fwd.back());
    return fwd.back() - ref_fwd.back();
  }

  Expression PartiallySupervisedCRFLoss(vector<Expression>& xins,
                            const vector<pair<int,int>>& yz,  // .first = y, .second = duration (z)
                            ComputationGraph& cg,
                            int max_seg_len = 0) {

    // I keep the z here although I am not going to use it
    int len = xins.size();
    unordered_map<PKey, Expression> p_map = ConstructSegmentMap(xins, cg, max_seg_len);

    vector<Expression> fwd(len+1);  // fwd trellis for model
    vector<Expression> f;

    vector<vector<Expression>> fwd_iq(len+1);
    for (int ind = 0; ind < fwd_iq.size(); ++ind){
      // if yz.size() == 2, then we have end state q = 2 means already consume 2 labels
      // therefore here, should resize to yz.size() + 1 so that fwd_iq[len][2] is the end state
      fwd_iq[ind].resize(yz.size()+1);
    }

    vector<vector<Expression>> fq;
    // cg.forward();
    // cerr << "here" << endl;
    // cerr << len << endl;

    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      // cg.forward();
      // cerr << "j: " << j << endl;
      // fwd[j] is the total unnoramlized probability for all segmentations / labels
      // ending (and including) the symbol at position j-1
      f.clear(); // f stores all additive contributions to item [j]
      fq.clear();
      fq.resize(yz.size() + 1);
      for(auto e : fq){
        e.clear();
      }

      for (unsigned tag = 0; tag < td.size(); ++tag) {
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          PKey key_tuple(i,j,tag);
          assert(p_map.find (key_tuple) != p_map.end());
          Expression p = p_map.at(key_tuple);

          if (i == 0) { // fwd[0] is the path up and including -1, so it's the empty set, i.e., its probability is 1
            f.push_back(p);
          } else {
            f.push_back(p + fwd[i]);
          }

          for (int q = 0; (q < yz.size()+1 && q <= j) ; ++q){
            if (i == 0){
              if(tag == yz[q-1].first && q == 1){
                // q can be 0 now and must be 0 now
                // cerr << "j: " << j << " i: " << i << " q: " << q << endl;
                fq[q].push_back(p);
                // cg.forward();
              }
            } else {
              if(tag == yz[q-1].first && q > 1 && q-1 <= i ) {
                // if q == 0, it's not possible because i > 0, i already consume at least one lable
                // how to default value 0
                // note fwd_iq[max_seg_len+1] can't only consume 1 tag
                // basically, (i-1)/max_seg_len < q-1
                if ((i-1)/max_seg_len < q-1){
                  fq[q].push_back(p + fwd_iq[i][q-1]);
                }
                // else this item got cut
              }
            }
          }
        }
      }
      fwd[j] = logsumexp(f);
      for (int q = 0; q < yz.size()+1 && q <= j; ++q){
        if(fq[q].size() > 0){
          fwd_iq[j][q] = logsumexp(fq[q]);
        }
      }
    }

    return fwd.back() - fwd_iq.back().back();
  }

  Expression SumParts(int len,
                      const vector<pair<int,int>>& yz,
                      unordered_map<PKey, Expression>& p_map,
                      ComputationGraph& cg,
                      int max_seg_len = 0
                      ) {
    vector<Expression> useful_parts;
    unsigned cur = 0;
    for (unsigned ri = 0; ri < yz.size(); ++ri) {
      assert(cur < len);
      int y = yz[ri].first;
      int dur = yz[ri].second;
      if (max_seg_len && dur > max_seg_len) {
        cerr << "max_seg_len=" << max_seg_len << " but reference duration is " << dur << endl;
        abort();
      }
      unsigned j = cur + dur;
      assert(j <= len);

      PKey key_tuple(cur,j,y);
      assert(p_map.find (key_tuple) != p_map.end());
      Expression p = p_map.at(key_tuple);
      useful_parts.push_back(p);
      cur = j;
    }
    assert(cur == len);
    assert(useful_parts.size() > 0);
    Expression s = sum(useful_parts);
    return s;
  }

    // return Expression of total loss
  Expression SupervisedHingeLoss(vector<Expression>& xins,
                            const vector<pair<int,int>>& yz,  // .first = y, .second = duration (z)
                            ComputationGraph& cg,
                            int max_seg_len = 0) {
    int len = xins.size();
    unordered_map<PKey, Expression> p_map = ConstructSegmentMap(xins, cg, max_seg_len);
    vector<Expression> tempvec;
    Expression gold_score = SumParts(len, yz, p_map, cg, max_seg_len);
    tempvec.push_back(gold_score);
    vector<pair<int,int>> yz_pred;
    PureDecode(len, p_map, cg, yz_pred, max_seg_len);
    Expression predict_score = SumParts(len, yz_pred, p_map, cg, max_seg_len);
    tempvec.push_back(predict_score);
    Expression loss = hinge(concatenate(tempvec), (unsigned)0, (float)10);
    return loss;
  }

  void PureDecode(int len,
                  unordered_map<PKey,Expression>& p_map,
                  ComputationGraph& cg,
                  vector<pair<int, int>> &yz_pred,
                  int max_seg_len = 0){
    yz_pred.clear();
    // vector<Expression> fwd(len+1);  // fwd trellis for model
    // vector<Expression> f;
    vector<float> v_fwd(len+1);
    vector<float> v_f;

    vector<tuple<int, int, unsigned>> ijt; // ijt stores positions where we get the max (i.e. i, j and tag)
    vector<tuple<int, int, unsigned>> it; // it stores positions where we get the max ending at j
    it.push_back(make_tuple(0,0,0)); // push one to make the index consistent, now it[j] means j rather than j+1
    // careful: in the other algorithms spans are [i,j], here they are [i,j)
    for (int j = 1; j <= len; ++j) {
      // fwd[j] is the total unnoramlized probability for all segmentations / labels
      // ending (and including) the symbol at position j-1
      // f.clear(); // f stores all additive contributions to item [j]
      v_f.clear();
      ijt.clear();
      for (unsigned tag = 0; tag < td.size(); ++tag) {
        const int i_start = max_seg_len ? max(0, j - max_seg_len) : 0;
        for (int i = i_start; i < j; ++i) {  // i is the starting position
          PKey key_tuple(i,j,tag);
          Expression p = p_map.at(key_tuple);
          float v_p = as_scalar(cg.get_value(p.i));
          if (i == 0) { // fwd[0] is the path up and including -1, so it's the empty set, i.e., its probability is 1
            // f.push_back(p);
            v_f.push_back(v_p);
          } else {
            // f.push_back(p + fwd[i]);
            v_f.push_back(v_p + v_fwd[i]);
          }
          ijt.push_back(make_tuple(i, j, tag));
        }
      }
      // cerr << "size of vector f at j = " << j << " equals " << f.size() << endl;
      unsigned max_ind = 0;
      // auto max_val = as_scalar(cg.get_value(f[0].i));
      auto max_val = v_f[0];
      // for(auto ind = 1; ind < f.size(); ++ind){
      for(auto ind = 1; ind < v_f.size(); ++ind){
        // auto val = as_scalar(cg.get_value(f[ind].i));
        auto val = v_f[ind];
        // cerr << val << endl;
        if (max_val < val){
          max_val = val;
          max_ind = ind;
        }
      }
      // fwd[j] = f[max_ind];
      v_fwd[j] = v_f[max_ind];
      it.push_back(ijt[max_ind]);
      // cerr << "max value = \t" << as_scalar(cg.get_value(fwd[j])) << endl;
    }
    // for(auto j = 1; j <= len; ++j) {
    //   cerr << j << "\t" << std::get<0>(it[j]) << "\t" << std::get<1>(it[j]) << "\t" << std::get<2>(it[j]) << endl;
    // }
    auto cur_j = len;
    vector<tuple<int, int, unsigned>> pred;
    while(cur_j > 0){
      auto cur_i = std::get<0>(it[cur_j]);
      pred.push_back(make_tuple(cur_i, cur_j, std::get<2>(it[cur_j])));
      yz_pred.push_back(make_pair(std::get<2>(it[cur_j]), cur_j - cur_i));
      cur_j = cur_i;
    }
    std::reverse(pred.begin(),pred.end());
    std::reverse(yz_pred.begin(),yz_pred.end());
  }

  void ViterbiDecode(vector<Expression>& xins,
                           const vector<pair<int,int>>& yz_gold,  // .first = y, .second = duration (z)
                           ComputationGraph& cg,
                           vector<pair<int, int>> &yz_pred,
                           int max_seg_len = 0) {

    int len = xins.size();
    unordered_map<PKey, Expression> p_map = ConstructSegmentMap(xins, cg, max_seg_len);
    PureDecode(len, p_map, cg, yz_pred, max_seg_len);
    return;
  }

  Parameter p_d2h1, p_y2h1, p_fwd2h1, p_rev2h1, p_h1b;
  Parameter p_h12h2, p_h2b;
  Parameter p_h22o, p_ob;
  // used in the context awareness
  Parameter p_c_start, p_c_end;
  Parameter p_cf2h1, p_ce2h1;
};

// a a 0 1 a ||| O:1 O:1 N:2 O:1
pair<vector<int>,vector<pair<int,int>>> ParseTrainingInstance(const std::string& line, dynet::Dict& d, dynet::Dict& td, bool test_only = false) {
  std::istringstream in(line);
  std::string word;
  std::string sep = "|||";
  vector<int> x;
  vector<pair<int,int>> yz;
  while(1) {
    in >> word;
    if (!test_only){
      if (word == sep) break;
    }else{
      if (word == sep) break;
      if (!in) break;
    }
    x.push_back(d.convert(word));
  }
  if(!test_only){
    while(1) {
      in >> word;
      if (!in) break;
      size_t p = word.rfind(':');
      if (p == string::npos || p == 0 || p == (word.size() - 1)) {
        cerr << "mal-formed label: " << word << endl;
        abort();
      }
      int y = td.convert(word.substr(0, p));
      int z = atoi(word.substr(p+1).c_str());
      if (z > DATA_MAX_SEG_LEN){
        DATA_MAX_SEG_LEN = z;
      }
      yz.push_back(make_pair(y,z));
    }
  }
  return make_pair(x, yz);
}

bool inline check_max_seg(const vector<pair<int,int>>& yz, int max_seg_len = 0){
  if (max_seg_len == 0) return true;
  for (unsigned ri = 0; ri < yz.size(); ++ri) {
    int dur = yz[ri].second;
    if (max_seg_len && dur > max_seg_len) {
      cerr << "SKIP: max_seg_len=" << max_seg_len << " but reference duration is " << dur << endl;
      return false;
    }
  }
  return true;
}

double evaluate(vector<vector<pair<int,int>>>& yz_preds,
                vector<vector<pair<int,int>>>& yz_golds,
                dynet::Dict& d,
                dynet::Dict& td){
  assert(yz_preds.size() == yz_golds.size());
  int p_correct = 0;
  int r_correct = 0;
  int p_w_t_correct = 0;
  int r_w_t_correct = 0;
  int p_total = 0;
  int r_total = 0;
  int p_w_t_total = 0;
  int r_w_t_total = 0;
  int tag_o = ner_tagging ? td.convert("O") : -1;
  for (unsigned int i = 0; i < yz_preds.size(); i++){
    // for sentence i
    std::set<pair<int,int>> gold;
    std::set<pair<int,int>> pred;
    std::set<tuple<int,int,int>> gold_w_tag;
    std::set<tuple<int,int,int>> pred_w_tag;
    vector<pair<int,int>>& yz_pred = yz_preds[i];
    vector<pair<int,int>>& yz_gold = yz_golds[i];
    int cur = 0;
    for(auto e : yz_pred){

      if (ner_tagging && (e.first == tag_o) ){

      }else{
        pred.insert(make_pair(cur, cur + e.second));
        pred_w_tag.insert(make_tuple(e.first, cur, cur + e.second));
      }
      cur += e.second;
    }
    cur = 0;
    for(auto e : yz_gold){
      if (ner_tagging && (e.first == tag_o) ){

      }else{
        gold.insert(make_pair(cur, cur + e.second));
        gold_w_tag.insert(make_tuple(e.first, cur, cur + e.second));
      }
      cur += e.second;
    }

    for (auto e : pred){
      if(gold.find(e) != gold.end()){
        p_correct++;
      }
      p_total++;
    }
    for (auto e : gold){
      if(pred.find(e) != pred.end()){
        r_correct++;
      }
      r_total++;
    }

    for (auto e : pred_w_tag){
      if(gold_w_tag.find(e) != gold_w_tag.end()){
        p_w_t_correct++;
      }
      p_w_t_total++;
    }
    for (auto e : gold_w_tag){
      if(pred_w_tag.find(e) != pred_w_tag.end()){
        r_w_t_correct++;
      }
      r_w_t_total++;
    }
  }
  double p = (double)(p_correct) / (double)(p_total);
  double r = (double)(r_correct) / (double)(r_total);
  double f = 2.0 * ((p * r) / (p + r));
  cerr << "seg: p: " << p << "\tr: " << r << "\tf: " << f << endl;

  p = (double)(p_w_t_correct) / (double)(p_w_t_total);
  r = (double)(r_w_t_correct) / (double)(r_w_t_total);
  f = 2.0 * ((p * r) / (p + r));
  cerr << "tag: p: " << p << "\tr: " << r << "\tf: " << f << endl;
  return f;

}

void test_only(SegmentalRNN<LSTMBuilder>& segrnn,
          vector<pair<vector<int>,vector<pair<int,int>>>>& test_set,
          int max_seg_len = 0)
{
  for (auto& sent : test_set) {
    ComputationGraph cg;
    vector<pair<int, int>> yz_pred;
    vector<Expression> xins = segrnn.ConstructInput(sent.first, cg);
    segrnn.ViterbiDecode(xins, sent.second, cg, yz_pred, max_seg_len);
    unsigned int i;
    for(i = 0; i < yz_pred.size()-1; ++i){
      auto pred = yz_pred[i];
      cout << segrnn.td.convert(pred.first) << ":" << pred.second;
      if (i != yz_pred.size()-1) {
        cout << " ";
      }
    }
    cout << endl;
  }
}

void read_file(string file_path,
                     dynet::Dict& d,
                     dynet::Dict& td,
                     vector<pair<vector<int>,vector<pair<int,int>>>>& read_set,
                     bool test_only = false)
{
  read_set.clear();
  string line;
  cerr << "Reading data from " << file_path << "...\n";
  {
    ifstream in(file_path);
    assert(in);
    while(getline(in, line)) {
      read_set.push_back(ParseTrainingInstance(line, d, td, test_only));
    }
  }
  cerr << "Reading data from " << file_path << " finished \n";
}

void save_models(string model_file_prefix,
                    dynet::Dict& d,
                    dynet::Dict& td,
                    ParameterCollection& model){
  cerr << "saving models..." << endl;

  const string f_name = model_file_prefix + ".params";
  ofstream out(f_name);
  boost::archive::text_oarchive oa(out);
  oa << model;
  out.close();

  const string f_d_name = model_file_prefix + ".dict";
  ofstream out_d(f_d_name);
  boost::archive::text_oarchive oa_d(out_d);
  oa_d << d;
  out_d.close();

  const string f_td_name = model_file_prefix + ".tdict";
  ofstream out_td(f_td_name);
  boost::archive::text_oarchive oa_td(out_td);
  oa_td << td;
  out_td.close();
  cerr << "saving models finished" << endl;
}

void load_models(string model_file_prefix,
                 ParameterCollection& model){
  cerr << "loading models..." << endl;

  string fname = model_file_prefix + ".params";
  ifstream in(fname);
  boost::archive::text_iarchive ia(in);
  ia >> model;
  in.close();

  cerr << "loading models finished" << endl;
}

void load_dicts(string model_file_prefix,
                 dynet::Dict& d,
                 dynet::Dict& td)
{
  cerr << "loading dicts..." << endl;
  string f_d_name = model_file_prefix + ".dict";
  ifstream in_d(f_d_name);
  boost::archive::text_iarchive ia_d(in_d);
  ia_d >> d;
  in_d.close();

  string f_td_name = model_file_prefix + ".tdict";
  ifstream in_td(f_td_name);
  boost::archive::text_iarchive ia_td(in_td);
  ia_td >> td;
  in_td.close();
  cerr << "loading dicts finished" << endl;
}

unsigned int edit_distance(const std::string& s1, const std::string& s2)
{
  const std::size_t len1 = s1.size(), len2 = s2.size();
  std::vector<std::vector<unsigned int>> d(len1 + 1, std::vector<unsigned int>(len2 + 1));

  d[0][0] = 0;
  for(unsigned int i = 1; i <= len1; ++i) d[i][0] = i;
    for(unsigned int i = 1; i <= len2; ++i) d[0][i] = i;
      for(unsigned int i = 1; i <= len1; ++i)
        for(unsigned int j = 1; j <= len2; ++j)
          d[i][j] = std::min({ d[i - 1][j] + 1, d[i][j - 1] + 1, d[i - 1][j - 1] + (s1[i - 1] == s2[j - 1] ? 0 : 1) });
        return d[len1][len2];
}

double evaluate_partial(vector<vector<pair<int,int>>>& yz_preds,
                       vector<vector<pair<int,int>>>& yz_golds,
                       dynet::Dict& d,
                       dynet::Dict& td){
  assert(yz_preds.size() == yz_golds.size());

  int total_length_gold = 0;
  int total_length_pred = 0;
  int total_editing_distance = 0;

  for (int i = 0; i < yz_preds.size(); i++){
    // for sentence i
    vector<pair<int,int>>& yz_pred = yz_preds[i];
    vector<pair<int,int>>& yz_gold = yz_golds[i];
    std::string pred_s = "";
    for (auto e : yz_pred){
      pred_s.append(td.convert(e.first));
      total_length_pred++;
    }
    std::string gold_s = "";
    for (auto e : yz_gold){
      gold_s.append(td.convert(e.first));
      total_length_gold++;
    }
    auto dis = edit_distance(pred_s, gold_s);
    total_editing_distance += dis;

  }
  cerr << "total_editing_distance: " << total_editing_distance << "\ttotal_length_gold: " << total_length_gold << "\ttotal_length_pred: " << total_length_pred << endl;
  double score = (((float) total_editing_distance) * 2.0)/((float)(total_length_gold + total_length_pred));
  cerr << "score: " << score << endl;
  return score;
}

double predict_and_evaluate(SegmentalRNN<LSTMBuilder>& segrnn,
                            const vector<pair<vector<int>,vector<pair<int,int>>>>& input_set,
                            bool partial,
                            int max_seg_len,
                            string set_name = "DEV"
                            ){
  vector<vector<pair<int,int>>> yz_preds;
  vector<vector<pair<int,int>>> yz_golds;
  for (auto& sent : input_set) {
    ComputationGraph cg;
    vector<pair<int, int>> yz_pred;
    vector<Expression> xins = segrnn.ConstructInput(sent.first, cg);
    segrnn.ViterbiDecode(xins, sent.second, cg, yz_pred, max_seg_len);
    yz_golds.push_back(sent.second);
    yz_preds.push_back(yz_pred);
  }

  double f1 = evaluate(yz_preds, yz_golds, segrnn.d, segrnn.td);
  double f2 = evaluate_partial(yz_preds, yz_golds, segrnn.d, segrnn.td);

  double f = partial ? f2 : f1;
  cerr << set_name << endl;
  return f;
}

int main(int argc, char** argv) {
  dynet::initialize(argc, argv);
  int test_max_seg_len;
  int max_consider_sentence_len;
  unsigned dev_every_i_reports;
  po::options_description desc("Allowed options");
  desc.add_options()
      ("help", "produce help message")
      ("train", po::bool_switch()->default_value(false), "the training mode")
      ("hinge", po::bool_switch()->default_value(false), "training use the hinge loss")
      ("partial", po::bool_switch()->default_value(false), "the partially supervised training mode, where we don't have z")
      ("train_max_seg_len", po::value<int>(), "in the partially supervised training mode, the max len we consider")
      ("load_original_model", po::bool_switch()->default_value(false), "continuing the training by loading the model, only valid during training")
      ("max_consider_sentence_len", po::value<int>(&max_consider_sentence_len)->default_value(1000))
      ("dropout_rate", po::value<float>(), "dropout rate, also indicts using dropout during training")
      ("evaluate_test", po::bool_switch()->default_value(false), "evaluate test set every training iteration")
      ("test", po::bool_switch()->default_value(false), "the test mode")
      ("test_max_seg_len", po::value<int>(&test_max_seg_len)->default_value(30))
      ("dev_every_i_reports", po::value<unsigned>(&dev_every_i_reports)->default_value(1000))
      ("train_file", po::value<string>(), "path of the train file")
      ("dev_file", po::value<string>(), "path of the dev file")
      ("test_file", po::value<string>(), "path of the test file")
      ("model_file_prefix", po::value<string>(), "prefix path of the model files (and dictionaries)")
      ("upe", po::value<string>(), "use pre-trained word embeding")
  ;

  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);

  if (vm.count("help")) {
      cerr << desc << "\n";
      return 1;
  }

  if(vm["train"].as<bool>()){
    if (vm.count("upe")){
      cerr << "using pre-trained embeding from " << vm["upe"].as<string>() << endl;
      use_pretrained_embeding = true;
      pretrained_embeding = vm["upe"].as<string>();
    }else{
      use_pretrained_embeding = false;
      cerr << "not using pre-trained embeding" << endl;
    }

    if (vm.count("dropout_rate")){
      use_dropout = true;
      dropout_rate = vm["dropout_rate"].as<float>();
      cerr << "using dropout training, dropout rate: " << dropout_rate << endl;
    }else{
      use_dropout = false;
    }

    // create two dictionaries
    dynet::Dict d;
    dynet::Dict td;
    vector<pair<vector<int>,vector<pair<int,int>>>> training, dev, test;
    read_file(vm["train_file"].as<string>(), d, td, training);

    d.freeze();  // no new word types allowed
    td.freeze(); // no new tag types allowed
    d.set_unk("<UNK>"); // set UNK to allow the unseen character in the dev and test set

    read_file(vm["dev_file"].as<string>(), d, td, dev);
    if (vm["evaluate_test"].as<bool>()){
      read_file(vm["test_file"].as<string>(), d, td, test);
    }

    ParameterCollection model;
    // auto trainer = new SimpleSGDTrainer(model);
    std::unique_ptr<Trainer> trainer(new AdamTrainer(model, 0.0005, 0.01, 0.9999, 1e-8));
    int max_seg_len = DATA_MAX_SEG_LEN + 1;
    if(vm.count("train_max_seg_len")){
      max_seg_len = vm["train_max_seg_len"].as<int>();
    }

    cerr << "set max_seg_len = " << max_seg_len << endl;

    SegmentalRNN<LSTMBuilder> segrnn(model, d, td);

    if(vm["load_original_model"].as<bool>()){
      load_models(vm["model_file_prefix"].as<string>(), model);
    }

    double f_best = 0;
    unsigned report_every_i = 10;

    unsigned si = training.size();
    vector<unsigned> order(training.size());
    for (unsigned i = 0; i < order.size(); ++i) order[i] = i;
    int report = 0;
    unsigned lines = 0;
    while(1) {
      Timer iteration("completed in");
      double loss = 0;
      unsigned ttags = 0;
      double correct = 0;
      for (unsigned i = 0; i < report_every_i; ++i) {
        if (si == training.size()) {
          si = 0;
          cerr << "**SHUFFLE\n";
          shuffle(order.begin(), order.end(), *rndeng);
        }

        // build graph for this instance
        ComputationGraph cg;
        auto& sent = training[order[si]];
        ++si;
        if(check_max_seg(sent.second, max_seg_len)){
          if (sent.first.size() > max_consider_sentence_len){
            cerr << "skip a sentence because its length " << sent.first.size() << " > " << max_consider_sentence_len << endl;
            continue;
          }
          vector<Expression> xins = segrnn.ConstructInput(sent.first, cg);
          Expression loss_expr;
          if(vm["partial"].as<bool>()){
              if(vm["hinge"].as<bool>()){
                cerr << "Hingle Loss for partially supervised setting is not avaiable at this moment." << endl;
                abort();
              }else{
                loss_expr = segrnn.PartiallySupervisedCRFLoss(xins, sent.second, cg, max_seg_len);
              }
          }else{
            if(vm["hinge"].as<bool>()){
              loss_expr = segrnn.SupervisedHingeLoss(xins, sent.second, cg, max_seg_len);
            }else{
              loss_expr = segrnn.SupervisedCRFLoss(xins, sent.second, cg, max_seg_len);
            }
          }
          ttags += sent.second.size();
          loss += as_scalar(cg.forward(loss_expr));
          cg.backward(loss_expr);
          trainer->update(1.0);
        }
        ++lines;
      }
      trainer->status();
      cerr << " E = " << (loss / ttags) << " ppl=" << exp(loss / ttags) << " (acc=" << (correct / ttags) << ") ";
      report++;
      if (report % dev_every_i_reports == 0) {
        double f = predict_and_evaluate(segrnn, dev, vm["partial"].as<bool>(), max_seg_len);
        if (f > f_best) {
          f_best = f;
          save_models(vm["model_file_prefix"].as<string>(), d, td, model);
        }
        if (vm["evaluate_test"].as<bool>()){
          predict_and_evaluate(segrnn, test, vm["partial"].as<bool>(), max_seg_len, "TEST");
        }
      }
    }
  }else if(vm["test"].as<bool>()){
    use_pretrained_embeding = false;
    use_dropout = false;
    ParameterCollection model;
    dynet::Dict d;
    dynet::Dict td;
    load_dicts(vm["model_file_prefix"].as<string>(), d, td);
    SegmentalRNN<LSTMBuilder> segrnn(model, d, td);
    load_models(vm["model_file_prefix"].as<string>(), model);
    vector<pair<vector<int>,vector<pair<int,int>>>> test;
    read_file(vm["test_file"].as<string>(), d, td, test, true);

    test_only(segrnn, test);
  }

}


#include "oaz/games/connect_four.hpp"

#include <algorithm>
#include <boost/multi_array.hpp>
#include <string>

#include "gmock/gmock.h"
#include "gtest/gtest.h"

using namespace oaz::games;
using namespace testing;

TEST(InstantiationTest, Default) { ConnectFour game; }

TEST(GetAvailableMoves, Default) {
  ConnectFour game;
  std::vector<size_t> available_moves;
  game.GetAvailableMoves(&available_moves);
  ASSERT_THAT(available_moves, ElementsAre(0, 1, 2, 3, 4, 5, 6));
}

TEST(PlayTest, VerticalVictory) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("0103040");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(1, game.GetScore());
}

TEST(PlayTest, VerticalVictoryPlayer2) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("60103040");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(-1, game.GetScore());
}

TEST(PlayTest, HorizontalVictory) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("051625364");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(1, game.GetScore());
}

TEST(PlayTest, HorizontalVictoryPlayer2) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("6051625364");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(-1, game.GetScore());
}

TEST(PlayTest, FirstDiagonalVictory) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("12234334544");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(1, game.GetScore());
}

TEST(PlayTest, FirstDiagonalVictoryPlayer2) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("612234334544");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(-1, game.GetScore());
}

TEST(PlayTest, SecondDiagonalVictory) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("54432332122");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(1, game.GetScore());
}

TEST(PlayTest, SecondDiagonalVictoryPlayer2) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("654432332122");
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(-1, game.GetScore());
}

TEST(PlayTest, TieTest) {
  ConnectFour game;
  ASSERT_FALSE(game.IsFinished());
  game.PlayFromString("021302130213465640514455662233001144552636");
  std::vector<size_t> available_moves;
  game.GetAvailableMoves(&available_moves);
  ASSERT_THAT(available_moves, ElementsAre());
  ASSERT_TRUE(game.IsFinished());
  ASSERT_EQ(0, game.GetScore());
}

TEST(Clone, Default) {
  ConnectFour game;
  game.PlayFromString("021302130213465640514455662233001144552636");
  std::unique_ptr<Game> clone = game.Clone();
  ConnectFour* clone_ptr = dynamic_cast<ConnectFour*>(clone.get());
  ASSERT_TRUE(game == *clone_ptr);
}

TEST(GetCurrentPlayer, Default) {
  ConnectFour game;
  ASSERT_EQ(game.GetCurrentPlayer(), 0);
  game.PlayMove(0);
  ASSERT_EQ(game.GetCurrentPlayer(), 1);
}

TEST(ClassMethods, Default) {
  ConnectFour game;
  Game* game_ptr = &game;
  ASSERT_EQ(game_ptr->ClassMethods().GetMaxNumberOfMoves(), 7);
  ASSERT_EQ(game_ptr->ClassMethods().GetBoardShape()[0], 6);
  ASSERT_EQ(game_ptr->ClassMethods().GetBoardShape()[1], 7);
  ASSERT_EQ(game_ptr->ClassMethods().GetBoardShape()[2], 2);
}

TEST(WriteStateToTensorMemory, Default) {
  ConnectFour game;
  game.PlayMove(0);
  game.PlayMove(5);
  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());
  for (size_t i = 0; i != 6; ++i)
    for (size_t j = 0; j != 7; ++j)
      if (i == 0 && j == 0)
        ASSERT_EQ(tensor[i][j][0], 1.);
      else
        ASSERT_EQ(tensor[i][j][0], 0.);
  for (size_t i = 0; i != 6; ++i)
    for (size_t j = 0; j != 7; ++j)
      if (i == 0 && j == 5)
        ASSERT_EQ(tensor[i][j][1], 1.);
      else
        ASSERT_EQ(tensor[i][j][1], 0.);
}

TEST(WriteCanonicalStateToTensorMemory, Default) {
  ConnectFour game;
  game.PlayMove(0);
  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteCanonicalStateToTensorMemory(tensor.origin());
  for (size_t i = 0; i != 6; ++i)
    for (size_t j = 0; j != 7; ++j) {
      if (i == 0 && j == 0)
        ASSERT_EQ(tensor[i][j][1], 1.);
      else
        ASSERT_EQ(tensor[i][j][1], 0.);
      ASSERT_EQ(tensor[i][j][0], 0.);
    }
}

TEST(InitialiseFromState, Default) {
  ConnectFour game;
  boost::multi_array<float, 3> test_board(boost::extents[6][7][2]);
  test_board[0][0][0] = 1.0;
  test_board[1][0][1] = 1.0;
  game.InitialiseFromState(test_board.origin());

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());
  for (size_t i = 0; i != 6; ++i)
    for (size_t j = 0; j != 7; ++j) {
      if (i == 0 && j == 0)
        ASSERT_EQ(tensor[i][j][0], 1.0);
      else if (i == 1 && j == 0)
        ASSERT_EQ(tensor[i][j][1], 1.);
      else {
        ASSERT_EQ(tensor[i][j][0], 0.);
        ASSERT_EQ(tensor[i][j][1], 0.);
      }
    }
}

TEST(InitialiseFromState, CheckBoardCopy) {
  ConnectFour game;
  game.PlayFromString("0510055");

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());

  ConnectFour game2;
  game2.InitialiseFromState(tensor.origin());
  ASSERT_TRUE(game == game2);
}

TEST(InitialiseFromState, CheckBoardCopy2) {
  ConnectFour game;
  game.PlayFromString("0510055");

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());

  ConnectFour game2;
  game2.InitialiseFromState(tensor.origin());

  game.PlayMove(5);

  ASSERT_FALSE(game == game2);
}

TEST(InitialiseFromState, CheckGameInProgressInitialisation) {
  ConnectFour game;
  game.PlayFromString("0510055");

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());

  ConnectFour game2;
  game2.PlayFromString("0123313");
  game2.InitialiseFromState(tensor.origin());
  ASSERT_TRUE(game == game2);
}

TEST(InitialiseFromCanonicalState, Default) {
  ConnectFour game;
  boost::multi_array<float, 3> test_board(boost::extents[6][7][2]);
  test_board[0][0][0] = 1.0;
  test_board[1][0][1] = 1.0;
  test_board[2][0][0] = 1.0;
  game.InitialiseFromCanonicalState(test_board.origin());

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteCanonicalStateToTensorMemory(tensor.origin());
  for (size_t i = 0; i != 6; ++i)
    for (size_t j = 0; j != 7; ++j) {
      if (i == 0 && j == 0)
        ASSERT_EQ(tensor[i][j][0], 1.0f);
      else if (i == 1 && j == 0)
        ASSERT_EQ(tensor[i][j][1], 1.0f);
      else if (i == 2 && j == 0)
        ASSERT_EQ(tensor[i][j][0], 1.0f);
      else {
        ASSERT_EQ(tensor[i][j][0], 0.0f);
        ASSERT_EQ(tensor[i][j][1], 0.0f);
      }
    }
}

TEST(InitialiseFromCanonicalState, CheckBoardCopy) {
  ConnectFour game;
  game.PlayFromString("0510055");

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteCanonicalStateToTensorMemory(tensor.origin());

  ConnectFour game2;
  game2.InitialiseFromCanonicalState(tensor.origin());

  ASSERT_TRUE(game == game2);
}

TEST(InitialiseFromCanonicalState, CheckBoardCopy2) {
  ConnectFour game;
  game.PlayFromString("0510055");

  boost::multi_array<float, 3> tensor(boost::extents[6][7][2]);
  game.WriteStateToTensorMemory(tensor.origin());

  ConnectFour game2;
  game2.InitialiseFromCanonicalState(tensor.origin());

  ASSERT_FALSE(game == game2);
}

TEST(GameMap, Instantiation) {
  ConnectFour game;
  std::unique_ptr<oaz::games::Game::GameMap> game_map(
      game.ClassMethods().CreateGameMap());
}

TEST(GameMap, GetNotInDB) {
  ConnectFour game;
  std::unique_ptr<oaz::games::Game::GameMap> game_map(
      game.ClassMethods().CreateGameMap());
  size_t index;
  ASSERT_FALSE(game_map->Get(game, &index));
}

TEST(GameMap, GetInDB) {
  ConnectFour game;
  std::unique_ptr<oaz::games::Game::GameMap> game_map(
      game.ClassMethods().CreateGameMap());
  game_map->Insert(game, 1ll);

  size_t index = 0;
  ASSERT_TRUE(game_map->Get(game, &index));
  ASSERT_EQ(index, 1);
}

TEST(GameMap, InsertAlreadyInDB) {
  ConnectFour game;
  std::unique_ptr<oaz::games::Game::GameMap> game_map(
      game.ClassMethods().CreateGameMap());

  ASSERT_EQ(game_map->GetSize(), 0);
  game_map->Insert(game, 1ll);
  ASSERT_EQ(game_map->GetSize(), 1);
  game_map->Insert(game, 1ll);
  ASSERT_EQ(game_map->GetSize(), 1);
}


// ************************
// ************************
// Edited by CCJ:
// evaldisp_lib.cpp,
// which can be called
// by python code;
// ************************
// ************************

// evaluate disparity map
// simple version for SDK
// supports upsampling of disp map if GT has higher resolution

// DS 7/2/2014
// 10/14/2014 changed computation of average error
// 1/27/2015 added clipping of valid (non-INF) disparities to [0 .. maxdisp]
//    in fairness to those methods that do not utilize the given disparity range
//    (maxdisp is specified at disp resolution, NOT GT resolution)

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <stdio.h>
#include <iostream>
#include <stdlib.h>
#include <math.h>
#include <vector>
#include "imageLib/imageLib.h"
#include <stdint.h> /*for int32_t*/

//added by CCJ for Python + CPP coding;
#include <boost/python.hpp>
#include "boost/python/extract.hpp"
//#include "boost/python/numeric.hpp"
/* Updated by CCJ for Boost version 1.65:
 * Boost 1.65 removes boost/python/numeric.hpp
 */
// BOOST_LIB_VERSION and BOOST_VERSION are defined in this header file;
#include <boost/version.hpp>
#if BOOST_VERSION >= 106300 // >= 1.63.0
   #include "boost/python/numpy.hpp"
   namespace np = boost::python::numpy;
#else
   #include "boost/python/numeric.hpp"
	 namespace np = boost::python::numeric;
#endif

#include <numpy/ndarrayobject.h>
using namespace boost::python;

// see Fastest way to check if a file exist using standard C++/C++11/C?
// at https://stackoverflow.com/questions/12774207/fastest-way-to-check-if-a-file-exist-using-standard-c-c11-c;
#include <sys/stat.h>
inline bool exists_file_test(const std::string & name){
	struct stat buffer;
	return (stat (name.c_str(), &buffer) == 0);
}

const int verbose = 0;

/* The `occlusion mask` for the left image is given as a file "mask0nocc.png":
 * - Pixels without ground truth have the color (0, 0, 0).
 * - Pixels which are only observed by the left image have the color (128, 128, 128).
 * - Pixels which are observed by both images have the color (255, 255, 255).
 * - For the "non-occluded" evaluation, the evaluation is limited to the pixels observed by both images.
*/

void evaldisp(
		CFloatImage disp,
		CFloatImage gtdisp,
		CByteImage mask,
		float badthresh,
		int maxdisp,
		int rounddisp,
		float * statistics
		){

    CShape sh = gtdisp.Shape();
    CShape sh2 = disp.Shape();
    CShape msh = mask.Shape();

    int width = sh.width, height = sh.height;
    int width2 = sh2.width, height2 = sh2.height;
    int scale = width / width2;
		//std::cout << "scale = " << scale << "\n";

    if ((!(scale == 1 || scale == 2 || scale == 4))
				|| (scale * width2 != width)
				|| (scale * height2 != height)){
			printf("   disp size = %4d x %4d\n", width2, height2);
			printf("GT disp size = %4d x %4d\n", width,  height);
			throw CError("GT disp size must be exactly 1, 2, or 4 * disp size");
    }

    int usemask = (msh.width > 0 && msh.height > 0);
    if (!usemask)
			throw CError("No mask image's been read yet!\n");
    if (usemask && (msh != sh))
			throw CError("mask image must have same size as GT!\n");

		// all region;
    int n_all = 0;
    int bad_all = 0;
    int invalid_all = 0;
    float err_all = 0;
		// non-occluded region;
    int invalid_noc = 0;
		int n_noc = 0;
    int bad_noc = 0;
    float err_noc = 0;
		// updated for mae and rmse error metric, on 2019/08/31
    float err2_all = 0; // square
    float err2_noc = 0;

    for (int y = 0; y < height; y++) {
			for (int x = 0; x < width; x++) {
				float gt = gtdisp.Pixel(x, y, 0);
				if (gt == INFINITY) // unknown
					continue;
				float d = scale * disp.Pixel(x / scale, y / scale, 0);
	      int valid = (d != INFINITY);
				if (valid) {
					float maxd = scale * maxdisp; // max disp range
		      d = __max(0, __min(maxd, d)); // clip disps to max disp range
				}
	      if (valid && rounddisp){
					d = round(d);
				}
				float err = fabs(d - gt);

	      if (mask.Pixel(x,y,0) == 0){ // no ground truth;
					continue;
				}

				if (mask.Pixel(x, y, 0) == 255){ // non-occluded region;
					n_noc++;
					if (valid) {
						err_noc += err;
						err2_noc += err*err;

						if (err > badthresh) {
							bad_noc++;
						}
					}
					else {// invalid (i.e. hole in sparse disp map)
						invalid_noc++;
					}
				}

				n_all++;
				if(valid){
					err_all += err;
					err2_all += err * err;

					if (err > badthresh){
						bad_all++;
					}
				}
				else {// invalid (i.e. hole in sparse disp map)
					invalid_all++;
				}
			}/*end of width x*/
		}/*end of height y*/

		float badpercent_all =  (float)bad_all / (float)n_all;
    float avgErr_all = err_all / (float)(n_all - invalid_all); // CHANGED 10/14/2014 -- was: serr / n
		float badpercent_noc = (float) bad_noc / (float) n_noc;
		float avgErr_noc = err_noc /(float)(n_noc - invalid_noc);
		float rmse_all = sqrt(err2_all /(float)(n_all - invalid_all));
		float rmse_noc = sqrt(err2_noc /(float)(n_noc - invalid_noc));
#if 0
    float invalidpercent_all =  (float)invalid_all/(float)n_all;
    float totalbadpercent_all =  (float)(bad_all+invalid_all)/(float)n_all;
		printf("n_all = %d, n_noc = %d\n", n_all, n_noc);
    printf("%4.1f%(N-rate-all)  %6.2f%(bad-%2.1f-all)  %6.2f%(invalid-all)  %6.2f%(totalBad-all)  %6.2f(avgErr-all)\n",
				100.0*n_all/(width * height), 100.0*badpercent_all, badthresh, 100.0*invalidpercent_all,
				100.0*totalbadpercent_all, avgErr_all);
    float invalidpercent_noc =  (float)invalid_noc / (float)n_noc;
    float totalbadpercent_noc =  (float)(bad_noc+invalid_noc)/ (float)n_noc;
    printf("%4.1f%(N-rate-noc)  %6.2f%(bad-%2.1f-noc)  %6.2f%(invalid-noc)  %6.2f%(totalBad-noc)  %6.2f(avgErr-noc)\n",
				100.0*n_noc/(width * height), 100.0*badpercent_noc, badthresh, 100.0*invalidpercent_noc,
				100.0*totalbadpercent_noc, avgErr_noc);
#endif

		statistics[0] = badpercent_all;
		statistics[1] = avgErr_all;

		statistics[2] = badpercent_noc;
		statistics[3] = avgErr_noc;

		statistics[4] = rmse_all;
		statistics[5] = rmse_noc;
}

// C++ code
typedef std::vector<std::string> StringList;
class IMG_NAMES{
	private:
		StringList myvec;
	public:
		// constructor
		IMG_NAMES(){
			this -> myvec = StringList(0);
		}
	void set_img_names(boost::python::list & ns){
		for(int i = 0; i < len(ns); ++i){
			this -> myvec.push_back(boost::python::extract<std::string>(ns[i]));
		}
	}

	void show_img_names(){
		std:: cout << "Images :\n";
		for(auto i: this-> myvec){
			std::cout << i << ", ";
		}
		std:: cout << "\n";
	}

	StringList get_img_names(){
		return this -> myvec;
	}
};

PyObject * eval_disp_mbv3_1_img(
		const string & dispname,                // your predicted disparity results dir;
	  const string & disp_gt_trainingF_dir, 	// ground truth disparity dir to trainingF;
		float badthresh,                        // e.g., == 1.0 for Half resolution;
		const int &  ndisp,                      // max_disp in groud truth disparity;
		const int & rounddisp,
    const std::string & img_name           // 1 image to evaluate;
		){

		npy_intp * dims = new npy_intp[1];
		dims[0] = 4;
		PyObject * errs_A = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
	  //std::cout << "new PyObject errs_A, img_N = " << img_N << std::endl;
		float * err_result = static_cast<float*>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(errs_A)));

		//std::string dispname = disp_dir + img_name + "/" + disp_name; //"disp0.pfm"
		std::string gtdispname = disp_gt_trainingF_dir + img_name + "/disp0GT.pfm";
		// updated for middlebury additional data, due to the lack of mask0nocc.png file;
		std::string maskname = disp_gt_trainingF_dir + img_name + "/mask0nocc.png";
		bool readMask = true;
		if (!exists_file_test(maskname)){
			std::cout << "No mask0noc.png exists!\n";
			readMask = false;
		}

		//std::cout << "maxdisp = " << ndisp << ", "<< dispname << ", " << gtdispname << ", " << maskname << "\n";

		CFloatImage disp, gtdisp, gtdisp1;
		ReadImageVerb(disp, dispname.c_str(), verbose);
		ReadImageVerb(gtdisp, gtdispname.c_str(), verbose);
		CByteImage mask;
		if (readMask){
			//std::cout << "read mask0noc.png!\n";
			ReadImageVerb(mask, maskname.c_str(), verbose);
		}

		float statistics[4]= {-1.0, -1.0, -1.0, -1.0 };
		evaldisp(disp, gtdisp, mask, badthresh, ndisp, rounddisp, statistics);
		err_result[0] =  statistics[0];
		err_result[1] =  statistics[1];
		err_result[2] =  statistics[2];
		err_result[3] =  statistics[3];

		// release the memory
		disp.DeAllocate();
		gtdisp.DeAllocate();
		gtdisp1.DeAllocate();
		if (maskname.c_str())
			mask.DeAllocate();
		//exit(0);
		return errs_A;
}

PyObject * eval_disp_mbv3(
		const string & disp_dir, // your prediction disparity results dir;
	  const string & disp_gt_trainingF_dir, 	// ground truth disparity dir to trainingF;
		float badthresh, // e.g., == 1.0 for Half resolution;
		const string & disp_name, // e.g., == "_disp0" or "_disp0_post", or "_rf_disp0PKLS", and so on;
		PyObject * ndisps, // max_disp for each groud truth disparity
		PyObject * rounddisps, // disparity is int or float, specified in calib.txt files;
    boost::python::list img_names // img_names to evaluate;
		){

		/*vector<std::string> mb_v3_imgs_1 = {
			"Adirondack",
			"ArtL",
			"Jadeplant",
			"Motorcycle",
			"MotorcycleE",
			"Piano",
			"PianoL",
			"Pipes",
			"Playroom",
			"Playtable",
		 	"PlaytableP",
			"Recycle",
			"Shelves",
			"Teddy",
			"Vintage",
		};
		*/
    IMG_NAMES mb_v3;
		mb_v3.set_img_names(img_names);
		//mb_v3.show_img_names();
		vector<std::string> mb_v3_imgs = mb_v3.get_img_names();

    PyArrayObject* ndispsA = reinterpret_cast<PyArrayObject*>(ndisps);
    PyArrayObject* rounddispsA = reinterpret_cast<PyArrayObject*>(rounddisps);
		//int32_t * p_ndisps = reinterpret_cast<int32_t*>(PyArray_DATA(ndispsA));
		int * p_ndisps = reinterpret_cast<int*>(PyArray_DATA(ndispsA));
		int * p_rounddisps = reinterpret_cast<int*>(PyArray_DATA(rounddispsA));

		const int img_N = mb_v3_imgs.size();
		npy_intp * dims = new npy_intp[1];
		dims[0] = (4+2)*img_N;
		PyObject * errs_A = PyArray_SimpleNew(1, dims, NPY_FLOAT32);
	  //std::cout << "new PyObject errs_A, img_N = " << img_N << std::endl;
		float * err_result = static_cast<float*>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(errs_A)));
		for (int i = 0; i < img_N; i++){
			std::string dispname = disp_dir + mb_v3_imgs[i] + "/" + disp_name + ".pfm";
			// check if a file exists or not;
			// > see: https://stackoverflow.com/questions/12774207/fastest-way-to-check-if-a-file-exist-using-standard-c-c11-c;
		  if (!exists_file_test(dispname)){
				//std::cout << "Not exist : " << dispname << ", changed it to : ";

				//dispname = disp_dir + mb_v3_imgs[i] + "_" + disp_name;
				//Updated: delete "_" here;
				dispname = disp_dir + mb_v3_imgs[i] + disp_name + ".pfm";
				//std::cout << dispname << " ";
				if (!exists_file_test(dispname)){
					std::cout << "loading " << dispname << " ... but failed!\n";
				}
			}

			std::string gtdispname = disp_gt_trainingF_dir + mb_v3_imgs[i] + "/disp0GT.pfm";
			std::string maskname = disp_gt_trainingF_dir + mb_v3_imgs[i] + "/mask0nocc.png";
			bool readMask = true;
			// updated for middlebury additional data, due to the lack of mask0nocc.png file;
			if (!exists_file_test(maskname)){
			  std::cout << "No mask0noc.png exists!\n";
				readMask = false;
			}
	    int maxdisp = p_ndisps[i];
	    int rounddisp = p_rounddisps[i];
      //std::cout << "maxdisp = " << maxdisp << ", rounddisp = " << rounddisp << ", "<< dispname << ", " << gtdispname << ", " << maskname << "\n";

	    CFloatImage disp, gtdisp, gtdisp1;
	    ReadImageVerb(disp, dispname.c_str(), verbose);
	    ReadImageVerb(gtdisp, gtdispname.c_str(), verbose);
	    CByteImage mask;
	    if (readMask){
				//std::cout << "read mask0noc.png!\n";
				ReadImageVerb(mask, maskname.c_str(), verbose);
			}

      //float statistics[4]= {-1.0, -1.0, -1.0, -1.0 };
			// updated for rmse metric on 2019/08/31;
	    float statistics[4 + 2]= {-1.0, -1.0, -1.0, -1.0, -1.0, -1.0 };
			evaldisp(disp, gtdisp, mask, badthresh, maxdisp, rounddisp, statistics);
#if 0
			printf("Processing: %10s, %6.2f%(bad-%2.1f-all), %6.2f(avgErr-all), %6.2f%(bad-%2.1f-noc), %6.2f(avgErr-nov)\n",
					mb_v3_imgs[i].c_str(), statistics[0]*100.0, badthresh,
					statistics[1], statistics[2]*100.0, badthresh, statistics[3]);
#endif
			err_result[6*i]     =  statistics[0];
			err_result[6*i + 1] =  statistics[1];
			err_result[6*i + 2] =  statistics[2];
			err_result[6*i + 3] =  statistics[3];
			err_result[6*i + 4] =  statistics[4];
			err_result[6*i + 5] =  statistics[5];

			// release the memory
			disp.DeAllocate();
			gtdisp.DeAllocate();
			gtdisp1.DeAllocate();
	    if (maskname.c_str())
				mask.DeAllocate();

			//exit(0);

		}/*end of each image*/
		return errs_A;
}

BOOST_PYTHON_MODULE(libevaldisp_mbv3){
	/* for Boost <= version 1.63*/
	//numeric::array::set_module_and_type("numpy", "ndarray");
  /* for Boost > version 1.63*/
	//np::initialize();
#if BOOST_VERSION >= 106300 // >= 1.63.0
	np::initialize();
#else
	np::array::set_module_and_type("numpy", "ndarray");
#endif
	def("evaluate_mbv3", eval_disp_mbv3);
	def("evaluate_mbv3_1_img", eval_disp_mbv3_1_img);
	/* Error: return-statement with a value, in function returning 'void' [-fpermissive]
	 *        #define NUMPY_IMPORT_ARRAY_RETVAL NULL
	 * > See solution: https://github.com/numpy/numpy/issues/10486
	 * > 1) Solution1: Okay, so the issue occurs only on py2 + py3c (the initialization function is nonstandard, and has py3 semantics). Solution appears to be to use `import_array1()`;
	 * > 2) Solution2: Or call _import_array(), which allows you more control than just `return`;
	 */
  //import_array(); // work well for Python2.7;
	import_array1(); // work well for python3.7;
	//_import_array(); // work well for python3.7;
}


/*
  Copyright (c) 2012,2013,2014 Matthew H. Reilly (kb1vc)
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are
  met:

  Redistributions of source code must retain the above copyright
  notice, this list of conditions and the following disclaimer.
  Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in
  the documentation and/or other materials provided with the
  distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "ReSampler.hxx"
#include <string.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <boost/format.hpp>

static void bindump(char * fn, std::complex<float> * buf, unsigned int num_elts, int append = 0) __attribute__ ((unused));
static void bindump(char * fn, std::complex<float> * buf, unsigned int num_elts, int append)
{
  FILE * of;
  if(append != 0) {
    of = fopen(fn, "a");
  }
  else {
    of = fopen(fn, "w");
  }
  unsigned int i;
  for(i = 0; i < num_elts; i++) {
    fprintf(of, "%d %g %g\n", i + (append * num_elts), buf[i].real(), buf[i].imag());
  }
  fclose(of);
}

static void bindump_2(char * fn, std::complex<float> * buf, unsigned int num_elts, unsigned int num_cols) __attribute__ ((unused));
static void bindump_2(char * fn, std::complex<float> * buf, unsigned int num_elts, unsigned int num_cols)
{
  FILE * of;
  of = fopen(fn, "w");
  unsigned int i, j;
  for(i = 0; i < num_elts; i++) {
    fprintf(of, "%d ", i);
    for(j = 0; j < num_cols; j++) {
      fprintf(of, "%g %g ", buf[i * num_cols + j].real(), buf[i * num_cols + j].imag());
    }
    fprintf(of, "\n");
  }
  fclose(of);
}

static unsigned int ipow(unsigned int a, unsigned int b)
{
  unsigned int ret = 1;
  unsigned int i ;

  for(i = 0; i < b; i++) {
    ret *= a;
  }

  return ret;
}

SoDa::ReSampler::ReSampler(unsigned int interpolate_ratio,
			   unsigned int decimate_ratio,
			   unsigned int _inlen,
			   unsigned int _filter_len)
{
  // save the parameters
  inlen = _inlen;
  iM = interpolate_ratio;
  dN = decimate_ratio;
  filter_len = _filter_len;
  Q = filter_len;
  M = _inlen;

  // give preferences to convenient powers of two.
  // but let's search for the nearest solution that is
  // a multiple of 2^a * 3^b * 5^c where b and c are in the range 0..3
  // and a is in the range 1..16
  unsigned int N_guess, N_best, E_best;
  N_best = 2;
  E_best = 0x80000000;
  unsigned int i;
  for(i = 1; i <= 0xff; i++) {
    unsigned int a = i & 0xf;
    unsigned int b = (i >> 4) & 0x3;
    unsigned int c = (i >> 6) & 0x3;
    N_guess = ipow(2, a) * ipow(3, b) * ipow(5, c);
    if(N_guess >= (M + Q - 1)) {
      unsigned slop = N_guess - (M + Q - 1);
      if(slop < E_best) {
	N_best = N_guess;
	E_best = slop;
      }
    }
  }
  N = N_best;

  // now that we have N, we can back-calculate Q.
  Q = N - M;

  filter_len = Q;

  tail_index = M - (Q - 1);

  // create the buffers.
  // a copy of the input sequence
  inbuf = (std::complex<float> *) fftwf_alloc_complex(N);
  // the transformed input
  in_fft = (std::complex<float> *) fftwf_alloc_complex(N);

  // zero out the whole of the input buffer
  for(i = 0; i < N; i++) inbuf[i] = std::complex<float>(0.0,0.0);

  // the filter bank is really iM interleaved filters
  // but we can't seem to make that work right now, so we're
  // backing off to a simpler way.
  c_filt = new std::complex<float>*[iM];
  for(i = 0; i < iM; i++) {
    c_filt[i] = (std::complex<float> *) fftwf_alloc_complex(N);
  }

  // the upsampled result
  filt_fft = (std::complex<float> *) fftwf_alloc_complex(iM*N);
  // the inverse transform of the upsampled result
  interp_res = (std::complex<float> *) fftwf_alloc_complex(iM*N);

  // the first plan is pretty simple.
  in_fft_plan = fftwf_plan_dft_1d(N,
				  (fftwf_complex *) inbuf,
				  (fftwf_complex *) in_fft,
				  FFTW_FORWARD, FFTW_ESTIMATE);

  // now create the inverting plans
  int n[1];
  n[0] = N;
  mid_ifft_plan = fftwf_plan_many_dft(1, n, iM,
				      (fftwf_complex *) filt_fft, NULL, iM, 1,
				      (fftwf_complex *) interp_res, NULL, iM, 1,
				      FFTW_BACKWARD, FFTW_ESTIMATE);

  // now create the polyphase filter banks.
  CreateFilter(filter_len);

  transform_gain = 1.0 / ((float) N);
}

void SoDa::ReSampler::CreateFilter(unsigned int filter_len)
{
  // filter_len must be odd
  int fN = filter_len;
  if((fN & 1) == 0) fN--;

  unsigned int i, j, k;

  // Use the plan of fischer in "The mkfilter Digital Filter Generation Program" mkshape...
  // This is borrowed quite heavily and with tremendous gratitude from
  // http://www-users.cs.york.ac.uk/~fisher   Thanks, Prof. Fisher...
  float alpha, beta;
  float cutdown = ((float) ((iM > dN) ? iM : dN));
  beta = 1.0 / (2.2 * cutdown); // This is the transition freq -- a little shy of 0.5/cutdown
  alpha = 0.1; // this is a shape factor

  // create a frequency domain image of the filter.
  float f1 = (1.0 - alpha) * beta;
  float f2 = (1.0 + alpha) * beta;
  float tau = 0.5 / alpha;

  // now create the filter banks.
  std::complex<float> * filt_FD = (std::complex<float>*) fftwf_alloc_complex(N);
  std::complex<float> * filt_td = (std::complex<float>*) fftwf_alloc_complex(N);
  std::complex<float> * filt_td2 = (std::complex<float>*) fftwf_alloc_complex(N);

  float f;
  float f_incr = 1.0 / ((float) (N));
  unsigned int hlim = (N);
  float gain_corr = 1.0 / ((float) hlim);
  //  gain_corr = gain_corr * gain_corr;
  for(i = 0, f = 0.0; i <= hlim / 2; i++, f += f_incr) {
    if(f <= f1) {
      filt_FD[i] = std::complex<float>(1.0, 0.0);
    }
    else if (f <= f2) {
      float h = 0.5 * (1.0 + cos(M_PI * tau * (f - f1) / beta));
      filt_FD[i] = std::complex<float>(h, 0.0);
    }
    else {
      filt_FD[i] = std::complex<float>(0.0, 0.0);
    }

    filt_FD[i] *= tau * gain_corr;

    if(i != 0) filt_FD[hlim - i] = filt_FD[i];
  }

  // then ifft it to time domain
  fftwf_plan filt_ifft_plan = fftwf_plan_dft_1d(hlim,
						(fftwf_complex *) filt_FD,
						(fftwf_complex *) filt_td,
						FFTW_BACKWARD, FFTW_ESTIMATE);
  fftwf_execute(filt_ifft_plan);

  float sum = 0.0;
  for(i = 0; i < hlim; i++) sum += filt_td[i].real();

  float sgain = 1.0 / sum;
  for(i = 0; i < hlim; i++) filt_td[i] *= sgain;

  fftwf_destroy_plan(filt_ifft_plan);

  // now truncate the filter
  // We've got an image that is symmetric around 0.
  for(i = 0; i < N; i++) {
    filt_td2[i] = std::complex<float>(0.0,0.0);
  }

  // shift it up to be symmetric around filt_len / 2
  for(i = 0, j = N - (filter_len/2); j < N; i++, j++) {
    filt_td2[i] = filt_td[j];
  }
  for(j = 0; j < filter_len / 2; j++, i++) {
    filt_td2[i] = filt_td[j];
  }

  //  bindump("filt_td2.dat", filt_td2, N);

  // and fft back to frequency domain for each of the iM filters.
  for(i = 0; i < iM; i++) {
    // grab the slice of the filter that we need to transform.

    for(j = i, k = 0; j < filter_len; k++, j += iM) {
      filt_td[k] = filt_td2[j];
    }
    for(; k < N; k++) filt_td[k] = std::complex<float>(0.0,0.0);
    fftwf_plan filt_fft_plan = fftwf_plan_dft_1d(N,
						 (fftwf_complex *) filt_td,
						 (fftwf_complex *) c_filt[i],
						 FFTW_FORWARD, FFTW_ESTIMATE);
    fftwf_execute(filt_fft_plan);
    fftwf_destroy_plan(filt_fft_plan);
  }

  // now free up buffers
  fftwf_free(filt_td2);
  fftwf_free(filt_td);
  fftwf_free(filt_FD);
}

unsigned int SoDa::ReSampler::apply(std::complex<float> * in,
				    std::complex<float> * out,
				    float gain)
{
  unsigned int i, j;
  // copy the input buffer
  memcpy(&(inbuf[Q-1]), in, sizeof(std::complex<float>) * M);

  // transform the input
  fftwf_execute(in_fft_plan); // , (fftwf_complex *) in, (fftwf_complex *) in_fft);

  // now multiply by the M filters
  for(i = 0; i < N; i++) {
    for(j = 0; j < iM; j++) {
      filt_fft[j + i*iM] = in_fft[i] * c_filt[j][i];
    }
  }
  // do M inverse transforms on the filter result ...
  // to produce M interleaved time domain vectors.
  fftwf_execute(mid_ifft_plan);

  // save the last bits of the input buffer to the Q-1 side of the FFT input vector
  // Do this now incase inbuf and outbuf are the same buffers.  (we may
  // over-write the outbuf with the downsample at the bottom.... )
  memcpy(inbuf, &(in[tail_index]), sizeof(std::complex<float>) * (Q-1));

  // now build the output buffer.

  // now downsample by a factor of dN.
  int idx = 0;
  unsigned int mctr = 0;
  unsigned int out_lim = (M * iM) / dN;
  std::complex<float> * ir_clean = &(interp_res[(Q-1)]);
  // std::cerr << " iM, dN, M, N, out_lim = "
  // 	    << iM << " " << dN << " " << M << " " << N << " " << out_lim << std::endl;
  for(j = 0; j < out_lim; j++) {
    out[j] = ir_clean[mctr + idx * iM] * gain * transform_gain * ((float) iM);
    mctr += dN;
    while(mctr >= iM) {
      mctr -= iM;
      idx += 1;
    }
  }
  //  bindump("out.dat", out, j);

  return 0;
}

unsigned int SoDa::ReSampler::apply(float * in,
				    float * out,
				    float gain)
{
  unsigned int i, j;

  for(i = 0; i < M; i++) {
    inbuf[i+(Q-1)] = std::complex<float>(in[i], 0.0);
  }

  // transform the input
  fftwf_execute(in_fft_plan); // , (fftwf_complex *) in, (fftwf_complex *) in_fft);

  // now multiply by the M filters
  for(i = 0; i < N; i++) {
    for(j = 0; j < iM; j++) {
      filt_fft[j + i*iM] = in_fft[i] * c_filt[j][i];
    }
  }
  // do M inverse transforms on the filter result ...
  // to produce M interleaved time domain vectors.
  fftwf_execute(mid_ifft_plan);

  // save the last bits of the input buffer to the Q-1 side of the FFT input vector
  // Do this now incase inbuf and outbuf are the same buffers.  (we may
  // over-write the outbuf with the downsample at the bottom.... )
  for(i = 0; i < (Q-1); i++) {
    inbuf[i] = std::complex<float>(in[i + tail_index], 0.0);
  }

  // now build the output buffer.

  // now downsample by a factor of dN.
  int idx = 0;
  unsigned int mctr = 0;
  unsigned int out_lim = (M * iM) / dN;
  std::complex<float> * ir_clean = &(interp_res[(Q-1)]);
  for(j = 0; j < out_lim; j++) {
    out[j] = ir_clean[mctr + idx * iM].real() * gain * transform_gain * ((float) iM);
    mctr += dN;
    while(mctr >= iM) {
      mctr -= iM;
      idx += 1;
    }
  }

  return 0;
}