//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#include <boost/simd/function/ifrexp.hpp>
#include <boost/simd/constant/nbmantissabits.hpp>
#include <boost/simd/constant/halfeps.hpp>
#include <scalar_test.hpp>

namespace bs = boost::simd;
namespace bd = boost::dispatch;

STF_CASE_TPL("Check basic behavior of ifrexp", STF_IEEE_TYPES)
{
  STF_EXPR_IS ( (bs::ifrexp(T(0)))
              , (std::pair<T,bd::as_integer_t<T,signed>>)
              );

  auto p = bs::ifrexp(T(1));
  STF_EQUAL(p.first  , T(0.5));
  STF_EQUAL(p.second , T(1));
}

STF_CASE_TPL("Check behavior of ifrexp on Zero", STF_IEEE_TYPES)
{
  auto r = bs::ifrexp(T(0));

  STF_EQUAL (r.first , T(0));
  STF_EQUAL (r.second, T(0));
  STF_EQUAL (ldexp(r.first,r.second), T(0));
}

STF_CASE_TPL("Check behavior of ifrexp on Valmax", STF_IEEE_TYPES)
{
  auto r = bs::ifrexp(bs::Valmax<T>());

  STF_ULP_EQUAL (r.first , T(1)-bs::Halfeps<T>(), 1);
  STF_EQUAL     (r.second, bs::Limitexponent<T>());
  STF_EQUAL     (ldexp(r.first,r.second),bs::Valmax<T>());
}


// In KAI C++ 3.2c, restrict causes problems for copy propagation.
// Temporary kludge is to disable use of the restrict keyword.

#define BZ_DISABLE_RESTRICT

#include <blitz/vector.h>
#include <blitz/array.h>
#include <blitz/rand-uniform.h>
#include <blitz/benchext.h>

#ifdef BENCHMARK_VALARRAY
#include <valarray>
#endif

BZ_USING_NAMESPACE(blitz)

#ifdef BZ_FORTRAN_SYMBOLS_WITH_TRAILING_UNDERSCORES
 #define fdaxpy   fdaxpy_
 #define f90daxpy f90daxpy_
 #define fidaxpy  fidaxpy_
 #define fidaxpyo fidaxpyo_
#endif

#ifdef BZ_FORTRAN_SYMBOLS_CAPS
 #define fdaxpy   FDAXPY
 #define f90daxpy F90DAXPY
 #define fidaxpy  FIDAXPY
 #define fidaxpyo FIDAXPYO
#endif

extern "C" {
  void fdaxpy(const int& N, const double& da, double* x,
    const int& xstride, const double* y, const int& ystride);

  void f90daxpy(const double& a, double* x,
    const double* y, const int& length, const int& iters);

  void fidaxpy(const double& a, double* x, const double* y,
    const int& length, const int& iters);

  void fidaxpyo(const double& a, double* x, const double* y,
    const int& length, const int& iters);
}

void daxpyVectorVersion(BenchmarkExt<int>& bench, double a, double b);
void daxpyArrayVersion(BenchmarkExt<int>& bench, double a);
void daxpyF77Version(BenchmarkExt<int>& bench, double a);
void daxpyBLASVersion(BenchmarkExt<int>& bench, double a);
void daxpyF90Version(BenchmarkExt<int>& bench, double a);

#ifdef BENCHMARK_VALARRAY
void daxpyValarrayVersion(BenchmarkExt<int>& bench, double a);
#endif

int main()
{

#ifdef BENCHMARK_VALARRAY
    int numBenchmarks = 6;
#else
    int numBenchmarks = 5;
#endif

    BenchmarkExt<int> bench("DAXPY Benchmark", numBenchmarks);

    const int numSizes = 19;
    bench.setNumParameters(numSizes);
    bench.setRateDescription("Mflops/s");

    Vector<int> parameters(numSizes);
    Vector<long> iters(numSizes);
    Vector<double> flops(numSizes);

    for (int i=0; i < numSizes; ++i)
    {
        parameters[i] = pow(10.0, (i+1)/4.0);
        iters[i] = 50000000L / parameters[i];
        if (iters[i] < 2)
            iters[i] = 2;
        flops[i] = 2 * parameters[i] * 2;
    }

    bench.setParameterVector(parameters);
    bench.setIterations(iters);
    bench.setOpsPerIteration(flops);

    bench.beginBenchmarking();

    float a = .398498293819823;

    daxpyVectorVersion(bench, a, -a);
    daxpyArrayVersion(bench, a);
    daxpyF77Version(bench, a);
    daxpyBLASVersion(bench, a);
    daxpyF90Version(bench, a);

#ifdef BENCHMARK_VALARRAY
    daxpyValarrayVersion(bench, a);
#endif

    bench.endBenchmarking();

    bench.saveMatlabGraph("daxpy2.m");

    return 0;
}

void initializeRandomDouble(double* data, int numElements, int stride = 1)
{
    static Random<Uniform> rnd;

    for (int i=0; i < numElements; ++i)
        data[i*stride] = rnd.random();
}

template<class T>
void initializeArray(T& array, int numElements)
{
    static Random<Uniform> rnd;

    for (size_t i=0; i < numElements; ++i)
        array[i] = rnd.random();
}

void daxpyVectorVersion(BenchmarkExt<int>& bench, double a, double b)
{
    bench.beginImplementation("Vector<T>");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "Vector<T>: N = " << N << endl;
        cout.flush();

        long iters = bench.getIterations();

        Vector<double> x(N), y(N);
        initializeRandomDouble(x.data(), N);
        initializeRandomDouble(y.data(), N);

        bench.start();
        for (long i=0; i < iters; ++i)
        {
            y += a * x;
            y += b * x;
        }
        bench.stop();
    }

    bench.endImplementation();
}

void daxpyArrayVersion(BenchmarkExt<int>& bench, double a)
{
    bench.beginImplementation("Array<T,1>");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "Array<T,1>: N = " << N << endl;
        cout.flush();

        long iters = bench.getIterations();

        Array<double,1> x(N), y(N);
        initializeRandomDouble(x.data(), N);
        initializeRandomDouble(y.data(), N);

        double b = - a;

        bench.start();
        for (long i=0; i < iters; ++i)
        {
            y += a * x;
            y += b * x;
        }
        bench.stop();
    }

    bench.endImplementation();
}

void daxpyF77Version(BenchmarkExt<int>& bench, double a)
{
    bench.beginImplementation("Fortran 77");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "Fortran 77: N = " << N << endl;
        cout.flush();

        int iters = bench.getIterations();

        double* x = new double[N];
        double* y = new double[N];
        initializeRandomDouble(x, N);
        initializeRandomDouble(y, N);

        bench.start();
        fidaxpy(a, x, y, N, iters);
        bench.stop();

        delete [] x;
        delete [] y;
    }

    bench.endImplementation();
}

void daxpyBLASVersion(BenchmarkExt<int>& bench, double a)
{
    bench.beginImplementation("Fortran BLAS");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "Fortran BLAS: N = " << N << endl;
        cout.flush();

        int iters = bench.getIterations();

        double* x = new double[N];
        double* y = new double[N];
        initializeRandomDouble(x, N);
        initializeRandomDouble(y, N);

        int xstride = 1, ystride = 1;
        double b = - a;

        bench.start();
        for (long i=0; i < iters; ++i)
        {
            fdaxpy(N, a, x, xstride, y, ystride);
            fdaxpy(N, b, x, xstride, y, ystride);
        }
        bench.stop();

        delete [] x;
        delete [] y;
    }

    bench.endImplementation();
}

void daxpyF90Version(BenchmarkExt<int>& bench, double a)
{
    bench.beginImplementation("Fortran 90");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "Fortran 90: N = " << N << endl;
        cout.flush();

        int iters = bench.getIterations();

        double* x = new double[N];
        double* y = new double[N];
        initializeRandomDouble(x, N);
        initializeRandomDouble(y, N);

        bench.start();
        f90daxpy(a, x, y, N, iters);
        bench.stop();

        delete [] x;
        delete [] y;
    }

    bench.endImplementation();
}

#ifdef BENCHMARK_VALARRAY
void daxpyValarrayVersion(BenchmarkExt<int>& bench, double a)
{
    bench.beginImplementation("valarray<T>");

    while (!bench.doneImplementationBenchmark())
    {
        int N = bench.getParameter();

        cout << "valarray<T>: N = " << N << endl;
        cout.flush();

        long iters = bench.getIterations();

        valarray<double> x(N), y(N);
        initializeArray(x, N);
        initializeArray(y, N);

        double b = - a;

        bench.start();
        for (long i=0; i < iters; ++i)
        {
            y += a * x;
            y += b * x;
        }
        bench.stop();
    }

    bench.endImplementation();
}
#endif


#include <cv.h>
#include <highgui.h>
#include <cstdlib>
#include <ctime>
#include <iostream>
#include <sstream>
#include <limits>
#include <utility>
#include "otherheaders.h"

/// TODO: @Yupeng added
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <Eigen/Sparse>

#include <MatOp/SparseCholesky.h>
#include <MatOp/SparseGenMatProd.h>
#include <SymGEigsSolver.h>

// using namespace Eigen;
typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::Triplet<double, double> T;

using namespace cv;
using namespace std;

/* sort pair<int,int> based on first value */
typedef std::pair<int, int> mypair;
bool comparatorPair(const mypair& l, const mypair& r) {
    return l.first < r.first;
}

