#include <opencv2/opencv.hpp>
#include <string>
#include <chrono>
#include <Eigen/Core>
#include <Eigen/Dense>

using namespace std;
using namespace cv;

string file_1 = "../LK1.png";
string file_2 = "../LK2.png";

class OpticalFlowTracker {
public:
    OpticalFlowTracker(
            const Mat &img1_,
            const Mat &img2_,
            const vector<KeyPoint> &kp1_,
            vector<KeyPoint> &kp2_,
            vector<bool> &success_,
            bool inverse_ = true, bool has_initial_ = false) :
            img1(img1_), img2(img2_), kp1(kp1_), kp2(kp2_), success(success_), inverse(inverse_),
            has_initial(has_initial_) {}

    void calculateOpticalFlow(const Range &range);

private:
    const Mat &img1;
    const Mat &img2;
    const vector<KeyPoint> &kp1;
    vector<KeyPoint> &kp2;
    vector<bool> &success;
    bool inverse = true;
    bool has_initial = false;
};

void OpticalFlowSingleLevel(
        const Mat &img1,
        const Mat &img2,
        const vector<KeyPoint> &kp1,
        vector<KeyPoint> &kp2,
        vector<bool> &success,
        bool inverse = false,
        bool has_initial_guess = false
);

void OpticalFlowMultiLevel(
        const Mat &img1,
        const Mat &img2,
        const vector<KeyPoint> &kp1,
        vector<KeyPoint> &kp2,
        vector<bool> &success,
        bool inverse = false
);

inline float GetPixelValue(const cv::Mat &img, float x, float y) {
    // boundary check
    if (x < 0) x = 0;
    if (y < 0) y = 0;
    if (x >= img.cols - 1) x = img.cols - 2;
    if (y >= img.rows - 1) y = img.rows - 2;

    float xx = x - floor(x);
    float yy = y - floor(y);
    int x_a1 = std::min(img.cols - 1, int(x) + 1);
    int y_a1 = std::min(img.rows - 1, int(y) + 1);

    return (1 - xx) * (1 - yy) * img.at<uchar>(y, x)
           + xx * (1 - yy) * img.at<uchar>(y, x_a1)
           + (1 - xx) * yy * img.at<uchar>(y_a1, x)
           + xx * yy * img.at<uchar>(y_a1, x_a1);
}

int main(int argc, char **argv) {

    // images, note they are CV_8UC1, not CV_8UC3
    Mat img1 = imread(file_1, 0);
    Mat img2 = imread(file_2, 0);

    // key points, using GFTT here.
    vector<KeyPoint> kp1;
    Ptr<GFTTDetector> detector = GFTTDetector::create(500, 0.01, 20); // maximum 500 keypoints
    detector->detect(img1, kp1);

    // now lets track these key points in the second image
    // first use single level LK in the validation picture
    vector<KeyPoint> kp2_single;
    vector<bool> success_single;
    OpticalFlowSingleLevel(img1, img2, kp1, kp2_single, success_single);

    // then test multi-level LK
    vector<KeyPoint> kp2_multi;
    vector<bool> success_multi;
    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    OpticalFlowMultiLevel(img1, img2, kp1, kp2_multi, success_multi, true);
    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
    auto time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
    cout << "optical flow by gauss-newton: " << time_used.count() << endl;

    // use opencv's flow for validation
    vector<Point2f> pt1, pt2;
    for (auto &kp: kp1) pt1.push_back(kp.pt);
    vector<uchar> status;
    vector<float> error;
    t1 = chrono::steady_clock::now();
    cv::calcOpticalFlowPyrLK(img1, img2, pt1, pt2, status, error);
    t2 = chrono::steady_clock::now();
    time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
    cout << "optical flow by opencv: " << time_used.count() << endl;

    // plot the differences of those functions
    Mat img2_single;
    cv::cvtColor(img2, img2_single, COLOR_GRAY2BGR);
    for (int i = 0; i < kp2_single.size(); i++) {
        if (success_single[i]) {
            cv::circle(img2_single, kp2_single[i].pt, 2, cv::Scalar(0, 250, 0), 2);
            cv::line(img2_single, kp1[i].pt, kp2_single[i].pt, cv::Scalar(0, 250, 0));
        }
    }

    Mat img2_multi;
    cv::cvtColor(img2, img2_multi, COLOR_GRAY2BGR);
    for (int i = 0; i < kp2_multi.size(); i++) {
        if (success_multi[i]) {
            cv::circle(img2_multi, kp2_multi[i].pt, 2, cv::Scalar(0, 250, 0), 2);
            cv::line(img2_multi, kp1[i].pt, kp2_multi[i].pt, cv::Scalar(0, 250, 0));
        }
    }

    Mat img2_CV;
    cv::cvtColor(img2, img2_CV, COLOR_GRAY2BGR);
    for (int i = 0; i < pt2.size(); i++) {
        if (status[i]) {
            cv::circle(img2_CV, pt2[i], 2, cv::Scalar(0, 250, 0), 2);
            cv::line(img2_CV, pt1[i], pt2[i], cv::Scalar(0, 250, 0));
        }
    }

    cv::imshow("tracked single level", img2_single);
    cv::imshow("tracked multi level", img2_multi);
    cv::imshow("tracked by opencv", img2_CV);
    cv::waitKey(0);

    return 0;
}

void OpticalFlowSingleLevel(
        const Mat &img1,
        const Mat &img2,
        const vector<KeyPoint> &kp1,
        vector<KeyPoint> &kp2,
        vector<bool> &success,
        bool inverse, bool has_initial) {
    kp2.resize(kp1.size());
    success.resize(kp1.size());
    OpticalFlowTracker tracker(img1, img2, kp1, kp2, success, inverse, has_initial);
    parallel_for_(Range(0, kp1.size()),
                  std::bind(&OpticalFlowTracker::calculateOpticalFlow, &tracker, placeholders::_1));
}

void OpticalFlowTracker::calculateOpticalFlow(const Range &range) {
    // parameters
    int half_patch_size = 4;
    int iterations = 10;
    for (size_t i = range.start; i < range.end; i++) {
        auto kp = kp1[i];
        double dx = 0, dy = 0; // dx,dy need to be estimated
        if (has_initial) {
            dx = kp2[i].pt.x - kp.pt.x;
            dy = kp2[i].pt.y - kp.pt.y;
        }

        double cost = 0, lastCost = 0;
        bool succ = true; // indicate if this point succeeded

        // Gauss-Newton iterations
        Eigen::Matrix2d H = Eigen::Matrix2d::Zero();    // hessian
        Eigen::Vector2d b = Eigen::Vector2d::Zero();    // bias
        Eigen::Vector2d J;  // jacobian
        for (int iter = 0; iter < iterations; iter++) {
            if (inverse == false) {
                H = Eigen::Matrix2d::Zero();
                b = Eigen::Vector2d::Zero();
            } else {
                // only reset b
                b = Eigen::Vector2d::Zero();
            }

            cost = 0;

            // compute cost and jacobian
            for (int x = -half_patch_size; x < half_patch_size; x++)
                for (int y = -half_patch_size; y < half_patch_size; y++) {
                    double error = GetPixelValue(img1, kp.pt.x + x, kp.pt.y + y) -
                                   GetPixelValue(img2, kp.pt.x + x + dx, kp.pt.y + y + dy);;  // Jacobian
                    if (inverse == false) {
                        J = -1.0 * Eigen::Vector2d(
                                0.5 * (GetPixelValue(img2, kp.pt.x + dx + x + 1, kp.pt.y + dy + y) -
                                       GetPixelValue(img2, kp.pt.x + dx + x - 1, kp.pt.y + dy + y)),
                                0.5 * (GetPixelValue(img2, kp.pt.x + dx + x, kp.pt.y + dy + y + 1) -
                                       GetPixelValue(img2, kp.pt.x + dx + x, kp.pt.y + dy + y - 1))
                        );
                    } else if (iter == 0) {
                        // in inverse mode, J keeps same for all iterations
                        // NOTE this J does not change when dx, dy is updated, so we can store it and only compute error
                        J = -1.0 * Eigen::Vector2d(
                                0.5 * (GetPixelValue(img1, kp.pt.x + x + 1, kp.pt.y + y) -
                                       GetPixelValue(img1, kp.pt.x + x - 1, kp.pt.y + y)),
                                0.5 * (GetPixelValue(img1, kp.pt.x + x, kp.pt.y + y + 1) -
                                       GetPixelValue(img1, kp.pt.x + x, kp.pt.y + y - 1))
                        );
                    }
                    // compute H, b and set cost;
                    b += -error * J;
                    cost += error * error;
                    if (inverse == false || iter == 0) {
                        // also update H
                        H += J * J.transpose();
                    }
                }

            // compute update
            Eigen::Vector2d update = H.ldlt().solve(b);

            if (std::isnan(update[0])) {
                // sometimes occurred when we have a black or white patch and H is irreversible
                cout << "update is nan" << endl;
                succ = false;
                break;
            }

            if (iter > 0 && cost > lastCost) {
                break;
            }

            // update dx, dy
            dx += update[0];
            dy += update[1];
            lastCost = cost;
            succ = true;

            if (update.norm() < 1e-2) {
                // converge
                break;
            }
        }

        success[i] = succ;

        // set kp2
        kp2[i].pt = kp.pt + Point2f(dx, dy);
    }
}

void OpticalFlowMultiLevel(
        const Mat &img1,
        const Mat &img2,
        const vector<KeyPoint> &kp1,
        vector<KeyPoint> &kp2,
        vector<bool> &success,
        bool inverse) {

    // parameters
    int pyramids = 4;
    double pyramid_scale = 0.5;
    double scales[] = {1.0, 0.5, 0.25, 0.125};

    // create pyramids
    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    vector<Mat> pyr1, pyr2; // image pyramids
    for (int i = 0; i < pyramids; i++) {
        if (i == 0) {
            pyr1.push_back(img1);
            pyr2.push_back(img2);
        } else {
            Mat img1_pyr, img2_pyr;
            cv::resize(pyr1[i - 1], img1_pyr,
                       cv::Size(pyr1[i - 1].cols * pyramid_scale, pyr1[i - 1].rows * pyramid_scale));
            cv::resize(pyr2[i - 1], img2_pyr,
                       cv::Size(pyr2[i - 1].cols * pyramid_scale, pyr2[i - 1].rows * pyramid_scale));
            pyr1.push_back(img1_pyr);
            pyr2.push_back(img2_pyr);
        }
    }
    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
    auto time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
    cout << "build pyramid time: " << time_used.count() << endl;

    // coarse-to-fine LK tracking in pyramids
    vector<KeyPoint> kp1_pyr, kp2_pyr;
    for (auto &kp:kp1) {
        auto kp_top = kp;
        kp_top.pt *= scales[pyramids - 1];
        kp1_pyr.push_back(kp_top);
        kp2_pyr.push_back(kp_top);
    }

    for (int level = pyramids - 1; level >= 0; level--) {
        // from coarse to fine
        success.clear();
        t1 = chrono::steady_clock::now();
        OpticalFlowSingleLevel(pyr1[level], pyr2[level], kp1_pyr, kp2_pyr, success, inverse, true);
        t2 = chrono::steady_clock::now();
        auto time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
        cout << "track pyr " << level << " cost time: " << time_used.count() << endl;

        if (level > 0) {
            for (auto &kp: kp1_pyr)
                kp.pt /= pyramid_scale;
            for (auto &kp: kp2_pyr)
                kp.pt /= pyramid_scale;
        }
    }

    for (auto &kp: kp2_pyr)
        kp2.push_back(kp);
}


