/* Copyright (c) 2012, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See LICENSE.txt or
 * http://www.opensource.org/licenses/mit-license.php */

#pragma once

#include "hdp_base.hpp"
#include "hdp_var_base.hpp"

#include "random.hpp"
#include "baseMeasure.hpp"
#include "probabilityHelpers.hpp"

#include <stddef.h>
#include <stdint.h>
#include <typeinfo>

#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/beta.hpp>
#include <armadillo>

using namespace std;
using namespace arma;

/*
 * this one assumes that the number of words per document is
 * smaller than the dictionary size
 *
 * http://en.wikipedia.org/wiki/Virtual_inheritance
 */
template <class U>
class HDP_var: public HDP<U>, public virtual HDP_var_base
{
  public:

    HDP_var(const BaseMeasure<U>& base, double alpha, double omega)
      : HDP_var_base(0,0,0), HDP<U>(base, alpha, omega)
    {};

    ~HDP_var()
    {};

    // interface mainly for python
    uint32_t addDoc(const Mat<U>& x_i)
    {
      uint32_t x_ind = HDP<U>::addDoc(x_i);
      // TODO: potentially slow
      // add the index of the added x_i
      mInd2Proc.resize(mInd2Proc.n_elem+1);
      mInd2Proc[mInd2Proc.n_elem-1] = x_ind;
      return x_ind;
    };

    /*
     * Initializes the corpus level parameters mA and mLambda according
     * to Blei's Stochastic Variational paper
     *
     * @param D is the assumed number of documents for init
     */
    void initCorpusParams(uint32_t Nw, uint32_t K, uint32_t T, uint32_t D)
    {

      mT = T;
      mK = K;
      mNw = Nw;

      mA.ones(K,2);
      mA.col(1) *= HDP<U>::mOmega;

      // initialize lambda
      HDP<U>::mLambda.init(HDP<U>::mH0,K); // initialize the priors of the topics from the base measure.

//      HDP<U>::mLambda.zeros(K,Nw);
//      GammaRnd gammaRnd(1.0, 1.0);
//      for (uint32_t k=0; k<K; ++k){
//        for (uint32_t w=0; w<Nw; ++w) HDP<U>::mLambda(k,w) = gammaRnd.draw();
//        HDP<U>::mLambda.row(k) += ((Dir*)(&mH0))->mAlphas;
//      }
    };

    /* From: Online Variational Inference for the HDP
     * method for "one shot" computation without storing data in this class
     *  Nw: number of different words
     *  kappa=0.9: forgetting rate
     *  uint32_t T=10; // truncation on document level
     *  uint32_t K=100; // truncation on corpus level
     *  S = batch size
     */
    void densityEst(const vector<Mat<U> >& x, uint32_t Nw,
        double kappa, uint32_t K, uint32_t T, uint32_t S)
    {
      cout<<"densityEstimate with: K="<<K<<"; T="<<T<<"; kappa="<<kappa<<"; Nw="<<Nw<<"; S="<<S<<endl;

      HDP<U>::mX = x;
      uint32_t D=HDP<U>::mX.size();
      cout<<"D="<<D<<endl;
      cout<<"mX[0].shape= "<<HDP<U>::mX[0].n_rows<<"x"<<HDP<U>::mX[0].n_cols<<endl;

      mInd2Proc = linspace<Row<uint32_t> >(0,D-1,D);

      mT = T;
      mK = K;
      mNw = Nw;

      initCorpusParams(mNw,mK,mT,D);
      cout<<"Init of corpus params done"<<endl;
      Row<uint32_t> ind = updateEst_batch(mInd2Proc,mZeta,mPhi,mGamma,mA,HDP<U>::mLambda,mPerp,HDP<U>::mOmega,kappa,S,true);
//      cout<<"mPhi -> D="<<mPhi.size()<<endl;
//      cout<<"mPhi -> D="<<HDP_var_base::mPhi.size()<<endl;
//      cout<<"mPerp="<<mPerp.t()<<endl;

      Mat<double> pi(D,T);
      Mat<double> sigPi(D,T+1);
      Mat<uint32_t> c(D,T);
      getDocTopics(pi,sigPi,c);
//      cout<<"c:"<<c<<endl;

      mInd2Proc.set_size(0); // all processed

    };

    /*
     * compute density estimate based on data previously fed into the class using addDoc
     */
    bool densityEst(uint32_t Nw, double kappa, uint32_t K, uint32_t T, uint32_t S)
    {
      if(HDP<U>::mX.size() > 0)
      {
        densityEst(HDP<U>::mX,Nw,kappa,K,T,S);
        //TODO: return p_d(x)
        return true;
      }else{
        return false;
      }
    };

    /*
     * updated the estimate using newly added docs in mX
     * "newly added": the ones that are indicated in mInd2Proc
     */
    bool updateEst_batch(double kappa, uint32_t S)
    {
      uint32_t Db=mInd2Proc.n_elem;
      if (Db >0){
        cout<<"updatedEstimate with: K="<<mK<<"; T="<<mT<<"; kappa="<<kappa<<"; Nw="<<mNw<<"; S="<<S<<endl;
        vector<Mat<double> > zeta; // will get resized accordingly inside updateEst_batch
        vector<Mat<double> > phi;
        vector<Mat<double> > gamma;
        Col<double> perp;

        Row<uint32_t> ind = updateEst_batch(mInd2Proc,zeta,phi,gamma,mA,HDP<U>::mLambda,perp,HDP<U>::mOmega,kappa,S);

        mZeta.resize(mZeta.size()+Db);
        mPhi.resize(mPhi.size()+Db);
        mGamma.resize(mGamma.size()+Db);
        mPerp.resize(mPerp.n_elem+Db);
        mPerp.rows(mPerp.n_elem-Db,mPerp.n_elem-1) = perp;
        for (uint32_t i=0; i<ind.n_elem; ++i){
          mZeta[ind[i]] = zeta[i];
          mPhi[ind[i]] = phi[i];
          mGamma[ind[i]] = gamma[i];
        }

        mInd2Proc.set_size(0); // all processed

        return true;
      }else{
        cout<<"add more documents before starting to process"<<endl;
      }
      return false;
    };

    /*
     * after an initial densitiy estimate has been made using densityEst()
     * can use this to update the estimate with information from additional x
     */
    bool  updateEst(const Mat<U>& x, double kappa)
    {
      if (HDP<U>::mX.size() > 0 && HDP<U>::mX.size() == mPhi.size()) { // this should indicate that there exists an estimate already
        uint32_t N = x.n_cols;
        uint32_t T = mT;
        uint32_t K = mK;
        HDP<U>::mX.push_back(x);
        mZeta.push_back(Mat<double>(T,K));
        mPhi.push_back(Mat<double>(N,T));
        //    mZeta.set_size(T,K);
        //    mPhi.set_size(N,T);
        mGamma.push_back(Mat<double>(T,2));
        uint32_t d = HDP<U>::mX.size()-1;
        mPerp.resize(d+1);

        if(updateEst(HDP<U>::mX[d],mZeta[d],mPhi[d],mGamma[d],mA,HDP<U>::mLambda,HDP<U>::mOmega,d,kappa))
        {
          mPerp[d] = 0.0;
          for (uint32_t i=0; i < HDP<U>::mX_ho.size(); ++i)
          {
            Row<double> logP=logP_w(HDP<U>::mX[d],mPhi[d], mZeta[d], mGamma[d], HDP<U>::mLambda);
            mPerp[d] += HDP<U>::perplexity(HDP<U>::mX_ho[i], logP);
            //mPerp[d] += perplexity(mX_ho[i], mZeta[d], mPhi[d], mGamma[d], HDP<U>::mLambda);
          }
          mPerp[d] /= double(HDP<U>::mX_ho.size());
          cout<<"Perplexity="<<mPerp[d]<<endl;
          return true;
        }else{
          return false;
        }
      }else{
        return false;
      }
    };

    /*
     *
     */
    bool updateEst(const Mat<U>& x, Mat<double>& zeta, Mat<double>& phi, Mat<double>& gamma, Mat<double>& a, DistriContainer<U>& lambda, double omega, uint32_t d, double kappa)
    {
      uint32_t D = d+1; // assume that doc d is appended to the end
//      uint32_t Nw = lambda.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      Mat<double> eLogBeta(mK,x.n_cols); //TODO
      Col<double> digam_lamb_sum(mK);

      compElogBeta(eLogBeta,  lambda, x);

      Col<double> eLogSig_a(K);
      compElogSig(eLogSig_a, a);

      //    cout<<"---------------- Document "<<d<<" N="<<N<<" -------------------"<<endl;
      initZeta(zeta,eLogBeta);
      initPhi(phi,zeta,eLogBeta);

      if(!is_finite(zeta))
      {
        cout<<"updateEst::a="<<size(a)<<endl<<a;
        cout<<"updateEst::lambda="<<lambda[0]->asRow();
        cout<<"updateEst::x="<<size(x)<<endl<<x;
        cout<<"zeta_init="<<zeta<<endl;
        cout<<"phi_init="<<phi<<endl;
        cout<<"eLogBeta="<<eLogBeta<<endl;
        exit(0);
      }

      // ------------------------ doc level updates --------------------
      bool converged = false;
      Col<double> eLogSig_gam(T);
      Mat<double> gamma_prev(T,2);
      gamma_prev.ones();

      uint32_t o=0;
      while(!converged){
        //      cout<<"-------------- Iterating local params #"<<o<<" -------------------------"<<endl;
        updateGamma(gamma,phi);

        if (!is_finite(gamma)){
          cout<<"gamma="<<gamma;
          cout<<"phi="<<phi;
          exit(1);
        }

        compElogSig(eLogSig_gam,gamma); // precompute

        updateZeta(zeta,phi,eLogSig_a,eLogBeta);
        updatePhi(phi,zeta,eLogSig_gam,eLogBeta);

        converged = (accu(gamma_prev != gamma))==0 || o>60 ;
        gamma_prev = gamma;
        ++o;

        if(!is_finite(zeta))
        {
          cout<<"o="<<o<<endl;
          cout<<"zeta="<<zeta.row(0)<<" |.|="<<sum(exp(zeta.row(0)))<<endl;
          cout<<"phi="<<phi.row(0)<<" |.|="<<sum(exp(phi.row(0)))<<endl;
          exit(0);
        }
      }

      //    cout<<" --------------------- natural gradients --------------------------- "<<endl;
      //    cout<<"\tD="<<D<<" omega="<<omega<<endl;
      DistriContainer<U> d_lambda(HDP<U>::mH0,K);
      Mat<double> d_a(K,2);
      computeNaturalGradients(d_lambda, d_a, zeta, phi, omega, D, x);

      cout<<"update::d_lambda:"<<d_lambda.toMat().rows(0,5);
      cout<<"update::lambda:"<<lambda.toMat().rows(0,5);

      //    cout<<" ------------------- global parameter updates: ---------------"<<endl;
      double ro = exp(-kappa*log(1+double(d+1)));
      //    cout<<"\tro="<<ro<<endl;
      for (uint32_t k=0; k<mK; ++k)
        lambda[k]->fromRow( (1.0-ro)*lambda[k]->asRow() + ro* d_lambda[k]->asRow());

      //lambda = (1.0-ro)*lambda + ro*d_lambda;
      a = (1.0-ro)*a+ ro*d_a;
      return true;
    };

