#include"main.hpp"
#include"Option.hpp"
#include"Mcmc.hpp"
#include"Pre.hpp"
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/random/gamma_distribution.hpp>

boost::mt19937 r;

const double sigma2_y0 = 10000;

double quantile(const vector<double>& data,const double f)
{
    vector<double> sorted_data(data);
    sort(sorted_data.begin(),sorted_data.end());
    const double index = f * (data.size() - 1) ;
    const size_t lhs = static_cast<size_t>(index) ;
    const double delta = index - lhs ;
    double result;
    if (data.size() == 0) throw runtime_error("quantile");
    if (lhs == data.size() - 1) {
        result = sorted_data.at(lhs) ;
    } else {
        result = (1 - delta) * sorted_data.at(lhs) + delta * sorted_data.at(lhs + 1) ;
    }
    return result ;
}

bool Ey_fn(const int p, const int j, const vector<double>& y0i,const vector<double>& Ri, vector<double>& Ey_ij,const Mcmc& mc)
{
    Ey_ij.at(0) = y0i.at(j);
    for (int t=1;t<mc.op().nt();t++) {
        Ey_ij.at(t) = Ey_ij.at(t-1)
            + (mc.op().timep().at(t)-mc.op().timep().at(t-1))*(exp(mc.xx().at(p).at(j*mc.op().nt()+t-1))*Ri.at(t-1)-Ey_ij.at(t-1)*(1-Ri.at(t-1)));
        if (Ey_ij.at(t)<=0) return true;
    }
    return false;
}

double sum_sq(const int p, const int j, const vector<double>& y0i, const vector<double>& Ri,const Mcmc& mc)
{
    double sum2=0;
    vector<double> Ey_ij(mc.op().nt());
    if (Ey_fn(p,j,y0i,Ri,Ey_ij,mc)) return NAN;
    for (int t=0; t<mc.op().nt(); t++) for (unsigned q=0;q<mc.yy().at(p).size();q++) {
        sum2 += pow(mc.yy().at(p).at(q).at(j*mc.op().nt()+t)-log(Ey_ij.at(t)),2);
    }
    return sum2;
}

double sum_sq(const int p, const vector<double>& y0i, const vector<double>& Ri,const Mcmc& mc,const int s=0)
{
    double sum2=0;
    for (int j=0; j<mc.op().nrep(); j++) {
        bool flag=false;
        for (unsigned q=0;q<mc.yy().at(p).size();q++) if (obs(mc.yy().at(p).at(q).at(j*mc.op().nt()))) {
            flag=true;
            break;
        }
        if (flag) {
            vector<double> Ey_ij(mc.op().nt());
            if (Ey_fn(p,j,y0i,Ri,Ey_ij,mc)) return NAN;
            for (int t=s; t<mc.op().nt(); t++) for (unsigned q=0;q<mc.yy().at(p).size();q++) {
                sum2 += pow(mc.yy().at(p).at(q).at(j*mc.op().nt()+t)-log(Ey_ij.at(t)),2);
            }
        }
    }
    return sum2;
}

boost::random::normal_distribution<> ugaussian(0,1);
boost::uniform_real<> uniform(0,1);
void Mcmc::update_y0(const int p)
{
    for (int j=0;j<op().nrep();j++) {
        bool flag=false;
        for (unsigned q=0;q<yy().at(p).size();q++) if (obs(yy().at(p).at(q).at(j*op().nt()))) {
            flag=true;
            break;
        }
        if (flag) {
            vector<double> newy0p(y0().at(p));
            newy0p.at(j) *= exp(ugaussian(r));
            double logr = -sum_sq(p,j,newy0p,RR().at(p),*this)/tau2().at(p)/2;
            logr -= -sum_sq(p,j,y0().at(p),RR().at(p),*this)/tau2().at(p)/2;
            logr += -pow(log(newy0p.at(j)),2)/2/sigma2_y0;
            logr -= -pow(log(y0().at(p).at(j)),2)/2/sigma2_y0;
            if (logr>0 or uniform(r)<exp(logr)) {
                d_y0.at(p).at(j) = newy0p.at(j);
            }
        }
    }
}

void Mcmc::update_R(const int p)
{
    vector<double> newRp;
    int s=-1;
    for (int t=0;t<op().nt()-1;t++) {
        if (CPS().at(p).at(t)==1) {
            s = t+1;
            newRp = RR().at(p);
            newRp.at(t) = inv_logit(logit(RR().at(p).at(t))+ugaussian(r));
        } else {
            newRp.at(t) = newRp.at(t-1);
        }
        if (CPS().at(p).at(t+1)==0) continue;
        double logr = -sum_sq(p,y0().at(p),newRp,*this,s)/tau2().at(p)/2;
        logr -= -sum_sq(p,y0().at(p),RR().at(p),*this,s)/tau2().at(p)/2;
        logr += log(pq(newRp.at(t)))-log(pq(RR().at(p).at(t)));
        if (logr>0 or uniform(r)<exp(logr)) {
            d_RR.at(p) = newRp;
        }
    }
}

void Mcmc::update_CPS(const int p)
{
    const vector<double>& Rp=RR().at(p);
    for (int cp=1;cp<op().nt()-1;cp++) {
        vector<double> newRp(Rp);
        int l_time_index=-1, r_time_index=-1;
        for (int t=cp-1;l_time_index<0;t--) if (CPS().at(p).at(t)==1) l_time_index=t;
        for (int t=cp+1;r_time_index<0;t++) if (CPS().at(p).at(t)==1) r_time_index=t;
        const double l_time = op().timep().at(cp)-op().timep().at(l_time_index);
        const double r_time = op().timep().at(r_time_index)-op().timep().at(cp);
        const double t_time = op().timep().at(r_time_index)-op().timep().at(l_time_index);
        double logr=0;
        if (CPS().at(p).at(cp)==0) {
            const double u = uniform(r);
            const double newRleft = inv_logit(logit(Rp.at(cp))+r_time/t_time*logit(u));
            const double newRright = inv_logit(logit(Rp.at(cp))-l_time/t_time*logit(u));
            for (int t=l_time_index;t<cp;t++) newRp.at(t) = newRleft;
            for (int t=cp;t<r_time_index;t++) newRp.at(t) = newRright;
            if (d_prior) logr += d_varphi.at(p).at(cp);
            //if (d_prior) logr += varphi_fn(p,cp,*this,d_mo);
            logr += 2*log(newRleft*(1-newRright)+newRright*(1-newRleft))-log(pq(Rp.at(cp)));
        } else {
            const double mergedR = inv_logit((l_time*logit(Rp.at(cp-1))+r_time*logit(Rp.at(cp)))/t_time);
            for (int t=l_time_index;t<r_time_index;t++) newRp.at(t)=mergedR;
            if (d_prior) logr -= d_varphi.at(p).at(cp);
            //if (d_prior) logr -= varphi_fn(p,cp,*this,d_mo);
            logr -= 2*log(Rp.at(cp-1)*(1-Rp.at(cp))+Rp.at(cp)*(1-Rp.at(cp-1)))-log(pq(mergedR));
        }
        logr += -sum_sq(p,y0().at(p),newRp,*this,l_time_index+1)/tau2().at(p)/2;
        logr -= -sum_sq(p,y0().at(p),Rp,*this,l_time_index+1)/tau2().at(p)/2;
        if (logr>0 or uniform(r)<exp(logr)) {
            d_RR.at(p) = newRp;
            d_CPS.at(p).at(cp)=1-CPS().at(p).at(cp);
        }
    }
}

void Mcmc::tau_hyperparams()
{

        d_tau2.assign(yy().size(),0);
            for (unsigned p=0;p<yy().size();p++) {
                for (unsigned q=0; q<yy().at(p).size(); q++) {
                    for (int j=0;j<op().nrep();j++) {
                        vector<double> mu(op().nt()-1);
                        for (int t=0;t<op().nt()-1;t++) {
                            mu.at(t)=yy().at(p).at(q).at(j*op().nt()+t)+yy().at(p).at(q).at(j*op().nt()+t+1);
                            mu.at(t)/=2;
                        }
                        for (int t=1;t<op().nt()-1;t++) {
                            d_tau2.at(p)+=.5*pow(yy().at(p).at(q).at(j*op().nt()+t)-mu.at(t-1),2);
                            d_tau2.at(p)+=.5*pow(yy().at(p).at(q).at(j*op().nt()+t)-mu.at(t),2);
                        }
                        d_tau2.at(p)+=pow(yy().at(p).at(q).at(j*op().nt()+0)-mu.at(0),2);
                        const int t=op().nt()-2;
                        d_tau2.at(p)+=pow(yy().at(p).at(q).at(j*op().nt()+t)-mu.at(t),2);
                    }
                }
                d_tau2.at(p) /= op().nc()*yy().at(p).size();
            }
            const double m=quantile(d_tau2,.1);
            for (unsigned p=0;p<yy().size();p++) d_tau2.at(p) += m;
        ofstream var_ofs("varc.txt");
        copy(tau2().begin(),tau2().end(),ostream_iterator<double>(var_ofs,"\n"));
        const double mom1=accumulate(tau2().begin(),tau2().end(),0.)/tau2().size();
        double mom2=0;
        for (vector<double>::const_iterator it=tau2().begin();it!=tau2().end();it++) mom2 += (*it)*(*it);
        mom2 /= tau2().size();
        d_a_tau=(2*mom2-mom1*mom1)/(mom2-mom1*mom1);//MOM estimates
        d_b_tau=mom1*mom2/(mom2-mom1*mom1);         //MOM estimates
}

