/*!
* Copyright 2007  Technical University of Catalonia
*
* Use, modification and distribution is subject to the Boost Software
* License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
* http://www.boost.org/LICENSE_1_0.txt)
*
*  Authors: Dmitry Bufistov
*           Andrey Parfenov
*/
#include <boost/graph/howard_cycle_ratio.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/property_iter_range.hpp>
#include <boost/graph/random.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

using namespace boost;
typedef adjacency_list<listS, listS, directedS, property<vertex_index_t, int, property<boost::vertex_name_t, std::string> >,
        property<edge_weight_t, double, property<edge_weight2_t, double, property<edge_index_t, int> > > > grap_real_t;

template <typename TGraph>
void gen_rand_graph(TGraph& g, size_t nV, size_t nE)
{
        g.clear();
        boost::mt19937 rng;
        boost::generate_random_graph(g, nV, nE, rng, true, true);
        boost::uniform_real<> ur(-1,10);
        boost::variate_generator<boost::mt19937&, boost::uniform_real<> >       ew1rg(rng, ur);
        randomize_property<edge_weight_t>(g, ew1rg);
        boost::uniform_int<> uint(1,5);
        boost::variate_generator<boost::mt19937&, boost::uniform_int<> >        ew2rg(rng, uint);
        randomize_property<edge_weight2_t>(g, ew2rg);
}

int main(int argc, char* argv[])
{
        const double epsilon = 0.000000001;
        double min_cr, max_cr; ///Minimum and maximum cycle ratio
        typedef std::vector<graph_traits<grap_real_t>::edge_descriptor> ccReal_t;
        ccReal_t cc; ///For storing critical edges

        grap_real_t tgr;
        property_map<grap_real_t, vertex_index_t>::type vim = get(vertex_index, tgr);
        property_map<grap_real_t, edge_weight_t>::type ew1m = get(edge_weight, tgr);
        property_map<grap_real_t, edge_weight2_t>::type ew2m = ew2m;

        gen_rand_graph(tgr, 1000, 300000);
        std::cout << "Vertices number: " << num_vertices(tgr) << '\n';
        std::cout << "Edges number: " << num_edges(tgr) << '\n';
        int i = 0;
        BGL_FORALL_VERTICES(vd, tgr, grap_real_t) put(vertex_index, tgr, vd, i++); ///Initialize vertex index property
        boost::posix_time::ptime        st = boost::posix_time::microsec_clock::local_time();
        max_cr = maximum_cycle_ratio(tgr, get(vertex_index, tgr), get(edge_weight, tgr), get(edge_weight2, tgr));
        std::cout << "Maximum cycle ratio is " << max_cr << '\n';
        std::cout << "Run time of the maximum_cycle_ratio() is " << to_simple_string(boost::posix_time::microsec_clock::local_time() - st) << '\n';

        ///One way to get the "good" value of the plus_infinity parameter
        double pl_infnt = double(*std::max_element(get_property_iter_range(tgr, edge_weight).first, get_property_iter_range(tgr, edge_weight).second)) /
                *std::min_element(get_property_iter_range(tgr, edge_weight2).first, get_property_iter_range(tgr, edge_weight2).second);
        std::cout << "Set infinity for minimum_cycle_ratio() call to " << pl_infnt << '\n';
        i = 0;
        BGL_FORALL_EDGES(ed, tgr, grap_real_t) put(edge_index, tgr, ed, i++); ///Initialize edge index property
        min_cr = minimum_cycle_ratio(tgr, get(vertex_index, tgr), get(edge_weight, tgr), get(edge_weight2, tgr), get(edge_index, tgr), &cc, pl_infnt);
        std::cout << "Minimal cycle ratio is " << min_cr << '\n';
        std::pair<double, double> cr(.0,.0);
        std::cout << "\nCritical cycle is:\n";
        for (ccReal_t::iterator itr = cc.begin(); itr != cc.end(); ++itr)
        {
                cr.first += get(edge_weight, tgr, *itr); cr.second += get(edge_weight2, tgr, *itr);
                std::cout << "(" << get(vertex_index, tgr, source(*itr, tgr)) << "," << get(vertex_index, tgr, target(*itr, tgr)) << ") ";
        }
        std::cout << '\n';
        assert(std::abs(cr.first / cr.second - min_cr) < epsilon);

        return 0;
}


/*!
 * Distance.h
 *
 *  Created on: 2012/12/29
 *  \author Hiroki Sudo
 *
 */

#ifndef DISTANCE_H_
#define DISTANCE_H_

#include <vector>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/shared_ptr.hpp>
#include "Element.h"

/*!
 *
 */
namespace {
/*!
 * Distance
 */
namespace Distance {

/*!
 * ElementクラスのVectorインスタンスを引数に2つ取り、それぞれの編集距離(Levenstein距離)を計算します。
 * 返値は引数に渡されるVectorインスタンスの要素数で割られるため、0~1の値を取ります。
 * 値が大きければ距離は離れていて、小さければ距離が近いと言う意味になります。計算量はO((n*m)^(1/2))
 */
template<class E>
double levenstein(std::vector<E> ex1, std::vector<E> ex2) {

	boost::numeric::ublas::matrix<int> matrix(0, 0);
	int col_size, row_size;
	int cost_delta, cost1, cost2, cost3, cost;

	if (ex1.size() == 0 && ex2.size() != 0)
		return ex2.size();
	if (ex1.size() != 0 && ex2.size() == 0)
		return ex1.size();
	if (ex1.size() == 0 && ex2.size() == 0)
		return 0;

	row_size = ex1.size() + 1;
	col_size = ex2.size() + 1;

	matrix.resize(row_size, col_size);

	for (int j = 0; j < col_size; j++) {
		matrix(0, j) = j;
	}
	for (int i = 0; i < row_size; i++) {
		matrix(i, 0) = i;
	}

	for (int index_ex1 = 0; index_ex1 < ex1.size(); index_ex1++) {
		for (int index_ex2 = 0; index_ex2 < ex2.size(); index_ex2++) {

			if (ex1[index_ex1].is_sym() && ex2[index_ex2].is_sym()) {
				cost_delta = ex1[index_ex1] == ex2[index_ex2] ? 0 : 1;
			} else if (ex1[index_ex1].is_cat() && ex2[index_ex2].is_cat()) {
				cost_delta = ex1[index_ex1].obj == ex2[index_ex2].obj ? 0 : 1;
			} else {
				cost_delta = 1;
			}

			int x, y;
			y = index_ex1 + 1;
			x = index_ex2 + 1;
			cost1 = matrix(y - 1, x) + 1;
			cost2 = matrix(y, x - 1) + 1;
			cost3 = matrix(y - 1, x - 1) + cost_delta;

			cost = cost1;
			cost = cost2 < cost ? cost2 : cost;
			cost = cost3 < cost ? cost3 : cost;

			matrix(y, x) = cost;
		}
	}

	int len;
	if (row_size > col_size)
		len = row_size;
	else
		len = col_size;

	return (matrix(matrix.size1() - 1, matrix.size2() - 1))
			/ (static_cast<double>(len));
}

/*!
 * Vectorインスタンスの要素数で割らないため，値は0~無限となるLevenstein距離
 */
template<class E>
double levenstein2(std::vector<E> ex1, std::vector<E> ex2) {

	boost::numeric::ublas::matrix<int> matrix(0, 0);
	int col_size, row_size;
	int cost_delta, cost1, cost2, cost3, cost;

	if (ex1.size() == 0 && ex2.size() != 0)
		return ex2.size();
	if (ex1.size() != 0 && ex2.size() == 0)
		return ex1.size();
	if (ex1.size() == 0 && ex2.size() == 0)
		return 0;

	row_size = ex1.size() + 1;
	col_size = ex2.size() + 1;

	matrix.resize(row_size, col_size);

	for (int j = 0; j < col_size; j++) {
		matrix(0, j) = j;
	}
	for (int i = 0; i < row_size; i++) {
		matrix(i, 0) = i;
	}

	for (int index_ex1 = 0; index_ex1 < ex1.size(); index_ex1++) {
		for (int index_ex2 = 0; index_ex2 < ex2.size(); index_ex2++) {

			if (ex1[index_ex1].is_sym() && ex2[index_ex2].is_sym()) {
				cost_delta = ex1[index_ex1] == ex2[index_ex2] ? 0 : 1;
			} else if (ex1[index_ex1].is_cat() && ex2[index_ex2].is_cat()) {
				cost_delta = ex1[index_ex1].obj == ex2[index_ex2].obj ? 0 : 1;
			} else {
				cost_delta = 1;
			}

			int x, y;
			y = index_ex1 + 1;
			x = index_ex2 + 1;
			cost1 = matrix(y - 1, x) + 1;
			cost2 = matrix(y, x - 1) + 1;
			cost3 = matrix(y - 1, x - 1) + cost_delta;

			cost = cost1;
			cost = cost2 < cost ? cost2 : cost;
			cost = cost3 < cost ? cost3 : cost;

			matrix(y, x) = cost;
		}
	}

	int len;
	if (row_size > col_size)
		len = row_size;
	else
		len = col_size;

	return (matrix(matrix.size1() - 1, matrix.size2() - 1));
}

template <class E>
int snake(int k, int y, std::vector<E >& ary1, std::vector<E >& ary2, int m, int n) {
	int x;
	x = y - k;
	while (x < m && y < n && ary1[x] == ary2[y]) {
		x++;
		y++;
	}

	return y;
}
/*!
 * O(np)法の実装
 */
template<class E>
double onp_lv(std::vector<E>& ary1, std::vector<E>& ary2, double limit = -1) {
	std::vector<E> temp;
	std::vector<int> fp;
	int size = 0;
	int delta = 0;
	int offset = 0;
	int n = 0;
	int m = 0;
	int p;
	int k;
	int i;

	if (ary1.size() > ary2.size()) {
		temp = ary1;
		ary1 = ary2;
		ary2 = temp;
	}

	m = ary1.size();
	n = ary2.size();

	offset = m + 1;
	delta = n - m;
	size = m + n + 3;

	if (limit == -1) {
		limit = 2 * m + delta;
	}

	fp = std::vector<int>(size, -1);
	p = -1;
	k = 0;

	while (fp[delta + offset] < n) {
		p++;

		if (2 * p + delta >= (limit) * ary2.size() ) {
			//std::cerr << "Cutted" << std::endl;
			return 1.0;
		}

		k = -p;
		while (k < delta) {
			fp[k + offset] = fp[k - 1 + offset] + 1 > fp[k + 1 + offset] ?
				snake(k, fp[k - 1 + offset] + 1, ary1, ary2, m, n) :
				snake(k, fp[k + 1 + offset], ary1, ary2, m, n);
			k++;
		}

		k = delta + p;
		while (k > delta) {
			fp[k + offset] = fp[k - 1 + offset] + 1 > fp[k + 1 + offset] ?
				snake(k, fp[k - 1 + offset] + 1, ary1, ary2, m, n) :
				snake(k, fp[k + 1 + offset], ary1, ary2, m, n);
			k--;
		}

		fp[delta + offset] = fp[delta - 1 + offset] + 1 > fp[delta + 1 + offset] ?
				snake(delta, fp[delta - 1 + offset] + 1, ary1, ary2, m,	n) :
				snake(delta, fp[delta + 1 + offset], ary1, ary2, m, n);
	}

	return (delta + 2 * p)/(static_cast<double>(ary2.size()));
}

template<class E>
double hamming(std::vector<E> ex1, std::vector<E> ex2){

    if(ex1.size()!=ex2.size()){
        std::cerr <<  "Happened error on hamming distance method." << std::endl;
        std::exit(0);
    }
    typename std::vector<E>::iterator ex1_it;

    ex1_it = ex1.begin();
    int cnt = 0;
    double ham_sum=0;
    for(; ex1_it != ex1.end(); ex1_it++) {

        if((*ex1_it)!=ex2[cnt])
            ham_sum+=1;

        cnt+=1;
    }
    return (ham_sum/(double)cnt);
}

} /*Distance*/
}
#endif /* DISTANCE_H_ */