void spectralPb(std::vector<Mat*>& sPb, Mat* mPb, Size orig_size, std::string outFile, int nvec) {
    int ty = mPb->rows;
    int tx = mPb->cols;
    /*
    Mat A1  = mPb->t();
    int ty = A1.rows;
    int tx = A1.cols;
    */

    /*
    cout << "========Printing mPb ============" <<endl;
    double* ptrData;
    for (int i = 0; i < mPb->rows; i++) {
        ptrData = (double*) mPb->ptr(i);
        for (int j = 0; j < mPb->cols; j++) {
            cout << ptrData[j] << ",";
        }
        cout << endl;
    }
    cout << "========Ended Printing mPb ============" <<endl;
    */

    Mat l1 = Mat::zeros(ty + 1, tx, CV_64F);
    Mat l2 = Mat::zeros(ty, tx + 1, CV_64F);

    mPb->copyTo(l1(Rect(0, 1, tx, ty)));
    mPb->copyTo(l2(Rect(1, 0, tx, ty)));

    log_info("Setting the sparse symmetric affinity matrix");
    vector<double> val;
    vector<int> I;      /// FIXME: long int or long long int??
    vector<int> J;      /// FIXME: long int or long long int??
    vector<double> valD;

    time_t now1, now2, now3, now4, now5, now5a, now5b, now6, now7, now8, now9, now10, now11;
    time(&now1);
    buildW(valD, val, I, J, l1, l2);    /// FIXME: check the output order of val I and J
    time(&now2);

    log_info("At the end of buildW: %ld", now2 - now1);

    /// TODO: @Yupeng uncomment @Manu's display codes
    /*
    cout << "=====Display matrix =======" << endl;
    for (int i = 0 ;i < val.size(); i++) {
        cout << "("<<I.at(i)+1<<"," << J.at(i)+1<<") =  " << val.at(i) << endl;
    }
    cout << "=====Display ended =======" << endl;
    */

    int nnz = val.size();
    int nnzD1 = valD.size();

    /// FIXME: @Yupeng added
    log_debug("nnz = %d, nnzD1 = %d", nnz, nnzD1);
    log_debug("ty = %d, tx = %d", ty, tx);

    assert(val.size() == I.size() && I.size() == J.size());

    int wy = (int) *max_element(I.begin(), I.end());
    time(&now3);
    log_info("Time for one max: %ld", now3 - now2);
    int wx = (int) *max_element(J.begin(), J.end());

    // cout << "wy = " << wy << "wx = " << wx << endl;
    int wymin = (int)*min_element(I.begin(), I.end());
    int wxmin = (int)*min_element(J.begin(), J.end());
    // cout << "wymin = " << wymin << "wxmin = " << wxmin << endl;
    int n = (wy == wx) ? wy + 1 : max(wy, wx) + 1;      /*  added one because indices start from zero */

    double valmin = (double)*min_element(val.begin(), val.end());
    double valmax = (double)*max_element(val.begin(), val.end());

    time(&now4);
    log_info("Time for all comparisons: %ld", now4 - now2);
    // cout << "valmax = " << valmax << " valmin = " << valmin << endl;

    vector<int> pcol(n + 1, -1);
    int colInd = 0;
    pcol.at(colInd) = 0;
    for (int i = 1; i < nnz; i++) {
        if (I.at(i) != I.at(i - 1)) pcol.at(++colInd) = i;
    }
    pcol.at(++colInd) = nnz;

    time(&now5a);

    log_debug("Time pcol: %ld", now5a - now4);

    /// FIXME: part below commented because of more efficient version above
    /*
    vector<int> pcol_old(n + 1, -1);
    std::vector<int>::iterator it;
    for (int i = 0; i < n; i++) {
        fs << "orig_size" << orig_size;
        it = find(I.begin(), I.end(), i);
        if (it != I.end()) {
            pcol_old.at(i) = it - I.begin();
        }
    }
    pcol_old.at(n) = nnz;
    */

    time(&now5);
    log_debug("Time pcol_old: %ld", now5 - now5a);

    time(&now5b);
    log_debug("Comparison time: %ld", now5b - now5);

    int nnzD = n;
    std::vector<double> valDmW(nnz);
    for (int i = 0; i < nnz; i++) {
        if (I.at(i) != J.at(i)) {
            valDmW.at(i) = -val.at(i);
        }
        else {
            /// FIXME: if any element changes to zero
            valDmW.at(i) = valD.at(I.at(i)) - val.at(i);
        }
    }

    time(&now6);
    log_info("Time pcol: %ld", now6 - now5);
    std::vector<int> IB;
    IB.reserve(n);
    int cnt = 0;
    std::generate_n(std::back_inserter(IB), n, [cnt]() mutable { return cnt++; });

    time(&now7);
    log_info("Time generate_n: %ld", now7 - now6);
    // cout << IB.at(0) << "," << IB.at(n - 1) << "," << nnzD << endl;

    std::vector<int> JB(IB);
    JB.push_back(nnzD);

    std::vector<double> valB(nnzD);
    for (int i = 0; i < nnzD; i++) {
        valB.at(i) = 0;
        if (pcol.at(i) != -1)
            for (int j = pcol[i]; j < pcol[i + 1]; j++) {
                valB.at(i) += val.at(j);
            }
    }

    /*
    cout << "=====Display matrix =======" << endl;
    for (int i = 0; i < valD.size(); i++) {
        cout << "("<<IB.at(i)+1<<"," << JB.at(i)+1<<") =  " << valD.at(i) << endl;
    }
    cout << "=====Display ended =======" << endl;
    */

    time(&now8);

    /// FIXME: @Yupeng fixed
    log_info("Before Spectra: %ld", now8 - now7);
    log_info("Overall orig reported: %ld", now8 - now2);
    log_info("Setting up matrices D and D-W for Spectra");

    time(&now9);

    /// FIXME: @Yupeng output for debugging
    log_info("OutFile: %s | ty (rows) = %d, tx (cols) = %d", outFile.c_str(), ty, tx);

    /* eigenvector computation using Spectra library begins */
    double* p_eigVals;
    double* p_eigVec;

    int nD = ty * tx;       /* nrows * ncols */
    Eigen::SparseMatrix<double> D2mW2(nD, nD);
    std::vector<T> D2mW2_tripletList;
    D2mW2_tripletList.reserve(nnz);
    for (int i = 0; i < nnz; i++) {
        /* I is rowIdx, J is colIdx, condition "tempCol >= row" for push_back of I, J, val */
        D2mW2_tripletList.push_back(T(J.at(i), I.at(i), valDmW.at(i)));
    }
    D2mW2.setFromTriplets(D2mW2_tripletList.begin(), D2mW2_tripletList.end());
    D2mW2.makeCompressed();

    Eigen::SparseMatrix<double> D2(nD, nD);
    std::vector<T> D2_tripletList;
    D2_tripletList.reserve(nD);
    for (int i = 0; i < nD; i++) {
        D2_tripletList.push_back(T(i, i, valD.at(i)));
    }
    D2.setFromTriplets(D2_tripletList.begin(), D2_tripletList.end());
    D2.makeCompressed();

    Spectra::SparseGenMatProd<double> opA(D2mW2);
    Spectra::SparseCholesky<double> opB(D2);

    /// FIXME: @Yupeng maxLabel is something got after watershed
    // int ncv = (maxLabel < 2*nvec) ? maxLabel : 2*nvec;
    log_debug("Spectra nvec: %d", nvec);
    int ncv = 2 * nvec;     /*  nev < ncv <= n (size of matrix) */

    Spectra::SymGEigsSolver<double, Spectra::SMALLEST_MAGN, Spectra::SparseGenMatProd<double>,
                            Spectra::SparseCholesky<double>, Spectra::GEIGS_CHOLESKY>
    geigs(&opA, &opB, nvec, ncv);

    geigs.init();
    int nconv = geigs.compute();

    Eigen::VectorXd evalues;
    Eigen::MatrixXd evecs;
    if (geigs.info() == Spectra::SUCCESSFUL) {
        evalues = geigs.eigenvalues();      /* nvec */
        evecs = geigs.eigenvectors();       /* (rows, cols) = (nD, nvec) */
    }
    else {
        log_error("Geigs is failing");
    }
    log_debug("Rows: %ld, Cols: %ld", evecs.rows(), evecs.cols());
    stringstream sStream;
    sStream << evalues;
    log_debug("Generalized eigenvalues found: %s", sStream.str().c_str());
    sStream.str("");
    sStream << evecs.topRows(10);
    log_debug("Generalized eigenvectors found: %s", sStream.str().c_str());

    log_debug("n: %d", n);
    log_debug("Rows: %ld, Cols: %ld", evecs.rows(), evecs.cols());
    log_debug("VectorXd size: %ld", evalues.size());

    /* data interface conversion */
    p_eigVals = new double[nvec];
    for (int i = 0; i < nvec; i++) p_eigVals[i] = evalues(i);
    p_eigVec = new double[n * nvec];    /// TODO: not sure about the conflict between n and nD
    for (int i = 0; i < nvec; i++) {
        for (int j = 0; j < nD; j++) {
            p_eigVec[i * nD + j] = evecs(j, i);
        }
    }

    /* eigenvector computation using Spectra library ends */

    std::vector<Mat*> vect(nvec);
    vect.at(0) = new Mat(ty, tx, CV_64FC1, Scalar::all(0));

    double* data;
    double minVal, maxVal, minXi, maxXi;
    Point minLoc, maxLoc;
    std::string imageName;
    double alpha, beta;
    Mat destIm(ty, tx, CV_64FC1);
    for (int i = 1; i < nvec; i++) {    /* excluding i = 0 */
        vect.at(i) = new Mat(ty, tx, CV_64FC1);
        data = (double*)vect.at(i)->data;
        for (int j = 0; j < n; j++) {
            data[j] = p_eigVec[i * n + j];
        }
        minMaxLoc(*vect.at(i), &minVal, &maxVal, &minLoc, &maxLoc);
        *(vect.at(i)) -= minVal;
        *(vect.at(i)) /= maxVal - minVal;   /// FIXME: check order of precedence
        minMaxLoc(*vect.at(i), &minVal, &maxVal, &minLoc, &maxLoc);
    }

    /* OE parameters */
    int hil = 0;
    int deriv = 1;
    int support = 3;
    double sigma = 1.0;
    int nOrient = 8;
    double dtheta = PI / nOrient;
    vector<int> ch_per = {3, 2, 1, 0, 7, 6, 5, 4};
    double theta = 0.0;
    sPb.resize(nOrient);
    /* initilaize sPb */
    for (int i = 0; i < nOrient; i++) {
        sPb.at(i) = new Mat(ty, tx, CV_64FC1, Scalar::all(0));
    }

    log_info("Filtering the sPb values");

    for (int i = 1; i < nvec; i++) {
        if (p_eigVals[i] >
            /// TODO: check what happens for zero of smallest eigenvalue
            std::numeric_limits<double>::epsilon()) {
            log_debug("Eigenvalue %d: %f", i, p_eigVals[i]);
            *(vect.at(i)) /= sqrt(p_eigVals[i]);
            for (int o = 0; o < nOrient; o++) {
                double theta = dtheta * static_cast<double>(o);
                /// FIXME: hil not handled and also not required
                Mat f = oeFilter(sigma, support, theta, deriv, hil);
                Mat destIm;
                filter2D(*vect.at(i), destIm, -1, f, Point(-1, -1), 0, BORDER_REFLECT_101);
                *(sPb.at(ch_per.at(o))) += abs(destIm);
            }
        }
    }

    /// FIXME: @Yupeng comment
    /*
    double alpha, beta;
    for (int i = 0; i < nOrient; i++) {
        // cout << "here_seg" << endl;
        imageName = outFile + "_sPb_i_" + to_string(i) + ".png";
        minMaxLoc(*sPb.at(i), &minXi, &maxXi);
        alpha = 255.0 / (maxXi - minXi);
        beta = -minXi * 255.0 / (maxXi - minXi);
        Mat destIm;
        sPb.at(i)->convertTo(destIm, CV_8U, alpha, beta);
        imwrite(imageName, destIm);
    }
    cout << "here seg 1" << endl;
    */

    /// FIXME: @Yupeng potential memory leak problems
    delete[] p_eigVals;
    delete[] p_eigVec;
}


#ifndef UTIL_HPP
#define UTIL_HPP

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

#include <time.h>
#include <math.h>
#include <ctype.h>
#include <stdlib.h>
#include <fstream>

// Google Log hack
#define LR LOG(INFO)<<"(rank:"<<ctx_->rank<<") "

// Container hack
#define RANGE(x) ((x).begin()), ((x).end())
#define SUM(x)   (std::accumulate(RANGE(x), .0))

// Random hack
#include <chrono>
#include <random>
#define CLOCK (std::chrono::system_clock::now().time_since_epoch().count())
static std::mt19937 _rng(CLOCK);
static std::uniform_real_distribution<double> _unif01;
static std::normal_distribution<double> _stdnormal;

struct ThreadRNG {
  ThreadRNG() : rng_(CLOCK) {}
  std::mt19937 rng_;
  std::uniform_real_distribution<double> unif01_;
  std::normal_distribution<double> stdnormal_;
};

// Eigen
#define EIGEN_INITIALIZE_MATRICES_BY_ZERO
#define EIGEN_DEFAULT_IO_FORMAT \
        Eigen::IOFormat(StreamPrecision,1," "," ","","","[","]")
#include <Eigen/Dense>
using EMatrix = Eigen::MatrixXd;
using EVector = Eigen::VectorXd;
using EArray  = Eigen::ArrayXd;
using EMAtrix = Eigen::ArrayXXd;
using EMatrixMap = Eigen::Map<EMatrix>;
using EVectorMap = Eigen::Map<EVector>;
using EArrayMap  = Eigen::Map<EArray>;

template<typename T>
using TMatrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

template<typename T>
static void save_matrix(const TMatrix<T>& mat, std::string file) {
  std::ofstream fout(file); CHECK(fout.good());
  Eigen::IOFormat space_sep(Eigen::StreamPrecision,0," ","\n","","","","");
  fout << mat.format(space_sep);
  fout.close();
}

template<typename T>
static void load_matrix(TMatrix<T> *mat, std::string file) {
  std::ifstream fin(file); CHECK(fin.good());
  int nrows = 0;
  std::string line;
  while (std::getline(fin, line)) ++nrows;
  fin.clear(); fin.seekg(0); // rewind
  T val;
  std::vector<T> flat; // row major
  while (not fin.eof()) {
    fin >> val;
    flat.push_back(val);
  }
  fin.close();
  mat->resize(flat.size() / nrows, nrows);
  memcpy(mat->data(), flat.data(), sizeof(double) * flat.size());
  mat->transposeInPlace();
}

template<typename T>
static void rowwise_max_ind(const TMatrix<T>& mat, TMatrix<int> *ind) {
  int nrows = mat.rows();
  ind->resize(nrows, 1);
  for (int i = 0; i < nrows; ++i) {
    mat.row(i).maxCoeff(ind->data() + i);
  }
}

// Multivariate Normal with mean = covariance * v; covariance = inv(precision)
static EVector draw_mvgaussian(const EMatrix& precision, const EVector& v) {
  Eigen::LLT<EMatrix> chol;
  chol.compute(precision);
  EVector alpha = chol.matrixL().solve(v); // alpha = L' * mu
  EVector mu = chol.matrixU().solve(alpha);
  EVector z(v.size());
  for (int i = 0; i < v.size(); ++i)
    z(i) = _stdnormal(_rng);
  EVector x = chol.matrixU().solve(z);
  return mu + x;
}

// Get monotonic time in seconds from a starting point
static double get_time() {
  struct timespec start;
  clock_gettime(CLOCK_MONOTONIC, &start);
  return (start.tv_sec + start.tv_nsec/1000000000.0);
}

class Timer {
public:
  void   tic() { start_ = get_time(); }
  double toc() { double ret = get_time() - start_; time_ += ret; return ret; }
  double get() { return time_; }
private:
  double time_  = .0;
  double start_ = get_time();
};

// Google flags hack
static void print_help() {
  fprintf(stderr, "Program Flags:\n");
  std::vector<google::CommandLineFlagInfo> all_flags;
  google::GetAllFlags(&all_flags);
  for (const auto& flag : all_flags) {
    if (flag.filename.find("src/") != 0) // HACK: filter out built-in flags
      fprintf(stderr,
              "-%s: %s (%s, default:%s)\n",
              flag.name.c_str(),
              flag.description.c_str(),
              flag.type.c_str(),
              flag.default_value.c_str());
  }
  exit(1);
}

// Google flags hack
static void print_flags() {
  LOG(INFO) << "---------------------------------------------------------------------";
  std::vector<google::CommandLineFlagInfo> all_flags;
  google::GetAllFlags(&all_flags);
  for (const auto& flag : all_flags) {
    if (flag.filename.find("src/") != 0) // HACK: filter out built-in flags
      LOG(INFO) << flag.name << ": " << flag.current_value;
  }
  LOG(INFO) << "---------------------------------------------------------------------";
}

// Faster strtol without error checking.
static long int strtol(const char *nptr, char **endptr) {
  // Skip spaces
  while (isspace(*nptr)) ++nptr;
  // Sign
  bool is_negative = false;
  if (*nptr == '-') {
    is_negative = true;
    ++nptr;
  } else if (*nptr == '+') {
    ++nptr;
  }
  // Go!
  long int res = 0;
  while (isdigit(*nptr)) {
    res = (res * 10) + (*nptr - '0');
    ++nptr;
  }
  if (endptr != NULL) *endptr = (char *)nptr;
  if (is_negative) return -res;
  return res;
}