    /*
     * Updates the estimate using mini batches
     * @param ind_x indices of docs to process within docs mX. !these are assumed to be in order!
     * @param sameIndAsX == true -> zeta,phi,gamma have same indices as x (ind_x). This typically happens for the initial batch update. Setting this to true eliminates the need of reordering the results afterwords to match the indices of the docs x.
     * @return the randomly shuffled indices to show how the data was processed -> this allows association of zetas, phis and gammas with docs in mX
     */
    Row<uint32_t> updateEst_batch(const Row<uint32_t>& ind_x, vector<Mat<double> >& zeta, vector<Mat<double> >& phi, vector<Mat<double> >& gamma, Mat<double>& a, DistriContainer<U>& lambda, Col<double>& perp, double omega, double kappa, uint32_t S, bool sameIndAsX=false)
    {
      uint32_t d_0 = min(ind_x); // thats the doc number that we start with -> needed for ro computation; assumes that all indices in mX prior to d_0 have already been processed.
      uint32_t D= max(ind_x)+1; // D is the maximal index of docs that we are processing +1

      Row<uint32_t> ind = shuffle(ind_x,1);
//        cout<<"ind_x: "<<ind_x.cols(0,S)<<endl;
//        cout<<"ind  : "<<ind.cols(0,S)<<endl;

      zeta.resize(ind.n_elem,Mat<double>(mT,mK));
      phi.resize(ind.n_elem);
      gamma.resize(ind.n_elem,Mat<double>(mT,2));
      perp.zeros(ind.n_elem);

      for (uint32_t dd=0; dd<ind.n_elem; dd += S)
      {
        DistriContainer<U> db_lambda(HDP<U>::mH0,mK);

        Mat<double> db_a(mK,2);
        db_a.zeros();

        Col<double> eLogSig_a(mK);
        compElogSig(eLogSig_a, a);

#pragma omp parallel for schedule(dynamic)
        for (uint32_t db=dd; db<min(dd+S,ind.n_elem); db++)
        {
          uint32_t d=ind[db];
          const Mat<U>& x_d = HDP<U>::mX[d];
          uint32_t dout=sameIndAsX?d:db;
          uint32_t N=x_d.n_cols;
          //      cout<<"---------------- Document "<<d<<" N="<<N<<" -------------------"<<endl;

          cout<<"-- db="<<db<<" d="<<d<<" N="<<N<<endl;

          Mat<double> eLogBeta(mK,x_d.n_cols);
          Col<double> digam_lamb_sum(mK);
          compElogBeta(eLogBeta, lambda, x_d);

          //Mat<double> zeta(T,K);
          phi[dout].resize(N,mT);
          initZeta(zeta[dout],eLogBeta);
          initPhi(phi[dout],zeta[dout],eLogBeta);

//            cout<<"zeta_init="<<zeta[dout]<<endl;
//            cout<<"phi_init="<<phi[dout]<<endl;
//            cout<<"eLogBeta="<<eLogBeta<<endl;
//            if(!is_finite(zeta[dout]))
//              exit(0);

          //cout<<" ------------------------ doc level updates --------------------"<<endl;
          //Mat<double> gamma(T,2);
          Col<double> eLogSig_gam(mT);
          Mat<double> gamma_prev(mT,2);
          gamma_prev.ones();
          gamma_prev.col(1) += HDP<U>::mAlpha;
          bool converged = false;
          uint32_t o=0;
          while(!converged){
//            cout<<"-------------- Iterating local params #"<<o<<" -------------------------"<<endl;
            updateGamma(gamma[dout],phi[dout]);

            if (!is_finite(gamma[dout])){
              cout<<"gamma="<<gamma[dout];
              cout<<"phi="<<phi[dout];
              exit(1);
            }

            compElogSig(eLogSig_gam,gamma[dout]); // precompute

            updateZeta(zeta[dout],phi[dout],eLogSig_a,eLogBeta);
            updatePhi(phi[dout],zeta[dout],eLogSig_gam,eLogBeta);

            converged = (accu(gamma_prev != gamma[dout]))==0 || o>30 ;
            gamma_prev = gamma[dout];
            ++o;

//            cout<<"o="<<o<<endl;
//            cout<<"zeta="<<zeta[dout].row(0)<<" |.|="<<sum(zeta[dout].row(0))<<endl;
//            cout<<"phi="<<phi[dout].row(0)<<" |.|="<<sum(phi[dout].row(0))<<endl;
//            if(!is_finite(zeta[dout]))
//              exit(0);
          }

          DistriContainer<U> d_lambda(HDP<U>::mH0,mK); // batch updates
          Mat<double> d_a(mK,2);
//          cout<<" --------------------- natural gradients dout="<< dout<<" dd="<< dd<<" --------------------------- "<<endl;
          computeNaturalGradients(d_lambda, d_a, zeta[dout], phi[dout], HDP<U>::mOmega, D, x_d);
#pragma omp critical
          {
            cout<<"update_batch::d_lambda:"<<endl<<d_lambda.toMat().rows(0,5);
            for (uint32_t k=0; k<d_lambda.size(); ++k){
              db_lambda[k]->fromRow(db_lambda[k]->asRow() + d_lambda[k]->asRow());
            }
            db_a += d_a;
          }
        }
        //for (uint32_t k=0; k<K; ++k)
        //  cout<<"delta lambda_"<<k<<" min="<<min(d_lambda.row(k))<<" max="<< max(d_lambda.row(k))<<" #greater 0.1="<<sum(d_lambda.row(k)>0.1)<<endl;
        // ----------------------- update global params -----------------------
        uint32_t t=dd+d_0; // d_0 is the timestep of the the first index to process; dd is the index in the current batch
        uint32_t bS = min(S,ind.n_elem-dd); // necessary for the last batch, which migth not form a complete batch
        //TODO: what is the time dd? d_0 needed?
        double ro = exp(-kappa*log(1+double(t)+double(bS)/2.0)); // as "time" use the middle of the batch
//        cout<<" -- global parameter updates t="<<t<<" bS="<<bS<<" ro="<<ro<<endl;
//        cout<<"d_a="<<db_a<<endl;

//        cout<<"dLambda"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<db_lambda[k]->asRow();
//
//        cout<<"Before"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<lambda[k]->asRow();

        cout<<"update_batch::db_lambda:"<<endl<<db_lambda.toMat().rows(0,5);
        for (uint32_t k=0; k<db_lambda.size(); ++k)
          lambda[k]->fromRow((1.0-ro)*lambda[k]->asRow() + (ro/S)*db_lambda[k]->asRow()); //TODO: doies this make sense for NIW prior???
        cout<<"update_batch::lambda(after):"<<endl<<lambda.toMat().rows(0,5);

//        cout<<"After"<<endl;
//        for (uint32_t k=0; k<10; ++k)
//          cout<<lambda[k]->asRow();

        //lambda = (1.0-ro)*lambda + (ro/S)*db_lambda;
        a = (1.0-ro)*a + (ro/S)*db_a;
        cout<<"update_batch::lambda:"<<lambda.toMat().rows(0,5);

        perp[dd+bS/2] = 0.0;
        if (HDP<U>::mX_te.size() > 0) {
          cout<<"computing "<<HDP<U>::mX_te.size()<<" perplexities"<<endl;
#pragma omp parallel for schedule(dynamic)
          for (uint32_t i=0; i < HDP<U>::mX_te.size(); ++i)
          {
            //cout<<"mX_te: "<< mX_te[i].n_rows << "x"<< mX_te[i].n_cols<<endl;
            //cout<<"mX_ho: "<< mX_ho[i].n_rows << "x"<< mX_ho[i].n_cols<<endl;
            //TODO: these subfunctions work on the member variables!!! dont do that...
            double perp_i =  perplexity(HDP<U>::mX_te[i],HDP<U>::mX_ho[i],dd+bS/2+1,ro); //perplexity(mX_ho[i], mZeta[d], mPhi[d], mGamma[d], lambda);
            //cout<<"perp_"<<i<<"="<<perp_i<<endl;
#pragma omp critical
            {
              perp[dd+bS/2] += perp_i;
            }
          }
          perp[dd+bS/2] /= double(HDP<U>::mX_te.size());
          cout<<"Perplexity="<<perp[dd+bS/2]<<endl;
        }
      }
      cout<<"perp="<<perp.t()<<endl;
      return ind;
    };

    // compute the perplexity of a given document split into x_test (to find a topic model for the doc) and x_ho (to evaluate the perplexity)
    double perplexity(const Mat<U>& x_te, const Mat<U>& x_ho, uint32_t d, double kappa=0.75)
    {
      if (HDP<U>::mX.size() > 0 && HDP<U>::mX.size() == mPhi.size()) { // this should indicate that there exists a estimate already

        uint32_t N = x_te.n_cols;
        uint32_t T = mT;
        uint32_t K = mK;

        Mat<double> zeta(T,K);
        Mat<double> phi(N,T);
        Mat<double> gamma(T,2);
        //uint32_t d = mX.size()-1;

        Mat<double> a(mA);// DONE: make deep copy here!

        DistriContainer<U> lambda(HDP<U>::mLambda);
        double omega = HDP<U>::mOmega;

        cout<<"perplexity::lambda:"<<lambda.toMat().rows(0,5);
        cout<<"perplexity::mLambda:"<<HDP<U>::mLambda.toMat().rows(0,5);

        cout<<"updating copied model with x"<<endl;
        updateEst(x_te,zeta,phi,gamma,a,lambda,omega,d,kappa);
//        cout<<" lambda.shape="<<lambda.size()<<endl;
        cout<<"computing perplexity under updated model"<<endl;
        //TODO: compute probabilities then use that to compute perplexity

        cout<<"x_te: "<<size(x_te);
        Mixture<U> mix = docMixture(phi, zeta, gamma, lambda);
        return HDP<U>::perplexity(x_ho, mix);
        //return perplexity(x_ho, zeta, phi, gamma, lambda);
      }else{
        return 1.0/0.0;
      }
    };

    Mixture<U> docMixture(uint32_t d) const {
      return docMixture(mPhi[d],mZeta[d],mGamma[d],HDP<U>::mLambda);
    };
    Mixture<U> docMixture(const Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& gamma, const DistriContainer<U>& lambda) const
    {
      Col<double> pi;
      Col<double> sigPi;
      Col<uint32_t> c;
      getDocTopics(pi,sigPi,c,gamma,zeta);
      //cout<<"getDocTopics done"<<endl;
      //cout<<"c="<<c.t()<<size(c);
      Col<uint32_t> z(mNw);
      getWordTopics(z, phi);
      //cout<<"getWordTopics done"<<endl;
      //cout<<"z="<<z.t()<<size(z);

      DistriContainer<U> beta;
      HDP<U>::getCorpTopics(beta,lambda);
//      cout<<"getCorpTopics done"<<endl;
      //cout<<"beta:\t"<<size(beta);

      Col<uint32_t> c_u = unique(c);
      Row<double> ps(c_u.n_elem); // proportions in  the mixture
      DistriContainer<U> beta_d(c_u.n_elem);
      for (uint32_t i=0; i< c_u.n_elem; ++i){
        beta_d[i] = beta[c_u(i)]->getCopy();
        ps[i] = sum(sigPi.elem(find(c == c_u(i) )));
      }
//      cout<<"Mixture done"<<endl;
      return Mixture<U>(beta_d,ps);
    };

