// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

/*
 *   File "gmm_b_tapenade_generated.c" is generated by Tapenade 3.14 (r7259) from this file.
 *   To reproduce such a generation you can use Tapenade CLI
 *   (can be downloaded from http://www-sop.inria.fr/tropics/tapenade/downloading.html)
 *
 *   Firstly, add a type declaration of Wishart to the content of this file (Tapenade can't process a file
 *   with unknown types). You can both take this declaration from the file "<repo root>/src/cpp/shared/defs.h" or
 *   copypaste the following lines removing asterisks:
 *
 *   typedef struct
 *   {
 *       double gamma;
 *       int m;
 *   } Wishart;
 *
 *   After Tapenade CLI installing use the next command to generate a file:
 *
 *      tapenade -b -o gmm_tapenade -head "gmm_objective(err)/(alphas means icf)" gmm.c
 *
 *   This will produce a file "gmm_tapenade_b.c" which content will be the same as the content of "gmm_b_tapenade_generated.c",
 *   except one-line header and a Wishart typedef (which should be removed). Moreover a log-file "gmm_tapenade_b.msg" will be produced.
 *
 *   NOTE: the code in "gmm_b_tapenade_generated.c" is wrong and won't work.
 *         REPAIRED SOURCE IS STORED IN THE FILE "gmm_b.c".
 *         You can either use diff tool or read "gmm_b.c" header to figure out what changes was performed to fix the code.
 *
 *   NOTE: you can also use Tapenade web server (http://tapenade.inria.fr:8080/tapenade/index.jsp)
 *         for generating but result can be slightly different.
 */
#include "../adbench/gmm.h"
extern "C" {
#include "gmm.h"

/* ==================================================================== */
/*                                UTILS                                 */
/* ==================================================================== */

// This throws error on n<1
double arr_max(int n, double const* x)
{
    int i;
    double m = x[0];
    for (i = 1; i < n; i++)
    {
        if (m < x[i])
        {
            m = x[i];
        }
    }

    return m;
}

// sum of component squares
double sqnorm(int n, double const* x)
{
    int i;
    double res = x[0] * x[0];
    for (i = 1; i < n; i++)
    {
        res = res + x[i] * x[i];
    }

    return res;
}

// out = a - b
void subtract(
    int d,
    double const* x,
    double const* y,
    double* out
)
{
    int id;
    for (id = 0; id < d; id++)
    {
        out[id] = x[id] - y[id];
    }
}

double log_sum_exp(int n, double const* x)
{
    int i;
    double mx = arr_max(n, x);
    double semx = 0.0;

    for (i = 0; i < n; i++)
    {
        semx = semx + exp(x[i] - mx);
    }

    return log(semx) + mx;
}

__attribute__((const))
double log_gamma_distrib(double a, double p)
{
    int j;
    double out = 0.25 * p * (p - 1) * log(PI);

    for (j = 1; j <= p; j++)
    {
        out = out + lgamma(a + 0.5 * (1 - j));
    }

    return out;
}

/* ======================================================================== */
/*                                MAIN LOGIC                                */
/* ======================================================================== */

double log_wishart_prior(
    int p,
    int k,
    Wishart wishart,
    double const* sum_qs,
    double const* Qdiags,
    double const* icf
)
{
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p * (p + 1) / 2;

    double C = n * p * (log(wishart.gamma) - 0.5 * log(2)) - log_gamma_distrib(0.5 * n, p);

    double out = 0;
    for (ik = 0; ik < k; ik++)
    {
        double frobenius = sqnorm(p, &Qdiags[ik * p]) + sqnorm(icf_sz - p, &icf[ik * icf_sz + p]);
        out = out + 0.5 * wishart.gamma * wishart.gamma * (frobenius) - wishart.m * sum_qs[ik];
    }

    return out - k * C;
}

void preprocess_qs(
    int d,
    int k,
    double const* icf,
    double* sum_qs,
    double* Qdiags
)
{
    int ik, id;
    int icf_sz = d * (d + 1) / 2;
    for (ik = 0; ik < k; ik++)
    {
        sum_qs[ik] = 0.;
        for (id = 0; id < d; id++)
        {
            double q = icf[ik * icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik * d + id] = exp(q);
        }
    }
}

void Qtimesx(
    int d,
    double const* Qdiag,
    double const* ltri, // strictly lower triangular part
    double const* x,
    double* out
)
{
    int i, j;
    for (i = 0; i < d; i++)
    {
        out[i] = Qdiag[i] * x[i];
    }

    //caching lparams as scev doesn't replicate index calculation
    // todo note changing to strengthened form
    //int Lparamsidx = 0;
    for (i = 0; i < d; i++)
    {
    	int Lparamsidx = i*(2*d-i-1)/2;
        for (j = i + 1; j < d; j++)
        {
            // and this x
            out[j] = out[j] + ltri[Lparamsidx] * x[i];
            Lparamsidx++;
        }
    }
}

void gmm_objective(
    int d,
    int k,
    int n,
    double const* __restrict alphas,
    double const* __restrict means,
    double const* __restrict icf,
    double const* __restrict x,
    Wishart wishart,
    double* __restrict err
)
{
    #define int int64_t
    int ix, ik;
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);
    int icf_sz = d * (d + 1) / 2;

    double* Qdiags = (double*)malloc(d * k * sizeof(double));
    double* sum_qs = (double*)malloc(k * sizeof(double));
    double* xcentered = (double*)malloc(d * sizeof(double));
    double* Qxcentered = (double*)malloc(d * sizeof(double));
    double* main_term = (double*)malloc(k * sizeof(double));

    preprocess_qs(d, k, icf, &sum_qs[0], &Qdiags[0]);

    double slse = 0.;
    for (ix = 0; ix < n; ix++)
    {
        for (ik = 0; ik < k; ik++)
        {
            subtract(d, &x[ix * d], &means[ik * d], &xcentered[0]);
            Qtimesx(d, &Qdiags[ik * d], &icf[ik * icf_sz + d], &xcentered[0], &Qxcentered[0]);
            // two caches for qxcentered at idx 0 and at arbitrary index
            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5 * sqnorm(d, &Qxcentered[0]);
        }

        // storing cmp for max of main_term
        // 2 x (0 and arbitrary) storing sub to exp
        // storing sum for use in log
        slse = slse + log_sum_exp(k, &main_term[0]);
    }

    //storing cmp of alphas
    double lse_alphas = log_sum_exp(k, alphas);

    *err = CONSTANT + slse - n * lse_alphas + log_wishart_prior(d, k, wishart, &sum_qs[0], &Qdiags[0], icf);

    free(Qdiags);
    free(sum_qs);
    free(xcentered);
    free(Qxcentered);
    free(main_term);
    #undef int
}

extern int enzyme_const;
extern int enzyme_dup;
extern int enzyme_dupnoneed;
void __enzyme_autodiff(...) noexcept;

// *      tapenade -b -o gmm_tapenade -head "gmm_objective(err)/(alphas means icf)" gmm.c
void dgmm_objective(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {
    __enzyme_autodiff(
            gmm_objective,
            enzyme_const, d,
            enzyme_const, k,
            enzyme_const, n,
            enzyme_dup, alphas, alphasb,
            enzyme_dup, means, meansb,
            enzyme_dup, icf, icfb,
            enzyme_const, x,
            enzyme_const, wishart,
            enzyme_dupnoneed, err, errb);
}

}