//
// Add --log_level=message to see the messages!
//
#include <boost/test/unit_test.hpp>

#include <sstream>

#include "core/functions.h"
#include "core/utils.h"
#include "core/random.h"
#include "core/training.h"
#include "layers/convlayer.h"

#include "test_utils.h"
#include "timer.h"
#include "test_layers.h"

using namespace std;
using namespace boost;
using namespace boost::unit_test;
using namespace yann;
using namespace yann::test;

struct ConvolutionalLayerTestFixture
{
  ConvolutionalLayerTestFixture()
  {

  }
  ~ConvolutionalLayerTestFixture()
  {

  }

  inline MatrixSize find_max_pos(const RefConstVector & vv)
  {
    MatrixSize pos = 0;
    vv.maxCoeff(&pos);
    return pos;
  }

  void conv_perf_test(const MatrixSize & size, const MatrixSize & filter_size, const size_t & epochs)
  {
      BOOST_TEST_MESSAGE("*** ConvOp Performance test with"
          << " size=" << size
          << ", filter_size=" << filter_size
          << ", epochs=" << epochs
      );

      Matrix input(size, size);
      Matrix filter(filter_size, filter_size);
      Matrix output(size + filter_size, size + filter_size);
      {
        Timer timer("Random generation");
        unique_ptr<RandomGenerator> gen = RandomGenerator::normal_distribution(0, 1);
        gen->generate(input);
        gen->generate(filter);
      }
      output.resize(ConvolutionalLayer::get_conv_output_rows(size, filter_size),
                    ConvolutionalLayer::get_conv_output_cols(size, filter_size));
      {
        Timer timer("Test plus_conv");
        for(auto ii = epochs; ii > 0; --ii) {
          ConvolutionalLayer::plus_conv(input, filter, output);
        }
      }
      output.resize(ConvolutionalLayer::get_full_conv_output_rows(size, filter_size),
                    ConvolutionalLayer::get_full_conv_output_cols(size, filter_size));
      {
        Timer timer("Test full_conv");
        for(auto ii = epochs; ii > 0; --ii) {
          ConvolutionalLayer::full_conv(input, filter, output);
        }
      }
  }

  void conv_perf_batch_test(const MatrixSize & batch_size, const MatrixSize & size, const MatrixSize & filter_size, const size_t & epochs)
  {
      BOOST_TEST_MESSAGE("*** ConvOp Performance batch test with"
          << " batch_size=" << batch_size
          << ", size=" << size
          << ", filter_size=" << filter_size
          << ", epochs=" << epochs
      );
      VectorBatch input, output;
      Matrix filter(filter_size, filter_size);

      resize_batch(input, batch_size, size * size);
      {
        Timer timer("Random generation");
        unique_ptr<RandomGenerator> gen = RandomGenerator::normal_distribution(0, 1);
        gen->generate(input);
        gen->generate(filter);
      }
      resize_batch(output, batch_size, ConvolutionalLayer::get_conv_output_size(size, size, filter_size));
      {
        Timer timer("Test plus_conv");
        for(auto ii = epochs; ii > 0; --ii) {
          ConvolutionalLayer::plus_conv(input, size, size, filter, output);
        }
      }
      resize_batch(output, batch_size, ConvolutionalLayer::get_full_conv_output_size(size, size, filter_size));
      {
        Timer timer("Test full_conv");
        for(auto ii = epochs; ii > 0; --ii) {
          ConvolutionalLayer::full_conv(input, size, size, filter, output);
        }
      }
  }
}; // struct ConvolutionalLayerTestFixture

BOOST_FIXTURE_TEST_SUITE(ConvolutionalLayerTest, ConvolutionalLayerTestFixture);

BOOST_AUTO_TEST_CASE(IO_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer IO test ...");

  const MatrixSize input_cols = 5;
  const MatrixSize input_rows = 3;
  const MatrixSize filter_size = 2;
  auto one = make_unique<ConvolutionalLayer>(input_cols, input_rows, filter_size);
  one->init(Layer::InitMode_Random, boost::none);

  BOOST_TEST_MESSAGE("ConvolutionalLayer before writing to file: " << "\n" << *one);
  ostringstream oss;
  oss << *one;
  BOOST_CHECK(!oss.fail());

  auto two = make_unique<ConvolutionalLayer>(input_cols, input_rows, filter_size);
  std::istringstream iss(oss.str());
  iss >> *two;
  BOOST_CHECK(!iss.fail());
  BOOST_TEST_MESSAGE("ConvolutionalLayer after loading from file: " << "\n" << *two);

  BOOST_CHECK(one->is_equal(*two, TEST_TOLERANCE));
}

