#include "polyFFT.h"
#include <cryptoTools/Common/Log.h>
#include <cryptoTools/Common/Timer.h>
#include <cryptoTools/Crypto/RandomOracle.h>
#include <vector>
#include <thread>
#include <NTL/BasicThreadPool.h>
#include <omp.h>
//#define fftTest

namespace osuCrypto
{
#define LEFT(X) (2*X+1)
#define RIGHT(X) (2*X+2)
#define PAPA(X) ((X-1)/2)

	void print_poly(ZZ_pX& P)
	{
		long degree = deg(P);
		if (-1 == degree) {
			cout << "0";
			return;
		}
		for (long i = 0; i <= degree; i++) {
			cout << coeff(P, i);
			if (i == 1)
				cout << "X";
			else if (i > 1)
				cout << "X^" << i;
			if (i < degree) {
				cout << " + ";
			}
		}
		//    cout << endl << "random poly:" << endl << P << endl;
	}

	void build_tree(ZZ_pX* tree, ZZ_p* points, unsigned int tree_size, int numThreads, ZZ &prime) {

		ZZ_p negated;
		unsigned int point_index;

		//build all leaves whose index starts at treesize/2 => tree[i]= x-xi
		for (u32 i = tree_size / 2; i< tree_size; i++) {
			point_index = i - (tree_size - 1) / 2;
			NTL::negate(negated, points[point_index]); //get -xi
			SetCoeff(tree[i], 0, negated);
			SetCoeff(tree[i], 1, 1);
		}

		vector<vector<int>> subs(numThreads);
		generateSubTreeArrays(subs, tree_size / 2, numThreads - 1);

		vector<thread> threads(numThreads);

		for (int t = 0; t < numThreads; t++)
			threads[t] = thread(&buildSubTree, tree, ref(subs[t]), ref(prime));

		for (int t = 0; t < numThreads; t++)
			threads[t].join();

		vector<thread> threadsLastPart(numThreads * 2);

		for (int i = log2(numThreads) - 1; i >= 1; i--) {
			//cout<<"layer "<<i<< " : " <<endl;
			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index to process is "<< j<<endl;

				threadsLastPart[j] = thread(&buildTreeSpecific, tree, j, ref(prime));
			}

			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index join to process is "<< j<<endl;
				threadsLastPart[j].join();
			}
		}

		tree[0] = tree[LEFT(0)] * tree[RIGHT(0)];

	}

	void generateSubTreeArrays(vector<vector<int>> &subArrays, int totalNodes, int firstIndex) {

		int numOfSubTrees = subArrays.size();

		//generate subtree array
		int maxElement = totalNodes;

		for (int i = 0; i < numOfSubTrees; i++) {
			subArrays[i].resize(maxElement);//this is overkill, but will be reduces later
			subArrays[i][0] = firstIndex + i;
		}

		for (int i = 0; i < numOfSubTrees; i++) {
			for (int j = 1; j < totalNodes; j++) {

				if (j % 2 == 1)
					subArrays[i][j] = LEFT(subArrays[i][PAPA(j)]);
				else
					subArrays[i][j] = RIGHT(subArrays[i][PAPA(j)]);

				if (subArrays[i][j] >= maxElement) {
					subArrays[i].resize(j);
					//cout << "j is in first loop " << j << endl;
					break;

				}

				//cout << subArrays[i][j] << "-";
			}
			//cout << "size of "<<i<<" array" << subArrays[i].size() << endl;
			//cout << "-------------------" << endl;

		}

	}

	void buildSubTree(ZZ_pX* tree, vector<int> &subTree, ZZ &prime) {

		ZZ_p::init(prime);
		int index;
		for (int i = subTree.size() - 1; i >= 0; i--) {
			index = subTree[i];

#ifdef fftTest
			FFTMul(tree[index], tree[LEFT(index)], tree[RIGHT(index)]);
#else
			tree[index] = tree[LEFT(index)] * tree[RIGHT(index)];
#endif // fftTest

			//

		}
	}

	void interSubTree(ZZ_pX* temp, ZZ_pX* M, vector<int> &subTree, ZZ &prime) {

		//cout<<"Thread indices";

		ZZ_p::init(prime);
		int index;
		for (int i = subTree.size() - 1; i >= 0; i--) {
			index = subTree[i];
			//cout<<"--"<<index;
#ifdef fftTest
			ZZ_pX tempfft1;
			FFTMul(tempfft1, temp[LEFT(index)], M[RIGHT(index)]);

			ZZ_pX tempfft2;
			FFTMul(tempfft2, temp[RIGHT(index)], M[LEFT(index)]);

			temp[index] = tempfft1 + tempfft2;
#else
			temp[index] = temp[LEFT(index)] * M[RIGHT(index)] + temp[RIGHT(index)] * M[LEFT(index)];
#endif // fftTest

		}

	}

	void evalSubTree(ZZ_pX* reminders, ZZ_pX* tree, vector<int> &subTree, ZZ &prime) {

		//cout<<"Thread indices";
		ZZ_p::init(prime);
		int index;

		unsigned int i = 1;
		for (i = 0; i < subTree.size(); i++) {
			index = subTree[i];
			//        cout << "i="<<i <<": ";
			//
#ifdef fftTest
			FFTRem(reminders[index], reminders[PAPA(index)], tree[index]);
#else
			reminders[index] = reminders[PAPA(index)] % tree[index];
#endif
		}

	}

	void test_tree(ZZ_pX& final_polynomial, ZZ_p* points, unsigned int npoints) {
		ZZ_p result;
		bool error = false;
		for (unsigned int i = 0; i < npoints; i++) {
			result = eval(final_polynomial, points[i]);
			if (0 != result) {
				cout << "FATAL ERROR: polynomials tree is incorrect!" << endl;
				error = true;
				break;
			}
		}
		if (!error)
			cout << "polynomials tree is correct." << endl;
	}

	void evaluate_main(ZZ_pX& P, ZZ_pX* tree, ZZ_pX* reminders, unsigned int tree_size, ZZ_p* results) {

		reminders[0] = P%tree[0];

		unsigned int i = 1;
		for (; i < tree_size / 2; i++) {
			reminders[i] = reminders[PAPA(i)] % tree[i];
		}

		unsigned int result_index;
		for (; i < tree_size; i++) {
			reminders[i] = reminders[PAPA(i)] % tree[i];
			result_index = i - (tree_size - 1) / 2;
			results[result_index] = coeff(reminders[i], 0);
		}
	}

	void evaluate(ZZ_pX& P, ZZ_pX* tree, ZZ_pX* reminders, unsigned int tree_size, ZZ_p* results, int numThreads, ZZ &prime) {

		auto begin1 = steady_clock::now();
		//set the reminder of the root
		//reminders[0] = P%tree[0];
		reminders[0] = P;
		reminders[1] = P;
		reminders[2] = P;

		auto end1 = steady_clock::now();
		//cout << "eval - not paralleled "<<" threads " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		//    print_poly(tree[0]);
		//
		//    cout<<endl<<"---Reminder"<<endl;
		//    print_poly(reminders[0]);
		//
		//    cout<<endl<<"---P"<<endl;
		//    print_poly(P);

		begin1 = steady_clock::now();
		vector<vector<int>> subs(numThreads);
		generateSubTreeArrays(subs, tree_size, numThreads - 1);
		end1 = steady_clock::now();
		//cout << "eval - generate sub trees: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		begin1 = steady_clock::now();
		vector<thread> threadsFirstPart(numThreads * 2);

		for (int i = 2; i < log2(numThreads); i++) {
			//cout<<"layer "<<i<< " : " <<endl;
			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index to process is "<< j<<endl;

				threadsFirstPart[j] = thread(&evalReminder, tree, reminders, j, ref(prime));

			}

			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index join to process is "<< j<<endl;
				threadsFirstPart[j].join();
			}
		}

		end1 = steady_clock::now();
		//cout << "eval - first part for max "<< numThreads/2<<" threads " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		begin1 = steady_clock::now();
		vector<thread> threads(numThreads);

		for (int t = 0; t < numThreads; t++) {

			threads[t] = thread(&evalSubTree, reminders, tree, ref(subs[t]), ref(prime));

		}

		for (int t = 0; t < numThreads; t++) {
			threads[t].join();
		}

		end1 = steady_clock::now();
		//cout << "eval - thread part for "<< numThreads<<" threads " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		begin1 = steady_clock::now();
		unsigned int result_index;
		for (int i = tree_size / 2; i < tree_size; i++) {
			result_index = i - (tree_size - 1) / 2;
			results[result_index] = coeff(reminders[i], 0);
		}

		end1 = steady_clock::now();
		//cout << "eval - assign last part: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

	}

	void evalReminder(ZZ_pX *tree, ZZ_pX *reminders, int i, ZZ &prime) {
		ZZ_p::init(prime);
#ifdef fftTest
		FFTRem(reminders[i], reminders[PAPA(i)], tree[i]);
#else
		reminders[i] = reminders[PAPA(i)] % tree[i];
#endif

	}

	void interSpecific(ZZ_pX *temp, ZZ_pX *M, int i, ZZ &prime) {
		ZZ_p::init(prime);

#ifdef fftTest
		ZZ_pX tempfft1;
		FFTMul(tempfft1, temp[LEFT(i)], M[RIGHT(i)]);

		ZZ_pX tempfft2;
		FFTMul(tempfft2, temp[RIGHT(i)], M[LEFT(i)]);

		temp[i] = tempfft1 + tempfft2;
#else
		temp[i] = temp[LEFT(i)] * M[RIGHT(i)] + temp[RIGHT(i)] * M[LEFT(i)];
#endif // fftTest

	}

	void buildTreeSpecific(ZZ_pX *tree, int i, ZZ &prime) {
		ZZ_p::init(prime);

#ifdef fftTest
		FFTMul(tree[i], tree[LEFT(i)], tree[RIGHT(i)]);
#else
		tree[i] = tree[LEFT(i)] * tree[RIGHT(i)];
#endif
	}

	void test_evaluate(ZZ_pX& P, ZZ_p* points, ZZ_p* results, unsigned int npoints) {
		bool error = false;
		for (unsigned int i = 0; i < npoints; i++) {
			ZZ_p y = eval(P, points[i]);
			if (y != results[i]) {
				cout << "y=" << y << " and results[i]=" << results[i] << endl;
				error = true;
			}
		}
		if (error)
			cout << "ERROR: evaluation results do not match real evaluation!" << endl;
		else
			cout << "All evaluation results computed correctly!" << endl;
	}

	void multipoint_evaluate_zp(ZZ_pX& P, ZZ_p* x, ZZ_p* y, long degree, int numThreads, ZZ &prime)
	{
		//    cout << "P:" <<endl; print_poly(P); cout << endl;
		// we want to evaluate P on 'degree+1' values.
		ZZ_pX* p_tree = new ZZ_pX[degree * 2 + 1];
		steady_clock::time_point begin1 = steady_clock::now();
		build_tree(p_tree, x, degree * 2 + 1, numThreads, prime);
		steady_clock::time_point end1 = steady_clock::now();
		//test_tree(p_tree[0], x, degree+1);

		ZZ_pX* reminders = new ZZ_pX[degree * 2 + 1];
		steady_clock::time_point begin2 = steady_clock::now();
		evaluate(P, p_tree, reminders, degree * 2 + 1, y, numThreads, prime);
		chrono::steady_clock::time_point end2 = steady_clock::now();

		//cout << "Building tree: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;
		//cout << "Evaluating points: " << duration_cast<milliseconds>(end2 - begin2).count() << " ms" << endl;
		//cout << "Total: " << duration_cast<milliseconds>(end1 - begin1).count()+ duration_cast<milliseconds>(end2 - begin2).count() << " ms" << endl;

		//test_evaluate(P,x,y,10);
	}

	void test_multipoint_eval_zp(ZZ prime, long degree, int numThreads)
	{
		// init underlying prime field
		ZZ_p::init(ZZ(prime));

		// the given polynomial
		ZZ_pX P;
		random(P, degree + 1);
		SetCoeff(P, degree, random_ZZ_p());

		// evaluation points:
		ZZ_p* x = new ZZ_p[degree + 1];
		ZZ_p* y = new ZZ_p[degree + 1];

		for (unsigned int i = 0; i <= degree; i++) {
			random(x[i]);
		}

		multipoint_evaluate_zp(P, x, y, degree, numThreads, prime);
	}

	/*
	* expects an "empty" polynomial 'resultP'
	*/
	void iterative_interpolate_zp_main(ZZ_pX& resultP, ZZ_pX* temp, ZZ_p* y, ZZ_p* a, ZZ_pX* M, unsigned int tree_size)
	{
		int i = tree_size - 1;
		ZZ_p inv_a;
		unsigned int y_index;
		for (; i >= tree_size / 2; i--) {
			y_index = i - (tree_size - 1) / 2;
			inv(inv_a, a[y_index]); // inv_a = 1/a[y_index]
			SetCoeff(temp[i], 0, y[y_index] * inv_a);
		}

		for (; i >= 0; i--) {
			temp[i] = temp[LEFT(i)] * M[RIGHT(i)] + temp[RIGHT(i)] * M[LEFT(i)];
		}

		resultP = temp[LEFT(0)] * M[RIGHT(0)] + temp[RIGHT(0)] * M[LEFT(0)];
	}

	void iterative_interpolate_zp(ZZ_pX& resultP, ZZ_pX* temp, ZZ_p* y, ZZ_p* a, ZZ_pX* M, unsigned int tree_size, int numThreads, ZZ &prime)
	{
		unsigned int i = tree_size - 1;
		ZZ_p inv_a;
		unsigned int y_index;
		for (; i >= tree_size / 2; i--) {
			y_index = i - (tree_size - 1) / 2;
			inv(inv_a, a[y_index]); // inv_a = 1/a[y_index]
			SetCoeff(temp[i], 0, y[y_index] * inv_a);
		}

		auto begin1 = steady_clock::now();
		vector<vector<int>> subs(numThreads);
		generateSubTreeArrays(subs, tree_size / 2, numThreads - 1);
		auto end1 = steady_clock::now();
		//cout << "inter - generate sub trees: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		begin1 = steady_clock::now();
		vector<thread> threads(numThreads);

		for (int t = 0; t < numThreads; t++) {

			threads[t] = thread(&interSubTree, temp, M, ref(subs[t]), ref(prime));

		}

		for (int t = 0; t < numThreads; t++) {
			threads[t].join();
		}

		end1 = steady_clock::now();
		//cout << "interpolate - threads part: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		//    begin1 = steady_clock::now();
		//    for(int i=numThreads-2; i>0; i--)
		//        temp[i] = temp[LEFT(i)] * M[RIGHT(i)] + temp[RIGHT(i)] * M[LEFT(i)] ;
		//
		//    end1 = steady_clock::now();
		//    cout << "inter - first part SERIAL "<< duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;
		//

		begin1 = steady_clock::now();
		vector<thread> threadsLastPart(numThreads * 2);

		for (int i = log2(numThreads) - 1; i >= 1; i--) {
			//cout<<"layer "<<i<< " : " <<endl;
			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index to process is "<< j<<endl;

				threadsLastPart[j] = thread(&interSpecific, temp, M, j, ref(prime));
			}

			for (int j = pow(2, i) - 1; j <= pow(2, i + 1) - 2; j++) {
				//cout<<" index join to process is "<< j<<endl;
				threadsLastPart[j].join();
			}
		}

		end1 = steady_clock::now();
		//cout << "inter - last part for max "<< numThreads/2<<" threads " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

		begin1 = steady_clock::now();
		//last iteration
		resultP = temp[LEFT(0)] * M[RIGHT(0)] + temp[RIGHT(0)] * M[LEFT(0)];

		end1 = steady_clock::now();
		//cout << "interpolate - last iteration for root: " << duration_cast<milliseconds>(end1 - begin1).count() << " ms" << endl;

	}

	void interpolate_zp(ZZ_pX& resultP, ZZ_p* x, ZZ_p* y, long degree, int numThreads, ZZ &prime)
	{
		system_clock::time_point begin[4];
		system_clock::time_point end[4];
		ZZ_pX *M = new ZZ_pX[degree * 2 + 1];;
		ZZ_p *a = new ZZ_p[degree + 1];;

		prepareForInterpolate(x, degree, M, a, numThreads, prime);

		//now we can apply the formula
		ZZ_pX* temp = new ZZ_pX[degree * 2 + 1];
		iterative_interpolate_zp(resultP, temp, y, a, M, degree * 2 + 1, numThreads, prime);

		//cout << "Interpolation: " << duration_cast<milliseconds>(end[4] - begin[4]).count() << " ms" << endl;
		//cout << "Total: " << duration_cast<milliseconds>(end[1]-begin[1] + end[2]-begin[2] + end[3]-begin[3] + end[4]-begin[4]).count() << " ms" << endl;
	}

	void prepareForInterpolate(ZZ_p *x, long degree, ZZ_pX *M, ZZ_p *a, int numThreads, ZZ &prime) {

		system_clock::time_point begin[4];
		system_clock::time_point end[4];
		begin[1] = system_clock::now();
		build_tree(M, x, degree * 2 + 1, numThreads, prime);
		end[1] = system_clock::now();

		ZZ_pX D;
		//we construct a preconditioned global structure for the a_k for all 1<=k<=(degree+1)ZZ_pX D;
		begin[2] = system_clock::now();
		diff(D, M[0]);
		end[2] = system_clock::now();

		//evaluate d(x) to obtain the results in the array a
		ZZ_pX* reminders = new ZZ_pX[degree * 2 + 1];
		begin[3] = system_clock::now();
		evaluate(D, M, reminders, degree * 2 + 1, a, numThreads, prime);
		end[3] = system_clock::now();
		//    test_evaluate(D,x,a,degree+1);

		/*cout << "Building tree: " << duration_cast<milliseconds>(end[1] - begin[1]).count() << " ms" << endl;
		cout << "Differentiate: " << duration_cast<milliseconds>(end[2] - begin[2]).count() << " ms" << endl;
		cout << "Evaluate diff: " << duration_cast<milliseconds>(end[3] - begin[3]).count() << " ms" << endl;*/

	}

	void test_interpolation_result_zp(ZZ_pX& P, ZZ_p* x, ZZ_p* y, long degree)
	{
		cout << "Testing result polynomial" << endl;
		ZZ_p res;
		for (long i = 0; i < degree + 1; i++) {
			eval(res, P, x[i]);
			if (res != y[i]) {
				cout << "Error! x = " << x[i] << ", y = " << y[i] << ", res = " << res << endl;
				return;
			}
		}
		cout << "Polynomial is interpolated correctly!" << endl;
	}

	void test_interpolate_zp(ZZ prime, long degree, int numThreads)
	{
		// init underlying prime field
		ZZ_p::init(ZZ(prime));

		// interpolation points:
		ZZ_p* x = new ZZ_p[degree + 1];
		ZZ_p* y = new ZZ_p[degree + 1];
		for (unsigned int i = 0; i <= degree; i++) {
			random(x[i]);
			random(y[i]);
		}

		ZZ_pX P;
		interpolate_zp(P, x, y, degree, numThreads, prime);
		//cout << "P: "; print_poly(P); cout << endl;
		//test_interpolation_result_zp(P, x, y, degree);
	}

	void BytesToZZ_px(unsigned char *bytesArr, ZZ_pX& poly, long numOfElements, long sizeOfElement) {

		//turn each byte to zz_p element in a vector

		vec_ZZ repFromBytes;
		repFromBytes.SetLength(numOfElements);

		for (int i = 0; i < numOfElements; i++) {

			ZZ zz;

			//translate the bytes into a ZZ element
			ZZFromBytes(zz, bytesArr + i*sizeOfElement, sizeOfElement);

			repFromBytes[i] = zz;
		}

		//turn the vec_zzp to the polynomial

		poly = to_ZZ_pX(to_vec_ZZ_p(repFromBytes));

	}
	void ZZ_pxToBytes(ZZ_pX& poly, unsigned char *bytesArr, long numOfElements, long sizeOfElement) {

		//get the zz_p vector

		for (int i = 0; i < numOfElements; i++) {

			BytesFromZZ(bytesArr + i*sizeOfElement, rep(poly.rep[i]), sizeOfElement);
		}

	}
}