// ! Tapenade
extern "C" {

#include <adBuffer.h>

/*
  Differentiation of arr_max in reverse (adjoint) mode:
   gradient     of useful results: *x arr_max
   with respect to varying inputs: *x
   Plus diff mem management of: x:in

 ====================================================================
                                UTILS
 ==================================================================== */
// This throws error on n<1
void arr_max_b(int n, const double *x, double *xb, double arr_maxb) {
    int i;
    double m = x[0];
    double mb = 0.0;
    int branch;
    double arr_max;
    for (i = 1; i < n; ++i)
        if (m < x[i]) {
            m = x[i];
            pushControl1b(1);
        } else
            pushControl1b(0);
    mb = arr_maxb;
    for (i = n-1; i > 0; --i) {
        popControl1b(&branch);
        if (branch != 0) {
            xb[i] = xb[i] + mb;
            mb = 0.0;
        }
    }
    xb[0] = xb[0] + mb;
}

/* ====================================================================
                                UTILS
 ==================================================================== */
// This throws error on n<1
double arr_max_nodiff(int n, const double *x) {
    int i;
    double m = x[0];
    for (i = 1; i < n; ++i)
        if (m < x[i])
            m = x[i];
    return m;
}

/*
  Differentiation of sqnorm in reverse (adjoint) mode:
   gradient     of useful results: *x sqnorm
   with respect to varying inputs: *x
   Plus diff mem management of: x:in
*/
// sum of component squares
void sqnorm_b(int n, const double *x, double *xb, double sqnormb) {
    int i;
    double res = x[0]*x[0];
    double resb = 0.0;
    double sqnorm;
    resb = sqnormb;
    for (i = n-1; i > 0; --i)
        xb[i] = xb[i] + 2*x[i]*resb;
    xb[0] = xb[0] + 2*x[0]*resb;
}

// sum of component squares
double sqnorm_nodiff(int n, const double *x) {
    int i;
    double res = x[0]*x[0];
    for (i = 1; i < n; ++i)
        res = res + x[i]*x[i];
    return res;
}

/*
  Differentiation of subtract in reverse (adjoint) mode:
   gradient     of useful results: *out *y
   with respect to varying inputs: *out *y
   Plus diff mem management of: out:in y:in
*/
// out = a - b
void subtract_b(int d, const double *x, const double *y, double *yb, double *
        out, double *outb) {
    int id;
    for (id = d-1; id > -1; --id) {
        yb[id] = yb[id] - outb[id];
        outb[id] = 0.0;
    }
}

// out = a - b
void subtract_nodiff(int d, const double *x, const double *y, double *out) {
    int id;
    for (id = 0; id < d; ++id)
        out[id] = x[id] - y[id];
}

/*
  Differentiation of log_sum_exp in reverse (adjoint) mode:
   gradient     of useful results: *x log_sum_exp
   with respect to varying inputs: *x
   Plus diff mem management of: x:in
*/
void log_sum_exp_b(int n, const double *x, double *xb, double log_sum_expb) {
    int i;
    double mx;
    double mxb;
    double tempb;
    double log_sum_exp;
    mx = arr_max_nodiff(n, x);
    double semx = 0.0;
    double semxb = 0.0;
    for (i = 0; i < n; ++i)
        semx = semx + exp(x[i] - mx);
    semxb = log_sum_expb/semx;
    mxb = log_sum_expb;
    for (i = n-1; i > -1; --i) {
        tempb = exp(x[i]-mx)*semxb;
        xb[i] = xb[i] + tempb;
        mxb = mxb - tempb;
    }
    arr_max_b(n, x, xb, mxb);
}

double log_sum_exp_nodiff(int n, const double *x) {
    int i;
    double mx;
    mx = arr_max_nodiff(n, x);
    double semx = 0.0;
    for (i = 0; i < n; ++i)
        semx = semx + exp(x[i] - mx);
    return log(semx) + mx;
}

double log_gamma_distrib_nodiff(double a, double p) {
    int j;
    /* TFIX */
    double out = 0.25*p*(p-1)*log(PI);
    double arg1;
    float result1;
    for (j = 1; j < p+1; ++j) {
        arg1 = a + 0.5*(1-j);
        result1 = lgamma(arg1);
        out = out + result1;
    }
    return out;
}

/*
  Differentiation of log_wishart_prior in reverse (adjoint) mode:
   gradient     of useful results: log_wishart_prior
   with respect to varying inputs: *Qdiags *sum_qs *icf
   Plus diff mem management of: Qdiags:in sum_qs:in icf:in

 ========================================================================
                                MAIN LOGIC
 ======================================================================== */
void log_wishart_prior_b(int p, int k, Wishart wishart, const double *sum_qs,
        double *sum_qsb, const double *Qdiags, double *Qdiagsb, const double *
        icf, double *icfb, double log_wishart_priorb) {
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p*(p+1)/2;
    double C;
    float arg1;
    double result1;
    double out = 0;
    double outb = 0.0;
    double log_wishart_prior;
    for (ik = 0; ik < k; ++ik) {
        double frobenius;
        double result1;
        int arg1;
        double result2;
    }
    outb = log_wishart_priorb;
    for (ik = 0; ik < k * p; ik++) /* TFIX */
        Qdiagsb[ik] = 0.0;
    for (ik = 0; ik < k; ik++) /* TFIX */
        sum_qsb[ik] = 0.0;
    for (ik = 0; ik < k * icf_sz; ik++) /* TFIX */
        icfb[ik] = 0.0;
    for (ik = k-1; ik > -1; --ik) {
        double frobenius;
        double frobeniusb;
        double result1;
        double result1b;
        int arg1;
        double result2;
        double result2b;
        frobeniusb = wishart.gamma*wishart.gamma*0.5*outb;
        sum_qsb[ik] = sum_qsb[ik] - wishart.m*outb;
        result1b = frobeniusb;
        result2b = frobeniusb;
        arg1 = icf_sz - p;
        sqnorm_b(arg1, &(icf[ik*icf_sz + p]), &(icfb[ik*icf_sz + p]), result2b
                );
        sqnorm_b(p, &(Qdiags[ik*p]), &(Qdiagsb[ik*p]), result1b);
    }
}

/* ========================================================================
                                MAIN LOGIC
 ======================================================================== */
double log_wishart_prior_nodiff(int p, int k, Wishart wishart, const double *
        sum_qs, const double *Qdiags, const double *icf) {
    int ik;
    int n = p + wishart.m + 1;
    int icf_sz = p*(p+1)/2;
    double C;
    float arg1;
    double result1;
    arg1 = 0.5*n;
    result1 = log_gamma_distrib_nodiff(arg1, p);
    C = n*p*(log(wishart.gamma)-0.5*log(2)) - result1;
    double out = 0;
    for (ik = 0; ik < k; ++ik) {
        double frobenius;
        double result1;
        int arg1;
        double result2;
        result1 = sqnorm_nodiff(p, &(Qdiags[ik*p]));
        arg1 = icf_sz - p;
        result2 = sqnorm_nodiff(arg1, &(icf[ik*icf_sz + p]));
        frobenius = result1 + result2;
        out = out + 0.5*wishart.gamma*wishart.gamma*frobenius - wishart.m*
            sum_qs[ik];
    }
    return out - k*C;
}

/*
  Differentiation of preprocess_qs in reverse (adjoint) mode:
   gradient     of useful results: *Qdiags *sum_qs *icf
   with respect to varying inputs: *icf
   Plus diff mem management of: Qdiags:in sum_qs:in icf:in
*/
void preprocess_qs_b(int d, int k, const double *icf, double *icfb, double *
        sum_qs, double *sum_qsb, double *Qdiags, double *Qdiagsb) {
    int ik, id;
    int icf_sz = d*(d+1)/2;
    for (ik = 0; ik < k; ++ik)
        for (id = 0; id < d; ++id) {
            double q = icf[ik*icf_sz + id];
            pushReal8(q);
        }
    for (ik = k-1; ik > -1; --ik) {
        for (id = d-1; id > -1; --id) {
            double q;
            double qb = 0.0;
            popReal8(&q);
            qb = exp(q)*Qdiagsb[ik*d+id];
            Qdiagsb[ik*d + id] = 0.0;
            qb = qb + sum_qsb[ik];
            icfb[ik*icf_sz + id] = icfb[ik*icf_sz + id] + qb;
        }
        sum_qsb[ik] = 0.0;
    }
}

void preprocess_qs_nodiff(int d, int k, const double *icf, double *sum_qs,
        double *Qdiags) {
    int ik, id;
    int icf_sz = d*(d+1)/2;
    for (ik = 0; ik < k; ++ik) {
        sum_qs[ik] = 0.;
        for (id = 0; id < d; ++id) {
            double q = icf[ik*icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik*d + id] = exp(q);
        }
    }
}

/*
  Differentiation of Qtimesx in reverse (adjoint) mode:
   gradient     of useful results: *out *Qdiag *x *ltri
   with respect to varying inputs: *out *Qdiag *x *ltri
   Plus diff mem management of: out:in Qdiag:in x:in ltri:in
*/
void Qtimesx_b(int d, const double *Qdiag, double *Qdiagb, const double *ltri,
        double *ltrib, const double *x, double *xb, double *out, double *outb)
{
    // strictly lower triangular part
    int i, j;
    int adFrom;
    int Lparamsidx = 0;
    for (i = 0; i < d; ++i) {
        adFrom = i + 1;
        for (j = adFrom; j < d; ++j)
            Lparamsidx++;
        pushInteger4(adFrom);
    }
    for (i = d-1; i > -1; --i) {
        popInteger4(&adFrom);
        for (j = d-1; j > adFrom-1; --j) {
            --Lparamsidx;
            ltrib[Lparamsidx] = ltrib[Lparamsidx] + x[i]*outb[j];
            xb[i] = xb[i] + ltri[Lparamsidx]*outb[j];
        }
    }
    for (i = d-1; i > -1; --i) {
        Qdiagb[i] = Qdiagb[i] + x[i]*outb[i];
        xb[i] = xb[i] + Qdiag[i]*outb[i];
        outb[i] = 0.0;
    }
}

void Qtimesx_nodiff(int d, const double *Qdiag, const double *ltri, const
        double *x, double *out) {
    // strictly lower triangular part
    int i, j;
    for (i = 0; i < d; ++i)
        out[i] = Qdiag[i]*x[i];
    int Lparamsidx = 0;
    for (i = 0; i < d; ++i)
        for (j = i+1; j < d; ++j) {
            out[j] = out[j] + ltri[Lparamsidx]*x[i];
            Lparamsidx++;
        }
}

/*
  Differentiation of gmm_objective in reverse (adjoint) mode:
   gradient     of useful results: *err
   with respect to varying inputs: *err *means *icf *alphas
   RW status of diff variables: *err:in-out *means:out *icf:out
                *alphas:out
   Plus diff mem management of: err:in means:in icf:in alphas:in
*/
void gmm_objective_b(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {
    int ix, ik;
    /* TFIX */
    const double CONSTANT = -n*d*0.5*log(2*PI);
    int icf_sz = d*(d+1)/2;
    double *Qdiags;
    double *Qdiagsb;
    double result1;
    double result1b;
    int ii1;
    Qdiagsb = (double *)malloc(d*k*sizeof(double));
    for (ii1 = 0; ii1 < d*k; ++ii1)
        Qdiagsb[ii1] = 0.0;
    Qdiags = (double *)malloc(d*k*sizeof(double));
    double *sum_qs;
    double *sum_qsb;
    sum_qsb = (double *)malloc(k*sizeof(double));
    for (ii1 = 0; ii1 < k; ++ii1)
        sum_qsb[ii1] = 0.0;
    sum_qs = (double *)malloc(k*sizeof(double));
    double *xcentered;
    double *xcenteredb;
    xcenteredb = (double *)malloc(d*sizeof(double));
    for (ii1 = 0; ii1 < d; ++ii1)
        xcenteredb[ii1] = 0.0;
    xcentered = (double *)malloc(d*sizeof(double));
    double *Qxcentered;
    double *Qxcenteredb;
    Qxcenteredb = (double *)malloc(d*sizeof(double));
    for (ii1 = 0; ii1 < d; ++ii1)
        Qxcenteredb[ii1] = 0.0;
    Qxcentered = (double *)malloc(d*sizeof(double));
    double *main_term;
    double *main_termb;
    main_termb = (double *)malloc(k*sizeof(double));
    for (ii1 = 0; ii1 < k; ++ii1)
        main_termb[ii1] = 0.0;
    main_term = (double *)malloc(k*sizeof(double));
    preprocess_qs_nodiff(d, k, icf, &(sum_qs[0]), &(Qdiags[0]));
    double slse = 0.;
    double slseb = 0.0;
    for (ix = 0; ix < n; ++ix)
        for (ik = 0; ik < k; ++ik) {
            pushReal8Array(xcentered, d); /* TFIX */
            subtract_nodiff(d, &(x[ix*d]), &(means[ik*d]), &(xcentered[0]));
            pushReal8Array(Qxcentered, d); /* TFIX */
            Qtimesx_nodiff(d, &(Qdiags[ik*d]), &(icf[ik*icf_sz + d]), &(
                           xcentered[0]), &(Qxcentered[0]));
            result1 = sqnorm_nodiff(d, &(Qxcentered[0]));
            pushReal8(main_term[ik]);
            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5*result1;
        }
    double lse_alphas;
    double lse_alphasb;
    slseb = *errb;
    lse_alphasb = -(n*(*errb));
    result1b = *errb;
    *errb = 0.0;
    log_wishart_prior_b(d, k, wishart, &(sum_qs[0]), &(sum_qsb[0]), &(Qdiags[0
                        ]), &(Qdiagsb[0]), icf, icfb, result1b);
    for (ii1 = 0; ii1 < k; ii1++) /* TFIX */
        alphasb[ii1] = 0.0;
    log_sum_exp_b(k, alphas, alphasb, lse_alphasb);
    for (ii1 = 0; ii1 < d * k; ii1++) /* TFIX */
        meansb[ii1] = 0.0;
    for (ix = n-1; ix > -1; --ix) {
        result1b = slseb;
        log_sum_exp_b(k, &(main_term[0]), &(main_termb[0]), result1b);
        for (ik = k-1; ik > -1; --ik) {
            popReal8(&(main_term[ik]));
            alphasb[ik] = alphasb[ik] + main_termb[ik];
            sum_qsb[ik] = sum_qsb[ik] + main_termb[ik];
            result1b = -(0.5*main_termb[ik]);
            main_termb[ik] = 0.0;
            sqnorm_b(d, &(Qxcentered[0]), &(Qxcenteredb[0]), result1b);
            popReal8Array(Qxcentered, d); /* TFIX */
            Qtimesx_b(d, &(Qdiags[ik*d]), &(Qdiagsb[ik*d]), &(icf[ik*icf_sz +
                      d]), &(icfb[ik*icf_sz + d]), &(xcentered[0]), &(
                      xcenteredb[0]), &(Qxcentered[0]), &(Qxcenteredb[0]));
            popReal8Array(xcentered, d); /* TFIX */
            subtract_b(d, &(x[ix*d]), &(means[ik*d]), &(meansb[ik*d]), &(
                       xcentered[0]), &(xcenteredb[0]));
        }
    }
    preprocess_qs_b(d, k, icf, icfb, &(sum_qs[0]), &(sum_qsb[0]), &(Qdiags[0])
                    , &(Qdiagsb[0]));
    free(main_term);
    free(main_termb);
    free(Qxcentered);
    free(Qxcenteredb);
    free(xcentered);
    free(xcenteredb);
    free(sum_qs);
    free(sum_qsb);
    free(Qdiags);
    free(Qdiagsb);
}
}

//! Adept
#include <adept_source.h>
#include <adept.h>
#include <adept_arrays.h>
using adept::adouble;
using adept::aVector;

namespace adeptTest {

    // out = a - b
template<typename T1, typename T2, typename T3>
void subtract(int d,
    const T1* const x,
    const T2* const y,
    T3* out)
{
    for (int id = 0; id < d; id++)
    {
        out[id] = x[id] - y[id];
    }
}

template<typename T>
T sqnorm(int n, const T* const x)
{
    T res = x[0] * x[0];
    for (int i = 1; i < n; i++)
        res = res + x[i] * x[i];
    return res;
}

// This throws error on n<1
template<typename T>
T arr_max(int n, const T* const x)
{
    T m = x[0];
    for (int i = 1; i < n; i++)
    {
        if (m < x[i])
            m = x[i];
    }
    return m;
}

    template<typename T>
void gmm_objective(int d, int k, int n, const T* const alphas, const T* const means,
    const T* const icf, const double* const x, Wishart wishart, T* err);

// split of the outer loop over points
template<typename T>
void gmm_objective_split_inner(int d, int k,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err);
// other terms which are outside the loop
template<typename T>
void gmm_objective_split_other(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    Wishart wishart,
    T* err);

template<typename T>
T logsumexp(int n, const T* const x);

// p: dim
// k: number of components
// wishart parameters
// sum_qs: k sums of log diags of Qs
// Qdiags: d*k
// icf: (p*(p+1)/2)*k inverse covariance factors
template<typename T>
T log_wishart_prior(int p, int k,
    Wishart wishart,
    const T* const sum_qs,
    const T* const Qdiags,
    const T* const icf);

template<typename T>
void preprocess_qs(int d, int k,
    const T* const icf,
    T* sum_qs,
    T* Qdiags);

template<typename T>
void Qtimesx(int d,
    const T* const Qdiag,
    const T* const ltri, // strictly lower triangular part
    const T* const x,
    T* out);

////////////////////////////////////////////////////////////
//////////////////// Definitions ///////////////////////////
////////////////////////////////////////////////////////////

template<typename T>
T logsumexp(int n, const T* const x)
{
    T mx = arr_max(n, x);
    T semx = 0.;
    for (int i = 0; i < n; i++)
    {
        semx = semx + exp(x[i] - mx);
    }
    return log(semx) + mx;
}

template<typename T>
T log_wishart_prior(int p, int k,
    Wishart wishart,
    const T* const sum_qs,
    const T* const Qdiags,
    const T* const icf)
{
    int n = p + wishart.m + 1;
    int icf_sz = p * (p + 1) / 2;

    double C = n * p * (log(wishart.gamma) - 0.5 * log(2)) - log_gamma_distrib(0.5 * n, p);

    T out = 0;
    for (int ik = 0; ik < k; ik++)
    {
        T frobenius = sqnorm(p, &Qdiags[ik * p]) + sqnorm(icf_sz - p, &icf[ik * icf_sz + p]);
        out = out + 0.5 * wishart.gamma * wishart.gamma * (frobenius)
            -wishart.m * sum_qs[ik];
    }

    return out - k * C;
}

template<typename T>
void preprocess_qs(int d, int k,
    const T* const icf,
    T* sum_qs,
    T* Qdiags)
{
    int icf_sz = d * (d + 1) / 2;
    for (int ik = 0; ik < k; ik++)
    {
        sum_qs[ik] = 0.;
        for (int id = 0; id < d; id++)
        {
            T q = icf[ik * icf_sz + id];
            sum_qs[ik] = sum_qs[ik] + q;
            Qdiags[ik * d + id] = exp(q);
        }
    }
}

template<typename T>
void Qtimesx(int d,
    const T* const Qdiag,
    const T* const ltri, // strictly lower triangular part
    const T* const x,
    T* out)
{
    for (int id = 0; id < d; id++)
        out[id] = Qdiag[id] * x[id];

    int Lparamsidx = 0;
    for (int i = 0; i < d; i++)
    {
        for (int j = i + 1; j < d; j++)
        {
            out[j] = out[j] + ltri[Lparamsidx] * x[i];
            Lparamsidx++;
        }
    }
}

template<typename T>
void gmm_objective(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err)
{
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);
    int icf_sz = d * (d + 1) / 2;

    vector<T> Qdiags(d * k);
    vector<T> sum_qs(k);
    vector<T> xcentered(d);
    vector<T> Qxcentered(d);
    vector<T> main_term(k);

    preprocess_qs(d, k, icf, &sum_qs[0], &Qdiags[0]);

    T slse = 0.;
    for (int ix = 0; ix < n; ix++)
    {
        for (int ik = 0; ik < k; ik++)
        {
            subtract(d, &x[ix * d], &means[ik * d], &xcentered[0]);
            Qtimesx(d, &Qdiags[ik * d], &icf[ik * icf_sz + d], &xcentered[0], &Qxcentered[0]);

            main_term[ik] = alphas[ik] + sum_qs[ik] - 0.5 * sqnorm(d, &Qxcentered[0]);
        }
        slse = slse + logsumexp(k, &main_term[0]);
    }

    T lse_alphas = logsumexp(k, alphas);

    *err = CONSTANT + slse - n * lse_alphas;

    *err = *err + log_wishart_prior(d, k, wishart, &sum_qs[0], &Qdiags[0], icf);
}

template<typename T>
void gmm_objective_split_inner(int d, int k,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    const double* const x,
    Wishart wishart,
    T* err)
{
    int icf_sz = d * (d + 1) / 2;

    T* Ldiag = new T[d];
    T* xcentered = new T[d];
    T* mahal = new T[d];
    T* lse = new T[k];

    for (int ik = 0; ik < k; ik++)
    {
        int icf_off = ik * icf_sz;
        T sumlog_Ldiag(0.);
        for (int id = 0; id < d; id++)
        {
            sumlog_Ldiag = sumlog_Ldiag + icf[icf_off + id];
            Ldiag[id] = exp(icf[icf_off + id]);
        }

        for (int id = 0; id < d; id++)
        {
            xcentered[id] = x[id] - means[ik * d + id];
            mahal[id] = Ldiag[id] * xcentered[id];
        }
        int Lparamsidx = d;
        for (int i = 0; i < d; i++)
        {
            for (int j = i + 1; j < d; j++)
            {
                mahal[j] = mahal[j] + icf[icf_off + Lparamsidx] * xcentered[i];
                Lparamsidx++;
            }
        }
        T sqsum_mahal(0.);
        for (int id = 0; id < d; id++)
        {
            sqsum_mahal = sqsum_mahal + mahal[id] * mahal[id];
        }

        lse[ik] = alphas[ik] + sumlog_Ldiag - 0.5 * sqsum_mahal;
    }

    *err = logsumexp(k, lse);

    delete[] mahal;
    delete[] xcentered;
    delete[] Ldiag;
    delete[] lse;
}

template<typename T>
void gmm_objective_split_other(int d, int k, int n,
    const T* const alphas,
    const T* const means,
    const T* const icf,
    Wishart wishart,
    T* err)
{
    const double CONSTANT = -n * d * 0.5 * log(2 * PI);

    T lse_alphas = logsumexp(k, alphas);

    T* sum_qs = new T[k];
    T* Qdiags = new T[d * k];
    preprocess_qs(d, k, icf, sum_qs, Qdiags);
    *err = CONSTANT - n * lse_alphas + log_wishart_prior(d, k, wishart, sum_qs, Qdiags, icf);
    delete[] sum_qs;
    delete[] Qdiags;
}

};