/*
* Copyright (c) by CryptoLab inc.
* This program is licensed under a
* Creative Commons Attribution-NonCommercial 3.0 Unported License.
* You should have received a copy of the license along with this
* work.  If not, see <http://creativecommons.org/licenses/by-nc/3.0/>.
*/
#include "TestScheme.h"

#include <NTL/BasicThreadPool.h>
#include <NTL/ZZ.h>
#include <complex>

#include "Ciphertext.h"
#include "EvaluatorUtils.h"
#include "Ring.h"
#include "Scheme.h"
#include "SchemeAlgo.h"
#include "SecretKey.h"
#include "StringUtils.h"
#include "TimeUtils.h"
#include "SerializationUtils.h"

using namespace std;
using namespace NTL;

//----------------------------------------------------------------------------------
//   STANDARD TESTS
//----------------------------------------------------------------------------------

void TestScheme::testEncrypt(long logq, long logp, long logn) {
	cout << "!!! START TEST ENCRYPT !!!" << endl;
	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	long n = (1 << logn);
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	Ciphertext cipher;

	timeutils.start("Encrypt");
	scheme.encrypt(cipher, mvec, n, logp, logq);
	timeutils.stop("Encrypt");

	timeutils.start("Decrypt");
	complex<double>* dvec = scheme.decrypt(secretKey, cipher);
	timeutils.stop("Decrypt");

	StringUtils::compare(mvec, dvec, n, "val");

	cout << "!!! END TEST ENCRYPT !!!" << endl;
}

void TestScheme::testEncryptSingle(long logq, long logp) {
	cout << "!!! START TEST ENCRYPT SINGLE !!!" << endl;
	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	complex<double> mval = EvaluatorUtils::randomComplex();
	Ciphertext cipher;

	timeutils.start("Encrypt Single");
	scheme.encryptSingle(cipher, mval, logp, logq);
	timeutils.stop("Encrypt Single");

	complex<double> dval = scheme.decryptSingle(secretKey, cipher);

	StringUtils::compare(mval, dval, "val");

	cout << "!!! END TEST ENCRYPT SINGLE !!!" << endl;
}

void TestScheme::testAdd(long logq, long logp, long logn) {
	cout << "!!! START TEST ADD !!!" << endl;

	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	long n = (1 << logn);
	complex<double>* mvec1 = EvaluatorUtils::randomComplexArray(n);
	complex<double>* mvec2 = EvaluatorUtils::randomComplexArray(n);
	complex<double>* madd = new complex<double>[n];

	for(long i = 0; i < n; i++) {
		madd[i] = mvec1[i] + mvec2[i];
	}

	Ciphertext cipher1, cipher2;
	scheme.encrypt(cipher1, mvec1, n, logp, logq);
	scheme.encrypt(cipher2, mvec2, n, logp, logq);

	timeutils.start("Addition");
	scheme.multAndEqual(cipher1, cipher2);
	timeutils.stop("Addition");

	complex<double>* dadd = scheme.decrypt(secretKey, cipher1);

	//StringUtils::showVec(mvec1, n);
	//StringUtils::showVec(mvec2, n);
	//<StringUtils::showVec(madd, n);

	StringUtils::compare(madd, dadd, n, "add");

	cout << "!!! END TEST ADD !!!" << endl;
}

void TestScheme::testMult(long logq, long logp, long logn) {
	cout << "!!! START TEST MULT !!!" << endl;

	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	long n = (1 << logn);
	complex<double>* mvec1 = EvaluatorUtils::randomComplexArray(n);
	complex<double>* mvec2 = EvaluatorUtils::randomComplexArray(n);
	complex<double>* mmult = new complex<double>[n];

	StringUtils::showVec(mvec1, n);
	StringUtils::showVec(mvec2, n);

	for(long i = 0; i < n; i++) {
		mmult[i] = mvec1[i] * mvec2[i];
	}

	Ciphertext cipher1, cipher2;
	scheme.encrypt(cipher1, mvec1, n, logp, logq);
	scheme.encrypt(cipher2, mvec2, n, logp, logq);

	timeutils.start("Multiplication");
	scheme.multAndEqual(cipher1, cipher2);
	timeutils.stop("Multiplication");

	complex<double>* dmult = scheme.decrypt(secretKey, cipher1);

	StringUtils::showVec(mmult, n);

	StringUtils::compare(mmult, dmult, n, "mult");

	cout << "!!! END TEST MULT !!!" << endl;
}

void TestScheme::testimult(long logq, long logp, long logn) {
	cout << "!!! START TEST i MULTIPLICATION !!!" << endl;

	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	long n = (1 << logn);

	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	complex<double>* imvec = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		imvec[i].real(-mvec[i].imag());
		imvec[i].imag(mvec[i].real());
	}

	StringUtils::showVec(mvec, n);
	StringUtils::showVec(imvec, n);

	Ciphertext cipher;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Multiplication by i");
	scheme.imultAndEqual(cipher);
	timeutils.stop("Multiplication by i");

	complex<double>* idvec = scheme.decrypt(secretKey, cipher);

	StringUtils::compare(imvec, idvec, n, "imult");

	cout << "!!! END TEST i MULTIPLICATION !!!" << endl;
}

//----------------------------------------------------------------------------------
//   ROTATE & CONJUGATE
//----------------------------------------------------------------------------------

void TestScheme::testRotateFast(long logq, long logp, long logn, long logr) {
	cout << "!!! START TEST ROTATE FAST !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	long n = (1 << logn);
	long r = (1 << logr);
	scheme.addLeftRotKey(secretKey, r);
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	Ciphertext cipher;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Left Rotate Fast");
	scheme.leftRotateFastAndEqual(cipher, r);
	timeutils.stop("Left Rotate Fast");

	complex<double>* dvec = scheme.decrypt(secretKey, cipher);

	StringUtils::showVec(mvec, n);

	EvaluatorUtils::leftRotateAndEqual(mvec, n, r);

	StringUtils::showVec(mvec, n);
	//StringUtils::showVec(mvec, n);
	//StringUtils::showVec(dvec, n);
	StringUtils::compare(mvec, dvec, n, "rot");

	cout << "!!! END TEST ROTATE BY POWER OF 2 BATCH !!!" << endl;
}

void TestScheme::testConjugate(long logq, long logp, long logn) {
	cout << "!!! START TEST CONJUGATE !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	scheme.addConjKey(secretKey);

	long n = (1 << logn);

	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	complex<double>* mvecconj = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mvecconj[i] = conj(mvec[i]);
	}

	Ciphertext cipher;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Conjugate");
	scheme.conjugateAndEqual(cipher);
	timeutils.stop("Conjugate");

	complex<double>* dvecconj = scheme.decrypt(secretKey, cipher);
	StringUtils::showVec(mvec, n);
	StringUtils::showVec(mvecconj, n);
	StringUtils::compare(mvecconj, dvecconj, n, "conj");

	cout << "!!! END TEST CONJUGATE !!!" << endl;
}

//----------------------------------------------------------------------------------
//   POWER & PRODUCT TESTS
//----------------------------------------------------------------------------------

void TestScheme::testPowerOf2(long logq, long logp, long logn, long logdeg) {
	cout << "!!! START TEST POWER OF 2 !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	long degree = 1 << logdeg;
	complex<double>* mvec = new complex<double>[n];
	//cout << "mvec and mpow 2^" << logdeg << endl;
	complex<double>* mpow = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mvec[i] = EvaluatorUtils::randomCircle();
		mpow[i] = pow(mvec[i], degree);
	}
	//StringUtils::showVec(mvec, n);
	//StringUtils::showVec(mpow, n);

	Ciphertext cipher, cpow;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Power of 2");
	algo.powerOf2(cpow, cipher, logp, logdeg);
	timeutils.stop("Power of 2");

	complex<double>* dpow = scheme.decrypt(secretKey, cpow);
	StringUtils::compare(mpow, dpow, n, "pow2");

	cout << "!!! END TEST POWER OF 2 !!!" << endl;
}

//-----------------------------------------

void TestScheme::testPower(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST POWER !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomCircleArray(n);
	cout << "mvec and mpow ^"<< degree << endl;
	StringUtils::showVec(mvec, n);

	complex<double>* mpow = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mpow[i] = pow(mvec[i], degree);
	}
	StringUtils::showVec(mpow, n);

	Ciphertext cipher, cpow;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Power");
	algo.power(cpow, cipher, logp, degree);
	timeutils.stop("Power");

	complex<double>* dpow = scheme.decrypt(secretKey, cpow);
	StringUtils::compare(mpow, dpow, n, "pow");

	cout << "!!! END TEST POWER !!!" << endl;
}

//----------------------------------------------------------------------------------
//   FUNCTION TESTS
//----------------------------------------------------------------------------------

void TestScheme::testInverse(long logq, long logp, long logn, long steps) {
	cout << "!!! START TEST INVERSE !!!" << endl;

	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomCircleArray(n, 0.1);
	//cout << "mvec and minv" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* minv = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		minv[i] = 1. / mvec[i];
	}
	//StringUtils::showVec(minv, n);

	Ciphertext cipher, cinv;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start("Inverse");
	algo.inverse(cinv, cipher, logp, steps);
	timeutils.stop("Inverse");

	complex<double>* dinv = scheme.decrypt(secretKey, cinv);
	StringUtils::compare(minv, dinv, n, "inv");

	cout << "!!! END TEST INVERSE !!!" << endl;
}

void TestScheme::testLogarithm(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST LOGARITHM !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n, 0.1);
	//cout << "mvec and mlog" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* mlog = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mlog[i] = log(mvec[i] + 1.);
	}
	//StringUtils::showVec(mlog, n);

	Ciphertext cipher, clog;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(LOGARITHM);
	algo.function(clog, cipher, LOGARITHM, logp, degree);
	timeutils.stop(LOGARITHM);

	complex<double>* dlog = scheme.decrypt(secretKey, clog);
	StringUtils::compare(mlog, dlog, n, LOGARITHM);

	cout << "!!! END TEST LOGARITHM !!!" << endl;
}

void TestScheme::testExponent(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST EXPONENT !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	//cout << "mvec and mexp" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* mexp = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mexp[i] = exp(mvec[i]);
	}
	//StringUtils::showVec(mexp, n);

	Ciphertext cipher, cexp;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(EXPONENT);
	algo.function(cexp, cipher, EXPONENT, logp, degree);
	timeutils.stop(EXPONENT);

	complex<double>* dexp = scheme.decrypt(secretKey, cexp);
	StringUtils::compare(mexp, dexp, n, EXPONENT);

	cout << "!!! END TEST EXPONENT !!!" << endl;
}

