/*
This is SCRIMP++, as published by Zhua, Yeh, Zimmerman et al. at https://sites.google.com/site/scrimpplusplus/
It is provided as a baseline for our work and republished with kind permission of Prof. Eamonn Keogh
All rights belong the original authors and they shall be asked for licensing, if required.
Few modifications to the code were made to adapt it in our framework

Details of the SCRIMP++ algorithm can be found at:
(author information ommited for ICDM review),
"SCRIMP++: Motif Discovery at Interactive Speeds", submitted to ICDM 2018.
*/
#include <stdio.h>
#include <stdlib.h>
#include <fftw3.h>
#include <math.h>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <limits>
#include <vector>
#include <algorithm>
#include <string>
#include <sstream>
#include <timing.h>
#include <logging.hpp>

#include <ScrimpppOrig.hpp>
#include <boost/filesystem.hpp>

using namespace matrix_profile;

static FactoryRegistration<ScrimpppOrig> s_origRegistration("scrimppp_orig");
static const int notification_interval_iter = 10000;

void ScrimpppOrig::initialize(const Scrimppp_params &params) {

}

void ScrimpppOrig::compute_matrix_profile(const Scrimppp_params& params)
{
	// start timer
	//time_t tstart, tend;
	//tstart = time(0);
	Timepoint tstart, tend;
	Timespan time_elapsed;

	// read time series and subsequence length (windowSize).
	std::fstream timeSeriesFile(params.time_series_filename, std::ios_base::in);

	int windowSize = params.query_window_len;
	int stepSize = floor(params.prescrimp_stride*windowSize);

	const std::string preoutfilename = params.output_filename + "_prescrimp";

	if (!timeSeriesFile.is_open())
	{
		throw std::runtime_error("Could not open input file");
	}

	std::vector<double> A;
	double tempval;
	int timeSeriesLength = 0;
	while (timeSeriesFile >> tempval)
	{
		A.push_back(tempval);
		timeSeriesLength++;
	}
	EXEC_INFO( "loaded time series of length " << timeSeriesLength );

	if (timeSeriesLength < windowSize) {
		throw std::runtime_error("ERROR: Time series is shorter than the window length, can not proceed");
	}

	timeSeriesFile.close();

	// set exclusion zone
	int exclusionZone = windowSize / 4;

	// set Matrix Profile Length
	int ProfileLength = timeSeriesLength - windowSize + 1;

	// preprocess, statistics, get the mean and standard deviation of every subsequence in the time series
	double* ACumSum = new double[timeSeriesLength];
	ACumSum[0] = A[0];
	for (int i = 1; i < timeSeriesLength; i++)
		ACumSum[i] = A[i] + ACumSum[i - 1];
	double* ASqCumSum = new double[timeSeriesLength];
	ASqCumSum[0] = A[0] * A[0];
	for (int i = 1; i < timeSeriesLength; i++)
		ASqCumSum[i] = A[i] * A[i] + ASqCumSum[i - 1];
	double* ASum = new double[ProfileLength];
	ASum[0] = ACumSum[windowSize - 1];
	for (int i = 0; i < timeSeriesLength - windowSize; i++)
		ASum[i + 1] = ACumSum[windowSize + i] - ACumSum[i];
	double* ASumSq = new double[ProfileLength];
	ASumSq[0] = ASqCumSum[windowSize - 1];
	for (int i = 0; i < timeSeriesLength - windowSize; i++)
		ASumSq[i + 1] = ASqCumSum[windowSize + i] - ASqCumSum[i];
	double* AMean = new double[ProfileLength];
	for (int i = 0; i < ProfileLength; i++)
		AMean[i] = ASum[i] / windowSize;
	double* ASigmaSq = new double[ProfileLength];
	for (int i = 0; i < ProfileLength; i++)
		ASigmaSq[i] = ASumSq[i] / windowSize - AMean[i] * AMean[i];
	double* ASigma = new double[ProfileLength];
	for (int i = 0; i < ProfileLength; i++)
		ASigma[i] = sqrt(ASigmaSq[i]);
	delete [] ACumSum;
	delete [] ASqCumSum;
	delete [] ASum;
	delete [] ASumSq;
	delete [] ASigmaSq;

	//Initialize Matrix Profile and Matrix Profile Index
	double* profile = new double[ProfileLength];
	int* profileIndex = new int[ProfileLength];
	for (int i=0; i<ProfileLength; i++)
	{
		profile[i]=std::numeric_limits<double>::infinity();
		profileIndex[i]=0;
	}
	tstart = get_cur_time();

	//int fftsize = pow(2,ceil(log2(timeSeriesLength)));
	int fftsize = timeSeriesLength; //fftsize must be at least 2*windowSize
	fftsize = fftsize > 2 * windowSize ? fftsize : 2 * windowSize;
	EXEC_TRACE ( "length of fft input: " << fftsize );

	/*******************************PreSCRIMP***************************************/

	fftw_plan plan;
	fftw_complex* ATime = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);
	fftw_complex* AFreq = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);

	for (int i = 0; i < fftsize; i++)
	{
		ATime[i][1] = 0;
		if (i < timeSeriesLength)
			ATime[i][0] = A[i];
		else
			ATime[i][0] = 0;
	}

	plan = fftw_plan_dft_1d(fftsize, ATime, AFreq, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(plan);
	fftw_free(ATime);

	fftw_complex* queryTime = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);
	fftw_complex* queryFreq = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);
	fftw_complex* AQueryTime = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);
	fftw_complex* AQueryFreq = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * fftsize);

	//Sample subsequences with a fixed stepSize, then random shuffle their computation order
	std::vector<int> idx;
	for (int i = 0; i < timeSeriesLength - windowSize + 1; i += stepSize)
		idx.push_back(i);
	std::random_shuffle(idx.begin(), idx.end());

	double* query = new double[windowSize];

	for (int idx_i = 0; idx_i < idx.size(); idx_i++)
	{
		int i = idx[idx_i];
		for (int j = 0; j < windowSize; j++)
		{
			query[j] = A[i + j];
		}
		double queryMean = AMean[i];
		double queryStd = ASigma[i];

		for (int j = 0; j < fftsize; j++)
		{
			queryTime[j][1] = 0;

			if (j < windowSize)
				queryTime[j][0] = query[windowSize - j - 1];
			else
				queryTime[j][0] = 0;
		}

		plan = fftw_plan_dft_1d(fftsize, queryTime, queryFreq, FFTW_FORWARD, FFTW_ESTIMATE);
		fftw_execute(plan);

		for (int j = 0; j < fftsize; j++)
		{
			AQueryFreq[j][0] = AFreq[j][0] * queryFreq[j][0] - AFreq[j][1] * queryFreq[j][1];
			AQueryFreq[j][1] = AFreq[j][1] * queryFreq[j][0] + AFreq[j][0] * queryFreq[j][1];
		}

		plan = fftw_plan_dft_1d(fftsize, AQueryFreq, AQueryTime, FFTW_BACKWARD, FFTW_ESTIMATE);
		fftw_execute(plan);

		int exclusionZoneStart = i - exclusionZone;
		int exclusionZoneEnd = i + exclusionZone;
		double minimumDistance = std::numeric_limits<double>::infinity();
		int minimumDistanceIndex;
		for (int j = 0; j < timeSeriesLength - windowSize + 1; j++)
		{
			double distance;
			if ((j > exclusionZoneStart) && (j < exclusionZoneEnd))
				distance = std::numeric_limits<double>::infinity();
			else
			{
				distance = 2 * (windowSize - (AQueryTime[windowSize + j - 1][0] / fftsize - windowSize * AMean[j] * queryMean) / (ASigma[j] * queryStd));
			}

			if (distance < minimumDistance)
			{
				minimumDistance = distance;
				minimumDistanceIndex = j;
			}

			if (distance < profile[j])
			{
				profile[j] = distance;
				profileIndex[j] = i;
			}
		}
		profile[i] = minimumDistance;
		profileIndex[i] = minimumDistanceIndex;

		int j = profileIndex[i];
		double lastz = (windowSize - profile[i] / 2) * (ASigma[j] * ASigma[i]) + windowSize * AMean[j] * AMean[i];
		double lastzz = lastz;
		double distance;
		for (int k = 1; k < stepSize && i + k < timeSeriesLength - windowSize + 1 && j + k < timeSeriesLength - windowSize + 1; k++)
		{
			lastz = lastz - A[i + k - 1] * A[j + k - 1] + A[i + k + windowSize - 1] * A[j + k + windowSize - 1];
			distance = 2 * (windowSize - (lastz - windowSize * AMean[j + k] * AMean[i + k]) / (ASigma[j + k] * ASigma[i + k]));
			if (distance < profile[i + k])
			{
				profile[i + k] = distance;
				profileIndex[i + k] = j + k;
			}
			if (distance < profile[j + k])
			{
				profile[j + k] = distance;
				profileIndex[j + k] = i + k;
			}
		}
		lastz = lastzz;
		for (int k = 1; k < stepSize && i - k >= 0 && j - k >= 0; k++)
		{
			lastz = lastz - A[i - k + windowSize] * A[j - k + windowSize] + A[i - k] * A[j - k];
			distance = 2 * (windowSize - (lastz - windowSize * AMean[j - k] * AMean[i - k]) / (ASigma[j - k] * ASigma[i - k]));
			if (distance < profile[i - k])
			{
				profile[i - k] = distance;
				profileIndex[i - k] = j - k;
			}
			if (distance < profile[j - k])
			{
				profile[j - k] = distance;
				profileIndex[j - k] = i - k;
			}
		}
	}

	fftw_destroy_plan(plan);
	fftw_free(AFreq);
	fftw_free(queryTime);
	fftw_free(queryFreq);
	fftw_free(AQueryTime);
	fftw_free(AQueryFreq);
	delete[] query;

	tend = get_cur_time();
	time_elapsed = tend - tstart;

	// output
	EXEC_INFO("finished prescrimp")
	PERF_LOG( "Time for PreSCRIMP: " << std::setprecision(std::numeric_limits<double>::digits10 + 2) << time_elapsed );

	std::fstream preprofileOutFile(preoutfilename.c_str(), std::ios_base::out);

	// Write PreSCRIMP Matrix Profile and Matrix Profile Index to file.
	for (int i = 0; i < timeSeriesLength - windowSize + 1; i++)
	{
		preprofileOutFile << std::setprecision(std::numeric_limits<double>::digits10 + 2) << sqrt(abs(profile[i])) << " " << std::setprecision(std::numeric_limits<int>::digits10 + 1) << profileIndex[i] << std::endl;
	}

	preprofileOutFile.close();

	/******************** SCRIMP ********************/

	//Random shuffle the computation order of the diagonals of the distance matrix
std::srand(1);//TODO: remove. Introduced for constistent evaluation order among all algorithms during debugging!
    idx.clear();
	for (int i = exclusionZone+1; i < ProfileLength; i++)
		idx.push_back(i);
	std::random_shuffle(idx.begin(), idx.end());

	double* dotproduct = new double[timeSeriesLength];

	//iteratively evaluate the diagonals of the distance matrix
	for (int ri = 0; ri < idx.size(); ri++)
	    {
		//select a random diagonal
		int diag = idx[ri];

		//calculate the dot product of every two time series values that ar diag away
		for (int j=diag; j < timeSeriesLength; j++)
			dotproduct[j]=A[j]*A[j-diag];

		//evaluate the fist distance value in the current diagonal
		double distance;
		double lastz=0; //the dot product of a subsequence
		for (int k = 0; k < windowSize; k++)
			lastz += dotproduct[k+diag];

		//j is the column index, i is the row index of the current distance value in the distance matrix
		int j=diag, i=j-diag;

		//evaluate the distance based on the dot product
		distance = 2 * (windowSize - (lastz - windowSize * AMean[j] * AMean[i]) / (ASigma[j] * ASigma[i]));

		//update matrix profile and matrix profile index if the current distance value is smaller
		if (distance < profile[j])
		{
			profile[j] = distance;
			profileIndex [j] = i;
		}
		if (distance < profile[i])
		{
			profile[i] = distance;
			profileIndex [i] = j;
		}

//std::cout << "diag " << diag << " lastz " << lastz << std::endl;
		//evaluate the second to the last distance values along the diagonal and update the matrix profile/matrix profile index.
		for (j=diag+1; j<ProfileLength; j++)
		{
			i=j-diag;
			lastz = lastz + dotproduct[j+windowSize-1] - dotproduct [j-1];
			distance = 2 * (windowSize - (lastz - windowSize * AMean[j] * AMean[i]) / (ASigma[j] * ASigma[i]));

//std::cout << "eval i: " << i << " j: " << j << " lastz" << lastz << std::endl;
			if (distance < profile[j])
			{
				profile[j] = distance;
				profileIndex [j] = i;
			}
			if (distance < profile[i])
			{
				profile[i] = distance;
				profileIndex [i] = j;
			}
		}

		//Show time per 10000 iterations
		if ( (ri+1) % notification_interval_iter == 0)
		{
			tend = get_cur_time();
			time_elapsed = tend - tstart;
			//std::cout << "Time spent: " << std::setprecision(std::numeric_limits<double>::digits10 + 2) << difftime(time(0), tstart) << " seconds." << std::endl;
			PERF_TRACE( "completed " << notification_interval_iter << " iterations in: " << time_elapsed );
		}

		//The following commented section is to produce provisional results. Basically, if you would like to enable interrupt and look at the current matrix profile/matrix profile index, you can uncomment this section, and revise line 182 according to the intterupt mechanism you're using.

		/*if (interrupt_detected) //revise this line to enable your interrupt mechanism
		{
			std::fstream prov_profileOutFile(outfilename_provisional.c_str(), std::ios_base::out);

			// Write Current Matrix Profile and Matrix Profile Index to file.
			for (int k = 0; k < timeSeriesLength - windowSize + 1; k++)
				prov_profileOutFile << std::setprecision(std::numeric_limits<double>::digits10 + 2) << sqrt(abs(profile[k])) << " " << std::setprecision(std::numeric_limits<int>::max()) << profileIndex[k] << std::endl;
			prov_profileOutFile.close();
		}
		*/
	}

	// end timer
	//tend = time(0);
	tend = get_cur_time();
	time_elapsed = tend - tstart;

	PERF_LOG ( "total computation time: " << time_elapsed << " seconds." );
	const double triang_len = ProfileLength-exclusionZone;
	PERF_LOG ( "throughput computations: " << triang_len * triang_len / get_seconds(time_elapsed) << " matrix entries/second");
	EXEC_INFO( "Writing result to file");

	if (params.output_filename.empty()) {
		throw std::runtime_error("Empty output file name specified!");
	}
	std::fstream profileOutFile(params.output_filename.c_str(), std::ios_base::out);

	// Write final Matrix Profile and Matrix Profile Index to file.
	for (int i = 0; i < timeSeriesLength - windowSize + 1; i++)
	{
		profile[i] = sqrt(abs(profile[i]));
		profileOutFile << std::setprecision(std::numeric_limits<double>::digits10 + 2) << profile[i] << " " << std::setprecision(std::numeric_limits<int>::digits10+1) << profileIndex[i] << std::endl;
	}

	profileOutFile.close();

	delete [] dotproduct;
	delete [] AMean;
	delete [] ASigma;
	delete [] profile;
	delete [] profileIndex;
}