void adept_dgmm_objective(int d, int k, int n, const double *alphas, double *
        alphasb, const double *means, double *meansb, const double *icf,
        double *icfb, const double *x, Wishart wishart, double *err, double *
        errb) {

  int icf_sz = d*(d + 1) / 2;
  int Jrows = 1;
  int Jcols = (k*(d + 1)*(d + 2)) / 2;

  adept::Stack stack;
  adouble *aalphas = new adouble[k];
  adouble *ameans = new adouble[d*k];
  adouble *aicf = new adouble[icf_sz*k];

      adept::set_values(aalphas, k, alphas);
      adept::set_values(ameans, d*k, means);
      adept::set_values(aicf, icf_sz*k, icf);

      stack.new_recording();
      adouble aerr;

      adeptTest::gmm_objective(d, k, n, aalphas, ameans,
          aicf, x, wishart, &aerr);
      aerr.set_gradient(1.); // only one J row here

      stack.compute_adjoint();

      adept::get_gradients(aalphas, k, alphasb);
      adept::get_gradients(ameans, d*k, meansb);
      adept::get_gradients(aicf, icf_sz*k, icfb);

  delete[] aalphas;
  delete[] ameans;
  delete[] aicf;
}


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>, Randi Cabezas <rcabezas@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#include <iostream>
#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE distributions test
#include <boost/test/unit_test.hpp>

#include <dpMM/dir.hpp>
#include <dpMM/cat.hpp>
#include <dpMM/niw.hpp>
#include <dpMM/iw.hpp>
#include <dpMM/normal.hpp>

#include <omp.h>

//#include <dpMM/matrix.h>
//#include <dpMM/mex.h>
//#include <dpMM/helperMEX.h>
//#include <dpMM/debugMEX.h>

#include "gsl/gsl_rng.h"
#include "gsl/gsl_randist.h"
#include "gsl/gsl_permutation.h"
#include "gsl/gsl_cdf.h"

//#include "dpmmSubclusters/common.h"
//
//#include "dpmmSubclusters/niw_sampled.h"
////#include "dpmmSubclusters/cluster_sampledT.cpp"
//#include "dpmmSubclusters/linkedList.cpp"
//
////#include "dpmmSubclusters/reduction_array.h"
////#include "dpmmSubclusters/reduction_array2.h"
//#include "dpmmSubclusters/linear_algebra.h"
////#include "dpmmSubclusters/sample_categorical.h"
////#include "dpmmSubclusters/niw_sampled.h"

#include <dpMM/vmfPriorFull.hpp>
#include <dpMM/vmf.hpp>

using std::cout;
using std::endl;

BOOST_AUTO_TEST_CASE( dir_test)
{
  cout<<"----------------------- dir ----------------------"<<endl;
  VectorXd alpha(3);
  alpha << 1.0,100.0,1.0;

  boost::mt19937 rndGen(1);
  Dir<Cat<double>,double> dir(alpha,&rndGen);
  VectorXd piPdf = dir.sample().pdf();

  BOOST_CHECK_EQUAL(piPdf.size(),alpha.size());

  cout<<"-- sampling a bit"<<endl;
  cout<<"alpha="<<alpha.transpose()<<endl;
  for(uint32_t t=0; t<10; ++t)
    cout<<"piPdf="<<dir.sample().pdf().transpose()<<endl;

  boost::mt19937 rndGen2(1);
  Dir<Catd,double> dir2(alpha,&rndGen2);
  VectorXd piPdf2 = dir2.sample().pdf();
  BOOST_CHECK_EQUAL(piPdf,piPdf2);

  Dir<Catd,double> dir3(dir);
  VectorXd piPdf3 = dir3.sample().pdf();
  BOOST_CHECK_EQUAL(piPdf.size(),piPdf3.size());

  VectorXf alphaf = alpha.cast<float>();
  Dir<Catf,float> dirf(alphaf,&rndGen);
}

BOOST_AUTO_TEST_CASE( cat_test)
{
  cout<<"----------------------- cat ----------------------"<<endl;
  VectorXd pdf(3);
  pdf << 0.5,0.25,0.25;

  boost::mt19937 rndGen(1);
  Cat<double> cat(pdf,&rndGen);
  double z1 = cat.sample();
  boost::mt19937 rndGen2(1);
  Cat<double> cat2(pdf,&rndGen2);
  double z2 = cat2.sample();
  BOOST_CHECK_EQUAL(z1,z2);

  Cat<double> cat3(cat);
  double z3 = cat3.sample();

  BOOST_CHECK_EQUAL(cat.pdf_.size(), cat3.pdf_.size());

  VectorXu z(1000);
  cat.sample(z);
  cout<<"-- sampling a bit"<<endl;
  cout<<"pdf="<<pdf.transpose()<<endl;
  cout<<"z="<<z.transpose()<<endl;

  VectorXd pdfEmp = Cat<double>(z,&rndGen).pdf();
  BOOST_CHECK_EQUAL(pdfEmp.size(), cat.pdf().size());
  cout<<"pdf from counts="<<pdfEmp.transpose()<<endl;

  VectorXf pdff = pdf.cast<float>();
  Cat<float> catf(pdff,&rndGen);
}

BOOST_AUTO_TEST_CASE(iw_test)
{
  cout<<"----------------------- iw ----------------------"<<endl;

  MatrixXd Delta(3,3);
  Delta << 1.0,0.0,0.0,
        0.0,1.0,0.0,
        0.0,0.0,1.0;
  double nu = 100.0;

  boost::mt19937 rndGen(1);
  IW<double> iw(Delta,nu,&rndGen);

  for(uint32_t t=0; t<10; ++t)
  {
    MatrixXd Sigma = iw.sample();
    cout<<"Delta=\n"<<iw.Delta_<<endl;
    cout<<"Sigma=\n"<<Sigma<<endl;
    cout<<"logPdf="<<iw.logPdf(Sigma)<<endl;
  }

}