void TestScheme::testSqrt4(long logq, long logp, long logn) {
	cout << "!!! START TEST SQRT4 !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n, 20); // get a rdm vector in [0, 20]^n
	complex<double>* msqrt = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mvec[i]  = complex<double>(mvec[i].real(), 0);
		msqrt[i] = complex<double>(sqrt(abs(mvec[i].real())), sqrt(abs(mvec[i].imag())));
	}
	StringUtils::showVec(mvec, n);

	Ciphertext cipher, csqrt;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(SQRT4);
	algo.function(csqrt, cipher, SQRT4, logp, 5);
	timeutils.stop(SQRT4);

	complex<double>* dsqrt = scheme.decrypt(secretKey, csqrt);
	//StringUtils::showVec(dsqrt, n);
	StringUtils::compare(msqrt, dsqrt, n, SQRT4);

	cout << "!!! END TEST SQRT4 !!!" << endl;
}

void TestScheme::testSqrt8(long logq, long logp, long logn) {
	cout << "!!! START TEST SQRT8 !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n, 12); // get a rdm vector in [0, 12]^n
	complex<double>* msqrt = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mvec[i]  = complex<double>(mvec[i].real(), 0.);
		msqrt[i] = complex<double>(sqrt(abs(mvec[i].real())), sqrt(abs(mvec[i].imag())));
	}

	Ciphertext cipher, csqrt;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(SQRT8);
	algo.function(csqrt, cipher, SQRT8, logp, 9);
	timeutils.stop(SQRT8);

	complex<double>* dsqrt = scheme.decrypt(secretKey, csqrt);
	StringUtils::compare(msqrt, dsqrt, n, SQRT8);

	cout << "!!! END TEST SQRT8 !!!" << endl;
}

void TestScheme::testAllSqrt(long logq, long logp, long logn) {
	cout << "!!! START TEST SQRT8 !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n, 10); // get a rdm vector in [0, 12]^n
	complex<double>* msqrt = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mvec[i]  = complex<double>(mvec[i].real(), 0.);
		msqrt[i] = complex<double>(sqrt(abs(mvec[i].real())), sqrt(abs(mvec[i].imag())));
	}

	Ciphertext cipher, csqrt, csqrt2;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(SQRT8);
	algo.function(csqrt, cipher, SQRT8, logp, 9);
	timeutils.stop(SQRT8);

	timeutils.start(SQRT8);
	Ciphertext cipher2;

	scheme.square(cipher2, cipher);
	//cout << "c2_logQ = "<<cipher2.logq << endl;

	scheme.reScaleByAndEqual(cipher2, logp); // cipher2.logq : logq - logp
	//cout << "c2_logQ = "<<cipher2.logq << endl;

	Ciphertext cipher4;
	scheme.square(cipher4, cipher2);
	//cout << "c4_logQ = "<<cipher4.logq << endl;
	scheme.reScaleByAndEqual(cipher4, logp); // cipher4.logq : logq -2logp
	//cout << "c4_logQ = "<<cipher4.logq << endl;

	Ciphertext cipher8;
	scheme.square(cipher8, cipher4);
	//cout << "c8_logQ = "<<cipher8.logq << endl;
	scheme.reScaleByAndEqual(cipher8, logp); // cipher4.logq : logq -2logp
	//cout << "c8_logQ = "<<cipher8.logq << endl;

	RR c;
	c =  45./144; //a0/a1
	Ciphertext cipher01;
	scheme.addConst(cipher01, cipher, c, logp);
	//cout << "cip1_logQ = "<<cipher01.logq << endl;

	c = 288./323; //a1
	scheme.multByConstAndEqual(cipher01, c, logp);
	//cout << "cip1_logQ = "<<cipher01.logq << endl;
	scheme.reScaleByAndEqual(cipher01, logp);
	//cout << "cip1_logQ = "<<cipher01.logq << endl;

	c = -125./24; //a2/a3
	Ciphertext cipher23;
	scheme.addConst(cipher23, cipher, c, logp);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;

	c = 44352./1009375;//a3
	scheme.multByConstAndEqual(cipher23, c, logp);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;
	scheme.reScaleByAndEqual(cipher23, logp);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;

	scheme.multAndEqual(cipher23, cipher2);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;
	scheme.reScaleByAndEqual(cipher23, logp);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;

	scheme.addAndEqual(cipher23, cipher01);
	//cout << "cip23_logQ = "<<cipher23.logq << endl;

	c = -5625./392; //a4/a5
	Ciphertext cipher45;
	scheme.addConst(cipher45, cipher, c, logp);
	//cout << "cip45_logQ = "<<cipher45.logq << endl;

	c = 224224./630859375; //a5
	scheme.multByConstAndEqual(cipher45, c, logp);
	//cout << "cip45_logQ = "<<cipher45.logq << endl;
	scheme.reScaleByAndEqual(cipher45, logp);
	//cout << "cip45_logQ = "<<cipher45.logq << endl;

	c = -15925./352; // a6/a7
	scheme.addConstAndEqual(cipher, c, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	c = 25344./78857421875; // a7
	scheme.multByConstAndEqual(cipher, c, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;
	scheme.reScaleByAndEqual(cipher, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	scheme.multAndEqual(cipher, cipher2);
	//cout << "cip_logQ = "<<cipher.logq << endl;
	scheme.reScaleByAndEqual(cipher, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	scheme.modDownByAndEqual(cipher45, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;
	scheme.addAndEqual(cipher, cipher45);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	scheme.multAndEqual(cipher, cipher4);
	//cout << "cip_logQ = "<<cipher.logq << endl;
	scheme.reScaleByAndEqual(cipher, logp);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	scheme.modDownByAndEqual(cipher23, logp);
	scheme.addAndEqual(cipher, cipher23);
	//cout << "cip_logQ = "<<cipher.logq << endl;

	c = -1716./579833984375; //a8
	scheme.multByConstAndEqual(cipher8, c, logp);
	scheme.reScaleByAndEqual(cipher8, logp);
	scheme.addAndEqual(cipher, cipher8);
	timeutils.stop(SQRT8);

	complex<double>* dsqrt = scheme.decrypt(secretKey, csqrt);
	complex<double>* dsqrt2 = scheme.decrypt(secretKey, cipher);
	StringUtils::compare(msqrt, dsqrt, n, SQRT8);
	cout<<" "<<endl;
	StringUtils::compare(msqrt, dsqrt2, n, SQRT8);

	cout << "!!! END TEST SQRT8 !!!" << endl;
}

void TestScheme::testExponentLazy(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST EXPONENT LAZY !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	//cout << "mvec and mexp" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* mexp = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		mexp[i] = exp(mvec[i]);
	}
	//StringUtils::showVec(mexp, n);

	Ciphertext cipher, cexp;
	scheme.encrypt(cipher, mvec, n, logp, logQ);

	timeutils.start(EXPONENT + " lazy");
	algo.functionLazy(cexp, cipher, EXPONENT, logp, degree);
	timeutils.stop(EXPONENT + " lazy");

	complex<double>* dexp = scheme.decrypt(secretKey, cexp);
	StringUtils::compare(mexp, dexp, n, EXPONENT);

	cout << "!!! END TEST EXPONENT LAZY !!!" << endl;
}

//-----------------------------------------

void TestScheme::testSigmoid(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST SIGMOID !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;

	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	//cout << "mvec and msig" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* msig = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		msig[i] = exp(mvec[i]) / (1. + exp(mvec[i]));
	}
	//StringUtils::showVec(msig, n);

	Ciphertext cipher, csig;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(SIGMOID);
	algo.function(csig, cipher, SIGMOID, logp, degree);
	timeutils.stop(SIGMOID);

	complex<double>* dsig = scheme.decrypt(secretKey, csig);
	StringUtils::compare(msig, dsig, n, SIGMOID);

	cout << "!!! END TEST SIGMOID !!!" << endl;
}

void TestScheme::testSigmoidLazy(long logq, long logp, long logn, long degree) {
	cout << "!!! START TEST SIGMOID LAZY !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = 1 << logn;
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(n);
	//cout << "mvec and msig" << endl;
	//StringUtils::showVec(mvec, n);
	complex<double>* msig = new complex<double>[n];
	for (long i = 0; i < n; ++i) {
		msig[i] = exp(mvec[i]) / (1. + exp(mvec[i]));
	}
	//StringUtils::showVec(msig, n);

	Ciphertext cipher, csig;
	scheme.encrypt(cipher, mvec, n, logp, logq);

	timeutils.start(SIGMOID + " lazy");
	algo.functionLazy(csig, cipher, SIGMOID, logp, degree);
	timeutils.stop(SIGMOID + " lazy");

	complex<double>* dsig = scheme.decrypt(secretKey, csig);
	StringUtils::compare(msig, dsig, n, SIGMOID);

	cout << "!!! END TEST SIGMOID LAZY !!!" << endl;
}

void TestScheme::testDistance(long logq, long logp, long logn, complex<double>* mvec, complex<double>* mbin, complex<double>* mDes, int points, int norm, double realDistance) {
	cout << "!!! START DISTANCE !!!" << endl;
	srand(time(NULL));
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);
	SchemeAlgo algo(scheme);

	long n = (1 << logn);
	long r = (1 << 0);
	scheme.addLeftRotKey(secretKey, r);
	long r2 = n/2;
	scheme.addLeftRotKey(secretKey, r2);

	timeutils.start(ENC);
	//* Encrypt the plaintext
	Ciphertext cipher;
	scheme.encrypt(cipher, mvec, n, logp, logq);
	timeutils.stop(ENC);

	timeutils.start(DISTANCE);
	//* Create the difference vector
	Ciphertext diff;
	scheme.leftRotateFast(diff, cipher, r);
	scheme.subAndEqual(cipher, diff);

	//* Distroy elements in slot point and n/2+point
	Ciphertext cbin;
	scheme.encrypt(cbin, mbin, n, logp, logq);//encode(plain, vals, n, logp, logq);
	scheme.multAndEqual(cipher, cbin);
	scheme.reScaleByAndEqual(cipher, logp);

	Plaintext plnBin;

	//* Square each slot
	scheme.squareAndEqual(cipher);
	scheme.reScaleByAndEqual(cipher, logp);

	//* Sum subdistances
	Ciphertext subd;
	scheme.leftRotateFast(subd, cipher, r2);
	scheme.addAndEqual(cipher, subd);
	complex<double>* tmp10 = scheme.decrypt(secretKey, cipher);

	// Check range for square root
	for (long i = 0; i < points-1; ++i) {
			if (tmp10[i].real()>9){
			cout<<"distance "<<i<<" out of sqrt range : "<<tmp10[i].real()<<endl;
		}
			if (tmp10[i].real()<0.2){
			cout<<"distance "<<i<<" out of sqrt range inf : "<<tmp10[i].real()<<endl;
		}
	}

	//* Take square root of each slot
	Ciphertext csqrt;
	algo.function(csqrt, cipher, SQRT8, logp, 9); //SQRT9

	///* Show The subdistance vector
	complex<double>* dm = scheme.decrypt(secretKey, csqrt);

	///* Remove non necessary slots before summation
	Ciphertext cDes;
	scheme.encrypt(cDes, mDes, n, logp, logq);
	scheme.multAndEqual(csqrt, cDes);
	scheme.reScaleByAndEqual(csqrt, logp);

	///* Sum to slot 0
	for (long j =0; j< log(points-1)/log(2); j++) {
		Ciphertext rot;
		scheme.addLeftRotKey(secretKey, pow(2,j));
		scheme.leftRotateFast(rot, csqrt, pow(2,j));
		scheme.addAndEqual(csqrt, rot);
		//cout << "sum logq = "<<csqrt.logq << endl;
	}
	timeutils.stop(DISTANCE);

	//cout << "sum logq = "<<csqrt.logq << endl;

	// DECRYPTION
	timeutils.start(DEC);
	///* Decrypt
	complex<double>* decDir = scheme.decrypt(secretKey, csqrt);
	timeutils.stop(DEC);
	cout<<"distance = "<<decDir[0].real()*norm<<endl;

	double distance = 0;
	for (long i = 0; i < points-1; ++i) {
		distance += dm[i].real();
	}

	// Print for Latex
	// log n  | logq  |  logp  | RealDist | Distance  | error | t_i
	double result = decDir[0].real()*norm;
	cout<< logn <<" & "<< logq <<" & "<< logp <<" & "<< points <<" & "<<realDistance <<" & "<< decDir[0].real()*norm <<" & "<< distance*norm <<" & "<< floor((result*100./realDistance -100)*100)/100;//<<" & "<< floor(tenc/10)/100 <<" & "<< floor(tdist/10)/100 <<" & " << floor(tdec)/1000<<" & "<<counter;

}

void TestScheme::testWriteAndRead(long logq, long logp, long logSlots) {
	cout << "!!! START TEST WRITE AND READ !!!" << endl;

	cout << "!!! END TEST WRITE AND READ !!!" << endl;
}

void TestScheme::testBootstrap(long logq, long logQ, long logp, long logSlots, long logT) {
	/**
	 * Testing bootstrapping procedure for single real value
	 * number of modulus bits up: depends on parameters
	 * @param[in] logN: input parameter for Params class
	 * @param[in] logq: log of initial modulus
	 * @param[in] logQ: input parameter for Params class
	 * @param[in] logSlots: log of number of slots
	 * @param[in] nu: auxiliary parameter, corresonds to message bits (message bits is logq - nu)
	 * @param[in] logT: auxiliary parameter, corresponds to number of iterations in removeIpart (num of iterations is logI + logT)
	 * testBootstrap(long logq, long logQ, long logp, long logn, long logT); as of Nov18

	 */

	cout << "!!! START TEST BOOTSTRAP !!!" << endl;

	srand(time(NULL));
	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	timeutils.start("Key generating");
	scheme.addBootKey(secretKey, logSlots, logq + 4);
	timeutils.stop("Key generated");

	long slots = (1 << logSlots);
	complex<double>* mvec = EvaluatorUtils::randomComplexArray(slots);

	Ciphertext cipher;
	scheme.encrypt(cipher, mvec, slots, logp, logq);

	cout << "cipher logq before: " << cipher.logq << endl;

	scheme.modDownToAndEqual(cipher, logq);
	scheme.normalizeAndEqual(cipher);
	cipher.logq = logQ;
	cipher.logp = logq + 4;

	Ciphertext rot;
	timeutils.start("SubSum");
	for (long i = logSlots; i < logNh; ++i) {
		scheme.leftRotateFast(rot, cipher, (1 << i));
		scheme.addAndEqual(cipher, rot);
	}
	scheme.divByPo2AndEqual(cipher, logNh);
	timeutils.stop("SubSum");

	timeutils.start("CoeffToSlot");
	scheme.coeffToSlotAndEqual(cipher); // Issue here with default parameters
	timeutils.stop("CoeffToSlot");

	timeutils.start("EvalExp");
	scheme.evalExpAndEqual(cipher, logT);
	timeutils.stop("EvalExp");

	timeutils.start("SlotToCoeff");
	scheme.slotToCoeffAndEqual(cipher);
	timeutils.stop("SlotToCoeff");

	cipher.logp = logp;
	cout << "cipher logq after: " << cipher.logq << endl;

	complex<double>* dvec = scheme.decrypt(secretKey, cipher);

	StringUtils::compare(mvec, dvec, slots, "boot");

	cout << "!!! END TEST BOOTSRTAP !!!" << endl;
}

void TestScheme::testBootstrapSingleReal(long logq, long logQ, long logp, long logT) {
	cout << "!!! START TEST BOOTSTRAP SINGLE REAL !!!" << endl;

	srand(time(NULL));
//	SetNumThreads(8);
	TimeUtils timeutils;
	Ring ring;
	SecretKey secretKey(ring);
	Scheme scheme(secretKey, ring);

	timeutils.start("Key generating");
	scheme.addBootKey(secretKey, 0, logq + 4);
	timeutils.stop("Key generated");

	cout << "key ok" << endl;

	double mval = EvaluatorUtils::randomReal();

	Ciphertext cipher;
	scheme.encryptSingle(cipher, mval, logp, logq);

	cout << "cipher logq before: " << cipher.logq << endl;
	scheme.modDownToAndEqual(cipher, logq);
	scheme.normalizeAndEqual(cipher);
	cipher.logq = logQ;

	Ciphertext rot, cconj;
	timeutils.start("SubSum");
	for (long i = 0; i < logNh; ++i) {
		scheme.leftRotateFast(rot, cipher, 1 << i);
		scheme.addAndEqual(cipher, rot);
	}
	scheme.conjugate(cconj, cipher);
	scheme.addAndEqual(cipher, cconj);
	scheme.divByPo2AndEqual(cipher, logN);
	timeutils.stop("SubSum");

	timeutils.start("EvalExp");
	scheme.evalExpAndEqual(cipher, logT);
	timeutils.stop("EvalExp");

	cout << "cipher logq after: " << cipher.logq << endl;

	cipher.logp = logp;
	complex<double> dval = scheme.decryptSingle(secretKey, cipher);

	StringUtils::compare(mval, dval.real(), "boot");

	cout << "!!! END TEST BOOTSRTAP SINGLE REAL !!!" << endl;
}

void TestScheme::test() {
}


#include <tiny_math_types.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

typedef tiny::MathTypes<double>  MT;
typedef MT::value_traits         VT;

BOOST_AUTO_TEST_SUITE(tiny_degrees_and_radians);

BOOST_AUTO_TEST_CASE(conversion)
{
  double const degrees180 = 180.0;
  double const radians180 = VT::convert_to_radians(degrees180);

  BOOST_CHECK_CLOSE( radians180, VT::pi(), 0.1 );

  double const radiansPI = VT::pi();
  double const degreesPI = VT::convert_to_degrees(radiansPI);

  BOOST_CHECK_CLOSE( degreesPI, 180.0, 0.1 );

}

BOOST_AUTO_TEST_SUITE_END();


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

  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#include <boost/simd/function/scalar/log10.hpp>
#include <boost/simd/function/std.hpp>
#include <scalar_test.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/minf.hpp>
#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/mone.hpp>
#include <boost/simd/constant/zero.hpp>
#include <boost/simd/constant/mzero.hpp>
#include <boost/simd/constant/ten.hpp>
#include <boost/simd/constant/three.hpp>

STF_CASE_TPL (" log10 std",  STF_IEEE_TYPES)
{
  namespace bs = boost::simd;
  namespace bd = boost::dispatch;
  using bs::log10;

  using r_t = decltype(log10(T()));

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

  // specific values tests
#ifndef BOOST_SIMD_NO_INVALIDS
  STF_ULP_EQUAL(bs::std_(log10)(bs::Inf<T>()), bs::Inf<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Minf<T>()), bs::Nan<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Nan<T>()), bs::Nan<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Mone<T>()), bs::Nan<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Zero<T>()), bs::Minf<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Mzero<T>()), bs::Minf<r_t>(), 0);