/***********************************************************************
This file is part of the librjmcmc project source files.

Copyright : Institut Geographique National (2008-2012)
Contributors : Mathieu Brédif, Olivier Tournaire, Didier Boldo
email : librjmcmc@ign.fr

This software is a generic C++ library for stochastic optimization.

This software is governed by the CeCILL license under French law and
abiding by the rules of distribution of free software. You can use,
modify and/or redistribute the software under the terms of the CeCILL
license as circulated by CEA, CNRS and INRIA at the following URL
"http://www.cecill.info".

As a counterpart to the access to the source code and rights to copy,
modify and redistribute granted by the license, users are provided only
with a limited warranty and the software's author, the holder of the
economic rights, and the successive licensors have only limited liability.

In this respect, the user's attention is drawn to the risks associated
with loading, using, modifying and/or developing or reproducing the
software by the user in light of its specific status of free software,
that may mean that it is complicated to manipulate, and that also
therefore means that it is reserved for developers and experienced
professionals having in-depth computer knowledge. Users are therefore
encouraged to load and test the software's suitability as regards their
requirements in conditions enabling the security of their systems and/or
data to be ensured and, more generally, to use and operate it in the
same conditions as regards security.

The fact that you are presently reading this means that you have had
knowledge of the CeCILL license and that you accept its terms.

***********************************************************************/

#ifndef RJMCMC_RASTER_VARIATE_HPP
#define RJMCMC_RASTER_VARIATE_HPP

#include <boost/random/uniform_real.hpp>
#include <vector>
#include <algorithm>

namespace rjmcmc {

    // this variate generates a point in [0,1]^N with a piecewise uniform density where the pieces are given by a uniform rectangular voxel grid.
    // int *size gives the number of elements of the grid in each dimension
    // T *pdf gives the unnormalized pdf values of each voxel as a N dimensional array of size (size)

    template<int N>
    class raster_variate
    {
        typedef boost::uniform_real<> rand_type;
        mutable rand_type m_rand;

    public:
        typedef double value_type;
        enum { dimension = N };
        template<typename Engine, typename OutputIterator>
        inline double operator()(Engine& e, OutputIterator it) const {
            double x = m_rand(e);
            int offset = int(std::upper_bound(m_cdf.begin()+1,m_cdf.end(),x)-(m_cdf.begin()+1));
            double pdf = (m_cdf[offset+1]-m_cdf[offset])*m_totsize;
            for(int i=0; i<N; ++i)
            {
                int ix = offset % m_size[i];
                *it++ = (ix+m_rand(e))/m_size[i];
                offset /= m_size[i];
            }
            return pdf;
        }
        template<typename InputIterator>
        inline double pdf(InputIterator it) const {
            int offset = 0;
            int stride = 1;
            for(int i=0; i<N; ++i)
            {
                double x = *it++;
                if(x<0. || x>=1.) return 0.;
                int ix = int(x*m_size[i]);
                offset += stride*ix;
                stride *= m_size[i];
            }
            return (m_cdf[offset+1]-m_cdf[offset])*m_totsize;
        }
        template<typename T>
        raster_variate(T* pdf, int *size) {
            m_size.resize(N);
            m_totsize = 1;
            for(int i=0; i<N; ++i) { m_totsize *= size[i]; m_size[i] = size[i]; }
            m_sum = 0;
            m_cdf.resize(m_totsize+1);
            m_cdf[0] = 0.;
            for(int i=0; i<m_totsize; ++i) m_sum = m_cdf[i+1] = m_sum + pdf[i]; // assert(pdf[i]>=0)
            for(int i=0; i<m_totsize; ++i) m_cdf[i+1]/=m_sum;
        }
    private:
        std::vector<double> m_cdf;
        std::vector<int> m_size;
        int m_totsize;
        double m_sum;
    };

}; // namespace rjmcmc

#endif // RJMCMC_RASTER_VARIATE_HPP


/*
* Copyright 2019 © Centre Interdisciplinaire de développement en Cartographie des Océans (CIDCO), Tous droits réservés
*/

#ifndef GEOREFERENCING_HPP
#define GEOREFERENCING_HPP

#include <Eigen/Dense>
#include "math/CoordinateTransform.hpp"
#include "Raytracing.hpp"

/*!
* \brief Georeferencing class
* \author Guillaume Labbe-Morissette, Jordan McManus, Emile Gagne
* \date October 2, 2018, 9:39 AM
*/
class Georeferencing{
public:
  /**
  * Georeferences a ping
  *
  * @param georeferencedPing georeferenced ping in vector form
  * @param attitude the attitude of the ship in the IMU frame
  * @param position the position of the ship in the TRF
  * @param ping the ping of the georeference in the sonar frame
  * @param svp the SoundVelocityProfile
  * @param leverArm vector from the position reference point (PRP) to the acoustic center
  *
  */
  virtual void georeference(Eigen::Vector3d & georeferencedPing,Attitude & attitude,Position & position,Ping & ping,SoundVelocityProfile & svp,Eigen::Vector3d & leverArm,Eigen::Matrix3d & boresight){};
};

/*!
* \brief TRF Georeferencing class
*
* Extends Georeferencing class
*/
class GeoreferencingTRF : public Georeferencing{
public:

  /**
  * Georeferences a ping in the TRF
  *
  * @param georeferencedPing vector of a ping georeferenced
  * @param attitude the attitude of the ship in the IM frame
  * @param position the position of the ship in the TRF
  * @param ping the ping of the georeference in the sonar frame
  * @param svp the sound velocity profile
  * @param leverArm vector from the position reference point (PRP) to the acoustic center
  *
  */
  void georeference(Eigen::Vector3d & georeferencedPing,Attitude & attitude,Position & position,Ping & ping,SoundVelocityProfile & svp,Eigen::Vector3d & leverArm,Eigen::Matrix3d & boresight) {
    //Compute transform matrixes
    Eigen::Matrix3d ned2ecef;
    CoordinateTransform::ned2ecef(ned2ecef,position);

    Eigen::Matrix3d imu2ned;
    CoordinateTransform::getDCM(imu2ned,attitude);

    //Convert position to ECEF
    Eigen::Vector3d positionECEF;
    CoordinateTransform::getPositionECEF(positionECEF,position);

    //Convert ping to ECEF
    Eigen::Vector3d pingVector;
    Raytracing::rayTrace(pingVector,ping,svp);

    Eigen::Vector3d pingECEF = ned2ecef * (imu2ned * boresight * pingVector);

    //Convert lever arm to ECEF
    Eigen::Vector3d leverArmECEF =  ned2ecef * (imu2ned * leverArm);

    //Compute total ECEF vector

    georeferencedPing = positionECEF + pingECEF + leverArmECEF;
  }
};

/*!
* \brief LGF Georeferencing class
*/
class GeoreferencingLGF : public Georeferencing{
public:

  /**
  * Georeferences a ping in the LGF (NED)
  *
  * @param georeferencedPing vector of a ping georeferenced
  * @param attitude the attitude of the ship in the IM frame
  * @param position the position of the ship in the TRF
  * @param ping the ping of the georeference in the sonar frame
  * @param svp the sound velocity profile
  * @param leverArm vector from the position reference point (PRP) to the acoustic center
  *
  */
  void georeference(Eigen::Vector3d & georeferencedPing,Attitude & attitude,Position & position,Ping & ping,SoundVelocityProfile & svp,Eigen::Vector3d & leverArm,Eigen::Matrix3d & boresight) {
    Eigen::Matrix3d imu2ned;
    CoordinateTransform::getDCM(imu2ned,attitude);

    //Center position wrt centroid
    Position pos(
      position.getTimestamp(),
      position.getLatitude() 		- centroid->getLatitude(),
      position.getLongitude()		- centroid->getLongitude(),
      position.getEllipsoidalHeight()	- centroid->getEllipsoidalHeight()
    );

    //Convert position's geographic coordinates to ECEF, and then from ECEF to NED
    Eigen::Vector3d positionECEF;
    CoordinateTransform::getPositionECEF(positionECEF,pos);
    Eigen::Vector3d positionNED = ecef2ned * positionECEF;

    //Convert ping to NED
    Eigen::Vector3d pingVector;
    Raytracing::rayTrace(pingVector,ping,svp);

    Eigen::Vector3d pingNED = imu2ned * boresight * pingVector;

    //Convert lever arm to NED
    Eigen::Vector3d leverArmNED =  imu2ned * leverArm;

    //Compute total NED vector

    georeferencedPing = positionNED + pingNED + leverArmNED;
  }

  /**
  * Sets centroid and inits ECEF 2 NED matrix
  */
  void setCentroid(Position * centroid){
    this->centroid = centroid;
    CoordinateTransform::ned2ecef(ecef2ned,*this->centroid);
    ecef2ned.transposeInPlace();
  }

  /**
  *  Get a pointer to the centroid
  */

  Position * getCentroid(){ return centroid;};

private:
  Position * centroid = NULL; //in geographic coordinates
  Eigen::Matrix3d ecef2ned;
};

#endif


#define EIGEN_USE_MKL_ALL
#define EIGEN_VECTORIZE_SSE4_2

#include <cstdio>
#include <iostream>
#include <fstream>
#include <ctime>
#include <cmath>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <gflags/gflags.h>
#include <boost/math/special_functions/bessel.hpp>

#include "frpca/frpca.h"
#include "frpca/matrix_vector_functions_intel_mkl.h"
#include "frpca/matrix_vector_functions_intel_mkl_ext.h"

using namespace std;
using namespace Eigen;
using namespace boost;

const float EPS = 0.00000000001f;
typedef Eigen::SparseMatrix<float, Eigen::RowMajor> SMatrixXf;

DEFINE_string(filename, "data/PPI.ungraph", "Filename for edgelist file.");
DEFINE_string(emb1, "sparse.emb", "Filename for svd results.");
DEFINE_string(emb2, "spectral.emb", "Filename for svd results.");
DEFINE_int32(num_node, 3890, "Number of node in the graph.");
DEFINE_int32(num_rank, 128, "Embedding dimension.");
DEFINE_int32(num_step, 10, "Number of order for recursion.");
DEFINE_int32(num_iter, 5, "Number of iter in randomized svd.");
DEFINE_int32(num_thread, 10, "Number of threads.");
DEFINE_double(theta, 0.5, "Parameter of ProNE");
DEFINE_double(mu, 0.1, "Parameter of ProNE");

SMatrixXf readGraph(string filename, int num_node){
    SMatrixXf A(num_node, num_node);
    typedef Eigen::Triplet<float> T;
    vector<T> tripletList;
    ifstream fin(filename.c_str());
    while (1)
    {
        string x, y;
        if (!(fin >> x >> y))
            break;
        int a = atoi(x.c_str()), b = atoi(y.c_str());
        if (a==b) continue;
        tripletList.push_back(T(a, b, 1));
        tripletList.push_back(T(b, a, 1));
    }
    A.setFromTriplets(tripletList.begin(), tripletList.end());
    return A;
}

SMatrixXf l1Normalize(SMatrixXf & mat){
    SMatrixXf mat2(mat.rows(), mat.cols());
    for (int k=0; k<mat.outerSize(); ++k){
        int num_neighbor = mat.row(k).sum();
        for (SMatrixXf::InnerIterator it(mat,k); it; ++it)
            mat2.insert(k, it.col()) = it.value()/num_neighbor;
    }
    return mat2;
}