BOOST_AUTO_TEST_CASE( gauss_test)
{
  cout<<"----------------------- gauss ----------------------"<<endl;

  MatrixXd Sigma(3,3);
  Sigma << 1.0,0.0,0.0,
        0.0,1.0,0.0,
        0.0,0.0,1.0;
  VectorXd mu(3);
  mu << 1.0,1.0,1.0;

  boost::mt19937 rndGen(1);
  Normal<double> normal(mu,Sigma,&rndGen);

  VectorXd x(3);
  x << 1.0,1.0,1.0;
  cout << normal.logPdf(x) <<endl;
  x << 1.0,1.0,0.0;
  cout << normal.logPdf(x) <<endl;
  double a = normal.logPdf(x);
  x << 0.0,1.0,1.0;
  BOOST_CHECK_EQUAL(a,normal.logPdf(x));
  cout << normal.logPdf(x) <<endl;
  x << 0.0,0.0,0.0;
  cout << normal.logPdf(x) <<endl;
  BOOST_CHECK(a>normal.logPdf(x));

  Sigma << 5.0,0.0,0.0,
        0.0,1.0,0.0,
        0.0,0.0,0.1;
  mu << 1.0,1.0,0.0;
  Normal<double> normalA(mu,Sigma,&rndGen);
  for (uint32_t i=0; i<100; ++i)
  {
    x = normalA.sample();
    cout<<"x="<<x.transpose()<<" logPdf="<<normalA.logPdf(x)<<endl;
  }

  cout<<"small variance and far away data"<<endl;
  MatrixXd Sigma2D(2,2);
  VectorXd mu2D(2);
  VectorXd x2D(2);
  Sigma2D << 1.0e-12,0.0,
        0.0,1.0e-12;
  mu2D << 0.0,0.0;
  x2D << M_PI*0.5,0.0; // 90 degree away
  Normal<double> normalS(mu2D,Sigma2D,&rndGen);
  cout<<" logPDf "<<normalS.logPdf(x2D)<<endl;

// check the sufficient statistics machinery
  VectorXd SS(13);
  SS <<                      120,
         -18.3223543707169,
          5.02645593390708,
       -0.0572298238878314,
          2.79777673293449,
        -0.767573133190367,
       0.00873808845092194,
        -0.767573133190367,
         0.210748312679343,
      -0.00239897955468155,
       0.00873808845092194,
      -0.00239897955468155,
      2.73160817868802e-05;

  double count = SS(0);
  Matrix<double,Dynamic,1> mean(3);
  if(count>0)
	  mean = SS.middleRows(1,3)/count;
  else
	  mean = Matrix<double,Dynamic,1>::Zero(3); //this should not matter since everything gets multiplied by 0 counts

  cout<<"SS "<<SS.transpose()<<endl;
  cout<<"count "<<count<<endl;
  cout<<"xSum "<<SS.middleRows(1,3).transpose()<<endl;
  cout<<"mean "<<mean.transpose()<<endl;

  double* datPtr = const_cast<double*>(&(SS.data()[(3+1)]));
  Matrix<double,Dynamic,Dynamic> scatter =
    Map<Matrix<double,Dynamic,Dynamic> >(datPtr,3,3);
  cout<<"outerSum "<<scatter<<endl;
  scatter -= (mean*mean.transpose())*count;
  cout<<"scatter "<<endl<<scatter<<endl;

  Sigma << 0.0314229552991203,      -0.00429368708799009,
        -0.000853036009326598,
        -0.00429368708799009 ,       0.0292337474264411,
        -5.82413154649128e-05,
        -0.000853036009326598 ,    -5.82413154649128e-05,
        0.000564454579403038;
  cout<<"Sigma "<<endl<<Sigma<<endl;
  mu <<  -0.04866953611, -0.02217239974,-0.02938455595 ;
  cout<<"mu "<<mu.transpose()<<endl;
  Normal<double> g(mu, Sigma,&rndGen);
  cout<<"logPdf      : "<< g.logPdf(scatter,mean,count)<<endl;
  cout<<"logPdfSlower: "<< g.logPdfSlower(scatter,mean,count)<<endl;
  cout<<"logDetSigma : "<< g.logDetSigma()<<endl;

}

BOOST_AUTO_TEST_CASE( vMF_test)
{
  cout<<"----------------------- vMF ----------------------"<<endl;

  VectorXd m0(3);
  m0 << 0.0,0.0,1.0;
  double t0 = 0.01;
  double a0 = 2.0;
  double b0 = 1.7;

  boost::mt19937 rndGen(1);
  vMFpriorFull<double> vMFprior(m0,t0,a0,b0,&rndGen);

  MatrixXd x(3,100);
  for(uint32_t i=0; i<100; ++i)
    x(0,i) = 1.0;
  VectorXu z(100); z.fill(0);
  uint32_t k = 0;

  vMFprior.getSufficientStatistics(x,z,k);

  for(uint32_t i = 0; i < 10;++i)
  {
    cout<<" -- "<<endl;
    vMF<double> vmf =  vMFprior.sample();
    vmf.print();
//    for(uint32_t j=0;j<10; ++j)
//    {
     vmf = vMFprior.sampleFromPosterior(vmf);
    vmf.print();
//      cout<<" j="<<j<<" logPdf = "<<vMFprior.logPdf
//    }
  }

}

//BOOST_AUTO_TEST_CASE( niw_test) { cout<<"----------------------- niw
//----------------------"<<endl;
//
//  uint32_t D =3;
//  MatrixXd Delta(D,D);
//  Delta << 1.0,0.0,0.0,
//        0.0,1.0,0.0,
//        0.0,0.0,1.0;
//  VectorXd theta(D);
//  theta << 1.0,1.0,1.0;
//  double nu = 100.0;
//  double kappa = 100.0;
//
//  boost::mt19937 rndGen(1);
//  NIW<double> niw(Delta,theta,nu,kappa,&rndGen);
//
//  for(uint32_t t=0; t<10; ++t)
//  {
//    Normal<double> Norm = niw.sample();
//	Norm.print();
//    cout<<"logPdf="<<niw.logPdf(Norm)<<endl;
//  }
//  cout<<" comparing NIW against Jasons implementation ------"<<endl;
//  niw.print();
//
//  uint32_t N=80;
//  uint32_t K=1;
//  VectorXu z(N);
//  shared_ptr<MatrixXd> spx(new MatrixXd(D,N));
//  cout<<"true mus:"<<endl<<sampleClusters<double>(*spx, z, K)<<endl;
//
//  double count = N;
//  MatrixXd Outer(D,D); Outer.setZero();
//  VectorXd sum(D); sum.setZero();
//  for(uint32_t i=0; i<N; ++i)
//  {
//    sum += spx->col(i);
//    Outer += spx->col(i)*spx->col(i).transpose();
//  }
//  cout<<"Outer"<<endl<<Outer<<endl;
//  cout<<"sum"<<endl<<sum<<endl;
//  cout<<"1/N*sum*sum.T"<<endl<<sum*sum.transpose()/count<<endl;
//  cout<<"sum*sum.T"<<endl<<sum*sum.transpose()<<endl;
//  MatrixXd Scatter = Outer - sum*sum.transpose() / count;
//  VectorXd mean = sum/count;
//  cout<<"Scatter"<<endl<<Scatter<<endl;
//
//  niw.scatter() = Scatter;
//  niw.mean() = mean;
//  niw.count() = count;
//
//  cout<<"--- Julians posterior:"<<endl;
//  niw.posterior().print();
//
//  cout<<"--- Julians posterior direct:"<<endl;
//  niw.posterior(*spx,z,0).print();
//
//
//  MatrixXd DeltaOverNu(Delta);
//  DeltaOverNu /= nu;
//  VectorXd thetaOverKappa(theta);
////  thetaOverKappa /= kappa;
//  niw_sampled niwJason(3,kappa,nu,thetaOverKappa.data(),DeltaOverNu.data());
//  niwJason.set_stats(count,sum.data(),Outer.data());
//  niwJason.update_posteriors();
//
//  Map<MatrixXd> DeltaPost(niwJason.Delta,D,D);
//  Map<VectorXd> thetaPost(niwJason.theta,D);
//  cout<<"--- Jason posterior:"<<endl;
//  cout<<"nu="<<niwJason.nu<<" kappa="<<niwJason.kappa<<endl;
//  cout<<"delta"<<endl<<DeltaPost<<endl;
//  cout<<"theta "<<thetaPost.transpose()<<endl;
//  cout<<"--- Jasons adapted to julians:"<<endl;
//  cout<<"delta"<<endl<<DeltaPost*niwJason.nu<<endl;
//  cout<<"theta "<<thetaPost.transpose()*niwJason.kappa<<endl;
//
//  cout<<" ----- marginal probability of data under NIW"<<endl;
//  cout<<" Julian: "<<niw.logPdfMarginalized()<<endl;
//  cout<<" Jason:  "<<niwJason.data_loglikelihood_marginalized()<<endl;
//
//  BOOST_CHECK(fabs(niw.logPdfMarginalized() - niwJason.data_loglikelihood_marginalized())<1e-3);
//
//  cout<<"------------------ merge test -----------------------"<<endl;
//
//  NIW<double> niwB(Delta,theta,nu,kappa,&rndGen);
//
//  double countB = 1000;
//  MatrixXd OuterB(D,D);
//  OuterB<< 10,0,0,
//          0,100,0,
//          0,0,10;
//  VectorXd sumB(D);
//    sumB << 1,10,0;
//  MatrixXd ScatterB = OuterB ;//- count*sum*sum.transpose();
//  VectorXd meanB = sumB/countB;
//  cout<<"Scatter"<<endl<<ScatterB<<endl;
//
//  niwB.scatter() = ScatterB;
//  niwB.mean() = meanB;
//  niwB.count() = countB;
//
//  NIW<double>* niwBmerged = niwB.merge(niw);
//  cout<<"--- Julians posterior after merge:"<<endl;
//  niwBmerged->posterior().print();
//
//  MatrixXd DeltaOverNuB(Delta);
//  DeltaOverNuB /= nu;
//  VectorXd thetaOverKappaB(theta);
////  thetaOverKappa /= kappa;
//  niw_sampled niwJasonB(3,kappa,nu,thetaOverKappaB.data(),DeltaOverNuB.data());
//  niwJasonB.set_stats(countB,sumB.data(),OuterB.data());
//  niwJasonB.update_posteriors();
//  Map<MatrixXd> DeltaPostB(niwJasonB.Delta,D,D);
//  Map<VectorXd> thetaPostB(niwJasonB.theta,D);
//
//  niwJasonB.merge_with(niwJason, false);
////  niwJasonB.update_posteriors();
//  cout<<"--- Jason posterior after merge:"<<endl;
//  cout<<"nu="<<niwJasonB.nu<<" kappa="<<niwJasonB.kappa<<endl;
//  cout<<"delta"<<endl<<DeltaPostB<<endl;
//  cout<<"theta "<<thetaPostB.transpose()<<endl;
//  cout<<"--- Jasons adapted to julians:"<<endl;
//  cout<<"delta"<<endl<<DeltaPostB*niwJasonB.nu<<endl;
//  cout<<"theta "<<thetaPostB.transpose()*niwJasonB.kappa<<endl;
//
//  cout<<" ------------------- sampling ----------------------------"<<endl;
//  for(uint32_t t=0; t<5; ++t)
//  {
//    cout<<"Jason:"<<endl;
//    niwJasonB.sample();
//    Map<MatrixXd> jasonCov(niwJasonB.param.cov,D,D);
//    Map<VectorXd> jasonMean(niwJasonB.param.mean,D);
//    cout<<jasonMean.transpose()<<endl;
//    cout<<jasonCov<<endl;
//    cout<<"Julian:"<<endl;
//    Normal<double> normB = niwB.posterior().sample();
//    normB.print();
//
//    //cout<<normB.mu_.transpose()<<endl;
//    //cout<<normB.Sigma_<<endl;
//  }
//}


#define BOOST_TEST_MODULE TEST_PARAMS

#include <dynet/dynet.h>
#include <dynet/expr.h>
#include <dynet/model.h>
#include <boost/test/unit_test.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>
#include <iostream>
#include <fstream>

#include <stdexcept>

using namespace dynet;
using namespace dynet::expr;
using namespace std;

struct ParamsTest {
    ParamsTest() {
        // initialize if necessary
        if (default_device == nullptr) {
            for (auto x : {"ParamsTest", "--dynet-mem", "512"}) {
                av.push_back(strdup(x));
            }
            char **argv = &av[0];
            int argc = av.size();
            dynet::initialize(argc, argv);
        }
        gain = 2.0;
        epsilon = 1e-6;
        d = dynet::Dim({10, 10});
    }
    ~ParamsTest() {
        for (auto x : av) free(x);
    }

    float gain, epsilon;
    dynet::Dim d;
    std::vector<char*> av;
};

// define the test suite
BOOST_FIXTURE_TEST_SUITE(params_test, ParamsTest);

BOOST_AUTO_TEST_CASE( init_saxe ) {
    dynet::Model mod;
    // Random orthogonal matrix scaled by gain
    dynet::Parameter saxe_p = mod.add_parameters({10, 10}, ParameterInitSaxe(gain));
    // gain^2 * identity matrix
    dynet::Parameter identity_p = mod.add_parameters({10, 10}, ParameterInitIdentity());
    // Initialize graph
    dynet::ComputationGraph cg;
    dynet::Expression saxe = dynet::parameter(cg, saxe_p);
    dynet::Expression identity = dynet::parameter(cg, identity_p);
    // check that the matrix is indeed orthogonal
    dynet::Expression diff_expr_left = dynet::squared_norm(dynet::transpose(saxe) * saxe - (gain * gain) * identity);
    dynet::Expression diff_expr_right = dynet::squared_norm(saxe * dynet::transpose(saxe) - (gain * gain) * identity);
    float diff = dynet::as_scalar(cg.forward((diff_expr_left + diff_expr_right) / 2.0));
    // Leave a margin of error of epsilon=10^-6
    BOOST_CHECK_LT(diff, epsilon);
}