    /* Probability distribution over the words in document d
     *
     * TODO: so is that here not some MAP or ML estimate?!
     */
    Row<double> logP_w(uint32_t d) const {
//      cout<<"mX.size="<<mX.size()<<endl;
      return logP_w(HDP<U>::mX[d],mPhi[d],mZeta[d],mGamma[d],HDP<U>::mLambda);
    };
    /*
     * log probability using the samples x (not using sufficient statistics -> x is just a list of words)
     */
    Row<double> logP_w(const Mat<U>& x, const Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& gamma, const DistriContainer<U>& lambda) const
    {
      Row<double> p(x.n_cols);
      p.zeros();
//      cout<<"x:\t"<<size(x);
//      cout<<"phi:\t"<<size(phi);
//      cout<<"zeta:\t"<<size(zeta);
//      cout<<"gamma:\t"<<size(gamma);
//      cout<<"lambda:\t"<<size(lambda);

      Col<double> pi;
      Col<double> sigPi;
      Col<uint32_t> c;
      getDocTopics(pi,sigPi,c,gamma,zeta);
      //cout<<"getDocTopics done"<<endl;
      //cout<<"c="<<c.t()<<size(c);
      Col<uint32_t> z(mNw);
      getWordTopics(z, phi);
      //cout<<"getWordTopics done"<<endl;
      //cout<<"z="<<z.t()<<size(z);

      DistriContainer<U> beta;
      HDP<U>::getCorpTopics(beta,lambda);
      //cout<<"getCorpTopics done"<<endl;
      //cout<<"beta:\t"<<size(beta);

      for (uint32_t i=0; i<x.n_cols; ++i){
        //cout<<"z_"<<i<<"="<<z[i]<<endl;
        p[i] = beta[c[z[i]]]->logP(x.col(i));// logCat(x[i], beta.row( c[ z[i] ]));
        //cout<<"p_"<<x[i]<<"="<<p[x[i]]<<endl;
      }
//      for (uint32_t w=0; w<mNw; ++w)
//        p[w] = p[w]==0.0?-1e10:p[w];

      //cout<<"p="<<p<<endl;
      return p;
    };

  protected:

    Row<uint32_t> mInd2Proc; // indices of docs that have not been processed

  private:

    /*
     * precompute necessary digamma function values, because these are slowing the whole algorithm down
     * all the update methods for zeta and phi need these values very often! I can precumpute these once after updating the global parameters (and hence lambda)
     */
    void compElogBeta(Mat<double>& eLogBeta, const DistriContainer<U>& lambda, const Mat<U>& x_d) const
    {
      eLogBeta.zeros(mK,x_d.n_cols);

      Col<double> digam_lamb_sum(mK);
    //  Mat<uint32_t> x_u = unique(x_d); // cannot do this trick anymore since x_d are continuous for the NIW base measure
//      cout<<"x_d="<<x_d<<endl;
//      cout<<"x_u="<<x_u<<endl;

      for (uint32_t i = 0; i < x_d.n_cols ; i++) {
        for (uint32_t k = 0; k < mK; k++) {
           eLogBeta(k,i) = lambda[k]->Elog(x_d.col(i)); // E[log beta] computation in paper
        }
      }
//      for (uint32_t k = 0; k < mK; k++) {
//        digam_lamb_sum(k) = digamma(sum(lambda.row(k)));
//      }
//      for (uint32_t i = 0; i < x_u.n_elem ; i++) {
//        for (uint32_t k = 0; k < mK; k++) {
//           eLogBeta(k,x_u(i)) = digamma(lambda(k,x_u(i))) - digam_lamb_sum(k);
//        }
//      }
    }

    void compElogSig(Col<double>& eLogSig, const Mat<double>& a) const
    {
      for (uint32_t k=0; k<a.n_rows; ++k){
        eLogSig(k) = digamma(a(k,0)) - digamma(a(k,0) + a(k,1));
        for (uint32_t l=0; l<k; ++l)
          eLogSig(k) += digamma(a(l,1)) - digamma(a(l,0) + a(l,1));
      }
    }

    void initZeta(Mat<double>& zeta, const Mat<double>& eLogBeta)
    {
      uint32_t N = eLogBeta.n_cols; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;
      //cerr<<"\tinit zeta"<<endl;
      for (uint32_t i=0; i<T; ++i) {
        for (uint32_t k=0; k<K; ++k) {
          zeta(i,k)=0.0;
          for (uint32_t n=0; n<N; ++n) {
            //if(i==0 && k==0) cout<<zeta(i,k)<<" -> ";
            zeta(i,k) += eLogBeta(k,n); //ElogBeta(lambda, k, x_d(n));
          }
        }
        normalizeLogDistribution(zeta.row(i));
        //cout<<" normalized="<<zeta(0,0)<<endl;
      }
      //cerr<<"zeta>"<<endl<<zeta<<"<zeta"<<endl;
      //cerr<<"normalization check:"<<endl<<sum(zeta,1).t()<<endl; // sum over rows
    };

    void initPhi(Mat<double>& phi, const Mat<double>& zeta, const Mat<double>& eLogBeta)
    {
      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;
      //cout<<"\tinit phi"<<endl;
      for (uint32_t n=0; n<N; ++n){
        for (uint32_t i=0; i<T; ++i) {
          phi(n,i)=0.0;
          for (uint32_t k=0; k<K; ++k) {
            phi(n,i)+=zeta(i,k)* eLogBeta(k,n); // ElogBeta(lambda, k, x_d(n));
          }
        }
        normalizeLogDistribution(phi.row(n));
      }
      //cerr<<"phi>"<<endl<<phi<<"<phi"<<endl;
    };

    void updateGamma(Mat<double>& gamma, const Mat<double>& phi)
    {
      uint32_t N = phi.n_rows;
      uint32_t T = phi.n_cols;

      gamma.ones();
      gamma.col(1) *= HDP<U>::mAlpha;
      for (uint32_t i=0; i<T; ++i)
      {
        for (uint32_t n=0; n<N; ++n){
          gamma(i,0) += phi(n,i);
          for (uint32_t j=i+1; j<T; ++j) {
            gamma(i,1) += phi(n,j);
          }
        }
      }
      //cout<<gamma.t()<<endl;
    };

    void updateZeta(Mat<double>& zeta, const Mat<double>& phi, const Col<double>& eLogSig_a, const Mat<double>& eLogBeta)
    {
//      assert(x_d.n_rows == 1);

      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      for (uint32_t i=0; i<T; ++i){
        //zeta(i,k)=0.0;
        for (uint32_t k=0; k<K; ++k) {
          zeta(i,k) = eLogSig_a(k); //ElogSigma(a,k);
          //cout<<zeta(i,k)<<endl;
          for (uint32_t n=0; n<N; ++n){
            zeta(i,k) += phi(n,i)* eLogBeta(k,n); //ElogBeta(lambda,k,x_d(n));
          }
        }
        normalizeLogDistribution(zeta.row(i));
      }
    }

    void updatePhi(Mat<double>& phi, const Mat<double>& zeta, const Col<double>& eLogSig_gam, const Mat<double>& eLogBeta)
    {
//      assert(x_d.n_rows == 1);

      uint32_t N = phi.n_rows; // x_d.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

      for (uint32_t n=0; n<N; ++n){
        //phi(n,i)=0.0;
        for (uint32_t i=0; i<T; ++i) {
          phi(n,i) = eLogSig_gam(i); //ElogSigma(gamma,i);
          for (uint32_t k=0; k<K; ++k) {
            phi(n,i) += zeta(i,k)* eLogBeta(k,n); //ElogBeta(lambda,k,x_d(n)) ;
          }
        }
        normalizeLogDistribution(phi.row(n));
      }
    }

    void computeNaturalGradients(DistriContainer<U>& d_lambda, Mat<double>& d_a, const Mat<double>& zeta, const Mat<double>&  phi, double omega, uint32_t D, const Mat<U>& x_d)
    {
//      uint32_t N = x_d.n_cols;
//      uint32_t Nw = d_lambda.n_cols;
      uint32_t T = zeta.n_rows;
      uint32_t K = zeta.n_cols;

//      d_lambda.init(HDP<U>::mH0,mK);

//      d_lambda.zeros();
      d_a.zeros();
      for (uint32_t k=0; k<K; ++k)
      { // for all K corpus level topics
//        cout<<zeta.col(k).t()<<endl;
//        cout<<phi;
        d_lambda[k]->posteriorHDP_var(zeta.col(k),phi,D,x_d);

//        for (uint32_t i=0; i<T; ++i)
//        {
//          Row<double> _lambda(Nw); _lambda.zeros();
//          for (uint32_t n=0; n<N; ++n){
//            _lambda(x_d(n)) += phi(n,i);
//          }
//          d_lambda.row(k) += zeta(i,k) * _lambda;
//        }
//        d_lambda.row(k) = D*d_lambda.row(k);
//        //cout<<"lambda-nu="<<d_lambda[k].t()<<endl;
//        d_lambda.row(k) += ((Dir*)(&mH0))->mAlphas;
//        //cout<<"lambda="<<d_lambda[k].t()<<endl;

        for (uint32_t i=0; i<T; ++i)
        {
          d_a(k,0) += zeta(i,k);
          for (uint32_t l=k+1; l<K; ++l) {
            d_a(k,1) += zeta(i,l);
          }
        }
        d_a(k,0) = D*d_a(k,0)+1.0;
        d_a(k,1) = D*d_a(k,1)+omega;
      }
    }

    //bool normalizeLogDistribution(Row<double>& r)
    bool normalizeLogDistribution(arma::subview_row<double> r) const
    {
      // known as the log sum exp trick!
//      cout<<" r="<<r<<endl;
      double maxR = as_scalar(max(r));
//      cout<<"  maxR="<<maxR<<endl;
//      cout<<"  exp(r-maxR)="<<exp(r-maxR)<<endl;

      r -= maxR + log(sum(exp(r-maxR)));
      r=exp(r);

//      cout<<" r="<<r<<endl;
      return true;

//      double minR = as_scalar(min(r));
//      cout<<" minR="<<minR<<endl;
//      if(minR > -100.0) {
//        cout<<" logDenom="<<sum(exp(r),1)<<endl;
//        double denom = as_scalar(sum(exp(r),1));
//        cout<<" logDenom="<<denom<<endl;
//        r -= log(denom); // avoid division by 0
//        cout<<" r - logDenom="<<r<<endl;
//        r = exp(r);
//        r /= sum(r);
//        cout<<" exp(r - logDenom)="<<r<<endl;
//        return true;
//      }else{ // cannot compute this -> set the smallest r to 1.0 and the rest to 0
//        double maxR = as_scalar(max(r));
//        cout<<"maxR="<<maxR<<" <-" <<arma::max(r) <<endl;
//        uint32_t kMax=as_scalar(find(r==maxR,1));
//        cout<<"maxR="<<maxR<<" kMax="<<kMax<<" <-" <<arma::max(r) <<endl;
//        r.zeros();
//        r(kMax)=1.0;
//        cout<<" r ="<<r<<endl;
//        return false;
//      }
    }

};