#endif


/**
 * vector_mutator_test.cc
 * Copyright 2016 John Lawson
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "vector_mutator.h"

#include "gtest/gtest.h"
#include "qv/quiver_matrix.h"

#include <armadillo>

namespace refl {
namespace {
using Matrix = arma::Mat<int>;
}
TEST(VecMut, A2) {
	Matrix a = { { 2, 1 }, { 1, 2 } };
	cluster::QuiverMatrix q("{ { 0 1 } { -1 0 } }");
	VectorMutator vm(q, a);
	arma::mat vecs = { { 0, 1 }, { 1, 0 } };
	arma::mat res(2, 2);
	vm.mutate(vecs, 0, res);
	arma::mat exp = { { 0, 1 }, { 1, 0 } };
	EXPECT_TRUE(arma::all(arma::all(exp == res)));

	arma::mat res2(2, 2);
	vm.mutate(vecs, 1, res2);
	arma::mat exp2 = { { -1, 1 }, { 1, 0 } };
	EXPECT_TRUE(arma::all(arma::all(exp2 == res2)));
}
TEST(VecMut, A3) {
	Matrix a = { { 2, 1, -1 }, { 1, 2, 1 }, { -1, 1, 2 } };
	cluster::QuiverMatrix q("{ { 0 1 0 } { -1 0 1 } { 0 -1 0 } }");
	VectorMutator vm(q, a);
	arma::mat vecs = { { 1, 0, 0 }, { 0, 1, 0 }, { 0, 0, 1 } };
	arma::mat res(3, 3);
	vm.mutate(vecs, 0, res);
	arma::mat exp = vecs;
	EXPECT_TRUE(arma::all(arma::all(exp == res)));

	arma::mat res2(3, 3);
	vm.mutate(vecs, 1, res2);
	arma::mat exp2 = { { 1, 0, 0 }, { -1, 1, 0 }, { 0, 0, 1 } };
	EXPECT_TRUE(arma::all(arma::all(exp2 == res2)));

	arma::mat res3(3, 3);
	vm.mutate(vecs, 2, res3);
	arma::mat exp3 = { { 1, 0, 0 }, { 0, 1, 0 }, { 0, -1, 1 } };
	EXPECT_TRUE(arma::all(arma::all(exp2 == res2)));
}
TEST(VecMut, A3signs) {
	Matrix a = { { 2, -1, -1 }, { -1, 2, -1 }, { -1, -1, 2 } };
	cluster::QuiverMatrix q("{ { 0 1 0 } { -1 0 1 } { 0 -1 0 } }");
	VectorMutator vm(q, a);
	arma::mat vecs = { { 1, 0, 0 }, { 0, 1, 0 }, { 0, 0, 1 } };
	arma::mat res(3, 3);
	vm.mutate(vecs, 0, res);
	arma::mat exp = vecs;
	EXPECT_TRUE(arma::all(arma::all(exp == res)));

	arma::mat res2(3, 3);
	vm.mutate(vecs, 1, res2);
	arma::mat exp2 = { { 1, 0, 0 }, { 1, 1, 0 }, { 0, 0, 1 } };
	EXPECT_TRUE(arma::all(arma::all(exp2 == res2)));

	arma::mat res3(3, 3);
	vm.mutate(vecs, 2, res3);
	arma::mat exp3 = { { 1, 0, 0 }, { 0, 1, 0 }, { 0, 1, 1 } };
	EXPECT_TRUE(arma::all(arma::all(exp2 == res2)));
}
}


//
// Created by kellerberrin on 19/4/21.
//

#include "kol_OntologyTypes.h"
#include "kol_InformationAncestorMean.h"

#include "kol_SetUtilities.h"
#include "kol_Accumulators.h"

#include <boost/accumulators/statistics/max.hpp>

namespace kol = kellerberrin::ontology;

//! A method for calculating the shared infromation between two concepts.
/*!
  This method returns the shared information between two concepts.
*/
double kol::InformationAncestorMean::sharedInformation(const std::string &termA, const std::string &termB) const  {
  // return 0 for any terms not in the datbase
  if (not ic_map_ptr_->validateTerms(termA, termB)) {

    return 0.0;

  }

  Accumulators::MeanAccumulator meanIC;

  OntologySetType<std::string> ancestorsA = graph_ptr_->getSelfAncestorTerms(termA);
  OntologySetType<std::string> ancestorsB = graph_ptr_->getSelfAncestorTerms(termB);

  OntologySetType<std::string> sharedAncestors = SetUtilities::setIntersection(ancestorsA, ancestorsB);

  for (auto const &term : sharedAncestors) {

    meanIC(ic_map_ptr_->termInformation(term));

  }

  return Accumulators::extractMean(meanIC);

}


/*******************************************************************************
 * Copyright (c) 2018-, UT-Battelle, LLC.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the MIT License
 * which accompanies this distribution.
 *
 * Contributors:
 *   Alexander J. McCaskey - initial API and implementation
 *   Thien Nguyen - implementation
 *******************************************************************************/
#include "qsim_utils.hpp"
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <Eigen/QR>
#include <cassert>
namespace qcor {
namespace QuaSiMo {
std::shared_ptr<CostFunctionEvaluator>
getEvaluator(Operator *observable, const HeterogeneousMap &params) {
  // If an evaluator was provided explicitly:
  if (params.pointerLikeExists<CostFunctionEvaluator>("evaluator")) {
    return xacc::as_shared_ptr(
        params.getPointerLike<CostFunctionEvaluator>("evaluator"));
  }

  // Cost Evaluator was provided by name:
  if (params.stringExists("evaluator")) {
    return getObjEvaluator(observable, params.getString("evaluator"));
  }

  // No specific evaluator/evaluation method was requested,
  // use the default one (partial tomography based).
  return getObjEvaluator(observable);
}

PronyResult pronyFit(const std::vector<std::complex<double>> &in_signal) {
  assert(!in_signal.empty());

  // Returns the Hankel matrix constructed from the first column c, and
  // (optionally) the last row r.
  auto hankelMat =
      [](const std::vector<std::complex<double>> &c,
         const std::vector<std::complex<double>> &r) -> Eigen::MatrixXcd {
    Eigen::MatrixXcd result = Eigen::MatrixXcd::Zero(c.size(), r.size());
    const auto m = c.size();
    for (int i = 0; i < result.rows(); ++i) {
      for (int j = 0; j < result.cols(); ++j) {
        // H(i,j) = c(i+j),  i+j < m;
        if (i + j < m) {
          result(i, j) = c[i + j];
        } else {
          // H(i,j) = r(i+j+1-m),  otherwise
          result(i, j) = r[i + j + 1 - m];
        }
      }
    }

    return result;
  };

  // Pythonic vector slice:
  auto vectorSlice = [](const std::vector<std::complex<double>> &in_vec,
                        int in_start, int in_end) {
    std::vector<std::complex<double>> result;
    auto startIter = in_vec.begin() + in_start;
    auto endIter =
        in_end > 0 ? in_vec.begin() + in_end : (in_vec.end() + in_end);
    result.assign(startIter, endIter);
    return result;
  };

  const size_t num_freqs = in_signal.size() / 2;
  auto hankel0 = hankelMat(vectorSlice(in_signal, 0, num_freqs),
                           vectorSlice(in_signal, num_freqs - 1, -1));
  auto hankel1 = hankelMat(vectorSlice(in_signal, 1, num_freqs + 1),
                           vectorSlice(in_signal, num_freqs, 0));
  // std::cout << "Hankel Matrix 0: \n" << hankel0 << "\n";
  // std::cout << "Hankel Matrix 1: \n" << hankel1 << "\n";

  hankel0.transposeInPlace();
  hankel1.transposeInPlace();
  Eigen::MatrixXcd shift_matrix =
      Eigen::MatrixXcd::Zero(hankel0.cols(), hankel1.cols());

  for (size_t i = 0; i < hankel1.cols(); ++i) {
    Eigen::VectorXcd shift_matrix_col =
        hankel0.fullPivHouseholderQr().solve(hankel1.col(i));
    shift_matrix.col(i) = shift_matrix_col;
  }
  // std::cout << "Shift matrix: \n" << shift_matrix << "\n";

  shift_matrix.transposeInPlace();
  Eigen::ComplexEigenSolver<Eigen::MatrixXcd> s(shift_matrix);
  auto phases = s.eigenvalues();
  Eigen::MatrixXcd generation_matrix =
      Eigen::MatrixXcd::Zero(in_signal.size(), phases.size());
  for (int i = 0; i < generation_matrix.rows(); ++i) {
    for (int j = 0; j < generation_matrix.cols(); ++j) {
      generation_matrix(i, j) = std::pow(phases[j], i);
    }
  }

  auto signalData = in_signal;
  Eigen::VectorXcd signal =
      Eigen::Map<Eigen::VectorXcd>(signalData.data(), signalData.size());
  Eigen::VectorXcd amplitudes =
      generation_matrix.fullPivHouseholderQr().solve(signal);
  assert(phases.size() == amplitudes.size());
  // std::cout << "Amplitude:\n" << amplitudes << "\n";
  // std::cout << "Phases:\n" << phases << "\n";
  PronyResult finalResult;
  for (size_t i = 0; i < phases.size(); ++i) {
    finalResult.emplace_back(std::make_pair(amplitudes(i), phases(i)));
  }

  // Sort by amplitude:
  std::sort(finalResult.begin(), finalResult.end(),
            [](const auto &a, const auto &b) {
              return std::abs(a.first) < std::abs(b.first);
            });

  return finalResult;
}
} // namespace QuaSiMo
} // namespace qcor


#include "TriangleMuscleConstraint.h"
#include <Eigen/SVD>
#include <Eigen/Sparse>
#include <Eigen/Geometry>
#include <iostream>

using namespace FEM;
TriangleMuscleConstraint::
TriangleMuscleConstraint(double stiffness,const Eigen::Vector2d& fiber_direction,int i0,int i1,int i2,double area,const Eigen::Matrix2d& invDm)
	:Constraint(stiffness),mFiberDirection(fiber_direction),
	mi0(i0),mi1(i1),mi2(i2),mArea(area),mInvDm(invDm),mActivationLevel(0.0)
{

}
int
TriangleMuscleConstraint::
GetDof()
{
	return 1;
}

void
TriangleMuscleConstraint::
EvaluateJMatrix(int index, std::vector<Eigen::Triplet<double>>& J_triplets)
{
	Eigen::MatrixXd Ai(3,9);

	Eigen::Vector2d v = mInvDm*mFiberDirection;

	Ai<<
		-v[0]-v[1],0,0,v[0],0,0,v[1],0,0,
		0,-v[0]-v[1],0,0,v[0],0,0,v[1],0,
		0,0,-v[0]-v[1],0,0,v[0],0,0,v[1];

	Eigen::MatrixXd MuAiT = mStiffness*mArea*Ai.transpose();

	int idx[3] = {mi0,mi1,mi2};

	for(int i=0;i<3;i++)
	{
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*(index)+0, MuAiT(3*i+0,0)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*(index)+1, MuAiT(3*i+0,1)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*(index)+2, MuAiT(3*i+0,2)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*(index)+0, MuAiT(3*i+1,0)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*(index)+1, MuAiT(3*i+1,1)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*(index)+2, MuAiT(3*i+1,2)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*(index)+0, MuAiT(3*i+2,0)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*(index)+1, MuAiT(3*i+2,1)));
		J_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*(index)+2, MuAiT(3*i+2,2)));
	}
}
void
TriangleMuscleConstraint::
EvaluateLMatrix(std::vector<Eigen::Triplet<double>>& L_triplets)
{
	Eigen::MatrixXd Ai(3,9);

	Eigen::Vector2d v = mInvDm*mFiberDirection;
	Ai<<
		-v[0]-v[1],0,0,v[0],0,0,v[1],0,0,
		0,-v[0]-v[1],0,0,v[0],0,0,v[1],0,
		0,0,-v[0]-v[1],0,0,v[0],0,0,v[1];

	Eigen::MatrixXd MuAiTAi = mStiffness*mArea*((Ai.transpose())*Ai);

	int idx[3] = {mi0,mi1,mi2};
	for(int i =0;i<3;i++)
	{
		for(int j=0;j<3;j++)
		{
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*idx[j]+0, MuAiTAi(3*i+0, 3*j+0)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*idx[j]+1, MuAiTAi(3*i+0, 3*j+1)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+0, 3*idx[j]+2, MuAiTAi(3*i+0, 3*j+2)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*idx[j]+0, MuAiTAi(3*i+1, 3*j+0)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*idx[j]+1, MuAiTAi(3*i+1, 3*j+1)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+1, 3*idx[j]+2, MuAiTAi(3*i+1, 3*j+2)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*idx[j]+0, MuAiTAi(3*i+2, 3*j+0)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*idx[j]+1, MuAiTAi(3*i+2, 3*j+1)));
			L_triplets.push_back(Eigen::Triplet<double>(3*idx[i]+2, 3*idx[j]+2, MuAiTAi(3*i+2, 3*j+2)));
		}
	}
}
void
TriangleMuscleConstraint::
EvaluateDVector(const Eigen::VectorXd& x)
{
	Eigen::Vector3d x0(x.segment<3>(mi0*3));

	Eigen::Matrix32d Ds, P;
	Ds.col(0) = x.segment<3>(mi1*3) - x0;
	Ds.col(1) = x.segment<3>(mi2*3) - x0;

	P.col(0) = Ds.col(0).normalized();
	P.col(1) = (Ds.col(1)-Ds.col(1).dot(P.col(0))*P.col(0)).normalized();

	Eigen::Matrix2d F = P.transpose()*Ds*mInvDm;

	md = (1.0-mActivationLevel)*P*F*mFiberDirection;
}
void
TriangleMuscleConstraint::
GetDVector(int& index,Eigen::VectorXd& d)
{
	d.segment<3>(3*(index)) = md;
	index++;
}