BOOST_AUTO_TEST_SUITE_END()


// ROS includes
#include "ros/ros.h"
#include "ros/assert.h"
#include "dynamic_reconfigure/server.h"
#include "create_driver/vicon_driver.h"
#include "geometry_msgs/Twist.h"

// Library includes
#include <string>
#include <vector>
#include <map>
#include <Eigen/Core>
#include <Eigen/LU>

// Local Package includes
#include "to_velocity_tracking/ControlVariablesConfig.h"
#include "to_velocity_tracking/ControlMsgs.h"

using namespace std;
using namespace Eigen;

// Constants
const double PI = 3.1415926535;
const double wmax = 8;

// Dynamic reconfigure variables
double k = 0.5; // linear velocity gain
double ktheta = 5.0; // angular velocity gain
bool power = true;
double epsilon = 0.25;
bool u_switch = true;

void reconfigureCallback(to_velocity_tracking::ControlVariablesConfig &config, uint32_t level) {
	k = config.k;
	ktheta = config.ktheta;
	power = config.power;
	epsilon = config.epsilon;
	u_switch = config.u_switch;
}

int main(int argc, char **argv)
{
	// ROS Initalization
	ros::init(argc, argv, "control");

	ros::NodeHandle n;
	ros::NodeHandle private_n("~");

	// ROS Parameters

	// List of robot names
	vector<string> robot_names;
	XmlRpc::XmlRpcValue robot_list;
	private_n.getParam("robot_list", robot_list);
	ROS_ASSERT(robot_list.getType() == XmlRpc::XmlRpcValue::TypeArray);

	for (int i = 0; i < robot_list.size(); i++)
	{
		ROS_ASSERT(robot_list[i].getType() == XmlRpc::XmlRpcValue::TypeString);
		robot_names.push_back(static_cast<string>(robot_list[i]));
	}

	// Number of robots
	const int num_robots = robot_names.size();

	// Robot name
	string robot_name;
	private_n.param<std::string>("robotname", robot_name, "defaultname");

	int index = find(robot_names.begin(), robot_names.end(), robot_name) - robot_names.begin();

	// ROS Dynamic Reconfigure

	dynamic_reconfigure::Server<to_velocity_tracking::ControlVariablesConfig> server;
  	dynamic_reconfigure::Server<to_velocity_tracking::ControlVariablesConfig>::CallbackType f;
  	f = boost::bind(&reconfigureCallback, _1, _2);
 	server.setCallback(f);

	// ROS Subscribers
	map<string, create_driver::ViconStream> vicon;
	vector<string>::iterator name_it;
	vector<ros::Subscriber> sub;
	for (name_it = robot_names.begin(); name_it != robot_names.end(); name_it++)
	{
		vicon.insert(pair<string, create_driver::ViconStream>(*name_it, create_driver::ViconStream()));
		sub.push_back(n.subscribe("/" + *name_it + "/tf", 10, &create_driver::ViconStream::callback, &vicon[*name_it]));
	}

	// ROS Publishers
	ros::Publisher vel = n.advertise<geometry_msgs::Twist>("cmd_vel", 10);

	ros::Publisher dia = n.advertise<to_velocity_tracking::ControlMsgs>("dia", 10);

	// ROS loop
	ros::Rate loop_rate(250); // 250 Hz

	while (ros::ok())
	{
		// Retrieve Vicon Data
		ros::spinOnce();

		/* An instance of this code is run on each robot in the formation.
		 *
		 * You can access each robot's position and oreintation through the following statements:
		 * vicon[robot_names[i]].x()
		 * vicon[robot_names[i]].y()
		 * vicon[robot_names[i]].theta()
		 *
		 * The index of the robot running the instance of the code is contain in the "index"
		 * variable. For example, if you wanted to get the distance of this robot from all of its
		 * neighbors you could run the following:
		 * double dist[num_robots];
		 * for (int i = 0; i < num_robots; i++) {
		 * 		dist[i] = sqrt(pow(vicon[robot_names[i]].x() - vicon[robot_names[index]].x(),2)
		 * 				     + pow(vicon[robot_names[i]].y() - vicon[robot_names[index]].y(),2));
		 * }
		 */
		double ux = 0;
		double uy = 0;

		if (u_switch) {
			ux = 1;
			uy = 1;
		} else {
			ux = -1;
			uy = -1;
		}

		double thetad = atan2(uy, ux);
		double theta = vicon[robot_name].theta();
		double etheta = theta - thetad;

		double v = k*1000*cos(etheta)*sqrt(pow(ux,2) + pow(uy,2)); // linear velocity output
		double w = -ktheta*etheta; // angular velocity output

		if (fabs(etheta) > epsilon){
			v = 0;
			w = (etheta > 0) ? -wmax : wmax;
		}

		if ( power ) {
		// Send wheel velocities to driver
			geometry_msgs::Twist msg;
			msg.linear.x = v;
			msg.linear.y = 0;
			msg.linear.z = 0;
			msg.angular.x = 0;
			msg.angular.y = 0;
			msg.angular.z = w;
			vel.publish(msg);
		} else {
			geometry_msgs::Twist msg;
			msg.linear.x = 0;
			msg.linear.y = 0;
			msg.linear.z = 0;
			msg.angular.x = 0;
			msg.angular.y = 0;
			msg.angular.z = 0;
			vel.publish(msg);
		}

		to_velocity_tracking::ControlMsgs dmsg;
		dmsg.theta = theta;

		dia.publish(dmsg);

		loop_rate.sleep();
	}

	return 0;
}


//===----------------------------------------------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is dual licensed under the MIT and the University of Illinois Open
// Source Licenses. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//  Adaptation to Boost of the libcxx
//  Copyright 2010 Vicente J. Botet Escriba
//  Distributed under the Boost Software License, Version 1.0.
//  See http://www.boost.org/LICENSE_1_0.txt

// duration
// Test nested types

// typedef Rep rep;
// typedef Period period;

#include <boost/chrono/duration.hpp>
#include <boost/type_traits.hpp>
#if !defined(BOOST_NO_STATIC_ASSERT)
#define NOTHING ""
#endif

typedef boost::chrono::duration<long, boost::ratio<3, 2> > D;
BOOST_CHRONO_STATIC_ASSERT((boost::is_same<D::rep, long>::value), NOTHING, ());
BOOST_CHRONO_STATIC_ASSERT((boost::is_same<D::period, boost::ratio<3, 2> >::value), NOTHING, ());


#ifndef CUAUV_SIM_ENTITY_H
#define CUAUV_SIM_ENTITY_H

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

namespace cuauv {
namespace fishbowl {

typedef Eigen::DiagonalMatrix<double, 3> inertia_tensor;

/**
 * Entity represents an entity in a World.
 */
class entity {
public:
    /**
     * Constructs a new entity.
     *
     * @param m Mass.
     * @param r Radius.
     * @param I Inertia tensor, in the body frame.
     * @param btom_rq Body to model frame rotation quaternion. Undefined behavior if not normalized.
     */
    entity(double m, double r, const inertia_tensor& I, const Eigen::Quaterniond& btom_rq);

    double get_m() const; //!< Mass.
    double get_r() const; //!< Radius.
    inertia_tensor get_I() const; //!< Inertia tensor, in the body frame.

    Eigen::Vector3d x { 0, 0, 0 }; //!< Position.
    Eigen::Vector3d v { 0, 0, 0 }; //!< Translational velocity, in the world frame.
    Eigen::Vector3d a { 0, 0, 0 }; //!< Translational acceleration in the previous timestep, in the world frame.

    Eigen::Vector3d xp { 0, 0, 0 }; //!< Position, in the previous timestep. Initially 0.

    /**
     * Orientation, of the body relative to the world frame.
     * To convert to the orientation of the model relative to the world frame, use
     * <tt>q * body_rm_inv</tt>
     * Behavior is undefined if q is not a unit quaternion.
     */
    Eigen::Quaterniond q { 1, 0, 0, 0 };
    Eigen::Vector3d w { 0, 0, 0 }; //!< Angular velocity, in the body frame.
    Eigen::Vector3d t { 0, 0, 0 }; //<! Torque, in the body frame, in the previous timestep.

    bool corporeal = true; //!< True iff forces and engines should affect this entity.

    /**
     * The rotation matrix from the body frame to the model frame.
     * We precompute the matrix because it is more efficient for rotation
     * operations, according to Eigen's documentation.
     * @see http://eigen.tuxfamily.org/dox/classEigen_1_1Quaternion.html
     */
    Eigen::Matrix3d get_btom_rm() const;
    Eigen::Matrix3d get_mtob_rm() const;

    // Utility functions.

    inertia_tensor get_Ir() const; //!< The reciprocal of the inertia tensor, in the body frame.
    Eigen::Quaterniond get_model_q() const; //!< Orientation, of the model relative to the world frame.

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW

private:
    double m;
    double r;
    inertia_tensor I;
    inertia_tensor Ir;
    Eigen::Quaterniond btom_rq;
    Eigen::Matrix3d btom_rm;
    Eigen::Matrix3d mtob_rm;
    Eigen::Quaterniond model_q;
};

} // namespace fishbowl
} // namespace cuauv

#endif // CUAUV_SIM_ENTITY_H


#ifndef _CNF_SIMPLE_HPP
#define _CNF_SIMPLE_HPP

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <gmpxx.h>
#include <string>
#include <fstream>
#include <cstdlib>
#include <array>
#include <tuple>
#include <ctime>
#include <boost/multiprecision/gmp.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/tokenizer.hpp>
#include <boost/algorithm/string.hpp>

// debug parameter
#define DEBUG 1

// type definitions for matrix, sparse matrix and sparse vector
typedef std::set<boost::multiprecision::mpz_int> s_vector;
typedef std::map<int, s_vector> s_matrix;
typedef std::vector<std::vector<int>> matrix;

// node structure for the Tree structure
struct Node {
  int              idx;   // index (line of CNF matrix)
  int              value; // value (0 or 1)
  std::vector<int> alpha; // vector of alpha_i for the node
  Node            *left;  // left child
  Node            *right; // right child
};

// Tree structure used by compute_all_h_fasy
struct Tree {
  Node *root;                  // root node
  std::string id_cnf;          // the id of the problem (typically the filename)

  Tree(std::string an_id);
  ~Tree();
  void destroy_tree(Node *leaf);
  void create_tree(s_matrix &H, int n_litterals);
  int create_children(Node *cur_node, s_matrix &H, int n_litterals);
};

// split a string into a 2-uple (id, s_vector)
std::tuple<int, s_vector> split(const std::string &msg, const std::string &separators);

// removing a file
void remove_file(std::string filename);

// create a cnf matrix from an input file
matrix read_input(std::string filename);

// create a sparse vector from an input file
std::tuple<int, s_vector> read_input_vector(std::string filename);

// create a sparse matrix from an input  file
s_matrix read_input_matrix(std::string filename);

// compute the matrics F, G and H for a clause of a cnf matrix
std::tuple<int, s_matrix, s_matrix, s_matrix> compute_FGH(matrix const &matrix_cnf, int clause);

// given the indexes of 3 litterals, compute the associated index
boost::multiprecision::mpz_int compute_index3(int index_1, int index_2, int index_3, int n_litterals);

// compute a given row of matrix D telling wich alpha_k are present the row (OK!)
std::vector<int> D_fnct(boost::multiprecision::mpz_int row, int n_litterals);

// add an element to a binary matrix/vector
void add_elem_vector(s_vector &v, boost::multiprecision::mpz_int val);
void add_elem_matrix(s_matrix &matrix, int row, boost::multiprecision::mpz_int col);

// printing sparse matrices and (sparse) vectors
void print_matrix_s(s_matrix const &M);
void print_vector_s(s_vector const &v);
void print_vector(std::vector<int> const &v, unsigned int start = 0);

// saving a sparse matrix/vector in a file
void save_matrix_s(s_matrix const &M, std::string filename);
void save_vector(std::vector<int> const &v, unsigned int start = 0, std::string filename = "sol");

// simplifies a binary matrix/vector
s_matrix simplify_matrix_sparse(s_matrix &matrix_in, int n_litterals, int cur_litteral);
s_vector simplify_vector_sparse(s_matrix &matrix_in, s_vector &vector_in, int cur_row, int n_litterals);

// create a matrix from a vector
s_matrix generate_matrix(s_vector &v, int row);

// remove a row in the matrix M
void remove_row_in_M(s_matrix &M, int row);