///=======================================================================
// Copyright 2015-2020 Clemson University
// Authors: Bradley S. Meyer
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
#ifndef RANK_SPANNING_BRANCHINGS_HPP
#define RANK_SPANNING_BRANCHINGS_HPP

/*
 * Rank spanning branchings
 *         Camerini et al Algorithm
 *
 * Requirement:
 *      directed graph with single root vertex
 */

#include <vector>
#include <limits>

#include <boost/concept_check.hpp>
#include <boost/config.hpp>
#include <boost/foreach.hpp>
#include <boost/graph/properties.hpp>
#include <boost/graph/filtered_graph.hpp>
#include <boost/graph/named_function_params.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/graph/depth_first_search.hpp>
#include <boost/heap/fibonacci_heap.hpp>
#include <boost/pending/disjoint_sets.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/optional.hpp>
#include <boost/unordered_set.hpp>

using namespace boost;

namespace rsb {

  namespace detail {

    typedef adjacency_list < vecS, vecS, bidirectionalS,
      no_property, no_property > BranchingGraph;

    typedef graph_traits<BranchingGraph>::vertex_descriptor BranchingVertex;

    // Structure to store edges and compare them by weight.

    template <typename Edge, typename WeightMap, typename Compare>
    struct EdgeNode
    {
      Edge                                                            edge;
      typename property_traits<WeightMap>::value_type                 weight;
      Compare                                                         compare;
      EdgeNode(){}
      EdgeNode(
        const Edge& e,
        const typename property_traits<WeightMap>::value_type & w,
        Compare c ) :
        edge( e ), weight( w ), compare( c ) {}
      bool operator<( EdgeNode const & rhs ) const
      { return compare( weight, rhs.weight ); }
    };

    // Insert edges from graph into queue, taking into account constraints.

    template <typename Graph, typename Edge, typename WeightMap,
              typename Compare, typename MergablePriorityQueueMap>
    bool
    insert_edges
    (
      const Graph& g,
      WeightMap& weight_map,
      Compare& comp,
      MergablePriorityQueueMap& in_edges,
      const unordered_set<Edge>& include_edges,
      const unordered_set<Edge>& exclude_edges
    )
    {
      typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
      unordered_set<Vertex> include_vertices, in_set, out_set;
      typename unordered_set<Vertex>::iterator it;

      // Insert vertices in sets to check for spanning branching.

      BGL_FORALL_VERTICES_T( v, g, Graph )
      {
        in_set.insert( v );
        out_set.insert( v );
      }

      // Insert edges that must be present and note the invertex.  Remove
      // vertices from relevant sets.

      BOOST_FOREACH( const Edge& e, include_edges )
      {
        include_vertices.insert( target( e, g ) );
        in_edges[target( e, g )].push(
          EdgeNode<Edge,WeightMap,Compare>( e, get( weight_map, e ), comp )
        );
        it = in_set.find( target( e, g ) );
        if( it != in_set.end() ) { in_set.erase( it ); }
        it = out_set.find( source( e, g ) );
        if( it != out_set.end() ) { out_set.erase( it ); }
      }

      // Insert edges, but not edges to be excluded or into vertices that
      // already have in edges.  Remove vertices from relevant sets.

      BOOST_FOREACH( const Edge &e, edges(g) )
      {
        if(
          include_vertices.find( target(e, g) ) == include_vertices.end()
          &&
          exclude_edges.find( e ) == exclude_edges.end()
        )
        {
          if( source(e, g) != target(e, g) )
          {
            in_edges[target(e,g)].push(
              EdgeNode<Edge,WeightMap,Compare>( e, get( weight_map, e ), comp )
            );
            it = in_set.find( target( e, g ) );
            if( it != in_set.end() ) { in_set.erase( it ); }
            it = out_set.find( source( e, g ) );
            if( it != out_set.end() ) { out_set.erase( it ); }
          }
        }
      }

      // Check for correct number of roots.
      //
      if( in_set.size() != 1 )
        return false;   // Zero roots or more than one root.
      else if( out_set.find( *in_set.begin() ) != out_set.end() )
        return false;   // Root is isolated.
      else
        return true;

    }

    // Retrieve the path back from leaf to root in cycle branching.

    void
    find_back_path(
      BranchingGraph& cycle_branching,
      std::vector<BranchingVertex>& bv
    )
    {

      BGL_FORALL_INEDGES( bv[0], e, cycle_branching, BranchingGraph )
      {
        bv.insert( bv.begin(), source( e, cycle_branching ) );
        find_back_path( cycle_branching, bv );
      }

    }

    // Expand cycles.

    template <typename Graph, typename Edge, typename IndexMap,
              typename WeightMap, typename Compare>
    void
    expand(
      const Graph& g,
      IndexMap& v_id,
      BranchingGraph& cycle_branching,
      unordered_set<BranchingVertex>& root_set,
      std::vector<EdgeNode< Edge, WeightMap, Compare> >& beta
    )
    {

      BOOST_FOREACH( BranchingVertex v, root_set )
      {
        if(
          in_degree( v, cycle_branching ) == 0 &&
          out_degree( v, cycle_branching ) != 0
        )
        {
          std::vector<BranchingVertex> bv;
          bv.push_back( v_id[target( beta[v].edge, g )] );
          find_back_path( cycle_branching, bv );
          for( std::size_t i = 0; i < bv.size() - 1; i++ )
          {
            beta[bv[i+1]] = beta[bv[i]];
            clear_vertex( bv[i], cycle_branching );
          }
        }
      }

      // Remove isolated vertices.

      std::vector<BranchingVertex> vertices_to_remove_from_set;

      BOOST_FOREACH( BranchingVertex v, root_set )
      {
        if(
           in_degree( v, cycle_branching ) == 0 &&
           out_degree( v, cycle_branching ) == 0
        )
        {
          vertices_to_remove_from_set.push_back( v );
        }
      }

      BOOST_FOREACH( BranchingVertex v, vertices_to_remove_from_set )
      {
        root_set.erase( v );
      }

    }

    // Camerini et al. BEST routine.

    template <typename EdgeNodeType, typename MergablePriorityQueue,
              typename Graph, typename Edge, typename IndexMap,
              typename WeightMap, typename Rank, typename Pred,
              typename Compare>
    void
    best_spanning_branching( const Graph& g,
                             unordered_set<Edge>& branching,
                             IndexMap& v_id,
                             WeightMap& weight_map,
                             Compare& comp,
                             Rank rank,
                             Pred pred1,
                             Pred pred2,
                             unordered_set<Edge>& include_edges,
                             unordered_set<Edge>& exclude_edges
                           )
    {

      // Define types.

      typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
      typedef typename graph_traits<Graph>::vertices_size_type vertex_idx_t;

      typedef std::map<Vertex, EdgeNodeType> exit_map_t;

      // Create various objects.  Note in particular the two disjoint
      // sets.  The set of weakly connected components is used to determine
      // cycles.  The set of strongly connected components is used to
      // represent supervertices (condensed cycles).

      unordered_set<Vertex> unvisited_vertex_set;

      Vertex root_vertex;

      BranchingGraph cycle_branching;

      vertex_idx_t n = num_vertices(g);

      std::map<Vertex, MergablePriorityQueue> in_edges;

      // Create disjoint sets.  The set of weakly connected components is
      // used to determine cycles.  The set of strongly connected components
      // is used to represent supervertices (condensed cycles).

      disjoint_sets<Rank, Pred>
        weak_cc( rank, pred1 ), strong_cc( rank, pred2 );

      if(
        !insert_edges(
          g,
          weight_map,
          comp,
          in_edges,
          include_edges,
          exclude_edges
        )
      ) return;

      std::vector<EdgeNodeType> beta( 2 * n );

      std::map<Vertex, vertex_idx_t> parent;

      BGL_FORALL_VERTICES_T( v, g, Graph )
      {
        weak_cc.make_set( v );
        strong_cc.make_set( v );
        parent[v]  = v_id[v];
        add_vertex( cycle_branching );
        unvisited_vertex_set.insert( v );
      }

      while( !unvisited_vertex_set.empty() )
      {

        typename unordered_set<Vertex>::iterator it =
          unvisited_vertex_set.begin();

        if( in_edges[*it].empty() )
        {
          root_vertex = *it;
	  unvisited_vertex_set.erase( it );
        }
        else
        {

	  EdgeNodeType critical_edge_node = in_edges[*it].top();

	  beta[parent[*it]] = critical_edge_node;

          // Done with this vertex.

	  unvisited_vertex_set.erase( it );

          // Check for cycle and, if present, condense.

	  if(
	    weak_cc.find_set( source( critical_edge_node.edge, g ) ) !=
	    weak_cc.find_set( target( critical_edge_node.edge, g ) )
	  )
	  {
	    weak_cc.union_set(
              source( critical_edge_node.edge, g ),
              target( critical_edge_node.edge, g )
	    );
	  }
	  else
	  {
	    BranchingVertex v_new = add_vertex( cycle_branching );

	    EdgeNodeType least_costly_edge_node = critical_edge_node;

	    boost::unordered_set<Vertex> cycle_vertex_set;

	    for(
              Vertex v = source( critical_edge_node.edge, g );
	      cycle_vertex_set.find( strong_cc.find_set( v ) ) ==
                cycle_vertex_set.end();
	      v = source( beta[parent[strong_cc.find_set( v )]].edge, g )
	    )
	    {
              Vertex u = strong_cc.find_set( v );
	      cycle_vertex_set.insert( u );
	      add_edge( v_new, parent[u], cycle_branching );
              if(
                comp( beta[parent[u]].weight, least_costly_edge_node.weight )
              )
              {
		least_costly_edge_node = beta[parent[u]];
              }
	    }

            Vertex new_repr = *cycle_vertex_set.begin();

	    BOOST_FOREACH( Vertex u, cycle_vertex_set )
	    {
              strong_cc.link( u, new_repr );
              new_repr = strong_cc.find_set( new_repr );
            }

	    BOOST_FOREACH( Vertex v, cycle_vertex_set )
	    {
              exit_map_t v_exit;
              BOOST_FOREACH( EdgeNodeType en, in_edges[v] )
              {
                if( strong_cc.find_set( source( en.edge, g ) ) != new_repr )
                {
                  en.weight += least_costly_edge_node.weight -
                      beta[parent[v]].weight;
                  Vertex u = strong_cc.find_set( source( en.edge, g ) );
                  if( v_exit.find(u) != v_exit.end() )
                  {
                    if( comp( v_exit[u].weight, en.weight ) )
                    {
                      v_exit[u] = en;
                    }
                  }
                  else
                  {
                    v_exit[u] = en;
                  }
                }
	      }
              MergablePriorityQueue tmp_queue;
              BOOST_FOREACH( typename exit_map_t::value_type& t, v_exit )
              {
                tmp_queue.push( t.second );
              }
              in_edges[v].swap( tmp_queue );
	    }

	    BOOST_FOREACH( Vertex v, cycle_vertex_set )
	    {
              if( v != new_repr ) in_edges[new_repr].merge( in_edges[v] );
            }

	    unvisited_vertex_set.insert( new_repr );

            parent[new_repr] = v_new;

	  }

        }

      }

      // Create a set containing possible roots of the cycle branching.
      // In each cycle expansion, remove isolated vertices of the
      // cycle branching to avoid considering them in subsequent cycle
      // expansions.

      unordered_set<BranchingVertex> root_set;

      BGL_FORALL_VERTICES( u, cycle_branching, BranchingGraph )
      {
        root_set.insert( u );
      }

      while( !root_set.empty() )
      {
        expand( g, v_id, cycle_branching, root_set, beta );
      }

      BGL_FORALL_VERTICES_T( v, g, Graph )
      {
        if( v != root_vertex )
        {
          branching.insert( beta[v_id[v]].edge );
        }
      }

    }