#endif
  STF_ULP_EQUAL(bs::std_(log10)(bs::One<T>()), bs::Zero<r_t>(), 0);
  STF_ULP_EQUAL(bs::std_(log10)(bs::Two<T>()), T(0.301029995663981195213738894724), 0);
}


#include "core/algebra/EncVec.hpp"
#include "core/algebra/EncMat.hpp"
#include "HElib/FHE.h"
#include "HElib/EncryptedArray.h"
#include "core/coda.hpp"
#include <NTL/vector.h>
#include <NTL/mat_ZZ.h>
int test1(core::sk_ptr sk, core::pk_ptr pk, const EncryptedArray *ea) {
    core::EncVec encVec(pk);
    NTL::vec_ZZ slots;
    slots.SetLength(4);
    for (int i = 0; i < slots.length(); i++)
        slots[i] = NTL::to_ZZ(i + 1);
    encVec.pack(slots);

    long width = 4;
    long idx = 2;
    auto tmp = encVec.replicate(idx, width);
    NTL::vec_ZZ _slots;
    tmp.unpack(_slots, sk, true); // [5, 5, 5, 5]
    std::cout << slots << " " << _slots << std::endl;
    if (_slots.length() != width)
        return -1;
    for (auto t : _slots) {
        if (t != slots[idx]) return -1;
    }

    core::EncMat encMat(pk);
    NTL::mat_ZZ matrix;
    matrix.SetDims(3, 4);
    for (int r = 0; r < matrix.NumRows(); r++)
        for (int c = 0; c < matrix.NumCols(); c++)
            matrix.put(r, c, NTL::to_ZZ(1));
    encMat.pack(matrix);

    core::EncMat encMat2(pk);
    NTL::mat_ZZ matrix2;
    matrix2.SetDims(4, 5);
    for (int r = 0; r < matrix2.NumRows(); r++)
        for (int c = 0; c < matrix2.NumCols(); c++)
            matrix2.put(r, c, NTL::to_ZZ(1));
    encMat2.pack(matrix2);

    encMat.dot(encMat2);
    NTL::mat_ZZ _mat;
    encMat.unpack(_mat, sk, false);
    NTL::mat_ZZ zM, zM2, zM3;
    matrix *= matrix2;
    if (matrix != _mat)
        return -1;
    return 0;
}

int test2(core::sk_ptr sk, core::pk_ptr pk, const EncryptedArray *ea) {
    long p = pk->getContext().alMod.getPPowR();
    NTL::mat_ZZ X;
    X.SetDims(4, 4);
    NTL::ZZ range = NTL::to_ZZ(3);
    for (int r = 0; r < X.NumRows(); r++) {
        for (int c = r; c < X.NumCols(); c++) {
            auto rnd = NTL::RandomBnd(range);
            X.put(r, c, rnd);
            X.put(c, r, rnd);
        }
    }

    NTL::vec_ZZ v;
    v.SetLength(4);
    for (int c = 0; c < v.length(); c++)
        v[c] = NTL::RandomBnd(range);

    core::EncMat encMat(pk);
    core::EncVec encVec(pk);
    encMat.pack(X);
    encVec.pack(v);

    NTL::ZZ_p::init(NTL::to_ZZ(p));
    NTL::mat_ZZ_p Xp;
    NTL::vec_ZZ_p Vp;
    NTL::conv(Xp, X);
    NTL::conv(Vp, v);
    long T = 4;
    for (int t = 0; t < T; t++) {
        encVec = encMat.sym_dot(encVec); // only works for symmetric
        NTL::mul(Vp, Xp, Vp);
    }

    NTL::vec_ZZ_p modV;
    encVec.unpack(v, sk);
    NTL::conv(modV, v);
    if (modV != Vp)
        return -1;
    return 0;
}

int test3(core::sk_ptr sk, core::pk_ptr pk, const EncryptedArray *ea) {
    long p = pk->getContext().alMod.getPPowR();
    NTL::mat_ZZ X, Y;
    X.SetDims(4, 4);
    Y.SetDims(4, 4);
    NTL::ZZ range = NTL::to_ZZ(3);
    for (int r = 0; r < X.NumRows(); r++) {
        for (int c = 0; c < X.NumCols(); c++) {
            X.put(r, c, NTL::RandomBnd(range));
            Y.put(r, c, NTL::RandomBnd(range));
        }
    }

    NTL::ZZ_p::init(NTL::to_ZZ(p));
    NTL::mat_ZZ_p Xp, Yp;
    NTL::conv(Xp, X);
    NTL::conv(Yp, Y);

    core::EncMat encMat(pk), encMat2(pk);
    encMat.pack(X);
    encMat2.pack(Y);

    long T = 4;
    for (int t = 0; t < T; t++) {
        encMat.dot(encMat2);
        NTL::mul(Xp, Xp, Yp);
    }
    encMat.unpack(X, sk);
    NTL::conv(Yp, X);
    if (Xp != Yp)
        return -1;
    return 0;
}