//==================================================================================================
/*!
  @file

    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_TANH_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_TANH_HPP_INCLUDED

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

  /*!

    @ingroup group-hyperbolic
    This function object returns the hyperbolic tangent: \f$\sinh(x)/\cosh(x)\f$.

    @see sinh, cosh, sech, csch, sinhcosh

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

    @par Example:

      @snippet tanh.cpp tanh

    @par Possible output:

      @snippet tanh.txt tanh
  **/
  IEEEValue tanh(IEEEValue const& x);
} }
#endif

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

#endif


#include <Kokkos_Core.hpp>
#include <iostream>

#include <boost/program_options.hpp>

#define PI 3.1415926535897932384

typedef int int_t;

//#include "Tile3D.h"
KOKKOS_FUNCTION double left_flux(double left, double self, double cfl)
{
    return -cfl * left + cfl * self;
}

KOKKOS_FUNCTION double right_flux(double self, double right, double cfl)
{
    return -cfl * right + cfl * self;
}

KOKKOS_FUNCTION double stencil6p(double self, double xm, double xp, double ym,
    double yp, double zm, double zp, double cfl)
{
    return (1.0 - 6.0 * cfl) * self + cfl * (xm + xp + ym + yp + zm + zp);
}

int main(int argc, char* argv[])
{
    double cfl = 0.1;    // must be less than 1/6

    namespace bpo = boost::program_options;
    bpo::options_description desc("ABFT3D");

    desc.add_options()(
        "xsize", bpo::value<int>()->default_value(98), "X Dimension");
    desc.add_options()(
        "ysize", bpo::value<int>()->default_value(98), "Y Dimension");
    desc.add_options()(
        "zsize", bpo::value<int>()->default_value(98), "Z Dimension");

    bpo::variables_map vm;

    // Setup commandline arguments
    bpo::store(bpo::parse_command_line(argc, argv, desc), vm);
    bpo::notify(vm);

    int xsize = vm["xsize"].as<int>();
    int ysize = vm["ysize"].as<int>();
    int zsize = vm["zsize"].as<int>();

    bool do_check_sum = true;

    Kokkos::initialize(argc, argv);
    {
        using range_policy = Kokkos::RangePolicy<>;

        Kokkos::View<double***> stencil_0(
            "data1", xsize + 2, ysize + 2, zsize + 2);
        Kokkos::View<double***> stencil_1(
            "data2", xsize + 2, ysize + 2, zsize + 2);

        // Soft copy
        auto stencil_old = stencil_0;
        auto stencil_new = stencil_1;
        auto stencil_tmp = stencil_new;
        Kokkos::View<double*, Kokkos::MemoryTraits<Kokkos::Atomic>> checksum(
            "checksum", 1);

        Kokkos::View<double*, Kokkos::DefaultHostExecutionSpace> hchecksum(
            "h_checksum", 1);
        hchecksum[0] = 0.;

        Kokkos::deep_copy(checksum, hchecksum);

        // Initialize stencil
        Kokkos::parallel_for(
            "init", range_policy(0, ysize + 2), KOKKOS_LAMBDA(int j) {
                for (int k = 0; k < zsize + 2; ++k)
                {
                    for (int i = 0; i < xsize + 2; ++i)
                    {
                        stencil_old(i, j, k) = std::sin(1.0 * PI * (double) i /
                            (double) (xsize + 1));    //*
                        stencil_new(i, j, k) = stencil_old(i, j, k);
                    }
                }
            });
        Kokkos::fence();

        Kokkos::Timer timer;

        //
        // How to catch the error ... order matters.  where does the error come from?
        // Physical location could matter (plus interaction with OS)
        //

        // Original checksum
        if (do_check_sum)
        {
            Kokkos::parallel_for(
                "init_checksum", range_policy(1, xsize + 1),
                KOKKOS_LAMBDA(int i) {
                    for (int j = 1; j <= ysize; ++j)
                    {
                        for (int k = 1; k <= zsize; ++k)
                        {
                            checksum[0] += stencil_old(i, j, k);
                        }
                    }
                });
            Kokkos::fence();
        }

        Kokkos::deep_copy(hchecksum, checksum);

        std::cout << "Initial Checksum " << std::scientific << hchecksum[0]
                  << std::endl;

        double chk_ = 0;
        for (int its = 0; its < 10; ++its)
        {
            if (do_check_sum)
            {
                // Deduct the checksum from the boundary
                Kokkos::parallel_for(
                    "first_loop", range_policy(1, zsize + 1),
                    KOKKOS_LAMBDA(int k) {
                        for (int j = 1; j <= ysize; ++j)
                        {
                            checksum[0] -= left_flux(stencil_old(0, j, k),
                                               stencil_old(1, j, k), cfl) +
                                right_flux(stencil_old(xsize, j, k),
                                    stencil_old(xsize + 1, j, k), cfl);
                        }
                    });
                Kokkos::fence();

                Kokkos::parallel_for(
                    "second_loop", range_policy(1, zsize + 1),
                    KOKKOS_LAMBDA(int k) {
                        for (int i = 1; i <= xsize; ++i)
                        {
                            checksum[0] -= left_flux(stencil_old(i, 0, k),
                                               stencil_old(i, 1, k), cfl) +
                                right_flux(stencil_old(i, ysize, k),
                                    stencil_old(i, ysize + 1, k), cfl);
                        }
                    });
                Kokkos::fence();

                Kokkos::parallel_for(
                    "third_loop", range_policy(1, ysize + 1),
                    KOKKOS_LAMBDA(int j) {
                        for (int i = 1; i <= xsize; ++i)
                        {
                            checksum[0] -= left_flux(stencil_old(i, j, 0),
                                               stencil_old(i, j, 1), cfl) +
                                right_flux(stencil_old(i, j, zsize),
                                    stencil_old(i, j, zsize + 1), cfl);
                        }
                    });
                Kokkos::fence();
            }

            // Apply stencil
            // Replace the loops by parallel for
            //
            Kokkos::parallel_reduce(
                "stencil_op", range_policy(1, xsize + 1),
                KOKKOS_LAMBDA(int i, double& chk_) {
                    for (int j = 1; j <= ysize; ++j)
                    {
                        for (int k = 1; k <= zsize; ++k)
                        {
                            //std::cout << "Old " << i << " " << j << " " << k << ": " <<  stencil_old(i,j,k) << std::endl;
                            stencil_new(i, j, k) = stencil6p(
                                stencil_old(i, j, k), stencil_old(i - 1, j, k),
                                stencil_old(i + 1, j, k),
                                stencil_old(i, j - 1, k),
                                stencil_old(i, j + 1, k),
                                stencil_old(i, j, k - 1),
                                stencil_old(i, j, k + 1), cfl);

                            chk_ += stencil_new(i, j, k);
                        }
                    }
                },
                chk_);
            Kokkos::fence();

            // chk_ = 0;

            if (do_check_sum)
            {
                Kokkos::deep_copy(hchecksum, checksum);

                // Testing the checksum
                auto error = std::abs((chk_ - hchecksum[0]));

                std::cout << "New Checksum " << std::scientific << chk_
                          << " and analytical checksum " << hchecksum[0]
                          << " with Error: " << error << std::endl;

                if (error > 1e-2)
                {
                    std::cout << "Failed " << error << std::endl;
                }

                // Alternate stencil assignment
                if (its % 2 == 0)
                {
                    stencil_old = stencil_1;
                    stencil_new = stencil_0;
                }
                else
                {
                    stencil_old = stencil_0;
                    stencil_new = stencil_1;
                }

                chk_ = 0.;
            }
        }

        double time_taken = timer.seconds();

        std::cout << "Executed in: " << time_taken << std::endl;
    }

    Kokkos::finalize();
    return 0;
}


/**
 * \file libsanm/tensor_polymat.cpp
 * This file is part of SANM, a symbolic asymptotic numerical solver.
 */

#include "libsanm/stl.h"
#include "libsanm/tensor_impl_helper.h"

#include <Eigen/Dense>
#include <complex>
#include <functional>
#include <mutex>
#include <span>

using namespace sanm;

namespace {

using cfp_t = std::complex<fp_t>;

/*!
 * \brief compute the DFT of given coeffs using the FFT algorithm
 *
 * Evaluate the polynomial defined by the \p coeffs at the points \f$\omega_n^0,
 * \ldots, \omega_n^{n-1}\f$ where \f$n\f$ is \p nr_term and must be a power of
 * two.
 *
 * \return the real part and the imaginary part
 */
std::pair<TensorArray, TensorArray> fft(const TensorArray& coeffs,
                                        size_t nr_term) {
    sanm_assert(!(nr_term & (nr_term - 1)));
    if (coeffs.size() == 1) {
        return {TensorArray(nr_term, coeffs[0]),
                TensorArray(nr_term, coeffs[0].fill_with(0))};
    }
    if (nr_term == 1) {
        TensorND sum = coeffs[0];
        for (size_t i = 1; i < coeffs.size(); ++i) {
            sum += coeffs[i];
        }
        return {{sum}, {sum.fill_with(0)}};
    }

    TensorArray coeffs_even, coeffs_odd;
    coeffs_even.reserve((coeffs.size() + 1) / 2);
    coeffs_odd.reserve(coeffs.size() / 2);
    for (size_t i = 0; i < coeffs.size(); i += 2) {
        coeffs_even.emplace_back(coeffs[i]);
        if (i + 1 < coeffs.size()) {
            coeffs_odd.emplace_back(coeffs[i + 1]);
        }
    }

    auto result_even = fft(coeffs_even, nr_term / 2),
         result_odd = fft(coeffs_odd, nr_term / 2);

    TensorArray ret_real(nr_term), ret_imag(nr_term);
    for (size_t si = 0; si < nr_term / 2; ++si) {
        const TensorND& a_real = result_even.first[si];
        const TensorND& a_imag = result_even.second[si];
        const TensorND& b_real = result_odd.first[si];
        const TensorND& b_imag = result_odd.second[si];

        for (size_t i : {si, si + nr_term / 2}) {
            TensorND& y_real = ret_real[i];
            TensorND& y_imag = ret_imag[i];

            fp_t angle = fp_t(i) * (M_PI * 2) / fp_t(nr_term),
                 x_real = std::cos(angle), x_imag = std::sin(angle);

            // y = a + x * b
            y_real.set_shape(coeffs[0].shape());
            y_imag.set_shape(coeffs[0].shape());
            as_vector_w(y_real) = as_vector_r(a_real) +
                                  as_vector_r(b_real) * x_real -
                                  as_vector_r(b_imag) * x_imag;
            as_vector_w(y_imag) = as_vector_r(a_imag) +
                                  as_vector_r(b_imag) * x_real +
                                  as_vector_r(b_real) * x_imag;
        }
    }
    return {std::move(ret_real), std::move(ret_imag)};
}

size_t next_pow2(size_t x) {
    size_t y = 1;
    while (y < x) {
        y <<= 1;
    }
    return y;
}

TensorND compute_polymat_det_coeff_with_fft(const TensorArray& coeffs,
                                            size_t nr_term_pow2,
                                            size_t target_order) {
    std::pair<TensorArray, TensorArray> polymat_dft;
    {
        SANM_SCOPED_PROFILER("polymat_det_fft");
        polymat_dft = fft(coeffs, nr_term_pow2);
    }

    const size_t batch = coeffs[0].shape(0);
    const Eigen::Index mdim = coeffs[0].shape(1);
    ScopedAllowMalloc allow_malloc;
    Eigen::Matrix<cfp_t, Eigen::Dynamic, Eigen::Dynamic> eigmat(mdim, mdim);

    TensorND ret{TensorShape{batch, 1}};
    auto ret_ptr = ret.woptr();
    for (size_t ib = 0; ib < batch; ++ib) {
        // do the inverse dft to solve the target coefficient
        cfp_t accum = 0;
        for (size_t i = 0; i < nr_term_pow2; ++i) {
            EigenMatDyn mreal{const_cast<fp_t*>(polymat_dft.first[i].ptr()) +
                                      ib * mdim * mdim,
                              mdim, mdim},
                    mimag{const_cast<fp_t*>(polymat_dft.second[i].ptr()) +
                                  ib * mdim * mdim,
                          mdim, mdim};
            eigmat.real() = mreal;
            eigmat.imag() = mimag;
            cfp_t dfti = eigmat.determinant();
            fp_t angle =
                    -(M_PI * 2) * fp_t(i * target_order) / fp_t(nr_term_pow2);
            accum += dfti * cfp_t{std::cos(angle), std::sin(angle)};
        }
        accum /= fp_t(nr_term_pow2);
        sanm_assert(std::fabs(accum.imag()) <
                            1e-4 * std::max<fp_t>(1, std::fabs(accum.real())),
                    "IDFT not real: real=%g imag=%g", accum.real(),
                    accum.imag());
        ret_ptr[ib] = accum.real();
    }

    return ret;
}

TensorArray transpose_coeffs(const TensorArray& coeffs) {
    TensorArray ret;
    ret.resize(coeffs.size());
    for (size_t i = 0; i < coeffs.size(); ++i) {
        TensorND& dst = ret[i];
        const TensorND& src = coeffs[i];
        sanm_assert(src.rank() == 3);
        size_t n = src.shape(0), m0 = src.shape(1), m1 = src.shape(2);
        dst.set_shape({m0, m1, n});
        EigenMatDyn mdst{dst.woptr(), static_cast<Eigen::Index>(n),
                         static_cast<Eigen::Index>(m0 * m1)},
                msrc{const_cast<fp_t*>(src.ptr()),
                     static_cast<Eigen::Index>(m0 * m1),
                     static_cast<Eigen::Index>(n)};
        mdst = msrc.transpose();
    }
    return ret;
}

using EigenVecArr = std::span<EigenVec>;

void conv(EigenVecArr dst, EigenVecArr x, EigenVecArr y) {
    for (EigenVec& i : dst) {
        i.setZero();
    }
    for (size_t i = 0; i < x.size(); ++i) {
        for (size_t j = 0; j < y.size() && i + j < dst.size(); ++j) {
            dst[i + j].array() += x[i].array() * y[j].array();
        }
    }
}

void conv_k(size_t k, EigenVec& dst, EigenVecArr x, EigenVecArr y) {
    bool first = true;
    for (size_t i = std::max(0, int(k) + 1 - int(y.size()));
         i < x.size() && i <= k; ++i) {
        if (first) {
            dst = x[i].array() * y[k - i].array();
            first = false;
        } else {
            dst.array() += x[i].array() * y[k - i].array();
        }
    }
    if (first) {
        dst.setZero();
    }
}

class EigenVecArrStorage : public ObjArray<EigenVec> {
public:
    explicit EigenVecArrStorage(size_t size)
            : ObjArray<EigenVec>(size, nullptr, 0, 1) {}
};

/*!
 * \brief Compute a single term in the expansion of determinant
 *
 * The result is the product of mat[i][row_indices[i]]
 * Result is negated if row_indices[0]<0 or row_indices[1]<0
 *
 * \param coeffs_trans transposed coefficients in (m, m, batch) shape
 * \param k target term order
 */
class DetSingleTermCompute {
    TensorArray m_coeffs_trans;
    const size_t m_batch;
    const size_t m_k;
    std::unique_ptr<fp_t[]> m_buf_storage;
    TensorND m_ret;
    EigenVecArrStorage m_buf0, m_buf1, m_opr0, m_opr1;

public:
    DetSingleTermCompute(const TensorArray& coeffs, size_t k)
            : m_coeffs_trans{transpose_coeffs(coeffs)},
              m_batch{m_coeffs_trans[0].shape(2)},
              m_k{k},
              m_buf_storage{new fp_t[2 * (k + 1) * m_batch]},
              m_ret{TensorShape{m_batch, 1}},
              m_buf0{k + 1},
              m_buf1{k + 1},
              m_opr0{coeffs.size()},
              m_opr1{coeffs.size()} {
        auto ptr = m_buf_storage.get();
        for (size_t i = 0; i <= k; ++i) {
            reset(m_buf0[i], ptr + (i * 2) * m_batch, m_batch);
            reset(m_buf1[i], ptr + (i * 2 + 1) * m_batch, m_batch);
        }
    }