#include <iostream>
#include <iomanip>
#include <algorithm>
#include <functional>
#include <ctime>
#include <cassert>

#include <boost/timer.hpp>
#include <boost/random.hpp>
#include <boost/lambda/lambda.hpp>

#include "StructuredPerceptron.h"

using namespace boost::lambda;

namespace Grante {

StructuredPerceptron::StructuredPerceptron(FactorGraphModel* fg_model,
	bool do_averaging, bool verbose)
	: ParameterEstimationMethod(fg_model), do_averaging(do_averaging),
		verbose(verbose), lh(fg_model)
{
	// Initialize parameter order
	const std::vector<FactorType*>& factor_types = fg_model->FactorTypes();
	for (std::vector<FactorType*>::const_iterator fti = factor_types.begin();
		fti != factor_types.end(); ++fti) {
		parameter_order.push_back((*fti)->Name());
	}
}

StructuredPerceptron::~StructuredPerceptron() {
}

void StructuredPerceptron::AddPrior(const std::string& factor_type,
	Prior* prior) {
	// Priors are not supported for Perceptron training
	assert(0);
}

double StructuredPerceptron::Train(double conv_tol, unsigned int max_epochs) {
	// Initialize random number generator
	size_t N = training_data.size();
	assert(N > 0);
	boost::mt19937 rgen(static_cast<const boost::uint32_t>(std::time(0))+1);
	boost::uniform_int<unsigned int> rdestd(0, static_cast<unsigned int>(N-1));
	boost::variate_generator<boost::mt19937,
		boost::uniform_int<unsigned int> > rand_n(rgen, rdestd);
	std::cout << "Perceptron training with " << N << " training instances"
		<< std::endl;

	// For all epochs
	assert(max_epochs > 0);
	boost::timer total_timer;
	unsigned int total_iter = 0;
	unsigned int violation_count = 0;
	for (unsigned int epoch = 0; max_epochs == 0 || epoch < max_epochs; ++epoch) {
		violation_count = 0;
		for (size_t n = 0; n < N; ++n) {
			// Sample an instance id uniformly at random
			unsigned int sample_id = rand_n();

			ClearParameterGradient();
			bool is_violated = ProcessSample(sample_id);
			if (is_violated) {
				UpdateFactorWeights();
				violation_count += 1;
			}

			// Average parameters
			total_iter += 1;
			if (do_averaging) {
				UpdateAveragedParameters(static_cast<double>(total_iter-1)/
					static_cast<double>(total_iter),
					1.0/static_cast<double>(total_iter));
			}
		}

		if (verbose && (epoch % 20 == 0)) {
			std::cout << std::endl;
			std::cout << "  iter     time      violations" << std::endl;
		}
		if (verbose) {
			std::ios_base::fmtflags original_format = std::cout.flags();
			std::streamsize original_prec = std::cout.precision();

			// Iteration
			std::cout << std::setiosflags(std::ios::left)
				<< std::setiosflags(std::ios::adjustfield)
				<< std::setw(6) << epoch << "  ";
			// Total runtime
			std::cout << std::setiosflags(std::ios::left)
				<< std::resetiosflags(std::ios::scientific)
				<< std::setiosflags(std::ios::fixed)
				<< std::setiosflags(std::ios::adjustfield)
				<< std::setprecision(1)
				<< std::setw(6) << total_timer.elapsed() << "s  ";
			std::cout << std::resetiosflags(std::ios::fixed);

			// Objective function
			std::cout << std::resetiosflags(std::ios::scientific)
				<< std::setprecision(5)
				<< std::setiosflags(std::ios::left)
				<< std::setiosflags(std::ios::showpos)
				<< std::setw(7) << violation_count << " of " << N;
			std::cout << std::endl;

			std::cout.precision(original_prec);
			std::cout.flags(original_format);
		}
		if (violation_count == 0)
			break;
	}

	// Use averaged parameters
	if (do_averaging) {
		assert(total_iter >= 1);
		SetFactorWeights();
	}

	return (violation_count);
}

bool StructuredPerceptron::ProcessSample(unsigned int sample_id) {
	FactorGraph* ts_fg = training_data[sample_id].first;
	const FactorGraphObservation* ts_obs = training_data[sample_id].second;
	InferenceMethod* ts_inf = inference_methods[sample_id];
	assert(ts_fg != 0);
	assert(ts_obs != 0);
	assert(ts_inf != 0);

	// Compute forward map: parameters (changed) to energies
	ts_fg->ForwardMap();

	// Compute: -E(y_n;x_n,w)
	/*double obj_truth = */
	lh.ComputeObservationEnergy(ts_fg, ts_obs, parameter_gradient, -1.0);

	// Compute: min_y E(y;x_n,w)
	std::vector<unsigned int> y_star;
	/*double obj_star = */
	ts_inf->MinimizeEnergy(y_star);

	//  ii) Compute gradient (already of negative sign)
	lh.ComputeObservationEnergy(ts_fg, y_star, parameter_gradient, 1.0);
	ts_inf->ClearInferenceResult();
#if 0
	std::cout << "perc, sample " << sample_id << ", true E " << obj_truth
		<< ", y* E " << obj_star << ", viol " << -(obj_truth - obj_star)
		<< std::endl;
#endif

	// Equal energy does not imply no violation (eg. w=0).  Therefore, check
	// violation
	if (ts_obs->Type() == FactorGraphObservation::ExpectationType)
		return (true);

	const std::vector<unsigned int>& true_state = ts_obs->State();
	assert(true_state.size() == y_star.size());
	for (unsigned int vi = 0; vi < true_state.size(); ++vi)
		if (true_state[vi] != y_star[vi])
			return (true);

	return (false);
}

void StructuredPerceptron::UpdateFactorWeights() {
	for (std::vector<std::string>::const_iterator
		ft_name = parameter_order.begin();
		ft_name != parameter_order.end(); ++ft_name) {
		// Get factor type
		FactorType* ft = fg_model->FindFactorType(*ft_name);
		std::transform(parameter_gradient[*ft_name].begin(),
			parameter_gradient[*ft_name].end(), ft->Weights().begin(),
			ft->Weights().begin(), std::plus<double>());
	}
}

void StructuredPerceptron::UpdateAveragedParameters(double old_factor,
	double new_factor) {
	for (std::vector<std::string>::const_iterator
		ft_name = parameter_order.begin();
		ft_name != parameter_order.end(); ++ft_name) {
		// Get factor type
		FactorType* ft = fg_model->FindFactorType(*ft_name);

		// First iteration: old average is zero
		if (parameter_averaged[*ft_name].empty()) {
			parameter_averaged[*ft_name].resize(ft->Weights().size());
			std::fill(parameter_averaged[*ft_name].begin(),
				parameter_averaged[*ft_name].end(), 0.0);
		}

		// Total average equation:
		//   v_0 = 0,
		//   v_t = ((L-1)/L) v_{t-1} + (1/L) w_L.
		std::transform(ft->Weights().begin(), ft->Weights().end(),
			parameter_averaged[*ft_name].begin(),
			parameter_averaged[*ft_name].begin(),
			old_factor * _2 + new_factor * _1);
	}
}

void StructuredPerceptron::SetFactorWeights() {
	for (std::vector<std::string>::const_iterator
		ft_name = parameter_order.begin();
		ft_name != parameter_order.end(); ++ft_name) {
		// Get factor type
		FactorType* ft = fg_model->FindFactorType(*ft_name);
		std::copy(parameter_averaged[*ft_name].begin(),
			parameter_averaged[*ft_name].end(), ft->Weights().begin());
	}
}

void StructuredPerceptron::ClearParameterGradient() {
	parameter_gradient.clear();
	for (std::vector<std::string>::const_iterator
		ft_name = parameter_order.begin();
		ft_name != parameter_order.end(); ++ft_name) {
		FactorType* ft = fg_model->FindFactorType(*ft_name);
		parameter_gradient[*ft_name] =
			std::vector<double>(ft->WeightDimension(), 0.0);
	}
}

}