MatrixXf & l2Normalize(MatrixXf & mat){
    for (int i = 0; i < mat.rows(); ++i){
        float ssn = sqrt(mat.row(i).squaredNorm());
        if (ssn < EPS) ssn = EPS;
        mat.row(i) = mat.row(i) / ssn;
      }
    return mat;
}

SMatrixXf & validate(SMatrixXf & mat){
    for (int k=0; k<mat.outerSize(); ++k)
          for (SMatrixXf::InnerIterator it(mat,k); it; ++it)
              if (it.value() <=0)
                mat.coeffRef(k, it.col()) = 1;
    return mat;
}

SMatrixXf & smfLog(SMatrixXf & mat){
    for (int k=0; k<mat.outerSize(); ++k)
          for (SMatrixXf::InnerIterator it(mat,k); it; ++it)
              mat.coeffRef(it.row(), it.col()) = log(it.value());
    return mat;
}

float bessel(int a, float b){
    return boost::math::cyl_bessel_i(a, b);
}

MatrixXf getEmbbeddingViaDenseSvd(MatrixXf &data, int rank){
    Eigen::BDCSVD<Eigen::MatrixXf> svdOfC(data, Eigen::ComputeThinU);
    MatrixXf emb = svdOfC.matrixU() * svdOfC.singularValues().cwiseSqrt().asDiagonal();
    emb = l2Normalize(emb);
    return emb;
}

MatrixXf runFrPCA(SMatrixXf & input, int rank, int iter)
{
    int m = input.rows(), nnz = input.nonZeros();
    mat_coo *A = coo_matrix_new(m, m, nnz);
    A->nnz = nnz;
    int i=0;
    for (int k=0; k<input.outerSize(); ++k)
        for (SMatrixXf::InnerIterator it(input,k); it; ++it)
          {
            A->rows[i] = k+1;
            A->cols[i] = it.col()+1;
            A->values[i] = it.value();
            i += 1;
          }
    cout << "read matrix done..." <<endl;
    // coo_matrix_print(A);

    //transform it to CSR format
    mat_csr* D = csr_matrix_new();
    csr_init_from_coo(D, A);
    coo_matrix_delete(A);

    //the test for frPCA
    mat *U = matrix_new(m, rank);
    mat *S = matrix_new(rank, 1);
    mat *V = matrix_new(m, rank);
    frPCA(D, &U, &S, &V, rank, iter);

    // matrix_print(U);
    // matrix_print(S);

    MatrixXf emb = MatrixXf::Random(m, rank);
    for (int i=0; i<m; i++)
        for (int j=0; j<rank; j++)
            emb(i, j) = matrix_get_element(U,i,j) *  sqrt(matrix_get_element(S,j,0));
    cout << "matrix decomposition done" <<endl;
    return emb;
}

MatrixXf getSparseEmbedding(SMatrixXf & A, int rank, int num_iter){
    time_t t1 = time(NULL);
    int row = A.rows(), col = A.cols();
    SMatrixXf B = l1Normalize(A);
    SMatrixXf C = B.transpose();
    SMatrixXf D(col, col), E(row, col), F(row, col);
    for (int i = 0; i < row; ++i){
        D.insert(i, i) = pow(C.row(i).sum(), 0.75);
    }

    D = D / D.sum();
    E = A * D;

    B = validate(B);
    E = validate(E);

    B = smfLog(B);
    E = smfLog(E);
    F = B - E;
    cout << "preprocess time: "<< (time(NULL) - t1 + 0.0) << endl;
    cout << "number of nnz: "<< F.nonZeros() <<endl;

    MatrixXf emb = runFrPCA(F, rank, num_iter);

    emb = l2Normalize(emb);
    return emb;
}

MatrixXf getSpectralEmbedding(SMatrixXf & A, MatrixXf & a, int step, float theta, float mu){
    time_t t1 = time(NULL);
    cout << "Chebyshev series --------------- " << endl;
    if (step==1) return a;
    int num_node = a.rows(), rank = a.cols();
    SMatrixXf I(num_node, num_node);
    for (int i = 0; i < num_node; ++i)
        I.insert(i, i) = 1;
    A = A + I;
    SMatrixXf B = l1Normalize(A);
    SMatrixXf L = I - B;
    SMatrixXf M = L - mu * I;

    MatrixXf Lx0 = a;
    MatrixXf Lx1 = M * a, Lx2;
    Lx1 = 0.5 * M * Lx1 - a;

    MatrixXf conv = bessel(0, theta)* Lx0;
    conv -= 2 * bessel(1, theta)* Lx1;
    for(int i=2; i<step; i++){
        Lx2 = M * Lx1;
        Lx2 = (M * Lx2 - 2 * Lx1) - Lx0;

        if (i % 2 == 0)
            conv += 2 * bessel(i, theta) * Lx2;
        else
            conv -= 2 * bessel(i, theta) * Lx2;
        Lx0 = Lx1;
        Lx1 = Lx2;
        cout << "Bessell time: " << i <<"\t"<< (time(NULL) - t1 + 0.0) << endl;
    }
    MatrixXf emb = A * (a - conv);
    cout << "Chebyshev time: "<< (time(NULL) - t1 + 0.0) << endl;

    // time_t t2 = time(NULL);
    // MatrixXf emb = getEmbbeddingViaDenseSvd(emb, rank);
    // cout << "dense svd time: "<< (time(NULL) - t2 + 0.0) << endl;
    emb = l2Normalize(emb);
    return emb;
}

void saveEmbedding(MatrixXf &data, string output){
    int m = data.rows(), d = data.cols();
    FILE *emb = fopen(output.c_str(), "wb");
    fprintf(emb, "%d %d\n", m, d);
    for (int i = 0; i < m; i++)
    {
        fprintf(emb, "%d", i);
        for (int j = 0; j < d; j++)
            fprintf(emb, " %f", data(i, j));
        fprintf(emb, "\n");
    }
    fclose(emb);
}

int main(int argc, char** argv)
{
    gflags::ParseCommandLineFlags(&argc, &argv, true);
    Eigen::setNbThreads(FLAGS_num_thread);

    time_t t1 = time(NULL);
    SMatrixXf A = readGraph(FLAGS_filename, FLAGS_num_node);

    MatrixXf feature = getSparseEmbedding(A, FLAGS_num_rank, FLAGS_num_iter);
    time_t t2 = time(NULL);
    cout << "Running time of get sparse embedding: " << (t2 - t1 + 0.0) << endl;

    MatrixXf embedding = getSpectralEmbedding(A, feature, FLAGS_num_step, FLAGS_theta, FLAGS_mu);
    time_t t3 = time(NULL);
    cout << "Running time of get spectral embedding: " << (t3 - t2 + 0.0)  << endl;
    cout << "Running time of ProNE: " << (t3 - t1 + 0.0) << endl;

    saveEmbedding(feature, FLAGS_emb1);
    saveEmbedding(embedding, FLAGS_emb2);
    cout << "Embedding save done " << endl;

}


#include <iostream>
#include <fstream>
#include <iomanip>
#include <math.h>
#include <vector>
#include <random>
#include <thread>
#include <boost/multi_array.hpp>
#include <boost/python.hpp>
#include <boost/python/numpy.hpp>

using namespace std;
namespace p = boost::python;
namespace np = boost::python::numpy;

#include "Layers.cpp"

// precision for floating numbers written in files
constexpr int prec_write_file = 16;

// define a standard normal distribution
std::random_device rd;
std::mt19937 gen(rd());
std::normal_distribution<> dis(0.,1.);

// uniform dictribution
std::uniform_real_distribution<> uni(0.,1.);

struct double_and_int{
	double d = 0.;
	int i = 0;
};

class CNN3{

	private:

		int num_CLs; // number of convolution layers
		int num_FCs; // number of FullCon layers
		int img_h_i; // image height
		int img_w_i; // image width
		int img_h; // image height before the fully connected layers
		int img_w; // image width before the fully connected layers
		int num_images; // number of images
		int num_labels; // number of different possible labels
		int n_channels; // number of channels

		vector<int> CL_size_filters; // Convolution layers: filter sizes
		vector<int> CL_num_filters; // Convolution layers: numbers of filters
		vector<int> MP_size; // Maxpool layers: pool sizes
		vector<int> FC_size; // FullCon layers: number of neurons

		// These four vectors will contain the layers
		vector<ConvLayer> CLs; // Convolution layers
		vector<ReLU> RLUs; // ReLU layers
		vector<MaxPool> MPs; // Maxpool layers
		vector<FullCon> FCs; // FullCon layers
		vector<SoftMax> SMs; // Softmax layers

	public:

		CNN3(){
			np::initialize(); // required to create numpy arrays (otherwise leads to segmentation faults)
		}

		CNN3(
			int img_w_, // image width
			int img_h_, // image height
			int n_channels_, // number of channels
			p::list& CL_size_filters_, // list of filter sizes
			p::list& CL_num_filters_, // list of numbers of filters
			p::list& MP_size_, // list of pool sizes
			p::list& FC_size_, // list of FullCon sizes
			int num_labels_ // number of different possible labels
		)
		{
			// initialization
			img_w_i = img_w_;
			img_h_i = img_h_;
			img_w = img_w_;
			img_h = img_h_;
			num_labels = num_labels_;
			n_channels = n_channels_;
			CL_size_filters = int_list_to_vector(CL_size_filters_);
			CL_num_filters = int_list_to_vector(CL_num_filters_);
			MP_size = int_list_to_vector(MP_size_);
			FC_size = int_list_to_vector(FC_size_);
			num_CLs = len(CL_size_filters_);
			num_FCs = len(FC_size_);

			num_images = n_channels; // tracks the number of images

			// build the layers
			for(int i=0; i<num_CLs; i++){
				CLs.push_back(ConvLayer(CL_size_filters[i], num_images, CL_num_filters[i], gen, dis));
				num_images = CL_num_filters[i];
				img_w = img_w + 1 - CL_size_filters[i];
				img_h = img_h + 1 - CL_size_filters[i];
				RLUs.push_back(ReLU());
				MPs.push_back(MaxPool(MP_size[i]));
				img_w = (int) img_w / MP_size[i];
				img_h = (int) img_h / MP_size[i];
			}
			long int n_inputs = num_images*img_h*img_w;
			for(int i=0; i<num_FCs; i++){
				int n_neurons = FC_size[i];
				FCs.push_back(FullCon(n_inputs, n_neurons, gen, dis));
				n_inputs = n_neurons;
			}
			SMs.push_back(SoftMax(n_inputs, num_labels, gen, dis));

			np::initialize(); // required to create numpy arrays (otherwise leads to segmentation faults)
		}

		// save the CNN parameters to a file
		void save(char* filename){
			ofstream file;
			file.open(filename);
			file << fixed << setprecision(prec_write_file);
			file << num_CLs << sep_val << num_FCs << sep_val << n_channels << sep_val << img_w_i << sep_val << img_h_i << sep_val << img_w << sep_val << img_h << sep_val << num_images << sep_line;
			save_vector(CL_size_filters, file);
			save_vector(CL_num_filters, file);
			save_vector(MP_size, file);
			for(int i=0; i<num_CLs; i++){
				CLs[i].save(file);
				RLUs[i].save(file);
				MPs[i].save(file);
			}
			for(int i=0; i<num_FCs; i++){
				FCs[i].save(file);
			}
			SMs[0].save(file);
			file.close();
		}

		// load the CNN parameters from a file
		void load(char* filename){
			ifstream file;
			file.open(filename);
			char c;
			file >> num_CLs >> c >> num_FCs >> c >> n_channels >> c >> img_w_i >> c >> img_h_i >> c >> img_w >> c >> img_h >> c >> num_images >> c;
			CL_size_filters = load_vector<int>(file);
			CL_num_filters = load_vector<int>(file);
			MP_size = load_vector<int>(file);
			CLs.clear();
			RLUs.clear();
			CLs.clear();
			SMs.clear();
			for(int i=0; i<num_CLs; i++){
				CLs.push_back(ConvLayer());
				CLs[i].load(file);
				RLUs.push_back(ReLU());
				RLUs[i].load(file);
				MPs.push_back(MaxPool());
				MPs[i].load(file);
			}
			for(int i=0; i<num_FCs; i++){
				FCs.push_back(FullCon());
				FCs[i].load(file);
			}
			SMs.push_back(SoftMax());
			SMs[0].load(file);
			file.close();
            num_labels = SMs[0].output.size();
		}

		// forward pass
		vector<double> forward(d3_array_type &input, double p_dropout = 0.){

			// number and dimensions of images
			int nim = input.shape()[0];
			int h = input.shape()[1];
			int w = input.shape()[2];
			if(h != img_h_i || w != img_w_i || nim != n_channels){
				cout << "\nInvalid input dimensions!\n" << endl;
			}
			for(int i=0; i<num_CLs; i++){
				d3_array_type output1 = CLs[i].forward(input);
				input.resize(boost::extents[output1.shape()[0]][output1.shape()[1]][output1.shape()[2]]);
				input = output1;

				d3_array_type output2 = MPs[i].forward(input);
				input.resize(boost::extents[output2.shape()[0]][output2.shape()[1]][output2.shape()[2]]);
				input = output2;

				d3_array_type output3 = RLUs[i].forward(input);
				input.resize(boost::extents[output3.shape()[0]][output3.shape()[1]][output3.shape()[2]]);
				input = output3;
			}

			vector<double> input_vec;
			auto input_shape = input.shape();
			for(int i=0; i<num_images; i++){
				for(int j=0; j<img_h; j++){
					for(int k=0; k<img_w; k++){
						input_vec.push_back(input[i][j][k]);
					}
				}
			}

			for(int i=0; i<num_FCs; i++){
				input_vec = FCs[i].forward(input_vec, gen, uni, p_dropout);
			}

			return SMs[0].forward(input_vec);
		}