    // Depth-first visitor to set up pre-order and post-order maps.

    template<typename OrderMap>
    class dfs_order_visitor:public default_dfs_visitor
    {

      public:
	dfs_order_visitor(
	  OrderMap& pr,
	  OrderMap& po
	) :  m_pr( pr ), m_po( po ) { td = 0; tf = 0; }

	template<typename Vertex, typename Graph>
	void discover_vertex(const Vertex u, const Graph & g)
	{
	  m_pr[u] = td++;
	}

	template <typename Vertex, typename Graph>
	void finish_vertex(const Vertex u, const Graph & g)
	{
	  m_po[u] = tf++;
	}

	OrderMap& m_pr;
	OrderMap& m_po;

      private:
	std::size_t td;
	std::size_t tf;

    };

    // The ancestor checker.  A vertex u is a proper ancestor of
    // vertex v (in the parent branching) if pr[u] < pr[v] and po[u] > po[v].

    template<typename OrderMap>
    struct ancestor_checker
    {
      OrderMap& m_pr;
      OrderMap& m_po;
      ancestor_checker( OrderMap& pr, OrderMap& po ) :
	m_pr( pr ), m_po( po ) {}

      template<typename Vertex>
      bool operator()( const Vertex& v1, const Vertex& v2 ) const
      { return ( m_pr[v1] < m_pr[v2] && m_po[v1] > m_po[v2] ); }
    };

    // Create a new branching from an input edge set.

    template<typename Graph, typename Edge, typename IndexMap>
    BranchingGraph
    create_branching_graph_from_edge_set(
      Graph& g,
      IndexMap& v_id,
      const unordered_set<Edge>& branching
    )
    {

      BranchingGraph new_branching;

      for( size_t i = 0; i < num_vertices( g ); i++ )
      {
        add_vertex( new_branching );
      }

      BOOST_FOREACH( const Edge& e, branching )
      {
        add_edge(
          v_id[source( e, g )], v_id[target( e, g )], new_branching
        );
      }

      return new_branching;

    }

    // Camerini et al. SEEK routine.  Find the in edge that, when removed,
    //  gives the next best branching for a vertex and the weight difference.

    template <typename Graph, typename Edge, typename IndexMap,
              typename MergablePriorityQueue, typename EdgeNodeType,
              typename Compare, typename OrderMap, typename OptionalWeight>
    void
    seek_next_edge_weight_diff(
      Graph& g,
      MergablePriorityQueue& in_edges,
      IndexMap& v_id,
      ancestor_checker<OrderMap>& is_ancestor,
      const unordered_set<Edge>& branching,
      EdgeNodeType& b,
      Compare& comp,
      Edge& return_edge,
      OptionalWeight& delta
    )
    {

      if( branching.find( b.edge ) != branching.end() )
      {
        for(
          typename MergablePriorityQueue::ordered_iterator ei =
            in_edges.ordered_begin();
          ei != in_edges.ordered_end();
          ei++
        )
        {
          if( (*ei).edge != b.edge )
          {
            if(
              !is_ancestor(
                v_id[target( b.edge, g )],
                v_id[source( (*ei).edge, g )]
              )
            )
            {
              if( !delta || comp( b.weight - (*ei).weight, delta.get() ) )
              {
                delta = b.weight - (*ei).weight;
                return_edge = b.edge;
                break;
              }
            }
          }
        }
      }
    }

    // Camerini et al. NEXT routine.  Find the edge that, when removed,
    // gives the next best branching and the resulting weight difference.

    template <typename EdgeNodeType, typename MergablePriorityQueue,
              typename Graph, typename Edge, typename IndexMap,
              typename WeightMap, typename Rank, typename Pred,
              typename Compare>
    boost::optional<
      std::pair<Edge, typename property_traits<WeightMap>::value_type>
    >
    next_spanning_branching( const Graph& g,
                             const unordered_set<Edge>& branching,
                             IndexMap& v_id,
                             WeightMap& weight_map,
                             Compare& comp,
                             Rank rank,
                             Pred pred1,
                             Pred pred2,
                             unordered_set<Edge>& include_edges,
                             unordered_set<Edge>& exclude_edges
                           )
    {

      typedef typename graph_traits<Graph>::vertex_descriptor Vertex;

      typedef std::map<Vertex, EdgeNodeType> exit_map_t;

      typedef typename property_traits<WeightMap>::value_type weight_t;

      Edge return_edge;

      boost::optional<weight_t> delta;

      EdgeNodeType b;

      // Create various objects.

      boost::optional<
        std::pair<Edge, typename property_traits<WeightMap>::value_type>
      > next_edge_and_weight_delta;

      unordered_set<Vertex> unvisited_vertex_set;

      std::map<Vertex, MergablePriorityQueue> in_edges;

      std::map<Vertex, EdgeNodeType> max_e;

      // Create disjoint sets.  The set of weakly connected components is
      // used to determine cycles.  The set of strongly connected components
      // is used to represent supervertices (condensed cycles).

      disjoint_sets<Rank, Pred>
        weak_cc( rank, pred1 ), strong_cc( rank, pred2 );

      // Create a branching graph to check whether a vertex is an ancestor of
      // another vertex in the branching.

      BranchingGraph ancestor_branching =
        create_branching_graph_from_edge_set( g, v_id, branching );

      // Insert edges.

      if(
        !insert_edges(
          g,
          weight_map,
          comp,
          in_edges,
          include_edges,
          exclude_edges
        )
      )
      {
        return next_edge_and_weight_delta;
      }

      // Initialize data structures and find start vertex for depth
      // first search.

      BranchingVertex start_vertex;

      BGL_FORALL_VERTICES_T( v, g, Graph )
      {
        weak_cc.make_set( v );
        strong_cc.make_set( v );
        unvisited_vertex_set.insert( v );
        if( in_edges[v].empty() ) { start_vertex = v_id[v]; }
      }

      // Create the ancestor checker from ancestor_branching.

      typedef std::map<BranchingVertex, std::size_t> OrderMap;

      OrderMap pr;
      OrderMap po;

      dfs_order_visitor<OrderMap> vis(pr, po);

      depth_first_search(
        ancestor_branching, visitor(vis).root_vertex( start_vertex )
      );

      ancestor_checker<OrderMap> is_ancestor(pr, po);

      // Main loop.

      while( !unvisited_vertex_set.empty() )
      {

        typename unordered_set<Vertex>::iterator it =
          unvisited_vertex_set.begin();

        if( in_edges[*it].empty() )
        {
	  unvisited_vertex_set.erase( it );
        }
        else
        {

          // Get largest in edge with ties solved in favor of edges in
          // input branching.

          for(
            typename MergablePriorityQueue::ordered_iterator ei =
              in_edges[*it].ordered_begin();
            ei != in_edges[*it].ordered_end();
            ei++
          )
          {
            if( comp( (*ei).weight, in_edges[*it].top().weight ) )
              break;
            b = *ei;
            if( branching.find( b.edge ) != branching.end() )
              break;
          }

	  max_e[*it] = b;

          // Seek next edge weight difference.

          seek_next_edge_weight_diff(
            g, in_edges[*it], v_id, is_ancestor, branching, b, comp,
            return_edge, delta
          );

          // Done with this vertex.

	  unvisited_vertex_set.erase( it );

          // Check for cycle and, if present, condense.

	  if(
	    weak_cc.find_set( source( b.edge, g ) ) !=
	    weak_cc.find_set( target( b.edge, g ) )
	  )
	  {
	    weak_cc.union_set( source( b.edge, g ), target( b.edge, g ) );
	  }
	  else
	  {

	    EdgeNodeType least_costly_edge_node = b;

	    boost::unordered_set<Vertex> cycle_vertex_set;

	    for(
              Vertex v = source( b.edge, g );
	      cycle_vertex_set.find( strong_cc.find_set( v ) ) ==
                cycle_vertex_set.end();
	      v = source( max_e[strong_cc.find_set( v )].edge, g )
	    )
	    {
              Vertex u = strong_cc.find_set( v );
	      cycle_vertex_set.insert( u );
	      if( comp( max_e[u].weight, least_costly_edge_node.weight ) )
              {
		least_costly_edge_node = max_e[u];
              }
	    }

            Vertex new_repr = *cycle_vertex_set.begin();

	    BOOST_FOREACH( Vertex v, cycle_vertex_set )
	    {
              strong_cc.link( v, new_repr );
              new_repr = strong_cc.find_set( new_repr );
            }

            // Adjust arc weights and remove parallel arcs.  Keep the
            // the largest weight in arc and the largest viable alternative
            // arc from each source outside the cycle.  Make sure that
            // an arc from the branching is among the added edges, if present.

	    BOOST_FOREACH( Vertex v, cycle_vertex_set )
	    {
              std::vector<exit_map_t> v_exit(3);
              BOOST_FOREACH( EdgeNodeType en, in_edges[v] )
              {
                if( strong_cc.find_set( source( en.edge, g ) ) != new_repr )
                {
                  en.weight += least_costly_edge_node.weight - max_e[v].weight;
                  Vertex u = strong_cc.find_set( source( en.edge, g ) );
                  if( branching.find( en.edge ) != branching.end() )
                  {
                    v_exit[0][u] = en;
                  }
                  else
                  {
                    if( v_exit[1].find(u) != v_exit[1].end() )
                    {
                      if( comp( v_exit[1][u].weight, en.weight ) )
                      {
                        if(
                          !is_ancestor(
                             v_id[target( v_exit[1][u].edge, g )],
                             v_id[source( v_exit[1][u].edge, g )]
                          )
                        )
                        {
                          v_exit[2][u] = v_exit[1][u];
                        }
                        v_exit[1][u] = en;
                      }
                      else if(
                        v_exit[2].find(u) == v_exit[2].end() ||
                        comp( v_exit[2][u].weight, en.weight )
                      )
                      {
                        if(
                          !is_ancestor(
                            v_id[target( en.edge, g )],
                            v_id[source( en.edge, g )]
                          )
                        )
                        {
                          v_exit[2][u] = en;
                        }
                      }
                    }
                    else
                    {
                      v_exit[1][u] = en;
                    }
                  }
                }
	      }
              MergablePriorityQueue tmp_queue;
              BOOST_FOREACH( exit_map_t& exit_map, v_exit )
              {
                BOOST_FOREACH( typename exit_map_t::value_type& t, exit_map )
                {
                  tmp_queue.push( t.second );
                }
              }
              in_edges[v].swap( tmp_queue );
	    }

	    BOOST_FOREACH( Vertex v, cycle_vertex_set )
	    {
              if( v != new_repr ) in_edges[new_repr].merge( in_edges[v] );
            }

	    unvisited_vertex_set.insert( new_repr );

	  }

        }

      }

      if( delta )
      {
        next_edge_and_weight_delta = std::make_pair( return_edge, delta.get() );
      }

      return next_edge_and_weight_delta;

    }