#ifndef ASLAM_BACKEND_SPARSE_MATRIX_FUNCTIONS_HPP
#define ASLAM_BACKEND_SPARSE_MATRIX_FUNCTIONS_HPP

#include <sparse_block_matrix/linear_solver.h>
#include <boost/shared_ptr.hpp>
#include <eigen3/Eigen/Dense>

namespace aslam {
  namespace backend {

    void applySchurComplement(sparse_block_matrix::SparseBlockMatrix<Eigen::MatrixXd>& H,
                              const Eigen::VectorXd& e,
                              double lambda,
                              int marginalizedStartingBlock,
                              bool doLevenberg,
                              sparse_block_matrix::SparseBlockMatrix<Eigen::MatrixXd>& A,
                              std::vector<Eigen::MatrixXd>& invVi,
                              Eigen::VectorXd& b);

    void buildDsi(int i,
                  sparse_block_matrix::SparseBlockMatrix<Eigen::MatrixXd>& H,
                  const Eigen::VectorXd& e, int marginalizedStartingBlock, const Eigen::MatrixXd& invVi,
                  const Eigen::VectorXd& dx, Eigen::VectorXd& outDsi);

  } // namespace backend
} // namespace aslam

#endif /* ASLAM_BACKEND_SPARSE_MATRIX_FUNCTIONS_HPP */