		// backpropagation
		void backprop(vector<double> d_L_d_out_i, double learn_rate){

			d3_array_type d_L_d_in(boost::extents[num_images][img_h][img_w]);

			vector<double> d_L_d_in_vec = SMs[0].backprop(d_L_d_out_i, learn_rate);

			for(int i=num_FCs-1; i>=0; i--){
				d_L_d_in_vec = FCs[i].backprop(d_L_d_in_vec, learn_rate);
			}

			for(int i=0; i<num_images; i++){
				for(int j=0; j<img_h; j++){
					for(int k=0; k<img_w; k++){
						d_L_d_in[i][j][k] = d_L_d_in_vec[i*img_h*img_w + j*img_w + k];
					}
				}
			}

			for(int i=num_CLs-1; i>=0; i--){
				d3_array_type d_L_d_in3 = RLUs[i].backprop(d_L_d_in);
				auto shape = d_L_d_in3.shape();
				d_L_d_in.resize(boost::extents[shape[0]][shape[1]][shape[2]]);
				d_L_d_in = d_L_d_in3;

				d3_array_type d_L_d_in2 = MPs[i].backprop(d_L_d_in);
				shape = d_L_d_in2.shape();
				d_L_d_in.resize(boost::extents[shape[0]][shape[1]][shape[2]]);
				d_L_d_in = d_L_d_in2;

				d3_array_type d_L_d_in1 = CLs[i].backprop(d_L_d_in, learn_rate);
				shape = d_L_d_in1.shape();
				d_L_d_in.resize(boost::extents[shape[0]][shape[1]][shape[2]]);
				d_L_d_in = d_L_d_in1;
			}
		}

		// loss function and accuracy (1 if correct answer, 0 otherwise)
		double_and_int loss_acc(vector<double> &output, int label){
			double_and_int results;
			results.d = -log(output[label]);
			results.i = 1;
			for(int i=0; i<output.size(); i++){
				if(output[i] > output[label]){
					results.i = 0;
				}
			}
			return results;
		}

		// Completes a training step on the image 'image' with label 'label'.
		// Returns the corss-entropy and accuracy.
		double_and_int train(d3_array_type image, int label, double learn_rate = 0.005, double p_dropout = 0.) {

			// forward pass
			vector<double> output_forward = forward(image, p_dropout);

			// gradient of the loss function with respect to the output
			vector<double> d_L_d_out;
			for(int i=0; i<num_labels; i++){
				d_L_d_out.push_back(0.);
			}
			d_L_d_out[label] = -1./output_forward[label];

			// backpropagation
			backprop(d_L_d_out, learn_rate);

			// return loss and accuracy
			return loss_acc(output_forward, label);
		}

		// full forward propagation - Python wrapper
		// input: 2d numpy array
		np::ndarray forward_python(np::ndarray image){
			d3_array_type input = d3_numpy_to_multi_array(image);
			return vector_to_numpy(forward(input));
		}

		// forward - return loss and accuracy - Python wrapper
		p::list forward_la_python(np::ndarray image, int label){
			d3_array_type input = d3_numpy_to_multi_array(image);
			vector<double> output = forward(input);
			double_and_int results = loss_acc(output, label);
			p::list results_p;
			results_p.append(results.d);
			results_p.append(results.i);
			return results_p;
		}

		// full backpropagation - Python wrapper
		void backprop_python(np::ndarray d_L_d_out, double learn_rate){
			backprop(numpy_to_vector(d_L_d_out), learn_rate);
		}

		// train - Python wrapper
		p::list train_python(np::ndarray image, int label, double learn_rate, double p_dropout) {
			double_and_int results;
			results = train(d3_numpy_to_multi_array(image), label, learn_rate, p_dropout);
			p::list results_p;
			results_p.append(results.d);
			results_p.append(results.i);
			return results_p;
		}

};

BOOST_PYTHON_MODULE(CNN3)
{
    p::class_<CNN3>("CNN3", p::init<int, int, int, p::list&, p::list&, p::list&, p::list&, int>())
		.def(p::init<>())
		.def("forward", &CNN3::forward_python)
		.def("backprop", &CNN3::backprop_python)
		.def("train", &CNN3::train_python)
		.def("save", &CNN3::save)
		.def("load", &CNN3::load)
		.def("forward_la", &CNN3::forward_la_python)
	;
}


#include "formant.h"
#include <Eigen/Dense>

using namespace Analysis::Formant;
using Analysis::FormantResult;

using namespace Eigen;

struct Analysis::Formant::KarmaState
{
    int numF;
    VectorXd y;
    MatrixXd F;
    MatrixXd Q;
    MatrixXd R;
    VectorXd m_up;
    VectorXd P_up;
};

constexpr int ncep = 15;
constexpr int numF = 3;

Karma::Karma()
    : state(new KarmaState)
{
    state->F.setIdentity(2 * numF, 2 * numF);

    state->Q.setZero(2 * numF, 2 * numF);
    state->Q.diagonal().head(numF).setConstant(320 * 320);
    state->Q.diagonal().tail(numF).setConstant(100 * 100);

    state->R.setZero(ncep, ncep);
    for (int i = 0; i < ncep; ++i) {
        state->R(i, i) = 1.0 / (double) (i + 1);
    }

    VectorXd x0(2 * numF);
    x0 << 500, 1500, 2500,
           80,  120,  160;

    state->m_up = x0;
    state->P_up = state->Q;
}

Karma::~Karma()
{
    delete state;
}

static VectorXd calcCepstrumCoefs(const double *lpc, int lpcOrder, int ncep);
static MatrixXd calcMatrixH(const VectorXd& m, int cepOrder, double Fs);
static VectorXd calcCepstrumMapping(const VectorXd& m, int cepOrder, double Fs);

FormantResult Karma::solve(const double *lpc, int lpcOrder, double sampleRate)
{
    auto& F  = state->F;
    auto  Ft = state->F.transpose();
    auto& Q  = state->Q;
    auto& R  = state->R;

    auto m_pred = F * state->m_up;
    auto P_pred = F * state->P_up * Ft + Q;

    auto H  = calcMatrixH(m_pred, ncep, sampleRate);
    auto Ht = H.transpose();

    auto S = H * P_pred * Ht + R;
    auto K = (P_pred * Ht) * S.colPivHouseholderQr().inverse();

    auto y_pred = calcCepstrumMapping(m_pred, ncep, sampleRate);
    auto y      = calcCepstrumCoefs(lpc, lpcOrder, ncep);

    state->m_up = m_pred + K * (y - y_pred);
    state->P_up = P_pred - K * H * P_pred;

    rpm::vector<FormantData> formants(numF);
    for (int i = 0; i < numF; ++i) {
        formants[i] = {
            .frequency = state->m_up(i),
            .bandwidth = state->m_up(numF + i),
        };
    }

    return { .formants = formants };
}

VectorXd calcCepstrumCoefs(const double *lpc, int lpcOrder, int ncep)
{
    VectorXd C(ncep);

    for (int n = 1; n <= ncep; ++n) {
        if (n == 1) {
            C(n - 1) = lpc[n - 1];
        }
        else if (n <= lpcOrder) {
            C(n - 1) = lpc[n - 1];
            for (int i = 1; i <= ncep - 1; ++i) {
                C(n - 1) += (double) i / (double) n * lpc[n - i - 1] * C(i - 1);
            }
        }
        else {
            C(n - 1) = 0.0;
            for (int i = ncep - lpcOrder; i <= ncep - 1; ++i) {
                C(n - 1) += (double) i / (double) n * lpc[n - i - 1] * C(i - 1);
            }
        }
    }

    return C;
}

MatrixXd calcMatrixH(const VectorXd &m, int cepOrder, double Fs)
{
    const int numF = m.size() / 2;
    auto freq = m.segment(0, numF);
    auto band = m.segment(numF, numF);

    MatrixXd H(cepOrder, 2 * numF);

    for (int i = 0; i < cepOrder; ++i) {
        for (int j = 0; j < numF; ++j) {
            const double bwTerm = exp((-M_PI * (i + 1) * band(j)) / Fs);

            H(i, j)        = -4.0 * M_PI / Fs * bwTerm * sin((2.0 * M_PI * (i + 1) * freq(j)) / Fs);
            H(i, numF + j) = -2.0 * M_PI / Fs * bwTerm * cos((2.0 * M_PI * (i + 1) * freq(j)) / Fs);
        }
    }

    return H;
}

VectorXd calcCepstrumMapping(const VectorXd& m, int cepOrder, double Fs)
{
    const int numF = m.size() / 2;
    auto freq = m.segment(0, numF);
    auto band = m.segment(numF, numF);

    VectorXd C_int(numF);
    VectorXd C(cepOrder);

    for (int i = 0; i < cepOrder; ++i) {
        for (int p = 0; p < numF; ++p) {
            const double bwTerm = (2.0 / (double) (i + 1)) * exp((-M_PI * (i + 1) * band(p)) / Fs);
            C_int(p) = bwTerm * cos((2.0 * M_PI * (i + 1) * freq(p)) / Fs);
        }
        C(i) = C_int.sum();
    }

    return C;
}


#pragma once

#include <armadillo>

class AmbientFluid;

/*!
 * Namespace for some utilities.
 */
namespace utils {

    /*!
     * \brief Calculate thermal conductivity of natural gas at given pressure.
     *
     * This is some kind of empirical relation or fitting, with an unknown source.
     * It should probably be replaced with something else. See for example
     * A simple correlation to estimate natural gas thermal conductivity (Azad Jarrahiana and Ehsan Heidaryan, Journal of Natural Gas Science and Engineering, Volume 18, May 2014).
     *
     * \param pressure Gas pressure [Pa].
     * \return Thermal conductivity [W/(m K)]
     */
    double calcGasThermalConductivity(const double pressure);

    /*!
     * \brief Calculate outer film coefficient for a given outer diameter and
     * AmbientFluid.
     *
     * This is just a wrapper around utils::calcOuterWallFilmCoefficient(const double, const double, const double, const double, const double, const double)
     *
     * \param diameter Outer diameter [m]
     * \param fluid AmbientFluid describing the fluid
     * \return Outer film coefficient [W/(m2 K)]
     */
    double calcOuterWallFilmCoefficient(
            const double diameter,
            const AmbientFluid& fluid
            );

    /*!
     * \brief Calculate the outer film coefficient for external flow normal to a
     * circular cylinder.
     *
     * Uses eq. 7.52 from Fundamentals of heat and mass transfer (7th Ed, 2011) (Bergman, Lavine, Incropera, DeWitt).
     *
     * \param diameter Outer diameter [m]
     * \param heatCapacityConstantPressure Fluid heat capacity (\f$c_p\f$) [J/(kg K)]
     * \param viscosity Fluid dynamic viscosity [Pa s] = [kg/m*s]
     * \param thermalConductivity Fluid thermal conductivity [W/(m K)]
     * \param density Fluid density [kg/m3]
     * \param velocity Fluid velocity [m/s]
     * \return Outer film coefficient [W/(m2 K)]
     */
    double calcOuterWallFilmCoefficient(
            const double diameter,
            const double heatCapacityConstantPressure = 4200, // [J/kg K]
            const double viscosity = 1.05/1000.0, // [Pa s] = [kg/m*s]
            const double thermalConductivity = 0.57, // [W/m K]
            const double density = 1020, // [kg/m3]
            const double velocity = 0.1 // [m/s]
            );

    /*!
     * \brief Calculate inner wall film coefficient for flow inside a cylinder.
     *
     * This uses the Dittus-Boelter equation at Reynolds numbers abve 1e4,
     * eq. 8.55 Fundamentals of heat and mass transfer (7th Ed, 2011) (Bergman, Lavine, Incropera, DeWitt)
     * at Reynolds number between 1e4 and 4e4, and returns 0 below 4e4.
     *
     * \param diameter Inner diameter [m]
     * \param fluidPressure Fluid pressure [Pa]
     * \param fluidReynoldsNumber Fluid Reynolds number [-]
     * \param fluidHeatCapacityConstantPressure Fluid heat capacity (\f$c_p\f$) [J/(kg K)]
     * \param fluidViscosity Fluid dynamic viscosity [Pa s] = [kg/m*s]
     * \return Inner film coefficient [W/(m2 K)]
     */
    double calcInnerWallFilmCoefficient(
            const double diameter,
            const double fluidPressure,
            const double fluidReynoldsNumber,
            const double fluidHeatCapacityConstantPressure,
            const double fluidViscosity);

    /*!
     * \brief Calculate equivalent burial layer thickness.
     *
     * For a pipeline buried in a medium at a given depth, this function
     * calculates the thickness of an equivalent cylinder shell of the same
     * medium around the pipeline which gives the same heat transfer between
     * the fluid in the pipeline and the ambient medium. The calculation is
     * based on equations in the documentation of the OLGA simulation software.
     *
     * Using the thickness from this function ensures the same results are
     * achieved with steady state and unsteady heat transfer models.
     *
     * \see utils::calcEquivalentBurialLayerRadius()
     *
     * \param innerDiameter Pipeline inner diameter [m]
     * \param wallThickness Pipeline wall thickness [m]
     * \param burialDepth Distance from top of pipe to top of burial medium [m]
     * \param burialMediumConductivity Thermal conductivity of burial medium [W/(m K)]
     * \return
     */
    double calcEquivalentBurialLayerWidth(
            const double innerDiameter,
            const double wallThickness,
            const double burialDepth,
            const double burialMediumConductivity = 2.0);