    TensorND operator()(std::span<const int> row_indices) {
        auto extract = [&row_indices, this, msize = m_coeffs_trans[0].shape(1)](
                               EigenVecArr dst, int r) {
            for (size_t i = 0; i < m_coeffs_trans.size(); ++i) {
                size_t c = std::abs(row_indices[r]);
                reset(dst[i],
                      m_coeffs_trans[i].ptr() + (r * msize + c) * m_batch,
                      m_batch);
            }
        };

        auto conv_ret = [this, &row_indices](EigenVecArr x, EigenVecArr y) {
            auto ret_vec = as_vector_w(m_ret);
            conv_k(m_k, ret_vec, x, y);
            if (row_indices[0] < 0 || row_indices[1] < 0) {
                m_ret.inplace_neg();
            }
            return m_ret;
        };

        extract(m_opr0, 0);
        extract(m_opr1, 1);
        if (row_indices.size() == 2) {
            return conv_ret(m_opr0, m_opr1);
        }

        EigenVecArrStorage *prod = &m_buf0, *prod_next = &m_buf1;
        conv(*prod, m_opr0, m_opr1);

        for (size_t i = 2; i + 1 < row_indices.size(); ++i) {
            extract(m_opr0, i);
            conv(*prod_next, *prod, m_opr0);
            std::swap(prod, prod_next);
        }
        extract(m_opr0, row_indices.size() - 1);
        return conv_ret(*prod, m_opr0);
    }
};

/*!
 * \brief get the terms in the expansion of determinant
 * \return vector of size m! * m; the sign of items i*m are the sign of the
 *      terms
 */
const std::vector<int>& get_det_terms(size_t m) {
    static std::mutex mutex;
    static std::vector<std::vector<int>> results{std::vector<int>{0}};
    static auto compute = [&](int size) {
        sanm_assert(results.size() == static_cast<size_t>(size - 1));
        std::vector<int>& cur = results.emplace_back();
        const auto& prev = results[size - 2];
        for (int i = 0; i < size; ++i) {
            for (size_t j = 0; j < prev.size();) {
                size_t r0 = cur.size();
                bool neg = i % 2;
                cur.push_back(i);
                for (int jdt = 0; jdt < size - 1; ++jdt) {
                    int p = prev[j + jdt];
                    if (p < 0) {
                        neg = !neg;
                        p = -p;
                    }
                    cur.push_back(p + (p >= i));
                }
                if (neg) {
                    if (cur[r0]) {
                        cur[r0] = -cur[r0];
                    } else {
                        cur[r0 + 1] = -cur[r0 + 1];
                    }
                }
                j += size - 1;
            }
        }
#if 0
        printf("det(%d):\n", size);
        for (size_t i = 0; i < cur.size(); ++i) {
            printf("%d", cur[i]);
            if ((i + 1) % size == 0) {
                printf("\n");
            } else {
                printf(" ");
            }
        }
#endif
    };

    sanm_assert(m >= 1);
    std::lock_guard<std::mutex> mutex_lg{mutex};
    if (m - 1 < results.size()) {
        return results[m - 1];
    }
    for (size_t i = results.size() + 1; i <= m; ++i) {
        compute(i);
    }
    return results[m - 1];
}

TensorND compute_polymat_det_coeff_by_expanding(const TensorArray& coeffs,
                                                size_t target_order) {
    DetSingleTermCompute tc{coeffs, target_order};
    size_t m = coeffs[0].shape(1);
    const auto& terms = get_det_terms(m);

    TensorND ret;
    for (size_t i = 0; i < terms.size(); i += m) {
        TensorND cur = tc({terms.data() + i, m});
        if (!i) {
            ret = cur;
        } else {
            ret += cur;
        }
    }
    return ret;
}
}  // anonymous namespace

TensorND sanm::compute_polymat_det_coeff(const TensorArray& coeffs,
                                         size_t order) {
    SANM_SCOPED_PROFILER("polymat_det");
    sanm_assert(!coeffs.empty() && coeffs[0].rank() == 3 &&
                coeffs[0].shape(1) == coeffs[0].shape(2));
    for (size_t i = 1; i < coeffs.size(); ++i) {
        sanm_assert(coeffs[i].shape() == coeffs[0].shape());
    }

    const size_t batch = coeffs[0].shape(0), mdim = coeffs[0].shape(1),
                 nr_term = (coeffs.size() - 1) * mdim + 1;
    sanm_assert(mdim >= 2);

    if (order >= nr_term) {
        return TensorND{TensorShape{batch, 1}}.fill_with_inplace(0);
    }
    if (order == 0) {
        return coeffs[0].batched_determinant();
    }
    if (order == 1) {
        TensorND ret{TensorShape{batch, 1}},
                src = coeffs[0].batched_cofactor() * coeffs[1];
        EigenMatDyn smat{const_cast<fp_t*>(src.ptr()),
                         static_cast<Eigen::Index>(mdim * mdim),
                         static_cast<Eigen::Index>(batch)};
        as_vector_w(ret) = smat.colwise().sum().transpose();
        return ret;
    }

    if (mdim <= 4) {
        return compute_polymat_det_coeff_by_expanding(coeffs, order);
    }

    return compute_polymat_det_coeff_with_fft(coeffs, next_pow2(nr_term),
                                              order);
}


/*=========================================================================

  Program:   Small Body Geophysical Analysis
  Module:    SBGATSphericalHarmo.hpp

  Class derived from VTK's vtkPolyDataAlgorithm by Benjamin Bercovici

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
/**
 @class  SBGATSphericalHarmo
 @author Benjamin Bercovici
 @author Jay McMahon
@date October 2018
 @brief  Computes/evaluates the outer spherical harmonics expansion of the exterior gravity
 field around a constant density polyhedron

 @details  Computes/evaluates the outer spherical harmonics expansion of the exterior gravity
field around a constant density polyhedron. Normalized or non-normalized coefficients can be computed.
The computed coefficients are completely independent of the mass and density of the considered object
as they are only a geometric construct, thanks to the constant-density assumption. This class will always use results expressed in `meters` as their distance unit (e.g accelerations in m/s^2, potentials in m^2/s^2,...) . Unit consistency is enforced through the use of the SetScaleMeters()
and SetScaleKiloMeters() method.

Adapted from the works of Yu Takahashi and Siamak Hesar by Benjamin Bercovici, University of Colorado Boulder
for more details, see
Werner, R. a. (1997).
Spherical harmonic coefficients for the potential of a constant-density polyhedron.
Computers & Geosciences,
23(10),
1071–1077.
https://doi.org/10.1016/S0098-3004(97)00110-6
@copyright MIT License, Benjamin Bercovici and Jay McMahon
*/

#ifndef SBGATSphericalHarmo_h
#define SBGATSphericalHarmo_h

#include <vtkFiltersCoreModule.h> // For export macro
#include <vtkPolyDataAlgorithm.h>
#include <armadillo>

class VTKFILTERSCORE_EXPORT SBGATSphericalHarmo : public vtkPolyDataAlgorithm{
public:
  /**
   * Constructs with initial values of zero.
   */
  static SBGATSphericalHarmo *New();

  vtkTypeMacro(SBGATSphericalHarmo,vtkPolyDataAlgorithm);
  void PrintSelf(std::ostream& os, vtkIndent indent) override;
  void PrintHeader(std::ostream& os, vtkIndent indent) override;
  void PrintTrailer(std::ostream& os, vtkIndent indent) override;

  /**
  Sets degree of spherical harmonics expansion
  @param deg degree of spherical harmonics expansion
  */
  void SetDegree(const unsigned int deg){
    this -> degree = deg;
    this -> degreeSet = true;
  }

  /**
  Sets reference radius in spherical harmonics expansion. Must be consistent
  with the units in which the shape coordinates are expressed
  @param ref_radius reference radius in spherical harmonics expansion
  */
  void SetReferenceRadius(const double ref_radius){
    this -> referenceRadius = ref_radius;
    this -> referenceRadiusSet = true;
  }

  /**
  Sets polyhedron density
  @param density bulk density of polyhedron (kg/m^3)
  */
  void SetDensity(const double density){
    this -> density = density;
    this -> densitySet = true;
  }

  /*
  Will return normalized coefficients next (default is normalized)
  */
  void IsNormalized(){
    this -> normalized = true;
  }

  /*
  Will return non-normalized coefficients next (default is normalized)
  */
  void IsNonNormalized(){
    this -> normalized = false;
  }

  /*
  Return Cnm array of coefficients ordered like in the exemple below where degree = 5

    1     0     0     0    0      0
  C_10  C_11  C_12  C_13  C_14  C_15
  C_20  C_21  C_22  C_23  C_24  C_25
  C_30  C_31  C_32  C_33  C_34  C_35
  C_40  C_41  C_42  C_43  C_44  C_45
  C_50  C_51  C_52  C_53  C_54  C_55

  @return Cnm array of coefficients
  */
  arma::mat GetCnm() {this -> Update(); return this -> Cnm;}

  /*
  Return Cnm array of coefficients ordered like in the exemple below where degree = 5

    1     0     0     0    0      0
  C_10  C_11  C_12  C_13  C_14  C_15
  C_20  C_21  C_22  C_23  C_24  C_25
  C_30  C_31  C_32  C_33  C_34  C_35
  C_40  C_41  C_42  C_43  C_44  C_45
  C_50  C_51  C_52  C_53  C_54  C_55

  @param[out] Cnm array of coefficients
  */
  void GetCnm(arma::mat & C_nm) {this -> Update(); C_nm = this -> Cnm;}

  /*

  Return Snm array of coefficients ordered like in the exemple below where degree = 5

    0     0     0     0    0      0
    0  S_11  S_12  S_13  S_14  S_15
    0  S_21  S_22  S_23  S_24  S_25
    0  S_31  S_32  S_33  S_34  S_35
    0  S_41  S_42  S_43  S_44  S_45
    0  S_51  S_52  S_53  S_54  S_55

  @return Snm array of coefficients
  */
  arma::mat GetSnm() {this -> Update(); return this -> Snm;}