BOOST_AUTO_TEST_CASE(ConvOp_Test)
{
  BOOST_TEST_MESSAGE("*** Convolutional operation test ...");

  const size_t image_size = 4;
  const size_t max_filter_size = 4;
  Matrix filter(max_filter_size, max_filter_size);
  Matrix input(image_size, image_size);
  Matrix expected(image_size, image_size);
  Matrix output(image_size, image_size);

  //////////////////////////////////////////////////////////////////////
  //
  // Test different filter sizes
  //
  input <<
      1,  2,  3,  4,
      5,  6,  7,  8,
      9, 10, 11, 12,
     13, 14, 15, 16;

  // 1x1
  filter.resize(1, 1);
  filter <<
      3;
  expected.resize(4, 4);
  expected <<
      3,  6,  9, 12,
     15, 18, 21, 24,
     27, 30, 33, 36,
     39, 42, 45, 48;
  output.resizeLike(expected);
  {
      // ensure we don't do allocations in eigen
      BlockAllocations block;
      ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x1
  filter.resize(2, 1);
  filter <<
      2,
      3;
  expected.resize(3, 4);
  expected <<
      17, 22, 27, 32,
      37, 42, 47, 52,
      57, 62, 67, 72;
  output.resizeLike(expected);
  {
    // ensure we don't do allocations in eigen
    BlockAllocations block;
    ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x2
  filter.resize(2, 2);
  filter <<
      1, 2,
      3, 4;
  expected.resize(3, 3);
  expected <<
      44,  54,  64,
      84,  94, 104,
     124, 134, 144;
  output.resizeLike(expected);
  {
    // ensure we don't do allocations in eigen
    BlockAllocations block;
    ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x3
  filter.resize(2, 3);
  filter <<
      1, 2, 3,
      4, 5, 6;
  expected.resize(3, 2);
  expected <<
      106, 127,
      190, 211,
      274, 295;
  output.resizeLike(expected);
  {
    // ensure we don't do allocations in eigen
    BlockAllocations block;
    ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 3x3
  filter.resize(3, 3);
  filter <<
      1, 2, 3,
      4, 5, 6,
      7, 8, 9;
  expected.resize(2, 2);
  expected <<
      348, 393,
      528, 573;
  output.resizeLike(expected);
  {
    // ensure we don't do allocations in eigen
    BlockAllocations block;
    ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 4x4
  filter.resize(4, 4);
  filter <<
      1,  2,  3,  4,
      5,  6,  7,  8,
      9, 10, 11, 12,
     13, 14, 15, 16;
  expected.resize(1, 1);
  expected <<
      1496;
  output.resizeLike(expected);
  {
    // ensure we don't do allocations in eigen
    BlockAllocations block;
    ConvolutionalLayer::plus_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

}

BOOST_AUTO_TEST_CASE(ConvOp_Batch_Test)
{
  BOOST_TEST_MESSAGE("*** Convolutional operation batch test ...");

  const size_t max_filter_size = 4;
  const MatrixSize image_size = 4;
  const MatrixSize input_size = image_size * image_size;
  const MatrixSize batch_size = 2;
  Matrix filter(max_filter_size, max_filter_size);
  VectorBatch input, output, expected;

  //////////////////////////////////////////////////////////////////////
  //
  // Test different filter sizes
  //
  resize_batch(input, batch_size, input_size);
  input << 1, 0, 1, 0,
           0, 1, 0, 0,
           1, 0, 1, 0,
           0, 0, 0, 0,
           ///////////
           0, 0, 0, 0,
           0, 1, 0, 0,
           0, 0, 1, 0,
           0, 0, 0, 0;

  // 1x1
  filter.resize(1, 1);
  filter << 1;
  resize_batch(expected, batch_size, 4 * 4);
  expected << 1, 0, 1, 0,
              0, 1, 0, 0,
              1, 0, 1, 0,
              0, 0, 0, 0,
              ///////////
              0, 0, 0, 0,
              0, 1, 0, 0,
              0, 0, 1, 0,
              0, 0, 0, 0;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x1
  filter.resize(2, 1);
  filter << 1,
            1;
  resize_batch(expected, batch_size, 3 * 4);
  expected << 1, 1, 1, 0,
              1, 1, 1, 0,
              1, 0, 1, 0,
              ///////////
              0, 1, 0, 0,
              0, 1, 1, 0,
              0, 0, 1, 0;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x2
  filter.resize(2, 2);
  filter << 1, 1,
            1, 0;
  resize_batch(expected, batch_size, 3 * 3);
  expected << 1, 2, 1,
              2, 1, 1,
              1, 1, 1,
              ////////
              0, 1, 0,
              1, 1, 1,
              0, 1, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 3x2
  filter.resize(3, 2);
  filter << 1, 1,
            1, 0,
            1, 0;
  resize_batch(expected, batch_size, 2 * 3);
  expected << 2, 2, 2,
              2, 1, 1,
              ////////
              0, 1, 1,
              1, 1, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 3x3
  filter.resize(3, 3);
  filter << 1, 1, 1,
            1, 0, 0,
            1, 0, 1;
  resize_batch(expected, batch_size, 2 * 2);
  expected << 4, 2,
              2, 1,
              /////
              1, 1,
              1, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 4x4
  filter.resize(4, 4);
  filter << 1, 1, 1, 1,
            1, 1, 0, 0,
            1, 0, 1, 0,
            1, 0, 0, 1;
  resize_batch(expected, batch_size, 1 * 1);
  expected << 5,
              //
              2;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::plus_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));
}

BOOST_AUTO_TEST_CASE(FullConvOp_Test)
{
  BOOST_TEST_MESSAGE("*** Full convolutional operation test ...");

  const size_t image_size = 3;
  const size_t max_filter_size = 2;
  Matrix filter(max_filter_size, max_filter_size);
  Matrix input(image_size, image_size);
  Matrix expected(image_size + max_filter_size, image_size + max_filter_size);
  Matrix output(image_size + max_filter_size, image_size + max_filter_size);

  //////////////////////////////////////////////////////////////////////
  //
  // Test different filter sizes
  //
  input.resize(image_size, image_size);
  input <<
      1, 2, 3,
      4, 5, 6,
      7, 8, 9;

  // 1x1
  filter.resize(1, 1);
  filter <<
      3;
  expected.resize(3, 3);
  expected <<
      3,  6,  9,
     12, 15, 18,
     21, 24, 27;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x2
  filter.resize(2, 2);
  filter <<
      1, 2,
      3, 4;
  expected.resize(4, 4);
  expected <<
      4, 11, 18, 9,
     18, 37, 47, 21,
     36, 67, 77, 33,
     14, 23, 26, 9;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x3
  filter.resize(2, 3);
  filter <<
      1, 2, 3,
      4, 5, 6;
  expected.resize(4, 5);
  expected <<
      6,  17,  32, 23, 12,
     27,  58,  91, 58, 27,
     54, 106, 154, 94, 42,
     21,  38,  50, 26, 9;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 3x3
  filter.resize(3, 3);
  filter <<
      1, 2, 3,
      4, 5, 6,
      7, 8, 9;
  expected.resize(5, 5);
  expected <<
      9,  26,  50,  38,  21,
     42,  94, 154, 106,  54,
     90, 186, 285, 186,  90,
     54, 106, 154,  94,  42,
     21,  38,  50,  26,   9;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));
}

BOOST_AUTO_TEST_CASE(FullConvOp_Batch_Test)
{
  BOOST_TEST_MESSAGE("*** Full convolutional operation batch test ...");

  const size_t max_filter_size = 4;
  const MatrixSize image_size = 3;
  const MatrixSize input_size = image_size * image_size;
  const MatrixSize batch_size = 3;
  Matrix filter(max_filter_size, max_filter_size);
  VectorBatch input, output, expected;

  //////////////////////////////////////////////////////////////////////
  //
  // Test different filter sizes
  //
  resize_batch(input, batch_size, input_size);
  input << 1, 0, 1,
           0, 1, 0,
           1, 0, 0,
           ///////
           0, 0, 0,
           0, 0, 0,
           0, 0, 0,
           ///////
           1, 1, 1,
           1, 1, 1,
           1, 1, 1;

  // 1x1
  filter.resize(1, 1);
  filter << 1;
  resize_batch(expected, batch_size, 3 * 3);
  expected << 1, 0, 1,
              0, 1, 0,
              1, 0, 0,
              ///////
              0, 0, 0,
              0, 0, 0,
              0, 0, 0,
              ///////
              1, 1, 1,
              1, 1, 1,
              1, 1, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x2
  filter.resize(2, 2);
  filter << 1, 1,
            1, 0;
  resize_batch(expected, batch_size, 4 * 4);
  expected << 0, 1, 0, 1,
              1, 1, 2, 1,
              0, 2, 1, 0,
              1, 1, 0, 0,
              //////////
              0, 0, 0, 0,
              0, 0, 0, 0,
              0, 0, 0, 0,
              0, 0, 0, 0,
              //////////
              0, 1, 1, 1,
              1, 3, 3, 2,
              1, 3, 3, 2,
              1, 2, 2, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 2x3
  filter.resize(2, 3);
  filter << 1, 1, 1,
            1, 0, 0;
  resize_batch(expected, batch_size, 4 * 5);
  expected << 0, 0, 1, 0, 1,
              1, 1, 2, 2, 1,
              0, 1, 2, 1, 0,
              1, 1, 1, 0, 0,
              /////////////
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              /////////////
              0, 0, 1, 1, 1,
              1, 2, 4, 3, 2,
              1, 2, 4, 3, 2,
              1, 2, 3, 2, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

  // 3x3
  filter.resize(3, 3);
  filter << 1, 1, 1,
            1, 0, 0,
            1, 0, 1;
  resize_batch(expected, batch_size, 5 * 5);
  expected << 1, 0, 2, 0, 1,
              0, 1, 1, 1, 1,
              2, 1, 3, 2, 1,
              0, 1, 2, 1, 0,
              1, 1, 1, 0, 0,
              /////////////
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              0, 0, 0, 0, 0,
              /////////////
              1, 1, 2, 1, 1,
              1, 1, 3, 2, 2,
              2, 3, 6, 4, 3,
              1, 2, 4, 3, 2,
              1, 2, 3, 2, 1;
  output.resizeLike(expected);
  {
     // ensure we don't do allocations in eigen
     BlockAllocations block;
     ConvolutionalLayer::full_conv(input, image_size, image_size, filter, output);
  }
  BOOST_CHECK(expected.isApprox(output, TEST_TOLERANCE));

}

BOOST_AUTO_TEST_CASE(Rotate180_Test)
{
  BOOST_TEST_MESSAGE("*** Rotate180 test ...");

  const MatrixSize input0_size = 2;
  Matrix input0(input0_size, input0_size);
  input0 << 1, 2,
            3, 4;
  Matrix expected0(input0_size, input0_size);
  expected0 << 4, 3,
               2, 1;
  Matrix output0(input0_size, input0_size);
  {
       // ensure we don't do allocations in eigen
       BlockAllocations block;
       ConvolutionalLayer::rotate180(input0, output0);
  }
  BOOST_CHECK(expected0.isApprox(output0, TEST_TOLERANCE));

  const MatrixSize input1_size = 3;
  Matrix input1(input1_size, input1_size);
  input1 << 1, 2, 3,
            4, 5, 6,
            7, 8, 9;
  Matrix expected1(input1_size, input1_size);
  expected1 << 9, 8, 7,
               6, 5, 4,
               3, 2, 1;
  Matrix output1(input1_size, input1_size);
  {
       // ensure we don't do allocations in eigen
       BlockAllocations block;
       ConvolutionalLayer::rotate180(input1, output1);
  }
  BOOST_CHECK(expected1.isApprox(output1, TEST_TOLERANCE));
}

BOOST_AUTO_TEST_CASE(FeedForward_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer FeedForward test ...");

  const MatrixSize image_size = 5;
  const MatrixSize input_size = image_size * image_size;
  const MatrixSize filter_size = 3;
  const MatrixSize batch_size = 2;
  const MatrixSize output_size = ConvolutionalLayer::get_conv_output_size(
      image_size, image_size, filter_size);

  Matrix ww(filter_size, filter_size);
  VectorBatch input, expected_output;
  resize_batch(input, batch_size, input_size);
  resize_batch(expected_output, batch_size, output_size);

  ww <<
      1, 0, 1,
      0, 1, 0,
      1, 0, 1;

  input <<
      1, 0, 1, 0, 0,
      0, 1, 0, 0, 0,
      1, 0, 1, 0, 0,
      0, 0, 0, 0, 0,
      0, 0, 0, 0, 1,
      /////////////
      0, 0, 1, 0, 1,
      0, 1, 0, 1, 0,
      1, 0, 1, 0, 1,
      0, 1, 0, 1, 0,
      1, 0, 1, 0, 0;

  expected_output <<
      5.5, 0.5, 2.5,
      0.5, 2.5, 0.5,
      2.5, 0.5, 2.5,
      /////////////
      4.5, 0.5, 5.5,
      0.5, 5.5, 0.5,
      5.5, 0.5, 4.5;

  auto layer = make_unique<ConvolutionalLayer>(image_size, image_size, filter_size);
  BOOST_CHECK(layer);
  layer->set_activation_function(make_unique<IdentityFunction>());
  layer->set_values(ww, 0.5);

  test_layer_feedforward(*layer, input, expected_output);
}

BOOST_AUTO_TEST_CASE(Backprop_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer backprop test ...");

  const MatrixSize image_size = 3;
  const MatrixSize input_size = image_size * image_size;
  const MatrixSize filter_size = 2;
  const MatrixSize output_size = ConvolutionalLayer::get_conv_output_size(
      image_size, image_size, filter_size);
  const double learning_rate = 0.1;
  const size_t epochs = 100;
  const MatrixSize batch_size = 2;

  Matrix ww(filter_size, filter_size);
  VectorBatch input, expected_output;
  resize_batch(input, batch_size, input_size);
  resize_batch(expected_output, batch_size, output_size);

  ww <<
      0, 0,
      0, 1;

  input <<
      0, 0, 0,
      0, 1, 1,
      0, 1, 0,
      ////////
      0, 0, 0,
      1, 1, 0,
      1, 0, 0;

  expected_output <<
      1, 0,
      0, 0,
      ////
      0, 0,
      1, 0;

  auto layer = make_unique<ConvolutionalLayer>(image_size, image_size, filter_size);
  BOOST_CHECK(layer);
  layer->set_activation_function(make_unique<IdentityFunction>());
  layer->set_values(ww, 0.0);

  test_layer_backprop(
      *layer,
      input,
      boost::none,
      expected_output,
      make_unique<QuadraticCost>(),
      learning_rate,
      epochs
  );
}

BOOST_AUTO_TEST_CASE(Backprop_OnVector_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer backprop on vector test ...");

  const MatrixSize shift = 4;
  const MatrixSize image_size = 3;
  const MatrixSize input_size = image_size * image_size;
  const MatrixSize filter_size = 2;
  const MatrixSize output_size = ConvolutionalLayer::get_conv_output_size(
      image_size, image_size, filter_size);
  const double learning_rate = 0.1;
  const size_t epochs = 100;
  const MatrixSize batch_size = 2;

  Matrix ww(filter_size, filter_size);
  Vector input_buffer(shift + batch_size * input_size);
  MapMatrix input(input_buffer.data() + shift, batch_size, input_size); // this assumes RowMajor layout
  VectorBatch expected_output;
  resize_batch(expected_output, batch_size, output_size);

  ww <<
      0, 0,
      0, 1;

  input_buffer <<
      1, 2, 3, 4, // shift
      0, 0, 0,
      0, 1, 1,
      0, 1, 0,
      ////////
      0, 0, 0,
      1, 1, 0,
      1, 0, 0;

  expected_output <<
      1, 0,
      0, 0,
      ////
      0, 0,
      1, 0;

  auto layer = make_unique<ConvolutionalLayer>(image_size, image_size, filter_size);
  BOOST_CHECK(layer);
  layer->set_activation_function(make_unique<IdentityFunction>());
  layer->set_values(ww, 0.0);

  test_layer_backprop(
      *layer,
      input,
      boost::none,
      expected_output,
      make_unique<QuadraticCost>(),
      learning_rate,
      epochs
  );
}

BOOST_AUTO_TEST_CASE(Training_WithIdentity_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer with Identity activation training test ...");

  const MatrixSize image_rows  = 6;
  const MatrixSize image_cols  = 3;
  const MatrixSize input_size  = image_rows * image_cols;
  const MatrixSize filter_size = 2;
  const MatrixSize output_size = ConvolutionalLayer::get_conv_output_size(image_rows, image_cols, filter_size);
  const MatrixSize batch_size  = 2;

  VectorBatch input, expected;
  resize_batch(input, batch_size, input_size);
  resize_batch(expected, batch_size, output_size);

  input <<
      0.01, 0.02, 0.11, 0.10, 0.13, 0.18,
      0.03, 0.04, 0.12, 0.09, 0.14, 0.17,
      0.05, 0.06, 0.07, 0.08, 0.15, 0.16,
      ///////////////////////////////////
      0.02, 0.06, 0.08, 0.10, 0.13, 0.14,
      0.03, 0.02, 0.07, 0.09, 0.16, 0.15,
      0.04, 0.05, 0.13, 0.11, 0.16, 0.18;

  expected <<
      0.347, 0.478, 0.308, 0.462, 0.378,
      0.492, 0.343, 0.415, 0.409, 0.480,
      //////////////////////////////////
      0.373, 0.446, 0.282, 0.368, 0.395,
      0.442, 0.332, 0.492, 0.429, 0.530;

  // create layer
  auto layer = make_unique<ConvolutionalLayer>(image_rows, image_cols, filter_size);
  BOOST_CHECK(layer);
  layer->set_activation_function(make_unique<IdentityFunction>());
  layer->init(Layer::InitMode_Zeros, boost::none);

  // test
  test_layer_training(
      *layer,
      input,
      expected,
      batch_size, // tests num
      make_unique<QuadraticCost>(),
      0.09, // learning rate
      8000  // epochs
  );
}

BOOST_AUTO_TEST_CASE(Training_WithSigmoid_Test)
{
  BOOST_TEST_MESSAGE("*** ConvolutionalLayer with Sigmoid activation training test ...");

  const MatrixSize image_rows  = 6;
  const MatrixSize image_cols  = 3;
  const MatrixSize input_size  = image_rows * image_cols;
  const MatrixSize filter_size = 2;
  const MatrixSize output_size = ConvolutionalLayer::get_conv_output_size(image_rows, image_cols, filter_size);
  const MatrixSize batch_size  = 2;

  VectorBatch input, expected;
  resize_batch(input, batch_size, input_size);
  resize_batch(expected, batch_size, output_size);

  input <<
      0.01, 0.02, 0.11, 0.10, 0.13, 0.18,
      0.03, 0.04, 0.12, 0.09, 0.14, 0.17,
      0.05, 0.06, 0.07, 0.08, 0.15, 0.16,
      ///////////////////////////////////
      0.02, 0.06, 0.08, 0.10, 0.13, 0.14,
      0.03, 0.02, 0.07, 0.09, 0.16, 0.15,
      0.04, 0.05, 0.13, 0.11, 0.16, 0.18;

  expected <<
      0.348, 0.478, 0.309, 0.462, 0.378,
      0.492, 0.341, 0.412, 0.409, 0.480,
      //////////////////////////////////
      0.372, 0.445, 0.287, 0.366, 0.395,
      0.440, 0.331, 0.494, 0.429, 0.532;

  // create layer
  auto layer = make_unique<ConvolutionalLayer>(image_rows, image_cols, filter_size);
  BOOST_CHECK(layer);
  layer->set_activation_function(make_unique<SigmoidFunction>());
  layer->init(Layer::InitMode_Zeros, boost::none);

  // test
  test_layer_training(
      *layer,
      input,
      expected,
      batch_size, // tests num
      make_unique<CrossEntropyCost>(),
      0.75,  // learning rate
      8000  // epochs
  );
}

////////////////////////////////////////////////////////////////////////////////////////////////
//
// perf conv tests
//
BOOST_AUTO_TEST_CASE(PerfConvTest, * disabled())
{
  conv_perf_test(
      1000, // size
      10,   // filter
      100   // epochs
  );
}

BOOST_AUTO_TEST_CASE(PerfBatchConvTest, * disabled())
{
  conv_perf_batch_test(
      10,   // batch
      1000, // size
      10,   // filter
      10    // epochs
   );
}

BOOST_AUTO_TEST_SUITE_END()


#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <random>
#include <Eigen/Cholesky>
#include <fstream>
#include <iomanip>
#include <srrg_system_utils/parse_command_line.h>

#include "srrg_solver/solver_core/instances.h"
#include "srrg_solver/solver_core/factor_graph.h"
#include "srrg_solver/variables_and_factors/types_3d/variable_se3_ad.h"
#include "srrg_solver/variables_and_factors/types_3d/variable_point3_ad.h"
#include "srrg_solver/variables_and_factors/types_3d/se3_pose_pose_geodesic_error_factor.h"
#include "srrg_solver/variables_and_factors/types_3d/se3_pose_point_offset_error_factor.h"
#include "srrg_solver/variables_and_factors/types_projective/se3_pose_point_omni_ba_error_factor.h"

#include "srrg_solver/variables_and_factors/types_2d/variable_se2_ad.h"
#include "srrg_solver/variables_and_factors/types_2d/variable_point2_ad.h"
#include "srrg_solver/variables_and_factors/types_2d/se2_pose_pose_geodesic_error_factor.h"
#include "srrg_solver/variables_and_factors/types_2d/se2_pose_point_error_factor.h"

#include "factor_noise_adder.h"

using namespace srrg2_core;
using namespace srrg2_solver;
using namespace std;

const std::string exe_name = "solver_app_noise_adder";
#define LOG std::cerr << exe_name << "|"

using StringNoiseAdderMap=std::map<std::string, FactorNoiseAdderPtr >;

default_random_engine rnd_gen;
std::normal_distribution<float> norm_gen;

StringNoiseAdderMap noise_adders;
void initNoiseAdders() {
  noise_adders.insert(std::make_pair("SE3PosePointOffsetErrorFactor", FactorNoiseAdderPtr(new FactorNoiseAdderEuclidean_<SE3PosePointOffsetErrorFactor>(rnd_gen, norm_gen))));

  noise_adders.insert(std::make_pair("SE3PosePoseGeodesicErrorFactor", FactorNoiseAdderPtr(new FactorNoiseAdderSE3Quat_<SE3PosePoseGeodesicErrorFactor>(rnd_gen, norm_gen))));

  noise_adders.insert(std::make_pair("SE2PosePointErrorFactor", FactorNoiseAdderPtr(new FactorNoiseAdderEuclidean_<SE2PosePointErrorFactor>(rnd_gen, norm_gen))));

  noise_adders.insert(std::make_pair("SE2PosePoseGeodesicErrorFactor", FactorNoiseAdderPtr(new FactorNoiseAdderSE2_<SE2PosePoseGeodesicErrorFactor>(rnd_gen, norm_gen))));

  noise_adders.insert(std::make_pair("Sim3PosePoseErrorFactorAD", FactorNoiseAdderPtr(new FactorNoiseAdderSim3_<Sim3PosePoseErrorFactorAD>(rnd_gen, norm_gen))));

  noise_adders.insert(std::make_pair("SE3PosePointOmniBAErrorFactor", FactorNoiseAdderPtr(new FactorNoiseAdderNormalize_<SE3PosePointOmniBAErrorFactor>(rnd_gen, norm_gen))));

}

std::string listNoiseAdders() {
  ostringstream os;
  for (const auto& it: noise_adders) {
    os << it.first << endl;
  }
  return os.str();
}

static const char* banner[] = {
  "adds noise to an (ideal) graph, based on the value of the information matrices",
  "usage: solver_app_noise_adder -i <input> -o <output>",
  0
};

// ia THE PROGRAM
int main(int argc, char** argv) {
  initNoiseAdders();
  ParseCommandLine cmd_line(argv, banner);
  ArgumentInt seed                    (&cmd_line, "s",    "seed",           "starting_seed", 0);
  ArgumentString output_file          (&cmd_line, "o",    "output-file",           "file where to save the output", "");
  ArgumentString input_file          (&cmd_line, "i",    "input-file",             "file where to read the input ", "");
  ArgumentFlag list_adders          (&cmd_line, "l",    "list-adders",             " prints a list of noise adders registered");
  cmd_line.parse();

  if (seed.isSet()){
    rnd_gen.seed(seed.value());
  }
  if (list_adders.isSet()) {
    cerr << "Adders: " << endl;
    cerr << listNoiseAdders() << endl;
    return 0;
  }
  if (! input_file.isSet()) {
    cerr << "no input file provided, aborting" << endl;
    return 0;
  }
  std::cerr << "loding file: [" << input_file.value() << "]... ";
  FactorGraphPtr graph = FactorGraph::read(input_file.value());
  std::cerr << "done, factors:" << graph->factors().size() << " vars: " << graph->variables().size() << std::endl;

  std::cerr << "adding noise...";
  for (auto it = graph->factors().begin();
       it != graph->factors().end(); ++it) {
    FactorBase* factor=it.value();

    auto adder_it=noise_adders.find(factor->className());
    if (adder_it==noise_adders.end()) {
      cerr << "unknown adder [" << factor->className() << "] skupping" << endl;
      continue;
    }
    FactorNoiseAdderPtr adder=adder_it->second;

    adder->addNoise(factor);
  }

  std::cerr << "done" << std::endl;

  if (! output_file.isSet()) {
    cerr << "no output file provided, skipping output" << endl;
    return 0;
  }

  std::cerr << "writing output to file [ " << output_file.value() << "]... ";
  graph->write(output_file.value());
  cerr << " done" << endl;
  return 0;
}


#include <Eigen/Core>
#include <igl/writeOBJ.h>
#include "platonic_solid.h"

int main(int argc, char **argv) {
    Eigen::MatrixXd TV, OV, CV, IV, DV;
    Eigen::MatrixXi TF, OF, CF, IF, DF;

    platonic_solid::tetrahedron(TV, TF);
    igl::writeOBJ("tetrahedron.obj", TV, TF);

    platonic_solid::octahedron(OV, OF);
    igl::writeOBJ("octahedron.obj", OV, OF);

    platonic_solid::cube(CV, CF);
    igl::writeOBJ("cube.obj", CV, CF);

    platonic_solid::icosahedron(IV, IF);
    igl::writeOBJ("icosahedron.obj", IV, IF);

    platonic_solid::dodecahedron(DV, DF);
    igl::writeOBJ("dodecahedron.obj", DV, DF);
}


#include <stan/math/rev/arr.hpp>
#include <gtest/gtest.h>
#include <boost/numeric/odeint.hpp>
#include <test/unit/math/rev/arr/functor/util.hpp>
#include <test/unit/math/prim/arr/functor/harmonic_oscillator.hpp>
#include <test/unit/math/prim/arr/functor/lorenz.hpp>
#include <iostream>
#include <sstream>
#include <vector>

template <typename F, typename T_y0, typename T_theta>
void sho_value_test(F harm_osc, std::vector<double>& y0, double t0,
                    std::vector<double>& ts, std::vector<double>& theta,
                    std::vector<double>& x, std::vector<int>& x_int) {
  using stan::math::promote_scalar;
  using stan::math::var;

  std::vector<std::vector<var> > ode_res_vd = stan::math::integrate_ode_rk45(
      harm_osc, promote_scalar<T_y0>(y0), t0, ts,
      promote_scalar<T_theta>(theta), x, x_int, 0);
  EXPECT_NEAR(0.995029, ode_res_vd[0][0].val(), 1e-5);
  EXPECT_NEAR(-0.0990884, ode_res_vd[0][1].val(), 1e-5);

  EXPECT_NEAR(-0.421907, ode_res_vd[99][0].val(), 1e-5);
  EXPECT_NEAR(0.246407, ode_res_vd[99][1].val(), 1e-5);
}

void sho_finite_diff_test(double t0) {
  using stan::math::var;
  harm_osc_ode_fun harm_osc;

  std::vector<double> theta;
  theta.push_back(0.15);

  std::vector<double> y0;
  y0.push_back(1.0);
  y0.push_back(0.0);

  std::vector<double> ts;
  for (int i = 0; i < 100; i++)
    ts.push_back(t0 + 0.1 * (i + 1));

  std::vector<double> x;
  std::vector<int> x_int;

  test_ode(harm_osc, t0, ts, y0, theta, x, x_int, 1e-8, 1e-4);

  sho_value_test<harm_osc_ode_fun, double, var>(harm_osc, y0, t0, ts, theta, x,
                                                x_int);
  sho_value_test<harm_osc_ode_fun, var, double>(harm_osc, y0, t0, ts, theta, x,
                                                x_int);
  sho_value_test<harm_osc_ode_fun, var, var>(harm_osc, y0, t0, ts, theta, x,
                                             x_int);
}

void sho_data_finite_diff_test(double t0) {
  using stan::math::var;
  harm_osc_ode_data_fun harm_osc;

  std::vector<double> theta;
  theta.push_back(0.15);

  std::vector<double> y0;
  y0.push_back(1.0);
  y0.push_back(0.0);

  std::vector<double> ts;
  for (int i = 0; i < 100; i++)
    ts.push_back(t0 + 0.1 * (i + 1));

  std::vector<double> x(3, 1);
  std::vector<int> x_int(2, 0);

  test_ode(harm_osc, t0, ts, y0, theta, x, x_int, 1e-8, 1e-4);

  sho_value_test<harm_osc_ode_data_fun, double, var>(harm_osc, y0, t0, ts,
                                                     theta, x, x_int);
  sho_value_test<harm_osc_ode_data_fun, var, double>(harm_osc, y0, t0, ts,
                                                     theta, x, x_int);
  sho_value_test<harm_osc_ode_data_fun, var, var>(harm_osc, y0, t0, ts, theta,
                                                  x, x_int);
}

TEST(StanAgradRevOde_integrate_ode_rk45, harmonic_oscillator_finite_diff) {
  sho_finite_diff_test(0);
  sho_finite_diff_test(1.0);
  sho_finite_diff_test(-1.0);

  sho_data_finite_diff_test(0);
  sho_data_finite_diff_test(1.0);
  sho_data_finite_diff_test(-1.0);
}

TEST(StanAgradRevOde_integrate_ode_rk45, lorenz_finite_diff) {
  lorenz_ode_fun lorenz;

  std::vector<double> y0;
  std::vector<double> theta;
  double t0;
  std::vector<double> ts;

  t0 = 0;

  theta.push_back(10.0);
  theta.push_back(28.0);
  theta.push_back(8.0 / 3.0);
  y0.push_back(10.0);
  y0.push_back(1.0);
  y0.push_back(1.0);

  std::vector<double> x;
  std::vector<int> x_int;

  for (int i = 0; i < 100; i++)
    ts.push_back(0.1 * (i + 1));

  test_ode(lorenz, t0, ts, y0, theta, x, x_int, 1e-8, 1e-1);
}


#include "signing.h"
#include <NTL/ZZ.h>
#include "generalhelpers.h"

namespace CryptoHelpers
{
    void ModifyHashToLong(const long long& hash, NTL::ZZ& longModifiedHash)
    {
        longModifiedHash = 0;
        auto rhash = hash;
        ReverseBytes(8, rhash);
        ConvertHexStringToLong("00FFFFFFFFFFFF00", longModifiedHash);

#pragma region dirtyhackforntl
        longModifiedHash *= NTL::power2_ZZ(16);
        longModifiedHash += ((rhash >> 48) & 0xFFFF);
        longModifiedHash *= NTL::power2_ZZ(16);
        longModifiedHash += ((rhash >> 32) & 0xFFFF);
        longModifiedHash *= NTL::power2_ZZ(16);
        longModifiedHash += ((rhash >> 16) & 0xFFFF);
        longModifiedHash *= NTL::power2_ZZ(16);
        longModifiedHash += ((rhash) & 0xFFFF);

#pragma endregion
    }

    bool CheckSignature(const long long& hash,const NTL::ZZ& signature,const NTL::ZZ& y,const NTL::ZZ& k)
    {
        NTL::ZZ a(0);
        NTL::ZZ p(0);
        ConvertHexStringToLong("9E93A4096E5416CED0242228014B67B5", a);
        ConvertHexStringToLong("AF5228967057FE1CB84B92511BE89A47", p);

        NTL::ZZ longModifiedHash;
        ModifyHashToLong(hash, longModifiedHash);
        NTL::ZZ tmp1(0);
        NTL::PowerMod(tmp1,signature, longModifiedHash, p);
        NTL::ZZ tmp2(0);
        NTL::ZZ tmp3(0);
        NTL::ZZ tmp4(0);
        NTL::PowerMod(tmp3,a, k, p);
        NTL::MulMod(tmp4, signature, tmp3, p);
        NTL::PowerMod(tmp2,y, (tmp4), p);

        if (tmp2 == tmp1)
        {
            return true;
        }
        return false;
    }

    bool Sign(const long long& hash, NTL::ZZ& signature, NTL::ZZ& y, NTL::ZZ& k, NTL::ZZ& g, NTL::ZZ& rnd, NTL::ZZ& z, NTL::ZZ& x)
    {
        NTL::ZZ longModifiedHash;
        ModifyHashToLong(hash, longModifiedHash);

        GetRandom(rnd);

        NTL::ZZ a(0);
        NTL::ZZ p(0);
        NTL::ZZ q(0);
        ConvertHexStringToLong("9E93A4096E5416CED0242228014B67B5", a);
        ConvertHexStringToLong("AF5228967057FE1CB84B92511BE89A47", p);
        ConvertHexStringToLong("57A9144B382BFF0E5C25C9288DF44D23", q);
        NTL::PowerMod(z,
            a, rnd, p
        );
        x = NTL::RandomBits_ZZ(127);
        y = NTL::PowerMod(a, x, p);
        k = ((rnd*longModifiedHash - x*z) * NTL::InvMod(longModifiedHash, q)) % q;
        g = (x*z * NTL::InvMod(longModifiedHash,q)) % q;
        NTL::PowerMod(signature,a, g, p);
        return true;
    }
}


// TRENTO: Reduced Thickness Event-by-event Nuclear Topology
// Copyright 2015 Jonah E. Bernhard, J. Scott Moreland
// MIT License

#include "../src/nucleon.h"

#include <cmath>

#include "catch.hpp"
#include "util.h"

#include "../src/eventqty.h"
#include "../src/nucleus.h"
#include "../src/random.h"
#include <boost/math/constants/constants.hpp>
#include <boost/multi_array.hpp>

#include <iostream>

using namespace trento;

TEST_CASE( "nucleon" ) {
  auto fluct = 1. + .5*random::canonical<>();
  auto xsec = 4. + 3.*random::canonical<>();
  auto width = .5 + .2*random::canonical<>();
  auto wsq = width*width;

  auto var_map = make_var_map({
      {"ncoll", false},
      {"fluctuation",   fluct},
      {"cross-section", xsec},
      {"nucleon-width", width},
      {"constit-width", width},
      {"constit-number", 1},
      {"columns", DefaultEventQuantityList},
  });

  NucleonCommon nc{var_map};

  Proton proton{};
  proton.sample_nucleons(0.);
  auto& nucleon = *proton.begin();
  while (!nc.participate(nucleon, nucleon)) {}

  // truncation radius
  auto R = 5*width;

  // random angle
  auto phi = math::double_constants::two_pi * random::canonical<>();
  auto cs = std::cos(phi);
  auto sn = std::sin(phi);

  // thickness function
  // check relative to zero
  auto tzero = nc.thickness(nucleon, 0., 0.);
  CHECK( nc.thickness(nucleon, width*cs, width*sn) == Approx(tzero*std::exp(-.5)) );

  // random point inside radius
  auto x = R * random::canonical<>();
  auto y = R * random::canonical<>();
  auto dsq = x*x + y*y;
  CHECK( nc.thickness(nucleon, x, y) == Approx(tzero*std::exp(-.5*dsq/wsq)).epsilon(1e-5).margin(1e-5) );

  // random point outside radius
  auto d = R * (1 + random::canonical<>());
  CHECK( nc.thickness(nucleon, d*cs, d*sn) == 0. );

  // fluctuations
  // just check they have unit mean -- the rest is handled by the C++ impl.
  auto total = 0.;
  auto n = 1e6;
  for (auto i = 0; i < static_cast<int>(n); ++i) {
    proton.sample_nucleons(0.);
    nucleon = *proton.begin();
    while (!nc.participate(nucleon, nucleon)) {}
    total += nc.thickness(nucleon, 0., 0.) * (2*M_PI*wsq);
  }

  auto mean = total/n;
  CHECK( mean == Approx(1.).epsilon(.003) );

  // must use a Nucleus to set Nucleon position
  // Proton conveniently sets a deterministic position
  // a mock class would be better but this works fine
  Proton A{}, B{};
  A.sample_nucleons(0.);
  B.sample_nucleons(0.);
  auto& nA = *A.begin();
  auto& nB = *B.begin();
  CHECK( nA.x() == 0. );
  CHECK( nA.y() == 0. );
  CHECK( nA.z() == 0. );
  CHECK( !nA.is_participant() );

  // wait until the nucleons participate
  while (!nc.participate(nA, nB)) {}
  CHECK( nA.is_participant() );
  CHECK( nB.is_participant() );

  // resampling nucleons resets participant state
  A.sample_nucleons(0.);
  CHECK( !nA.is_participant() );

  // test cross section
  // min-bias impact params
  auto bmax = nc.max_impact();
  CHECK( bmax == Approx(6*width) );

  auto nev = 1e6;
  auto count = 0;
  for (auto i = 0; i < static_cast<int>(nev); ++i) {
    auto b = bmax * std::sqrt(random::canonical<>());
    A.sample_nucleons(.5*b);
    B.sample_nucleons(-.5*b);
    if (nc.participate(nA, nB))
      ++count;
  }

  auto xsec_mc = M_PI*bmax*bmax * static_cast<double>(count)/nev;

  // precision is better than this, but let's be conservative
  CHECK( xsec_mc == Approx(xsec).epsilon(.02) );

  // impact larger than max should never participate
  auto b = bmax + random::canonical<>();
  A.sample_nucleons(.5*b);
  B.sample_nucleons(-.5*b);
  CHECK( !nc.participate(nA, nB) );

  // very large fluctuation parameters mean no fluctuations
  auto no_fluct_var_map = make_var_map({
      {"ncoll", false},
      {"fluctuation",   1e12},
      {"cross-section", xsec},
      {"nucleon-width", width},
      {"constit-width", width},
      {"constit-number", 1},
  });

  NucleonCommon no_fluct_nc{no_fluct_var_map};

  proton.sample_nucleons(0.);
  nucleon = *proton.begin();
  while (!no_fluct_nc.participate(nucleon, nucleon)) {}
  CHECK( no_fluct_nc.thickness(nucleon, 0, 0) == Approx(1/(2*M_PI*wsq)) );

  CHECK_THROWS_AS([]() {
    // nucleon width too small
    auto bad_var_map = make_var_map({
        {"ncoll", false},
        {"fluctuation",   1.},
        {"cross-section", 5.},
        {"nucleon-width", .1},
        {"constit-width", .1},
        {"constit-number", 1},
    });
    NucleonCommon bad_profile{bad_var_map};
  }(),
  std::domain_error);

  // test nucleon-nucleon cross section with random constituent substructure
  auto nucleon_width = 0.5;
  auto constituent_number = std::uniform_int_distribution<>{2, 6}(random::engine);
  auto constituent_width = .3 + .2*random::canonical<>();

  // Coarse-ish p+p grid.
  auto grid_max = 3.;
  auto grid_step = 0.1;
  auto grid_nsteps = 60;

  // constituent var map
  var_map = make_var_map({
      {"ncoll", false},
      {"grid-max", grid_max},
      {"grid-step", grid_step},
      {"fluctuation",   1e12},
      {"cross-section", xsec},
      {"nucleon-width", nucleon_width},
      {"constit-width", constituent_width},
      {"constit-number", constituent_number},
  });

  NucleonCommon nc_constituent{var_map};
  bmax = nc_constituent.max_impact();

  // inelastic collision counter
  count = 0;

  // measure the cross section with min bias p+p collisions
  for (auto i = 0; i < static_cast<int>(nev); ++i) {
    b = bmax * std::sqrt(random::canonical<>());
    A.sample_nucleons(.5*b);
    B.sample_nucleons(-.5*b);
    if (nc_constituent.participate(*A.begin(), *B.begin()))
      ++count;
  }

  // calculate Monte Carlo cross section
  xsec_mc = M_PI*bmax*bmax * static_cast<double>(count)/nev;

  // assert a rather loose tolerance for the numerical cross section finder
  CHECK( xsec_mc/xsec == Approx(1).epsilon(.02) );

  // verify root-mean-square nucleon width for nucleons with substructure
  double mean_square_width_num = 0.0;
  double mean_square_width_denom = 0.0;

  for (auto i = 0; i < 1000; ++i) {
    A.sample_nucleons(0.);
    auto nucleon = *A.begin();

    while (!nc_constituent.participate(nucleon, nucleon)) {}

    for (auto iy = 0; iy < grid_nsteps; ++iy) {
      for (auto ix = 0; ix < grid_nsteps; ++ix) {
        auto x = (ix + .5) * 2 * grid_max/grid_nsteps - grid_max;
        auto y = (iy + .5) * 2 * grid_max/grid_nsteps - grid_max;
        auto thick = nc_constituent.thickness(nucleon, x, y);
        mean_square_width_num += (x*x + y*y)*thick;
        mean_square_width_denom += 2*thick;
      }
    }
  }

  auto rms_width = std::sqrt(mean_square_width_num / mean_square_width_denom);

  CHECK( rms_width == Approx(nucleon_width).epsilon(.01).margin(.01) );

}


#include <iostream>
#include <unordered_map>

#include <NTL/ZZ.h>

#include "RSA.h"

using namespace std;
using namespace NTL;

void usage(char *progname) {
	cout << "This program returns the private exponent `d` of the "
		"public key <n, e2> given the secret key <n, e', d'>."
		<< endl;
	cout << "Usage: " << progname << " n e' d' e" << endl;
}

int main(int argc, char *argv[]) {
	if (argc != 5) { usage(argv[0]); return 3; }

	ZZ n, e1, d1, e2, d2, phi;
	n = conv<ZZ>(argv[1]);
	e1 = conv<ZZ>(argv[2]);
	d1 = conv<ZZ>(argv[3]);
	e2 = conv<ZZ>(argv[4]);

	RSAkey k1 = RSAkey(n, e1, d1);
	phi = k1.get_param("phi");

	// d2 = e2^(-1) (mod phi)
	InvMod(d2, e2, phi);

	RSAkey k2 = RSAkey(n, e2, d2);
	cout << k2.get_param("d") << endl;
}


//
// Timing Pinocchio
//
// Based on timings.cpp from Pinocchio
//

#include "pinocchio/spatial/fwd.hpp"
#include "pinocchio/spatial/se3.hpp"
#include "pinocchio/multibody/visitor.hpp"
#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/centroidal.hpp"
#include "pinocchio/algorithm/aba.hpp"
#include "pinocchio/algorithm/rnea.hpp"
#include "pinocchio/algorithm/cholesky.hpp"
#include "pinocchio/algorithm/jacobian.hpp"
#include "pinocchio/algorithm/center-of-mass.hpp"
#include "pinocchio/algorithm/compute-all-terms.hpp"
#include "pinocchio/algorithm/kinematics.hpp"
#include "pinocchio/parsers/urdf.hpp"
#include "pinocchio/parsers/sample-models.hpp"

#include <stdlib.h>
#include <iostream>
#include <fstream>
#include <sstream>

#include "csv_reader.h"
#include "tictoc_timer.h"

#include <Eigen/StdVector>
EIGEN_DEFINE_STL_VECTOR_SPECIALIZATION(Eigen::VectorXd)

int main(int argc, const char ** argv)
{
   using namespace Eigen;
   using namespace pinocchio;

   // Check arguments
   if (argc != 2)
   {
      std::cerr << "You have to specify a robot model. Choices are: iiwa, hyq, atlas" << std::endl;
      return 1;
   }

   // Set robot model
   bool floating_base;
   std::string robot_model(argv[1]);
   if (robot_model == "iiwa")
   {
      std::cout << "Robot Model = " << robot_model << std::endl;
      floating_base = false;
   }
   else if ((robot_model == "hyq")|(robot_model == "atlas"))
   {
      std::cout << "Robot Model = " << robot_model << std::endl;
      floating_base = true;
   }
   else
   {
      std::cerr << "Invalid robot model: " << robot_model << "\nChoices are: iiwa, hyq, atlas" << std::endl;
      return 2;
   }

   // Setup timer
   TicToc timer(TicToc::NS);
   const int NBT = 1000*100;

   // Import URDF model
   Model model;
   std::string urdf_filename;
   urdf_filename = RBD_BENCHMARKS_DIR"/description/urdf/"+robot_model+".urdf";
   if (floating_base)
      pinocchio::urdf::buildModel(urdf_filename,JointModelFreeFlyer(),model);
   else
      pinocchio::urdf::buildModel(urdf_filename,model);
   model.gravity = Eigen::Vector3d(0,0,0);
   Data data(model);
   int dof = model.nv;
   std::cout << "dof = " << dof << std::endl;

   // Import CSV inputs
   std::string input_filename;
   input_filename = RBD_BENCHMARKS_DIR"/csv/pinocchio/"+robot_model+"_inputs.csv";
   std::ifstream input_csv(input_filename.c_str());
   CSVRow row;
   std::vector<VectorXd> qs     (NBT, VectorXd::Zero(model.nq));
   std::vector<VectorXd> qdots  (NBT, VectorXd::Zero(model.nv));
   std::vector<VectorXd> qddots (NBT, VectorXd::Zero(model.nv));
   std::vector<VectorXd> taus   (NBT, VectorXd::Zero(model.nv));
   input_csv >> row;
   int tot_q, tot_qdot;
   tot_q    = model.nq;
   tot_qdot = model.nv;
   std::cout << "qs, qdots = " << tot_q << ", " << tot_qdot << std::endl;
   int col, start_col;
   for(int i=0;i<NBT;++i)
   {
      input_csv >> row;
      start_col = 0;                 // 0
      for(int j=0;j<tot_q;++j)
      {
            col = start_col+j;
            qs[i][j] = atof(row[col].c_str());
      }
      start_col = tot_q;               // 1xQ
      for(int j=0;j<tot_qdot;++j)
      {
            col = start_col+j;
            qdots[i][j] = atof(row[col].c_str());
      }
      start_col = tot_q+tot_qdot;      // 1xQ+1xQd
      for(int j=0;j<tot_qdot;++j)
      {
            col = start_col+j;
            qddots[i][j] = atof(row[col].c_str());
      }
      start_col = tot_q+2*tot_qdot;  // 1xQ+2xQd
      for(int j=0;j<tot_qdot;++j)
      {
            col = start_col+j;
            taus[i][j] = atof(row[col].c_str());
      }
   }

   // Initialize some output variables
   double time_data;

   std::cout << "--" << std::endl;

   // Dynamics Algorithms
   std::string algorithm_name;
   #ifdef RNEA_ALG
   algorithm_name = "rnea";
   std::cout << "RNEA" << std::endl;
   #elif  CRBA_ALG
   algorithm_name = "crba";
   std::cout << "CRBA" << std::endl;
   #elif  ABA_ALG
   algorithm_name = "aba";
   std::cout << "ABA" << std::endl;
   #else
   algorithm_name = "rnea";
   std::cout << "RNEA" << std::endl;
   #endif /* Dynamics Algorithms */

   // Write CSV outputs
   std::ofstream output_csv;
   std::string output_filename;
   std::stringstream num_inputs;
   num_inputs << NBT;
   output_filename = RESULTS_DIR"/terr_"+robot_model+"_"+num_inputs.str()+"_pinocchio_"+algorithm_name+".csv";
   output_csv.open (output_filename.c_str());

   SMOOTH(NBT)
   {
      if (_smooth != 0) output_csv << "\n";

      timer.tic();
      // Dynamics Algorithms
      #ifdef RNEA_ALG
      rnea(model,data,qs[_smooth],qdots[_smooth],qddots[_smooth]);
      #elif  CRBA_ALG
      crba(model,data,qs[_smooth]);
      #elif  ABA_ALG
      aba(model,data,qs[_smooth],qdots[_smooth],taus[_smooth]);
      #else
      rnea(model,data,qs[_smooth],qdots[_smooth],qddots[_smooth]);
      #endif /* Dynamics Algorithms */
      time_data = timer.toc(TicToc::NS);

      output_csv << time_data;
   }

   output_csv.close();

   std::cout << "--" << std::endl;

   return 0;
}


// Copyright PinaPL
//
// cell.cpp
// PinaPL
//
#include <math.h>
#include <Eigen/Dense>
#include <vector>
#include "weights.hpp"
#include "cell.hpp"
#include "functions.hpp"

Cell::Cell(Weights* weights) {
    this->weights = weights;
    this->reset();
}

void Cell::compute(Eigen::MatrixXd* input) {
/*    this->forget_gate_out =
        (this->weights->weight_in_forget_gate * input
        + this->weights->weight_st_forget_gate * previous_cell_state)
        .unaryExpr(&sigmoid); */
    this->inputs.push_back((*input));

    this->input_gate_out.push_back(
        (this->weights->weight_in_input_gate * (*input)
        + this->weights->weight_st_input_gate * this->cell_out.back()
        + this->weights->bias_input_gate).unaryExpr(&sigmoid));

    this->input_block_out.push_back(
        (this->weights->weight_in_input_block * (*input)
        + this->weights->weight_st_input_block * this->cell_out.back()
        + this->weights->bias_input_block).unaryExpr(&tanhyp));

    this->output_gate_out.push_back(
        (this->weights->weight_in_output_gate * (*input)
        + this->weights->weight_st_output_gate * this->cell_out.back()
        + this->weights->bias_output_gate).unaryExpr(&sigmoid));

    this->cell_state.push_back(
        (this->cell_state.back()
        + this->input_gate_out.back()
        .cwiseProduct(this->input_block_out.back())));

    this->cell_out.push_back(
        this->cell_state.back().unaryExpr(&tanhyp)
        .cwiseProduct(this->output_gate_out.back()));
}

Eigen::MatrixXd Cell::compute_gate_gradient(Eigen::MatrixXd* deltas, int time) {
    int output_size = this->weights->output_size;
    // Computes dy(t)
    delta_cell_out.push_back(
        (*deltas)
        + this->weights->weight_st_input_block * delta_input_block_out.back()
        + this->weights->weight_st_input_gate * delta_input_gate_out.back()
//      + this->weights->weight_st_forget_gate * delta_forget_gate_out.back()
        + this->weights->weight_st_output_gate * delta_output_gate_out.back() );

    // Computes do(t)
    delta_output_gate_out.push_back(delta_cell_out.back()
        .cwiseProduct(cell_state.at(time + 1).unaryExpr(&tanhyp))
        .cwiseProduct(output_gate_out.at(time + 1).cwiseProduct(
            Eigen::MatrixXd::Ones(output_size, 1)
            - output_gate_out.at(time + 1))));

    // Computes dc(t)
    delta_cell_state.push_back(
        delta_cell_out.back()
        .cwiseProduct(output_gate_out.at(time + 1))
        .cwiseProduct(cell_state.at(time + 1).unaryExpr(&tanh_derivative)));

    // Computes di(t)
    delta_input_gate_out.push_back(
        delta_cell_state.back()
        .cwiseProduct(input_block_out.at(time + 1))
        .cwiseProduct(input_gate_out.at(time + 1).cwiseProduct(
            Eigen::MatrixXd::Ones(output_size, 1)
            - input_gate_out.at(time + 1))) );

    // Computes dz(t)
    delta_input_block_out.push_back(
        delta_cell_state.back()
        .cwiseProduct(input_gate_out.at(time + 1))
        .cwiseProduct(input_block_out.at(time + 1).cwiseProduct(
            Eigen::MatrixXd::Ones(output_size, 1)
            - input_block_out.at(time + 1))) );

    // Computes dx(t)
    Eigen::MatrixXd delta_input =
    this->weights->weight_in_input_block.transpose()
      * delta_input_block_out.back()
    + this->weights->weight_in_input_gate.transpose()
      * delta_input_gate_out.back()
//  + this->weights->weight_in_input_block.transpose()
//    * delta_input_block_out.back()
    + this->weights->weight_in_output_gate.transpose()
      * delta_output_gate_out.back();

    return delta_input;
}

void Cell::compute_weight_gradient() {
    int last_item_index = this->inputs.size() - 1;
    // Computes dW
    for (int t = 0; t < last_item_index + 1; ++t) {
        // Computes dWz
        this->weights->delta_weight_in_input_block +=
            delta_input_block_out.at(last_item_index - t + 1)
            * inputs.at(t).transpose();
/*
        std::cout << "computeWG : "<< t << std::endl;
        std::cout << delta_input_gate_out.at(last_item_index - t + 1)
        << std::endl;
        std::cout << " * " << std::endl;
        std::cout << inputs.at(t).transpose() << std::endl;
        std::cout << " = " << std::endl;
        std::cout << delta_input_gate_out.at(last_item_index - t + 1)
        * inputs.at(t).transpose() << std::endl;
*/
        // Computes dWi
        this->weights->delta_weight_in_input_gate +=
            delta_input_gate_out.at(last_item_index - t + 1)
            * inputs.at(t).transpose();

        // Computes dWf
        /*
        this->weights->delta_weight_in_input_block +=
            delta_input_block_out.at(last_item_index - t + 1)
            * inputs.at(t).transpose(); */

        // Computes dWo
        this->weights->delta_weight_in_output_gate +=
            delta_output_gate_out.at(last_item_index - t + 1)
            * inputs.at(t).transpose();
    }
    // Computes dR
    for (int t = 0; t < last_item_index; ++t) {
        // Computes dRz
        this->weights->delta_weight_st_input_block +=
            delta_input_block_out.at(last_item_index - t)
            * cell_out.at(t + 1).transpose();

        // Computes dRi
        this->weights->delta_weight_st_input_gate +=
            delta_input_gate_out.at(last_item_index - t)
            * cell_out.at(t + 1).transpose();

        // Computes dRo
        this->weights->delta_weight_st_output_gate +=
            delta_output_gate_out.at(last_item_index - t)
            * cell_out.at(t + 1).transpose();
    }
    // Computes dB
    for (int t = 0; t < last_item_index + 1; ++t) {
        // Computes dBz
        this->weights->delta_bias_input_block +=
            delta_input_block_out.at(last_item_index - t + 1);
        // Computes dBi
        this->weights->delta_bias_input_gate +=
            delta_input_gate_out.at(last_item_index - t + 1);
        // Computes dBo
        this->weights->delta_bias_output_gate +=
            delta_output_gate_out.at(last_item_index - t + 1);
    }
}

void Cell::update_weights(double lambda) {
    this->weights->apply_gradient(lambda);
}

void Cell::reset() {
    int output_size = this->weights->output_size;

    this->inputs.clear();
    this->input_gate_out.clear();
    this->input_block_out.clear();
    this->output_gate_out.clear();
    this->cell_state.clear();
    this->cell_out.clear();

    this->delta_cell_out.clear();
    this->delta_output_gate_out.clear();
    this->delta_cell_state.clear();
    this->delta_input_gate_out.clear();
    this->delta_input_block_out.clear();

    this->input_gate_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->input_block_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->output_gate_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->cell_state.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->cell_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));

    this->delta_cell_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->delta_output_gate_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->delta_cell_state.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->delta_input_gate_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
    this->delta_input_block_out.push_back(
        Eigen::MatrixXd::Zero(output_size, 1));
}