    /*!
     * \brief calcEquivalentBurialLayerRadius
     *
     * For a pipeline buried in a medium at a given depth, this function
     * calculates the outer radius of an equivalent cylinder shell of the same
     * medium around the pipeline which gives the same heat transfer between
     * the fluid in the pipeline and the ambient medium. The calculation is
     * based on equations in the documentation of the OLGA simulation software.
     *
     * \param innerDiameter Pipeline inner diameter [m]
     * \param wallThickness Pipeline wall thickness [m]
     * \param burialDepth Distance from top of pipe to top of burial medium [m]
     * \param burialMediumConductivity Thermal conductivity of burial medium [W/(m K)]
     * \return Outer radius of equivalent cylinder shell of burial medium [m]
     */
    double calcEquivalentBurialLayerRadius(
            const double innerDiameter,
            const double wallThickness,
            const double burialDepth,
            const double burialMediumConductivity = 2.0);

    /*!
     * \brief Calculate logarithmically (log10) spaced cylinder shell widths.
     *
     * This function calculates the widths of nShells cylinder shells,
     * logarithmically (log10) spaced between innerRadius and outerRadius.
     *
     * \param innerRadius Inner radius [m]
     * \param outerRadius Outer radius [m]
     * \param nShells Number of cylinder shells
     * \return arma::vec of shell widths [m]
     */
    arma::vec calcLogSpacedShellWidths(
            const double innerRadius,
            const double outerRadius,
            const arma::uword nShells = 10);

    /*!
     * \brief Calculate the widths of equivalent burial cylinder shells.
     *
     * This is just a wrapper around utils::calcEquivalentBurialLayerRadius()
     * and utils::calcLogSpacedShellWidths() that divides the shell into
     * several logarithmically spaced shells.
     *
     * \see utils::calcEquivalentBurialLayerRadius()
     * \see utils::calcLogSpacedShellWidths()
     *
     * \param innerDiameter Inner diameter [m]
     * \param wallThickness Wall thickness [m]
     * \param burialDepth Distance from top of pipe to top of burial medium [m]
     * \param burialMediumConductivity Thermal conductivity of burial medium [W/(m K)]
     * \param nShells Number of shells
     * \return Widths of logarithmically spaced cylinder shells.
     */
    arma::vec calcEquivalentBurialLayerWidths(
            const double innerDiameter,
            const double wallThickness,
            const double burialDepth,
            const double burialMediumConductivity = 2.0,
            const arma::uword nShells = 10);
}


//==============================================================================
//         Copyright 2015 - J.T. Lapreste
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_TENPOWER_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_SCALAR_TENPOWER_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/tenpower.hpp>
#include <boost/dispatch/meta/as_floating.hpp>
#include <boost/simd/include/constants/one.hpp>
#include <boost/simd/include/constants/ten.hpp>
#include <boost/simd/include/functions/abs.hpp>
#include <boost/simd/include/functions/sqr.hpp>
#include <boost/simd/include/functions/scalar/is_odd.hpp>
#include <boost/simd/include/functions/scalar/rec.hpp>
#include <boost/dispatch/attributes.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT  ( tenpower_, tag::cpu_
                            , (A0)
                            , (scalar_< int_<A0> >)
                            )
  {
    typedef  typename dispatch::meta::as_floating<A0>::type result_type;

    BOOST_FORCEINLINE result_type operator()(A0 exp) const
    {
      result_type result = One<result_type>();
      result_type base = Ten<result_type>();
      bool neg = exp < 0;
      exp =  boost::simd::abs(exp);
      while(exp)
      {
        if (is_odd(exp)) result *= base;
        exp >>= 1;
        base = sqr(base);
      }
      return neg ? rec(result) : result;
    }
  };

  BOOST_DISPATCH_IMPLEMENT  ( tenpower_, tag::cpu_
                            , (A0)
                            , (scalar_< uint_<A0> >)
                            )
  {
    typedef  typename dispatch::meta::as_floating<A0>::type result_type;

    BOOST_FORCEINLINE result_type operator()(A0 exp) const
    {
      result_type result = One<result_type>();
      result_type base = Ten<result_type>();
      while(exp)
      {
        if (is_odd(exp)) result *= base;
        exp >>= 1;
        base = sqr(base);
      }
      return result;
    }
  };
} } }

#endif


// Copyright (c) 2020 Chris Richardson & Garth Wells
// FEniCS Project
// SPDX-License-Identifier:    MIT

#include "lattice.h"
#include "cell.h"
#include "elements/lagrange.h"
#include "quadrature.h"
#include <Eigen/Dense>

using namespace basix;

namespace
{
//-----------------------------------------------------------------------------
Eigen::ArrayXd warp_function(int n, Eigen::ArrayXd& x)
{
  [[maybe_unused]] auto [pts, wts]
      = quadrature::gauss_lobatto_legendre_line_rule(n + 1);
  wts.setZero();

  pts *= 0.5;
  for (int i = 0; i < n + 1; ++i)
    pts[i] += (0.5 - static_cast<double>(i) / static_cast<double>(n));

  FiniteElement L = create_dlagrange(cell::type::interval, n);
  Eigen::MatrixXd v = L.tabulate(0, x)[0];
  return v * pts.matrix();
}
//-----------------------------------------------------------------------------

} // namespace

//-----------------------------------------------------------------------------
Eigen::ArrayXXd lattice::create(cell::type celltype, int n,
                                lattice::type lattice_type, bool exterior)
{
  switch (celltype)
  {
  case cell::type::point:
    return Eigen::ArrayXXd::Zero(1, 1);
  case cell::type::interval:
  {
    if (n == 0)
      return Eigen::ArrayXXd::Constant(1, 1, 0.5);

    Eigen::ArrayXd x;
    if (exterior)
      x = Eigen::VectorXd::LinSpaced(n + 1, 0.0, 1.0);
    else
    {
      const double h = 1.0 / static_cast<double>(n);
      x = Eigen::VectorXd::LinSpaced(n - 1, h, 1.0 - h);
    }

    if (lattice_type == lattice::type::gll_warped)
      x += warp_function(n, x);

    return x;
  }
  case cell::type::quadrilateral:
  {
    if (n == 0)
      return Eigen::ArrayXXd::Constant(1, 2, 0.5);

    Eigen::ArrayXd r;
    if (exterior)
      r = Eigen::VectorXd::LinSpaced(n + 1, 0.0, 1.0);
    else
    {
      const double h = 1.0 / static_cast<double>(n);
      r = Eigen::VectorXd::LinSpaced(n - 1, h, 1.0 - h);
    }

    if (lattice_type == lattice::type::gll_warped)
      r += warp_function(n, r);

    const int m = r.size();
    Eigen::ArrayX2d x(m * m, 2);
    int c = 0;
    for (int j = 0; j < m; ++j)
      for (int i = 0; i < m; ++i)
        x.row(c++) << r[i], r[j];

    return x;
  }
  case cell::type::hexahedron:
  {
    if (n == 0)
      return Eigen::ArrayXXd::Constant(1, 3, 0.5);

    Eigen::ArrayXd r;
    if (exterior)
      r = Eigen::VectorXd::LinSpaced(n + 1, 0.0, 1.0);
    else
    {
      const double h = 1.0 / static_cast<double>(n);
      r = Eigen::VectorXd::LinSpaced(n - 1, h, 1.0 - h);
    }
    if (lattice_type == lattice::type::gll_warped)
      r += warp_function(n, r);

    const int m = r.size();
    Eigen::ArrayXXd x(m * m * m, 3);
    int c = 0;
    for (int k = 0; k < m; ++k)
      for (int j = 0; j < m; ++j)
        for (int i = 0; i < m; ++i)
          x.row(c++) << r[i], r[j], r[k];

    return x;
  }
  case cell::type::triangle:
  {
    if (n == 0)
      return Eigen::ArrayXXd::Constant(1, 2, 1.0 / 3.0);

    // Warp points: see Hesthaven and Warburton, Nodal Discontinuous Galerkin
    // Methods, pp. 175-180

    const int b = exterior ? 0 : 1;

    // Points
    Eigen::ArrayX2d p((n - 3 * b + 1) * (n - 3 * b + 2) / 2, 2);

    // Displacement from GLL points in 1D, scaled by 1/(r(1-r))
    Eigen::ArrayXd r = Eigen::VectorXd::LinSpaced(2 * n + 1, 0.0, 1.0);
    Eigen::ArrayXd wbar = warp_function(n, r);
    const auto s = r.segment(1, 2 * n - 1);
    wbar.segment(1, 2 * n - 1) /= s * (1 - s);

    int c = 0;
    for (int j = b; j < (n - b + 1); ++j)
    {
      for (int i = b; i < (n - b + 1 - j); ++i)
      {
        const int l = n - j - i;
        const double x = r[2 * i];
        const double y = r[2 * j];
        const double a = r[2 * l];
        p.row(c) << x, y;
        if (lattice_type == lattice::type::gll_warped)
        {
          p(c, 0) += x * (a * wbar(n + i - l) + y * wbar(n + i - j));
          p(c, 1) += y * (a * wbar(n + j - l) + x * wbar(n + j - i));
        }

        ++c;
      }
    }

    return p;
  }
  case cell::type::tetrahedron:
  {
    if (n == 0)
      return Eigen::ArrayXXd::Constant(1, 3, 0.25);

    const int b = exterior ? 0 : 1;
    Eigen::ArrayX3d p((n - 4 * b + 1) * (n - 4 * b + 2) * (n - 4 * b + 3) / 6,
                      3);
    Eigen::ArrayXd r = Eigen::VectorXd::LinSpaced(2 * n + 1, 0.0, 1.0);
    Eigen::ArrayXd wbar = warp_function(n, r);
    const auto s = r.segment(1, 2 * n - 1);
    wbar.segment(1, 2 * n - 1) /= s * (1 - s);
    int c = 0;
    for (int k = b; k < (n - b + 1); ++k)
    {
      for (int j = b; j < (n - b + 1 - k); ++j)
      {
        for (int i = b; i < (n - b + 1 - j - k); ++i)
        {
          const int l = n - k - j - i;
          const double x = r[2 * i];
          const double y = r[2 * j];
          const double z = r[2 * k];
          const double a = r[2 * l];
          p.row(c) << x, y, z;
          if (lattice_type == lattice::type::gll_warped)
          {
            const double dx = x
                              * (a * wbar(n + i - l) + y * wbar(n + i - j)
                                 + z * wbar(n + i - k));
            const double dy = y
                              * (a * wbar(n + j - l) + z * wbar(n + j - k)
                                 + x * wbar(n + j - i));
            const double dz = z
                              * (a * wbar(n + k - l) + x * wbar(n + k - i)
                                 + y * wbar(n + k - j));
            p(c, 0) += dx;
            p(c, 1) += dy;
            p(c, 2) += dz;
          }

          ++c;
        }
      }
    }

    return p;
  }
  case cell::type::prism:
  {
    if (n == 0)
    {
      Eigen::ArrayXXd x = Eigen::ArrayXXd::Constant(1, 3, 1.0 / 3.0);
      x(0, 2) = 0.5;
      return x;
    }

    const Eigen::ArrayXXd tri_pts
        = lattice::create(cell::type::triangle, n, lattice_type, exterior);
    const Eigen::ArrayXXd line_pts
        = lattice::create(cell::type::interval, n, lattice_type, exterior);

    Eigen::ArrayX3d x(tri_pts.rows() * line_pts.rows(), 3);
    x.leftCols(2) = tri_pts.replicate(line_pts.rows(), 1);
    for (int i = 0; i < line_pts.rows(); ++i)
      x.block(i * tri_pts.rows(), 2, tri_pts.rows(), 1) = line_pts(i, 0);
    return x;
  }
  case cell::type::pyramid:
  {
    if (n == 0)
    {
      Eigen::ArrayXXd x = Eigen::ArrayXXd::Constant(1, 3, 0.4);
      x(0, 2) = 0.2;
      return x;
    }
    else
    {
      const double h = 1.0 / static_cast<double>(n);

      // Interpolate warp factor along interval
      std::tuple<Eigen::ArrayXXd, Eigen::ArrayXd> pw
          = quadrature::gauss_lobatto_legendre_line_rule(n + 1);
      Eigen::VectorXd pts = std::get<0>(pw) * 0.5;
      for (int i = 0; i < n + 1; ++i)
        pts[i] += (0.5 - static_cast<double>(i) / static_cast<double>(n));
      FiniteElement L = create_dlagrange(cell::type::interval, n);

      // Get interpolated value at r in range [-1, 1]
      auto w = [&](double r) {
        Eigen::ArrayXd rr = Eigen::ArrayXd::Constant(1, 0.5 * (r + 1.0));
        Eigen::VectorXd v = L.tabulate(0, rr)[0].row(0);
        return v.dot(pts);
      };

      int b = (exterior == false) ? 1 : 0;
      n -= b * 3;
      int m = (n + 1) * (n + 2) * (2 * n + 3) / 6;
      Eigen::ArrayX3d points(m, 3);
      int c = 0;
      for (int k = 0; k < n + 1; ++k)
        for (int j = 0; j < n + 1 - k; ++j)
          for (int i = 0; i < n + 1 - k; ++i)
          {
            double x = h * (i + b);
            double y = h * (j + b);
            double z = h * (k + b);

            if (lattice_type == lattice::type::gll_warped)
            {
              // Barycentric coordinates of triangle in x-z plane
              const double l1 = x;
              const double l2 = z;
              const double l3 = 1 - x - z;
              // Barycentric coordinates of triangle in y-z plane
              const double l4 = y;
              const double l5 = z;
              const double l6 = 1 - y - z;

              // b1-b6 are the blending factors for each edge
              double b1, f1, f2;
              if (std::fabs(l1) < 1e-12)
              {
                b1 = 1.0;
                f1 = 0.0;
                f2 = 0.0;
              }
              else
              {
                b1 = 2.0 * l3 / (2.0 * l3 + l1) * 2.0 * l2 / (2.0 * l2 + l1);
                f1 = l1 / (l1 + l4);
                f2 = l1 / (l1 + l6);
              }

              // r1-r4 are the edge positions for each of the z>0 edges
              // calculated so that they use the barycentric coordinates
              // of the triangle, if the point lies on a triangular face.
              // f1-f4 are face selecting functions, which blend between
              // adjacent triangular faces
              const double r1 = (l2 - l3) * f1 + (l5 - l6) * (1 - f1);
              const double r2 = (l2 - l3) * f2 + (l5 - l4) * (1 - f2);

              double b2;
              if (std::fabs(l2) < 1e-12)
                b2 = 1.0;
              else
                b2 = 2.0 * l3 / (2.0 * l3 + l2) * 2.0 * l1 / (2.0 * l1 + l2);

              double b3, f3, f4;
              if (std::fabs(l3) < 1e-12)
              {
                b3 = 1.0;
                f3 = 0.0;
                f4 = 0.0;
              }
              else
              {
                b3 = 2.0 * l2 / (2.0 * l2 + l3) * 2.0 * l1 / (2.0 * l1 + l3);
                f3 = l3 / (l3 + l4);
                f4 = l3 / (l3 + l6);
              }

              const double r3 = (l2 - l1) * f3 + (l5 - l6) * (1.0 - f3);
              const double r4 = (l2 - l1) * f4 + (l5 - l4) * (1.0 - f4);

              double b4;
              if (std::fabs(l4) < 1e-12)
                b4 = 1.0;
              else
                b4 = 2 * l6 / (2.0 * l6 + l4) * 2.0 * l5 / (2.0 * l5 + l4);

              double b5;
              if (std::fabs(l5) < 1e-12)
                b5 = 1.0;
              else
                b5 = 2.0 * l6 / (2.0 * l6 + l5) * 2.0 * l4 / (2.0 * l4 + l5);

              double b6;
              if (std::fabs(l6) < 1e-12)
                b6 = 1.0;
              else
                b6 = 2.0 * l4 / (2.0 * l4 + l6) * 2.0 * l5 / (2.0 * l5 + l6);

              double dx = -b3 * b4 * w(r3) - b3 * b6 * w(r4) + b2 * w(l1 - l3);
              double dy = -b1 * b6 * w(r2) - b3 * b6 * w(r4) + b5 * w(l4 - l6);
              double dz = b1 * b4 * w(r1) + b1 * b6 * w(r2) + b3 * b4 * w(r3)
                          + b3 * b6 * w(r4);

              x += dx;
              y += dy;
              z += dz;
            }

            points.row(c++) << x, y, z;
          }

      return points;
    }
  }
  default:
    throw std::runtime_error("Unsupported cell for lattice");
  }
}

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