    // Class to filter graph.

    template<typename EdgeSet>
    class branching_filter
    {

      public:
        branching_filter(){}
        branching_filter( const EdgeSet * _es ) : p_es( _es ){}

        template <typename Edge>
        bool operator()( const Edge& e ) const
        {
          if( p_es->find( e ) != p_es->end() )
            return true;
          else
            return false;
        }

     private:
       const EdgeSet * p_es;

    };

    // Structure to store branchings and compare them by weight.

    template<typename Edge, typename WeightMap, typename Compare>
    struct BranchingEntry
    {
      Edge                                                        edge;
      typename property_traits<WeightMap>::value_type             weight;
      Compare                                                     compare;
      unordered_set<Edge>                                         branching;
      unordered_set<Edge>                                         include_edges;
      unordered_set<Edge>                                         exclude_edges;
      BranchingEntry(){}
      BranchingEntry(
	const typename property_traits<WeightMap>::value_type& w,
	const Edge& e,
	Compare& comp,
	const unordered_set<Edge>& b,
	const unordered_set<Edge>& include,
	const unordered_set<Edge>& exclude
      ) :
	  edge( e ), weight( w ), compare( comp ),
	  branching( b ), include_edges( include ),
	  exclude_edges( exclude ){}
      bool operator<( BranchingEntry const & rhs ) const
      { return compare( weight, rhs.weight ); }
    };

    // Compute the weight of a branching.

    template<typename WeightMap, typename Edge>
    typename property_traits<WeightMap>::value_type
    compute_branching_weight(
      WeightMap& w,
      const unordered_set<Edge>& branching
    )
    {

       typedef typename property_traits<WeightMap>::value_type weight_t;

       boost::optional<weight_t> weight;

       BOOST_FOREACH( const Edge& e, branching )
       {

         if( !weight )
         {
           weight = get( w, e );
         }
         else
         {
	   weight = weight.get() + get( w, e );
         }
       }

       return weight.get();

    }

    // Routine implementation.

    template <template<class...> class PriorityQueue,
              typename Graph, typename BranchingProcessor, typename IndexMap,
              typename WeightMap, typename Compare,
              typename Rank, typename Parent>
    void
    rank_spanning_branchings_impl( const Graph& g,
			           BranchingProcessor bp,
                                   IndexMap v_id,
                                   WeightMap w,
			           Compare comp,
                                   Rank rank,
                                   Parent pred1,
                                   Parent pred2
			         )
    {

      typedef typename graph_traits<Graph>::vertex_descriptor Vertex;
      typedef typename graph_traits<Graph>::edge_descriptor Edge;

      typedef typename property_traits<WeightMap>::value_type weight_t;

      typedef BranchingEntry<Edge, WeightMap, Compare> branching_entry_t;

      typedef EdgeNode<Edge, WeightMap, Compare> edge_node_t;

      typedef PriorityQueue<branching_entry_t> branching_queue_t;

      typedef PriorityQueue<edge_node_t> edge_node_queue_t;

      BOOST_CONCEPT_ASSERT(( VertexAndEdgeListGraphConcept<Graph> ));

      BOOST_CONCEPT_ASSERT(( ReadablePropertyMapConcept<IndexMap, Vertex> ));
      BOOST_CONCEPT_ASSERT(( ReadablePropertyMapConcept<WeightMap, Edge> ));

      BOOST_CONCEPT_ASSERT(( heap::PriorityQueue<branching_queue_t> ));
      BOOST_CONCEPT_ASSERT(( heap::MergablePriorityQueue<edge_node_queue_t> ));

      unordered_set<Edge> best_branching, empty_set;

      boost::optional<std::pair<Edge, weight_t> > next_edge_and_weight_delta;

      Edge e;

      branching_queue_t branching_queue;

      best_spanning_branching<edge_node_t, edge_node_queue_t>
                             ( g,
                               best_branching,
                               v_id,
                               w,
                               comp,
                               rank,
                               pred1,
                               pred2,
                               empty_set,
                               empty_set
                             );

      branching_filter<unordered_set<Edge> > filter1( &best_branching );
      filtered_graph<Graph, branching_filter<unordered_set<Edge> > >
        fg1( g, filter1 );

      if( !bp( fg1 ) ) return;

      next_edge_and_weight_delta =
        next_spanning_branching<edge_node_t, edge_node_queue_t>
                              ( g,
                                 best_branching,
                                  v_id,
                                  w,
                                  comp,
                                  rank,
                                  pred1,
                                  pred2,
                                  empty_set,
                                  empty_set
                               );

      if( next_edge_and_weight_delta )
      {
	branching_queue.push(
          branching_entry_t(
	    compute_branching_weight( w, best_branching ) -
              (next_edge_and_weight_delta.get()).second,
	    (next_edge_and_weight_delta.get()).first,
	    comp,
	    best_branching,
	    empty_set,
	    empty_set
	  )
	);
      }
      else
      {
	return;
      }

      while( !branching_queue.empty() )
      {

	unordered_set<Edge> branching;

	branching_entry_t P = branching_queue.top();

	branching_queue.pop();

	unordered_set<Edge> include_edges = P.include_edges;

	unordered_set<Edge> exclude_edges = P.exclude_edges;

	include_edges.insert( P.edge );

	exclude_edges.insert( P.edge );

	best_spanning_branching<edge_node_t, edge_node_queue_t>
                              ( g,
                                 branching,
                                 v_id,
                                 w,
                                 comp,
                                 rank,
                                 pred1,
                                 pred2,
                                 P.include_edges,
                                 exclude_edges
                               );

        branching_filter<unordered_set<Edge> > filter( &branching );
        filtered_graph<Graph, branching_filter<unordered_set<Edge> > >
          fg( g, filter );

        if( !bp( fg ) ) return;

	next_edge_and_weight_delta =
	  next_spanning_branching<edge_node_t, edge_node_queue_t>
                                ( g,
                                   P.branching,
                                   v_id,
                                   w,
                                   comp,
                                   rank,
                                   pred1,
                                   pred2,
                                   include_edges,
                                   P.exclude_edges
                                 );

	if( next_edge_and_weight_delta )
	{
	  branching_queue.push(
            branching_entry_t(
	      compute_branching_weight( w, P.branching ) -
                (next_edge_and_weight_delta.get()).second,
	      (next_edge_and_weight_delta.get()).first,
	      comp,
	      P.branching,
	      include_edges,
	      P.exclude_edges
	    )
	  );
	}

	next_edge_and_weight_delta =
	  next_spanning_branching<edge_node_t, edge_node_queue_t>
                                ( g,
                                   branching,
                                   v_id,
                                   w,
                                   comp,
                                   rank,
                                   pred1,
                                   pred2,
                                   P.include_edges,
                                   exclude_edges
                                 );

	if( next_edge_and_weight_delta )
	{
	  branching_queue.push(
            branching_entry_t(
	      P.weight - (next_edge_and_weight_delta.get()).second,
	      (next_edge_and_weight_delta.get()).first,
	      comp,
	      branching,
	      P.include_edges,
	      exclude_edges
	    )
	  );
	}

      }

    }

    template <template<class...> class PriorityQueue,
              typename Graph, typename BranchingProcessor, typename IndexMap,
              typename WeightMap, typename Compare>
    void
    rank_spanning_branchings_dispatch2( const Graph& g,
                                        BranchingProcessor bp,
                                        IndexMap id_map,
                                        WeightMap weight_map,
                                        Compare compare
                                      )
    {

      typename graph_traits<Graph>::vertices_size_type n = num_vertices(g);

      if( num_vertices( g ) == 0 ) return; // Nothing to do.

      typedef typename graph_traits<Graph>::vertices_size_type vertex_idx_t;
      typedef typename graph_traits<Graph>::vertex_descriptor Vertex;

      // Set up rank and parent for disjoint sets.

      std::vector<vertex_idx_t> rank( n );

      std::vector<Vertex> pred1( n ), pred2( n );

      rank_spanning_branchings_impl<PriorityQueue>(
        g,
        bp,
        id_map,
        weight_map,
        compare,
        make_iterator_property_map( rank.begin(), id_map, rank[0] ),
        make_iterator_property_map( pred1.begin(), id_map, pred1[0] ),
        make_iterator_property_map( pred2.begin(), id_map, pred2[0])
      );

    }

    template <template<class...> class PriorityQueue,
              typename Graph, typename BranchingProcessor, typename Compare,
              typename P, typename T, typename R>
    void rank_spanning_branchings_dispatch1(
      const Graph& g,
      BranchingProcessor bp,
      Compare compare,
      const bgl_named_params<P, T, R>& params
    )
    {

      detail::rank_spanning_branchings_dispatch2<PriorityQueue>(
        g,
        bp,
        choose_param(
          get_param( params, vertex_index_t()), get( vertex_index, g )
        ),
        choose_param(
          get_param( params, edge_weight_t()), get( edge_weight, g )
        ),
        compare
      );

    }

    template <template<class...> class PriorityQueue,
              typename Graph, typename BranchingProcessor,
              typename P, typename T, typename R>
    void rank_spanning_branchings_dispatch1(
      const Graph& g,
      BranchingProcessor bp,
      param_not_found,
      const bgl_named_params<P, T, R>& params
    )
    {

      typedef
        typename
        property_traits<
          typename property_map<Graph, edge_weight_t>::const_type
        >::value_type weight_t;

      BOOST_CONCEPT_ASSERT(( ComparableConcept<weight_t> ));

      detail::rank_spanning_branchings_dispatch2<PriorityQueue>(
        g,
        bp,
        choose_param(
          get_param( params, vertex_index_t()), get( vertex_index, g )
        ),
        choose_param(
          get_param( params, edge_weight_t()), get( edge_weight, g )
        ),
        std::less<weight_t>()
      );

    }

  } // namespace detail

  template <template<class...> class PriorityQueue = heap::fibonacci_heap,
            typename Graph, typename BranchingProcessor,
            typename P, typename T, typename R>
  void
  inline rank_spanning_branchings(
    const Graph& g,
    BranchingProcessor bp,
    const bgl_named_params<P, T, R>& params
  )
  {

    detail::rank_spanning_branchings_dispatch1<PriorityQueue>(
      g,
      bp,
      get_param( params, distance_compare_t() ),
      params
    );

  }

  template <template<class...> class PriorityQueue = heap::fibonacci_heap,
            typename Graph, typename BranchingProcessor>
  void
  inline rank_spanning_branchings( const Graph& g,
                            BranchingProcessor bp
                          )
  {

    bgl_named_params<int,int> params(0);
    rank_spanning_branchings<PriorityQueue>( g, bp, params );

  }

} // namespace rsb

#endif // RANK_SPANNING_BRANCHINGS_HPP


#pragma once

#include <map>

#include <cvode/cvode.h>               /* prototypes for CVODE fcts., consts.  */
#include <nvector/nvector_serial.h>    /* access to serial N_Vector            */
#include <sunmatrix/sunmatrix_dense.h> /* access to dense SUNMatrix            */
#include <sunlinsol/sunlinsol_dense.h> /* access to dense SUNLinearSolver      */
#include <sundials/sundials_types.h>   /* defs. of realtype, sunindextype      */

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