int test4(core::sk_ptr sk, core::pk_ptr pk, const EncryptedArray *ea) {
    NTL::mat_ZZ X, Y;
    X.SetDims(4, 4);
    NTL::ZZ range = NTL::to_ZZ(3);
    for (int r = 0; r < X.NumRows(); r++) {
        for (int c = 0; c < X.NumCols(); c++) {
            X.put(r, c, NTL::RandomBnd(range));
        }
    }

    {
        core::EncMat encMat(pk);
        encMat.add(X);
        encMat.unpack(Y, sk);
        if (X != Y)
            return -1;
        core::EncMat encMat2(pk);
        encMat2.sub(X);
        encMat2.unpack(Y, sk, true);
        auto tmp(X);
        NTL::negate(tmp, X);
        if (tmp != Y)
            return -1;
    } // Test empty matrix adds plain matrix
    {
        core::EncMat encMat(pk), encMat2(pk), encMat3(pk);
        encMat2.pack(X);
        encMat.add(encMat2);
        encMat.unpack(Y, sk);
        if (X != Y)
            return -1;
        encMat3.sub(encMat2);
        encMat3.unpack(Y, sk, true);
        auto tmp(X);
        NTL::negate(tmp, X);
        if (tmp != Y)
            return -1;
    } // Test empty matrix adds cipher matrix
    return 0;
}

int testIO(core::sk_ptr sk, core::pk_ptr pk, const EncryptedArray *ea) {
    core::EncVec encVec(pk);
    NTL::vec_ZZ slots;
    slots.SetLength(7);
    for (int i = 0; i < slots.length(); i++)
        slots[i] = NTL::to_ZZ(i + 1);
    encVec.pack(slots);
    {
        std::ofstream out("./tmp.ctx", std::ios::binary);
        if (!encVec.dump(out)) {
            out.close();
            return -1;
        }
        out.close();

        std::ifstream in("./tmp.ctx", std::ios::binary);
        core::EncVec encVec2(pk);
        if (!encVec2.restore(in)) {
            in.close();
            return -1;
        }
        in.close();

        NTL::vec_ZZ slots2;
        encVec2.unpack(slots2, sk);
        if (slots2 != slots)
            return -1;
    } // Test EncVec

    core::EncMat encMat(pk);
    NTL::mat_ZZ mat;
    mat.SetDims(7, 7);
    for (long r = 0; r < mat.NumRows(); r++)
        for (long c = 0; c < mat.NumCols(); c++) mat[r][c] = 1;
    encMat.pack(mat);
    {
        std::ofstream out("./tmp.ctx", std::ios::binary);
        if (!encMat.dump(out)) {
            out.close();
            return -1;
        }
        out.close();

        std::ifstream in("./tmp.ctx", std::ios::binary);
        core::EncMat encMat2(pk);
        if (!encMat2.restore(in)) {
            in.close();
            return -1;
        }
        in.close();

        NTL::mat_ZZ mat2;
        encMat2.unpack(mat2, sk);
        if (mat != mat2)
            return -1;
    } // Test EncMat
    return 0;
}

int main() {
    core::context_ptr context = std::make_shared<FHEcontext>(512, 8209, 1);
    buildModChain(*context, 14);
    core::sk_ptr sk = std::make_shared<FHESecKey>(*context);
    sk->GenSecKey(64);
    addSome1DMatrices(*sk);
    core::pk_ptr pk = std::make_shared<FHEPubKey>(*sk);
    auto ea = context->ea;
    if (test1(sk, pk, ea) != 0) {
        std::cout << "test1 fail\n";
        return -1;
    }
    if (test2(sk, pk, ea) != 0) {
        std::cout << "test2 fail\n";
        return -1;
    }
    if (test3(sk, pk, ea) != 0) {
        std::cout << "test3 fail\n";
        return -1;
    }
    if (test4(sk, pk, ea) != 0) {
        std::cout << "test4 fail\n";
        return -1;
    }
    if (testIO(sk, pk, ea) != 0) {
        std::cout << "test3 fail\n";
        return -1;
    }
    std::cout << "passed\n";
    return 0;
}


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include "Resize.hpp"

#include "TensorBufferArrayView.hpp"

#include <boost/numeric/conversion/cast.hpp>

#include <cmath>
#include <algorithm>

using namespace armnnUtils;

namespace armnn
{

namespace
{

inline float Lerp(float a, float b, float w)
{
    return w * b + (1.f - w) * a;
}

inline double EuclideanDistance(float Xa, float Ya, const unsigned int Xb, const unsigned int Yb)
{
    return std::sqrt(pow(Xa - boost::numeric_cast<float>(Xb), 2) + pow(Ya - boost::numeric_cast<float>(Yb), 2));
}

}// anonymous namespace

void Resize(Decoder<float>&   in,
            const TensorInfo& inputInfo,
            Encoder<float>&   out,
            const TensorInfo& outputInfo,
            DataLayoutIndexed dataLayout,
            armnn::ResizeMethod resizeMethod,
            bool alignCorners)
{
    // We follow the definition of TensorFlow and AndroidNN: the top-left corner of a texel in the output
    // image is projected into the input image to figure out the interpolants and weights. Note that this
    // will yield different results than if projecting the centre of output texels.

    const unsigned int batchSize = inputInfo.GetShape()[0];
    const unsigned int channelCount = inputInfo.GetShape()[dataLayout.GetChannelsIndex()];

    const unsigned int inputHeight = inputInfo.GetShape()[dataLayout.GetHeightIndex()];
    const unsigned int inputWidth = inputInfo.GetShape()[dataLayout.GetWidthIndex()];
    const unsigned int outputHeight = outputInfo.GetShape()[dataLayout.GetHeightIndex()];
    const unsigned int outputWidth = outputInfo.GetShape()[dataLayout.GetWidthIndex()];

    const unsigned int sizeOffset = resizeMethod == armnn::ResizeMethod::Bilinear && alignCorners ? 1 : 0;

    // How much to scale pixel coordinates in the output image, to get the corresponding pixel coordinates
    // in the input image.
    const float scaleY = boost::numeric_cast<float>(inputHeight - sizeOffset)
                       / boost::numeric_cast<float>(outputHeight - sizeOffset);
    const float scaleX = boost::numeric_cast<float>(inputWidth - sizeOffset)
                       / boost::numeric_cast<float>(outputWidth - sizeOffset);

    TensorShape inputShape =  inputInfo.GetShape();
    TensorShape outputShape =  outputInfo.GetShape();

    for (unsigned int n = 0; n < batchSize; ++n)
    {
        for (unsigned int c = 0; c < channelCount; ++c)
        {
            for (unsigned int y = 0; y < outputHeight; ++y)
            {
                // Corresponding real-valued height coordinate in input image.
                const float iy = boost::numeric_cast<float>(y) * scaleY;

                // Discrete height coordinate of top-left texel (in the 2x2 texel area used for interpolation).
                const float fiy = floorf(iy);
                const unsigned int y0 = boost::numeric_cast<unsigned int>(fiy);

                // Interpolation weight (range [0,1]).
                const float yw = iy - fiy;

                for (unsigned int x = 0; x < outputWidth; ++x)
                {
                    // Real-valued and discrete width coordinates in input image.
                    const float ix = boost::numeric_cast<float>(x) * scaleX;
                    const float fix = floorf(ix);
                    const unsigned int x0 = boost::numeric_cast<unsigned int>(fix);

                    // Interpolation weight (range [0,1]).
                    const float xw = ix - fix;

                    // Discrete width/height coordinates of texels below and to the right of (x0, y0).
                    const unsigned int x1 = std::min(x0 + 1, inputWidth - 1u);
                    const unsigned int y1 = std::min(y0 + 1, inputHeight - 1u);

                    float interpolatedValue;
                    switch (resizeMethod)
                    {
                        case armnn::ResizeMethod::Bilinear:
                        {
                            in[dataLayout.GetIndex(inputShape, n, c, y0, x0)];
                            float input1 = in.Get();
                            in[dataLayout.GetIndex(inputShape, n, c, y0, x1)];
                            float input2 = in.Get();
                            in[dataLayout.GetIndex(inputShape, n, c, y1, x0)];
                            float input3 = in.Get();
                            in[dataLayout.GetIndex(inputShape, n, c, y1, x1)];
                            float input4 = in.Get();

                            const float ly0 = Lerp(input1, input2, xw); // lerp along row y0.
                            const float ly1 = Lerp(input3, input4, xw); // lerp along row y1.
                            interpolatedValue = Lerp(ly0, ly1, yw);
                            break;
                        }
                        case armnn::ResizeMethod::NearestNeighbor:
                        {
                            // calculate euclidean distance to the 4 neighbours
                            auto distance00 = EuclideanDistance(fix, fiy, x0, y0);
                            auto distance01 = EuclideanDistance(fix, fiy, x0, y1);
                            auto distance10 = EuclideanDistance(fix, fiy, x1, y0);
                            auto distance11 = EuclideanDistance(fix, fiy, x1, y1);

                            auto minimum = std::min( { distance00, distance01, distance10, distance11 } );

                            unsigned int xNearest = 0;
                            unsigned int yNearest = 0;

                            if (minimum == distance00)
                            {
                               xNearest = x0;
                               yNearest = y0;
                            }
                            else if (minimum == distance01)
                            {
                                xNearest = x0;
                                yNearest = y1;
                            }
                            else if (minimum == distance10)
                            {
                                xNearest = x1;
                                yNearest = y0;
                            }
                            else if (minimum == distance11)
                            {
                                xNearest = x1;
                                yNearest = y1;
                            }
                            else
                            {
                                throw armnn::InvalidArgumentException("Resize Nearest Neighbor failure");
                            }

                            in[dataLayout.GetIndex(inputShape, n, c, yNearest, xNearest)];
                            interpolatedValue = in.Get();
                            break;
                        }
                        default:
                            throw armnn::InvalidArgumentException("Unknown resize method: " +
                                                                  std::to_string(static_cast<int>(resizeMethod)));
                    }
                    out[dataLayout.GetIndex(outputShape, n, c, y, x)];
                    out.Set(interpolatedValue);
                }
            }
        }
    }
}

} //namespace armnn


#include <iostream>
#include <vector>

#include <Eigen/Dense>
#include <eigen3/Eigen/Core>

using MatrixXd = Eigen::MatrixXd;
using VectorXd = Eigen::VectorXd;
using Vector2d = Eigen::Vector2d;
using Vector4d = Eigen::Vector4d;
using Vector6d = Eigen::Matrix<double, 6, 1>;
using Vector8d = Eigen::Matrix<double, 8, 1>;
using Matrix8d = Eigen::Matrix<double, 8, 8>;
using TimeStamp = double;
using TimeDuration = double;

constexpr double dt = 1.;

struct RocketState {
    // Using a single dt is a better option.
    // But it make index handling a little bit more complex.
    enum StateIndex : int16_t {
        i_position_x = 0,
        i_position_y,
        i_velocity_x,
        i_velocity_y,
        STATE_SIZE
    };

    explicit RocketState(const double px,
        const double py,
        const double vel_x,
        const double vel_y)
    {
        variables << px, py, vel_x, vel_y;
    }