#pragma once
#include <iostream>
#include <array>
#include <Eigen/Dense>
#include "Types/SFINAE.hpp"

///@file
///@brief Contains the class \ref mackey::Z_mod

namespace mackey {
	///////////////////////////////////
	///The class of Z/N coefficients where N is prime.

	///The operators are self explanatory. User must ensure that N is prime for division to work.
	/////////////////////////////////
	template<int64_t N, typename T = int64_t>
	class Z_mod {
	public:
		constexpr static int64_t order = N; ///<The N of Z/N
		T x; ///< A modulo N number
		Z_mod() : x(0) {} ///<Default value 0
		Z_mod(bool x) : x(x) {} ///<Initialize from 0,1
		Z_mod(int x); ///<Initialize from int
		Z_mod(int64_t x); ///<Initialize from 64bit int
		explicit operator char() const;
		explicit operator short() const;
		explicit operator int() const;
		explicit operator int64_t() const;
		explicit operator unsigned char() const;
		explicit operator unsigned short() const;
		explicit operator unsigned int() const;
		explicit operator uint64_t() const;
		Z_mod<N, T> operator +(Z_mod<N, T> b) const;
		Z_mod<N, T> operator -(Z_mod<N, T> b) const;
		Z_mod<N, T>& operator +=(Z_mod<N, T> b);
		Z_mod<N, T>& operator -=(Z_mod<N, T> b);
		Z_mod<N, T>& operator *=(Z_mod<N, T> b);
		Z_mod<N, T>& operator /=(Z_mod<N, T> b);
		bool operator ==(Z_mod<N, T> a) const;
		bool operator !=(Z_mod<N, T> a) const;
		bool operator <=(Z_mod<N, T> a) const; ///<Needed for Eigen pruning; standard order on 0,...,N-1
	};

	template<int64_t N, typename T>
	Z_mod<N, T> operator -(Z_mod<N, T> a);

	template<int64_t N, typename T>
	Z_mod<N, T> operator *(Z_mod<N, T> a, Z_mod<N, T> b); //Eigen needs this to be non member

	template<int64_t N, typename T>
	Z_mod<N, T> operator /(Z_mod<N, T> a, Z_mod<N, T> b);

	///The usual absolute value for integer and Z/N types
	template<typename T>
	T abs(T a);

	template<int64_t N, typename T> //Eigen needs this to be non member
	std::ostream& operator<<(std::ostream& out, const Z_mod<N, T> a);

	///The \f$\mathbf Z/2\f$ coefficients
	using Z2 = Z_mod<2, bool>;
}

///See the Eigen documentation for this.
namespace Eigen {
	using namespace mackey;

	///Specializing NumTraits to Z/nZ coefficients
	template<int64_t N, typename T>
	struct NumTraits<Z_mod<N, T>>
	{
		typedef Z_mod<N, T> Real;
		typedef Z_mod<N, T> Nested;
		typedef int Literal;
		enum {
			IsComplex = 0,
			IsInteger = 0,
			IsSigned = 0,
			RequireInitialization = 0,
			ReadCost = 1,
			AddCost = 1,
			MulCost = 1
		};
		static inline Z_mod<N, T> dummy_precision() { return Z_mod<N, T>(0); }
		static inline int digits10() { return 0; }
	};
}

#include "impl/Z_n.ipp"


#include <algorithm>
#include <cstdint>
#include <iostream>
#include <random>
#include <string>
#include <iomanip>
#include <chrono>

#include <Eigen/Eigen>
#include "sparse_gemm_problems.h"

using namespace Eigen;

#ifndef ITERS
#define ITERS 10
#endif

template<typename T, int U = ColMajor>
Matrix<T, Dynamic, Dynamic, U> generate_random_matrix(int rows, int cols) {
    return Matrix<T, Dynamic, Dynamic, U>::Random(rows, cols);
}

template<typename T>
Matrix<T, Dynamic, 1> generate_random_vector(int entries) {
    return Matrix<T, Dynamic, 1>::Random(entries);
}

template<typename T, int U>
std::tuple <int, int> time_sparse_gemv(const SparseMatrix<T, U>& sp_A, const Matrix<T, Dynamic, Dynamic, U>& A, const Matrix<T, Dynamic, 1>& B, Matrix<T, Dynamic, 1>& C) {

    // Try dense-dense multiplication
    auto start = std::chrono::steady_clock::now();
    for (int i = 0; i < ITERS; ++i) {
        C = A * B;
        C[0]++;  // dummy instruction to prevent optimizing away prev line
    }
    auto end = std::chrono::steady_clock::now();
    int d_time = static_cast<int>(std::chrono::duration<double, std::micro>(end - start).count() / ITERS);

    // Try sparse-dense multiplication
    start = std::chrono::steady_clock::now();
    for (int i = 0; i < ITERS; ++i) {
        C = sp_A * B;
        C[0]++;  // dummy instruction to prevent optimizing away prev line
    }
    end = std::chrono::steady_clock::now();
    int sp_time = static_cast<int>(std::chrono::duration<double, std::micro>(end - start).count() / ITERS);
    return std::tuple<int, int>(sp_time, d_time);
}

template<typename T, int U, int V, int W>
std::tuple <int, int> time_sparse_gemm(const SparseMatrix<T, U>& sp_A, const Matrix<T, Dynamic, Dynamic, U>& A, const Matrix<T, Dynamic, Dynamic, V>& B, Matrix<T, Dynamic, Dynamic, W>& C) {

    // Try dense-dense multiplication
    auto start = std::chrono::steady_clock::now();
    for (int i = 0; i < ITERS; ++i) {
        C = A * B;
        C(0,0)++;  // dummy instruction to prevent optimizing away prev line
    }
    auto end = std::chrono::steady_clock::now();
    int d_time = static_cast<int>(std::chrono::duration<double, std::micro>(end - start).count() / ITERS);

    // Try sparse-dense multiplication
    start = std::chrono::steady_clock::now();
    for (int i = 0; i < ITERS; ++i) {
        C = sp_A * B;
        C(0,0)++;  // dummy instruction to prevent optimizing away prev line
    }
    end = std::chrono::steady_clock::now();
    int sp_time = static_cast<int>(std::chrono::duration<double, std::micro>(end - start).count() / ITERS);
    return std::tuple<int, int>(sp_time, d_time);
}

template<typename T>
std::tuple<int, int> time_sparse_bench_helper(int m, int n, int k, float sparsity, std::default_random_engine & e, std::uniform_real_distribution<double> & rng) {

    const int U = RowMajor;

    // Note: We've determined empirically that B,C in ColMajor is best for
    // both sparse and dense gemm implementations.
    const int V = ColMajor;
    const int W = ColMajor;

    auto A = generate_random_matrix<T, U>(m, k);
    for (int j = 0; j < k; ++j) {
        for (int i = 0; i < m; ++i) {
            if (rng(e) < sparsity) {
                A(i, j) = 0;
            }
        }
    }
    SparseMatrix<T, U> sp_A = A.sparseView();

    if (n == 1) {
        auto B = generate_random_vector<T>(k);
        auto C = generate_random_vector<T>(m);
        return time_sparse_gemv<T, U>(sp_A, A, B, C);
    }
    else {
        auto B = generate_random_matrix<T, V>(k, n);
        auto C = generate_random_matrix<T, W>(m, n);
        return time_sparse_gemm<T, U, V, W>(sp_A, A, B, C);
    }
}

int main() {

    // Set up RNG
    std::random_device r;
    std::default_random_engine e(r());
    std::uniform_real_distribution<double> rng(0, 1);

    std::cout << std::setw(30) << "Times" << std::endl;
    std::cout << std::setfill('-') << std::setw(110) << "-" << std::endl;
    std::cout << std::setfill(' ');
    std::cout << "    m       n      k      a_t    b_t    sparsity  precision  sparse time (usec) dense time (usec)   speedup " << std::endl;

    std::vector<std::string> types = {"uint8_t", "float"};

    for (const auto &type_name : types) {

        for (const auto &problem : inference_device_set) {

            int m,n,k;
            bool a_t, b_t;
            float sparsity;

            std::tie(m, n, k, a_t, b_t, sparsity) = problem;

            std::cout << std::setw(7) << m;
            std::cout << std::setw(7) << n;
            std::cout << std::setw(7) << k;
            std::cout << std::setw(7) << a_t ? "t" : "n";
            std::cout << std::setw(7) << b_t ? "t" : "n";
            std::cout << std::setw(11) << sparsity;
            std::cout << std::setw(12) << type_name;

            int sp_time, d_time;

            if (type_name == "uint8_t") {
                std::tie(sp_time, d_time) = time_sparse_bench_helper<std::uint8_t>(m, n, k, sparsity, e, rng);
            } else if (type_name == "float") {
                std::tie(sp_time, d_time) = time_sparse_bench_helper<float>(m, n, k, sparsity, e, rng);
            } else {
                throw std::runtime_error("Unsupported type_name");
            }

            std::cout << std::setw(15) << sp_time;
            std::cout << std::setw(15) << d_time;
            std::cout << std::setw(20) << float(d_time)/sp_time;
            std::cout << std::endl;
        }
    }

    return 0;
}


#include <unistd.h>	// sleep
#include <armadillo>
#include "Functions.h"

using namespace arma;
using namespace std;