#include <Eigen/Cholesky>
#include <Eigen/Dense>
#include "Object.hpp"

const double PI = 3.141592653589793238;

void Object::internalProjection(const Eigen::Vector2d& obs) {
    HeapType heap(edges);
    for (size_t i = 0; i < edges.size(); i++) {
        if (edges[i].valid)
            heap.emplace(i);
    }
    while (heap.empty() == false) {
        size_t top = heap.top();
        heap.pop();
        for (size_t i = 0; i < edges.size(); i++) {
            Edge& this_edge = edges[top];
            Edge& eg = edges[i];
            if (eg.valid == false || &eg == &this_edge)
                continue;
            printf("Source : %lu, dst: %lu\n", top, i);
            projectEdge2Edge(this_edge, obs, eg, heap);
            printf("After Source : %lu, dst: %lu\n", top, i);
        }
    }
}

void Object::externalOcclusion(Object& obj, Eigen::Vector2d obs) {
    HeapType heap(edges);
    for (size_t i = 0; i < edges.size(); i++) {
        if (edges[i].valid)
            heap.emplace(i);
    }
    while (heap.empty() == false) {
        size_t top = heap.top();
        Edge& this_edge = edges[top];
        heap.pop();
        obj.externalProjector(this_edge, obs);
    }
}