/*
	Copyright (c) 2019,	Mobile Robots Laboratory:
	-Jan Wietrzykowski (jan.wietrzykowski@put.poznan.pl).
	Poznan University of Technology
	All rights reserved.

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

	1. Redistributions of source code must retain the above copyright notice,
	this list of conditions and the following disclaimer.

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

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

#include <iostream>
#include <thread>
#include <random>
#include <chrono>
#include <string>
#include <fstream>

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

#include "pgm/ParamEstSGD.h"
#include "pgm/Exceptions.h"
#include "pgm/Inference.h"

using namespace std;

std::vector<double> ParamEstSGD::estimate(std::shared_ptr<PgmCreator> pgmCreator,
										Params estParams,
										const std::vector<double>& initParams)
{
	static const int maxThreads = 4;

	std::default_random_engine randGen;

	int m = pgmCreator->getM();

	//generate random permutation
	vector<int> perm;
	for(int i = 0; i < m; ++i){
		perm.push_back(i);
	}
	for(int i = 0; i < m - 1; ++i){
		std::uniform_int_distribution<int> dist(i, m - 1);
		int swapIdx = dist(randGen);
		int tmp = perm[i];
		perm[i] = perm[swapIdx];
		perm[swapIdx] = tmp;
	}

	int mVal = estParams.valSetFrac * perm.size();
	vector<int> permVal = vector<int>(perm.begin(), perm.begin() + mVal);

	int mTrain = perm.size() - mVal;
	vector<int> permTrain = vector<int>(perm.begin() + mVal, perm.end());

	vector<double> curParams = initParams;

	vector<double> paramsGrad;

	int iter = 0;
	bool stopFlag = false;
	while(!stopFlag){
		cout << "iteration " << iter << endl;
		cout << "curParams = " << curParams << endl;

		double curLearnRate = 1 / ((estParams.sigma * estParams.sigma) * (estParams.learnRateM0 + iter));
		cout << "curLearnRate = " << curLearnRate << endl;

		//generate random permutation
		vector<int> curPerm = permTrain;

		for(int i = 0; i < mTrain - 1; ++i){
			std::uniform_int_distribution<int> dist(i, mTrain - 1);
			int swapIdx = dist(randGen);
			int tmp = curPerm[i];
			curPerm[i] = curPerm[swapIdx];
			curPerm[swapIdx] = tmp;
		}

		vector<double> lhoodNumers;
		vector<double> lhoodDenoms;
		vector<vector<double>> Efis;
		vector<vector<double>> Eds;

		//take first n samples
		evaluateSamples(vector<int>(curPerm.begin(), curPerm.begin() + estParams.n),
						maxThreads,
						pgmCreator,
						curParams,
						lhoodNumers,
						lhoodDenoms,
						Efis,
						Eds);

		vector<double> curGrad(curParams.size(), 0.0);
		// compute current gradient
		for(int i = 0; i < Efis.size(); ++i){
			for(int p = 0; p < curGrad.size(); ++p){
				curGrad[p] = Eds[i][p] - Efis[i][p];
			}
		}
		for(int p = 0; p < curGrad.size(); ++p){
			curGrad[p] /= Efis.size();
		}
		for(int p = 0; p < curGrad.size(); ++p){
			curGrad[p] -= curParams[p] / (estParams.sigma * estParams.sigma);
		}

		cout << "curGrad = " << curGrad << endl;

		// if it is first iteration and previous gradient is empty
		if(paramsGrad.empty()){
			paramsGrad = curGrad;
		}

		for(int p = 0; p < paramsGrad.size(); ++p){
			paramsGrad[p] = curLearnRate * (paramsGrad[p]*estParams.momentum + curGrad[p]);
		}

		for(int p = 0; p < curParams.size(); ++p){
			curParams[p] += paramsGrad[p];
		}

		//evaluate on validation set
		if(iter % 10 == 0){
			double scoreVal = 0.0;

			vector<double> lhoodNumers;
			vector<double> lhoodDenoms;
			vector<vector<double>> Efis;
			vector<vector<double>> Eds;

			//take first n samples
			evaluateSamples(permVal,
							maxThreads,
							pgmCreator,
							curParams,
							lhoodNumers,
							lhoodDenoms,
							Efis,
							Eds);

			for(int i = 0; i < lhoodNumers.size(); ++i){
				cout << "scoreVal += " << lhoodNumers[i] << " - " << lhoodDenoms[i] << endl;
				scoreVal += lhoodNumers[i] - lhoodDenoms[i];
			}
			scoreVal /= lhoodNumers.size();

			cout << endl << endl << "scoreVal = " << scoreVal << endl << endl << endl;

			//saving a snapshot
			{
				ofstream paramsFile(string("log/params_") + to_string(iter));
				boost::archive::text_oarchive paramsArch(paramsFile);
				paramsArch << curParams;
			}
			{
				ofstream gradFile(string("log/grad_") + to_string(iter));
				boost::archive::text_oarchive gradArch(gradFile);
				gradArch << curGrad;
			}
			{
				ofstream scoreFile(string("log/score"), ios::app);
				scoreFile << iter << " " << scoreVal << endl;
			}
		}

		//stop condition
		//TODO use better stop condition
		if(curLearnRate < 0.005){
			stopFlag = true;
		}

		iter++;
	}

	return curParams;
}

void ParamEstSGD::evaluateSamples(const std::vector<int>& sampIdxs,
							int maxThreads,
							std::shared_ptr<PgmCreator> pgmCreator,
							const std::vector<double>& curParams,
							std::vector<double>& lhoodNumers,
							std::vector<double>& lhoodDenoms,
							std::vector<std::vector<double>>& Efis,
							std::vector<std::vector<double>>& Eds)
{
	int threadCnt = 0;
	vector<shared_ptr<ParamEstSGD::EvaluateThread>> threads(sampIdxs.size());
	vector<shared_ptr<Pgm>> pgms(sampIdxs.size());
	vector<shared_ptr<vector<double>>> obsVecs(sampIdxs.size());
	vector<shared_ptr<vector<double>>> varVals(sampIdxs.size());
	vector<shared_ptr<vector<int>>> varIds(sampIdxs.size());
	vector<shared_ptr<vector<double>>> paramVals(sampIdxs.size());
	lhoodNumers.resize(sampIdxs.size());
	lhoodDenoms.resize(sampIdxs.size());
	Efis.resize(sampIdxs.size());
	Eds.resize(sampIdxs.size());

	//take first n elements from permutation
	for(int i = 0; i < sampIdxs.size(); ++i){
		while(threadCnt == maxThreads){
			endThreadIfAny(threadCnt,
							threads,
							pgms,
							obsVecs,
							varVals,
							varIds,
							paramVals);

			std::this_thread::sleep_for(chrono::milliseconds(50));
		}

		pgms[i] = shared_ptr<Pgm>(new Pgm);
		obsVecs[i] = shared_ptr<vector<double>>(new vector<double>());
		varVals[i] = shared_ptr<vector<double>>(new vector<double>());
		varIds[i] = shared_ptr<vector<int>>(new vector<int>());
		paramVals[i] = shared_ptr<vector<double>>(new vector<double>(curParams));
		lhoodNumers[i] = 0.0;
		lhoodDenoms[i] = 0.0;
		Efis[i] = vector<double>(curParams.size(), 0.0);
		Eds[i] = vector<double>(curParams.size(), 0.0);

		pgmCreator->create(sampIdxs[i],
						*pgms[i],
						*obsVecs[i],
						*varVals[i],
						*varIds[i]);

//		for(int j = 0; j < Efis.size(); ++j){
//			cout << "Efis[" << j << "].data() = " << Efis[j].data() << endl;
//			cout << "&Efis[" << j << "] = " << &Efis[j] << endl;
//		}

		threads[i] = shared_ptr<EvaluateThread>(new EvaluateThread(pgms[i],
												obsVecs[i],
												varVals[i],
												varIds[i],
												paramVals[i],
												lhoodNumers[i],
												lhoodDenoms[i],
												Efis[i],
												Eds[i]));
		++threadCnt;
	}

	while(threadCnt > 0){
		endThreadIfAny(threadCnt,
						threads,
						pgms,
						obsVecs,
						varVals,
						varIds,
						paramVals);
		std::this_thread::sleep_for(chrono::milliseconds(50));
	}
}

void ParamEstSGD::endThreadIfAny(int& threadCnt,
						std::vector<std::shared_ptr<ParamEstSGD::EvaluateThread>>& threads,
						std::vector<std::shared_ptr<Pgm>>& pgms,
						std::vector<std::shared_ptr<std::vector<double>>>& obsVecs,
						std::vector<std::shared_ptr<std::vector<double>>>& varVals,
						std::vector<std::shared_ptr<std::vector<int>>>& varIds,
						std::vector<std::shared_ptr<std::vector<double>>>& paramVals)
{
	for(int t = 0; t < threads.size(); ++t){
		if(threads[t]){
			if(threads[t]->hasFinished()){
				cout << "Ending thread " << t << endl;
				threads[t].reset();
				pgms[t].reset();
				obsVecs[t].reset();
				varVals[t].reset();
				varIds[t].reset();
				paramVals[t].reset();

				--threadCnt;
			}
		}
	}
}

ParamEstSGD::EvaluateThread::EvaluateThread(std::shared_ptr<Pgm> ipgm,
										std::shared_ptr<const std::vector<double>> iobsVec,
										std::shared_ptr<const std::vector<double>> ivarVals,
										std::shared_ptr<const std::vector<int>> ivarIds,
										std::shared_ptr<const std::vector<double>> iparamVals,
										double& ilhoodNumer,
										double& ilhoodDenom,
										std::vector<double>& iEfi,
										std::vector<double>& iEd)
		: pgm(ipgm),
		  obsVec(iobsVec),
		  varVals(ivarVals),
		  varIds(ivarIds),
		  paramVals(iparamVals),
		  lhoodNumer(ilhoodNumer),
		  lhoodDenom(ilhoodDenom),
		  Efi(iEfi),
		  Ed(iEd),
		  finishedFlag(false)
{
//	cout << "Efi.data() constr = " << Efi.data() << ", this = " << this << endl;
//	cout << "&Efi constr = " << &Efi << ", this = " << this << endl;
	runThread = thread(&ParamEstSGD::EvaluateThread::run, this);
}

ParamEstSGD::EvaluateThread::~EvaluateThread()
{
	if(runThread.joinable()){
		runThread.join();
	}
}

bool ParamEstSGD::EvaluateThread::hasFinished()
{
	return finishedFlag;
}

void ParamEstSGD::EvaluateThread::run()
{
	// inference without clamped variables - denominator and model expectation
	{
//		cout << "Efi.data() 1 = " << Efi.data() << ", this = " << this <<  endl;
//		cout << "&Efi 1 = " << &Efi << ", this = " << this << endl;

		vector<vector<double>> marg;
		vector<vector<vector<double>>> msgs;
		double logPartFunc;

//		cout << "Efi.data() 2 = " << Efi.data() << ", this = " << this << endl;
		bool calibrated = Inference::compMarginalsParam(*pgm, marg, logPartFunc, msgs, *paramVals, *obsVec);
//		cout << "Efi.data() 3 = " << Efi.data() << endl;

		if(!calibrated){
			cout << "Warning - inference without clamped variables didn't calibrate" << endl;
		}

		lhoodDenom = logPartFunc;
		cout << "loodDenom = " << lhoodDenom << endl;

		for(int c = 0; c < (int)pgm->constClusters().size(); ++c){
//			cout << "Cluster " << c << endl;
			std::shared_ptr<Cluster> curCluster = pgm->clusters()[c];
//			cout << "curCluster->id() = " << curCluster->id() << endl;
			const vector<vector<double> >& curClustMsgs = msgs[curCluster->id()];
//			cout << "curClustMsgs = " << curClustMsgs << endl;
			vector<double> curEfi = curCluster->compSumModelExpectation(curClustMsgs, *paramVals, *obsVec);
//			cout << "Efi.size() = " << Efi.size() << ", curEfi.size() = " << curEfi.size() << endl;

			for(int f = 0; f < (int)curCluster->feats().size(); ++f){
//				cout << "f = " << f << endl;
				int paramNum = curCluster->feats()[f]->paramNum();
//				cout << "paramNum = " << paramNum << endl;
				Efi[paramNum] += curEfi[f];
			}
		}
	}

	// clamping variables
//	cout << "clamping variables" << endl;
	int posRv = 0;
	for(int v = 0; v < varIds->size(); ++v){
		bool stopFlag = false;
		while(!stopFlag && posRv < pgm->randVars().size() - 1){
			if(pgm->randVars()[posRv]->id() < varIds->at(v)){
				++posRv;
			}
			else{
				stopFlag = true;
			}
		}
		if(pgm->randVars()[posRv]->id() == varIds->at(v)){
			pgm->randVars()[posRv]->makeObserved(varVals->at(v));
		}
		else{
			cout << "pgm->randVars()[posRv]->id() = " << pgm->randVars()[posRv]->id() << endl;
			cout << "varIds->at(v) = " << varIds->at(v) << endl;
			throw PGM_EXCEPTION("Rand var not found");
		}
	}
//	cout << "end clamping variables" << endl;

	// inference with clamped variables - numerator and distribution expectation
	{
		vector<vector<double>> margClamp;
		vector<vector<vector<double>>> msgsClamp;
		double logPartFuncClamp;

		bool calibratedClamp = Inference::compMarginalsParam(*pgm, margClamp, logPartFuncClamp, msgsClamp, *paramVals, *obsVec);

		if(!calibratedClamp){
			cout << "Warning - inference with clamped variables didn't calibrate" << endl;
		}

		lhoodNumer = logPartFuncClamp;
		cout << "loodNumer = " << lhoodNumer << endl;

		// a model expectation with clamped variables is a distribution expectation
		for(int c = 0; c < (int)pgm->constClusters().size(); ++c){
//			cout << "Cluster " << c << endl;
			std::shared_ptr<Cluster> curCluster = pgm->clusters()[c];
			const vector<vector<double> >& curClustMsgs = msgsClamp[curCluster->id()];
			vector<double> curEfi = curCluster->compSumModelExpectation(curClustMsgs, *paramVals, *obsVec);

			for(int f = 0; f < (int)curCluster->feats().size(); ++f){
				int paramNum = curCluster->feats()[f]->paramNum();
				Ed[paramNum] += curEfi[f];
			}
		}
	}

	finishedFlag = true;
}


#include "../../src/mesh/mesh_generator.h"
#include "../test_utilities.h"

#include <deal.II/base/types.h>
#include <deal.II/grid/grid_out.h>

template <int dim>
void SetupParameters (dealii::ParameterHandler &prm) {
  prm.declare_entry ("is mesh generated by deal.II", "true",
                     dealii::Patterns::Bool(), "");
  prm.declare_entry ("is mesh pin-resolved", "true",
                     dealii::Patterns::Bool(), "");
  prm.declare_entry ("have reflective BC", "true",
                     dealii::Patterns::Bool(), "");
  prm.declare_entry ("fuel rod radius", "0.4",
                     dealii::Patterns::Double(), "");
  prm.declare_entry ("fuel rod triangulation type", "simple",
                     dealii::Patterns::Selection("simple|composite"), "");
  prm.declare_entry ("uniform refinements", "0",
                     dealii::Patterns::Integer(), "");
  prm.declare_entry ("x, y, z max values of boundary locations", "3.0,3.0,3.0",
                     dealii::Patterns::List (dealii::Patterns::Double ()), "");
  prm.declare_entry ("number of cells for x, y, z directions", "3,3,3",
                     dealii::Patterns::List (dealii::Patterns::Integer ()), "");
  prm.declare_entry ("reflective boundary names", "xmin, ymax",
                     dealii::Patterns::List(dealii::Patterns::Anything ()), "");
  prm.enter_subsection ("material ID map");
  {
    std::string id_fname = SOURCE_DIR + std::string ("/matid.unstruct.")
        + std::to_string (dim) + std::string ("d");
    prm.declare_entry ("material id file name", id_fname,
                       dealii::Patterns::FileName(), "file name for material id map");
    std::string fuel_id_fname = SOURCE_DIR + std::string ("/fuelid.unstruct.")
        + std::to_string (dim) + std::string ("d");
    prm.declare_entry ("fuel pin material id file name", fuel_id_fname,
                       dealii::Patterns::FileName(), "file name for material id map");
  }
  prm.leave_subsection ();
}

template <int dim>
void Test (dealii::ParameterHandler &prm) {
  dealii::deallog.push (dealii::Utilities::int_to_string(dim)+"D");

  // triangulation for the grid in the scope of the test function
  dealii::parallel::distributed::Triangulation<dim> tria (MPI_COMM_WORLD);

  // make grid
  MeshGenerator<dim> meshie (prm);
  meshie.MakeGrid (tria);

  for (typename dealii::Triangulation<dim>::active_cell_iterator
       cell=tria.begin_active(); cell!=tria.end(); ++cell)
    dealii::deallog << "cell ID: " << cell->id()
                    << ":material ID::" << static_cast<int>(cell->material_id())
                    << "::center x::" << cell->center()[0]
                    << "::center y::" << cell->center()[1] << std::endl;

  std::ofstream out("unstruct_mesh_2d.vtk");
  dealii::GridOut grid_out;
  grid_out.write_vtk (tria, out);
}

int main (int argc, char *argv[]) {
  // initialize MPI and log and declare ParameterHandler object
  dealii::Utilities::MPI::MPI_InitFinalize mpi_initialization (argc, argv, 1);
  testing::MPILogInit init;
  dealii::ParameterHandler prm;

  dealii::deallog.push ("2D");
  // parameter processing
  SetupParameters<2> (prm);
  dealii::deallog.pop();

  // testing 2D
  Test<2> (prm);

  return 0;
}


// Copyright (c) 2015-2016, ETH Zurich, Wyss Zurich, Zurich Eye
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//     * Neither the name of the ETH Zurich, Wyss Zurich, Zurich Eye nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL ETH Zurich, Wyss Zurich, Zurich Eye BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include <cstdint>
#pragma diagnostic push
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
// Eigen 3.2.7 uses std::binder1st and std::binder2nd which are deprecated since c++11
// Fix is in 3.3 devel (http://eigen.tuxfamily.org/bz/show_bug.cgi?id=872).
#include <Eigen/Core>
#pragma diagnostic pop
#include <Eigen/StdVector>
#include <ze/common/config.hpp>

namespace ze {

//------------------------------------------------------------------------------
// Scalars and fp precision.
using size_t    = std::size_t;
using int8_t    = std::int8_t;
using int16_t   = std::int16_t;
using int64_t   = std::int64_t;
using uint8_t   = std::uint8_t;
using uint16_t  = std::uint16_t;
using uint32_t  = uint32_t;
using uint64_t  = std::uint64_t;
#ifdef ZE_SINGLE_PRECISION_FLOAT
using real_t = float;
#else
using real_t = double;
#endif

//------------------------------------------------------------------------------
// Typedefs of commonly used Eigen matrices and vectors.

// MatrixMN, MatrixN = MatrixNN, I_NxN, and Z_NxN, for M,N=1..9.
#define ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(SIZE, SUFFIX)            \
  using Matrix##SUFFIX = Eigen::Matrix<real_t, SIZE, SIZE>; \
  using Matrix1##SUFFIX = Eigen::Matrix<real_t, 1, SIZE>;   \
  using Matrix2##SUFFIX = Eigen::Matrix<real_t, 2, SIZE>;   \
  using Matrix3##SUFFIX = Eigen::Matrix<real_t, 3, SIZE>;   \
  using Matrix4##SUFFIX = Eigen::Matrix<real_t, 4, SIZE>;   \
  using Matrix5##SUFFIX = Eigen::Matrix<real_t, 5, SIZE>;   \
  using Matrix6##SUFFIX = Eigen::Matrix<real_t, 6, SIZE>;   \
  using Matrix7##SUFFIX = Eigen::Matrix<real_t, 7, SIZE>;   \
  using Matrix8##SUFFIX = Eigen::Matrix<real_t, 8, SIZE>;   \
  using Matrix9##SUFFIX = Eigen::Matrix<real_t, 9, SIZE>;   \
  using Matrix##SUFFIX##X = Eigen::Matrix<real_t, SIZE, Eigen::Dynamic>; \
  using MatrixX##SUFFIX = Eigen::Matrix<real_t, Eigen::Dynamic, SIZE>;   \
  static const Eigen::MatrixBase<Matrix##SUFFIX>::IdentityReturnType I_##SUFFIX##x##SUFFIX = Matrix##SUFFIX::Identity(); \
  static const Eigen::MatrixBase<Matrix##SUFFIX>::ConstantReturnType Z_##SUFFIX##x##SUFFIX = Matrix##SUFFIX::Zero()

ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(1,1);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(2,2);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(3,3);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(4,4);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(5,5);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(6,6);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(7,7);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(8,8);
ZE_MAKE_EIGEN_MATRIX_TYPEDEFS(9,9);

// Typedef arbitary length vector and arbitrary sized matrix.
using VectorX = Eigen::Matrix<real_t, Eigen::Dynamic, 1>;
using MatrixX = Eigen::Matrix<real_t, Eigen::Dynamic, Eigen::Dynamic>;
using VectorXi = Eigen::VectorXi;

// Commonly used fixed size vectors.
using Vector1 = Eigen::Matrix<real_t, 1, 1>;
using Vector2 = Eigen::Matrix<real_t, 2, 1>;
using Vector3 = Eigen::Matrix<real_t, 3, 1>;
using Vector4 = Eigen::Matrix<real_t, 4, 1>;
using Vector5 = Eigen::Matrix<real_t, 5, 1>;
using Vector6 = Eigen::Matrix<real_t, 6, 1>;
using Vector7 = Eigen::Matrix<real_t, 7, 1>;
using Vector8 = Eigen::Matrix<real_t, 8, 1>;
using Vector9 = Eigen::Matrix<real_t, 9, 1>;
using Vector2i = Eigen::Vector2i;

//------------------------------------------------------------------------------
// Feature containers.
using Keypoint    = Vector2;
using Bearing     = Vector3;
using Position    = Vector3;
using HomPosition = Vector4;
using Gradient    = Vector2;
using Seed        = Vector4;
using LineMeasurement = Vector3;
using Keypoints   = Matrix2X;
using Bearings    = Matrix3X;
using Positions   = Matrix3X;
using HomPositions = Matrix4X;
using Gradients   = Matrix2X;
using Seeds       = Matrix4X;
//! Normal vector on line end-points bearings, as explained in
//! ze_geometry/doc/line_parametrization.pdf
using LineMeasurements = Matrix3X;
using KeypointLevel = int8_t;
using KeypointType  = int8_t;
using KeypointIndex = uint16_t;
using KeypointLevels = Eigen::Matrix<KeypointLevel, Eigen::Dynamic, 1>;
using KeypointTypes  = Eigen::Matrix<KeypointType, Eigen::Dynamic, 1>;
using KeypointAngles = VectorX;
using KeypointScores = VectorX;
using KeypointSizes  = VectorX;
using KeypointIndices  = Eigen::Matrix<KeypointIndex, Eigen::Dynamic, 1>;
using Descriptors = Eigen::Matrix<uint8_t, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;

//------------------------------------------------------------------------------
// Inertial containers.
using ImuStamps = Eigen::Matrix<int64_t, Eigen::Dynamic, 1>;
using ImuAccGyrContainer = Matrix6X;
// Order: Accelerometer, Gyroscope
using ImuAccGyr = Vector6;

} // namespace ze


#pragma once

#include <crab/analysis/abs_transformer.hpp>
#include <crab/analysis/dataflow/liveness.hpp>
#include <crab/analysis/fwd_analyzer.hpp>
#include <crab/analysis/graphs/dominance.hpp>
#include <crab/cfg/cfg.hpp>
#include <crab/cfg/var_factory.hpp>
#include <crab/iterators/interleaved_fixpoint_iterator.hpp>

#include <boost/range/iterator_range.hpp>

#include <algorithm>
#include <memory>
#include <set>
#include <unordered_map>
#include <vector>

namespace crab {

namespace analyzer {

/**
 * Compute necessary preconditions by computing the abstract least
 * fixpoint of the reversed CFG. Optionally, at each step the
 * precondition can be refined with an invariant computed by a
 * forward analysis.
 **/