    RocketState() = default;

    static int variable_size()
    {
        return STATE_SIZE;
    }

    static int lcn_state_size()
    {
        return 2;
    }

    static int control_state_size()
    {
        // i_turning_rate,
        // i_acceleration,
        return 2;
    }

    Eigen::Map<Vector2d> lcn_state()
    {
        return Eigen::Map<Vector2d>(variables.data());
    }

    Eigen::Map<Vector2d> cntl_state()
    {
        return Eigen::Map<Vector2d>(variables.data() + 2);
    }

    Vector2d lcn_state() const
    {
        return variables.block<2, 1>(0, 0);
    }

    Vector2d cntl_state() const
    {
        return variables.block<2, 1>(2, 0);
    }

    Eigen::Map<Vector2d> position()
    {
        return Eigen::Map<Vector2d>(variables.data() + i_position_x);
    }

    Eigen::Map<Vector2d> velocity()
    {
        return Eigen::Map<Vector2d>(variables.data() + i_velocity_x);
    }

    Vector2d velocity() const
    {
        return { variables(i_velocity_x), variables(i_velocity_y) };
    }

    Vector4d variables = Vector4d::Zero();
};

struct RocketMotionModel {
    RocketState motion(const RocketState& state_in) const
    {
        auto state = state_in;
        // TODO: non-const getter
        RocketState next_state = state;

        next_state.velocity() = state.velocity();

        next_state.position()[0] = state.position()[0] + sin(state.velocity()[0]);
        next_state.position()[1] = state.position()[1] + state.velocity()[1];

        return next_state;
    }

    MatrixXd jacobian_wrt_state(const RocketState& state_in) const
    {
        using S = RocketState;
        const int residual_size = RocketState::STATE_SIZE;
        MatrixXd jacobi_wrt_s = MatrixXd::Identity(residual_size, residual_size);

        jacobi_wrt_s(S::i_position_x, S::i_velocity_x) = cos(state_in.velocity()[0]);
        jacobi_wrt_s(S::i_position_y, S::i_velocity_y) = 1.;

        return jacobi_wrt_s;
    }
};

#define PRINT_NAME_VAR(var) std::cout << #var << " :" << var << std::endl

using State = Vector2d;
using Control = Vector2d;

struct RocketLandingProblem {
    // Variables
    std::vector<RocketState> states = std::vector<RocketState>(11);

    // support data
    RocketState start_state = RocketState(0.1, 0.1, 0, 0);
    RocketState end_state = RocketState(2, 2, 0, 0);
};

struct QuadraticCost {
    QuadraticCost() = default;

    QuadraticCost(const VectorXd& meanIn, const MatrixXd& weightIn)
        : mean(meanIn)
        , weight(weightIn)
    {
    }

    double eval(const VectorXd& x) const
    {
        return (x - mean).transpose() * weight * (x - mean);
    }

    VectorXd mean;
    MatrixXd weight;
};

// du = K1 * dx + k2
struct FeedbackLaw {
    FeedbackLaw() = default;

    FeedbackLaw(MatrixXd K1, VectorXd k2)
        : K1(K1)
        , k2(k2)
    {
    }

    VectorXd apply(const VectorXd& x) const
    {
        return K1 * x + k2;
    }

    MatrixXd K1;
    VectorXd k2;
};

// The definition of states and controls is differient from previous section.
struct DDPStates {
    State last_state()
    {
        return states.at(num_states - 1);
    }

    int num_states = 0;
    std::vector<Control> controls;
    std::vector<State> states;
    std::vector<FeedbackLaw> feedbacks;

    State init_state;
    State target_state;
};

inline RocketState ddp_vars_to_rocket_state(const State& x, const Control& u)
{
    RocketState rstate;
    rstate.cntl_state() = u;
    rstate.lcn_state() = x;
    return rstate;
}

class DifferentialDynamicProgramming {
public:
    void solve(RocketLandingProblem& problem)
    {
        DDPStates ddp_states = initialize_ddp_state(problem);
        for (int iteration = 0; iteration < 10; ++iteration) {
            forward_pass(ddp_states);
            backward_pass(ddp_states);
            apply_control_law(ddp_states);
        }

        update_output(ddp_states, problem);
    }

private:
    DDPStates initialize_ddp_state(const RocketLandingProblem& problem)
    {
        DDPStates ddp_states;
        ddp_states.num_states = problem.states.size();
        num_states_ = ddp_states.num_states;
        ddp_states.init_state = problem.start_state.lcn_state();
        ddp_states.target_state = problem.end_state.lcn_state();

        PRINT_NAME_VAR(ddp_states.num_states);
        PRINT_NAME_VAR(ddp_states.init_state);
        PRINT_NAME_VAR(ddp_states.target_state);

        ddp_states.states.resize(ddp_states.num_states);
        ddp_states.controls.resize(ddp_states.num_states - 1);
        ddp_states.feedbacks.resize(ddp_states.num_states - 1);

        for (int i = 0; i < ddp_states.num_states; ++i) {
            ddp_states.states.at(i) = problem.states.at(i).lcn_state();
        }

        for (int i = 0; i < ddp_states.num_states - 1; ++i) {
            ddp_states.controls.at(i) = problem.states.at(i).cntl_state();
        }

        return ddp_states;
    }

    void update_output(const DDPStates& ddp_states, RocketLandingProblem& problem)
    {
        for (int i = 0; i < ddp_states.num_states; ++i) {
            problem.states.at(i).lcn_state() = ddp_states.states.at(i);
        }

        for (int i = 0; i < ddp_states.num_states - 1; ++i) {
            problem.states.at(i).cntl_state() = ddp_states.controls.at(i);
        }
    }

    void forward_pass(DDPStates& ddp_states)
    {
        ddp_states.states.resize(ddp_states.num_states);
        ddp_states.controls.resize(ddp_states.num_states - 1);

        ddp_states.states[0] = ddp_states.init_state;

        RocketMotionModel motion_model;
        for (int i = 0; i < ddp_states.num_states - 1; ++i) {
            RocketState current_full_state = ddp_vars_to_rocket_state(
                ddp_states.states.at(i), ddp_states.controls.at(i));

            RocketState next_full_state = motion_model.motion(current_full_state);
            ddp_states.states.at(i + 1) = next_full_state.lcn_state();

            std::cout << "idx forward iter: " << i << std::endl;
            std::cout << "control: " << ddp_states.controls.at(i).transpose() << std::endl;
            std::cout << "state: " << ddp_states.states.at(i).transpose() << std::endl;
            std::cout << "next state: " << ddp_states.states.at(i + 1).transpose() << std::endl;
        }
    }

    void backward_pass(DDPStates& ddp_states)
    {
        QuadraticCost marginal_cost = target_cost(ddp_states.target_state);

        current_best_cost_ = std::min(current_best_cost_, marginal_cost.eval(ddp_states.last_state()));
        std::cout << "target cost :" << marginal_cost.eval(ddp_states.last_state()) << std::endl;

        for (int i = ddp_states.num_states - 2; i >= 0; --i) {
            const State cur_state = ddp_states.states.at(i);
            const Control cur_control = ddp_states.controls.at(i);

            solve_backward_subproblem(cur_state, cur_control, marginal_cost, ddp_states.feedbacks[i]);
        }
    }

    void solve_backward_subproblem(const State& cur_state,
        const Control& cur_control,
        QuadraticCost& marginal_cost,
        FeedbackLaw& feedback_law)
    {
        const int state_size = 2;
        const int control_size = 2;

        MatrixXd lhs;
        VectorXd rhs;
        compute_backward_subproblem_normal_equation(marginal_cost, cur_state, cur_control, lhs, rhs);

        MatrixXd A1 = lhs.block(0, 0, state_size, state_size);
        MatrixXd A2 = lhs.block(0, state_size, state_size, control_size);
        MatrixXd A3 = lhs.block(state_size, 0, control_size, state_size);
        MatrixXd A4 = lhs.block(state_size, state_size, control_size, control_size);

        VectorXd b1 = rhs.block(0, 0, state_size, 1);
        VectorXd b2 = rhs.block(state_size, 0, control_size, 1);

        MatrixXd A4_inv = A4.inverse();

        // eliminate du
        MatrixXd lhs_xi = A1 - A2 * A4_inv * A3;
        VectorXd rhs_xi = b1 - A2 * A4_inv * b2;

        // solve PSD system by Cholesky
        VectorXd xi_star = lhs_xi.llt().solve(rhs_xi);

        marginal_cost = QuadraticCost(cur_state + xi_star, lhs_xi);
        // std::cout << "marginal_cost \n mean:" << marginal_cost.mean.transpose()
        //     << "\n weight:" << marginal_cost.weight << std::endl;
        feedback_law = FeedbackLaw(-A4_inv * A3, A4_inv * b2);
    }

    void compute_backward_subproblem_normal_equation(const QuadraticCost& marginal_cost,
        const State& cur_state,
        const Control& cur_control,
        MatrixXd& lhs,
        VectorXd& rhs)
    {
        const int state_size = 2;
        const int control_size = 2;
        const int total_size = state_size + control_size;

        RocketMotionModel motion_model;
        const RocketState rstate = ddp_vars_to_rocket_state(cur_state, cur_control);

        const MatrixXd jacobi_wrt_augment_state = motion_model.jacobian_wrt_state(rstate);
        MatrixXd jacobian_marginal_cost = jacobi_wrt_augment_state.block(0, 0, state_size, total_size);

        const RocketState predict_state = motion_model.motion(rstate);

        std::cout << "rstate.var:       " << rstate.variables << std::endl;
        std::cout << "predict_state.var:" << predict_state.variables << std::endl;
        std::cout << "jacobian_marginal_cost: " << jacobian_marginal_cost << std::endl;

        const VectorXd residual = predict_state.lcn_state() - marginal_cost.mean;

        lhs = jacobian_marginal_cost.transpose() * marginal_cost.weight * jacobian_marginal_cost;
        rhs = -jacobian_marginal_cost.transpose() * marginal_cost.weight * residual;

        // current cost (just regularizations for this problem)
        const double time_regularization = 1e-3 / num_states_;
        lhs(RocketState::i_velocity_x, RocketState::i_velocity_x) += 2 * time_regularization;
        rhs(RocketState::i_velocity_x) += -2 * time_regularization * cur_control(RocketState::i_velocity_x - 2);

        lhs(RocketState::i_velocity_y, RocketState::i_velocity_y) += 2 * time_regularization;
        rhs(RocketState::i_velocity_y) += -2 * time_regularization * cur_control(RocketState::i_velocity_y - 2);
    }

    QuadraticCost target_cost(const Vector2d& target_state)
    {
        QuadraticCost qcost;
        qcost.mean = target_state;
        MatrixXd jacobi = MatrixXd::Identity(2, 2);

        MatrixXd weight = MatrixXd::Identity(2, 2);
        Vector2d weight_diag;
        //  dt, x, y, vel_x, vel_y, heading
        weight_diag << 1e0, 1e0;
        weight.diagonal() = weight_diag;

        qcost.weight = jacobi.transpose() * weight * jacobi;

        return qcost;
    }