void Object::externalProjector(Edge& src, const Eigen::Vector2d& obs) {
    HeapType heap(edges);
    for (size_t i = 0; i < edges.size(); i++)
        heap.emplace(i);
    while (heap.empty() == false) {
        size_t top = heap.top();
        Edge& this_edge = edges[top];
        heap.pop();
        if (this_edge.valid == false) continue;
        projectEdge2Edge(src, obs, this_edge, heap);
    }
    bool valid_flag = false;
    for (size_t i = 0; i < edges.size(); i++) {
        if (edges[i].valid == true) {
            valid_flag = true;
            break;
        }
    }
    if (valid_flag == false)                    // 假如全部被遮挡，那么这个object就没有投影的必要了
        valid = false;
}

void Object::intialize(const std::vector<Eigen::Vector2d>& pts, const Eigen::Vector2d& obs) {
    edges.clear();
    Edge to_add;
    bool zero_pushed = false;
    for (size_t i = 1; i < pts.size(); i++) {
        Eigen::Vector2d vec = pts[i] - pts[i - 1];
        Eigen::Vector2d ctr_vec = (pts[i] + pts[i - 1]) / 2.0 - obs;
        Eigen::Vector2d norm = Eigen::Vector2d(-vec(1), vec(0));
        if (ctr_vec.dot(norm) < 0.0) {
            if (i == 1)
                zero_pushed = true;
            Eigen::Vector2d o2p = pts[i - 1] - obs;
            to_add.emplace_back();
            to_add.back() << pts[i - 1], atan2(o2p(1), o2p(0));
        } else {
            if (to_add.empty() == false) {
                Eigen::Vector2d o2p = pts[i - 1] - obs;
                to_add.emplace_back();
                to_add.back() << pts[i - 1], atan2(o2p(1), o2p(0));
                edges.push_back(to_add);
                to_add.reset();
            }
        }
    }
    Eigen::Vector2d vec = pts.front() - pts.back();
    Eigen::Vector2d ctr_vec = (pts.back() + pts.front()) / 2.0 - obs;
    Eigen::Vector2d norm = Eigen::Vector2d(-vec(1), vec(0));
    if (to_add.empty() == false) {
        Eigen::Vector2d o2p = pts.back() - obs;
        to_add.emplace_back();
        to_add.back() << pts.back(), atan2(o2p(1), o2p(0));
    }
    if (ctr_vec.dot(norm) < 0.0) {
        if (zero_pushed == true) {
            Edge& front = edges.front();
            if (to_add.empty() == true) {
                front.emplace_front();
                Eigen::Vector2d o2p = pts.back() - obs;
                front.front() << pts.back(), atan2(o2p(1), o2p(0));
            } else {
                for (Edge::const_reverse_iterator rit = to_add.crbegin(); rit != to_add.crend(); rit++)
                    front.push_front(*rit);
            }
        } else {
            if (to_add.size() == 0) {
                Eigen::Vector2d o2p = pts.back() - obs;
                to_add.emplace_back();
                to_add.back() << pts.back(), atan2(o2p(1), o2p(0));
            }
            Eigen::Vector2d o2p = pts.front() - obs;
            to_add.emplace_back();
            to_add.back() << pts.front(), atan2(o2p(1), o2p(0));
            edges.push_back(to_add);
        }
    } else if (to_add.empty() == false) {
        edges.push_back(to_add);
    }
    for (Edge& eg: edges) {
        for (const Eigen::Vector3d& pt: eg) {
            double dist = (pt.block<2, 1>(0, 0) - obs).norm();
            if (dist < eg.min_dist) eg.min_dist = dist;
        }
    }
    for (size_t i = 0; i < edges.size(); i++) {
        Edge& eg = edges[i];
        eg.valid = true;
        if (eg.min_dist < min_dist)
            min_dist = eg.min_dist;
    }
    valid = true;
}