#include <autodiff/forward/real.hpp>
#include <autodiff/forward/real/eigen.hpp>
#include "types.hh"
//#define Ith(v,i)    NV_Ith_S(v,i-1)         /* Ith numbers components 1..NEQ */
//#define IJth(A,i,j) SM_ELEMENT_D(A,i-1,j-1) /* IJth numbers rows,cols 1..NEQ */

using namespace types;

namespace cvode_wrapper
{

std::map<std::string_view, int> cv_lmm = {{"BDF", CV_BDF}, {"ADAMS", CV_ADAMS}};

struct cv_options
{
    // Integrator settings
    std::string step_method = "BDF";
    double      rtol        = 1.0E-4;   // relative tolerance
    double      atol        = 1.0E-14;   // absolute tolerance
    int         max_steps   = 500;      // max number of steps between outputs
    int         max_order   = 5;        // max order of linear multistep
    double      t0          = 0.0;      // inital time
    double      tf          = 1.0;      // final time
    double      step_min    = 0.0;      // minimum step size
    double      step_max    = 0.0;      // maximum step size (0 -> no max size)

    // Linear solver and preconditioner settings
    double      epslin      = 1.0E-6;    // linear solver tolerance factor

    // output options
    int         output_lvl  = 0;
};

static int check_retval(void *returnvalue, const char *funcname, int opt)
{
  int *retval;

  /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
  if (opt == 0 && returnvalue == NULL) {
    fmt::print(stderr, "\nSUNDIALS_ERROR: {} failed - returned NULL pointer\n\n",
	    funcname);
    return(1); }

  /* Check if retval < 0 */
  else if (opt == 1) {
    retval = (int *) returnvalue;
    if (*retval < 0) {
      fmt::print(stderr, "\nSUNDIALS_ERROR: {} failed with retval = %d\n\n",
	      funcname, *retval);
      return(1); }}

  /* Check if function returned NULL pointer - no memory allocated */
  else if (opt == 2 && returnvalue == NULL) {
    fmt::print(stderr, "\nMEMORY_ERROR: {} failed - returned NULL pointer\n\n",
	    funcname);
    return(1); }

  return(0);
}

/* compare the solution at the final time 4e10s to a reference solution computed
   using a relative tolerance of 1e-8 and absoltue tolerance of 1e-14 */
static int check_ans(N_Vector y, realtype t, realtype rtol, N_Vector atol)
{
  int      passfail=0;        /* answer pass (0) or fail (1) retval */
  N_Vector ref;               /* reference solution vector        */
  N_Vector ewt;               /* error weight vector              */
  realtype err;               /* wrms error                       */
  realtype ONE=RCONST(1.0);
  realtype ZERO=RCONST(0.0);

  /* create reference solution and error weight vectors */
  ref = N_VClone(y);
  ewt = N_VClone(y);

  /* set the reference solution data */
  NV_Ith_S(ref,0) = RCONST(5.2083495894337328e-08);
  NV_Ith_S(ref,1) = RCONST(2.0833399429795671e-13);
  NV_Ith_S(ref,2) = RCONST(9.9999994791629776e-01);

  /* compute the error weight vector, loosen atol */
  N_VAbs(ref, ewt);
  N_VLinearSum(rtol, ewt, RCONST(10.0), atol, ewt);
  if (N_VMin(ewt) <= ZERO) {
    fmt::print(stderr, "\nSUNDIALS_ERROR: check_ans failed - ewt <= 0\n\n");
    return(-1);
  }
  fmt::print(stdout, "\n{} {} {}\n", NV_Ith_S(ewt,0),NV_Ith_S(ewt,1),NV_Ith_S(ewt,2));
  N_VInv(ewt, ewt);

  /* compute the solution error */
  N_VLinearSum(ONE, y, -ONE, ref, ref);
  err = N_VWrmsNorm(ref, ewt);

  /* is the solution within the tolerances? */
  passfail = (err < ONE) ? 0 : 1;

  if (passfail) {
    fmt::print(stdout, "\nSUNDIALS_WARNING: check_ans error={}\n\n", err);
  }

  /* Free vectors */
  N_VDestroy(ref);
  N_VDestroy(ewt);

  return(passfail);
}

template<class System>
struct cvode_stepper
{
    // options
    cv_options options;

    // cvode context pointer
    void *cvode_mem;

    size_t N;           // system size
    N_Vector cv_y, abst;     // ongoing solution
    SUNMatrix A;        //
    SUNLinearSolver LS;

    cvode_stepper(cv_options opts) : options(std::move(opts)), cvode_mem(nullptr)
    {

    }

    ~cvode_stepper()
    {
        CVodeFree(&cvode_mem);
        N_VDestroy(cv_y); N_VDestroy(abst);
        SUNLinSolFree(LS);
        SUNMatDestroy(A);
    }

    auto copy_nvect(const N_Vector a, vector_t& b)
    {
        for(int i = 0; i < N; ++i)
        {
            b(i) = NV_Ith_S(a, i);
        }
    }

    auto initialize(vector_t& y0)
    {
        // set system size
        N = y0.rows();

        // allocate solution vector
        //y = N_VNew_Serial(N);

        //cv_y = N_VMake_Serial(N, static_cast<Eigen::VectorXd>(y0).data());
        cv_y = N_VMake_Serial(N, y0.data());
        //cv_y = N_VNew_Serial(N);

        // fill solution vector
        //copy_evect(y0, cv_y);

        // linear solver data
        // TODO: implement sparse
        A = SUNDenseMatrix(N, N);
        LS = SUNLinSol_Dense(cv_y, A);

        // init a new cvode context
        cvode_mem = CVodeCreate(cv_lmm[options.step_method]);
        // init integrator memory

        CVodeInit(cvode_mem, [](realtype t, N_Vector y, N_Vector ydot, void* user_data)
            {
                System* sys = static_cast<System*>(user_data);
                vector_map_t yref(NV_DATA_S(y), NV_LENGTH_S(y));
                vector_map_t fref(NV_DATA_S(ydot), NV_LENGTH_S(ydot));

                fref = sys->f(yref, t);

                return 0;

            },
        options.t0, cv_y);
        // set max step attempts
        CVodeSetMaxNumSteps(cvode_mem, options.max_steps);
        // set max/min step sizes
        CVodeSetMinStep(cvode_mem, options.step_min);
        CVodeSetMaxStep(cvode_mem, options.step_max);
        // set max order
        CVodeSetMaxOrd(cvode_mem, options.max_order);
        // set tolarances
        // TODO: allow vector tolerance
        abst = N_VNew_Serial(N);
        NV_Ith_S(abst, 0) = 1.0E-8; NV_Ith_S(abst, 1) = 1.0E-14; NV_Ith_S(abst, 2) = 1.0E-6;
        CVodeSVtolerances(cvode_mem, options.rtol, abst);
        //CVodeSStolerances(cvode_mem, options.rtol, options.atol);
        // set linear solver
        CVodeSetLinearSolver(cvode_mem, LS, A);
        // set jacobin
        CVodeSetJacFn(cvode_mem, [](realtype t, N_Vector y, N_Vector fy, SUNMatrix J,
            void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3)
            {
                System* sys = static_cast<System*>(user_data);

                vector_map_t yref(NV_DATA_S(y), NV_LENGTH_S(y));
                matrix_map_t Jref(SM_DATA_D(J), SM_ROWS_D(J), SM_COLUMNS_D(J));

                Jref = sys->J(yref, t);
                return 0;
            }
        );

    }

    auto print_sol(double sol_time)
    {
        fmt::print("At t = {} y = {} {} {}\n", sol_time, NV_Ith_S(cv_y,0), NV_Ith_S(cv_y,1),NV_Ith_S(cv_y,2));
    }

    auto step()
    {

    }

    auto check_retval(void *returnvalue, const char *funcname, int opt)
    {
        int *retval;

        /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
        if (opt == 0 && returnvalue == NULL) {
            fmt::print(stderr, "\nSUNDIALS_ERROR: {} failed - returned NULL pointer\n\n",
                funcname);
            return(1); }

        /* Check if retval < 0 */
        else if (opt == 1) {
            retval = (int *) returnvalue;
            if (*retval < 0) {
            fmt::print(stderr, "\nSUNDIALS_ERROR: {} failed with retval = %d\n\n",
                funcname, *retval);
            return(1); }}

        /* Check if function returned NULL pointer - no memory allocated */
        else if (opt == 2 && returnvalue == NULL) {
            fmt::print(stderr, "\nMEMORY_ERROR: {} failed - returned NULL pointer\n\n",
                funcname);
            return(1); }

        return(0);
    }
    auto letsgo()
    {
        auto iout = 0;
        auto tout = 0.4;
        double sol_time;
        while(1)
        {
            auto rstep = CVode(cvode_mem, tout, cv_y, &sol_time, CV_NORMAL);
            print_sol(sol_time);
            if (check_retval(&rstep, "CVode", 1)) break;
            if (rstep == CV_SUCCESS) {
                iout++;
                tout *= 10.0;
            }
        if (iout == 12) break;
        }
        check_ans(cv_y,sol_time,options.rtol,abst);
    }

};

}


//==============================================================================
//         Copyright 2014 - Jean-Thierry Lapresté
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2013   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 NT2_LINALG_FUNCTIONS_FACTORIZATIONS_SCHUR_HPP_INCLUDED
#define NT2_LINALG_FUNCTIONS_FACTORIZATIONS_SCHUR_HPP_INCLUDED

#include <nt2/include/functions/schur.hpp>
#include <boost/dispatch/attributes.hpp>
#include <nt2/core/container/table/table.hpp>
#include <nt2/include/functions/geesx.hpp>
#include <nt2/include/functions/geesx_no_w.hpp>
#include <nt2/include/functions/geesx1.hpp>
#include <nt2/include/functions/geesxw.hpp>
#include <nt2/include/functions/height.hpp>
#include <nt2/include/functions/isreal.hpp>
#include <nt2/include/functions/of_size.hpp>
#include <nt2/include/functions/real.hpp>
#include <nt2/include/functions/resize.hpp>
#include <nt2/include/functions/tie.hpp>
#include <nt2/include/functions/issquare.hpp>
#include <nt2/sdk/meta/as_real.hpp>
#include <nt2/sdk/complex/meta/is_complex.hpp>
#include <nt2/linalg/options.hpp>
#include <nt2/linalg/details/utility/lapack_verify.hpp>
#include <nt2/core/container/dsl/as_terminal.hpp>
#include <nt2/core/utility/assign_swap.hpp>
#include <boost/assert.hpp>
#include <boost/mpl/bool.hpp>
#include <boost/mpl/if.hpp>
#include <boost/type_traits/is_same.hpp>
#include <boost/dispatch/meta/strip.hpp>

namespace nt2 { namespace ext
{
  //============================================================================
  //SCHUR Scalar
  //============================================================================
  BOOST_DISPATCH_IMPLEMENT  ( schur_, tag::cpu_
                            , (A0)
                            , (scalar_<unspecified_<A0> >)
                            )
  {
    typedef A0 result_type;

    BOOST_FORCEINLINE result_type operator()(const A0& a0) const
    {
      return a0;
    }
  };