// copy a vector in a given row of M
void copy_vector_in_M(s_matrix &M, s_vector &v, int row);

// copy a given row of M in row
s_vector copy_row_in_v(s_matrix &M, int row);

// check if a sparse vector is empty
bool is_null(s_vector const &v);

// check if 2 vectors are equals component-wise
bool is_equal(std::vector<int> const &v_1, const std::vector<int> &v_2);

// return the number of ones in a given row of a sparse matrix
size_t count_ones_vector(s_matrix &matrix_in, int cur_row);

// compute the vectors I_0 and I_1 from I
std::tuple<s_vector, s_vector> compute_I0_I1(s_vector &I, int n_litterals, int alpha_k);

// from a given binary vector and the number of litterals, return the index in the D matrix
boost::multiprecision::mpz_int compute_index_vect2(std::vector<int> const &v, int n_litterals);

// adding two binary vectors
void add(s_vector &v_1, s_vector &v_2);
s_vector add_out(s_vector const &v_1, s_vector const &v_2);

// component wise binary vector operation (if v_2[1] then out[i] = 1 else out[i] = v_1[i])
std::vector<int> mult(std::vector<int> const &v_1, std::vector<int> const &v_2);

// adding two binary matrices
s_matrix add_matrix(s_matrix const &M_1, s_matrix const &M_2);

// multiply a vector with another one using the matrix D representing the binary vectors
s_vector mult_mat_vec(s_vector &v1, s_vector &v2, int n_litterals);

// return the maximum k of the alpha_i in a vector v
int get_highest_alpha_k(s_vector &v, int n_litterals);

// merges 2 clauses using the matrics H, F and G.
bool merge(s_matrix &H1, s_matrix &H2, s_matrix &F1, s_matrix &F2, s_matrix &G1,
           s_matrix &G2, std::vector<int> &row_constraint, int id_new_constraint,
           int *lev_rec, int n_litterals);

// given an index and the number of litterals, generate the associated binary vector alpha
std::vector<int> generate_alpha(int n_litterals, boost::multiprecision::mpz_int x);

// returns the output of H for a given vector alpha
std::vector<int> h_vector(s_matrix &H, int n_litterals, std::vector<int> &alpha);

// compute the value of a row of H when we replace the alpha_i by some values, i.e. for a given vector alpha
int compute_h_i(std::vector<int> &v_alpha, s_vector &row_H, int n_litterals);

// compute every solution of a matrix H
void compute_all_h(s_matrix &H, int n_litterals);
// efficiently compute every solution of a matrix H using a tree structure
void compute_all_h_fast(s_matrix &H, int n_litterals, std::string filename);

// deterniming a
void merge_matrices(std::string filename);

// mergin a CNF matrix H with a line (vector) of another CNF matrix
void merge_matrix_vector(std::string filename_matrix, std::string filename_vector);

// compute the max litteral in the remaining clauses
std::vector<int> compute_max_lit_remaining_clauses(matrix &input_matrix);

// print the alpha in a row of a matrix
void print_alphas(s_matrix &matrix_in, int cur_row, int n_litterals);

#endif


/** \file   cmr_t1_mapping.cpp
    \brief  Implement CMR T1 mapping for 2D acquisition
    \author Hui Xue
*/

#include "cmr_t1_mapping.h"
#include "log.h"

#include "hoNDArray_reductions.h"
#include "hoNDArray_elemwise.h"
#include "hoNDArray_math.h"

#include "simplexLagariaSolver.h"
#include "twoParaExpRecoveryOperator.h"
#include "curveFittingCostFunction.h"

#include <boost/math/special_functions/sign.hpp>

namespace Gadgetron {

template <typename T>
CmrT1SRMapping<T>::CmrT1SRMapping() : BaseClass()
{
    max_iter_ = 150;
    max_fun_eval_ = 1000;
    thres_fun_ = 1e-4;

    // maximal allowed T1
    max_map_value_ = 2500;
}

template <typename T>
CmrT1SRMapping<T>::~CmrT1SRMapping()
{
}

template <typename T>
void CmrT1SRMapping<T>::get_initial_guess(const VectorType& ti, const VectorType& yi, VectorType& guess)
{
    if (guess.size() != this->get_num_of_paras())
    {
        guess.resize(this->get_num_of_paras(), 0);
    }

    guess[0] = 500;
    guess[1] = 1200;

    // A
    if(!yi.empty()) guess[0] = *std::max_element(yi.begin(), yi.end());

    // T1
    if (!ti.empty()) guess[1] = ti[ti.size() / 2];
}

template <typename T>
void CmrT1SRMapping<T>::compute_map(const VectorType& ti, const VectorType& yi, const VectorType& guess, VectorType& bi, T& map_v)
{
    try
    {
        bi = guess;
        map_v = 0;

        typedef Gadgetron::twoParaExpRecoveryOperator< std::vector<T> > SignalType;
        typedef Gadgetron::leastSquareErrorCostFunction< std::vector<T> > CostType;

        // define solver
        Gadgetron::simplexLagariaSolver< VectorType, SignalType, CostType > solver;

        // define signal model
        SignalType t1_sr;

        // define cost function
        CostType lse;

        solver.signal_model_ = &t1_sr;
        solver.cf_ = &lse;

        solver.max_iter_ = max_iter_;
        solver.max_fun_eval_ = max_fun_eval_;
        solver.thres_fun_ = thres_fun_;

        solver.x_ = ti;
        solver.y_ = yi;

        solver.solve(bi, guess);

        if (bi[0] > 0 && bi[1] > 0)
        {
            map_v = bi[1];
            if (map_v >= max_map_value_) map_v = hole_marking_value_;
            if (map_v <= min_map_value_) map_v = hole_marking_value_;
        }
    }
    catch (...)
    {
        GADGET_THROW("Exceptions happened in CmrT1SRMapping<T>::compute_map(...) ... ");
    }
}

template <typename T>
void CmrT1SRMapping<T>::compute_sd(const VectorType& ti, const VectorType& yi, const VectorType& bi, VectorType& sd, T& map_sd)
{
    try
    {
        sd.clear();
        sd.resize(bi.size(), 0);

        map_sd = 0;

        typedef Gadgetron::twoParaExpRecoveryOperator< std::vector<T> > SignalType;
        SignalType t1_sr;

        // compute fitting values
        VectorType y;
        t1_sr.magnitude(ti, bi, y);

        // compute residual
        VectorType res(y), abs_res(y);

        size_t num = ti.size();
        size_t N = this->get_num_of_paras();

        size_t n;
        for (n = 0; n < num; n++)
        {
            res[n] = y[n] - yi[n];
            abs_res[n] = std::abs(res[n]);
        }

        hoNDArray<T> grad;
        grad.create(N, num);
        Gadgetron::clear(grad);

        VectorType gradVec(N);
        for (n = 0; n < num; n++)
        {
            t1_sr.gradient(ti[n], bi, gradVec);
            memcpy(grad.begin() + n*N, &gradVec[0], sizeof(T)*N);
        }

        GADGET_CATCH_THROW(this->compute_sd_impl(ti, yi, bi, abs_res, grad, sd));

        map_sd = sd[1];
        if (map_sd > max_map_value_) map_sd = this->hole_marking_value_;
    }
    catch (...)
    {
        GADGET_THROW("Exceptions happened in CmrT1SRMapping<T>::compute_map(...) ... ");
    }
}

template <typename T>
size_t CmrT1SRMapping<T>::get_num_of_paras() const
{
    return 2; // A and T1
}

// ------------------------------------------------------------
// Instantiation
// ------------------------------------------------------------

template class EXPORTCMR CmrT1SRMapping< float >;

}


//Author: Dr. Shantanu Shahane
#ifndef __coefficient_computations_H
#define __coefficient_computations_H
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <time.h>
#include <float.h>
#include <string.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include "class.hpp"
#include <unistd.h>
#include <limits.h>
#include <Eigen/Dense>
#include <unsupported/Eigen/SparseExtra>
#include <Eigen/SparseLU>
#include <Eigen/OrderingMethods>
#include <Eigen/Eigenvalues>
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Spectra/GenEigsSolver.h>
#include <Spectra/MatOp/SparseGenMatProd.h>
#include <Spectra/GenEigsRealShiftSolver.h>
#include <Spectra/MatOp/SparseGenRealShiftSolve.h>
using namespace std;

void shifting_scaling(vector<double> &vert, double *scale, int dim);

void shifting_scaling(double *xyz_interp, vector<double> &vert, double *scale, int dim);

void calc_PHS_RBF_grad_laplace_single_vert_A(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &A, double *scale);

void calc_PHS_RBF_grad_laplace_single_vert_grad_x_rhs(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &rhs, double *scale, vector<int> &central_vert_list);

void calc_PHS_RBF_interp_single_vert_rhs(double *xyz_interp, vector<double> &vert, PARAMETERS &parameters, Eigen::VectorXd &rhs);

void calc_PHS_RBF_grad_laplace_single_vert_grad_y_rhs(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &rhs, double *scale, vector<int> &central_vert_list);

void calc_PHS_RBF_grad_laplace_single_vert_grad_z_rhs(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &rhs, double *scale, vector<int> &central_vert_list);

void calc_PHS_RBF_grad_laplace_single_vert_laplacian_rhs(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &rhs, double *scale, vector<int> &central_vert_list);

double calc_PHS_RBF_grad_laplace_single_vert(vector<double> &vert, PARAMETERS &parameters, Eigen::MatrixXd &laplacian, Eigen::MatrixXd &grad_x, Eigen::MatrixXd &grad_y, Eigen::MatrixXd &grad_z, double *scale, vector<int> &central_vert_list);

void calc_cloud_points_slow(vector<vector<int>> &cloud_points, vector<double> &xyz_probe, POINTS &points, PARAMETERS &parameters);

void calc_cloud_points_slow_periodic_bc(vector<vector<int>> &cloud_points, vector<double> &xyz_probe, POINTS &points, PARAMETERS &parameters, vector<vector<int>> &periodic_bc_section);

Eigen::SparseMatrix<double, Eigen::RowMajor> calc_interp_matrix(vector<double> &xyz_probe, POINTS &points, PARAMETERS &parameters);

#endif


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

// With contributions from Cornelius Steinhardt

#ifndef MTL_MATRIX_GIVENS_INCLUDE
#define MTL_MATRIX_GIVENS_INCLUDE

#include <cmath>
#include <boost/numeric/mtl/utility/exception.hpp>
#include <boost/numeric/mtl/utility/irange.hpp>
#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/operation/householder.hpp>
#include <boost/numeric/mtl/operation/rank_one_update.hpp>
#include <boost/numeric/mtl/operation/trans.hpp>

namespace mtl { namespace matrix {

/// Given's transformator
/** Requires Hessenberg form, i.e. for transformations near the diagonal.
    For general form use qr_givens.
    \sa qr_givens. **/
template <typename Matrix>
class givens
{
    typedef typename Collection<Matrix>::value_type   value_type;
    typedef typename Collection<Matrix>::size_type    size_type;

  public:
    /// Re-set the rotation parameters \p a and \p b
    void set_rotation(value_type a, value_type b)
    {
	using std::abs;
	value_type zero= math::zero(a), one= math::one(b), t;

	if ( b == zero ) {
	    c= one; s= zero;
	} else if ( abs(b) > abs(a) ) {
	    t= -a / b;
	    s= one / sqrt(one + t*t);
	    c= s * t;
	} else {
	    t= -b / a;
	    c= one / sqrt(one + t*t);
	    s= c * t;
	}
	G= c, s,
	  -s, c;
    }

    /// Constructor takes %matrix \p H to be transformed and the rotation parameters \p a and \p b
    givens(Matrix& H, value_type a, value_type b) : H(H), G(2, 2)
    {	set_rotation(a, b);    }

    /// Given's transformation of \p H with \p G regarding column \p k
    Matrix& trafo(const Matrix& G, size_type k)
    {
	    irange r(k,k+2);
	    // trans(H[r][ind])*= G; H[ind][r]*= G; // most compact form but does not work yet

	    Matrix col_block(H[r][iall]), col_perm(trans(G) * col_block);
	    H[r][iall]= col_perm;
	    Matrix row_perm(H[iall][r] * G);
	    H[iall][r]= row_perm;

	    return H;
    }

    /// Given's transformation of \p H regarding column \p k
    Matrix& trafo(size_type k)
    {
	return trafo(G, k);
    }

  private:
    Matrix&    H, G;
    value_type c, s;
};

}// namespace matrix

namespace vector {

/// Given's transformator on %vector (swap a*line(k) with b*line(k+1) )
template <typename Vector>
class givens
{
    typedef typename Collection<Vector>::value_type   value_type;
    typedef typename Collection<Vector>::size_type    size_type;