Mcmc::Mcmc(const Pre& pr) : d_op(pr.op()),d_prior(false),d_varphi(vector<vector<double> >()),d_xx(pr.xx()),d_yy(pr.yy())//,d_nobs(pr.nobs())
{
    Construct();
}

Mcmc::Mcmc(const Pre& pr,const vector<vector<double> >& varphi) : d_op(pr.op()),d_prior(true),d_varphi(varphi),d_xx(pr.xx()),d_yy(pr.yy())//,d_nobs(pr.nobs())
{
    Construct();
}

void Mcmc::Construct()
{
    d_y0.assign(yy().size(),vector<double>(op().nrep()));
    for (unsigned p=0;p<yy().size();p++) for (int j=0;j<op().nrep();j++) {
        double sum=0;
        for (unsigned q=0;q<yy().at(p).size();q++) sum+=yy().at(p).at(q).at(j*op().nt());
        d_y0.at(p).at(j)=exp(sum/yy().at(p).size());
    }
    d_RR.assign(yy().size(),vector<double> (op().nt(),.9));
    d_CPS.assign(yy().size(),vector<int> (op().nt(),1));
    d_CPSum.assign(yy().size(),vector<int> (op().nt()));

    tau_hyperparams();
    ofstream y0_ofs("s_y0");
    ofstream RR_ofs("s_RR");
    ofstream tau2_ofs("s_tau2");
    if (not tau2_ofs) throw runtime_error("can't open s_tau2");
    ofstream CPS_ofs("s_CPS");
    ofstream xRyD_ofs("s_xRyD");
    ofstream loglike_ofs("s_loglike");

    for (int iter=0,ns=0; ns<op().ns(); iter++) {

        if (iter%100==0) cout<<'\r'<<iter<<flush;
        for (unsigned p=0; p<yy().size(); p++) {
            update_y0(p);
            update_R(p);
            update_CPS(p);
        }

        if (iter>op().nburn() and iter%op().nthin()==0) {
            ns++;

            for (unsigned p=0;p<yy().size();p++) {
                tau2_ofs<<d_tau2.at(p)<<'\t';
                for (int j=0;j<op().nrep();j++) y0_ofs<<d_y0.at(p).at(j)<<'\t';
                for (int t=0;t<op().nt()-1;t++) {
                    RR_ofs<<d_RR.at(p).at(t)<<'\t';
                }
                for (int t=1;t<op().nt()-1;t++) {
                    CPS_ofs<<d_CPS.at(p).at(t)<<'\t';
                    d_CPSum.at(p).at(t)+=d_CPS.at(p).at(t);
                }

            }

            y0_ofs<<'\n';
            tau2_ofs<<'\n';
            RR_ofs<<'\n';
            CPS_ofs<<'\n';

            double loglike = 0;
            for (unsigned p=0;p<yy().size();p++) for (unsigned q=0;q<yy().at(p).size();q++) for (int j=0;j<op().nrep();j++) {
                if (obs(yy().at(p).at(q).at(j*op().nt()))) {
                    vector<double> Ey_ij(op().nt());
                    if (Ey_fn(p,j,d_y0.at(p),d_RR.at(p),Ey_ij,*this)) throw runtime_error("fit");
                    for (int t=0;t<op().nt();t++) {
                        xRyD_ofs << log(Ey_ij.at(t)) << '\t';
                        //likelihood
                        loglike += -log(d_tau2.at(p))-pow(yy().at(p).at(q).at(j*op().nt()+t)-log(Ey_ij.at(t)),2)/d_tau2.at(p)/2;
                    }
                } else for (int t=0;t<op().nt();t++) xRyD_ofs << "NA\t";
            }
            xRyD_ofs<<'\n';
            loglike_ofs << loglike << '\n';
        }
    }
    cout<<'\n';
}


// Copyright (C) 1996-2016 by Synge Todo <wistaria@phys.s.u-tokyo.ac.jp>
//
// 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 MATH_POWER_HPP
#define MATH_POWER_HPP

#include <boost/type_traits/is_arithmetic.hpp>
#include <boost/utility/enable_if.hpp>
#include <complex>