    void apply_control_law(DDPStates& ddp_states)
    {
        std::vector<Control> new_controls = ddp_states.controls;
        std::vector<State> new_states = ddp_states.states;
        RocketMotionModel motion_model;

        double shooting_cost = std::numeric_limits<double>::max();

        PRINT_NAME_VAR(shooting_cost);
        PRINT_NAME_VAR(current_best_cost_);

        // DDP is HARD to converge!
        double step = 0.5;
        // for (; current_best_cost_ < shooting_cost && step > 1e-4; step *= 0.5) {
        for (int i = 0; i < ddp_states.num_states - 1; ++i) {
            const FeedbackLaw& feedback = ddp_states.feedbacks.at(i);
            Vector2d delta_state = new_states.at(i) - ddp_states.states.at(i);
            if (i == 0)
                assert(delta_state.sum() == 0);

            Vector2d delta_control = feedback.apply(delta_state);

            std::cout << "=== idex: " << i << std::endl;
            std::cout << "delta_state: " << delta_state.transpose() << std::endl;
            std::cout << "delta_control: " << delta_control.transpose() << std::endl;

            new_controls.at(i) = ddp_states.controls.at(i) + step * delta_control;

            const RocketState new_rstate = ddp_vars_to_rocket_state(new_states.at(i), new_controls.at(i));
            new_states.at(i + 1) = motion_model.motion(new_rstate).lcn_state();

            std::cout << "new_control: " << new_controls.at(i).transpose() << std::endl;
            std::cout << "cur_state:   " << new_states.at(i).transpose() << std::endl;
            std::cout << "new_state:   " << new_states.at(i + 1).transpose() << std::endl;
        }

        QuadraticCost marginal_cost = target_cost(ddp_states.target_state);
        double cur_cost = marginal_cost.eval(new_states.back());
        shooting_cost = std::min(shooting_cost, cur_cost);
        // }

        current_best_cost_ = std::min(current_best_cost_, shooting_cost);

        std::cout << "ddp backtracking step:" << step << std::endl;
        std::cout << "current_best_cost_:" << current_best_cost_ << std::endl;
        std::cout << "cur_cost:" << cur_cost << std::endl;

        ddp_states.states = new_states;
        ddp_states.controls = new_controls;
    }

    // current minmum target cost. Should use total cost.
    double current_best_cost_ = std::numeric_limits<double>::max();

    double num_states_ = -1;
};

// yimu@yimu-mate:~/Desktop/blog$ /usr/bin/python3 /home/yimu/Desktop/blog/yimu-blog/least_squares/ddp/ddp_optimization.py
// initial_state: [0.1 0.1]
// target_state: [2. 2.]
// num_states: 11
// final_state_end_cost: [[0.22966875]]
// final_state_end_cost: [[0.00102592]]
// final_state_end_cost: [[0.00737957]]
// final_state_end_cost: [[0.00287627]]
// final_state_end_cost: [[0.00056493]]
// final_state_end_cost: [[6.04254466e-05]]
// final_state_end_cost: [[1.60630207e-06]]
// final_state_end_cost: [[8.4323494e-07]]
// final_state_end_cost: [[1.20808104e-06]]
// final_state_end_cost: [[6.61618857e-07]]
// ----------------------------------
// new_controls:
//  [array([0.19102753, 0.18981443]), array([0.19110838, 0.18991752]), array([0.19117968, 0.19001093]), array([0.19123781, 0.19009055]), array([0.19127743, 0.19015008]), array([0.19129027, 0.19017957]), array([0.19126273, 0.19016206]), array([0.19117054, 0.19006576]), array([0.19096421, 0.18982076]), array([0.19051605, 0.18921421])]
// new_states:
//  [array([0.1, 0.1]), array([0.28986784, 0.28981443]), array([0.47981505, 0.47973195]), array([0.66983226, 0.66974288]), array([0.85990655, 0.85983343]), array([1.05001973, 1.04998351]), array([1.24014552, 1.24016308]), array([1.43024427, 1.43032515]), array([1.62025251, 1.6203909 ]), array([1.81005818, 1.81021167]), array([1.9994238 , 1.99942588])]

// ./test_ddp | grep cur_cost
// cur_cost:0.229777
// cur_cost:0.00101972
// cur_cost:0.00736767
// cur_cost:0.00287048
// cur_cost:0.000562644
// cur_cost:5.97038e-05
// cur_cost:1.5031e-06
// cur_cost:9.30061e-07
// cur_cost:1.31691e-06
// cur_cost:7.43505e-07
//
// new_control: 0.190882 0.189625
// cur_state:   0.1 0.1
// --
// new_control: 0.191034 0.189811
// cur_state:   0.289725 0.289625
// --
// new_control: 0.191173 0.189985
// cur_state:   0.479599 0.479436
// --
// new_control: 0.191293 0.190141
// cur_state:    0.66961 0.669421
// --
// new_control: 0.191385 0.190269
// cur_state:   0.859738 0.859562
// --
// new_control: 0.191435 0.190352
// cur_state:   1.04996 1.04983
// --
// new_control: 0.191417 0.190362
// cur_state:   1.24023 1.24018
// --
// new_control:  0.19128 0.190237
// cur_state:   1.43048 1.43055
// --
// new_control:  0.19091 0.189834
// cur_state:   1.62059 1.62078
// --
// new_control: 0.189926 0.188637
// cur_state:   1.81034 1.81062
// --
// new_control: 0.191024 0.189811
// cur_state:   0.1 0.1
// --
// new_control: 0.191105 0.189914
// cur_state:   0.289864 0.289811
// --
// new_control: 0.191176 0.190007
// cur_state:   0.479808 0.479725
// --
// new_control: 0.191234 0.190087
// cur_state:   0.669822 0.669733
// --
// new_control: 0.191274 0.190147
// cur_state:   0.859892  0.85982
// --
// new_control: 0.191287 0.190176
// cur_state:      1.05 1.04997
// --
// new_control: 0.191259 0.190159
// cur_state:   1.24012 1.24014
// --
// new_control: 0.191167 0.190062
// cur_state:   1.43022  1.4303
// --
// new_control: 0.190961 0.189817
// cur_state:   1.62022 1.62036
// --
// new_control: 0.190513 0.189211
// cur_state:   1.81003 1.81018

int main(int argc, char* argv[])
{
    RocketLandingProblem p;

    DifferentialDynamicProgramming solver;
    solver.solve(p);

    return 1;
}


/**TODO:  Add copyright*/

#define BOOST_TEST_MODULE ModelFile test suite
#include <boost/test/included/unit_test.hpp>
#include <SmartPeak/io/ModelFile.h>

using namespace SmartPeak;
using namespace std;

BOOST_AUTO_TEST_SUITE(ModelFile1)