  /*
  Return Snm array of coefficients ordered like in the exemple below where degree = 5

    0     0     0     0    0      0
    0  S_11  S_12  S_13  S_14  S_15
    0  S_21  S_22  S_23  S_24  S_25
    0  S_31  S_32  S_33  S_34  S_35
    0  S_41  S_42  S_43  S_44  S_45
    0  S_51  S_52  S_53  S_54  S_55

  @param[out] Snm array of coefficients
  */
  void GetSnm(arma::mat & S_nm) {this -> Update(); S_nm = this -> Snm;}

  /**
  Return the acceleration due to gravity at the specified point
  @param pos position at which the acceleration must be evaluated (meters)
  @return acceleration (m / s ^ 2)
  */
  arma::vec::fixed<3> GetAcceleration(const arma::vec::fixed<3> & pos);

  /**
  Evaluates the gravity gradient matrix (the partial derivative of the spherical
  harmonics acceleration with respect to the position vector) at the prescribed
  location. Note that this gravity gradient matrix is expressed in the body-fixed frame of the considered object.
  @param[in] pos position at which the gravity gradient matrix must be evaluated, expressed in the same frame/same unit L as
  the shape used to build the spherical harmonics expansion.
  @param[out] dAccdPos container holding the gravity gradient matrix (1 / s ^ 2)
  */
  void GetGravityGradientMatrix(const arma::vec::fixed<3> & pos,
    arma::mat::fixed<3,3> & dAccdPos);

  /**
  Evaluates the partial derivative of the spherical harmonics acceleration
  with respect to the spherical harmonics coefficients. The coefficients and the partials of the acceleration w/r to the coefficients
  are ordered like so:

  Cnm :
(C_00 == 1)  0     0    0     0     0
  C_10     C_11    0    0     0     0
  C_20     C_21  C_22   0     0     0
  C_30     C_31  C_32  C_33   0     0
  C_40     C_41  C_42  C_43  C_44   0
  C_50     C_51  C_52  C_53  C_54  C_55

  Snm :
  0    0     0     0    0      0
  0  S_11    0     0    0      0
  0  S_21  S_22    0    0      0
  0  S_31  S_32  S_33   0      0
  0  S_41  S_42  S_43  S_44    0
  0  S_51  S_52  S_53  S_54  S_55

  so

  partial_C = [ dA/C_00,dA/C_10,dA/C_11,dA/C_20,dA/C_21,dA/C_22,...]
  partial_S = [ dA/S_11,dA/C_21,dA/S_22,dA/S_31,dA/S_32,dA/S_33,...]

  @param[in] pos position at which the partial derivatives must be evaluated, expressed in the same frame/same unit as
  the shape used to build the spherical harmonics expansion.
  @param[out] partial_C container holding the partial derivative of the acceleration
  with respect to the Cnm spherical harmonic coefficients. If the degree/order is n, then there are (n+1) * (n + 2)/2 non-zero Cnm coefficients

  @param[out] partial_S container holding the partial derivative of the acceleration
  with respect to the Snm spherical harmonic coefficients. If the degree/order is n, then there are (n+1) * (n + 2)/2 - n non-zero Snm coefficients

  */
  void GetPartialHarmonics(const arma::vec::fixed<3> & pos,
    arma::mat & partial_C,
    arma::mat & partial_S);
  /**
  Sets the scale factor to 1, indicative that the polydata has its coordinates expressed in meters
  */
  void SetScaleMeters() { this -> scaleFactor = 1; this -> scaleFactorSet = true;}

  /**
  Sets the scale factor to 1000, indicative that the polydata has its coordinates expressed in kilometers
  */
  void SetScaleKiloMeters() { this -> scaleFactor = 1000; this -> scaleFactorSet = true;}

  /**
  Exports the computed spherical harmonics expansion to
  a JSON file. The saved fields are:
  - facets == number of facets
  - vertices == number of vertices
  - totalMass : {value, unit}
  - density : {value, unit}
  - reference_radius : {value, unit}
  - normalized == true if the coefficients are normalized
  - degree == degree of the spherical expansion
  - Cnm_coefs - vector of coefficients triplets {n,m,Cnm}
  - Snm_coefs - vector of coefficients triplets {n,m,Snm}
  @param path JSON file where the spherical harmonics model will be saved

  */
  void SaveToJson(std::string path) const;

  /**
  Loads a previously computed spherical harmonics expansion
  from a JSON file. Will set the appropriate fields in the SBGATSphericalHarmo object to
  allow calls to other methods.
  The loadable fields are:
  - facets == number of facets (not needed for evaluation)
  - vertices == number of vertices (not needed for evaluation)
  - totalMass : {value, unit}
  - density : {value, unit}
  - reference_radius : {value, unit}
  - normalized == true if the coefficients are normalized
  - degree == degree of the spherical expansion
  - Cnm_coefs - vector of coefficients triplets {n,m,Cnm}
  - Snm_coefs - vector of coefficients triplets {n,m,Snm}
  @param path JSON file storing the spherical harmonics model
  */
  void LoadFromJson(std::string path);

  /**
  Sets the Cnm coefficients. There is normally no need to use this method outside of Sbgat's tests
  @param[in] Cnm coefficients
  */
  void SetCnm(arma::mat Cnm){this ->Cnm =Cnm;}

  /**
  Sets the Snm coefficients. There is normally no need to use this method outside of Sbgat's tests
  @param[in] Snm coefficients
  */
  void SetSnm(arma::mat Snm){this ->Snm =Snm;}

protected:
  SBGATSphericalHarmo();
  ~SBGATSphericalHarmo() override;

  int RequestData(vtkInformation* request,
    vtkInformationVector** inputVector,
    vtkInformationVector* outputVector) override;

  arma::mat Cnm;
  arma::mat Snm;

  double referenceRadius;
  double density;
  double totalMass;
  double scaleFactor = 1;

  bool normalized;
  unsigned int degree;

  int n_facets;
  int n_vertices;

  bool degreeSet;
  bool densitySet;
  bool referenceRadiusSet;
  bool scaleFactorSet;
  bool setFromJSON;

private:
  SBGATSphericalHarmo(const SBGATSphericalHarmo&) = delete;
  void operator=(const SBGATSphericalHarmo&) = delete;
};

#endif


#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <complex>
#include "larsson_iccv19.h"
#include "kukelova_iccv13.h"
#include "../misc/univariate.h"
#include "../misc/distortion.h"

using namespace Eigen;
using namespace radialpose;
using std::complex;

static const double SMALL_NUMBER = 1e-8;
static const double DAMP_FACTOR = 1e-8;

static void linsolve_known_pose_dist(const Points2D &x, const Points3D &X, double t3, int Np, int Nd, double damp, Camera* camera)
{
	int n_pts = x.cols();
	int n_param = 1 + Np + Nd;

	int n_rows = n_pts;
	if (damp > 0)
		n_rows += n_param - 1;

	// System matrix for normal equations
	Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> A(n_rows, n_param);
	Eigen::Matrix<double, Eigen::Dynamic, 1> b(n_rows, 1);
	A.setZero();  b.setZero();

	double r_pow[4];
	for (int i = 0; i < n_pts; i++) {

		double uu, uv;
		if (std::abs(x(0, i)) < SMALL_NUMBER) {
			uv = x(1, i);
			uu = X(1, i) / (t3 + X(2, i));
		}
		else {
			uv = x(0, i);
			uu = X(0, i) / (t3 + X(2, i));
		}

		//double rd2 = x(0, i) * x(0, i) + x(1, i) * x(1, i);
		double ru2 = (X(0, i) * X(0, i) + X(1, i) * X(1, i)) / ((X(2, i) + t3) * (X(2, i) + t3));

		// compute powers
		r_pow[0] = ru2;            // r^2
		r_pow[1] = ru2 * ru2;      // r^4
		r_pow[2] = ru2 * r_pow[1]; // r^6
		r_pow[3] = ru2 * r_pow[2]; // r^8

		A(i, 0) = uu;
		for (int k = 0; k < Np; k++)
			A(i, 1 + k) = r_pow[k] * uu;
		for (int k = 0; k < Nd; k++)
			A(i, 1 + Np + k) = -uv * r_pow[k];
		b(i) = uv;
	}
	if (damp > 0) {
		for (int i = 1; i < n_param; i++)
			A(n_pts - 1 + i, i) += damp;
	}

	//Eigen::JacobiSVD<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>> svd(A);
	//std::cout << "A: " << A << "\n";
	//std::cout << "svd(A): " << svd.singularValues() << "\n";

	Eigen::Matrix<double, Eigen::Dynamic, 1> sol = A.fullPivHouseholderQr().solve(b);
	camera->focal = sol(0);
	for (int i = 1; i < n_param; i++) {
		if (i <= Np)
			camera->dist_params.push_back(sol(i) / camera->focal);
		else
			camera->dist_params.push_back(sol(i));
	}
}

static void linsolve_known_pose_undist(const Points2D &x, const Points3D &X, double t3, int Np, int Nd, double damp, Camera* camera)
{
	int n_pts = x.cols();
	int n_param = 1 + Np + Nd;

	int n_rows = n_pts;
	if (damp > 0)
		n_rows += n_param - 1;

	// System matrix for normal equations
	Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> A(n_rows, n_param);
	Eigen::Matrix<double, Eigen::Dynamic, 1> b(n_rows, 1);
	A.setZero();  b.setZero();

	double r_pow[4];
	for (int i = 0; i < n_pts; i++) {

		double uu, uv;
		if (std::abs(x(0, i)) < SMALL_NUMBER) {
			uv = x(1, i);
			uu = X(1, i) / (t3 + X(2, i));
		} else {
			uv = x(0, i);
			uu = X(0, i) / (t3 + X(2, i));
		}

		double rd2 = x(0, i) * x(0, i) + x(1, i) * x(1, i);

		// compute powers
		r_pow[0] = rd2;            // r^2
		r_pow[1] = rd2 * rd2;      // r^4
		r_pow[2] = rd2 * r_pow[1]; // r^6
		r_pow[3] = rd2 * r_pow[2]; // r^8

		A(i, 0) = uu;
		for (int k = 0; k < Np; k++)
			A(i, 1 + k) = -r_pow[k] * uv;
		for (int k = 0; k < Nd; k++)
			A(i, 1 + Np + k) = uu * r_pow[k];
		b(i) = uv;
	}
	if (damp > 0) {
		for (int i = 1; i < n_param; i++)
			A(n_pts - 1 + i, i) += damp;
	}

	Eigen::Matrix<double, Eigen::Dynamic, 1> sol = A.fullPivHouseholderQr().solve(b);
	camera->focal = sol(0);
	double f2 = camera->focal * camera->focal;
	double f2k = f2;
	// We have the real mu_k = sol * f^2k and lambda_k = sol * f^2k-1
	for (int i = 1; i < n_param; i++) {
		if (i == Np+1)
			f2k = camera->focal;
		camera->dist_params.push_back(sol(i) * f2k);
		f2k = f2k * f2;
	}
}

// Small refinement on the minimal sample. TODO: Move this to refinement.cc
template<int Np, int Nd>
void radial_refinement_dist(Camera &p, const Points2D &x, const Points3D &X, double damp_factor) {
	// It is assumed that X is already rotated by p
	constexpr int n_pts = std::max(2 + Np + Nd, 5);

	Matrix<double, 2 * n_pts, 2 + Np + Nd> J;
	Matrix<double, 2 * n_pts, 1> res;
	Matrix<double, 2 + Np + Nd, 1> dx;

	for (int i = 0; i < Np; ++i)
		p.dist_params[i] *= p.focal;

	for (int iter = 0; iter < 5; ++iter) {
		for (int i = 0; i < n_pts; ++i) {
			double d = X(2, i) + p.t(2);
			double d2 = d * d;
			double r2 = X.block<2, 1>(0, i).squaredNorm();

			double num = p.focal;
			double denom = d;

			double dnum_dt = 0.0;
			double ddenom_dt = 1.0;

			for (int k = 0; k < Np; k++) {
				double r2d2 = std::pow(r2 / d2, k + 1);

				num += p.dist_params[k] * r2d2;
				dnum_dt += - p.dist_params[k] * 2 * (k + 1) * r2d2 / d;

			}
			for (int k = 0; k < Nd; k++) {
				double r2d2 = std::pow(r2 / d2, k + 1);

				denom += d * p.dist_params[Np + k] * r2d2;
				ddenom_dt += -p.dist_params[Np + k] * (2 * k + 1)*r2d2;
			}

			double factor = num / denom;
			double dfactor_dt = (dnum_dt*denom - num * ddenom_dt) / (denom*denom);
			double dfactor_df = 1 / denom;

			res(2 * i + 0) = factor * X(0, i) - x(0, i);
			res(2 * i + 1) = factor * X(1, i) - x(1, i);

			J(2 * i + 0, 0) = dfactor_dt * X(0, i);
			J(2 * i + 1, 0) = dfactor_dt * X(1, i);
			J(2 * i + 0, 1) = dfactor_df * X(0, i);
			J(2 * i + 1, 1) = dfactor_df * X(1, i);

			for (int k = 0; k < Np; ++k) {
				double r2d2 = std::pow(r2 / d2, k + 1);
				double dfactor_dmu = r2d2 / denom;

				J(2 * i + 0, 2 + k) = dfactor_dmu * X(0, i);
				J(2 * i + 1, 2 + k) = dfactor_dmu * X(1, i);
			}

			for (int k = 0; k < Nd; ++k) {
				double r2d2 = std::pow(r2 / d2, k + 1);
				double dfactor_dlambda = -d * r2d2 * num / (denom*denom);
				J(2 * i + 0, 2 + Np + k) = dfactor_dlambda * X(0, i);
				J(2 * i + 1, 2 + Np + k) = dfactor_dlambda * X(1, i);
			}

		}

		if (res.norm() < SMALL_NUMBER)
			break;

		//std::cout << "res = " << res << "\n";
		//std::cout << "jac = " << J << "\n";

		Matrix<double, 2 + Np + Nd, 2 + Np + Nd> H = J.transpose()*J;
		H.diagonal().array() += 1e-6; // LM dampening
		Matrix<double, 2 + Np + Nd, 1> g = -J.transpose()*res;

		if (Nd > 0 && Np > 0 && damp_factor > 0) {
			// For rational models we add a small dampening factor
			H.template block<Np + Nd,Np + Nd>(2, 2).diagonal().array() += damp_factor;
			for(int i = 0; i < Np+Nd; i++)
				g(2+i) -= damp_factor * p.dist_params[i];
		}

		dx = H.ldlt().solve(g);

		p.t(2) += dx(0);
		p.focal += dx(1);
		for (int i = 0; i < Np; ++i)
			p.dist_params[i] += dx(2 + i);
		for (int i = 0; i < Nd; ++i)
			p.dist_params[i + Np] += dx(Np + 2 + i);

		if (dx.array().abs().maxCoeff() < SMALL_NUMBER)
			break;
	}

	for (int i = 0; i < Np; ++i)
		p.dist_params[i] /= p.focal;
}