template <typename CFG, typename AbsDom>
class necessary_preconditions_fixpoint_iterator
    : private ikos::interleaved_fwd_fixpoint_iterator<crab::cfg::cfg_rev<CFG>,
                                                      AbsDom> {

  typedef ikos::interleaved_fwd_fixpoint_iterator<crab::cfg::cfg_rev<CFG>,
                                                  AbsDom>
      fixpoint_iterator_t;
  typedef typename CFG::basic_block_label_t bb_label_t;
  typedef typename CFG::statement_t stmt_t;
  typedef std::unordered_map<bb_label_t, AbsDom> bb_abstract_map_t;
  typedef std::unordered_map<stmt_t *, AbsDom> pp_abstract_map_t;

  /// forward and backward transformers
  typedef intra_abs_transformer<AbsDom> abs_fwd_tr_t;
  typedef intra_necessary_preconditions_abs_transformer<AbsDom,
                                                        pp_abstract_map_t>
      abs_bwd_tr_t;

  CFG m_cfg;
  // postcondition (i.e, final states) that we want to propagate backwards
  AbsDom m_postcond;
  // necessary preconditions
  bb_abstract_map_t m_preconditions;
  // forward invariants
  bb_abstract_map_t m_invariants;
  // preconditions from good states, otherwise from bad states
  bool m_good_states;

  /**
   * Compute necessary preconditions for a basic block
   **/
  virtual AbsDom analyze(bb_label_t node, AbsDom &&precond) override {
    auto &bb = m_cfg.get_node(node);

    CRAB_LOG("backward-fixpoint", crab::outs() << "Post at "
                                               << cfg_impl::get_label_str(node)
                                               << ": " << precond << "\n");

    // invariants that hold at the entry of the block
    AbsDom invariant = m_invariants[node];
    // rebuild local invariants that hold at each program point.
    abs_fwd_tr_t F(invariant);
    pp_abstract_map_t pp_invariants;
    for (auto &s : boost::make_iterator_range(bb.begin(), bb.end())) {
      auto inv = F.get_abs_value();
      pp_invariants.insert(std::make_pair(&s, inv));
      CRAB_LOG("backward-fixpoint", crab::outs()
                                        << "\tRebuilding at statement " << s
                                        << " inv=" << inv << "\n");
      s.accept(&F);
    }

    CRAB_LOG("backward-fixpoint",
             crab::outs() << "Done forward propagation at each program point \n"
                          << "Starting backward propagation ... \n");

    // compute precondition at the entry of the block
    abs_bwd_tr_t B(std::move(precond), &pp_invariants, m_good_states);
    for (auto &s : boost::make_iterator_range(bb.rbegin(), bb.rend())) {
      s.accept(&B);
    }
    precond = std::move(B.preconditions());
    CRAB_LOG("backward-fixpoint", crab::outs() << "Pre at "
                                               << cfg_impl::get_label_str(node)
                                               << ": " << precond << "\n");
    return std::move(precond);
  }

  virtual void process_pre(bb_label_t /*node*/, AbsDom /*postcond*/) override {}

  /**
   *  Store necessary preconditions
   **/
  virtual void process_post(bb_label_t node, AbsDom precond) override {
    m_preconditions.insert(std::make_pair(node, precond));
  }

public:
  typedef typename fixpoint_iterator_t::wto_t wto_t;
  typedef bb_abstract_map_t precond_map_t;
  typedef typename precond_map_t::iterator iterator;
  typedef typename precond_map_t::const_iterator const_iterator;

  // This constructor computes necessary preconditions from error
  // states.
  necessary_preconditions_fixpoint_iterator(
      CFG cfg, AbsDom postcond,
      /* fixpoint parameters */
      unsigned int widening_delay = 1,
      unsigned int descending_iterations = UINT_MAX, size_t jump_set_size = 0)
      : fixpoint_iterator_t(crab::cfg::cfg_rev<CFG>(cfg), nullptr,
                            widening_delay, descending_iterations,
                            jump_set_size),
        m_cfg(cfg), m_postcond(postcond), m_good_states(false) {}

  // This constructor computes necessary preconditions from
  // safe/good (error) states if good_states is true (false).
  necessary_preconditions_fixpoint_iterator(
      CFG cfg, const wto_t *wto, AbsDom postcond, bool good_states,
      /* fixpoint parameters */
      unsigned int widening_delay = 1,
      unsigned int descending_iterations = UINT_MAX, size_t jump_set_size = 0)
      : fixpoint_iterator_t(crab::cfg::cfg_rev<CFG>(cfg), wto, widening_delay,
                            descending_iterations, jump_set_size),
        m_cfg(cfg), m_postcond(postcond), m_good_states(good_states) {}

  void run_backward() { this->run(m_postcond); }

  void run_backward(const std::unordered_map<typename CFG::basic_block_label_t,
                                             AbsDom> &fwd_invariants) {
    m_invariants.insert(fwd_invariants.begin(), fwd_invariants.end());
    run_backward();
  }

  iterator begin() { return m_preconditions.begin(); }

  iterator end() { return m_preconditions.end(); }

  const_iterator begin() const { return m_preconditions.begin(); }

  const_iterator end() const { return m_preconditions.end(); }

  // return the preconditions at basic block node
  AbsDom operator[](bb_label_t node) const {
    auto it = m_preconditions.find(node);
    if (it != m_preconditions.end())
      return it->second;
    else
      return AbsDom::top();
  }

  // clear preconditions and forward invariants (if any)
  void clear() {
    m_preconditions.clear();
    m_invariants.clear();
  }

  const wto_t &get_WTO() const { return this->get_wto(); }
};