void Object::makePolygons4Render(Eigen::Vector2d obs, std::vector<std::vector<cv::Point>>& polygons) const{
    for (const Edge& eg: edges) {
        if (eg.valid == false) continue;
        polygons.emplace_back();
        std::vector<cv::Point>& back = polygons.back();
        back.emplace_back(obs.x(), obs.y());
        for (const Eigen::Vector3d& p: eg) {
            back.emplace_back(p.x(), p.y());
        }
    }
}

void Object::projectEdge2Edge(const Edge& src, const Eigen::Vector2d& obs, Edge& dst, HeapType& heap) {
    bool pop_back_flag = true;
    Eigen::Vector2d fpt = src.front().block<2, 1>(0, 0) - obs, ept = src.back().block<2, 1>(0, 0) - obs;
    double f_ang = src.front().z(), e_ang = src.back().z();
    bool head_in_range = dst.angleInRange(f_ang), end_in_range = dst.angleInRange(e_ang), mid_in_range = true;
    if (src.size() > 2) {
        size_t mid_id = src.size() / 2;
        mid_in_range = dst.angleInRange(src[mid_id].z());
    }
    std::vector<Eigen::Vector3d> tasks;
    if (head_in_range && end_in_range) {
        if (mid_in_range) {                     // 加edge 打断
            LOG_ERROR_STREAM("Breaking edge!");
            if (breakEdge(fpt, ept, obs, dst, heap) == true)
                return;
        }                          // 特殊情况
        tasks.emplace_back(fpt.x(), fpt.y(), f_ang);
        tasks.emplace_back(ept.x(), ept.y(), e_ang);
    } else if (head_in_range) {
        tasks.emplace_back(fpt.x(), fpt.y(), f_ang);
    } else if (end_in_range) {
        tasks.emplace_back(ept.x(), ept.y(), e_ang);
        pop_back_flag = false;          // pop_front
    } else {                            // SRC的两个端点不在DST范围内（有可能完全遮挡的）
        const Eigen::Vector2d dst_f = dst.front().block<2, 1>(0, 0) - obs, dst_b = dst.back().block<2, 1>(0, 0) - obs;
        int f_id = src.rotatedBinarySearch(dst.front().z());
        if (f_id > 0) {
            Eigen::Vector3d tmp;
            bool dst_closer = rangeSuitable(src[f_id - 1], src[f_id], dst_f, obs, tmp);
            if (dst_closer == true) return;
        }
        int e_id = src.rotatedBinarySearch(dst.back().z());
        if (e_id > 0) {
            Eigen::Vector3d tmp;
            bool dst_closer = rangeSuitable(src[e_id - 1], src[e_id], dst_b, obs, tmp);
            if (dst_closer == true) return;        // 投影边的range更大
        }
        if (f_id <= 0 && e_id <= 0) return;          // 反向光线
        if (f_id > 0 && e_id > 0) {
            dst.valid = false;
            return;
        }
        if (f_id > 0) {
            tasks.emplace_back(ept.x(), ept.y(), e_ang);
            pop_back_flag = false;          // pop_front
        } else {
            tasks.emplace_back(fpt.x(), fpt.y(), f_ang);
        }
    }
    int task_id = 0;
    for (const Eigen::Vector3d& task: tasks) {
        if (task_id > 0) pop_back_flag = false;
        double angle = task.z();
        Eigen::Vector2d beam = task.block<2, 1>(0, 0);
        int id = dst.rotatedBinarySearch(angle);
        if (task_id == 0 && id <= 0) {
            LOG_ERROR_STREAM("Task is 0 and id <= 0");
            throw 0;
        } else if (task_id > 0 && id <= 0) break;
        Eigen::Vector3d intersect = Eigen::Vector3d::Zero();
        if (rangeSuitable(dst[id - 1], dst[id], beam, obs, intersect) == false) {
            return;
        }
        if (pop_back_flag == true) {            // 删除大角度
            int delete_cnt = static_cast<int>(dst.size()) - id;
            for (int i = 0; i < delete_cnt; i++)
                dst.pop_back();
            dst.emplace_back(intersect);
        } else {
            for (int i = 0; i < id; i++) {
                dst.pop_front();
            }
            dst.emplace_front(intersect);
        }
        for (const Eigen::Vector3d& pt: dst) {
            double dist = (pt.block<2, 1>(0, 0) - obs).norm();
            if (dist < dst.min_dist) dst.min_dist = dist;
        }
        task_id ++;
    }
}