  public:
    /// Constructor takes %vector \p v to be transformed and the rotation parameters \p a and \p b
    givens(Vector& v, value_type a, value_type b) : v(v), a(a), b(b)
    {  }

    /// Given's transformation of \p v with \p a and \p b regarding column \p k
    Vector& trafo(size_type k)
    {
	    value_type w1(0), w2(0);
	    w1= a*v[k] - b*v[k+1]; //given's rotation on solution
            w2= b*v[k] + a*v[k+1]; //rotation on vector
            v[k]= w1;
            v[k+1]= w2;

	    return v;
    }

  private:
    Vector&    v;
    value_type a, b;
};

}// namespace vector

} // namespace mtl

#endif // MTL_MATRIX_GIVENS_INCLUDE


#pragma once

#include <Eigen/Core>

#include <array>

namespace wmtk {
double harmonic_tet_energy(const std::array<double, 12>& T);
void harmonic_tet_jacobian(const std::array<double, 12>& T, Eigen::Vector3d& result_0);
} // namespace wmtk


// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#include <algorithm>
#include <iostream>

#include <Eigen/Dense>
#include <cmath>

#include "Open3D/Camera/PinholeCameraIntrinsic.h"
#include "Open3D/Geometry/Image.h"

namespace open3d {
namespace geometry {

std::shared_ptr<Image> Image::CreateDepthToCameraDistanceMultiplierFloatImage(
        const camera::PinholeCameraIntrinsic &intrinsic) {
    auto fimage = std::make_shared<Image>();
    fimage->Prepare(intrinsic.width_, intrinsic.height_, 1, 4);
    float ffl_inv[2] = {
            1.0f / (float)intrinsic.GetFocalLength().first,
            1.0f / (float)intrinsic.GetFocalLength().second,
    };
    float fpp[2] = {
            (float)intrinsic.GetPrincipalPoint().first,
            (float)intrinsic.GetPrincipalPoint().second,
    };
    std::vector<float> xx(intrinsic.width_);
    std::vector<float> yy(intrinsic.height_);
    for (int j = 0; j < intrinsic.width_; j++) {
        xx[j] = (j - fpp[0]) * ffl_inv[0];
    }
    for (int i = 0; i < intrinsic.height_; i++) {
        yy[i] = (i - fpp[1]) * ffl_inv[1];
    }
    for (int i = 0; i < intrinsic.height_; i++) {
        float *fp =
                (float *)(fimage->data_.data() + i * fimage->BytesPerLine());
        for (int j = 0; j < intrinsic.width_; j++, fp++) {
            *fp = sqrtf(xx[j] * xx[j] + yy[i] * yy[i] + 1.0f);
        }
    }
    return fimage;
}

std::shared_ptr<Image> Image::CreateWeightImage(
        const camera::PinholeCameraIntrinsic &intrinsic) const {

    auto output = std::make_shared<Image>();

    output->Prepare(intrinsic.width_, intrinsic.height_, 1, 4);
    auto focal_length = intrinsic.GetFocalLength();
    auto principal_point = intrinsic.GetPrincipalPoint();

    #ifdef _OPENMP
    #ifdef _WIN32
    #pragma omp parallel for schedule(static)
    #else
    #pragma omp parallel for collapse(2) schedule(static)
    #endif
    #endif
    for (int i = 0; i < output->height_; i++) {
        for (int j = 0; j < output->width_; j++) {
            float *p = output->PointerAt<float>(j, i);
            float *ip = PointerAt<float>(j, i);
            double weight = 0.0;

            if (*ip > 0) {
                if(i > 0 && j > 0 && i < output->height_-1 && j < output->width_-1){

                    // computing normalized vertex
                    double z = (double)(*ip);
                    double x = (j - principal_point.first) * z / focal_length.first;
                    double y =
                            (i - principal_point.second) * z / focal_length.second;
                    Eigen::Vector3d point = Eigen::Vector3d(x, y, z);
                    Eigen::Vector3d v_norm = point.normalized();

                    //computing normalized normal
                    float *dx1 = PointerAt<float>(j+1, i);
                    float *dx2 = PointerAt<float>(j-1, i);

                    float *dy1 = PointerAt<float>(j, i+1);
                    float *dy2 = PointerAt<float>(j, i-1);

                    double dzdx = (((double)*dx1 - (double)*dx2)/2.0)*1000.0;
                    double dzdy = (((double)*dy1 - (double)*dy2)/2.0)*1000.0;

                    Eigen::Vector3d normal = Eigen::Vector3d(-dzdx, -dzdy, 1.0);
                    Eigen::Vector3d n_norm = normal.normalized();

                    // Eigen::Vector3d captureDir = Eigen::Vector3d(0, 0, 1.0);

                    // double w1 = abs(captureDir.dot(n_norm));
                    // double w2 = abs(captureDir.dot(v_norm));

                    // weight = w1;
                    // weight = w2;
                    double w = abs(n_norm.dot(v_norm));
                    weight = w * w;

                    // // Adding gaussian weight centering at principle point
                    // double w = (double)output->width_/2.0;
                    // double h = (double)output->height_/2.0;
                    // double w_x = ((double)j - w)/w ;
                    // double w_y = ((double)i - h)/h;

                    // double d = sqrt(w_x * w_x + w_y * w_y);
                    // weight = exp(-((d*d)/(2.0))); // assuming mu = 0 and sigma = 1.0
                }
                else{
                    weight = 1.0f; // if this weight is set to 0, some of the points are missing in final integrated reconstruction. Therefore assigning it 1.0
                }

            }
            *p = (float)weight;
        }
    }
    return output;
}

std::shared_ptr<Image> Image::CreateFloatImage(
        Image::ColorToIntensityConversionType type /* = WEIGHTED*/) const {
    auto fimage = std::make_shared<Image>();
    if (IsEmpty()) {
        return fimage;
    }
    fimage->Prepare(width_, height_, 1, 4);
    for (int i = 0; i < height_ * width_; i++) {
        float *p = (float *)(fimage->data_.data() + i * 4);
        const uint8_t *pi =
                data_.data() + i * num_of_channels_ * bytes_per_channel_;
        if (num_of_channels_ == 1) {
            // grayscale image
            if (bytes_per_channel_ == 1) {
                *p = (float)(*pi) / 255.0f;
            } else if (bytes_per_channel_ == 2) {
                const uint16_t *pi16 = (const uint16_t *)pi;
                *p = (float)(*pi16);
            } else if (bytes_per_channel_ == 4) {
                const float *pf = (const float *)pi;
                *p = *pf;
            }
        } else if (num_of_channels_ == 3) {
            if (bytes_per_channel_ == 1) {
                if (type == Image::ColorToIntensityConversionType::Equal) {
                    *p = ((float)(pi[0]) + (float)(pi[1]) + (float)(pi[2])) /
                         3.0f / 255.0f;
                } else if (type ==
                           Image::ColorToIntensityConversionType::Weighted) {
                    *p = (0.2990f * (float)(pi[0]) + 0.5870f * (float)(pi[1]) +
                          0.1140f * (float)(pi[2])) /
                         255.0f;
                }
            } else if (bytes_per_channel_ == 2) {
                const uint16_t *pi16 = (const uint16_t *)pi;
                if (type == Image::ColorToIntensityConversionType::Equal) {
                    *p = ((float)(pi16[0]) + (float)(pi16[1]) +
                          (float)(pi16[2])) /
                         3.0f;
                } else if (type ==
                           Image::ColorToIntensityConversionType::Weighted) {
                    *p = (0.2990f * (float)(pi16[0]) +
                          0.5870f * (float)(pi16[1]) +
                          0.1140f * (float)(pi16[2]));
                }
            } else if (bytes_per_channel_ == 4) {
                const float *pf = (const float *)pi;
                if (type == Image::ColorToIntensityConversionType::Equal) {
                    *p = (pf[0] + pf[1] + pf[2]) / 3.0f;
                } else if (type ==
                           Image::ColorToIntensityConversionType::Weighted) {
                    *p = (0.2990f * pf[0] + 0.5870f * pf[1] + 0.1140f * pf[2]);
                }
            }
        }
    }
    return fimage;
}

template <typename T>
std::shared_ptr<Image> Image::CreateImageFromFloatImage() const {
    auto output = std::make_shared<Image>();
    if (num_of_channels_ != 1 || bytes_per_channel_ != 4) {
        utility::LogError(
                "[CreateImageFromFloatImage] Unsupported image format.");
    }

    output->Prepare(width_, height_, num_of_channels_, sizeof(T));
    const float *pi = (const float *)data_.data();
    T *p = (T *)output->data_.data();
    for (int i = 0; i < height_ * width_; i++, p++, pi++) {
        if (sizeof(T) == 1) *p = static_cast<T>(*pi * 255.0f);
        if (sizeof(T) == 2) *p = static_cast<T>(*pi);
    }
    return output;
}

template std::shared_ptr<Image> Image::CreateImageFromFloatImage<uint8_t>()
        const;
template std::shared_ptr<Image> Image::CreateImageFromFloatImage<uint16_t>()
        const;

ImagePyramid Image::CreatePyramid(size_t num_of_levels,
                                  bool with_gaussian_filter /*= true*/) const {
    std::vector<std::shared_ptr<Image>> pyramid_image;
    pyramid_image.clear();
    if ((num_of_channels_ != 1) || (bytes_per_channel_ != 4)) {
        utility::LogError("[CreateImagePyramid] Unsupported image format.");
    }

    for (size_t i = 0; i < num_of_levels; i++) {
        if (i == 0) {
            std::shared_ptr<Image> input_copy_ptr = std::make_shared<Image>();
            *input_copy_ptr = *this;
            pyramid_image.push_back(input_copy_ptr);
        } else {
            if (with_gaussian_filter) {
                // https://en.wikipedia.org/wiki/Pyramid_(image_processing)
                auto level_b = pyramid_image[i - 1]->Filter(
                        Image::FilterType::Gaussian3);
                auto level_bd = level_b->Downsample();
                pyramid_image.push_back(level_bd);
            } else {
                auto level_d = pyramid_image[i - 1]->Downsample();
                pyramid_image.push_back(level_d);
            }
        }
    }
    return pyramid_image;
}

}  // namespace geometry
}  // namespace open3d


#include <iostream>
#include <array>
#include <mtao/algebra/horner_evaluation.hpp>
#include "mtao/geometry/grid/grid_utils.h"
#include <Eigen/Dense>
#include <iterator>