/**
 * A refining forward-backward analyzer based on Cousot&Cousot
 * (JLP'92 and ASE'99, section 4).
 *
 * Do first a forward pass computing superset of reachable states
 * from initial states. Then do a backward pass, computing
 * superset of reachable meet with co-reachable from error
 * states. Repeat this until no more refinement or maximum of
 * iterations reached.
 *
 * The implementation follows mc_5 algorithm described in page 82,
 * ASE'99. It computes an over-approximation of the intersection
 * between the set of reachable states starting from the initial
 * states and the set of co-reachable states from error
 * states.
 *
 * This forward-backward refinement implementation produces the
 * same invariants than a classical forward analyzer. The
 * difference is that it can prove more assertions.
 **/
template <typename CFG, typename AbsDom> class intra_forward_backward_analyzer {
public:
  typedef CFG cfg_t;
  typedef typename CFG::basic_block_t bb_t;
  typedef typename CFG::basic_block_label_t bb_label_t;
  typedef typename CFG::statement_t stmt_t;
  typedef typename CFG::varname_t varname_t;
  typedef typename CFG::number_t number_t;
  // used for checkers
  typedef AbsDom abs_dom_t;
  typedef intra_abs_transformer<AbsDom> abs_tr_t;

private:
  typedef intra_fwd_analyzer<CFG, AbsDom> fwd_analyzer_t;
  typedef necessary_preconditions_fixpoint_iterator<CFG, AbsDom>
      bwd_fixpoint_iterator_t;
  typedef typename bwd_fixpoint_iterator_t::precond_map_t precond_map_t;
  typedef typename bwd_fixpoint_iterator_t::wto_t bwd_wto_t;
  typedef liveness<CFG> liveness_t;
  typedef std::unordered_map<bb_label_t, std::set<bb_label_t>> idom_tree_t;
  typedef typename bb_t::assert_t assert_t;
  typedef typename bb_t::bool_assert_t bool_assert_t;
  typedef typename bb_t::ptr_assert_t ptr_assert_t;

public:
  typedef typename fwd_analyzer_t::assumption_map_t assumption_map_t;
  typedef typename fwd_analyzer_t::invariant_map_t invariant_map_t;
  // bwd_wto_t and wto_t are different types because bwd_wto_t is
  // over the reversed CFG.
  typedef typename fwd_analyzer_t::wto_t wto_t;

private:
  // -- the cfg
  CFG m_cfg;
  // we keep the two wto's (from forward and reversed CFGs) to
  // avoid recompute them during the below iterative process.
  // Only the forward wto is exposed to outside clients.
  const wto_t *m_wto;
  const bwd_wto_t *m_b_wto;
  // keep track of which assertions cannot be proven.
  std::vector<std::pair<bb_label_t, stmt_t *>> m_unproven_assertions;
  // keep track of which assertions have been proven
  std::set<stmt_t *> m_proved_assertions;
  // keep the results of the first forward iteration.
  invariant_map_t m_pre_invariants;
  invariant_map_t m_post_invariants;

  void store_analysis_results(fwd_analyzer_t &f) {
    for (auto &kv : f.get_pre_invariants()) {
      m_pre_invariants.insert({kv.first, std::move(kv.second)});
    }

    for (auto &kv : f.get_post_invariants()) {
      m_post_invariants.insert({kv.first, std::move(kv.second)});
    }
  }

  void gather_assertions() {
    for (auto it = m_cfg.begin(), et = m_cfg.end(); it != et; ++it) {
      for (auto &s : *it) {
        if (s.is_assert() || s.is_bool_assert() || s.is_ptr_assert()) {
          m_unproven_assertions.push_back({it->label(), &s});
        }
      }
    }
  }

  // idom induces a tree each key-value pair means that key
  // (block) is an immediate dominator of each element in value
  // (set of basic blocks). TODO: caching
  bool dominates(bb_label_t u, bb_label_t v, const idom_tree_t &idom) {
    auto it = idom.find(u);
    if (it == idom.end()) {
      // u does not dominate any block
      return false;
    }

    // immediate dominated by u
    const std::set<bb_label_t> &idom_u = it->second;
    if (idom_u.count(v) > 0) {
      // u dominates v
      return true;
    }

    for (const bb_label_t w : idom_u) {
      // a successor of u dominates v
      if (dominates(w, v, idom)) {
        return true;
      }
    }

    return false;
  }

  /**
   * For each pair (bb, dom) in refined_assumptions we know that
   * dom is the meet of the over-approximated set of reachable
   * states from bb and the over-approximated co-reachable states
   * from the error states. Therefore, if dom is bottom then any
   * assertion dominated (in the sense of CFG dominance) by bb
   * must hold and hence, it is removed from
   * m_unproven_assertions and added to m_proved_assertions.
   **/
  void discharge_assertions(assumption_map_t &refined_assumptions,
                            const idom_tree_t &idom) {
    if (m_unproven_assertions.empty()) {
      return;
    }

    if (!idom.empty()) {
      // If it's not possible starting from the initial states to
      // reach a block B that leads to an assertion violation then
      // we can discharge all the assertions dominated by B.
      //
      // Once we know that an assertion must hold, we could add in
      // m_pre_invariants an extra constraint from the assertion
      // condition. The problem is that the condition might not be
      // expressed in the abstract domain (e.g., a disequality). So
      // that wouldn't be enough for the checker to discharge the
      // assertion, even if we know already the assertion must
      // hold. Instead, we tell explicitly the checker which
      // assertions were proved by calling safe_assertions method.

      for (auto it = m_cfg.begin(), et = m_cfg.end(); it != et; ++it) {
        bb_label_t n = it->label();
        auto jt = refined_assumptions.find(n);
        if (jt == refined_assumptions.end()) {
          continue;
        }
        if (jt->second.is_bottom()) {
          m_unproven_assertions.erase(
              std::remove_if(
                  m_unproven_assertions.begin(), m_unproven_assertions.end(),
                  [&n, &idom, this](std::pair<bb_label_t, stmt_t *> &kv) {
                    bb_label_t m = kv.first;
                    stmt_t *s = kv.second;
                    bool res = this->dominates(n, m, idom);
                    if (res) {
                      this->m_proved_assertions.insert(s);
                    }
                    return res;
                  }),
              m_unproven_assertions.end());
        }
      }
    } else {
      // if dominance information is not available then we only
      // discharge all assertions only if starting from the entry
      // block we know that we cannot violate them.
      auto it = refined_assumptions.find(m_cfg.entry());
      if (it->second.is_bottom()) {
        for (auto &kv : m_unproven_assertions) {
          m_proved_assertions.insert(kv.second);
        }
        m_unproven_assertions.clear();
      }
    }
  }

public:
  typedef typename bwd_fixpoint_iterator_t::iterator iterator;
  typedef typename bwd_fixpoint_iterator_t::const_iterator const_iterator;

  intra_forward_backward_analyzer(CFG cfg)
      : m_cfg(cfg), m_wto(nullptr), m_b_wto(nullptr) {}

  ~intra_forward_backward_analyzer() {
    if (m_wto)
      delete m_wto;
    if (m_b_wto)
      delete m_b_wto;
  }

  intra_forward_backward_analyzer(
      const intra_forward_backward_analyzer<CFG, AbsDom> &o) = delete;
  intra_forward_backward_analyzer<CFG, AbsDom> &
  operator=(const intra_forward_backward_analyzer<CFG, AbsDom> &o) = delete;

  /**
   * Perform the refining forward-backward loop.
   **/
  void run(AbsDom init_states,
           // behaves as a standard forward analysis
           bool only_forward,
           // assumptions
           const assumption_map_t &assumptions,
           // liveness information
           const liveness_t *live,
           // parameters for each forward or backward analysis
           unsigned int widening_delay = 1,
           unsigned int descending_iters = UINT_MAX, size_t jump_set_size = 0) {
    run(m_cfg.entry(), init_states, only_forward, assumptions, live,
        widening_delay, descending_iters, jump_set_size);
  }

  void run(bb_label_t entry, // only used for the forward pass.
           AbsDom init_states,
           // behaves as a standard forward analysis
           bool only_forward,
           // assumptions
           const assumption_map_t &assumptions,
           // liveness information
           const liveness_t *live,
           // parameters for each forward or backward analysis
           unsigned int widening_delay = 1,
           unsigned int descending_iters = UINT_MAX, size_t jump_set_size = 0) {

    CRAB_LOG("backward", crab::outs()
                             << "Initial states=" << init_states << "\n");

    if (!only_forward && !m_cfg.has_exit()) {
      CRAB_WARN("cannot run backward analysis because CFG has no exit block");
      only_forward = true;
    }

    crab::CrabStats::count("CombinedForwardBackward.invocations");

    // maximum number of refinement iterations
    const unsigned max_num_iters = 5;
    // number of refinement iterations
    unsigned iters = 0;
    // CFG assertions
    std::vector<std::pair<bb_label_t, stmt_t *>> assertions;
    // immediate dominance tree
    idom_tree_t idom_tree;

    crab::CrabStats::resume("CombinedForwardBackward.GatherAssertions");
    gather_assertions();
    CRAB_LOG("backward", crab::outs()
                             << "Found " << m_unproven_assertions.size()
                             << " assertions.\n";);
    crab::CrabStats::stop("CombinedForwardBackward.GatherAssertions");

    if (!m_unproven_assertions.empty() && !only_forward) {
      crab::CrabStats::resume("CombinedForwardBackward.DominatorTree");
      std::unordered_map<bb_label_t, bb_label_t> idom_map;
      crab::analyzer::graph_algo::dominator_tree(m_cfg, m_cfg.entry(),
                                                 idom_map);
      // build idom_tree
      for (auto &kv : idom_map) {
        if (kv.second == boost::graph_traits<CFG>::null_vertex()) {
          continue;
        }
        auto it = idom_tree.find(kv.second);
        if (it == idom_tree.end()) {
          std::set<bb_label_t> reachable({kv.first});
          idom_tree.insert({kv.second, reachable});
        } else {
          it->second.insert(kv.first);
        }
      }

      CRAB_LOG("backward", crab::outs() << "Computed dominance tree:\n";
               for (auto &kv
                    : idom_tree) {
                 crab::outs() << "\t" << cfg_impl::get_label_str(kv.first)
                              << " dominates={";
                 for (auto d : kv.second) {
                   crab::outs() << d << ";";
                 }
                 crab::outs() << "}\n";
               });
      crab::CrabStats::stop("CombinedForwardBackward.DominatorTree");
    }

    assumption_map_t refined_assumptions(assumptions.begin(),
                                         assumptions.end());
    while (true) {
      iters++;
      crab::CrabStats::count("CombinedForwardBackward.iterations");
      CRAB_VERBOSE_IF(1, get_msg_stream() << "Iteration " << iters << "\n"
                                          << "Started forward analysis.\n";);

      crab::CrabStats::resume("CombinedForwardBackward.ForwardPass");
      // run forward analysis refined with preconditions from error states
      fwd_analyzer_t F(m_cfg, init_states, live, m_wto, widening_delay,
                       descending_iters, jump_set_size);
      F.run(entry, refined_assumptions);

      if (iters == 1) {
        store_analysis_results(F);
      }
      crab::CrabStats::stop("CombinedForwardBackward.ForwardPass");

      CRAB_VERBOSE_IF(1, get_msg_stream() << "Finished forward analysis.\n";);

      CRAB_LOG(
          "backward", crab::outs() << "Forward analysis: \n";
          for (auto &kv
               : boost::make_iterator_range(F.pre_begin(), F.pre_end())) {
            crab::outs() << cfg_impl::get_label_str(kv.first) << ":\n"
                         << kv.second << "\n";
          } crab::outs()
          << "\n";);

      if (only_forward) {
        CRAB_VERBOSE_IF(1, get_msg_stream()
                               << "\nSkipped backward refinement.\n";);
        break;
      } else if (m_unproven_assertions.empty()) {
        CRAB_VERBOSE_IF(
            1, get_msg_stream()
                   << "\nNo assertions found: skipped backward refinement.\n";);
        break;
      }

      // reuse wto for next iteration
      if (iters == 1) {
        m_wto = new wto_t(F.get_wto());
      }

      CRAB_VERBOSE_IF(1, get_msg_stream() << "Started backward analysis.\n";);

      crab::CrabStats::resume("CombinedForwardBackward.BackwardPass");
      // run backward analysis computing necessary preconditions
      // refined with results from the forward analysis.
      AbsDom final_states = AbsDom::bottom();
      bwd_fixpoint_iterator_t B(m_cfg, m_b_wto,
                                // A final state is safe so here means bottom
                                final_states,
                                // negate assertions:
                                // preconditions from error states
                                false, widening_delay, descending_iters,
                                jump_set_size);

      const invariant_map_t &fwd_invariants = F.get_pre_invariants();
      crab::CrabStats::resume("CombinedForwardBackward.MinimizeInvariants");
      // Important for apron and elina domains.
      invariant_map_t minimized_fwd_invariants;
      for (auto &kv : fwd_invariants) {
        AbsDom dom(kv.second);
        dom.minimize();
        minimized_fwd_invariants.insert({kv.first, dom});
      }
      crab::CrabStats::stop("CombinedForwardBackward.MinimizeInvariants");
      B.run_backward(minimized_fwd_invariants);
      crab::CrabStats::stop("CombinedForwardBackward.BackwardPass");

      CRAB_VERBOSE_IF(1, get_msg_stream() << "Finished backward analysis.\n";);

      CRAB_LOG(
          "backward", crab::outs() << "Backward analysis:\n";
          for (auto &kv
               : boost::make_iterator_range(B.begin(), B.end())) {
            crab::outs() << cfg_impl::get_label_str(kv.first) << ":\n"
                         << kv.second << "\n";
          } crab::outs()
          << "\n");

      crab::CrabStats::resume("CombinedForwardBackward.CheckRefinement");
      assumption_map_t new_refined_assumptions;
      bool more_refinement = false;
      for (auto it = m_cfg.begin(), et = m_cfg.end(); it != et; ++it) {
        AbsDom x = refined_assumptions[it->label()];
        AbsDom y = B[it->label()];
        AbsDom x_narrowing_y = x && y;
        more_refinement |= (!(x <= x_narrowing_y));
        new_refined_assumptions.insert({it->label(), x_narrowing_y});
      }
      if (more_refinement) {
        refined_assumptions.clear();
        refined_assumptions.insert(new_refined_assumptions.begin(),
                                   new_refined_assumptions.end());
        CRAB_LOG("backward",
                 crab::outs()
                     << "Backward analysis refined forward analysis.\n");
      } else {
        CRAB_LOG(
            "backward",
            crab::outs()
                << "Backward analysis cannot refine more forward analysis.\n");
      }
      crab::CrabStats::stop("CombinedForwardBackward.CheckRefinement");

      if (!more_refinement || iters > max_num_iters) {
        if (more_refinement) {
          CRAB_LOG("backward",
                   crab::outs()
                       << "Limit of backward refinements already reached.\n");
        }

        // If refined_assumptions can prove that s is safe then
        // remove assertion s from m_unproven_assertions
        discharge_assertions(refined_assumptions, idom_tree);
        break;
      }

      // reuse wto for next iteration
      if (iters == 1) {
        m_b_wto = new bwd_wto_t(B.get_WTO());
      }

    } // end while true

    CRAB_VERBOSE_IF(1, get_msg_stream()
                           << "Combined forward+backward analysis done after "
                           << iters << " iterations.\n";);
  }

  // Return the invariants that hold at the entry of b
  AbsDom operator[](bb_label_t b) const { return get_pre(b); }

  // Return the invariants that hold at the entry of b
  AbsDom get_pre(bb_label_t b) const {
    auto it = m_pre_invariants.find(b);
    if (it == m_pre_invariants.end())
      return abs_dom_t::top();
    else
      return it->second;
  }

  // Return the invariants that hold at the exit of b
  AbsDom get_post(bb_label_t b) const {
    auto it = m_post_invariants.find(b);
    if (it == m_post_invariants.end())
      return abs_dom_t::top();
    else
      return it->second;
  }

  // Return the wto of the cfg
  const wto_t &get_wto() const {
    assert(m_wto);
    return *m_wto;
  }

  // clear internal state
  void clear() {
    m_pre_invariants.clear();
    m_post_invariants.clear();
  }

  /** Extra API for checkers **/

  CFG get_cfg(void) { return m_cfg; }

  std::shared_ptr<abs_tr_t> get_abs_transformer(AbsDom &&inv) {
    return std::make_shared<abs_tr_t>(std::move(inv));
  }

  void get_safe_assertions(std::set<const stmt_t *> &out) const {
    out.insert(m_proved_assertions.begin(), m_proved_assertions.end());
  }
};

} // namespace analyzer
} // namespace crab