template<int Np, int Nd>
void radial_refinement_undist(Camera &p, const Points2D &x, const Points3D &X, double damp_factor) {
	// It is assumed that X is already rotated by p
	constexpr int n_pts = std::max(2 + Np + Nd, 5);

	Matrix<double, 2 * n_pts, 2 + Np + Nd> J;
	Matrix<double, 2 * n_pts, 1> res;
	Matrix<double, 2 + Np + Nd, 1> dx;

	// Change of variables
	double f2 = p.focal * p.focal;
	double f2k = f2;
	for (int i = 0; i < Np; ++i) {
		p.dist_params[i] /= f2k;
		f2k *= f2;
	}
	f2k = p.focal;
	for (int i = 0; i < Nd; ++i) {
		p.dist_params[Np + i] /= f2k;
		f2k *= f2;
	}

	double r_pow[3];

	for (int iter = 0; iter < 5; ++iter) {
		for (int i = 0; i < n_pts; ++i) {
			double d = X(2, i) + p.t(2);
			double r2 = x.col(i).squaredNorm();

			double num = 1.0;
			double denom = p.focal;

			r_pow[0] = r2; // 2
			r_pow[1] = r2 * r2; // 4
			r_pow[2] = r2 * r_pow[1]; // 6

			for (int k = 0; k < Np; k++) {
				num += p.dist_params[k] * r_pow[k];
			}
			for (int k = 0; k < Nd; k++) {
				denom += p.dist_params[Np + k] * r_pow[k];
			}

			double factor = num / denom;

			double dfactor_df = -factor / denom;

			res(2 * i + 0) = factor * x(0, i) - X(0, i) / d;
			res(2 * i + 1) = factor * x(1, i) - X(1, i) / d;

			J(2 * i + 0, 0) = X(0, i) / (d*d); // t
			J(2 * i + 1, 0) = X(1, i) / (d*d);
			J(2 * i + 0, 1) = dfactor_df * x(0, i);
			J(2 * i + 1, 1) = dfactor_df * x(1, i);

			for (int k = 0; k < Np; ++k) {
				double dfactor_dmu = r_pow[k] / denom;
				J(2 * i + 0, 2 + k) = dfactor_dmu * x(0, i);
				J(2 * i + 1, 2 + k) = dfactor_dmu * x(1, i);
			}
			for (int k = 0; k < Nd; ++k) {
				double dfactor_dlambda = -r_pow[k] * num / (denom*denom);
				J(2 * i + 0, 2 + Np + k) = dfactor_dlambda * x(0, i);
				J(2 * i + 1, 2 + Np + k) = dfactor_dlambda * x(1, i);
			}
		}

		//std::cout << "res = " << res << "\n";
		//std::cout << "jac = " << J << "\n";
		if (res.norm() < SMALL_NUMBER)
			break;

		Matrix<double, 2 + Np + Nd, 2 + Np + Nd> H = J.transpose()*J;
		H.diagonal().array() += 1e-6; // LM dampening
		Matrix<double, 2 + Np + Nd, 1> g = -J.transpose()*res;

		if (Nd > 0 && Np > 0 && damp_factor > 0) {
			// For rational models we add a small dampening factor
			H.template block<Np + Nd,Np + Nd>(2, 2).diagonal().array() += damp_factor;
			for (int i = 0; i < Np + Nd; i++)
				g(2 + i) -= damp_factor * p.dist_params[i];
		}

		dx = H.ldlt().solve(g);

		p.t(2) += dx(0);
		p.focal += dx(1);
		for (int i = 0; i < Np; ++i)
			p.dist_params[i] += dx(2 + i);
		for (int i = 0; i < Nd; ++i)
			p.dist_params[i + Np] += dx(Np + 2 + i);

		if (dx.array().abs().maxCoeff() < SMALL_NUMBER)
			break;
	}

	// Revert change of variables
	f2 = p.focal * p.focal;
	f2k = f2;
	for (int i = 0; i < Np; ++i) {
		p.dist_params[i] *= f2k;
		f2k *= f2;
	}
	f2k = p.focal;
	for (int i = 0; i < Nd; ++i) {
		p.dist_params[Np + i] *= f2k;
		f2k *= f2;
	}
}

inline double simple_preconditioner(const Points2D &x, const Points3D &X) {
	// Simple preconditioner using the first two points
	double nx1 = x.col(0).squaredNorm();
	double nx2 = x.col(1).squaredNorm();
	double sx1 = X.block<2, 1>(0, 0).dot(x.col(0));
	double sx2 = X.block<2, 1>(0, 1).dot(x.col(1));
	return (X(2, 1)*nx2*sx1 - X(2, 0)*nx1*sx2) / (nx1 * sx2 - nx2 * sx1);
}

template<int Np, int Nd, bool DistortionModel>
int radialpose::larsson_iccv19::Solver<Np,Nd,DistortionModel>::solve(const Points2D& image_points, const Points3D& world_points, std::vector<Camera>* poses) const
{

	std::vector<Camera> initial_poses;
	std::vector<double> t3;

	if (use_radial_solver) {
		kukelova_iccv13::Radial1DSolver::p5p_radial_impl(image_points, world_points, &initial_poses);
	} else {
		initial_poses.push_back(Camera(Matrix3d::Identity(), Vector3d::Zero()));
	}
	Matrix<double, 3, Dynamic> X;

	for (int k = 0; k < initial_poses.size(); k++) {
		t3.clear();

		X = initial_poses[k].R * world_points;
		X.colwise() += initial_poses[k].t;

		//std::cout << "Initial pose " << k + 1 << "/" << initial_poses.size() << "\n";
		//std::cout << "R=" << initial_poses[k].R << "\nt=" << initial_poses[k].t << "\n";
		//std::cout << "X=" << X << "\n";

		double t0 = 0;
		if (use_precond) {
			t0 = simple_preconditioner(image_points, X);
			X.row(2).array() += t0;
		}

		solver_impl(image_points, X, &t3);

		for (int i = 0; i < t3.size(); ++i) {
			Camera pose;
			pose.R = initial_poses[k].R;
			pose.t = initial_poses[k].t;
			pose.t(2) = t3[i];

			if (DistortionModel) {
				linsolve_known_pose_dist(image_points, X, t3[i], Np, Nd, damp_factor, &pose);
				if(root_refinement)
					radial_refinement_dist<Np, Nd>(pose, image_points, X, damp_factor);
			} else {
				linsolve_known_pose_undist(image_points, X, t3[i], Np, Nd, damp_factor, &pose);
				if(root_refinement)
					radial_refinement_undist<Np, Nd>(pose, image_points, X, damp_factor);
			}

			if (pose.focal < 0) {
				// flipped solution
				pose.focal = -pose.focal;
				pose.R.row(0) = -pose.R.row(0);
				pose.R.row(1) = -pose.R.row(1);
				pose.t(0) = -pose.t(0);
				pose.t(1) = -pose.t(1);
			}

			// Revert precond
			pose.t(2) += t0;

			//std::cout << "solution[" << i << "], t3=" << pose.t(2) << "\n";
			poses->push_back(pose);
		}
	}
	return poses->size();
}

// Template instantiations
template class radialpose::larsson_iccv19::Solver<1, 0, true>;
template class radialpose::larsson_iccv19::Solver<2, 0, true>;
template class radialpose::larsson_iccv19::Solver<3, 0, true>;
template class radialpose::larsson_iccv19::Solver<3, 3, true>;
template class radialpose::larsson_iccv19::Solver<1, 0, false>;
//template class radialpose::larsson_iccv19::Solver<2, 0, false>;
//template class radialpose::larsson_iccv19::Solver<3, 0, false>;
//template class radialpose::larsson_iccv19::Solver<3, 3, false>;

template class radialpose::PoseEstimator<radialpose::larsson_iccv19::Solver<1, 0, true>>;
template class radialpose::PoseEstimator<radialpose::larsson_iccv19::Solver<2, 0, true>>;
template class radialpose::PoseEstimator<radialpose::larsson_iccv19::Solver<3, 0, true>>;
template class radialpose::PoseEstimator<radialpose::larsson_iccv19::Solver<3, 3, true>>;
template class radialpose::PoseEstimator<radialpose::larsson_iccv19::Solver<1, 0, false>>;

/*
 This is broken?
// These are implemented in larsson_iccv19_impl.cc
extern template int radialpose::larsson_iccv19::Solver<1, 0, true>::solver_impl(Eigen::Matrix<double, 2, Eigen::Dynamic>, Eigen::Matrix<double, 3, Eigen::Dynamic>, std::vector<double>*);
extern template int radialpose::larsson_iccv19::Solver<2, 0, true>::solver_impl(Eigen::Matrix<double, 2, Eigen::Dynamic>, Eigen::Matrix<double, 3, Eigen::Dynamic>, std::vector<double>*);
extern template int radialpose::larsson_iccv19::Solver<3, 0, true>::solver_impl(Eigen::Matrix<double, 2, Eigen::Dynamic>, Eigen::Matrix<double, 3, Eigen::Dynamic>, std::vector<double>*);
extern template int radialpose::larsson_iccv19::Solver<3, 3, true>::solver_impl(Eigen::Matrix<double, 2, Eigen::Dynamic>, Eigen::Matrix<double, 3, Eigen::Dynamic>, std::vector<double>*);
extern template int radialpose::larsson_iccv19::Solver<1, 0, false>::solver_impl(Eigen::Matrix<double, 2, Eigen::Dynamic>, Eigen::Matrix<double, 3, Eigen::Dynamic>, std::vector<double>*);
*/


#include <Eigen/Dense>
#include <iostream>

using namespace Eigen;
using namespace std;

int main() {
  ArrayXXf m(2, 2);

  // assign some values coefficient by coefficient
  m(0, 0) = 1.0;
  m(0, 1) = 2.0;
  m(1, 0) = 3.0;
  m(1, 1) = m(0, 1) + m(1, 0);

  // print values to standard output
  cout << m << endl << endl;

  // using the comma-initializer is also allowed
  m << 1.0, 2.0,
      3.0, 4.0;

  // print values to standard output
  cout << m << endl;
}


/*!
  \file gpp_hyperparameter_optimization_demo.cpp
  \rst
  ``moe/optimal_learning/cpp/gpp_hyperparameter_optimization_demo.cpp``

  This is a demo for the model selection (via hyperparameter optimization) capability
  present in this project.  These capabilities live in
  gpp_model_selection.

  In gpp_expected_improvement_demo, we choose the hyperparameters arbitrarily.  Here,
  we will walk through an example of how one would select hyperparameters for a given
  class of covariance function; here, SquareExponential will do.  This demo supports:

  1. User-specified training data
  2. Randomly generated training data (more automatic)

  More details on the second case:

  1. Choose a set of hyperparameters randomly: source covariance
  2. Build a fake\* training set by drawing from a GP with source covariance, at randomly
     chosen locations
     \* By defining OL_USER_INPUTS to 1, you can specify your own input data.
  3. Choose a new random set of hyperparameters and run hyperparameter optimization

     a. Show log likelihood using the optimized hyperparameters AND the source hyperparameters
     b. observe that with larger training sets, the optimized hyperparameters converge
        to the source values; but in smaller sets other optima may exist

  Further notes about [newton] optimization performance and robustness are spread throughout the
  demo code, placed near the function call/object construction that they are relevant to.

  Please read and understand gpp_expected_improvement_demo.cpp before going through
  this example.  In addition, understanding gpp_model_selection.hpp's
  file comments (as well as cpp for devs) is prerequisite.
\endrst*/

#include <cstdio>

#include <vector>

#include <boost/random/uniform_real.hpp>  // NOLINT(build/include_order)

#include "gpp_common.hpp"
#include "gpp_covariance.hpp"
#include "gpp_domain.hpp"
#include "gpp_logging.hpp"
#include "gpp_math.hpp"
#include "gpp_model_selection.hpp"
#include "gpp_optimizer_parameters.hpp"
#include "gpp_random.hpp"
#include "gpp_test_utils.hpp"

#define OL_USER_INPUTS 0

using namespace optimal_learning;  // NOLINT, i'm lazy in this file which has no external linkage anyway