namespace math {

using boost::is_arithmetic;
using boost::enable_if;
using boost::disable_if;

namespace detail {

template<typename T>
struct power_traits {
  typedef T power_type;
};

template<typename T>
struct power_traits<std::complex<T> > {
  typedef T power_type;
};

} // end namespace detail

//
// function power2 and p2
//

#ifndef BOOST_NO_SFINAE

template<typename T>
typename detail::power_traits<T>::power_type
power2(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
power2(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
p2(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power2(t); }

template<typename T>
typename detail::power_traits<T>::power_type
p2(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power2(t); }

#else

template<typename T>
typename detail::power_traits<T>::power_type
power2(T const& t) { return t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
p2(T const& t) { return power2(t); }

#endif

template<typename T>
typename detail::power_traits<std::complex<T> >::power_type
power2(std::complex<T> const& t) { return power2(real(t)) + power2(imag(t)); }

//
// function power3 and p3
//

#ifndef BOOST_NO_SFINAE

template<typename T>
typename detail::power_traits<T>::power_type
power3(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return t * t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
power3(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return t * t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
p3(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power3(t); }

template<typename T>
typename detail::power_traits<T>::power_type
p3(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power3(t); }

#else

template<typename T>
typename detail::power_traits<T>::power_type
power3(T const& t) { return t * t * t; }

template<typename T>
typename detail::power_traits<T>::power_type
p3(T const& t) { return power3(t); }

#endif

//
// function power4 and p4
//

#ifndef BOOST_NO_SFINAE

template<typename T>
typename detail::power_traits<T>::power_type
power4(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power2(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
power4(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power2(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
p4(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power4(t); }

template<typename T>
typename detail::power_traits<T>::power_type
p4(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power4(t); }

#else

template<typename T>
typename detail::power_traits<T>::power_type
power4(T const& t) { return power2(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
p4(T const& t) { return power4(t); }

#endif

//
// function power6 and p6
//

#ifndef BOOST_NO_SFINAE

template<typename T>
typename detail::power_traits<T>::power_type
power6(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power3(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
power6(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power3(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
p6(T t, typename enable_if<is_arithmetic<T> >::type* = 0) { return power6(t); }

template<typename T>
typename detail::power_traits<T>::power_type
p6(T const& t, typename disable_if<is_arithmetic<T> >::type* = 0) { return power6(t); }

#else

template<typename T>
typename detail::power_traits<T>::power_type
power6(T const& t) { return power3(power2(t)); }

template<typename T>
typename detail::power_traits<T>::power_type
p6(T const& t) { return power6(t); }

#endif

} // end namespace math

#endif // MATH_POWER_HPP


/*
	@author herumi

	tiny calculator 2
	This program generates a function to calc the value of
	polynomial given by user in run-time.
	use boost::spirit::qi
*/
#ifdef _WIN32
	#pragma warning(disable : 4127) // for boost(constant condition)
	#pragma warning(disable : 4512) // for boost
	#pragma warning(disable : 4819)
#endif
#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_container.hpp>
#include <boost/spirit/include/phoenix_bind.hpp>
#include <boost/timer.hpp>

#include <stdio.h>
#include <assert.h>
#include <string>
#include <vector>
#define XBYAK_NO_OP_NAMES
#include "xbyak/xbyak.h"

enum Operand {
	OpAdd,
	OpSub,
	OpMul,
	OpDiv,
	OpNeg,
	OpImm,
	OpVarX
};

struct Code {
	Operand op_;
	double val_;
	Code(Operand op)
		: op_(op)
		, val_(0)
	{
	}
	Code(double val)
		: op_(OpImm)
		, val_(val)
	{
	}
};

typedef std::vector<Code> CodeSet;

struct Vm {
	CodeSet code_;
	double operator()(double x) const
	{
		const size_t maxStack = 16;
		double stack[maxStack];
		double *p = stack;
		CodeSet::const_iterator pc = code_.begin();

		while (pc != code_.end()) {
			switch (pc->op_) {
			case OpVarX:
				*p++ = x;
				break;
			case OpImm:
				*p++ = pc->val_;
				break;
			case OpNeg:
				p[-1] = -p[-1];
				break;
			case OpAdd:
				--p;
				p[-1] += p[0];
				break;
			case OpSub:
				--p;
				p[-1] -= p[0];
				break;
			case OpMul:
				--p;
				p[-1] *= p[0];
				break;
			case OpDiv:
				--p;
				p[-1] /= p[0];
				break;
			}
			++pc;
			assert(p < stack + maxStack);
		}
		return p[-1];
	}
};

class Jit : public Xbyak::CodeGenerator {
private:
	enum {
		MAX_CONST_NUM = 32
	};
	MIE_ALIGN(16) double constTbl_[MAX_CONST_NUM];
	Xbyak::uint64_t negConst_;
	size_t constTblPos_;
#ifdef XBYAK32
	const Xbyak::Reg32& varTbl_;
	const Xbyak::Reg32& tbl_;
#else
	const Xbyak::Reg64& tbl_;
#endif
	int regIdx_;
public:
	/*
		double jit(double x);
		@note 32bit: x : [esp+4], return fp0
		      64bit: x [rcx](win), xmm0(gcc), return xmm0
	*/
	Jit()
		: negConst_(Xbyak::uint64_t(1) << 63)
		, constTblPos_(0)
#ifdef XBYAK32
		, varTbl_(eax)
		, tbl_(edx)
#elif defined(XBYAK64_WIN)
		, tbl_(rcx)
#else
		, tbl_(rdi)
#endif
		, regIdx_(-1)
	{
#ifdef XBYAK32
		lea(varTbl_, ptr [esp+4]);
#else
#ifdef XBYAK64_WIN
		movaps(ptr [rsp + 8], xm6); // save xm6, xm7
		movaps(ptr [rsp + 8 + 16], xm7);
#endif
		movaps(xm7, xm0); // save xm0
#endif
		mov(tbl_, (size_t)constTbl_);
	}
	void genPush(double n)
	{
		if (constTblPos_ >= MAX_CONST_NUM) throw;
		constTbl_[constTblPos_] = n;
		if (regIdx_ == 7) throw;
		movsd(Xbyak::Xmm(++regIdx_), ptr[tbl_ + constTblPos_ * sizeof(double)]);
		constTblPos_++;
	}
	void genVarX()
	{
#ifdef XBYAK32
		if (regIdx_ == 7) throw;
		movsd(Xbyak::Xmm(++regIdx_), ptr[varTbl_]);
#else
		if (regIdx_ == 6) throw;
		movsd(Xbyak::Xmm(++regIdx_), xm7);
#endif
	}
	void genAdd()
	{
		addsd(Xbyak::Xmm(regIdx_ - 1), Xbyak::Xmm(regIdx_)); regIdx_--;
	}
	void genSub()
	{
		subsd(Xbyak::Xmm(regIdx_ - 1), Xbyak::Xmm(regIdx_)); regIdx_--;
	}
	void genMul()
	{
		mulsd(Xbyak::Xmm(regIdx_ - 1), Xbyak::Xmm(regIdx_)); regIdx_--;
	}
	void genDiv()
	{
		divsd(Xbyak::Xmm(regIdx_ - 1), Xbyak::Xmm(regIdx_)); regIdx_--;
	}
	void genNeg()
	{
		xorpd(Xbyak::Xmm(regIdx_), ptr [tbl_ + MAX_CONST_NUM * sizeof(double)]);
	}
	void complete()
	{
#ifdef XBYAK32
		sub(esp, 8);
		movsd(ptr [esp], xm0);
		fld(qword [esp]);
		add(esp, 8);
#else
#ifdef XBYAK64_WIN
		movaps(xm6, ptr [rsp + 8]);
		movaps(xm7, ptr [rsp + 8 + 16]);
#endif
#endif
		ret();
	}
};

template<typename Iterator>
struct Parser : boost::spirit::qi::grammar<Iterator, boost::spirit::ascii::space_type> {
	boost::spirit::qi::rule<Iterator, boost::spirit::ascii::space_type> expression, term, factor;
	CodeSet& code_;
	Parser(CodeSet& code)
		: Parser::base_type(expression)
		, code_(code)
	{
		namespace qi = boost::spirit::qi;
		using namespace qi::labels;

		using boost::phoenix::ref;
		using boost::phoenix::push_back;

		expression = term >> *(('+' > term[push_back(ref(code_), OpAdd)])
							 | ('-' > term[push_back(ref(code_), OpSub)]));

		term = factor >> *(('*' > factor[push_back(ref(code_), OpMul)])
						 | ('/' > factor[push_back(ref(code_), OpDiv)]));

		factor = qi::double_[push_back(ref(code_), _1)]
				| qi::lit('x')[push_back(ref(code_), OpVarX)]
				| ('(' > expression > ')')
				| ('-' > factor[push_back(ref(code_), OpNeg)])
				| ('+' > factor);
	}
};

template<typename Iterator>
struct ParserJit : boost::spirit::qi::grammar<Iterator, boost::spirit::ascii::space_type> {
	boost::spirit::qi::rule<Iterator, boost::spirit::ascii::space_type> expression, term, factor;
	Jit code_;
	ParserJit()
		: ParserJit::base_type(expression)
	{
		namespace qi = boost::spirit::qi;
		using namespace qi::labels;

		using boost::phoenix::ref;
		using boost::phoenix::push_back;
		using boost::phoenix::bind;

		expression = term >> *(('+' > term[bind(&Jit::genAdd, ref(code_))])
							 | ('-' > term[bind(&Jit::genSub, ref(code_))]));

		term = factor >> *(('*' > factor[bind(&Jit::genMul, ref(code_))])
						 | ('/' > factor[bind(&Jit::genDiv, ref(code_))]));

		factor = qi::double_[bind(&Jit::genPush, ref(code_), _1)]
				| qi::lit('x')[bind(&Jit::genVarX, ref(code_))]
				| ('(' > expression > ')')
				| ('-' > factor[bind(&Jit::genNeg, ref(code_))])
				| ('+' > factor);
	}
};

template<class Func>
void Test(const char *msg, const Func& f)
{
	printf("%s:", msg);
	boost::timer t;
	double sum = 0;
	for (double x = 0; x < 1000; x += 0.0001) {
		sum += f(x);
	}
	printf("sum=%f, %fsec\n", sum, t.elapsed());
}

int main(int argc, char *argv[])
{
	if (argc < 2) {
		fprintf(stderr, "input formula\n");
		return 1;
	}
	const std::string str(argv[1]);

	try {
		Vm vm;
		Parser<std::string::const_iterator> parser(vm.code_);
		ParserJit<std::string::const_iterator> parserJit;

		const std::string::const_iterator end = str.end();

		std::string::const_iterator i = str.begin();
		if (!phrase_parse(i, end, parser, boost::spirit::ascii::space) || i != end) {
			puts("err 1");
			return 1;
		}
		printf("ret=%f\n", vm(2.3));

		i = str.begin();
		if (!phrase_parse(i, end, parserJit, boost::spirit::ascii::space) || i != end) {
			puts("err 2");
			return 1;
		}
		parserJit.code_.complete();
		double (*jit)(double) = parserJit.code_.getCode<double (*)(double)>();

		Test("VM ", vm);
		Test("JIT", jit);
	} catch (...) {
		fprintf(stderr, "err\n");
	}
}


///////////////////////////////////////////////////////////////////////////////
// BSD 3-Clause License
//
// Copyright (C) 2021, University of Edinburgh
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.
///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_CORE_NUMDIFF_RESIDUAL_HPP_
#define CROCODDYL_CORE_NUMDIFF_RESIDUAL_HPP_

#include <boost/function.hpp>
#include "crocoddyl/multibody/fwd.hpp"
#include "crocoddyl/core/residual-base.hpp"

namespace crocoddyl {

/**
 * @brief This class computes the numerical differentiation of a residual model.
 *
 * It computes Jacobian of the residual model via numerical differentiation, i.e., \f$\mathbf{R_x}\f$
 * and \f$\mathbf{R_u}\f$ which denote the Jacobians of the residual function
 * \f$\mathbf{r}(\mathbf{x},\mathbf{u})\f$.
 *
 * \sa `ResidualModelAbstractTpl()`, `calcDiff()`
 */
template <typename _Scalar>
class ResidualModelNumDiffTpl : public ResidualModelAbstractTpl<_Scalar> {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef _Scalar Scalar;
  typedef ResidualDataAbstractTpl<Scalar> ResidualDataAbstract;
  typedef ResidualModelAbstractTpl<Scalar> Base;
  typedef ResidualDataNumDiffTpl<Scalar> Data;
  typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;
  typedef MathBaseTpl<Scalar> MathBase;
  typedef typename MathBaseTpl<Scalar>::VectorXs VectorXs;
  typedef typename MathBaseTpl<Scalar>::MatrixXs MatrixXs;
  typedef boost::function<void(const VectorXs&, const VectorXs&)> ReevaluationFunction;

  /**
   * @brief Initialize the numdiff residual model
   *
   * @param model  Residual model that we want to apply the numerical differentiation
   */
  explicit ResidualModelNumDiffTpl(const boost::shared_ptr<Base>& model);

  /**
   * @brief Initialize the numdiff residual model
   */
  virtual ~ResidualModelNumDiffTpl();

  /**
   * @brief @copydoc Base::calc()
   */
  virtual void calc(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const VectorXs>& x,
                    const Eigen::Ref<const VectorXs>& u);

  /**
   * @brief @copydoc Base::calc(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const VectorXs>&
   * x)
   */
  virtual void calc(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const VectorXs>& x);

  /**
   * @brief @copydoc Base::calcDiff()
   */
  virtual void calcDiff(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const VectorXs>& x,
                        const Eigen::Ref<const VectorXs>& u);

  /**
   * @brief @copydoc Base::calcDiff(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const
   * VectorXs>& x)
   */
  virtual void calcDiff(const boost::shared_ptr<ResidualDataAbstract>& data, const Eigen::Ref<const VectorXs>& x);

  /**
   * @brief @copydoc Base::createData()
   */
  virtual boost::shared_ptr<ResidualDataAbstract> createData(DataCollectorAbstract* const data);

  /**
   * @brief Return the original residual model
   */
  const boost::shared_ptr<Base>& get_model() const;

  /**
   * @brief Return the disturbance value used by the numdiff routine
   */
  const Scalar get_disturbance() const;

  /**
   * @brief Modify the disturbance value used by the numdiff routine
   */
  void set_disturbance(const Scalar disturbance);

  /**
   * @brief Register functions that updates the shared data computed for a system rollout
   * The updated data is used to evaluate of the gradient and hessian.
   *
   * @param reevals are the registered functions.
   */
  void set_reevals(const std::vector<ReevaluationFunction>& reevals);

 protected:
  using Base::nu_;
  using Base::state_;
  using Base::unone_;

 private:
  /**
   * @brief Make sure that when we finite difference the residual model, the user
   * does not face unknown behaviour because of the finite differencing of a
   * quaternion around pi. This behaviour might occur if ResidualModelState and
   * FloatingInContact differential model are used together.
   *
   * For full discussions see issue
   * https://gepgitlab.laas.fr/loco-3d/crocoddyl/issues/139
   *
   * @param x is the state at which the check is performed.
   */
  void assertStableStateFD(const Eigen::Ref<const VectorXs>& /*x*/);

  boost::shared_ptr<Base> model_;              //!< Residual model hat we want to apply the numerical differentiation
  Scalar disturbance_;                         //!< Disturbance used in the numerical differentiation routine
  std::vector<ReevaluationFunction> reevals_;  //!< Functions that needs execution before calc or calcDiff
};

template <typename _Scalar>
struct ResidualDataNumDiffTpl : public ResidualDataAbstractTpl<_Scalar> {
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  typedef _Scalar Scalar;
  typedef MathBaseTpl<Scalar> MathBase;
  typedef ResidualDataAbstractTpl<Scalar> Base;
  typedef DataCollectorAbstractTpl<Scalar> DataCollectorAbstract;
  typedef ActivationDataAbstractTpl<Scalar> ActivationDataAbstract;
  typedef typename MathBaseTpl<Scalar>::VectorXs VectorXs;

  /**
   * @brief Initialize the numdiff residual data
   *
   * @tparam Model is the type of the `ResidualModelAbstractTpl`.
   * @param model is the object to compute the numerical differentiation from.
   */
  template <template <typename Scalar> class Model>
  explicit ResidualDataNumDiffTpl(Model<Scalar>* const model, DataCollectorAbstract* const shared_data)
      : Base(model, shared_data),
        dx(model->get_state()->get_ndx()),
        xp(model->get_state()->get_nx()),
        du(model->get_nu()),
        up(model->get_nu()) {
    dx.setZero();
    xp.setZero();
    du.setZero();
    up.setZero();

    const std::size_t& ndx = model->get_model()->get_state()->get_ndx();
    const std::size_t& nu = model->get_model()->get_nu();
    data_0 = model->get_model()->createData(shared_data);
    for (std::size_t i = 0; i < ndx; ++i) {
      data_x.push_back(model->get_model()->createData(shared_data));
    }
    for (std::size_t i = 0; i < nu; ++i) {
      data_u.push_back(model->get_model()->createData(shared_data));
    }
  }

  virtual ~ResidualDataNumDiffTpl() {}

  using Base::r;
  using Base::Ru;
  using Base::Rx;
  using Base::shared;

  VectorXs dx;  //!< State disturbance.
  VectorXs xp;  //!< The integrated state from the disturbance on one DoF "\f$ \int x dx_i \f$".
  VectorXs du;  //!< Control disturbance.
  VectorXs up;  //!< The integrated control from the disturbance on one DoF "\f$ \int u du_i = u + du \f$".
  boost::shared_ptr<Base> data_0;                //!< The data at the approximation point.
  std::vector<boost::shared_ptr<Base> > data_x;  //!< The temporary data associated with the state variation.
  std::vector<boost::shared_ptr<Base> > data_u;  //!< The temporary data associated with the control variation.
};

}  // namespace crocoddyl

/* --- Details -------------------------------------------------------------- */
/* --- Details -------------------------------------------------------------- */
/* --- Details -------------------------------------------------------------- */
#include "crocoddyl/core/numdiff/residual.hxx"

#endif  // CROCODDYL_CORE_NUMDIFF_RESIDUAL_HPP_


/*
 * Copyright (c) 2019 Opticks Team. All Rights Reserved.
 *
 * This file is part of Opticks
 * (see https://bitbucket.org/simoncblyth/opticks).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include <sstream>
#include <iomanip>
#include <cassert>

#include <boost/math/constants/constants.hpp>

#include "NGLMExt.hpp"
#include "NTris.hpp"

void NTris::add( const glm::vec3& a, const glm::vec3& b, const glm::vec3& c)
{
    unsigned i = verts.size(); verts.push_back(a);
    unsigned j = verts.size(); verts.push_back(b);
    unsigned k = verts.size(); verts.push_back(c);
    glm::uvec3 tri(i,j,k);
    tris.push_back(tri);
}

unsigned NTris::get_num_tri() const
{
    return tris.size();
}

unsigned NTris::get_num_vert() const
{
    return verts.size();
}

void NTris::get_vert( unsigned i, glm::vec3& v ) const
{
    v = verts[i] ;
}
void NTris::get_normal( unsigned /*i*/, glm::vec3& n ) const
{
    n.x = 0  ;
    n.y = 0  ;
    n.z = 0  ;
}
void NTris::get_uv( unsigned /*i*/, glm::vec3& uv ) const
{
    uv.x = 0  ;
    uv.y = 0  ;
    uv.z = 0  ;
}

void NTris::get_tri(unsigned i, glm::uvec3& t) const
{
    t = tris[i] ;
}
void NTris::get_tri(unsigned i, glm::uvec3& t, glm::vec3& a, glm::vec3& b, glm::vec3& c ) const
{
    t = tris[i] ;
    a = verts[t.x];
    b = verts[t.y];
    c = verts[t.z];
}

std::string NTris::brief() const
{
    std::stringstream ss ;

    ss << "NTris"
       << " nf " << std::setw(5) << tris.size()
       << " nv " << std::setw(5) << verts.size()
       ;

    return ss.str();
}

void NTris::dump(const char* msg) const
{
    std::cout << msg << " " << brief() << std::endl ;

    glm::uvec3 t ;
    glm::vec3 a ;
    glm::vec3 b ;
    glm::vec3 c ;

    unsigned ntri = get_num_tri();
    for(unsigned i=0 ; i < ntri ; i++)
    {
        get_tri(i, t, a,b,c );

        std::cout << " t " << std::setw(20) << glm::to_string(t)
                  << " a " << std::setw(20) << glm::to_string(a)
                  << " b " << std::setw(20) << glm::to_string(b)
                  << " c " << std::setw(20) << glm::to_string(c)
                  << std::endl
                  ;

    }

}

NTris* NTris::make_sphere( unsigned n_polar, unsigned n_azimuthal, float ctmin, float ctmax )
{

    /*

                                  t0             t1            ct0

            t = 0                  0             1/n_polar      +1

            t = n_polar - 1      1 - 1/n_polar     1            -1

                 -  ct0
                    max_exclude:  ct1 > zmax
                 -  ct1               - ct0
    zmax     +-----------------+        max_straddle
            /    -              \     - ct1
           /                     \
          /      -                \   - ct0
    zmin +-------------------------+    min_straddle
                 - ct0                - ct1
                   min_exclude:  ct0 < zmin
                 - ct1

    */

    assert(ctmax > ctmin && ctmax <= 1.f && ctmin >= -1.f);

    NTris* tris = new NTris ;

    float pi = boost::math::constants::pi<float>() ;

    for(unsigned t=0 ; t < n_polar ; t++)
    {
        double t0 = 1.0f*pi*float(t)/n_polar ;
        double t1 = 1.0f*pi*float(t+1)/n_polar ;

        double st0,st1,ct0,ct1 ;
        sincos_<double>(t0, st0, ct0 );
        sincos_<double>(t1, st1, ct1 );
        assert( ct0 > ct1 );

        bool max_exclude  = ct1 > ctmax ;
        bool max_straddle = ctmax <= ct0 && ctmax > ct1 ;
        bool min_straddle = ctmin <= ct0 && ctmin > ct1 ;
        bool min_exclude  = ct0 < ctmin ;

        if(max_exclude || min_exclude )
        {
            continue ;
        }
        else if(max_straddle)
        {
            sincos_<double>(acos(ctmax), st0, ct0 );
        }
        else if(min_straddle)
        {
            sincos_<double>(acos(ctmin), st1, ct1 );
        }

        for(unsigned p=0 ; p < n_azimuthal ; p++)
        {
            float p0 = 2.0f*pi*float(p)/n_azimuthal ;
            float p1 = 2.0f*pi*float(p+1)/n_azimuthal ;

            double sp0,sp1,cp0,cp1 ;
            sincos_<double>(p0, sp0, cp0 );
            sincos_<double>(p1, sp1, cp1 );

            glm::vec3 x00( st0*cp0, st0*sp0, ct0 );
            glm::vec3 x10( st0*cp1, st0*sp1, ct0 );

            glm::vec3 x01( st1*cp0, st1*sp0, ct1 );
            glm::vec3 x11( st1*cp1, st1*sp1, ct1 );

            if( t == 0 || t == n_polar - 1)
            {
                tris->add(x00,x01,x11);
            }
            else
            {
                tris->add(x00,x01,x10);
                tris->add(x10,x01,x11);
            }
        }
    }
    return tris ;
}


//! Implementation of the input_parser.

#include <iostream>
#include <sstream>
#include <fstream>
#include <string>
#include <vector>
#include <map>
#include <utility>
#include <stack>
#include <algorithm>
#include <Eigen/Dense>
#include "../include/ply.h"
#include "../include/input_parser.h"

//! The subscript of the laminate code. e.g. [0/45/45/90]8s4. 8 is the pre_count,
//! s is the has_symmetry, and 4 is the post_count.
struct SubscriptInfo {
    int pre_count;
    bool has_symmetry;
    int post_count;
};

//! Read material_data file and return the format into a map.
std::map<std::string, Properties>
    load_material_data(const std::string& filename);

//! Parse the laminate code and return a pair in which the first elemnt is the
//! vector containing ply angles, and the second element is the SubscriptInfo
//! of the given laminate.
std::pair<std::vector<double>, SubscriptInfo>
    laminate_code_parser(std::string laminate_code);

// Parse the trailing info of the laminate code and return the info into the
// SubscriptInfo struct.
SubscriptInfo subscript_parser(std::string subscript);

//! Convert the laminate information into a vector of ply objects.
std::vector<ply> build_laminate_vector(
    const std::pair<std::vector<double>, SubscriptInfo>& layout_info,
    std::vector<std::string>& ply_materials,
    const std::vector<double>& ply_thickness,
    const std::map<std::string, Properties>& material_map);

//! Strip the left and right square bracket of a string.
std::string strip_bracket(std::string input_str);

//! Convert the string inside bracket into vector of doubles.
std::vector<double> strs_to_vector(std::string code);

//! Convert the string inside bracket into vector of strings. Used for the
//! material data parsing.
std::vector<std::string> mat_str_to_vector(std::string strings);

using std::string;
using std::cin; using std::cout; using std::endl;
using std::vector; using std::map; using std::stack; using std::pair;

vector<string> read_composite_input(const string& filename) {
    vector<string> result;
    string line;
    std::ifstream file(filename);
    if (file.is_open()) {
        while (getline(file, line)) {
            if(line.find("[") != string::npos &&
                line.find("]") != string::npos) {
                string::size_type l_bracket_pos = line.find("[");
                result.push_back(line.substr(l_bracket_pos));
            }
        }
    }
    return result;
}

vector<ply> get_ply_vector(vector<string>& input_strings,
    const string& material_data_filename) {

        map<string, Properties> material_data =
            load_material_data(material_data_filename);

        pair<vector<double>, SubscriptInfo> layout_info =
            laminate_code_parser(input_strings[0]);

        string material_strings = strip_bracket(input_strings[1]);
        vector<string> ply_materials =
            mat_str_to_vector(material_strings);

        string thickness_strings = strip_bracket(input_strings[2]);

        vector<double> ply_thickness =
            strs_to_vector(thickness_strings);

        vector<ply> laminate = build_laminate_vector(layout_info,
            ply_materials, ply_thickness, material_data);

    return laminate;
    }

Eigen::Matrix<double, 6, 1> get_load_vector(string& input_string) {
    string load_string = strip_bracket(input_string);
    vector<double> load_stl_vector = strs_to_vector(load_string);
    Eigen::Matrix<double, 6, 1> load_vector(load_stl_vector.data());
    return load_vector;

}

double get_minimum_ply_thickness(string& thickness_strings_with_brackets) {
    string thickness_strings = strip_bracket(thickness_strings_with_brackets);
    vector<double> ply_thickness =
        strs_to_vector(thickness_strings);
    auto min_it = std::min_element(ply_thickness.begin(),
                                    ply_thickness.end());

    double minimum = *min_it;

    return minimum;
}

map<string, Properties> load_material_data(const string& filename) {
    map<string, Properties> data;
    string line;
    std::ifstream input_file(filename);
    if (input_file.is_open()) {
        string material_name;
        Properties p;
        input_file.ignore(500, '\n');  // Ignore the first line
        while (!input_file.eof()) {
            input_file >> material_name >> p.E1 >> p.E2 >> p.nu12 >> p.G12;
            if (!input_file) {
                cout << "Error: file corrupted. Lines After "
                    "corrupted line are not read."<< endl;
                break;
            }
            data.insert({material_name, p});
        }
    } else {
        cout << "Error: Cannot open file." ;
    }
    return data;
}

string strip_bracket(string input_str) {
    std::size_t left_bracket_pos = input_str.find("[");
    std::size_t right_bracket_pos = input_str.find("]");
    return input_str.substr(left_bracket_pos+1,
        right_bracket_pos-left_bracket_pos-1);

}

pair<vector<double>, SubscriptInfo> laminate_code_parser(string laminate_code) {
    std::size_t left_bracket_pos = laminate_code.find("[");
    std::size_t right_bracket_pos = laminate_code.find("]");
    pair<vector<double>, SubscriptInfo> result;
    if (left_bracket_pos == string::npos || right_bracket_pos == string::npos
        || left_bracket_pos > right_bracket_pos) {
            cout << "Invalid laminate code input." << endl;
        return result;
    }
    string theta_string = laminate_code.substr(left_bracket_pos+1,
        right_bracket_pos-left_bracket_pos-1);
    string subscript_string = laminate_code.substr(right_bracket_pos+1,
        laminate_code.size()-right_bracket_pos-1);
    result.first = strs_to_vector(theta_string);
    result.second = subscript_parser(subscript_string);
    return result;
}

vector<ply> build_laminate_vector(
    const pair<vector<double>, SubscriptInfo>& layout_info,
    vector<string>& ply_materials,
    const vector<double>& ply_thickness,
    const std::map<string, Properties>& material_map) {

        vector<double> theta_vec = layout_info.first;
        SubscriptInfo info = layout_info.second;

        vector<ply> laminate_vec;
        stack<ply> ply_stack;

        // It's not optimized in that we don't need the stack operation if
        // we don't have the symmetry subscript, but this arrangement
        // makes the code less redundant.
        for (int i = 1; i <= info.pre_count; i++) {
            auto theta_it = theta_vec.begin();
            auto mat_it = ply_materials.begin();
            auto t_it = ply_thickness.begin();
            for (; theta_it != theta_vec.end(); theta_it++, mat_it++, t_it++) {
                ply local_ply(*mat_it, material_map.at(*mat_it),
                    *theta_it, *t_it);
                laminate_vec.push_back(local_ply);
                ply_stack.push(local_ply); // for symmetric condition
            }
        }

        if (info.has_symmetry) {
            while (!ply_stack.empty()) {
                laminate_vec.push_back(ply_stack.top());
                ply_stack.pop();
            }
        }

        if (info.post_count > 1) {
            vector<ply> laminate_copy(laminate_vec);
            for (int i = 2; i <= info.post_count; i++) {
                for (auto it = laminate_copy.begin();
                    it != laminate_copy.end(); it++) {
                    laminate_vec.push_back(*it);
                }

            }
        }

    return laminate_vec;

}

std::vector<double> strs_to_vector(string theta_string) {
    std::replace(theta_string.begin(), theta_string.end(), '/', ' ');
    std::replace(theta_string.begin(), theta_string.end(), ',', ' ');
    std::istringstream theta_stream(theta_string);
    std::vector<double> theta_vec;
    double theta;
    while (theta_stream >> theta) {
        theta_vec.push_back(theta);
    }

    return theta_vec;
}

std::vector<string> mat_str_to_vector(string strings) {
    std::replace(strings.begin(), strings.end(), ',', ' ');
    std::istringstream stream(strings);
    std::vector<string> vec;
    string single_string;
    while (stream >> single_string) {
        vec.push_back(single_string);
    }
    return vec;
}

SubscriptInfo subscript_parser(string subscript) {
    SubscriptInfo info;
    std::size_t s_pos = subscript.find("s");
    if (s_pos != string::npos) {  // subscript contains an 's'
        info.has_symmetry = true;

        if (s_pos != 0) {  // subscript not at the beginning
            string pre_repetition_string = subscript.substr(0, s_pos);
            if (std::stoi(pre_repetition_string) > 0) {
                info.pre_count = std::stoi(pre_repetition_string);
            } else {
                info.pre_count = 1;
            }
        } else {
            info.pre_count = 1;
        }

        if (s_pos != subscript.size() - 1) {  // subscript not at the end
            string post_repetition_string =
                subscript.substr(s_pos+1, subscript.size());
            info.post_count = std::stoi(post_repetition_string);
        } else {
            info.post_count = 1;
        }

    } else {
        string pre_repetition_string = subscript.substr(0, subscript.size());
        info.pre_count = std::stoi(pre_repetition_string);
        info.has_symmetry = false;
        info.post_count = 0;
    }
    return info;
}


/**
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an "AS
 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 *  express or implied.  See the License for the specific language
 *  governing permissions and limitations under the License.
 *
 * For more about this software visit:
 *
 *      http://www.graphlab.ml.cmu.edu
 *
 */

/**
 * \file
 *
 * \brief The main file for the NMF matrix factorization algorithm.
 *
 */

#include <graphlab/util/stl_util.hpp>
#include <graphlab.hpp>
#include <graphlab/engine/gl3engine.hpp>
#include <Eigen/Dense>
#include "eigen_serialization.hpp"
#include <graphlab/macros_def.hpp>

#define VERTEX_DELTA_TASK_ID 0

typedef Eigen::VectorXd vec_type;
typedef Eigen::MatrixXd mat_type;

//when using negative node id range, we are not allowed to use
//0 and 1 so we add 2.
const static int SAFE_NEG_OFFSET=2;
static bool debug;
int iter = 0;
enum { PHASE1, PHASE2};
int phase = PHASE1;
double epsilon = 1e-16;

bool isuser(uint node){
  return ((int)node) >= 0;
}

/**
 * \ingroup toolkit_matrix_pvecization
 *
 * \brief the vertex data type which contains the latent pvec.
 *
 * Each row and each column in the matrix corresponds to a different
 * vertex in the SGD graph.  Associated with each vertex is a pvec
 * (vector) of latent parameters that represent that vertex.  The goal
 * of the SGD algorithm is to find the values for these latent
 * parameters such that the non-zero entries in the matrix can be
 * predicted by taking the dot product of the row and column pvecs.
 */
struct vertex_data {
  /**
   * \brief A shared "constant" that specifies the number of latent
   * values to use.
   */
  static size_t NLATENT;
  /** \brief The latent pvec for this vertex */
  vec_type pvec;

  int nupdates;

  /**
   * \brief Simple default constructor which randomizes the vertex
   *  data
   */
  vertex_data() { if (debug) pvec = vec_type::Ones(NLATENT); else randomize(); }
  /** \brief Randomizes the latent pvec */
  void randomize() { pvec.resize(NLATENT); pvec.setRandom(); }
  /** \brief Save the vertex data to a binary archive */
  void save(graphlab::oarchive& arc) const {
    arc << pvec;
  }
  /** \brief Load the vertex data from a binary archive */
  void load(graphlab::iarchive& arc) {
    arc >> pvec;
  }
}; // end of vertex data

std::size_t hash_value(vertex_data const& b) {
  return b.nupdates;
}

/**
 * \brief The edge data stores the entry in the matrix.
 *
 * In addition the edge data sgdo stores the most recent error estimate.
 */
struct edge_data : public graphlab::IS_POD_TYPE {
  /**
   * \brief The type of data on the edge;
   *
   * \li *Train:* the observed value is correct and used in training
   * \li *Validate:* the observed value is correct but not used in training
   * \li *Predict:* The observed value is not correct and should not be
   *        used in training.
   */
  enum data_role_type { TRAIN, VALIDATE, PREDICT  };

  /** \brief the observed value for the edge */
  float obs;

  /** \brief The train/validation/test designation of the edge */
  data_role_type role;

  /** \brief basic initialization */
  edge_data(float obs = 0, data_role_type role = PREDICT) :
    obs(obs), role(role) { }

}; // end of edge data

std::size_t hash_value(edge_data const& b) {
  return boost::hash_value(b.obs);
}

/**
 * \brief The graph type is defined in terms of the vertex and edge
 * data.
 */
typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;
typedef graphlab::gl3engine<graph_type> engine_type;

vec_type x1;
vec_type x2;
vec_type * px;

bool isuser_node(const graph_type::vertex_type& vertex){
  return isuser(vertex.id());
}

/**
 * \brief The graph loader function is a line parser used for
 * distributed graph construction.
 */
inline bool graph_loader(graph_type& graph,
                         const std::string& filename,
                         const std::string& line) {
  ASSERT_FALSE(line.empty());
  // Determine the role of the data
  edge_data::data_role_type role = edge_data::TRAIN;
  if(boost::ends_with(filename,".validate")) role = edge_data::VALIDATE;
  else if(boost::ends_with(filename, ".predict")) role = edge_data::PREDICT;
  // Parse the line
  std::stringstream strm(line);
  graph_type::vertex_id_type source_id(-1), target_id(-1);
  float obs(0);
  strm >> source_id >> target_id;

  // for test files (.predict) no need to read the actual rating value.
  if(role == edge_data::TRAIN || role == edge_data::VALIDATE) {
    strm >> obs;
  }
  target_id = -(graphlab::vertex_id_type(target_id + SAFE_NEG_OFFSET));

  // Create an edge and add it to the graph
  graph.add_edge(source_id, target_id, edge_data(obs, role));
  return true; // successful load
} // end of graph_loader

void vertex_delta(graph_type::vertex_type& vtx, const vec_type& delta) {
  if (delta.sum() != 0)
    vtx.data().pvec.array() *= delta.array() / px->array();
  for (uint i=0; i< vertex_data::NLATENT; i++)
    if (vtx.data().pvec[i] < epsilon)
      vtx.data().pvec[i] = epsilon;
}

void nmf_function(engine_type::context_type& context,
                  graph_type::edge_type& edge) {
  double pred = edge.source().data().pvec.dot(edge.target().data().pvec);
  if (pred == 0)
     logstream(LOG_FATAL)<<"Got into numerical error!" << std::endl;
  vec_type delta;
  delta = (phase == PHASE1 ? edge.target().data().pvec : edge.source().data().pvec) * edge.data().obs / pred;
  context.send_delta(VERTEX_DELTA_TASK_ID, phase == PHASE1 ? edge.source() : edge.target(), delta);
}

vec_type count_edges(const graph_type::edge_type& edge) {
  vec_type ret = vec_type::Zero(2);
  if (edge.data().role == edge_data::TRAIN){
    ret[0] = 1;
  }
  else if (edge.data().role == edge_data::VALIDATE){
    ret[1] = 1;
  }
  if (edge.data().obs < 0)
    logstream(LOG_FATAL)<<"Found a negative entry in matirx row " << edge.source().id() << " with value: " << edge.data().obs << std::endl;
  return ret;
}

void verify_rows(
    graph_type::vertex_type& vertex){
        if (isuser(vertex.id()) && vertex.num_out_edges() == 0)
          logstream(LOG_FATAL)<<"NMF algorithm can not work when the row " << vertex.id() << " of the matrix contains all zeros" << std::endl;
}

vec_type pre_iter( const graph_type::vertex_type & vertex){
  return vertex.data().pvec;
}

void sync_function(engine_type::context_type& context,
                   graph_type::vertex_type& vertex) {
  context.synchronize(vertex);
}

/**
 * \brief Given an edge compute the error associated with that edge
 */
double extract_l2_error(const graph_type::edge_type & edge) {
  double pred =
      edge.source().data().pvec.dot(edge.target().data().pvec);
  double rmse = (edge.data().obs - pred) * (edge.data().obs - pred);
  return rmse;
} // end of extract_l2_error

size_t vertex_data::NLATENT = 20;

int main(int argc, char** argv) {
  global_logger().set_log_level(LOG_INFO);
  global_logger().set_log_to_console(true);

  // Parse command line options -----------------------------------------------
  const std::string description =
      "Compute the ALS factorization of a matrix.";
  graphlab::command_line_options clopts(description);
  std::string input_dir;
  size_t interval = 0;
  size_t ITERATIONS = 10;
  std::string exec_type = "synchronous";
  clopts.attach_option("matrix", input_dir,
                       "The directory containing the matrix file");
  clopts.add_positional("matrix");
  clopts.attach_option("D", vertex_data::NLATENT,
                       "Number of latent parameters to use.");
  clopts.attach_option("interval", interval,
                       "The time in seconds between error reports");
  clopts.attach_option("iterations", ITERATIONS,
                       "number of NMF iterations");
  if(!clopts.parse(argc, argv) || input_dir == "") {
    std::cout << "Error in parsing command line arguments." << std::endl;
    clopts.print_description();
    return EXIT_FAILURE;
  }
  ///! Initialize control plain using mpi
  graphlab::mpi_tools::init(argc, argv);
  graphlab::distributed_control dc;

  dc.cout() << "Loading graph." << std::endl;
  graphlab::timer timer;
  graph_type graph(dc, clopts);
  graph.load(input_dir, graph_loader);
  dc.cout() << "Loading graph. Finished in "
            << timer.current_time() << std::endl;

  dc.cout() << "Finalizing graph." << std::endl;
  timer.start();
  graph.finalize();
  dc.cout() << "Finalizing graph. Finished in "
            << timer.current_time() << std::endl;

  if (!graph.num_edges() || !graph.num_vertices())
     logstream(LOG_FATAL)<< "Failed to load graph. Check your input path: " << input_dir << std::endl;

  dc.cout()
      << "========== Graph statistics on proc " << dc.procid()
      << " ==============="
      << "\n Num vertices: " << graph.num_vertices()
      << "\n Num edges: " << graph.num_edges()
      << "\n Num replica: " << graph.num_replicas()
      << "\n Replica to vertex ratio: "
      << float(graph.num_replicas())/graph.num_vertices()
      << "\n --------------------------------------------"
      << "\n Num local own vertices: " << graph.num_local_own_vertices()
      << "\n Num local vertices: " << graph.num_local_vertices()
      << "\n Replica to own ratio: "
      << (float)graph.num_local_vertices()/graph.num_local_own_vertices()
      << "\n Num local edges: " << graph.num_local_edges()
      //<< "\n Begin edge id: " << graph.global_eid(0)
        << "\n Edge balance ratio: "
        << float(graph.num_local_edges())/graph.num_edges()
        << std::endl;

  dc.cout() << "Creating engine" << std::endl;

  engine_type engine(dc, graph, clopts);
  vec_type edge_count = graph.map_reduce_edges<vec_type>(count_edges);
  dc.cout()<<"Training edges: " << edge_count[0] << " validation edges: " << edge_count[1] << std::endl;

  graphlab::vertex_set left = graph.select(isuser_node);
  graphlab::vertex_set right = ~left;
  graph.transform_vertices(verify_rows, left);

  engine.register_vertex_delta<vec_type>(VERTEX_DELTA_TASK_ID, vertex_delta);

  dc.cout() << "Running NMF" << std::endl;

  timer.start();
  for (size_t i = 0;i < ITERATIONS; ++i) {
    phase = PHASE1;
    x1 = graph.map_reduce_vertices<vec_type>(pre_iter,right);
    px = &x1;
    dc.cout() <<"x1 is: " << x1 << std::endl;

    engine.parfor_all_local_edges(nmf_function); //todo - only left
    engine.parfor_all_local_vertices(sync_function); //todo - only left
    engine.wait();

    phase = PHASE2;

    x2 = graph.map_reduce_vertices<vec_type>(pre_iter,left);
    px = &x2;
    dc.cout() <<"x2 is: " << x2 << std::endl;

    engine.parfor_all_local_edges(nmf_function); //todo only right
    engine.parfor_all_local_vertices(sync_function); //todo only right
    engine.wait();

    double rmse = graph.map_reduce_edges<double>(extract_l2_error);
    dc.cout() << "RMSE = " << sqrt(rmse / graph.num_edges()) << std::endl;

  }

  const double runtime = timer.current_time();
  dc.cout() << "----------------------------------------------------------"
            << std::endl
            << "Final Runtime (seconds):   " << runtime;

  // Compute the final training error -----------------------------------------
  dc.cout() << "Final error: " << std::endl;

  graphlab::mpi_tools::finalize();
  return EXIT_SUCCESS;
} // end of main


/**
 * @author Andre Anjos <andre.anjos@idiap.ch>
 * @date Sun  4 Mar 10:02:45 2012 CET
 *
 * @brief Implementation of the SVM training methods
 *
 * Copyright (C) 2011-2014 Idiap Research Institute, Martigny, Switzerland
 */

#include <bob.learn.libsvm/trainer.h>
#include <boost/format.hpp>
#include <boost/make_shared.hpp>
#include <bob.core/logging.h>

#ifdef BOB_DEBUG
//remove newline
#include <boost/algorithm/string/trim.hpp>
static std::string strip(const char* s) {
  std::string t(s);
  boost::algorithm::trim(t);
  return t;
}
#endif

static void debug_libsvm(const char* s) {
  TDEBUG1("[libsvm-" << LIBSVM_VERSION << "] " << strip(s));
}

bob::learn::libsvm::Trainer::Trainer(
    bob::learn::libsvm::machine_t machine_type,
    bob::learn::libsvm::kernel_t kernel_type,
    double cache_size,
    double eps,
    bool shrinking,
    bool probability
    )
{
  m_param.svm_type = machine_type;
  m_param.kernel_type = kernel_type;
  m_param.degree = 3;
  m_param.gamma = 0.;
  m_param.coef0 = 0.;
  m_param.cache_size = cache_size;
  m_param.eps = eps;
  m_param.C = 1;
  m_param.nu = 0.5;
  m_param.p = 0.1;
  m_param.shrinking = shrinking;
  m_param.probability = probability;

  //extracted from the data
  m_param.nr_weight = 0;
  m_param.weight_label = 0;
  m_param.weight = 0;
}

bob::learn::libsvm::Trainer::~Trainer() { }

/**
 * Erases an SVM problem:
 *
 * struct svm_problem {
 *   int l; //number of entries
 *   double* y; //labels
 *   svm_node** x; //each set terminated with a -1 index entry
 * };
 *
 * At svm-train the nodes for each entry are allocated globally, what is
 * probably more efficient from the allocation perspective. It still requires
 * libsvm to scan the data twice to understand how many nodes need to be
 * allocated globally.
 */
static void delete_problem(svm_problem* p) {
  delete[] p->y; //all labels
  delete[] p->x[0]; //all entries
  delete[] p->x; //entry pointers
  delete p;
}

/**
 * Allocates an svm_problem matrix
 */
static svm_problem* new_problem(size_t entries) {
  svm_problem* retval = new svm_problem;
  retval->l = (int)entries;
  retval->y = new double[entries];
  typedef svm_node* svm_node_ptr;
  retval->x = new svm_node_ptr[entries];
  for (size_t k=0; k<entries; ++k) retval->x[k] = 0;
  return retval;
}

/**
 * Converts the input arrayset data into an svm_problem matrix, used by libsvm
 * training routines. Updates "gamma" at the svm_parameter's.
 */
static boost::shared_ptr<svm_problem> data2problem
(const std::vector<blitz::Array<double, 2> >& data,
 const blitz::Array<double,1>& sub, const blitz::Array<double,1>& div,
 svm_parameter& param) {

  //counts the number of samples required
  size_t entries = 0;
  for (size_t k=0; k<data.size(); ++k)
    entries += data[k].extent(blitz::firstDim);

  //allocates the container that will represent the problem; at this stage, we
  //allocate entries for each vector, but not the space in which feature will
  //be put at. This will come next.
  boost::shared_ptr<svm_problem> problem(new_problem(entries),
      std::ptr_fun(delete_problem));

  //choose labels.
  if(param.svm_type==ONE_CLASS)
  {
    if ((data.size() != 1)) {
      boost::format m("Only support a singular entry for one class. Your are training ONE_CLASS svm classifier. You passed me a list of %d arraysets.");
      m % data.size();
      throw std::runtime_error(m.str());
    }
  }
  else {
    if ((data.size() <= 1) | (data.size() > 16)) {
      boost::format m("Only supports SVMs for binary or multi-class classification problems (up to 16 classes). You passed me a list of %d arraysets.");
      m % data.size();
      throw std::runtime_error(m.str());
    }
  }

  std::vector<double> labels;
  labels.reserve(data.size());
  if (data.size() == 1) {
    //oc-svm only support one class.
    labels.push_back(+1.);
  }
  else if (data.size() == 2) {
    //keep libsvm ordering
    labels.push_back(+1.);
    labels.push_back(-1.);
  }
  else { //data.size() == 3, 4, ..., 16
    for (size_t k=0; k<data.size(); ++k) labels.push_back(k+1);
  }

  //just count how many nodes we need; unfortunately we have no other choice
  //than doing a 2-pass instantiation here as libsvm has a very weird way to
  //optimize data access in which it requires all nodes to be allocated in a
  //single shot.
  size_t nodes = 0; //total number of nodes to be allocated
  blitz::Range all=blitz::Range::all();
  int n_features = data[0].extent(blitz::secondDim);
  blitz::Array<double,1> d(n_features); //for temporary feature manipulation

  for (size_t k=0; k<data.size(); ++k) {
    for (int i=0; i<data[k].extent(blitz::firstDim); ++i) {
      d = (data[k](i,all)-sub)/div; //eval and copy in 1 instruction
      for (int p=0; p<d.extent(blitz::firstDim); ++p) {
        if (d(p)) {
          ++nodes;
        }
      }
      ++nodes; //one extra for the termination node "index == -1"
    }
  }

  //allocates all the nodes, set first entry, a la libsvm
  svm_node* all_nodes = new svm_node[nodes];

  //iterates over each class data and fills the svm_node's
  int max_index = 0; //data width
  size_t sample = 0; //sample counter
  size_t node = 0; //node counter

  for (size_t k=0; k<data.size(); ++k) {
    for (int i=0; i<data[k].extent(blitz::firstDim); ++i) {
      problem->x[sample] = &all_nodes[node]; //setup current sample base pointer
      d = (data[k](i,all)-sub)/div; //eval and copy in 1 instruction
      for (blitz::sizeType p=0; p<d.size(); ++p) {
        if (d(p)) {
          int index = p+1; //starts indexing at 1
          all_nodes[node].index = index;
          all_nodes[node].value = d(p);
          if ( index > max_index ) max_index = index;
          ++node; //index within the current sample
        }
      }
      //marks end of sequence
      all_nodes[node].index = -1;
      all_nodes[node].value = 0;
      problem->y[sample] = labels[k];
      ++node;
      ++sample;
    }
  }

  //extracted from svm-train.c
  if (param.gamma == 0. && max_index > 0) {
    param.gamma = 1.0/max_index;
  }

  //do not support pre-computed kernels...
  if (param.kernel_type == PRECOMPUTED) {
    throw std::runtime_error("We currently dod not support PRECOMPUTED kernels in these bindings to libsvm");
  }

  return problem;
}

/**
 * A wrapper, to standardize the freeing of the svm_model
 */
static void svm_model_free(svm_model*& m) {
#if LIBSVM_VERSION >= 300
  svm_free_and_destroy_model(&m);
#else
  svm_destroy_model(m);
#endif
}

bob::learn::libsvm::Machine* bob::learn::libsvm::Trainer::train
(const std::vector<blitz::Array<double, 2> >& data,
 const blitz::Array<double,1>& input_subtraction,
 const blitz::Array<double,1>& input_division) const {

  //sanity check of input arraysets
  int n_features = data[0].extent(blitz::secondDim);

  for (size_t cl=0; cl<data.size(); ++cl) {
    if (data[cl].extent(blitz::secondDim) != n_features) {
      boost::format m("number of features (columns) of array for class %u (%d) does not match that of array for class 0 (%d)");
      m % cl % data[cl].extent(blitz::secondDim) % n_features;
      throw std::runtime_error(m.str());
    }
  }

  //converts the input arraysets into something libsvm can digest
  double save_gamma = m_param.gamma; ///< the next method may update it!
  boost::shared_ptr<svm_problem> problem =
    data2problem(data, input_subtraction, input_division,
        const_cast<svm_parameter&>(m_param) ///< temporary cast
        );

  //checks parametrization to make sure all is alright.
  const char* error_msg = svm_check_parameter(problem.get(), &m_param);

  if (error_msg) {
    const_cast<double&>(m_param.gamma) = save_gamma;
    boost::format m("libsvm-%d reports: %s");
    m % libsvm_version % error_msg;
    std::runtime_error(m.str());
  }

  //do the training, returns the new machine
#if LIBSVM_VERSION >= 291
  svm_set_print_string_function(debug_libsvm);
#else
  boost::format m("libsvm-%d does not support debugging stream setting");
  m % libsvm_version;
  debug_libsvm(m.str().c_str());
#endif
  boost::shared_ptr<svm_model> model(svm_train(problem.get(), &m_param),
      std::ptr_fun(svm_model_free));

  const_cast<double&>(m_param.gamma) = save_gamma;

  //save newly created machine to file, reload from there to get rid of memory
  //dependencies due to the poorly implemented memory model in libsvm
  boost::shared_ptr<svm_model> new_model =
    bob::learn::libsvm::svm_unpickle(bob::learn::libsvm::svm_pickle(model));

  auto retval = new bob::learn::libsvm::Machine(new_model);

  //sets up the scaling parameters given as input
  retval->setInputSubtraction(input_subtraction);
  retval->setInputDivision(input_division);

  return retval;
}

bob::learn::libsvm::Machine* bob::learn::libsvm::Trainer::train
(const std::vector<blitz::Array<double,2> >& data) const {
  int n_features = data[0].extent(blitz::secondDim);

  blitz::Array<double,1> sub(n_features);
  sub = 0.;
  blitz::Array<double,1> div(n_features);
  div = 1.;
  return train(data, sub, div);
}


/* Copyright (c) 2016 - 2020, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#include <ThermalOperator.hh>
#include <instantiation.hh>

#include <deal.II/base/index_set.h>
#include <deal.II/grid/filtered_iterator.h>
#include <deal.II/matrix_free/fe_evaluation.h>

namespace adamantine
{

template <int dim, int fe_degree, typename MemorySpaceType>
ThermalOperator<dim, fe_degree, MemorySpaceType>::ThermalOperator(
    MPI_Comm const &communicator,
    std::shared_ptr<MaterialProperty<dim>> material_properties)
    : _communicator(communicator), _material_properties(material_properties),
      _inverse_mass_matrix(
          new dealii::LA::distributed::Vector<double, MemorySpaceType>())
{
  _matrix_free_data.tasks_parallel_scheme =
      dealii::MatrixFree<dim, double>::AdditionalData::partition_color;
  _matrix_free_data.mapping_update_flags = dealii::update_gradients |
                                           dealii::update_JxW_values |
                                           dealii::update_quadrature_points;
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::reinit(
    dealii::DoFHandler<dim> const &dof_handler,
    dealii::AffineConstraints<double> const &affine_constraints,
    dealii::QGaussLobatto<1> const &quad)
{
  _matrix_free.reinit(dof_handler, affine_constraints, quad, _matrix_free_data);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::reinit(
    dealii::DoFHandler<dim> const &dof_handler,
    dealii::AffineConstraints<double> const &affine_constraints,
    dealii::QGauss<1> const &quad)
{
  _matrix_free.reinit(dof_handler, affine_constraints, quad, _matrix_free_data);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::
    compute_inverse_mass_matrix(
        dealii::DoFHandler<dim> const &dof_handler,
        dealii::AffineConstraints<double> const &affine_constraints)
{
  // Compute the inverse of the mass matrix
  dealii::QGaussLobatto<1> mass_matrix_quad(fe_degree + 1);
  dealii::MatrixFree<dim, double> mass_matrix_free;
  typename dealii::MatrixFree<dim, double>::AdditionalData mf_data;
  mf_data.tasks_parallel_scheme =
      dealii::MatrixFree<dim, double>::AdditionalData::partition_color;
  mf_data.mapping_update_flags = dealii::update_values |
                                 dealii::update_JxW_values |
                                 dealii::update_quadrature_points;

  mass_matrix_free.reinit(dof_handler, affine_constraints, mass_matrix_quad,
                          mf_data);
  mass_matrix_free.initialize_dof_vector(*_inverse_mass_matrix);
  dealii::VectorizedArray<double> one =
      dealii::make_vectorized_array(static_cast<double>(1.));
  dealii::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, double> fe_eval(
      mass_matrix_free);
  unsigned int const n_q_points = fe_eval.n_q_points;
  for (unsigned int cell = 0; cell < mass_matrix_free.n_macro_cells(); ++cell)
  {
    fe_eval.reinit(cell);
    for (unsigned int q = 0; q < n_q_points; ++q)
      fe_eval.submit_value(one, q);
    fe_eval.integrate(true, false);
    fe_eval.distribute_local_to_global(*_inverse_mass_matrix);
  }
  _inverse_mass_matrix->compress(dealii::VectorOperation::add);
  unsigned int const local_size = _inverse_mass_matrix->local_size();
  for (unsigned int k = 0; k < local_size; ++k)
  {
    if (_inverse_mass_matrix->local_element(k) > 1e-15)
      _inverse_mass_matrix->local_element(k) =
          1. / _inverse_mass_matrix->local_element(k);
    else
      _inverse_mass_matrix->local_element(k) = 0.;
  }
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::clear()
{
  _matrix_free.clear();
  _inverse_mass_matrix->reinit(0);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::vmult(
    dealii::LA::distributed::Vector<double, MemorySpaceType> &dst,
    dealii::LA::distributed::Vector<double, MemorySpaceType> const &src) const
{
  dst = 0.;
  vmult_add(dst, src);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::Tvmult(
    dealii::LA::distributed::Vector<double, MemorySpaceType> &dst,
    dealii::LA::distributed::Vector<double, MemorySpaceType> const &src) const
{
  dst = 0.;
  Tvmult_add(dst, src);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::vmult_add(
    dealii::LA::distributed::Vector<double, MemorySpaceType> &dst,
    dealii::LA::distributed::Vector<double, MemorySpaceType> const &src) const
{
  // Execute the matrix-free matrix-vector multiplication
  _matrix_free.cell_loop(&ThermalOperator::local_apply, this, dst, src);

  // Because cell_loop resolves the constraints, the constrained dofs are not
  // called they stay at zero. Thus, we need to force the value on the
  // constrained dofs by hand. The variable scaling is used so that we get the
  // right order of magnitude.
  // TODO: for now the value of scaling is set to 1
  double const scaling = 1.;
  std::vector<unsigned int> const &constrained_dofs =
      _matrix_free.get_constrained_dofs();
  for (auto &dof : constrained_dofs)
    dst.local_element(dof) += scaling * src.local_element(dof);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::Tvmult_add(
    dealii::LA::distributed::Vector<double, MemorySpaceType> &dst,
    dealii::LA::distributed::Vector<double, MemorySpaceType> const &src) const
{
  // The system of equation is symmetric so we can use vmult_add
  vmult_add(dst, src);
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::local_apply(
    dealii::MatrixFree<dim, double> const &data,
    dealii::LA::distributed::Vector<double, MemorySpaceType> &dst,
    dealii::LA::distributed::Vector<double, MemorySpaceType> const &src,
    std::pair<unsigned int, unsigned int> const &cell_range) const
{
  dealii::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, double> fe_eval(data);
  dealii::Tensor<1, dim> unit_tensor;
  for (unsigned int i = 0; i < dim; ++i)
    unit_tensor[i] = 1.;

  // Loop over the "cells". Note that we don't really work on a cell but on a
  // set of quadrature point.
  for (unsigned int cell = cell_range.first; cell < cell_range.second; ++cell)
  {
    // Reinit fe_eval on the current cell
    fe_eval.reinit(cell);
    // Store in a local vector the local values of src
    fe_eval.read_dof_values(src);
    // Evaluate only the function gradients on the reference cell
    fe_eval.evaluate(false, true);
    // Apply the Jacobian of the transformation, multiply by the variable
    // coefficients and the quadrature points
    for (unsigned int q = 0; q < fe_eval.n_q_points; ++q)
      fe_eval.submit_gradient(-_inv_rho_cp(cell, q) *
                                  _thermal_conductivity(cell, q) *
                                  fe_eval.get_gradient(q),
                              q);
    // Sum over the quadrature points.
    fe_eval.integrate(false, true);
    fe_eval.distribute_local_to_global(dst);
  }
}

template <int dim, int fe_degree, typename MemorySpaceType>
void ThermalOperator<dim, fe_degree, MemorySpaceType>::
    evaluate_material_properties(
        dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host> const
            &temperature)
{
  // Update the state of the materials
  _material_properties->update(_matrix_free.get_dof_handler(), temperature);

  unsigned int const n_cells = _matrix_free.n_macro_cells();
  dealii::FEEvaluation<dim, fe_degree, fe_degree + 1, 1, double> fe_eval(
      _matrix_free);
  _inv_rho_cp.reinit(n_cells, fe_eval.n_q_points);
  _thermal_conductivity.reinit(n_cells, fe_eval.n_q_points);
  for (unsigned int cell = 0; cell < n_cells; ++cell)
    for (unsigned int q = 0; q < fe_eval.n_q_points; ++q)
      for (unsigned int i = 0; i < _matrix_free.n_components_filled(cell); ++i)
      {
        typename dealii::DoFHandler<dim>::cell_iterator cell_it =
            _matrix_free.get_cell_iterator(cell, i);
        // Cast to Triangulation<dim>::cell_iterator to access the material_id
        typename dealii::Triangulation<dim>::active_cell_iterator cell_tria(
            cell_it);

        _thermal_conductivity(cell, q)[i] = _material_properties->get(
            cell_tria, StateProperty::thermal_conductivity);

        _inv_rho_cp(cell, q)[i] =
            1. / (_material_properties->get(cell_tria, StateProperty::density) *
                  _material_properties->get(cell_tria,
                                            StateProperty::specific_heat));
        _cell_it_to_mf_cell_map[cell_it] = std::make_pair(cell, i);
      }
}
} // namespace adamantine

INSTANTIATE_DIM_FEDEGREE_HOST(TUPLE(ThermalOperator))