/* Copyright (c) 2021, United States Government, as represented by the
 * Administrator of the National Aeronautics and Space Administration.
 *
 * All rights reserved.
 *
 * The "ISAAC - Integrated System for Autonomous and Adaptive Caretaking
 * platform" software is licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with the
 * License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */

// A little tool to help testing brightness correction

#include <gflags/gflags.h>
#include <glog/logging.h>

#include <dense_map_utils.h>

#include <TinyEXIF.h>

// OpenCV includes
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc.hpp>

// Eigen includes
#include <Eigen/Geometry>
#include <Eigen/Core>

#include <fstream>
#include <iostream>

void test_gamma() {
  std::string filename = "gamma.txt";
  std::cout << "Writing: " << filename << std::endl;

  std::ofstream fh(filename.c_str());
  fh.precision(17);

  int n = 1000;
  for (int col = 0; col < n; col++) {
    double x = static_cast<double>(col) / (n - 1.0);
    double y = gamma(x);
    double z = dense_map::inv_gamma(gamma(x));
    double w = dense_map::inv_gamma(x);

    fh << x << ' ' << y << ' ' << z << ' ' << w << std::endl;
  }
  fh.close();
}

int main(int argc, const char** argv) {
  // test_gamma();
  // return 0;

  std::vector<cv::Mat> images;
  std::vector<double> exposures, isos, products;
  std::vector<std::string> image_files;

  for (int it = 0; it < argc - 1; it++) {
    // read entire image file
    std::string image_file = argv[it + 1];
    image_files.push_back(image_file);
    std::ifstream file(image_file.c_str(), std::ifstream::in | std::ifstream::binary);
    file.seekg(0, std::ios::end);
    std::streampos length = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<uint8_t> data(length);
    file.read(reinterpret_cast<char*>(data.data()), length);

    // parse image EXIF metadata
    TinyEXIF::EXIFInfo imageEXIF(data.data(), length);
    if (imageEXIF.Fields) {
#if 1
      std::cout << "Image Description  " << imageEXIF.ImageDescription << "\n"
                << "Image Resolution   " << imageEXIF.ImageWidth << "x" << imageEXIF.ImageHeight << " pix\n"
                << "Camera Model       " << imageEXIF.Make << " - " << imageEXIF.Model << "\n"
                << "Focal Length       " << imageEXIF.FocalLength << " mm" << std::endl
                << "Exposure (seconds) " << imageEXIF.ExposureTime << std::endl
                << "ISO (int)          " << imageEXIF.ISOSpeedRatings << std::endl
                << "aperture           " << imageEXIF.ApertureValue << std::endl;
#endif

      cv::Mat image;
      images.push_back(image);
      images.back() = cv::imread(image_file);

      // scales.push_back(imageEXIF.ExposureTime * imageEXIF.ISOSpeedRatings);
      exposures.push_back(imageEXIF.ExposureTime);
      isos.push_back(imageEXIF.ISOSpeedRatings);
      products.push_back(imageEXIF.ExposureTime * imageEXIF.ISOSpeedRatings);
    } else {
      LOG(FATAL) << "Could not read info from image: " << image_file << std::endl;
    }
  }

  // return 0;

  if (exposures.empty()) LOG(FATAL) << "No exposures";

  double max_exposure = *std::max_element(exposures.begin(), exposures.end());
  double min_exposure = *std::min_element(exposures.begin(), exposures.end());
  std::cout << "min and max exposure " << min_exposure << ' ' << max_exposure << std::endl;

  double max_iso = *std::max_element(isos.begin(), isos.end());
  double min_iso = *std::min_element(isos.begin(), isos.end());
  std::cout << "min and max iso " << min_iso << ' ' << max_iso << std::endl;

  double max_product = *std::max_element(products.begin(), products.end());
  double min_product = *std::min_element(products.begin(), products.end());
  std::cout << "min and max product " << min_product << ' ' << max_product << std::endl;

  int num_images = images.size();
  for (int it = 0; it < num_images; it++) {
    // double scale = max_iso / isos[it];
    // double scale = max_exposure / exposures[it];
    double scale = max_product / products[it];

    // std::cout << std::endl;
    // std::cout << "input image is " << image_files[it] << std::endl;
    // std::cout << "will scale image in the inv gamma space " << it << " by "
    // << scale << std::endl;

    // Do brightness correction
    cv::Mat out_image;
    double max_iso_times_exposure = 5.2;
    dense_map::exposureCorrection(max_iso_times_exposure, isos[it], exposures[it], images[it], out_image);

    std::ostringstream oss;
    oss << "out_image_test4_" << it + 1000 << ".jpg";
    std::string out_name = oss.str();

    std::cout << "Writing " << out_name << std::endl;
    cv::imwrite(out_name, out_image);
  }

  return 0;
}