Model<float> makeModel1()
{
	/**
	* Directed Acyclic Graph Toy Network Model
	*/
	Node<float> i1, i2, h1, h2, o1, o2, b1, b2;
	Link l1, l2, l3, l4, lb1, lb2, l5, l6, l7, l8, lb3, lb4;
	Weight<float> w1, w2, w3, w4, wb1, wb2, w5, w6, w7, w8, wb3, wb4;
	Model<float> model1;

	// Toy network: 1 hidden layer, fully connected, DAG
	i1 = Node<float>("0", NodeType::input, NodeStatus::activated, std::make_shared<LinearOp<float>>(LinearOp<float>()), std::make_shared<LinearGradOp<float>>(LinearGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	i2 = Node<float>("1", NodeType::input, NodeStatus::activated, std::make_shared<LinearOp<float>>(LinearOp<float>()), std::make_shared<LinearGradOp<float>>(LinearGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	h1 = Node<float>("2", NodeType::hidden, NodeStatus::deactivated, std::make_shared<ReLUOp<float>>(ReLUOp<float>()), std::make_shared<ReLUGradOp<float>>(ReLUGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	h2 = Node<float>("3", NodeType::hidden, NodeStatus::deactivated, std::make_shared<ReLUOp<float>>(ReLUOp<float>()), std::make_shared<ReLUGradOp<float>>(ReLUGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	o1 = Node<float>("4", NodeType::output, NodeStatus::deactivated, std::make_shared<ReLUOp<float>>(ReLUOp<float>()), std::make_shared<ReLUGradOp<float>>(ReLUGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	o2 = Node<float>("5", NodeType::output, NodeStatus::deactivated, std::make_shared<ReLUOp<float>>(ReLUOp<float>()), std::make_shared<ReLUGradOp<float>>(ReLUGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	b1 = Node<float>("6", NodeType::bias, NodeStatus::activated, std::make_shared<LinearOp<float>>(LinearOp<float>()), std::make_shared<LinearGradOp<float>>(LinearGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));
	b2 = Node<float>("7", NodeType::bias, NodeStatus::activated, std::make_shared<LinearOp<float>>(LinearOp<float>()), std::make_shared<LinearGradOp<float>>(LinearGradOp<float>()), std::make_shared<SumOp<float>>(SumOp<float>()), std::make_shared<SumErrorOp<float>>(SumErrorOp<float>()), std::make_shared<SumWeightGradOp<float>>(SumWeightGradOp<float>()));

	// weights
	std::shared_ptr<WeightInitOp<float>> weight_init;
	std::shared_ptr<SolverOp<float>> solver;
	// weight_init.reset(new RandWeightInitOp(1.0)); // No random init for testing
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	w1 = Weight<float>("0", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
  w1.setWeight(1);
	w2 = Weight<float>("1", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
  w2.setWeight(2);
	w3 = Weight<float>("2", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
  w3.setWeight(3);
	w4 = Weight<float>("3", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	wb1 = Weight<float>("4", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	wb2 = Weight<float>("5", weight_init, solver);
	// input layer + bias
	l1 = Link("0", "0", "2", "0");
	l2 = Link("1", "0", "3", "1");
	l3 = Link("2", "1", "2", "2");
	l4 = Link("3", "1", "3", "3");
	lb1 = Link("4", "6", "2", "4");
	lb2 = Link("5", "6", "3", "5");
	// weights
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	w5 = Weight<float>("6", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	w6 = Weight<float>("7", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	w7 = Weight<float>("8", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	w8 = Weight<float>("9", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	wb3 = Weight<float>("10", weight_init, solver);
	weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
	solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
	wb4 = Weight<float>("11", weight_init, solver);
	// hidden layer + bias
	l5 = Link("6", "2", "4", "6");
	l6 = Link("7", "2", "5", "7");
	l7 = Link("8", "3", "4", "8");
	l8 = Link("9", "3", "5", "9");
	lb3 = Link("10", "7", "4", "10");
	lb4 = Link("11", "7", "5", "11");
	model1.setId(1);
	model1.setName("1");
	model1.addNodes({ i1, i2, h1, h2, o1, o2, b1, b2 });
	model1.addWeights({ w1, w2, w3, w4, wb1, wb2, w5, w6, w7, w8, wb3, wb4 });
	model1.addLinks({ l1, l2, l3, l4, lb1, lb2, l5, l6, l7, l8, lb3, lb4 });
	model1.setInputAndOutputNodes();
	return model1;
}

BOOST_AUTO_TEST_CASE(constructor)
{
  ModelFile<float>* ptr = nullptr;
  ModelFile<float>* nullPointer = nullptr;
  ptr = new ModelFile<float>();
  BOOST_CHECK_NE(ptr, nullPointer);
}

BOOST_AUTO_TEST_CASE(destructor)
{
  ModelFile<float>* ptr = nullptr;
	ptr = new ModelFile<float>();
  delete ptr;
}

BOOST_AUTO_TEST_CASE(storeModelDot)
{
	ModelFile<float> data;

	std::string filename = "ModelFileTest.gv";

  Model<float> model1 = makeModel1();
	data.storeModelDot(filename, model1);
}

BOOST_AUTO_TEST_CASE(loadModelCsv)
{
	ModelFile<float> data;
	Model<float> model_test;
	model_test.setId(1);
	model_test.setName("1");

	std::string filename_nodes = "ModelNodeFileTest.csv";
	std::string filename_links = "ModelLinkFileTest.csv";
	std::string filename_weights = "ModelWeightFileTest.csv";

  Model<float> model1 = makeModel1();
  model1.setInputAndOutputNodes();
	data.storeModelCsv(filename_nodes, filename_links, filename_weights, model1);

	data.loadModelCsv(filename_nodes, filename_links, filename_weights, model_test);
	BOOST_CHECK_EQUAL(model_test.getId(), model1.getId());
	BOOST_CHECK_EQUAL(model_test.getName(), model1.getName());
	BOOST_CHECK(model_test.getNodes() == model1.getNodes());
	BOOST_CHECK(model_test.getLinks() == model1.getLinks());
	//BOOST_CHECK(model_test.getWeights() == model1.getWeights());  // Broke
  BOOST_CHECK(model_test.getInputNodes().size() == model1.getInputNodes().size()); // Not sure why this fails
  BOOST_CHECK(model_test.getOutputNodes().size() == model1.getOutputNodes().size()); // Not sure why this fails
	//BOOST_CHECK(model_test == model1); // Not sure why this fails
}

BOOST_AUTO_TEST_CASE(loadModelBinary)
{
	ModelFile<float> data;
	Model<float> model_test;

	std::string filename = "ModelFileTest.binary";

  Model<float> model1 = makeModel1();
  model1.setInputAndOutputNodes();
	data.storeModelBinary(filename, model1);

	data.loadModelBinary(filename, model_test);
	BOOST_CHECK_EQUAL(model_test.getId(), model1.getId());
	BOOST_CHECK_EQUAL(model_test.getName(), model1.getName());
	BOOST_CHECK(model_test.getNodes() == model1.getNodes());
	BOOST_CHECK(model_test.getLinks() == model1.getLinks());
	BOOST_CHECK(model_test.getWeights() == model1.getWeights());
  BOOST_CHECK(model_test.getInputNodes().size() == model1.getInputNodes().size()); // Not sure why this fails
  BOOST_CHECK(model_test.getOutputNodes().size() == model1.getOutputNodes().size()); // Not sure why this fails
	//BOOST_CHECK(model_test == model1); // Not sure why this fails
}

BOOST_AUTO_TEST_CASE(loadWeightValuesBinary)
{
  // Store the binarized model
  ModelFile<float> data;
  std::string filename = "ModelFileTest.binary";
  Model<float> model1 = makeModel1();
  model1.setInputAndOutputNodes();
  data.storeModelBinary(filename, model1);

  // Read in the weight values
  std::map<std::string, std::shared_ptr<Weight<float>>> weights_test;
  for (int i = 0; i < 3; ++i) {
    auto weight_init = std::make_shared<ConstWeightInitOp<float>>(ConstWeightInitOp<float>(1.0));
    auto solver = std::make_shared<SGDOp<float>>(SGDOp<float>(0.01, 0.9));
    std::shared_ptr<Weight<float>> weight(new Weight<float>(
      std::to_string(i),
      weight_init,
      solver));
    weight->setModuleName(std::to_string(i));
    weight->setWeight(0);
    weights_test.emplace(weight->getName(), weight);
  }
  data.loadWeightValuesBinary(filename, weights_test);

  // Test that the weight values match
  for (int i = 0; i < 3; ++i) {
    BOOST_CHECK_EQUAL(model1.weights_.at(std::to_string(i))->getWeight(), weights_test.at(std::to_string(i))->getWeight());
    BOOST_CHECK(!weights_test.at(std::to_string(i))->getInitWeight());
  }
}

BOOST_AUTO_TEST_SUITE_END()


#define IGNORE

#include <boost/multiprecision/cpp_int.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix.hpp>

using mpi = boost::multiprecision::int128_t;
// using mpi = boost::multiprecision::cpp_int // 任意桁

using namespace boost::numeric;
using imatrix = ublas::matrix<mpi>;
// 行列積は ublas::prod
// imatrix mat(N, M);
// mat(N, M) = 2;


/*
MIT License

Copyright (c) 2019 Xiaohong Chen

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

#ifndef BOOST_UBLAS_PRECONDITIONER_HPP
#define BOOST_UBLAS_PRECONDITIONER_HPP

#include <boost/numeric/ublas/expression_types.hpp>
#include <boost/numeric/ublas/banded.hpp>
#include <boost/numeric/ublas/vector.hpp>

namespace boost { namespace numeric { namespace ublas {

template <class M>
class identity_precond {
public:
    identity_precond() {}
    identity_precond(const M&) {}

    template<class E>
    const vector_expression<E>& operator()(const vector_expression<E>& ve) const {
        return ve;
    }
};

template <class M>
class jacobi_precond {
public:
    typedef typename M::value_type value_type;
    typedef typename M::size_type size_type;
    jacobi_precond(const M& A)
        : diag_(A.size1()) {
        size_type size = A.size1();
        for (size_type i = 0; i < size; ++i) {
            diag_(i) = A(i, i);
        }
    }

    jacobi_precond(const jacobi_precond&) = default;
    jacobi_precond(jacobi_precond&&) = default;
    jacobi_precond& operator=(const jacobi_precond&) = default;
    jacobi_precond& operator=(jacobi_precond&&) = default;

    template<class E>
    decltype(auto) operator()(const vector_expression<E>& ve) const {
        return element_prod(ve, diag_);
    }

private:
    vector<value_type> diag_;
};

} // end namespace linear
} // end namespace solver
} // end namespace math

#endif


#include <stdio.h>
#include <fstream>
#include <string>
#include <time.h>
#include "Matrix.h"

// boost directives
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>

// // pybind11 directives for numpy arrays
// #include <pybind11/pybind11.h>
// #include <pybind11/numpy.h>

// current issue: cannot find <Python.h>, so we need to put this in -I flag
// This will be implemented in the Makefile

using namespace boost::numeric;

void read_from_file(ublas::matrix<double> &bmat, const std::string &fname)
{
    std::ifstream ifs{fname};

    if (!ifs)
    {
        throw std::runtime_error("Cannot open file!");
    }

    double val;

    std::string line;
    int i = 0; // row index

    while (std::getline(ifs, line))
    {                                // read each line in file
        int j = 0;                   // column index
        std::istringstream ss{line}; // stringstream of each row
        while (ss >> val)
        {
            bmat(i, j) = val; // set value
            ++j;              // increment column index
        }

        ++i; // increment row index
    }
}

// void read_from_nparray(ublas::matrix<double> &bmat, const pybind11::array_t<double>& arr)
// {
//         // request buffer info from numpy array
//     pybind11::buffer_info buf = arr.request();

//     // check if dimensions are equal
//     if (buf.shape[0] != bmat.size1() && buf.shape[1] != bmat.size2())
//     {
//         throw std::runtime_error("Dimensions are not equal!")
//     }
//     rsz = buf.shape[0];
//     csz = buf.shape[1];

//     // set a new pointer ptr to the buffer pointer
//     double *ptr = (double *)buf.ptr;

//     // set pointer to buffer pointer
//     for (int i = 0; i < rsz; ++i)
//     {
//         for (int j = 0; j < csz; ++j)
//         {
//             bmat(i,j) = ptr[i * csz + j];
//         }
//     }
// }

void test_performance(Matrix *myMat_ptr, ublas::matrix<double> bMat)
{
    // test performance
    int max_iter = 10000; // max iterations
    // empty arrays for mean computation
    // std::vector<double> my_normarr(max_iter);
    // std::vector<double> b_normarr(max_iter);
    // std::vector<double> my_timearr(max_iter);
    // std::vector<double> b_normarr(max_iter);
    double my_norm, b_norm, my_avgnorm, b_avgnorm;
    double my_time, b_time, my_avgtime, b_avgtime;

    for (int i = 0; i < max_iter; ++i)
    {
        clock_t my_t, b_t;
        // evaluate norm from matrix class
        my_t = clock();
        double my_n = myMat_ptr->norm();
        // myMat_ptr->norm() = my_n;
        my_t = clock() - my_t;
        // std::cout << "Norm from my library: " << my_n << std::endl
        //          << "Time from my library: " << (double) my_t / CLOCKS_PER_SEC << std::endl;

        // now from ublas
        b_t = clock();
        double b_n = norm_frobenius(bMat);
        b_t = clock() - b_t;
        // std::cout << "Norm from boost: " << b_n << std::endl
        //          << "Time from boost: " << (double) b_t / CLOCKS_PER_SEC << std::endl;

        // append the results
        my_norm += my_n;
        b_norm += b_n;
        my_time += ((double)my_t / CLOCKS_PER_SEC);
        b_time += ((double)b_t / CLOCKS_PER_SEC);
        // my_normarr.push_back(my_norm);
        // b_normarr.push_back(b_norm);
        // my_timearr.push_back((float) my_t / CLOCKS_PER_SEC);
        // b_timearr.push_back((float) b_t / CLOCKS_PER_SEC);
    }

    // evaluate the mean
    my_avgnorm = my_norm / max_iter;
    b_avgnorm = b_norm / max_iter;
    my_avgtime = my_time / max_iter;
    b_avgtime = b_time / max_iter;

    // print results
    printf("Results from performance benchmarking for %d-by-%d matrix:\n", bMat.size1(), bMat.size2());
    printf("Number of iterations for performance benchmark: %d\n", max_iter);
    std::cout << "Average norm evaluated from user-defined matrix: " << my_avgnorm << ", \t Average norm evaluated from Boost: " << b_avgnorm << std::endl
              << "Average time taken for evaluation of norm from user-defined matrix: " << my_avgtime << std::endl
              << "Average time taken for evaluation of norm from Boost: " << b_avgtime << std::endl
              << "Comparisons:" << std::endl
              << "Accuracy of user-defined matrix norm vs Boost norm: " << abs(b_avgnorm - my_avgnorm) << std::endl
              << "Ratio of user-defined matrix performance to Boost performance: " << my_avgtime / b_avgtime << std::endl;
}

int main()
try
{
    // initialize parameters
    // int rs{3};
    // int cs{4};
    // std::string fname{"data/small_data.tsv"};

    int rs{3000};
    int cs{4000};
    std::string fname{"data/large_data.tsv"};

    // create matrices
    Matrix *myMat_ptr = new Matrix{rs, cs, fname};
    ublas::matrix<double> bMat(rs, cs);
    read_from_file(bMat, fname);

    test_performance(myMat_ptr, bMat);

    return 0;
}
catch (std::exception &e)
{
    std::cerr << "error: " << e.what() << std::endl;
    return 1;
}

// {
//     // common parameters for both matrices
//     int rs{3};
//     int cs{4};
//     std::string fname{"data/small_data.txt"};

//     // // // common parameters for both matrices
//     // int rs {3000};
//     // int cs {4000};
//     // std::string fname {"data/large_data.txt"};

//     // create matrices and print them out
//     Matrix myMat{rs, cs, fname};
//     // Matrix *myMat = new Matrix {rs, cs, fname};
//     // double my_norm;
//     // myMat->norm() = my_norm;
//     // myMat.print_mat();

//     ublas::matrix<double> bMat(rs, cs);
//     read_from_file(bMat, fname);
//     // std::cout << bMat << '\n';

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

//     // now compute norm for both and print them out
//     // iterate 1e6 times and take average to determine average time
//     double my_norm = myMat.norm();
//     std::cout << "Norm from my library: " << my_norm << std::endl;

//     double b_norm = norm_frobenius(bMat);
//     std::cout << "Norm from boost: " << b_norm << std::endl;
// }