/* Description: Parallel Stochastic Iterative Hard Thresholding (StoIHT) algorithm
 to approximate the vector x from measurements u = A*x.
publication:  Linear Convergence of Stochastic Iterative Greedy Algorithms with Sparse Constraints
https://arxiv.org/abs/1407.0088
*/
vec parallel_Sto_IHT(const mat A, const vec y, const int sparsity, const vec prob_vec,
		const unsigned int max_iter, const double gamma, const double tol,
		unsigned int &num_iters, const simulation_parameters simulation_params){
	// signal parameters
	const unsigned int sig_dim = A.n_cols;
	const unsigned int meas_num = A.n_rows;
	const unsigned int num_block = prob_vec.n_elem;
	const unsigned int block_size = meas_num/num_block;

	// initialization of variables that are shared among cores
	vec x_hat(sig_dim,fill::zeros);	// estimation of the signal
	bool done = false;		// flag to check the convergence criteria
	unsigned int i = 0;		// total number of iterations

	// parallel section of the code starts here
	#pragma omp parallel num_threads(simulation_params.num_cores)
	{

	// initializaiotn of variables that are local to each core
	vec x_hat_local(sig_dim,fill::zeros);
	unsigned int selected_block,first_ind_block,last_ind_block;
	mat A_block;			// submatrix of A
	vec y_block,gradient,b;
	uvec sorted_ind,est_supp;

	// iterations to find the solutions
	while(!done){
		// master thread uses the tally vector to check the convergence criteria
		if (omp_get_thread_num() == 0 ){
			// check exit criteria
			if (norm (y - A*x_hat) < tol || i >= max_iter){
				done = true;	// the flag 'done' is shared among the cores
			}
			if (omp_get_num_threads()  > 1){
				continue;
			}
		}

		i++;
		// randomize
		selected_block = floor(randu()*num_block);
		first_ind_block = block_size*selected_block;
		last_ind_block = block_size*(selected_block+1)-1;

		// Proxy
		A_block = A.rows(first_ind_block,last_ind_block);
		y_block = y.subvec(first_ind_block,last_ind_block);
		#pragma omp critical
		{x_hat_local = x_hat;	}	// read the global estimate while memory is locked

		gradient = -2 * A_block.t() *	(y_block - A_block*x_hat_local);
		b = x_hat_local - gradient * gamma/(num_block*prob_vec(selected_block));

		// Identify
		sorted_ind = sort_index(abs(b),"descend");
		est_supp = sorted_ind(span(0,sparsity - 1));

		// Estimate
		x_hat_local.zeros();
		x_hat_local(est_supp) = b(est_supp);

		#pragma omp critical
		{x_hat = x_hat_local;}	//write the global estimate while memory is locked

	}
	}
	// parallel section of the code ends here
	num_iters = i;
	return x_hat;
}

/* Asynchronous StoIHT Iteration
Algorithm 2 in An Asynchronous Parallel Approach to Sparse Recovery
https://arxiv.org/abs/1701.03458*/

uvec Sto_IHT_async_iteration(vec &x_hat,const vec &tally,const mat &A,const vec &y,
 			const int sparsity, const vec prob_vec, const double gamma){
	// randomize
	const unsigned int num_block = prob_vec.n_elem;
	const unsigned int block_size = y.n_elem / num_block;
	const unsigned int selected_block = floor(randu()*num_block);
	const unsigned int first_ind_block = block_size*selected_block;
	const unsigned int last_ind_block = block_size*(selected_block+1)-1;

	// Proxy
	const mat A_block = A.rows(first_ind_block,last_ind_block);
	const vec y_block = y.subvec(first_ind_block,last_ind_block);
	const vec gradient = -2 * A_block.t() *	(y_block - A_block*x_hat);
	const vec b = x_hat - gradient * gamma/(num_block*prob_vec(selected_block));

	x_hat.zeros();  // this variable is local to each core. NOTE: passed by reference
	// Identify using b (local)
	uvec sorted_ind = sort_index(abs(b),"descend");
	const uvec est_supp_local = sorted_ind(span(0,sparsity - 1));
	// Estimate using b (local)
	x_hat(est_supp_local) = b(est_supp_local);

	// Identify using tally (collective)
	sorted_ind = sort_index(tally,"descend");
	const uvec est_supp_collective = sorted_ind(span(0,sparsity - 1));
	// Estimate using tally (collective)
	x_hat(est_supp_collective) = b(est_supp_collective);

	return 	est_supp_local;
}

void update_tally(vec &tally,const uvec est_supp_local,const uvec prev_est_supp,const unsigned int iter_local){
	/* update the tally score according the rules in:
	An Asynchronous Parallel Approach to Sparse Recovery
	https://arxiv.org/abs/1701.03458*/
	tally(est_supp_local) += iter_local;
	if (iter_local >= 2){
		tally(prev_est_supp) -= iter_local-1;
	}
	return;
}

/* Description: Parallel Stochastic Iterative Hard Thresholding (StoIHT) algorithm with tally score to approximate the vector x from measurements u = A*x.
Publication: An Asynchronous Parallel Approach to Sparse Recovery
https://arxiv.org/abs/1701.03458*/

vec tally_Sto_IHT(const mat &A, const vec &y, const int sparsity, const vec prob_vec,
		const unsigned int max_iter, const double gamma,const double tol,
		unsigned int &num_iters, const simulation_parameters simulation_params){
	uvec slow_cores;
	set_slow_cores(slow_cores, simulation_params);

	const unsigned int sig_dim = A.n_cols;

	// initialization of variables that are shared among cores
	vec tally(sig_dim,fill::zeros);		// vector of tally scores
	vec x_hat_total(sig_dim,fill::zeros);	// estimation of the signal
	bool done = false;			// flag to check the convergence criteria
	unsigned int i = 0;			// total number of iterations

	// parallel section of the code starts here
	#pragma omp parallel num_threads(simulation_params.num_cores)
	{
	//#pragma omp single // a single core executes the following line

	// initializaiotn of variables that are local to each core
	uvec prev_est_supp;			// estimated support in previous iteration
	vec x_hat_local(sig_dim,fill::zeros);  	// this is local to each core
	unsigned int iter_local = 0;		// number of iteration for this core

	// iterations to find the solutions
	while(!done){
		// master thread uses the tally vector to check the convergence criteria
		if (omp_get_thread_num() == 0){
			const uvec sorted_ind = sort_index(abs(tally),"descend");
			const uvec est_supp = sorted_ind(span(0,sparsity - 1));
			const mat A_supp = A.cols(est_supp);
			x_hat_local.zeros();
			x_hat_local(est_supp) = solve(A_supp,y);
			if (norm (y - A*x_hat_local) < tol || i >= max_iter){
				x_hat_total = x_hat_local;
				done = true;
			}
		}

		//slow cores sleep for  simulation_params.sleep_slow_cores microseconds
		if (any( slow_cores == omp_get_thread_num()) ){
			usleep(simulation_params.sleep_slow_cores);
		}

		i++;
		iter_local++;

		// update the local estimate of the support
		uvec est_supp_local;
        try{
        est_supp_local = Sto_IHT_async_iteration(x_hat_local, tally, A, y,
            sparsity, prob_vec, gamma);
        }
        catch(std::logic_error)
        {
            // algorithm did not converge
            i = max_iter;
            done = true;
        }
		// Update Tally
        update_tally(tally,est_supp_local,prev_est_supp,iter_local);

		prev_est_supp = est_supp_local;
	}
	}
	// parallel section of the code ends here
	//cout << endl;
	num_iters = i;
	return x_hat_total;
}


#include <iostream>
#include <vector>
#include <fstream>
#include <cmath>
#include <chrono>

using namespace std;

#include <boost/timer.hpp>

// for sophus
#include <sophus/se3.hpp>

using Sophus::SE3d;

// for eigen
#include <Eigen/Core>
#include <Eigen/Geometry>

using namespace Eigen;

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

#include "plot.h"

using namespace cv;

/**
 * Dataset from:
 *
 *   http://rpg.ifi.uzh.ch/datasets/remode_test_data.zip
 *
 * */

// ------------------------------------------------------------------
// parameters
const int boarder = 20;
const int width = 640;
const int height = 480;
const double fx = 481.2f;
const double fy = -480.0f;
const double cx = 319.5f;
const double cy = 239.5f;
const int ncc_window_size = 3;
const int ncc_area = (2 * ncc_window_size + 1) * (2 * ncc_window_size + 1);
const double min_cov = 0.1;
const double max_cov = 10;
const double epsilon = 1e-10;
// ------------------------------------------------------------------

inline double getBilinearInterpolatedValue_eigen(const Mat &img, const Eigen::Vector2d &pt) {
    uchar *d = &img.data[int(pt[1]) * img.step + int(pt[0])];
    double xx = pt[0] - floor(pt[0]);
    double yy = pt[1] - floor(pt[1]);
    return ((1 - xx) * (1 - yy) * double(d[0]) +
            xx * (1 - yy) * double(d[1]) +
            (1 - xx) * yy * double(d[img.step]) +
            xx * yy * double(d[img.step + 1])) / 255.0;
}

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

inline Vector3d px2cam(const Vector2d& px) {
    return Vector3d(
        (px(0, 0) - cx) / fx,
        (px(1, 0) - cy) / fy,
        1
    );
}

inline Vector2d cam2px(const Vector3d& p_cam) {
    return Vector2d(
        p_cam(0, 0) * fx / p_cam(2, 0) + cx,
        p_cam(1, 0) * fy / p_cam(2, 0) + cy
    );
}

inline bool inside(const Vector2d &pt) {
    return pt(0, 0) >= boarder && pt(1, 0) >= boarder
           && pt(0, 0) + boarder < width && pt(1, 0) + boarder <= height;
}

bool readDatasetFiles(
    const string &path,
    vector<string> &color_image_files,
    vector<SE3d> &poses,
    cv::Mat &ref_depth
);

void evaludateDepth(const Mat &depth_truth, const Mat &depth_estimate);
// ------------------------------------------------------------------

double ZNCC(const cv::Mat& im1, const Eigen::Vector2d& pt1, const cv::Mat& im2, Eigen::Vector2d& pt2)
{
    // no need to consider block partly outside because of boarder
    // std::vector<double> v1(ncc_area, 0.0), v2(ncc_area, 0.0); // much slower
    double v1[ncc_area], v2[ncc_area];
    double s1 = 0.0, s2 = 0.0;
    int idx = 0;
    for (int i = -ncc_window_size; i <= ncc_window_size; ++i)
    {
        for (int j = -ncc_window_size; j <= ncc_window_size; ++j)
        {
            double val_1 = static_cast<double>(im1.at<uchar>(pt1.y()+i, pt1.x()+j)) / 255;
            Eigen::Vector2d temp_p2 = pt2;
            temp_p2[0] += j;
            temp_p2[1] += i;
            double val_2 = getBilinearInterpolatedValue_eigen(im2, temp_p2);

            s1 += val_1;
            s2 += val_2;
            v1[idx] = val_1;
            v2[idx] = val_2;
            ++idx;
        }
    }

    double mean_1 = s1 / ncc_area;
    double mean_2 = s2 / ncc_area;

    double numerator = 0.0;
    double den1 = 0.0, den2 = 0.0;
    for (int i = 0; i < ncc_area; ++i)
    {
        double zv1 = v1[i] - mean_1;
        double zv2 = v2[i] - mean_2;
        numerator += zv1*zv2;
        den1 += zv1 * zv1;
        den2 += zv2 * zv2;
    }
    auto zncc =  numerator / (std::sqrt(den1 * den2 + epsilon));
    // std::cout << "zncc = " << zncc << "\n";
    return zncc;
}

bool epipolar_search(const cv::Mat& ref, const cv::Mat& cur, const Sophus::SE3d& Tcr, const Eigen::Vector2d& pt, double depth_mu, double depth_sigma2, Eigen::Vector2d& best_pc, Eigen::Vector2d& epipolar_dir)
{
    double depth_sigma = std::sqrt(depth_sigma2);
    double dmax = depth_mu + 3 * depth_sigma;
    double dmin = depth_mu - 3 * depth_sigma;
    dmin = std::max(0.1, dmin);

    Eigen::Vector3d pn((pt.x()-cx) / fx, (pt.y() - cy) / fy, 1.0);
    pn.normalize();
    Eigen::Vector3d P_max = pn * dmax;
    Eigen::Vector3d P_min = pn * dmin;
    Eigen::Vector3d P_mu = pn * depth_mu;

    Eigen::Vector2d pc_max = cam2px(Tcr * P_max);
    Eigen::Vector2d pc_min = cam2px(Tcr * P_min);
    Eigen::Vector2d pc_mu = cam2px(Tcr * P_mu);

    Eigen::Vector2d epipolar_line = pc_max - pc_min;
    epipolar_dir = epipolar_line.normalized();

    double step = 0.7;
    int nb_samples = std::ceil(epipolar_line.norm() / step);

    double half_range = 0.5 * epipolar_line.norm();
    if (half_range > 100) half_range = 100;

    Eigen::Vector2d p = pc_min;
    double best_zncc = -1.0;
    best_pc = pc_mu;

    // for (int i = 0; i < nb_samples; ++i)
    for (double l = -half_range; l<= half_range; l+= 0.7)
    {
        Eigen::Vector2d p = pc_mu + l * epipolar_dir;

        if (p.x() < boarder || p.x() >= width-boarder || p.y() < boarder || p.y() >= height-boarder)
            continue; // p is outside the cur image

        double zncc = ZNCC(ref, pt, cur, p);
        if (zncc > best_zncc)
        {
            best_zncc = zncc;
            best_pc = p;
        }

        // p += epipolar_dir * step;
    }

    // std::cout << best_zncc << "\n";
    if (best_zncc < 0.85)
        return false;
    else
        return true;
}