  BOOST_DISPATCH_IMPLEMENT  ( schur_, tag::cpu_
                            , (A0)(A1)
                            , (scalar_<unspecified_<A0> >)
                              (unspecified_<A1>)
                            )
  {
    typedef A0                                                      otype_t;
    typedef typename nt2::meta::as_real<A0>::type                   rtype_t;
    typedef typename meta::strip<A1>::type                            opt_t;
    typedef typename boost::mpl::if_ < boost::is_same< opt_t, nt2::policy<ext::real_> >
                              , rtype_t
                              , otype_t
                              >::type                           result_type;

    BOOST_FORCEINLINE result_type operator()(const A0& a0, const A1&) const
    {
      return eval(a0, opt_t());
    }
  private :
    template < class T >
    BOOST_FORCEINLINE result_type eval( const A0& a0, const T& ) const
    {
      return a0;
    }

    BOOST_FORCEINLINE result_type eval( const A0& a0, nt2::policy<ext::real_> const & ) const
    {
      BOOST_ASSERT_MSG(isreal(a0),
                       "all input matrix elements are to be real to support"
                       "'real_' option with complex type input");
      return real(a0);
    }

  };

  //============================================================================
  //SCHUR
  //============================================================================
  BOOST_DISPATCH_IMPLEMENT  ( schur_, tag::cpu_
                            , (A0)(N0)(A1)(N1)
                            , ((node_<A0, nt2::tag::schur_
                                    , N0, nt2::container::domain
                                      >
                              ))
                              ((node_<A1, nt2::tag::tie_
                                    , N1, nt2::container::domain
                                     >
                              ))
                            )
  {
    typedef void  result_type;
    typedef typename boost::proto::result_of::child_c<A0&,0>::value_type child0;
    typedef typename child0::value_type                                  type_t;
    typedef typename nt2::meta::as_complex<type_t>::type                ctype_t;
    typedef typename nt2::meta::as_real<type_t>::type                   rtype_t;
    typedef rtype_t T;
    typedef nt2::memory::container<tag::table_,  type_t, nt2::_2D>   o_semantic;
    typedef nt2::memory::container<tag::table_, rtype_t, nt2::_2D>   r_semantic;
    typedef nt2::memory::container<tag::table_, ctype_t, nt2::_2D>   c_semantic;

    BOOST_FORCEINLINE result_type operator()( A0& a0, A1& a1 ) const
    {
      BOOST_ASSERT_MSG(issquare(boost::proto::child_c<0>(a0)),"first input must be square");
      eval(a0, a1, N0(), N1());
    }
  private:
    //==========================================================================
    /// INTERNAL ONLY - T = SCHUR(A)
    // returns the schur matrix
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<1> const&
              , boost::mpl::long_<1> const&
              ) const
    {
      NT2_AS_TERMINAL_INOUT(o_semantic, a, boost::proto::child_c<0>(a0), boost::proto::child_c<0>(a1));
      NT2_LAPACK_VERIFY(nt2::geesx(boost::proto::value(a), type_t(0)));
      boost::proto::child_c<0>(a1) = a;
    }

    //==========================================================================
    /// INTERNAL ONLY - T = SCHUR(A, cmplx_/real_)
    // returns T in its real or complex form or only the eigenvalues
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<2> const&
              , boost::mpl::long_<1> const&
              ) const
    {
       eval1_2(a0, a1,
               boost::proto::value(boost::proto::child_c<1>(a0)));
    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval1_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::real_>
                 ) const
    {
      BOOST_ASSERT_MSG(isreal(boost::proto::child_c<0>(a0)),
                       "all input matrix elements are to be real to support"
                       "'real_' option with complex type input");
      NT2_AS_TERMINAL_INOUT(r_semantic, t
                           , real(boost::proto::child_c<0>(a0))
                           , boost::proto::child_c<0>(a1));
      // Here one cannot be sure that boost::proto::child_c<0>(a1) is real typed
      // as the 'real'schur decomposition can be put in a complex table
      // so we cannot pass it to geesx that needs a real table
      NT2_LAPACK_VERIFY(nt2::geesx(boost::proto::value(t), rtype_t(0)));
      assign_swap(boost::proto::child_c<0>(a1), t);
    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval1_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::eigs_>
                 ) const
    {
      nt2::container::table<type_t> work;
      NT2_AS_TERMINAL_INOUT(o_semantic, t
                           , boost::proto::child_c<0>(a0), work);
      NT2_AS_TERMINAL_OUT  (c_semantic, w
                           , boost::proto::child_c<0>(a1));
      NT2_LAPACK_VERIFY( nt2::geesxw(boost::proto::value(t),
                                     boost::proto::value(w)));
      w.resize(of_size(height(t), 1));
      assign_swap(boost::proto::child_c<0>(a1), w);
   }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval1_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::cmplx_>
                 ) const
    {
      typedef typename meta::is_complex<type_t>::type is_cmplx_t;
      eval1_2c(a0, a1, is_cmplx_t());
    }

    BOOST_FORCEINLINE
    void eval1_2c ( A0& a0, A1& a1
                 , boost::mpl::false_
                 ) const
    {
      // Here boost::proto::child_c<0>(a0) is real and geesx has to receive a complex table
      NT2_AS_TERMINAL_INOUT(c_semantic, t, boost::proto::child_c<0>(a0), boost::proto::child_c<0>(a1));
      NT2_LAPACK_VERIFY(nt2::geesx(boost::proto::value(t), ctype_t(0)));
      assign_swap(boost::proto::child_c<0>(a1), t);
    }

    BOOST_FORCEINLINE
    void eval1_2c ( A0& a0, A1& a1
                 , boost::mpl::true_
                 ) const
    {
      // Here  boost::proto::child_c<0>(a0) is complex
      NT2_AS_TERMINAL_INOUT(c_semantic, t
                           , boost::proto::child_c<0>(a0), boost::proto::child_c<0>(a1));
      NT2_LAPACK_VERIFY(nt2::geesx(boost::proto::value(t), ctype_t(0)));
      assign_swap(boost::proto::child_c<0>(a1), t);
     }

    //==========================================================================
    /// INTERNAL ONLY - [U, T]= SCHUR(A)
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<1> const&
              , boost::mpl::long_<2> const&
              ) const
    {
      NT2_AS_TERMINAL_INOUT(o_semantic, t
                           , boost::proto::child_c<0>(a0), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (o_semantic, u
                           , boost::proto::child_c<0>(a1));
      size_t n = height(t);
      u.resize(of_size(n, n));
      NT2_LAPACK_VERIFY(nt2::geesx_no_w( boost::proto::value(t)
                                       , boost::proto::value(u)
                                   ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
    }

    //==========================================================================
    /// INTERNAL ONLY - [T, U]= SCHUR(A, sort_/cmplx_/real_ )
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<2> const&
              , boost::mpl::long_<2> const&
              ) const
    {
      eval2_2( a0, a1
             , boost::proto::value(boost::proto::child_c<1>(a0)));

    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval2_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::real_>
                 ) const
    {
      BOOST_ASSERT_MSG(isreal(boost::proto::child_c<0>(a0)),
                       "all input matrix elements are to be real to support "
                       "'real_' option with complex type input");
      NT2_AS_TERMINAL_INOUT(r_semantic, t
                           , real(boost::proto::child_c<0>(a0)), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (r_semantic, u
                           , boost::proto::child_c<0>(a1));
      size_t n = height(t);
      u.resize(of_size(n, n));
      NT2_LAPACK_VERIFY(nt2::geesx_no_w( boost::proto::value(t)
                                       , boost::proto::value(u)
                                   ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval2_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::cmplx_>
                 ) const
    {
      NT2_AS_TERMINAL_INOUT(c_semantic, t
                           , boost::proto::child_c<0>(a0), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (c_semantic, u
                           , boost::proto::child_c<0>(a1));
      size_t n = nt2::height(t);
      u.resize(of_size(n, n));
      NT2_LAPACK_VERIFY(nt2::geesx_no_w( boost::proto::value(t)
                                       , boost::proto::value(u)
                                   ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
    }

    //==========================================================================
    /// INTERNAL ONLY - [U, T, W]= SCHUR(A)
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<1> const&
              , boost::mpl::long_<3> const&
              ) const
    {
      NT2_AS_TERMINAL_INOUT(o_semantic, t
                           , boost::proto::child_c<0>(a0), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (o_semantic, u
                           , boost::proto::child_c<0>(a1));
      NT2_AS_TERMINAL_OUT  (c_semantic, w
                           , boost::proto::child_c<2>(a1));
      size_t n =   height(t);
      u.resize(of_size(n, n));
      w.resize(of_size(n, 1));
      NT2_LAPACK_VERIFY(nt2::geesx( boost::proto::value(t)
                                  , boost::proto::value(w)
                                  , boost::proto::value(u)
                              ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
      assign_swap(boost::proto::child_c<2>(a1), w);

    }
    //==========================================================================
    /// INTERNAL ONLY - [U, T, W]= SCHUR(A, real_/complex_)
    BOOST_FORCEINLINE
    void eval ( A0& a0, A1& a1
              , boost::mpl::long_<2> const&
              , boost::mpl::long_<3> const&
              ) const
    {
      eval3_2(a0, a1,
              boost::proto::value(boost::proto::child_c<1>(a0)));
    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval3_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::cmplx_>
                 ) const
    {
      NT2_AS_TERMINAL_INOUT(c_semantic, t
                           , boost::proto::child_c<0>(a0), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (c_semantic, u
                           , boost::proto::child_c<0>(a1));
      NT2_AS_TERMINAL_OUT  (c_semantic, w
                           , boost::proto::child_c<2>(a1));
      size_t n = nt2::height(t);
      u.resize(of_size(n, n));
      w.resize(of_size(n, 1));
      NT2_LAPACK_VERIFY(nt2::geesx( boost::proto::value(t)
                                  , boost::proto::value(w)
                                  , boost::proto::value(u)
                                  ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
      assign_swap(boost::proto::child_c<2>(a1), w);
    }

    //==========================================================================
    /// INTERNAL ONLY
    BOOST_FORCEINLINE
    void eval3_2 ( A0& a0, A1& a1
                 , nt2::policy<ext::real_>
                 ) const
    {
      BOOST_ASSERT_MSG(isreal(boost::proto::child_c<0>(a0)),
                       "all input matrix elements are to be real to support "
                       "'real_' option with complex type input");
      NT2_AS_TERMINAL_INOUT(r_semantic, t
                           , real(boost::proto::child_c<0>(a0)), boost::proto::child_c<1>(a1));
      NT2_AS_TERMINAL_OUT  (r_semantic, u
                           , boost::proto::child_c<0>(a1));
      NT2_AS_TERMINAL_OUT  (c_semantic, w
                           , boost::proto::child_c<2>(a1));
      size_t n = nt2::height(t);
      u.resize(of_size(n, n));
      w.resize(of_size(n, 1));
      NT2_LAPACK_VERIFY(nt2::geesx( boost::proto::value(t)
                                  , boost::proto::value(w)
                                  , boost::proto::value(u)
                                  ));
      assign_swap(boost::proto::child_c<0>(a1), u);
      assign_swap(boost::proto::child_c<1>(a1), t);
      assign_swap(boost::proto::child_c<2>(a1), w);
    }
  };

} }

#endif