bool Object::breakEdge(Eigen::Vector2d b1, Eigen::Vector2d b2, Eigen::Vector2d obs, Edge& dst, HeapType& heap) {
    std::array<Eigen::Vector2d, 2> task = {b1, b2};
    std::vector<Eigen::Vector3d> crs;
    std::array<size_t, 2> ids = {0, 0};
    for (size_t i = 0; i < 2; i++) {
        const Eigen::Vector2d& beam = task[i];
        double this_angle = atan2(beam(1), beam(0));
        int id = dst.rotatedBinarySearch(this_angle);
        Eigen::Vector3d intersect = Eigen::Vector3d::Zero();
        if (rangeSuitable(dst[id - 1], dst[id], beam, obs, intersect) == false) return true;
        crs.emplace_back(intersect);
        ids[i] = dst.size() - static_cast<size_t>(id);
    }
    if (ids[0] < ids[1]) return false;
    Edge new_edge;
    size_t add_cnt = 0;
    for (Edge::const_reverse_iterator rit = dst.crbegin(); rit != dst.crend() && add_cnt < ids[1]; rit++, add_cnt++)
        new_edge.push_front(*rit);
    new_edge.push_front(crs[1]);
    new_edge.initWithObs(obs);
    for (size_t i = 0; i < ids[0]; i++)
        dst.pop_back();
    dst.push_back(crs[0]);
    for (const Eigen::Vector3d& pt: dst) {
        double dist = (pt.block<2, 1>(0, 0) - obs).norm();
        if (dist < dst.min_dist) dst.min_dist = dist;
    }
    edges.push_back(new_edge);
    heap.emplace(edges.size() - 1);              // 可能不太安全
    return true;
}