void update_depth_filter(const Eigen::Vector2d& pr, const Eigen::Vector2d& pc, const Sophus::SE3d& Tcr, const Eigen::Vector2d& epipolar_dir, cv::Mat& depth, cv::Mat& cov2)
{
    Sophus::SE3d Trc = Tcr.inverse();

    Eigen::Vector3d fr = px2cam(pr);
    fr.normalize();
    Eigen::Vector3d fc = px2cam(pc);
    fc.normalize();
    Eigen::Vector3d f2 = Trc.so3() * fc;
    Eigen::Vector3d trc = Trc.translation();

    // Solve the system of equation for triangulating depth
    Eigen::Matrix2d A;
    Eigen::Vector2d b;
    A(0, 0) = fr.dot(fr);
    A(0, 1) = -fr.dot(f2);
    A(1, 0) = f2.dot(fr);
    A(1, 1) = -f2.dot(f2);
    b[0] = fr.dot(trc);
    b[1] = f2.dot(trc);
    Eigen::Vector2d res = A.inverse() * b;
    Eigen::Vector3d P1 = fr * res[0];
    Eigen::Vector3d P2 = trc + fc * res[1];
    Eigen::Vector3d P_est = (P1 + P2) * 0.5;
    double depth_obs = P_est.norm(); //depth obs

    // Estimate depth uncertainty
    Eigen::Vector3d P = fr * depth_obs;
    Eigen::Vector3d a = P - trc;
    Eigen::Vector3d t = trc.normalized();
    double alpha = std::acos(fr.dot(t));
    double beta = std::acos(a.normalized().dot(-t));
    Eigen::Vector2d pc2 = pc + epipolar_dir;
    Eigen::Vector3d fc2 = px2cam(pc2);
    fc2.normalize();
    double beta_2 = std::acos(fc2.dot(-t));
    double gamma = M_PI - alpha - beta_2;
    double d_noise = trc.norm() * std::sin(beta_2) / std::sin(gamma); // sinus law
    double sigma_obs = depth_obs - d_noise;
    double sigma2_obs = sigma_obs * sigma_obs; // sigma2 obs

    // Depth fusion
    double d = depth.at<double>(static_cast<int>(pr.y()), static_cast<int>(pr.x()));
    double sigma2 = cov2.at<double>(static_cast<int>(pr.y()), static_cast<int>(pr.x()));

    double d_fused = (sigma2_obs * d + sigma2 * depth_obs) / (sigma2 + sigma2_obs);
    double sigma2_fused = (sigma2 * sigma2_obs) / (sigma2 + sigma2_obs);

    depth.at<double>(static_cast<int>(pr.y()), static_cast<int>(pr.x())) = d_fused;
    cov2.at<double>(static_cast<int>(pr.y()), static_cast<int>(pr.x())) = sigma2_fused;
}

void update(const cv::Mat& ref, const cv::Mat& cur, const Sophus::SE3d& Tcr, cv::Mat &depth, cv::Mat &cov2)
{
    Eigen::Vector2d pc;
    Eigen::Vector2d epipolar_dir;
    for (int j = boarder; j < width-boarder; ++j)
    {
        for (int i = boarder; i < height-boarder; ++i)
        {
            double depth_mu = depth.at<double>(i, j);
            double depth_sigma2 = cov2.at<double>(i, j);
            if (depth_sigma2 < min_cov || depth_sigma2 > max_cov)
                continue;
            Eigen::Vector2d pr(j, i);
            bool found = epipolar_search(ref, cur, Tcr, pr, depth_mu, depth_sigma2, pc, epipolar_dir);
            if (!found)
                continue;

            // showEpipolarMatch(ref, cur, pr, pc);

            update_depth_filter(pr, pc, Tcr, epipolar_dir, depth, cov2);
        }
    }
    // std::cout << depth << "\n";

}

int main(int argc, char **argv) {
    if (argc != 2) {
        cout << "Usage: dense_mapping path_to_test_dataset" << endl;
        return -1;
    }

    // Read dataset
    vector<string> color_image_files;
    vector<SE3d> poses_TWC;
    Mat ref_depth;
    bool ret = readDatasetFiles(argv[1], color_image_files, poses_TWC, ref_depth);
    if (ret == false) {
        cout << "Reading image files failed!" << endl;
        return -1;
    }
    cout << "read total " << color_image_files.size() << " files." << endl;

    // Initial depth image
    Mat ref = imread(color_image_files[0], 0); // gray-scale image
    SE3d pose_ref_TWC = poses_TWC[0];
    double init_depth = 3.0;
    double init_cov2 = 3.0;
    Mat depth(height, width, CV_64F, init_depth);
    Mat depth_cov2(height, width, CV_64F, init_cov2);

    for (int index = 1; index < color_image_files.size(); index++) {
        cout << "*** loop " << index << " ***" << endl;
        Mat curr = imread(color_image_files[index], 0);
        if (curr.data == nullptr) continue;
        SE3d pose_curr_TWC = poses_TWC[index];
        SE3d pose_T_C_R = pose_curr_TWC.inverse() * pose_ref_TWC;   // T_C_W * T_W_R = T_C_R
        chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
        update(ref, curr, pose_T_C_R, depth, depth_cov2);
        chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
        auto time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
        std::cout << "Time used: " << time_used.count() << "s\n";
        evaludateDepth(ref_depth, depth);
        plotDepth(ref_depth, depth);
        plotCur(curr);
    }

    cout << "estimation returns, saving depth map ..." << endl;
    imwrite("depth.png", depth);
    cout << "done." << endl;

    return 0;
}

bool readDatasetFiles(
    const string &path,
    vector<string> &color_image_files,
    std::vector<SE3d> &poses,
    cv::Mat &ref_depth) {
    ifstream fin(path + "/first_200_frames_traj_over_table_input_sequence.txt");
    if (!fin) return false;

    while (!fin.eof()) {
        // 数据格式：图像文件名 tx, ty, tz, qx, qy, qz, qw ，注意是 TWC 而非 TCW
        string image;
        fin >> image;
        double data[7];
        for (double &d:data) fin >> d;

        color_image_files.push_back(path + string("/images/") + image);
        poses.push_back(
            SE3d(Quaterniond(data[6], data[3], data[4], data[5]),
                 Vector3d(data[0], data[1], data[2]))
        );
        if (!fin.good()) break;
    }
    fin.close();

    // load reference depth
    fin.open(path + "/depthmaps/scene_000.depth");
    ref_depth = cv::Mat(height, width, CV_64F);
    if (!fin) return false;
    for (int y = 0; y < height; y++)
        for (int x = 0; x < width; x++) {
            double depth = 0;
            fin >> depth;
            ref_depth.ptr<double>(y)[x] = depth / 100.0;
        }

    return true;
}

void evaludateDepth(const Mat &depth_truth, const Mat &depth_estimate) {
    double ave_depth_error = 0;
    double ave_depth_error_sq = 0;
    int cnt_depth_data = 0;
    for (int y = boarder; y < depth_truth.rows - boarder; y++)
        for (int x = boarder; x < depth_truth.cols - boarder; x++) {
            double error = depth_truth.ptr<double>(y)[x] - depth_estimate.ptr<double>(y)[x];
            ave_depth_error += error;
            ave_depth_error_sq += error * error;
            cnt_depth_data++;
        }
    ave_depth_error /= cnt_depth_data;
    ave_depth_error_sq /= cnt_depth_data;

    cout << "Average squared error = " << ave_depth_error_sq << ", average error: " << ave_depth_error << endl;
}


/** @file
 *****************************************************************************

 Arithmetic in the finite field Fp, for prime p of fixed length using
 NTL as the backend.

 *****************************************************************************
 * @author     Samir Menon, Brennan Shacklett, and David J. Wu
 * @copyright  MIT license (see LICENSE file)
 *****************************************************************************/

#ifndef NTLFP_TCC_
#define NTLFP_TCC_

#include <cassert>
#include <cstdlib>
#include <cmath>
#include <NTL/ZZ.h>

#include <libff/algebra/fields/fp_aux.tcc>
#include <libff/algebra/fields/field_utils.hpp>

namespace libsnark {

template<unsigned long modulus>
NTLFp_model<modulus>::NTLFp_model()
{
    NTL::ZZ_p::init(mod_zz());
}

template<unsigned long modulus>
NTLFp_model<modulus>::NTLFp_model(long x)
{
    NTL::ZZ_p::init(mod_zz());
    this->value = x;
}

template <unsigned long modulus>
NTLFp_model<modulus>::NTLFp_model(const NTLFp_model &other)
{
    NTL::ZZ_p::init(mod_zz());
    this->value = other.value;
}

template<unsigned long modulus>
bool NTLFp_model<modulus>::operator==(const NTLFp_model& other) const
{
    return (this->value == other.value);
}

template<unsigned long modulus>
bool NTLFp_model<modulus>::operator!=(const NTLFp_model& other) const
{
    return (this->value != other.value);
}

template<unsigned long modulus>
bool NTLFp_model<modulus>::is_zero() const
{
    return NTL::IsZero(this->value);
}

template<unsigned long modulus>
void NTLFp_model<modulus>::print() const
{
    std::cout << *this;
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::zero()
{
    NTL::ZZ_p::init(mod_zz());
    NTLFp_model<modulus> z(NTL::ZZ_p::zero());
    return z;
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::one()
{
    NTL::ZZ_p::init(mod_zz());
    NTLFp_model<modulus> o(NTL::ZZ_p::zero() + 1);
    return o;
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator+=(const NTLFp_model<modulus>& other)
{
    NTL::ZZ_p::init(mod_zz());
    this->value += other.value;
    return *this;
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator-=(const NTLFp_model<modulus>& other)
{
    NTL::ZZ_p::init(mod_zz());
    this->value -= other.value;
    return *this;
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator*=(const NTLFp_model<modulus>& other)
{
    NTL::ZZ_p::init(mod_zz());
    this->value *= other.value;
    return *this;
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator^=(const NTLFp_model<modulus>& other)
{
    return this^=other.as_long();
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator^=(const libff::bigint<1>& pwr)
{
    return this^=pwr.as_ulong();
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::operator^=(const unsigned long pwr)
{
    NTL::ZZ_p::init(mod_zz());
    this->value = NTL::power(this->value, pwr);
    return (*this);
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator+(const NTLFp_model<modulus>& other) const
{
    NTLFp_model<modulus> r(*this);
    return (r += other);
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator-(const NTLFp_model<modulus>& other) const
{
    NTLFp_model<modulus> r(*this);
    return (r -= other);
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator*(const NTLFp_model<modulus>& other) const
{
    NTLFp_model<modulus> r;
    NTL::mul(r.value, value, other.value);
    r.value = value * other.value;
    return r;
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator^(const NTLFp_model<modulus>& other) const
{
    NTLFp_model<modulus> r(*this);
    return (r ^= other);
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator^(const unsigned long pwr) const
{
    NTLFp_model<modulus> r(*this);
    return (r ^= pwr);
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator^(const libff::bigint<1>& pwr) const
{
    NTLFp_model<modulus> r(*this);
    return (r ^= pwr.as_ulong());
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::operator-() const
{
    NTLFp_model<modulus> r(modulus - this->value);
    return r;
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::squared() const
{
    NTLFp_model<modulus> r(*this);
    return (r *= r);
}

template<unsigned long modulus>
NTLFp_model<modulus>& NTLFp_model<modulus>::invert()
{
    NTL::ZZ_p inverse = 1 / this->value;;
    this->value = inverse;
    return *this;
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::inverse() const
{
    NTLFp_model<modulus> r(*this);
    return r.invert();
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::random_element()
{
    NTL::ZZ_p::init(mod_zz());
    NTLFp_model<modulus> r;
    r.value = NTL::ZZ_p(NTL::RandomBnd(modulus));
    return r;
}

template<unsigned long modulus>
void NTLFp_model<modulus>::get_s_and_t(unsigned long& s, unsigned long& t)
{
    s = modulus - 1;
    t = 0;
    while (s % 2 == 0) {
        s /= 2;
        t++;
    }
}

template<unsigned long modulus>
NTLFp_model<modulus> NTLFp_model<modulus>::sqrt() const
{
    NTL::ZZ_p::init(mod_zz());
    NTL::ZZ_p sqrt = NTL::to_ZZ_p(NTL::SqrRootMod(NTL::rep(this->value), mod_zz()));
    NTLFp_model<modulus> root;
    root.value = sqrt;

    return root;
}

template<unsigned long modulus>
std::ostream& operator<<(std::ostream &out, const NTLFp_model<modulus> &p)
{
    out << p.value;
    return out;
}

template<unsigned long modulus>
std::istream& operator>>(std::istream &in, NTLFp_model<modulus> &p)
{
    in >> p.value;
    return in;
}

} // libsnark

#endif // NTLFP_TCC_


#include <iostream>
#include <boost/math/special_functions/daubechies_scaling.hpp>
#include <boost/math/special_functions/chebyshev_transform.hpp>

template<typename Real, int p>
void bootstrap()
{
    std::cout << "Computing phi. . .\n";
    auto phi = boost::math::daubechies_scaling<Real, p>();
    std::cout << "Computing Chebyshev transform of phi.\n";
    auto cheb = boost::math::chebyshev_transform(phi, phi.support().first, phi.support().second);
    std::cout << "Number of coefficients = " << cheb.coefficients().size() << "\n";
}

int main()
{
    bootstrap<long double, 9>();
}