int main() {
    std::array<int,3> coeffs{{1,1,0}};
    for(int i = 0; i < 10; ++i) {
        std::cout << mtao::algebra::horner_evaluate(i,coeffs) << std::endl;
    }
    {
        std::array<int,2> coeffs{{3,4}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                std::cout << mtao::algebra::horner_rowminor_index(std::array<int,2>{{i,j}},coeffs) << " ";
            }
            std::cout << std::endl;
        }
    }
    std::cout << std::endl << std::endl;
    {
        std::array<int,2> coeffs{{3,4}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                std::cout << mtao::algebra::horner_rowmajor_index(std::array<int,2>{{i,j}},coeffs) << " ";
            }
            std::cout << std::endl;
        }
    }
    {
        std::array<int,3> coeffs{{3,4,5}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                for(int k = 0; k < coeffs[2]; ++k) {
                    std::cout << mtao::algebra::horner_rowminor_index(std::array<int,3>{{i,j,k}},coeffs) << " ";
                }
                std::cout << std::endl;
            }
            std::cout << std::endl;
        }
    }
    std::cout << std::endl << std::endl;
    {
        std::array<int,3> coeffs{{3,4,5}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                for(int k = 0; k < coeffs[2]; ++k) {
                    std::cout << mtao::algebra::horner_rowmajor_index(std::array<int,3>{{i,j,k}},coeffs) << " ";
                }
                std::cout << std::endl;
            }
            std::cout << std::endl;
        }
    }

    {
        Eigen::Matrix<int,3,4,Eigen::ColMajor> A;
        for(int i = 0; i < A.rows(); ++i) {
            for(int j = 0; j < A.cols(); ++j) {
                A(i,j) = &A(i,j) - A.data();

            }
        }
        std::array<int,2> coeffs{{3,4}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                std::cout << mtao::algebra::horner_rowminor_index(std::array<int,2>{{i,j}},coeffs) << ":" << A(i,j) << " ";
            }
            std::cout << std::endl;
        }
    }
    std::cout << std::endl << std::endl;
    {
        Eigen::Matrix<int,3,4,Eigen::RowMajor> A;
        for(int i = 0; i < A.rows(); ++i) {
            for(int j = 0; j < A.cols(); ++j) {
                A(i,j) = &A(i,j) - A.data();

            }
        }
        std::array<int,2> coeffs{{3,4}};
        for(int i = 0; i < coeffs[0]; ++i) {
            for(int j = 0; j < coeffs[1]; ++j) {
                std::cout << mtao::algebra::horner_rowmajor_index(std::array<int,2>{{i,j}},coeffs)  << ":" << A(i,j) << " ";
            }
            std::cout << std::endl;
        }
    }

    auto arrstr = [](auto&& idx) {
        using ArrType = mtao::types::remove_cvref_t<decltype(idx)>;
        using VT = typename ArrType::value_type;
        std::stringstream ss;

        std::copy(idx.begin(),idx.end(),std::ostream_iterator<VT>(ss,","));
        return ss.str();

    };

    using namespace mtao::algebra;
    auto test_inversing = [&](auto&& coeffs) {
        std::cout << "Minor" << coeffs.size() << std::endl;
        mtao::geometry::grid::utils::multi_loop(coeffs,[&](auto&& ij) {
                int idx = horner_rowminor_index(ij,coeffs);
                auto arr = horner_rowminor_inverse_index(idx,coeffs);
                std::cout << arrstr(ij) << " => " << idx << " => " << arrstr(arr) << std::endl;
                assert(ij == arr);
                });
        std::cout << std::endl;
        std::cout << "Major" << coeffs.size() << std::endl;
        mtao::geometry::grid::utils::multi_loop(coeffs,[&](auto&& ij) {
                int idx = horner_rowmajor_index(ij,coeffs);
                auto arr = horner_rowmajor_inverse_index(idx,coeffs);
                std::cout << arrstr(ij) << " => " << idx << " => " << arrstr(arr) << std::endl;
                assert(ij == arr);
                });
    };
    test_inversing(std::array<int,2>{{3,4}});
    std::cout << std::endl;
    std::cout << std::endl;
    test_inversing(std::array<int,3>{{3,4,5}});
    std::cout << std::endl;
    std::cout << std::endl;
    test_inversing(std::array<int,4>{{3,4,5,7}});

}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_GENERIC_CORRECT_FMA_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_GENERIC_CORRECT_FMA_HPP_INCLUDED
#include <boost/simd/arithmetic/functions/correct_fma.hpp>
#include <boost/simd/include/functions/simd/two_prod.hpp>
#include <boost/simd/include/functions/simd/two_add.hpp>
#include <boost/simd/include/functions/simd/plus.hpp>
#include <boost/simd/include/functions/simd/multiplies.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::correct_fma_, tag::cpu_
                                   , (A0)
                                   , (generic_< integer_<A0> >)
                                     (generic_< integer_<A0> >)
                                     (generic_< integer_<A0> >)
                                   )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(3)
    {
      return a0*a1+a2;
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::correct_fma_, tag::cpu_
                                   , (A0)
                                   , (generic_< floating_<A0> >)
                                     (generic_< floating_<A0> >)
                                     (generic_< floating_<A0> >)
                                   )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(3)
    {
      result_type p, rp, s, rs;
      two_prod(a0, a1, p, rp);
      two_add(p, a2, s, rs);
      return s+(rp+rs);
    }
  };

} } }

#endif


#include <iostream>
//#include <unordered_map>
#include <vector>
#include <algorithm>

// Check for alternative Multiprecision Implementation from boost
#if defined(CPP_INT)
    // alternative cpp_int (~4x slower than GMP)
    // no external linking required
    // use compile command with "-DCPP_INT"
    #include <boost/multiprecision/cpp_int.hpp>
    using boost_mp_int = boost::multiprecision::cpp_int;
#elif defined(TOM_INT)
    // alternative libtommath (~50x slower than GMP)
    // use compile command with "-DTOM_INT -ltommath"
    #include <boost/multiprecision/tommath.hpp>
    using boost_mp_int = boost::multiprecision::tom_int;
#else
    // Default GMP (fastest)
    // use compile command with "-lgmp"
    #include <boost/multiprecision/gmp.hpp>
    using boost_mp_int = boost::multiprecision::mpz_int;
#endif

#include <limits>
#include <cstdlib>
#include <cassert>
using namespace std;
using namespace boost::multiprecision;

// Efficient Multiplicative Persistence checker by Hans Loeblich
// This represents large numbers as a collection of digit counts
// for [0:9]
// Multiplying digits is done by taking each digit to the power of its
// given count, and multiplying the resulting powers together.
//
// Powers are calculated and cached up to the max specified lengths
// before doing any checks.  This makes each persistence check
// boil down to some fast lookups of powers, and then a few
// multiplications depending on which digits are present.
//
// Inspired by the NumberPhile video about 277777788888899
// https://www.youtube.com/watch?v=Wim9WJeDTHQ

class digit_bag {
public:
    typedef array<vector<boost_mp_int>, 10> pow_cache_t;
    typedef vector<size_t> rule_prefix_t; // the prefix for a rule
    typedef vector<size_t> rule_tail_digits_t; // list of possible digits for tail
    typedef pair<rule_prefix_t,rule_tail_digits_t> rule_t;
    typedef vector<rule_t> rules_t;
    typedef array<size_t, 10> digits_t;

    digit_bag(size_t length) : size(length), rule_it(rules.begin()), digits(), result_str(length+1,'\0') {
        init_rule(*rule_it);
    }
    digit_bag(digit_bag&&) = delete;
    digit_bag& operator=(digit_bag&&) = delete;
    digit_bag(const digit_bag&) = delete;
    digit_bag& operator=(const digit_bag&) = delete;

    // use a fast lookup table for powers of each digit up to max length
    // ignore 0 and 1
    static void init_cache(size_t max) {
        boost_mp_int power;
        for(size_t i = 2; i < 10; ++i) {
            boost_mp_int base = i;
            for(size_t p = 0; p <= max; ++p) {
                power = pow(base, p);
                cache[i].push_back(power);
            }
        }
    }

    // when iterating to a new rule,
    // initialize digits based on lowest number in the rule
    inline void init_rule(const rule_t& rule) {
        // reset digit counts
        for(auto &d : digits) d = 0;
        // initialize head
        for(const size_t di : rule.first) digits[di]++;
        // set remaining digits to first option in tail_digits
        digits[rule.second[0]] += size - rule.first.size();
    }

    // check every combination of given length
    // starting from digits already in bag
    pair<digits_t, size_t> check_all() {
        digits_t current_max;
        size_t current_max_count = 0;
        do {
            size_t count = persistence();
            if (count > current_max_count) {
                current_max_count = count;
                current_max = digits;
            }
        } while(next());
        return make_pair(current_max,current_max_count);
    }

    // check persitence of current digits
    // assumes input is already length > 1
    size_t persistence() const {
        result = 1;
        for (size_t i=2; i<10; ++i) {
            if (digits[i]) {
                result *= cache[i][digits[i]];
            }
        }
        size_t mult_count = 1;

        while (true) {
            digits_t new_digits{0,0,0,0,0,0,0,0,0,0};
            constexpr const size_t CH_ZERO = '0';
            // convert result to string and count digits of each type
        #if !defined(CPP_INT) && !defined(TOM_INT)
            // Minor optimization (~2%) for GMP string conversion
            mpz_get_str(&result_str[0], 10, result.backend().data());
            if (result_str[1] == '\0') { // single digit result
                return mult_count;
            }
            for(size_t ch : result_str) {
                if (ch == '\0') break;
                size_t di = ch - CH_ZERO; // turn character into numeric digit
                if (di == 0) {
                    return mult_count+1;
                }
                ++new_digits[di];
            }
        #else
            // Generic std::string check
            result_str = result.str();
            if (result_str.size() == 1) { // single digit result
                return mult_count;
            }
            for(size_t ch : result_str) {
                size_t di = ch-'0'; // turn character into numeric digit
                if (di == 0) {
                    return mult_count+1;
                }
                ++new_digits[di];
            }
        #endif
            // we can't check this inside the loop in case there is a zero also,
            // we would return 1 higher than expected
            if (new_digits[5] && (new_digits[2] || new_digits[4] || new_digits[6] || new_digits[8])) {
                return mult_count+2;
            }

            result = 1;
            for (size_t i=2; i<10; ++i) {
                if (new_digits[i]) {
                    result *= cache[i][new_digits[i]];
                }
            }
            ++mult_count;
        }
    }

    // get the next set of digits, in ascending order based on rules
    inline bool next() noexcept {
        const rule_prefix_t &head = rule_it->first;
        const rule_tail_digits_t &tail_digits = rule_it->second;
        const size_t last_digit = tail_digits[tail_digits.size()-1];
        if (digits[last_digit] == size - head.size()) {
            rule_it++;
            if (rule_it == rules.end()) return false;
            init_rule(*rule_it);
            return true;
        } else {
            const size_t tds = tail_digits.size();
            if (tds >= 2) {
                for(size_t ti=tds-2; ti>=0; --ti) {
                    const size_t di = tail_digits[ti];
                    if (digits[di]) {
                        const size_t di2 = tail_digits[ti+1];
                        digits[di]--;
                        digits[di2]++;
                        for(size_t ji2=ti+2; ji2<tds; ++ji2) {
                            const size_t di3 = tail_digits[ji2];
                            digits[di2] += digits[di3];
                            digits[di3] = 0;
                        }
                        return true;
                    }
                }
                //assert(false && "where the digits go?");
            } else {
                rule_it++;
                if (rule_it == rules.end()) return false;
                init_rule(*rule_it);
                return true;
            }
        }
    }

    inline const digits_t& getDigits() const {
        return digits;
    }

private:
    static const rules_t rules;
    static pow_cache_t cache;

    size_t size;
    rules_t::const_iterator rule_it;
    digits_t digits;
    mutable string result_str;
    mutable boost_mp_int result;
};

ostream& operator<<(ostream &out, const digit_bag::digits_t &digits)
{
    out << "{ " << digits[0];
    for (size_t i = 1; i<10; ++i) cout << ", " << digits[i];
    cout << " }";
    return out;
}

// https://oeis.org/A003001
// Summarizing, a term a(n) for n > 2 consists of 7's, 8's and 9's
// with a prefix of one of the following sets of digits:
// {{}, {2}, {3}, {4}, {6}, {2,6}, {3,5}, {5, 5,...}}
// [Amended by Kohei Sakai, May 27 2017]

// Rules are listed to generate numbers in ascending order
const digit_bag::rules_t digit_bag::rules{
    {{2,6},{7,8,9}},
    {{2},  {7,8,9}},
    {{3,5},{7,8,9}},
    {{3},  {7,8,9}},
    {{4},  {7,8,9}},
    {{5},  {5,7,9}},
    {{6},  {7,8,9}},
    {{ },  {7,8,9}}
};
digit_bag::pow_cache_t digit_bag::cache;

int main(int argc, char** argv) {
    if (argc < 3 || argc > 4) {
        cout << "Search for the smallest numbers with the highest multiplicative persistence for base10 numbers with lengths in the range of [START:END)\n";
        cout << "   Or in other words: START < log10(number) < END\n";
        cout << "usage: multper START END [MAX]\n";
        cout << "   START   range lower bound(inclusive).  START must be >= 2\n";
        cout << "   END     range upper bound(exclusive)\n";
        cout << "   MAX     (optional) sets the maximum persistence to treat specially(default=0 regardless of START length)\n";
        cout << "     Any number found with persistence greater than MAX will print a line starting with \"NEW MAX\", with empty lines before/after for visibility.\n";
        cout << "     Otherwise print a single line for each length indicating the best match (smallest number with highest persistence for that length)\n";
        return 1;
    }

    const size_t START = stoul(argv[1]);
    assert(START>=2 && "");
    const size_t END = stoul(argv[2]);
    size_t max_count = 0;
    if (argc == 4) {
        max_count = stoul(argv[3]);
    }

    //cout << "Initializing lookup..."; cout.flush();
    digit_bag::init_cache(END);
    //cout << "Done." << endl;

    for(size_t length=START; length<END; ++length) {
        digit_bag bag(length);
        pair<digit_bag::digits_t, size_t> current_max = bag.check_all();
        if (current_max.second > max_count) {
            max_count = current_max.second;
            cout << endl << "NEW MAX " << max_count << " persistence for " << length << " digits: " << current_max.first << endl << endl;
        } else {
            cout << current_max.second << " best persistence for " << length << " digits: " << current_max.first << endl;
        }
    }
    return 0;
}