void Object::visualizeEdges(cv::Mat& src, cv::Point obs) const{
    if (valid == false)
        return;
    int cnt = -1;
    for (const Edge& eg: edges) {
        cnt++;
        if (eg.valid == false) continue;
        for (size_t i = 1; i < eg.size(); i++) {
            cv::line(src, cv::Point(eg[i - 1].x(), eg[i - 1].y()), cv::Point(eg[i].x(), eg[i].y()), cv::Scalar(0, 255, 255), 3);
        }
        char str[8];
        snprintf(str, 8, "%d:%lu", cnt, eg.size());
		cv::putText(src, str, cv::Point(eg.front().x(), eg.front().y()) + cv::Point(10, 10),
					cv::FONT_HERSHEY_PLAIN, 1, cv::Scalar(0, 255, 0));
        cv::circle(src, cv::Point(eg.front().x(), eg.front().y()), 3, cv::Scalar(255, 255, 0), -1);
        cv::circle(src, cv::Point(eg.back().x(), eg.back().y()), 3, cv::Scalar(255, 0, 255), -1);
    }
    cv::circle(src, obs, 4, cv::Scalar(0, 255, 0), -1);
}

Eigen::Vector3d Object::getIntersection(
    const Eigen::Vector2d& vec,
    const Eigen::Vector3d& _p1,
    const Eigen::Vector3d& _p2,
    const Eigen::Vector2d& obs
) {
    const Eigen::Vector2d p1 = _p1.block<2, 1>(0, 0), p2 = _p2.block<2, 1>(0, 0);
    const Eigen::Vector2d vec_line = p2 - p1;
    Eigen::Matrix2d A = Eigen::Matrix2d::Zero();
    A << -vec(1), vec(0), -vec_line(1), vec_line(0);
    double b1 = Eigen::RowVector2d(-vec(1), vec(0)) * obs;
    double b2 = Eigen::RowVector2d(-vec_line(1), vec_line(0)) * p1;
    const Eigen::Vector2d b(b1, b2);
    double det = A(0, 0) * A(1, 1) - A(0, 1) * A(1, 0);
    if (std::abs(det) < 1e-5)
        return _p1;
    const Eigen::Vector2d pt = A.inverse() * b, new_vec = pt - obs;
    double angle = atan2(new_vec(1), new_vec(0));
    Eigen::Vector3d result;
    result << pt, angle;
    return result;                   // 解交点
}

bool Object::rangeSuitable(
    const Eigen::Vector3d& p1,
    const Eigen::Vector3d& p2,
    const Eigen::Vector2d& beam,
    const Eigen::Vector2d& obs,
    Eigen::Vector3d& intersect
) const {
    intersect = getIntersection(beam, p1, p2, obs);
    double range = (intersect.block<2, 1>(0, 0) - obs).norm();
    return beam.norm() < range;
}

void Object::anglePostCheck(const Eigen::Vector2d& obs) {
    size_t max_size = edges.size();
    for (size_t j = 0; j < max_size; j++) {
        Edge& eg = edges[j];
        if (eg.size() < 4) continue;
        double front_angle = eg.front().z();
        int size_1 = static_cast<int>(eg.size()) - 1, id = 2;
        bool break_judge = false;
        for (; id < size_1; id++) {
            if (eg.angleInRange(front_angle, id, id + 1) == false) continue;
            break_judge = true;
            break;
        }
        if (break_judge == false) continue;
        int pop_cnt = size_1 - id;
        Edge new_edge;
        for (int i = 0; i < pop_cnt; i++) {
            Eigen::Vector3d pt = eg.back();
            eg.pop_back();
            new_edge.emplace_front(pt);
        }
        new_edge.emplace_front(eg.back());
        new_edge.initWithObs(obs);
        edges.push_back(new_edge);
    }
}


#pragma once
#include <vector>
#include <memory>
#include <Eigen/Core>
#include <Eigen/Sparse>

#include "MACGrid.hpp"

typedef Eigen::Triplet<double> T;

class Simulator
{
public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  Simulator(std::shared_ptr<MACGrid> grids, double &time);
  ~Simulator();

  void update();

private:
  void setEmitterVelocity();
  void addSource();

  void resetForce();
  void calVorticity();
  void addForce();
  void calPressure();
  void applyPressureTerm();
  void advectVelocity();
  void advectScalar();

  std::shared_ptr<MACGrid> m_grids;
  double &m_time;

  // solver
  std::vector<T> tripletList;
  Eigen::ConjugateGradient<Eigen::SparseMatrix<double>, Eigen::Lower | Eigen::Upper> ICCG;

  Eigen::SparseMatrix<double, Eigen::RowMajor> A;
  Eigen::VectorXd b;
  Eigen::VectorXd x;
};


/*
Part of the Fluid Corpus Manipulation Project (http://www.flucoma.org/)
Copyright 2017-2019 University of Huddersfield.
Licensed under the BSD-3 License.
See license.md file in the project root for full license information.
This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 725899).
*/

#pragma once

#include "STFT.hpp"
#include "../util/AlgorithmUtils.hpp"
#include "../util/FluidEigenMappings.hpp"
#include "../util/Munkres.hpp"
#include "../util/OptimalTransport.hpp"
#include "../util/RTPGHI.hpp"
#include "../../data/TensorTypes.hpp"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <vector>

namespace fluid {
namespace algorithm {

class NMFMorph
{

public:
  using MatrixXd = Eigen::MatrixXd;

  void init(RealMatrixView W1, RealMatrixView W2, RealMatrixView H,
            index winSize, index fftSize, index hopSize, bool assign)
  {
    using namespace Eigen;
    using namespace _impl;
    using namespace std;
    mInitialized = false;
    mW1 = asEigen<Matrix>(W1).transpose();
    mH = asEigen<Matrix>(H);
    MatrixXd tmpW2 = asEigen<Matrix>(W2).transpose();
    ArrayXXd cost = ArrayXXd::Zero(mW1.cols(), tmpW2.cols());
    if (assign)
    {
      for (index i = 0; i < mW1.cols(); i++)
      {
        for (index j = 0; j < tmpW2.cols(); j++)
        {
          OptimalTransport tmpOT;
          tmpOT.init(mW1.col(i), tmpW2.col(j));
          if(!tmpOT.initialized()) return;
          cost(i, j) = tmpOT.mDistance;
        }
      }
      Munkres munk;
      munk.init(mW1.cols(), tmpW2.cols());
      ArrayXi result = ArrayXi::Zero(mW1.cols());
      munk.process(cost, result);
      mW2 = MatrixXd::Zero(tmpW2.rows(), tmpW2.cols());
      for (index i = 0; i < result.size(); i++)
      { mW2.col(i) = tmpW2.col(result(i)); }
    }
    else
    {
      mW2 = asEigen<Matrix>(W2).transpose();
    }
    mWindowSize = winSize;
    mFFTSize = fftSize;
    mHopSize = hopSize;
    mRTPGHI.init(fftSize);

    index rank = mW1.cols();
    mOT = std::vector<OptimalTransport>(asUnsigned(rank));
    for (index i = 0; i < rank; i++) { mOT[asUnsigned(i)].init(mW1.col(i), mW2.col(i)); }
    mPos = 0;
    mInitialized = true;
  }

  bool initialized() { return mInitialized; }

  void processFrame(ComplexVectorView v, double interpolation)
  {
    using namespace Eigen;
    using namespace _impl;
    MatrixXd W = MatrixXd::Zero(mW1.rows(), mW1.cols());
    for (int i = 0; i < W.cols(); i++)
    {
      ArrayXd out = ArrayXd::Zero(mW2.rows());
      mOT[asUnsigned(i)].interpolate(interpolation, out);
      W.col(i) = out;
    }

    VectorXd       hFrame = mH.col(mPos);
    VectorXd       frame = W * hFrame;
    RealVectorView mag1 = asFluid(frame);
    mRTPGHI.processFrame(mag1, v, mWindowSize, mFFTSize, mHopSize, 1e-6);
    mPos = (mPos + 1) % mH.cols();
  }

private:
  MatrixXd                      mW1;
  MatrixXd                      mW2;
  MatrixXd                      mH;
  index                         mWindowSize;
  index                         mHopSize;
  index                         mFFTSize;
  RTPGHI                        mRTPGHI;
  std::vector<OptimalTransport> mOT;
  int                           mPos{0};
  bool                          mInitialized;
};
} // namespace algorithm
} // namespace fluid