int main() {
  using DomainType = TensorProductDomain;
  using HyperparameterDomainType = TensorProductDomain;
  // here we set some configurable parameters
  // feel free to change them (and recompile) as you explore
  // comments next to each parameter will indicate its purpose and domain

  // the "spatial" dimension, aka the number of independent (experiment) parameters
  // i.e., this is the dimension of the points in points_sampled
  static const int dim = 3;  // > 0

  // number of points that we have already sampled; i.e., size of the training set
  static const int num_sampled = 100;  // >= 0
  // observe that as num_sampled increases, the optimal set of hyperparameters (via optimization) will approach
  // the set used to generate the input data (in the case of generating inputs randomly from a GP).  Don't try overly
  // large values or it will be slow; for reference 500 samples takes ~2-3 min on my laptop whereas 100 samples takes ~1s

  // the log likelihoods will also decrease in value since by adding more samples, we are more greatly restricting the GP
  // into ever-narrower sets of likely realizations

  UniformRandomGenerator uniform_generator(314);  // repeatable results
  // construct with (base_seed, thread_id) to generate a 'random' seed

  // specifies the domain of each independent variable in (min, max) pairs
  // set appropriately for user-specified inputs
  // mostly irrelevant for randomly generated inputs
  std::vector<ClosedInterval> domain_bounds = {
    {-1.5, 2.3},  // first dimension
    {0.1, 3.1},   // second dimension
    {1.7, 2.9}};  // third dimension
  DomainType domain(domain_bounds.data(), dim);

  // now we allocate point sets; ALL POINTS MUST LIE INSIDE THE DOMAIN!
  std::vector<double> points_sampled(num_sampled*dim);

  std::vector<double> points_sampled_value(num_sampled);

  // default to 0 noise
  std::vector<double> noise_variance(num_sampled, 0.0);  // each entry must be >= 0.0
  // choosing too much noise makes little sense: cannot make useful predicitions if data
  // is drowned out by noise
  // choosing 0 noise is dangerous for large problems; the covariance matrix becomes very
  // ill-conditioned, and adding noise caps the maximum condition number at roughly
  // 1.0/min(noise_variance)

  // covariance selection
  using CovarianceClass = SquareExponential;  // see gpp_covariance.hpp for other options

  // arbitrary hyperparameters used to generate data
  std::vector<double> hyperparameters_original(1 + dim);
  CovarianceClass covariance_original(dim, 1.0, 1.0);
  // CovarianceClass provides SetHyperparameters, GetHyperparameters to read/modify
  // hyperparameters later on
  // Generate hyperparameters randomly
  boost::uniform_real<double> uniform_double_for_hyperparameter(0.5, 1.5);
  FillRandomCovarianceHyperparameters(uniform_double_for_hyperparameter, &uniform_generator,
                                      &hyperparameters_original, &covariance_original);

  std::vector<ClosedInterval> hyperparameter_domain_bounds(covariance_original.GetNumberOfHyperparameters(), {1.0e-10, 1.0e10});
  HyperparameterDomainType hyperparameter_domain(hyperparameter_domain_bounds.data(),
                                                 covariance_original.GetNumberOfHyperparameters());

  // now fill data
#if OL_USER_INPUTS == 1
  // if you prefer, insert your own data here
  // requirements aka variables that must be set:
  // noise variance, num_sampled values: defaulted to 0; need to set this for larger data sets to deal with conditioning
  // points_sampled, num_sampled*dim values: the locations of already-sampled points; must be INSIDE the domain
  // points_sampled_value, num_sampled values: the function values at the already-sampled points
  // covariance_perturbed: a CovarianceClass object constructed with perturbed (from covariance_original) hyperparameters;
  //   must have decltype(covariance_perturbed) == decltype(covariance_original) for hyperparameter opt to make any sense

  // NOTE: the GP is 0-mean, so shift your points_sampled_value entries accordingly
  // e.g., if the samples are from a function with mean M, subtract it out
#else
  // generate GP inputs randomly

  // set noise
  std::fill(noise_variance.begin(), noise_variance.end(), 1.0e-1);  // arbitrary choice

  // use latin hypercube sampling to get a reasonable distribution of training point locations
  domain.GenerateUniformPointsInDomain(num_sampled, &uniform_generator, points_sampled.data());

  // build an empty GP: since num_sampled (last arg) is 0, none of the data arrays will be used here
  GaussianProcess gp_generator(covariance_original, points_sampled.data(), points_sampled_value.data(),
                               noise_variance.data(), dim, 0);
  // fill the GP with randomly generated data
  FillRandomGaussianProcess(points_sampled.data(), noise_variance.data(), dim, num_sampled,
                            points_sampled_value.data(), &gp_generator);

  // choose a random initial guess reasonably far away from hyperparameters_original
  // to find some optima (see WARNING2 below), it may be necessary to start with hyperparameters smaller than the originals or
  // of similar magnitude
  std::vector<double> hyperparameters_perturbed(covariance_original.GetNumberOfHyperparameters());
  CovarianceClass covariance_perturbed(dim, 1.0, 1.0);
  boost::uniform_real<double> uniform_double_for_wrong_hyperparameter(5.0, 12.0);
  FillRandomCovarianceHyperparameters(uniform_double_for_wrong_hyperparameter, &uniform_generator,
                                      &hyperparameters_perturbed, &covariance_perturbed);
#endif

  // log likelihood type selection
  using LogLikelihoodEvaluator = LogMarginalLikelihoodEvaluator;
  // log likelihood evaluator object
  LogLikelihoodEvaluator log_marginal_eval(points_sampled.data(), points_sampled_value.data(),
                                           noise_variance.data(), dim, num_sampled);

  int total_newton_errors = 0;  // number of newton runs that failed due to singular hessians
  int newton_max_num_steps = 500;  // max number of newton steps
  double gamma_newton = 1.05;  // newton diagonal dominance scale-down factor (see newton docs for details)
  double pre_mult_newton = 1.0e-1;  // newton diagonal dominance scaling factor (see newton docs for details)
  double max_relative_change_newton = 1.0;
  double tolerance_newton = 1.0e-11;
  NewtonParameters newton_parameters(1, newton_max_num_steps, gamma_newton, pre_mult_newton,
                                     max_relative_change_newton, tolerance_newton);

  // call newton to optimize hyperparameters
  // in general if this takes the full hyperparameter_max_num_steps iterations, something went wrong
  // newton's solution:
  std::vector<double> new_newton_hyperparameters(covariance_original.GetNumberOfHyperparameters());

  printf(OL_ANSI_COLOR_CYAN "ORIGINAL HYPERPARMETERS:\n" OL_ANSI_COLOR_RESET);
  printf("Original Hyperparameters:\n");
  PrintMatrix(hyperparameters_original.data(), 1, covariance_original.GetNumberOfHyperparameters());

  printf(OL_ANSI_COLOR_CYAN "NEWTON OPTIMIZED HYPERPARAMETERS:\n" OL_ANSI_COLOR_RESET);
  // run newton optimization
  total_newton_errors += NewtonHyperparameterOptimization(log_marginal_eval, covariance_perturbed,
                                                          newton_parameters, hyperparameter_domain,
                                                          new_newton_hyperparameters.data());
  // WARNING: the gradient of log marginal appears to go to 0 as you move toward infinity.  if you do not start
  // close enough to an optima or have overly aggressive diagonal dominance settings, newton will skip miss everything
  // going on locally and shoot out to these solutions.
  // Having hyperparameters = 1.0e10 is nonsense, and usually this problem is further signaled by a log marginal likelihood
  // that is POSITIVE (impossible since p \in [0,1], so log(p) \in (-\infty, 0])

  // Long-term, we should solve this problem by multistarting newton.  Additionally there will be some kind of "quick kill"
  // mechanism needed--when newton is wandering down the wrong path (or to an already-known solution?) we should detect it
  // and kill it quickly to keep cost low.
  // For now, just play around with different initial conditions or more conservative gamam settings.

  // WARNING2: for small num_sampled, it often appears that the solution becomes independent of one or more hyperparameters.
  // e.g., in 2D, we'd have an optimal "ridge."  Finding this robustly requires starting near it, so the random choice of
  // initial conditions can fail horribly in general.

  // WARNING3: if you choose large values of num_sampled (like 300), this can be quite slow; about 5min on my computer
  // sometimes the reason is that machine prescision prevents us from reaching the cutoff criterion:
  // norm_gradient_likelihood <= 1.0e-13 in NewtonHyperparameterOptimization() in gpp_model_selection...cpp
  // So you may need to relax this to 1.0e-10 or something so that we aren't just spinning wheels at almost-converged but
  // unable to actually move anywhere.

  printf("Result of newton:\n");
  PrintMatrix(new_newton_hyperparameters.data(), 1, covariance_original.GetNumberOfHyperparameters());

  if (total_newton_errors > 0) {
    printf("WARNING: %d newton runs exited due to singular Hessian matrices.\n", total_newton_errors);
  }

  printf(OL_ANSI_COLOR_CYAN "LOG LIKELIHOOD + GRADIENT AT NEWTON OPTIMIZED FINAL HYPERPARAMS:\n" OL_ANSI_COLOR_RESET);

  CovarianceClass covariance_final(dim, new_newton_hyperparameters[0], new_newton_hyperparameters.data() + 1);
  typename LogLikelihoodEvaluator::StateType log_marginal_state_newton_optimized_hyper(log_marginal_eval,
                                                                                covariance_final);
  double newton_log_marginal_opt = log_marginal_eval.ComputeLogLikelihood(log_marginal_state_newton_optimized_hyper);
  printf("newton optimized log marginal likelihood = %.18E\n", newton_log_marginal_opt);

  std::vector<double> grad_log_marginal_opt(covariance_final.GetNumberOfHyperparameters());
  log_marginal_eval.ComputeGradLogLikelihood(&log_marginal_state_newton_optimized_hyper,
                                             grad_log_marginal_opt.data());
  printf("grad log likelihood: ");
  PrintMatrix(grad_log_marginal_opt.data(), 1, covariance_final.GetNumberOfHyperparameters());

  printf(OL_ANSI_COLOR_CYAN "LOG LIKELIHOOD + GRADIENT AT ORIGINAL HYPERPARAMS:\n" OL_ANSI_COLOR_RESET);
  typename LogLikelihoodEvaluator::StateType log_marginal_state_original_hyper(log_marginal_eval,
                                                                               covariance_original);

  double original_log_marginal = log_marginal_eval.ComputeLogLikelihood(log_marginal_state_original_hyper);
  printf("original log marginal likelihood = %.18E\n", original_log_marginal);

  std::vector<double> original_grad_log_marginal(covariance_original.GetNumberOfHyperparameters());
  log_marginal_eval.ComputeGradLogLikelihood(&log_marginal_state_original_hyper,
                                             original_grad_log_marginal.data());
  printf("grad log likelihood: ");
  PrintMatrix(original_grad_log_marginal.data(), 1, covariance_original.GetNumberOfHyperparameters());

  return 0;
}  // end main


#include "mex.h"
#include <Eigen/Dense>
#include "re3q3/re3q3.h"

void printUsage() {
	mexPrintf("[x] = re3q3(C);\n");
	mexPrintf("    coefficient order is [ x^2, xy, xz, y^2, yz, z^2, x, y, z, 1.0 ]\n");
}

void mexFunction(int nlhs,mxArray *plhs[], int nrhs, const mxArray *prhs[]) {

	if (nrhs != 1) {
		printUsage();
		mexErrMsgTxt("Please, specify 3 x 10N coefficient matrix");
	}
	if (nlhs != 1) {
		printUsage();
		mexErrMsgTxt("Require one output argument.");
	}

	if ((mxGetM(prhs[0]) != 3) || (mxGetN(prhs[0]) % 10 != 0)) {
		printUsage();
		mexErrMsgTxt("One input 3 x 10N matrix is required.");
	}

	int n_instances = mxGetN(prhs[0]) / 10;

	double *data = mxGetPr(prhs[0]);

	Eigen::Matrix<double, 3, 10> coeffs;
	Eigen::Matrix<double, 3, 8> solutions;

	if (n_instances == 1) {
		// we are solving a single instance, output is 3xNsols
		coeffs = Eigen::Map<Eigen::Matrix<double, 3, 10>>(data);

		int n_sols = re3q3::re3q3(coeffs, &solutions);
		plhs[0] = mxCreateDoubleMatrix(3, n_sols, mxREAL);

		Eigen::Map<Eigen::MatrixXd> output_matrix = Eigen::Map<Eigen::MatrixXd>(mxGetPr(plhs[0]), 3, n_sols);
		output_matrix = solutions.block(0, 0, 3, n_sols);
	} else {
		// we are solving multiple instances, output is 3x8N
		plhs[0] = mxCreateDoubleMatrix(3, 8 * n_instances, mxREAL);
		double *output = mxGetPr(plhs[0]);

		Eigen::Map<Eigen::MatrixXd> output_matrix = Eigen::Map<Eigen::MatrixXd>(mxGetPr(plhs[0]), 3, 8*n_instances);
		output_matrix.setZero();

		for (int i = 0; i < n_instances; i++) {
			coeffs = Eigen::Map<Eigen::Matrix<double, 3, 10>>(data + 30 * i);
			solutions.setZero();
			int n_sols = re3q3::re3q3(coeffs, &solutions);
			output_matrix.block(0, i * 8, 3, n_sols) = solutions.block(0,0,3,n_sols);
		}
	}
}