/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2012 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef ALGORITHM_3_ERROR_ESTIMATOR_EE
#define ALGORITHM_3_ERROR_ESTIMATOR_EE

#include <algorithm>
#include <iostream>
#include <utility>
#include <vector>

#include <boost/timer/timer.hpp>
#include <boost/fusion/include/at_c.hpp>

#include "algorithm/newton_bridge.hh"
#include "algorithm/adaptationStrategy.hh"
#include "algorithm/errorDistribution.hh"
#include "linalg/jacobiPreconditioner.hh"
#include "fem/forEach.hh"
#include "fem/hierarchicErrorEstimator.hh"
#include "fem/variables.hh"
#include "utilities/enums.hh"
#include "errorEstimationTraits.hh"
#include "lagrangeLinearization.hh"

// forward declarations
namespace Dune
{
  struct InverseOperatorResult;
  template <class,int> class FieldVector;
}

namespace Kaskade
{
  struct Dummy{};

  template <template <class,class,class,bool> class Functional, class VariableSetDescription, class ExtensionVariableSetDescription, class ExtensionSpace,
  class NormFunctional, template <class> class RefinementStrategy = Adaptivity::ErrorEquilibration, bool lump=false, int components=1, class ReferenceSolution = Dummy, class ReferenceOperator=Dummy>
  class YetAnotherHBErrorEstimator : public ErrorEstimatorBase<VariableSetDescription,typename ErrorDistribution<NormFunctional,ExtensionVariableSetDescription>::AnsatzVars::VariableSet,ExtensionSpace,RefinementStrategy>
  {
    template <class AnsatzVars, class TestVars, class OriginVars> using EstimatorFunctional = Functional<AnsatzVars,TestVars,OriginVars,false>;
    template <class AnsatzVars, class TestVars, class OriginVars> using LumpedFunctional = Functional<AnsatzVars,TestVars,OriginVars,lump>;
    typedef ErrorEstimationTraits<EstimatorFunctional,VariableSetDescription,ExtensionVariableSetDescription> Traits;
    typedef ErrorEstimationTraits<LumpedFunctional,VariableSetDescription,ExtensionVariableSetDescription> LumpedTraits;
    typedef ErrorEstimatorBase<VariableSetDescription,typename ErrorDistribution<NormFunctional,ExtensionVariableSetDescription>::AnsatzVars::VariableSet,ExtensionSpace,RefinementStrategy> Base;
    using Base::squaredError;
    using Base::mde;
    using Base::errorDistribution;
    using Base::extensionSpace;
  public:
    // assemblers

    typedef typename Traits::Scalar Scalar;
    typedef typename VariableSetDescription::Grid Grid;
    static constexpr int dim = VariableSetDescription::Grid::dimension;

    typedef typename ExtensionVariableSetDescription::GridView::template Codim<0>::Iterator CellIterator;
    typedef typename ExtensionVariableSetDescription::GridView::template Codim<0>::Entity Cell;
    typedef ErrorDistribution<NormFunctional,ExtensionVariableSetDescription> EnergyError;
    typedef typename EnergyError::AnsatzVars::VariableSet ErrorRepresentation;

    YetAnotherHBErrorEstimator(NormFunctional& normFunctional_, VariableSetDescription& variableSetDescription_, ExtensionVariableSetDescription& extensionVariableSetDescription_,
        ExtensionSpace& extensionSpace_, Scalar fraction=0.7, bool verbose_=false)
    : Base(extensionSpace_,fraction,verbose), normFunctional(normFunctional_), variableSetDescription(variableSetDescription_), extensionVariableSetDescription(extensionVariableSetDescription_),
      squaredFraction(fraction*fraction), verbose(verbose_)
      {
        ass_EE.reset(new typename LumpedTraits::Assembler_EE(extensionVariableSetDescription.spaces));
        ass_HE.reset(new typename Traits::Assembler_HE(variableSetDescription.spaces));
        ass_EH.reset(new typename Traits::Assembler_EH(extensionVariableSetDescription.spaces));
        ass_HH.reset(new typename Traits::Assembler_HH(variableSetDescription.spaces));
      }

    virtual ~YetAnotherHBErrorEstimator(){}

    void operator()(AbstractLinearization const& lin, AbstractFunctionSpaceElement const& x_, AbstractFunctionSpaceElement const& dx_, int step, AbstractFunctionSpaceElement const&)
    {
      boost::timer::cpu_timer overallTimer;
      using boost::fusion::at_c;
      if(verbose) std::cout << "ERROR ESTIMATOR: Start." << std::endl;
      Bridge::Vector<typename Traits::VarSet> const& xl = dynamic_cast<const Bridge::Vector<typename Traits::VarSet>&>(x_);
      computeError(xl.get());
      Bridge::Vector<typename Traits::ExtensionVarSet> xe(extensionVariableSetDescription);
      boost::timer::cpu_timer atimer;
      assembleAt(xl,xe);
      if(verbose) std::cout << "ERROR ESTIMATOR: assembly time: " << boost::timer::format(atimer.elapsed())  << std::endl;
      Bridge::Vector<typename Traits::VarSet> const& dxl = dynamic_cast<const Bridge::Vector<typename Traits::VarSet>&>(dx_);

      /***************************************************************************************/
      // state equation -> DLY
      typename Traits::A_HH a_HH(*ass_HH,false);
      typename Traits::A_HE a_HE(*ass_HE,false);
      typename LumpedTraits::A_EE a_EE(*ass_EE,false);
      typename Traits::B_HH b_HH(*ass_HH,false);
      typename Traits::B_HE b_HE(*ass_HE,false);
      typename Traits::B_EH b_EH(*ass_EH,false);
      typename LumpedTraits::B_EE b_EE(*ass_EE,false);

      typename Traits::ExtensionVectorP rhsPE(ass_EE->template rhs<Traits::adjointId,Traits::adjointId+1>());
      rhsPE *= -1.0;
      typename Traits::VectorY dy(at_c<Traits::stateId>(dxl.get().data).coefficients()), solYH = dy; solYH *= 0;
      typename Traits::VectorU du(at_c<Traits::controlId>(dxl.get().data).coefficients());
      typename Traits::ExtensionVectorY solYE(Traits::ExtensionVectorY_Initializer::init(extensionVariableSetDescription));

      b_HE.applyscaleadd(-1.0,du,rhsPE);
      a_HE.applyscaleadd(-1.0,dy,rhsPE);

      JacobiPreconditionerForTriplets<Scalar,typename LumpedTraits::ExtensionVectorY,typename LumpedTraits::ExtensionVectorP> jacobi(a_EE);
      jacobi.apply(solYE,rhsPE);

      // error in adjoint equation
      typename Traits::Lyy_HH lyy_HH(*ass_HH,false);
      typename Traits::Lyy_HE lyy_HE(*ass_HE,false);
      typename Traits::Lyy_EH lyy_EH(*ass_EH,false);
      typename LumpedTraits::Lyy_EE lyy_EE(*ass_EE,false);
      TransposedOperator<typename Traits::A_HE> at_EH(a_HE);//(*ass_EH,false);

      typename Traits::VectorY rhsYH(ass_HH->template rhs<Traits::stateId,Traits::stateId+1>());
      rhsYH *= -1.0;
      typename Traits::ExtensionVectorY rhsYE(ass_EE->template rhs<Traits::stateId, Traits::stateId+1>());
      rhsYE *= -1.0;

      lyy_HH.applyscaleadd(-1.0,dy,rhsYH);
      lyy_HE.applyscaleadd(-1.0,dy,rhsYE);
      lyy_EH.applyscaleadd(-1.0,solYE,rhsYH);
      lyy_EE.applyscaleadd(-1.0,solYE,rhsYE);

      typename Traits::ExtensionVectorP solPE(Traits::ExtensionVectorP_Initializer::init(extensionVariableSetDescription));
      typename Traits::VectorP solPH(Traits::VectorP_Initializer::init(variableSetDescription));

      jacobi.apply(solPE,rhsYE);
      at_EH.applyscaleadd(-1.0,solPE,rhsYH);

      boost::timer::cpu_timer mgTimer;
      MultiGridSolver<Grid,components>( a_HH, boost::fusion::at_c<0>(ass_HH->spaces())->gridManager().grid(),
                                        typename MultiGridSolver<Grid,components>::Parameter(mgSteps,mgSmoothingSteps,relativeAccuracy), true ).apply(at_c<0>(solPH.data),
                                                                                at_c<0>(rhsYH.data));
      std::cout << "mg: " << boost::timer::format(mgTimer.elapsed()) << std::endl;

      // variational equality
      TransposedOperator<typename Traits::B_HH> bt_HH(b_HH);
      TransposedOperator<typename Traits::B_EH> bt_HE(b_EH);
      TransposedOperator<typename Traits::B_HE> bt_EH(b_HE);
      TransposedOperator<typename LumpedTraits::B_EE> bt_EE(b_EE);
      typename Traits::Luu_HH luu_HH(*ass_HH,false);
      typename Traits::Luu_EH luu_EH(*ass_EH,false);
      typename Traits::Luu_HE luu_HE(*ass_HE,false);
      typename LumpedTraits::Luu_EE luu_EE(*ass_EE,false);

      typename Traits::VectorU rhsUH(ass_HH->template rhs<Traits::controlId,Traits::controlId+1>());
      rhsUH *= -1;
      typename Traits::ExtensionVectorU rhsUE(ass_EE->template rhs<Traits::controlId,Traits::controlId+1>());
      rhsUE *= -1;

      luu_HH.applyscaleadd(-1.0,du,rhsUH);
      luu_HE.applyscaleadd(-1.0,du,rhsUE);

      bt_HH.applyscaleadd(-1.0,solPH,rhsUH);
      bt_HE.applyscaleadd(-1.0,solPH,rhsUE);
      bt_EH.applyscaleadd(-1.0,solPE,rhsUH);
      bt_EE.applyscaleadd(-1.0,solPE,rhsUE);

      typename Traits::VectorU solUH(Traits::VectorU_Initializer::init(variableSetDescription));
      typename Traits::ExtensionVectorU solUE(Traits::ExtensionVectorU_Initializer::init(extensionVariableSetDescription));

      JacobiPreconditionerForTriplets<Scalar,typename LumpedTraits::ExtensionVectorU,typename LumpedTraits::ExtensionVectorU>(luu_EE).apply(solUE,rhsUE);
      luu_EH.applyscaleadd(-1.0,solUE,rhsUH);
      boost::timer::cpu_timer chebTimer;
      ChebyshevPreconditioner<typename Traits::Luu_HH> cheb(luu_HH,chebySteps);
      cheb.initForMassMatrix_TetrahedralQ1Elements();
      cheb.apply(solUH,rhsUH);
      std::cout << "cheb: " << boost::timer::format(chebTimer.elapsed()) << std::endl;

      typename ExtensionVariableSetDescription::VariableSet errorEstimate_E(extensionVariableSetDescription);
      typename VariableSetDescription::VariableSet errorEstimate_H(variableSetDescription);
      solYH *= 0.0; // DLY in state equation -> no error transport into space of linear finite elements
      at_c<Traits::stateId>(errorEstimate_H.data) = at_c<0>(solYH.data);
      at_c<Traits::controlId>(errorEstimate_H.data) = at_c<0>(solUH.data);
      at_c<Traits::stateId>(errorEstimate_E.data) = at_c<0>(solYE.data);
      at_c<Traits::controlId>(errorEstimate_E.data) = at_c<0>(solUE.data);
//      std::string savefilename = createFileName("estimate_E",".vtu",false);
//      std::cout << "Error estimator: writing estimates" << std::endl;
//      writeVTKFile(extensionVariableSetDescription.gridView,errorEstimate_E,savefilename,IoOptions(),2);
//      std::cout << "Error estimator: writing lower order estimates" << std::endl;
//      writeVTKFile(variableSetDescription.gridView,errorEstimate_H,savefilename,IoOptions(),1);
      /***************************************************************************************/
      // Transfer error indicators to cells.
      auto const& is = extensionSpace.gridManager().grid().leafIndexSet();
      errorDistribution.clear();
      errorDistribution.resize(is.size(0),std::make_pair(0.0,0));

      EnergyError energyError(normFunctional,xl.get(),errorEstimate_H,errorEstimate_E);
      energyError.considerStateVariable(true);
      energyError.considerControlVariable(true);
      energyError.considerAdjointVariable(false);
      typedef VariationalFunctionalAssembler<LinearizationAt<EnergyError> > EnergyErrorAssembler;
      EnergyErrorAssembler eeAssembler(extensionSpace.gridManager(), energyError.getSpaces());

      eeAssembler.assemble(linearization(energyError,xl.get()), EnergyErrorAssembler::RHS);
      typename EnergyError::ErrorVector distError( eeAssembler.rhs() );
      mde.reset(new ErrorRepresentation(energyError.getVariableSetDescription()) );
      *mde = distError;

//      if(verbose)
//      {
//        std::string name = "errorDistribution_";
//        name += std::to_string(step);
//        writeVTKFile(mde->descriptions.gridView,*mde,name);
//      }

      CellIterator cend = extensionVariableSetDescription.gridView.template end<0>();
      for (CellIterator ci=extensionVariableSetDescription.gridView.template begin<0>(); ci!=cend; ++ci)
        errorDistribution[is.index(*ci)] = std::make_pair( fabs(boost::fusion::at_c<0>(mde->data).value(*ci,Dune::FieldVector<Scalar,dim>(0.3))) , is.index(*ci));

      squaredError = std::accumulate(errorDistribution.begin(), errorDistribution.end(), 0.0, ErrorEstimator_Detail::add);

      if(verbose) std::cout << "overall error estimation time: " << boost::timer::format(overallTimer.elapsed()) << std::endl;
    }

    template <typename... Args>
    void initFunctionals(const Args&... args)
    {
      F_HH.reset(new typename Traits::Functional_HH(args...));
      F_HE.reset(new typename Traits::Functional_HE(args...));
    }

    template <typename... Args>
    void initExtensionFunctionals(const Args&... args)
    {
      F_EH.reset(new typename Traits::Functional_EH(args...));
      F_EE.reset(new typename LumpedTraits::Functional_EE(args...));
    }

    void setReference(ReferenceOperator const& Aref, ReferenceSolution const& ref, typename Traits::Vector& refcv)
    {
      referenceSolution = &ref;
      referenceOperator = &Aref;
      referenceCoeffVec = &refcv;
    }

    void computeError(typename Traits::VarSet const& x)
    {
      if(referenceSolution == nullptr) return;
      ReferenceSolution tmp(*referenceSolution);
      interpolateGloballyFromFunctor<PlainAverage>(boost::fusion::at_c<0>(tmp.data),[&x](Cell const& cell, Dune::FieldVector<Scalar,dim> const& xLocal)
      {
        return boost::fusion::at_c<0>(x.data).value(cell.geometry().global(xLocal));
      });

      interpolateGloballyFromFunctor<PlainAverage>(boost::fusion::at_c<1>(tmp.data),[&x](Cell const& cell, Dune::FieldVector<Scalar,dim> const& xLocal)
      {
        return boost::fusion::at_c<1>(x.data).value(cell.geometry().global(xLocal));
      });
//      interpolateGloballyWeak<PlainAverage>(boost::fusion::at_c<2>(tmp.data),boost::fusion::at_c<2>(x.data));
      tmp -= *referenceSolution;

      boost::fusion::at_c<0>(referenceCoeffVec->data) = boost::fusion::at_c<0>(tmp.data).coefficients();
      boost::fusion::at_c<1>(referenceCoeffVec->data) = boost::fusion::at_c<1>(tmp.data).coefficients();
      // boost::fusion::at_c<2>(referenceCoeffVec->data) = boost::fusion::at_c<2>(tmp.data).coefficients();

      auto tmp2 = *referenceCoeffVec;
      referenceOperator->apply(*referenceCoeffVec,tmp2);
      std::cout << "REAL ERROR: " << (*referenceCoeffVec)*tmp2 << std::endl;
    }

  private:
    void assembleAt(const Bridge::Vector<typename Traits::VarSet>& xl, const Bridge::Vector<typename Traits::ExtensionVarSet>& xe)
    {
      ass_HH->assemble(linearization(*F_HH,xl.get()));
      ass_HE->assemble(linearization(*F_HE,xl.get()));
      ass_EH->assemble(linearization(*F_EH,xe.get()));
      ass_EE->assemble(linearization(*F_EE,xe.get()));
    }

    std::unique_ptr<typename Traits::Functional_HH> F_HH;
    std::unique_ptr<typename Traits::Functional_HE> F_HE;
    std::unique_ptr<typename Traits::Functional_EH> F_EH;
    std::unique_ptr<typename LumpedTraits::Functional_EE> F_EE;
    std::unique_ptr<typename Traits::Assembler_HH> ass_HH;
    //typename Traits::Assembler_HH* ass_HH;
    std::unique_ptr<typename Traits::Assembler_HE> ass_HE;
    std::unique_ptr<typename Traits::Assembler_EH> ass_EH;
    std::unique_ptr<typename LumpedTraits::Assembler_EE> ass_EE;
    NormFunctional& normFunctional;
    VariableSetDescription& variableSetDescription;
    ExtensionVariableSetDescription& extensionVariableSetDescription;
    Scalar squaredFraction;
    bool verbose;
    size_t chebySteps = 30, mgSteps = 25, mgSmoothingSteps = 20;
    Scalar relativeAccuracy = 1e-3;
    ReferenceSolution const* referenceSolution = nullptr;
    ReferenceOperator const* referenceOperator = nullptr;
    typename Traits::Vector* referenceCoeffVec = nullptr;
    //std::function<Scalar(typename Traits::VarSet)> compareError;
  };

  template <template <class,class,class,bool> class Functional, class VariableSetDescription, class ExtensionVariableSetDescription, class ExtensionSpace,
  class NormFunctional, template <class> class RefinementStrategy = Adaptivity::ErrorEquilibration, bool lump=false, int components=1>
  class YetAnotherHBErrorEstimator_Elasticity : public ErrorEstimatorBase<VariableSetDescription,typename ErrorDistribution<NormFunctional,ExtensionVariableSetDescription>::AnsatzVars::VariableSet,ExtensionSpace,RefinementStrategy>
  {
    template <class AnsatzVars, class TestVars, class OriginVars> using EstimatorFunctional = Functional<AnsatzVars,TestVars,OriginVars,false>;
    template <class AnsatzVars, class TestVars, class OriginVars> using LumpedFunctional = Functional<AnsatzVars,TestVars,OriginVars,lump>;
    typedef ErrorEstimationTraits<EstimatorFunctional,VariableSetDescription,ExtensionVariableSetDescription> Traits;
    typedef ErrorEstimationTraits<LumpedFunctional,VariableSetDescription,ExtensionVariableSetDescription> LumpedTraits;
    typedef ErrorEstimatorBase<VariableSetDescription,typename ErrorDistribution<NormFunctional,ExtensionVariableSetDescription>::AnsatzVars::VariableSet,ExtensionSpace,RefinementStrategy> Base;
    using Base::squaredError;
    using Base::mde;
    using Base::errorDistribution;
    using Base::extensionSpace;
  public:
    // assemblers

    typedef typename Traits::Scalar Scalar;
    typedef typename VariableSetDescription::Grid Grid;
    static constexpr int dim = VariableSetDescription::Grid::dimension;

    typedef typename ExtensionVariableSetDescription::GridView::template Codim<0>::Iterator CellIterator;
    typedef typename ExtensionVariableSetDescription::GridView::template Codim<0>::Entity Cell;
    typedef ErrorDistribution<NormFunctional,ExtensionVariableSetDescription> EnergyError;
    typedef typename EnergyError::AnsatzVars::VariableSet ErrorRepresentation;

    YetAnotherHBErrorEstimator_Elasticity(NormFunctional& normFunctional_, VariableSetDescription& variableSetDescription_, ExtensionVariableSetDescription& extensionVariableSetDescription_,
        ExtensionSpace& extensionSpace_, Scalar fraction=0.7, bool verbose_=false)
    : Base(extensionSpace_,fraction,verbose), normFunctional(normFunctional_), variableSetDescription(variableSetDescription_), extensionVariableSetDescription(extensionVariableSetDescription_),
      squaredFraction(fraction*fraction), verbose(verbose_)
      {
        ass_EE.reset(new typename LumpedTraits::Assembler_EE(extensionVariableSetDescription.spaces));
        ass_HE.reset(new typename Traits::Assembler_HE(variableSetDescription.spaces));
        ass_EH.reset(new typename Traits::Assembler_EH(extensionVariableSetDescription.spaces));
        ass_HH.reset(new typename Traits::Assembler_HH(variableSetDescription.spaces));
      }

    virtual ~YetAnotherHBErrorEstimator_Elasticity(){}

    void operator()(AbstractLinearization const& lin, AbstractFunctionSpaceElement const& x_, AbstractFunctionSpaceElement const& dx_, int step, AbstractFunctionSpaceElement const&)
    {
      boost::timer::cpu_timer overallTimer;
      using boost::fusion::at_c;
      if(verbose) std::cout << "ERROR ESTIMATOR: Start." << std::endl;
      Bridge::Vector<typename Traits::VarSet> const& xl = dynamic_cast<const Bridge::Vector<typename Traits::VarSet>&>(x_);
      Bridge::Vector<typename Traits::ExtensionVarSet> xe(extensionVariableSetDescription);
      boost::timer::cpu_timer atimer;
      assembleAt(xl,xe);
      if(verbose) std::cout << "ERROR ESTIMATOR: assembly time: " << boost::timer::format(atimer.elapsed())  << std::endl;
      Bridge::Vector<typename Traits::VarSet> const& dxl = dynamic_cast<const Bridge::Vector<typename Traits::VarSet>&>(dx_);

      /***************************************************************************************/
      // state equation -> DLY
      typename Traits::A_HH a_HH(*ass_HH,false);
      typename Traits::A_HE a_HE(*ass_HE,false);
      typename LumpedTraits::A_EE a_EE(*ass_EE,false);
      typename Traits::B_HH b_HH(*ass_HH,false);
      typename Traits::B_HE b_HE(*ass_HE,false);
      typename Traits::B_EH b_EH(*ass_EH,false);
      typename LumpedTraits::B_EE b_EE(*ass_EE,false);

      typename Traits::ExtensionVectorP rhsPE(ass_EE->template rhs<Traits::adjointId,Traits::adjointId+1>());
      rhsPE *= -1.0;
      typename Traits::VectorY dy(at_c<Traits::stateId>(dxl.get().data).coefficients()), solYH = dy; solYH *= 0;
      typename Traits::VectorU du(at_c<Traits::controlId>(dxl.get().data).coefficients());
      typename Traits::ExtensionVectorY solYE(Traits::ExtensionVectorY_Initializer::init(extensionVariableSetDescription));

      b_HE.applyscaleadd(-1.0,du,rhsPE);
      a_HE.applyscaleadd(-1.0,dy,rhsPE);

      JacobiPreconditionerForTriplets<Scalar,typename LumpedTraits::ExtensionVectorY,typename LumpedTraits::ExtensionVectorP> jacobi(a_EE);
      jacobi.apply(solYE,rhsPE);

      // error in adjoint equation
      typename Traits::Lyy_HH lyy_HH(*ass_HH,false);
      typename Traits::Lyy_HE lyy_HE(*ass_HE,false);
      typename Traits::Lyy_EH lyy_EH(*ass_EH,false);
      typename LumpedTraits::Lyy_EE lyy_EE(*ass_EE,false);
      typename Traits::Luy_HH luy_HH(*ass_HH,false);
      typename Traits::Luy_HE luy_HE(*ass_HE,false);
      typename Traits::Luy_EH luy_EH(*ass_EH,false);
      typename LumpedTraits::Luy_EE luy_EE(*ass_EE,false);
      TransposedOperator<typename Traits::Luy_HH> lyu_HH(luy_HH);
      TransposedOperator<typename Traits::Luy_EH> lyu_HE(luy_EH);
//      typename Traits::Lyu_HH lyu_HH(*ass_HH,false);
//      typename Traits::Lyu_HE lyu_HE(*ass_HE,false);
      TransposedOperator<typename Traits::A_HE> at_EH(a_HE);//(*ass_EH,false);

      typename Traits::VectorY rhsYH(ass_HH->template rhs<Traits::stateId,Traits::stateId+1>());
      rhsYH *= -1.0;
      typename Traits::ExtensionVectorY rhsYE(ass_EE->template rhs<Traits::stateId, Traits::stateId+1>());
      rhsYE *= -1.0;

      lyy_HH.applyscaleadd(-1.0,dy,rhsYH);
      lyy_HE.applyscaleadd(-1.0,dy,rhsYE);
      lyy_EH.applyscaleadd(-1.0,solYE,rhsYH);
      lyy_EE.applyscaleadd(-1.0,solYE,rhsYE);
      lyu_HH.applyscaleadd(-1.0,du,rhsYH);
      lyu_HE.applyscaleadd(-1.0,du,rhsYE);

      typename Traits::ExtensionVectorP solPE(Traits::ExtensionVectorP_Initializer::init(extensionVariableSetDescription));
      typename Traits::VectorP solPH(Traits::VectorP_Initializer::init(variableSetDescription));

      jacobi.apply(solPE,rhsYE);
      at_EH.applyscaleadd(-1.0,solPE,rhsYH);

      boost::timer::cpu_timer mgTimer;
      MultiGridSolver<Grid,components>( a_HH, boost::fusion::at_c<0>(ass_HH->spaces())->gridManager().grid(),
                                        typename MultiGridSolver<Grid,components>::Parameter(mgSteps,mgSmoothingSteps,relativeAccuracy), true ).apply(at_c<0>(solPH.data),
                                                                                at_c<0>(rhsYH.data));
      std::cout << "mg: " << boost::timer::format(mgTimer.elapsed()) << std::endl;

      // variational equality
      TransposedOperator<typename Traits::B_HH> bt_HH(b_HH);
      TransposedOperator<typename Traits::B_EH> bt_HE(b_EH);
      TransposedOperator<typename Traits::B_HE> bt_EH(b_HE);
      TransposedOperator<typename LumpedTraits::B_EE> bt_EE(b_EE);
      typename Traits::Luu_HH luu_HH(*ass_HH,false);
      typename Traits::Luu_EH luu_EH(*ass_EH,false);
      typename Traits::Luu_HE luu_HE(*ass_HE,false);
      typename LumpedTraits::Luu_EE luu_EE(*ass_EE,false);

      typename Traits::VectorU rhsUH(ass_HH->template rhs<Traits::controlId,Traits::controlId+1>());
      rhsUH *= -1;
      typename Traits::ExtensionVectorU rhsUE(ass_EE->template rhs<Traits::controlId,Traits::controlId+1>());
      rhsUE *= -1;

      luu_HH.applyscaleadd(-1.0,du,rhsUH);
      luu_HE.applyscaleadd(-1.0,du,rhsUE);

      bt_HH.applyscaleadd(-1.0,solPH,rhsUH);
      bt_HE.applyscaleadd(-1.0,solPH,rhsUE);
      bt_EH.applyscaleadd(-1.0,solPE,rhsUH);
      bt_EE.applyscaleadd(-1.0,solPE,rhsUE);

      luy_HH.applyscaleadd(-1.0,dy,rhsUH);
      luy_HE.applyscaleadd(-1.0,dy,rhsUE);
      luy_EH.applyscaleadd(-1.0,solYE,rhsUH);
      luy_EE.applyscaleadd(-1.0,solYE,rhsUE);

      typename Traits::VectorU solUH(Traits::VectorU_Initializer::init(variableSetDescription));
      typename Traits::ExtensionVectorU solUE(Traits::ExtensionVectorU_Initializer::init(extensionVariableSetDescription));

      JacobiPreconditionerForTriplets<Scalar,typename LumpedTraits::ExtensionVectorU,typename LumpedTraits::ExtensionVectorU>(luu_EE).apply(solUE,rhsUE);
      luu_EH.applyscaleadd(-1.0,solUE,rhsUH);
      boost::timer::cpu_timer chebTimer;
      ChebyshevPreconditioner<typename Traits::Luu_HH> cheb(luu_HH,chebySteps);
      cheb.initForMassMatrix_TetrahedralQ1Elements();
      cheb.apply(solUH,rhsUH);
      std::cout << "cheb: " << boost::timer::format(chebTimer.elapsed()) << std::endl;

      typename ExtensionVariableSetDescription::VariableSet errorEstimate_E(extensionVariableSetDescription);
      typename VariableSetDescription::VariableSet errorEstimate_H(variableSetDescription);
      solYH *= 0.0; // DLY in state equation -> no error transport into space of linear finite elements
      at_c<Traits::stateId>(errorEstimate_H.data) = at_c<0>(solYH.data);
      at_c<Traits::controlId>(errorEstimate_H.data) = at_c<0>(solUH.data);
      at_c<Traits::stateId>(errorEstimate_E.data) = at_c<0>(solYE.data);
      at_c<Traits::controlId>(errorEstimate_E.data) = at_c<0>(solUE.data);
//      std::string savefilename = createFileName("estimate_E",".vtu",false);
//      std::cout << "Error estimator: writing estimates" << std::endl;
//      writeVTKFile(extensionVariableSetDescription.gridView,errorEstimate_E,savefilename,IoOptions(),2);
//      std::cout << "Error estimator: writing lower order estimates" << std::endl;
//      writeVTKFile(variableSetDescription.gridView,errorEstimate_H,savefilename,IoOptions(),1);
      /***************************************************************************************/
      // Transfer error indicators to cells.
      auto const& is = extensionSpace.gridManager().grid().leafIndexSet();
      errorDistribution.clear();
      errorDistribution.resize(is.size(0),std::make_pair(0.0,0));

      EnergyError energyError(normFunctional,xl.get(),errorEstimate_H,errorEstimate_E);
      energyError.considerStateVariable(true);
      energyError.considerControlVariable(true);
      energyError.considerAdjointVariable(false);
      typedef VariationalFunctionalAssembler<LinearizationAt<EnergyError> > EnergyErrorAssembler;
      EnergyErrorAssembler eeAssembler(extensionSpace.gridManager(), energyError.getSpaces());

      eeAssembler.assemble(linearization(energyError,xl.get()), EnergyErrorAssembler::RHS);
      typename EnergyError::ErrorVector distError( eeAssembler.rhs() );
      mde.reset(new ErrorRepresentation(energyError.getVariableSetDescription()) );
      *mde = distError;

//      if(verbose)
//      {
//        std::string name = "errorDistribution_";
//        name += std::to_string(step);
//        writeVTKFile(mde->descriptions.gridView,*mde,name);
//      }

      CellIterator cend = extensionVariableSetDescription.gridView.template end<0>();
      for (CellIterator ci=extensionVariableSetDescription.gridView.template begin<0>(); ci!=cend; ++ci)
        errorDistribution[is.index(*ci)] = std::make_pair( fabs(boost::fusion::at_c<0>(mde->data).value(*ci,Dune::FieldVector<Scalar,dim>(0.3))) , is.index(*ci));

      squaredError = std::accumulate(errorDistribution.begin(), errorDistribution.end(), 0.0, ErrorEstimator_Detail::add);

      if(verbose) std::cout << "overall error estimation time: " << boost::timer::format(overallTimer.elapsed()) << std::endl;
    }

    template <typename... Args>
    void initFunctionals(const Args&... args)
    {
      F_HH.reset(new typename Traits::Functional_HH(args...));
      F_HE.reset(new typename Traits::Functional_HE(args...));
    }

        template <typename... Args>
    void initExtensionFunctionals(const Args&... args)
    {
      F_EH.reset(new typename Traits::Functional_EH(args...));
      F_EE.reset(new typename LumpedTraits::Functional_EE(args...));
    }

  private:
    void assembleAt(const Bridge::Vector<typename Traits::VarSet>& xl, const Bridge::Vector<typename Traits::ExtensionVarSet>& xe)
    {
      ass_HH->assemble(linearization(*F_HH,xl.get()));
      ass_HE->assemble(linearization(*F_HE,xl.get()));
      ass_EH->assemble(linearization(*F_EH,xe.get()));
      ass_EE->assemble(linearization(*F_EE,xe.get()));
    }

    std::unique_ptr<typename Traits::Functional_HH> F_HH;
    std::unique_ptr<typename Traits::Functional_HE> F_HE;
    std::unique_ptr<typename Traits::Functional_EH> F_EH;
    std::unique_ptr<typename LumpedTraits::Functional_EE> F_EE;
    std::unique_ptr<typename Traits::Assembler_HH> ass_HH;
    //typename Traits::Assembler_HH* ass_HH;
    std::unique_ptr<typename Traits::Assembler_HE> ass_HE;
    std::unique_ptr<typename Traits::Assembler_EH> ass_EH;
    std::unique_ptr<typename LumpedTraits::Assembler_EE> ass_EE;
    NormFunctional& normFunctional;
    VariableSetDescription& variableSetDescription;
    ExtensionVariableSetDescription& extensionVariableSetDescription;
    Scalar squaredFraction;
    bool verbose;
    size_t chebySteps = 30, mgSteps = 25, mgSmoothingSteps = 20;
    Scalar relativeAccuracy = 1e-3;
  };
}

#endif


#include <boost/numeric/interval/rounded_arith.hpp>


/*
 * Copyright 2018-, Olivier Stasse LAAS-CNRS
 *
 * This file is part of sot-torque-control.
 * sot-torque-control is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * as published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 * sot-torque-control is distributed in the hope that it will be
 * useful, but WITHOUT ANY WARRANTY; without even the implied warranty
 * of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.  You should
 * have received a copy of the GNU Lesser General Public License along
 * with sot-torque-control.  If not, see <http://www.gnu.org/licenses/>.
 */

#define EIGEN_NO_MALLOC
#include <Eigen/Dense>
#include <dynamic-graph/factory.h>
#include <sot/core/debug.hh>
#include <sot/torque_control/commands-helper.hh>
#include <sot/torque_control/ddp_pyrene_actuator_solver.hh>
#include <math.h>

#if DEBUG
#define ODEBUG(x) std::cout << x << std::endl
#else
#define ODEBUG(x)
#endif
#define ODEBUG3(x) std::cout << x << std::endl

#define DBGFILE "/tmp/debug-ddp_pyrene_actuator_solver.dat"

#define RESETDEBUG5() { std::ofstream DebugFile;  \
    DebugFile.open(DBGFILE,std::ofstream::out);   \
    DebugFile.close();}
#define ODEBUG5FULL(x) { std::ofstream DebugFile; \
    DebugFile.open(DBGFILE,std::ofstream::app);   \
    DebugFile << __FILE__ << ":"      \
              << __FUNCTION__ << "(#"     \
              << __LINE__ << "):" << x << std::endl;  \
    DebugFile.close();}
#define ODEBUG5(x) { std::ofstream DebugFile; \
    DebugFile.open(DBGFILE,std::ofstream::app); \
    DebugFile << x << std::endl;    \
    DebugFile.close();}

#define RESETDEBUG4()
#define ODEBUG4FULL(x)
#define ODEBUG4(x)

namespace dynamicgraph
{
namespace sot
{
namespace torque_control
{

namespace dynamicgraph = ::dynamicgraph;
using namespace dynamicgraph;
using namespace dynamicgraph::command;
using namespace Eigen;

#define ALL_INPUT_SIGNALS m_pos_desSIN << m_pos_joint_measureSIN << m_dx_joint_measureSIN << m_tau_desSIN

#define ALL_OUTPUT_SIGNALS  m_tauSOUT

/// Define EntityClassName here rather than in the header file
/// so that it can be used by the macros DEFINE_SIGNAL_**_FUNCTION.
typedef DdpPyreneActuatorSolver EntityClassName;

/* --- DG FACTORY ------------------------------------------------------- */
DYNAMICGRAPH_FACTORY_ENTITY_PLUGIN(DdpPyreneActuatorSolver, "DdpPyreneActuatorSolver");

DdpPyreneActuatorSolver::
DdpPyreneActuatorSolver(const std::string &name)
  : Entity(name),
    CONSTRUCT_SIGNAL_IN (pos_des,           dynamicgraph::Vector),
    CONSTRUCT_SIGNAL_IN (pos_joint_measure, dynamicgraph::Vector),
    CONSTRUCT_SIGNAL_IN (dx_joint_measure,  dynamicgraph::Vector),
    CONSTRUCT_SIGNAL_IN (tau_des,           dynamicgraph::Vector),
    CONSTRUCT_SIGNAL_OUT(tau,               dynamicgraph::Vector, ALL_INPUT_SIGNALS),
    m_dt(1e-3),
    m_T(50),
    m_stopCrit(1e-3),
    m_iterMax(10),
    m_solver(m_model, m_cost,
             DISABLE_FULLDDP,
             DISABLE_QPBOX)
    {

      RESETDEBUG5();
      Entity::signalRegistration( ALL_INPUT_SIGNALS << ALL_OUTPUT_SIGNALS );

      m_zeroState.setZero();

      /* Commands. */
      addCommand("init",
                 makeCommandVoid4(*this, &DdpPyreneActuatorSolver::param_init,
                                  docCommandVoid4("Initialize the DDP solver.",
                                      "Control timestep [s].",
                                      "Size of the preview window (in nb of samples)",
                                      "Max. nb. of iterations",
                                      "Stopping criteria")));
      addCommand("setTorqueLimit",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setTorqueLimit,
                            docCommandVoid1("Set the Torque limit.",
                                "Limit of the motor torque.")));
      addCommand("setJointLimit",
           makeCommandVoid2(*this, &DdpPyreneActuatorSolver::setJointLimit,
                            docCommandVoid2("Set the angular limits of the joint.",
                                "Upper limit",
                                "Lower limit.")));
      addCommand("setJointVelLimit",
           makeCommandVoid2(*this, &DdpPyreneActuatorSolver::setJointVelLimit,
                            docCommandVoid2("Set the angular velocity limits of the joint.",
                                "Upper limit",
                                "Lower limit.")));
      addCommand("setLoadParam",
           makeCommandVoid3(*this, &DdpPyreneActuatorSolver::setLoadParam,
                            docCommandVoid3("Setter of the Load parameters.",
                                "Mass of the load [g].",
                                "X coordinate of the Load",
                                "Y coordinate of the Load")));
      addCommand("setLoadMass",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setLoadMass,
                            docCommandVoid1("Set the Load mass.",
                                "Mass of the load [g].")));
      addCommand("removeLoad",
            makeCommandVoid0(*this, &DdpPyreneActuatorSolver::removeLoad,
                        docCommandVoid0("Remove the Load.")));

      addCommand("setCostGainState",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setCostGainState,
                            docCommandVoid1("Set the Gain of the state cost matrix.",
                                "Matrix of Gains.")));
      addCommand("setCostGainCommand",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setCostGainCommand,
                            docCommandVoid1("Set the Gain of the command cost matrix.",
                                "Matrix of Gains.")));
      addCommand("setCostGainStateConstraint",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setCostGainStateConstraint,
                            docCommandVoid1("Set the Gain of the constraints on the state.",
                                "Matrix of Gains.")));
      addCommand("setCostGainTorqueConstraint",
           makeCommandVoid1(*this, &DdpPyreneActuatorSolver::setCostGainTorqueConstraint,
                            docCommandVoid1("Set the Gain of the torque constraints.",
                                "Matrix of Gains.")));

      m_initSucceeded = true;
    }

/* --- COMMANDS ---------------------------------------------------------- */
DEFINE_SIGNAL_OUT_FUNCTION(tau, dynamicgraph::Vector)
{
    if (!m_initSucceeded)
    {
      SEND_WARNING_STREAM_MSG("Cannot compute signal tau before initialization!");
      return s;
    }

  /// ---- Get the information -----
  /// Desired position
  const dynamicgraph::Vector &
  pos_des = m_pos_desSIN(iter);
  /// Measured joint position
  const dynamicgraph::Vector &
  pos_joint_measure = m_pos_joint_measureSIN(iter);
  /// Measured joint speed
  const dynamicgraph::Vector &
  dx_joint_measure = m_dx_joint_measureSIN(iter);
  /// Desired torque
  const dynamicgraph::Vector &
  tau_des = m_tau_desSIN(iter);

  DDPSolver<double, 2, 1>::stateVec_t xinit, xDes;

  xinit << pos_joint_measure(0),
           dx_joint_measure(0);

  xDes << pos_des, 0.0;
  //ODEBUG(xDes);

  m_solver.initSolver(xinit, xDes);
  //ODEBUG("FirstInitSolver");

  /// --- Solve the DDP ---
  m_solver.solveTrajectory();
  //ODEBUG("Trajectory solved");

  /// --- Get the command ---
  DDPSolver<double, 2, 1>::traj lastTraj;
  lastTraj = m_solver.getLastSolvedTrajectory();
  //ODEBUG("getLastSolvedTrajectory");

  DDPSolver<double, 2, 1>::commandVecTab_t uList;
  uList = lastTraj.uList;

  s = m_previous_tau;
  double tau_ddp = uList[0](0,0);
  if (!isnan(tau_ddp))
  {
    s[25] = tau_ddp;
  }

  m_previous_tau = s;
  //ODEBUG(s);

  return s;
}

void DdpPyreneActuatorSolver::param_init(const double &timestep,
                                   const int &T,
                                   const int &nbItMax,
                                   const double &stopCriteria)
{
  if (!m_pos_desSIN.isPlugged())
    return SEND_MSG("Init failed: signal pos_des is not plugged", MSG_TYPE_ERROR);
  if (!m_pos_joint_measureSIN.isPlugged())
    return SEND_MSG("Init failed: signal pos_joint_measure is not plugged", MSG_TYPE_ERROR);
  if (!m_dx_joint_measureSIN.isPlugged())
    return SEND_MSG("Init failed: signal dx_joint_measure is not plugged", MSG_TYPE_ERROR);

  m_previous_tau.resize(32);
  m_previous_tau.setZero();
  m_T = T;
  m_dt = timestep;
  m_iterMax = nbItMax;
  m_stopCrit = stopCriteria;
  m_cost.setTauLimit(70.0);
  m_cost.setJointLimit(0.0, -2.35619449019);
  m_cost.setJointVelLimit(30.0, -30.0);
  m_solver.FirstInitSolver( m_zeroState, m_zeroState,
                            m_T , m_dt, m_iterMax, m_stopCrit);
}

void DdpPyreneActuatorSolver::setTorqueLimit(const double& tau)
{
  m_cost.setTauLimit(tau);
}

void DdpPyreneActuatorSolver::setJointLimit(const double& upperLim, const double& lowerLim)
{
  m_cost.setJointLimit(upperLim, lowerLim);
}

void DdpPyreneActuatorSolver::setJointVelLimit(const double& upperLim, const double& lowerLim)
{
  m_cost.setJointVelLimit(upperLim, lowerLim);
}

void DdpPyreneActuatorSolver::setLoadParam(const double& mass, const double& coordX, const double& coordY)
{
  m_model.setLoadParam(mass, coordX, coordY);
}

void DdpPyreneActuatorSolver::setLoadMass(const double& mass)
{
  m_model.setLoadMass(mass);
}

void DdpPyreneActuatorSolver::removeLoad()
{
  m_model.removeLoad();
}

void DdpPyreneActuatorSolver::setCostGainState(const dynamicgraph::Vector& Q)
{
  const CostFunction<double,2,1>::stateMat_t Q_new = Eigen::Map<const CostFunction<double,2,1>::stateMat_t, Eigen::Unaligned >(Q.data(),2,1);
  m_cost.setCostGainState(Q_new);
}

void DdpPyreneActuatorSolver::setCostGainStateConstraint(const dynamicgraph::Vector& W)
{
  const CostFunction<double,2,1>::stateMat_t W_new = Eigen::Map<const CostFunction<double,2,1>::stateMat_t, Eigen::Unaligned >(W.data(),2,1);
  m_cost.setCostGainStateConstraint(W_new);
}

void DdpPyreneActuatorSolver::setCostGainCommand(const dynamicgraph::Vector& R)
{
  const CostFunction<double,2,1>::commandMat_t R_new = Eigen::Map<const CostFunction<double,2,1>::commandMat_t, Eigen::Unaligned >(R.data(),1);
  m_cost.setCostGainCommand(R_new);
}

void DdpPyreneActuatorSolver::setCostGainTorqueConstraint(const dynamicgraph::Vector& P)
{
  const CostFunction<double,2,1>::commandMat_t P_new = Eigen::Map<const CostFunction<double,2,1>::commandMat_t, Eigen::Unaligned >(P.data(),1);
  m_cost.setCostGainTorqueConstraint(P_new);
}

void DdpPyreneActuatorSolver::display(std::ostream &os) const
{
  os << " T: " << m_T
     << " timestep: " << m_dt
     << " nbItMax: " << m_iterMax
     << " stopCriteria: " << m_stopCrit << std::endl;
}

} // namespace torque_control
} // namespace sot
} // namespace dynamicgraph


/* Copyright © 2019 Apple Inc. All rights reserved.
 *
 * Use of this source code is governed by a BSD-3-clause license that can
 * be found in the LICENSE.txt file or at
 * https://opensource.org/licenses/BSD-3-Clause
 */

#include <boost/gil/extension/numeric/resample.hpp>
#include <boost/gil/extension/numeric/sampler.hpp>

#include <toolkits/object_detection/one_shot_object_detection/util/mapping_function.hpp>
#include <toolkits/object_detection/one_shot_object_detection/util/superposition.hpp>

static const boost::gil::rgb8_pixel_t RGB_WHITE(255, 255, 255);
static const boost::gil::rgba8_pixel_t RGBA_WHITE(255, 255, 255, 0);

namespace turi {
namespace one_shot_object_detection {
namespace data_augmentation {

void superimpose_image(const boost::gil::rgb8_image_t::view_t &superimposed,
                       const boost::gil::rgba8_image_t::view_t &transformed,
                       const boost::gil::rgba8_image_t::view_t &background) {
  boost::gil::transform_pixels(
      transformed, background, superimposed,
      [](boost::gil::rgba8_ref_t a, boost::gil::rgba8_ref_t b) {
        boost::gil::red_t R;
        boost::gil::green_t G;
        boost::gil::blue_t B;
        boost::gil::alpha_t A;
        auto AoverB = [](size_t C_a, size_t C_b, size_t alpha_a,
                         size_t alpha_b) {
          return ((C_a * alpha_a + C_b * alpha_b * (1 - alpha_a)) /
                  (alpha_a + alpha_b * (1 - alpha_a)));
        };
        return boost::gil::rgb8_image_t::value_type(
            AoverB(get_color(a, R), get_color(b, R), get_color(a, A) / 255,
                   get_color(b, A) / 255),
            AoverB(get_color(a, G), get_color(b, G), get_color(a, A) / 255,
                   get_color(b, A) / 255),
            AoverB(get_color(a, B), get_color(b, B), get_color(a, A) / 255,
                   get_color(b, A) / 255));
      });
}

flex_image create_synthetic_image(
    const boost::gil::rgba8_image_t::const_view_t &starter_image_view,
    const boost::gil::rgb8_image_t::const_view_t &background_view,
    ParameterSampler &parameter_sampler) {
  boost::gil::rgba8_image_t background_rgba(
      boost::gil::rgba8_image_t::point_t(background_view.dimensions()));
  boost::gil::copy_and_convert_pixels(background_view, view(background_rgba));
  boost::gil::rgba8_image_t transformed(
      boost::gil::rgba8_image_t::point_t(background_view.dimensions()));
  fill_pixels(view(transformed), RGBA_WHITE);
  boost::gil::rgb8_image_t superimposed(
      boost::gil::rgba8_image_t::point_t(background_view.dimensions()));
  fill_pixels(view(superimposed), RGB_WHITE);
  Eigen::Matrix<float, 3, 3> M = parameter_sampler.get_transform().inverse();
  resample_pixels(starter_image_view, view(transformed), M,
                  boost::gil::nearest_neighbor_sampler());
  superimpose_image(view(superimposed), view(transformed),
                    view(background_rgba));
  return flex_image(superimposed);
}

}  // namespace data_augmentation
}  // namespace one_shot_object_detection
}  // namespace turi


#include "model/Model.hpp"
#include "utils/AminoAcid.hpp"
#include <Eigen/Geometry>
using namespace std;
using namespace mdk;

Model::Residue& Model::addResidue(Chain *chain) {
    auto& res = residues.emplace_back();
    res.idx = residues.size()-1;
    if (chain) {
        res.chainIdx = chain->idx;
        if (chain->start < 0) chain->start = res.idx;
        chain->end = res.idx+1;
    }
    else {
        res.chainIdx = -1;
    }
    ++n;
    return res;
}

Model::Chain& Model::addChain() {
    auto& chain = chains.emplace_back();
    chain.idx = chains.size()-1;
    chain.start = chain.end = -1;
    return chain;
}

Model::Contact &Model::addContact() {
    auto& cont = contacts.emplace_back();
    cont.idx = contacts.size()-1;
    return cont;
}

Model::StructuredPart& Model::addSP() {
    auto& sp = structuredParts.emplace_back();
    sp.idx = structuredParts.size()-1;
    return sp;
}

void Model::morphIntoLine() {
    double tether0 = 3.8*angstrom;

    Eigen::Vector3d dir = { 0.0, 0.0, tether0 };
    for (auto& res: residues) {
        res.r = res.idx * dir;
    }
}

static Eigen::Vector3d sampleBox(Eigen::AlignedBox3d const& box, Random& rand) {
    Eigen::Vector3d v;
    for (int i = 0; i < 3; ++i) {
        v[i] = rand.uniform(box.min()[i], box.max()[i]);
    }
    return v;
}

static Eigen::Vector3d sampleSphere(Random& rand) {
    Eigen::Vector3d v;
    for (int i = 0; i < 3; ++i) {
        v[i] = rand.normal();
    }
    return v.normalized();
}

/**
 * Find a vector perpendicular to \p v.
 * @param v
 * @return
 */
static Eigen::Vector3d findPerp(Eigen::Vector3d const& v) {
    auto perp1 = v.cross(Eigen::Vector3d::UnitX());
    auto perp2 = v.cross(Eigen::Vector3d::UnitY());
    auto& perp = perp1.norm() > perp2.norm() ? perp1 : perp2;
    return perp.normalized();
}

void Model::morphIntoSAW(Random& rand, bool useTop, double density,
        double minDist) {
    Eigen::Vector3d minCorner, maxCorner;
    Eigen::AlignedBox3d box;

    double tether0 = 3.8*angstrom;
    auto pairs = nonlocalPairs();
    auto state0 = residues;

    if (density > 0) {
        auto vol = residues.size() / (density * atom);
        auto a = 0.5 * pow(vol, 1.0/3.0);
        minCorner = -a * Eigen::Vector3d { 1.0, 1.0, 1.0 };
        maxCorner = -minCorner;
    }
    else {
        minCorner = maxCorner = Eigen::Vector3d { 0.0, 0.0, 0.0 };
    }
    box = Eigen::AlignedBox3d(minCorner, maxCorner);
    if (useTop) top.setCell(box.max() - box.min());

    for (int attempt = 0; attempt < 9000; ++attempt) {
        bool failed = false;

        for (auto const& chain: chains) {
            auto pos = sampleBox(box, rand);
            auto dir = sampleSphere(rand);

            for (int i = chain.start; i < chain.end; ++i) {
                residues[i].r = pos;
                pos += tether0 * dir;

                double spread = rand.uniform(M_PI/12.0, M_PI/2.0);
                auto spreadRot = Eigen::AngleAxisd(spread, findPerp(dir));

                double around = rand.uniform(-M_PI, M_PI);
                auto aroundRot = Eigen::AngleAxisd(around, dir);

                dir = aroundRot * spreadRot * dir;
                dir.normalize();
            }
        }

        for (auto& res: residues) {
            if (res.chainIdx < 0)
                res.r = sampleBox(box, rand);
        }

        for (auto const& [res1, res2]: pairs) {
            Eigen::Vector3d dx = res1->r - res2->r;
            if (useTop) top.fix(dx);

            if (dx.norm() < minDist) {
                failed = true;
                break;
            }
        }

        if (!failed) return;
    }

    residues = state0;
}

Vectors genSAW(int n, double bondLen, Random& rand) {
    vector<double> phi(n);   // around
    vector<double> theta(n); // spread

    phi[0] = 3.0 * M_PI / 2.;
    phi[1] = M_PI;
    theta[0] = 0.;
    theta[1] = rand.uniform(0.0, M_PI / 3.0);
    for (int i = 2; i + 1 < n; ++i) {
        phi[i] = rand.uniform(0.0, 2.0 * M_PI);
        theta[i] = rand.uniform(0.0, M_PI / 3.0);
    }

    double dir_theta = M_PI / 2.0 - acos(rand.uniform(-1.0, 1.0));
    double dir_phi = rand.uniform(0.0, 2.0 * M_PI);

    auto dir = Eigen::Vector3d{cos(dir_theta)*cos(dir_phi),
                               -cos(dir_theta)*sin(dir_phi),
                               sin(dir_theta)};

    auto rot_matrix = Eigen::AngleAxisd::Identity();

    Vector pos = Vector::Zero();

    Vectors res(n);
    res[0] = pos;
    for(int i = 1; i < n; i++) {
        rot_matrix = rot_matrix *
            Eigen::AngleAxisd(phi[i - 1], Eigen::Vector3d::UnitX()) *
            Eigen::AngleAxisd(theta[i - 1], Eigen::Vector3d::UnitZ());
        pos += rot_matrix * dir * bondLen;
        res[i] = pos;
    }

    auto flip = Eigen::AngleAxisd(M_PI, Eigen::Vector3d::UnitX());
    for (int i = 0; i < n; ++i) res[i] = flip * res[i];
    return res;
}

// SAW implementation that mirrors cg.f:confstart
// Assumes that each residue is in some chain
// and chains vector is sorted
void Model::legacyMorphIntoSAW(Random& rand, bool useTop, double density,
                               double minDist, bool nativeBondLen) {

    Eigen::Vector3d minCorner, maxCorner;
    Eigen::AlignedBox3d box;

    double tether0 = (float) 3.8*angstrom;
    if (nativeBondLen) {
        double avgBondLen = 0.0;
        for (Chain const& chain : chains) {
            for (int i = chain.start; i + 1 < chain.end; ++i) {
                avgBondLen += (residues[i + 1].r - residues[i].r).norm();
            }
        }
        // TODO: fix this bug in cg.f and then here
        avgBondLen /= (int) residues.size() - 1;
        tether0 = avgBondLen;
    }

    auto pairs = nonlocalPairs();
    auto state0 = residues;

    if (density > 0) {
        auto vol = residues.size() / (density * atom);
        auto a = 0.5 * pow(vol, 1.0/3.0);
        minCorner = -a * Eigen::Vector3d { 1.0, 1.0, 1.0 };
        maxCorner = -minCorner;
    }
    else {
        minCorner = maxCorner = Eigen::Vector3d { 0.0, 0.0, 0.0 };
    }
    box = Eigen::AlignedBox3d(minCorner, maxCorner);
    if (useTop) top.setCell(box.max() - box.min());

    for (auto& res: residues) {
        // this SAW implementation doesn't support residues without chains
        assert (res.chainIdx >= 0);
    }

    for (auto const& chain: chains) {
        int attemptsLeft = 9000;
        bool success = false;
        while (not success && attemptsLeft > 0) {
            attemptsLeft--;

            Vectors saw = genSAW(chain.end - chain.start, tether0, rand);
            auto displacement = sampleBox(box, rand);
            for (int i = chain.start; i < chain.end; ++i) {
                residues[i].r = saw[i - chain.start] + displacement;
            }
            success = true;

            double minDistSquared = minDist * minDist;
            for (int i = 0; i < chain.end && success; ++i) {
                for (int j = i + 3; j < chain.end; ++j) {
                    Eigen::Vector3d vec = residues[i].r - residues[j].r;
                    if (useTop) top.fix(vec);

                    if (vec.squaredNorm() < minDistSquared) {
                        success = false;
                        break;
                    }
                }
            }
        }

        assert (success /* morphing into SAW failed */);
    }
}

void Model::exactLegacySAW(Random &rand, bool useTop, double density,
    double minDist, double bond) {

    auto minDistSq = pow(minDist, 2.0);

    Eigen::AlignedBox3d box;
    if (density > 0) {
        auto startvolume = residues.size() / density;
        auto startboxsize = 0.5 * pow(startvolume, 1.0/3.0);
        box.min() = Vector::Constant(-startboxsize);
        box.max() = -box.min();

        if (useTop) {
            top.setCell(box.max() - box.min());
        }
    }
    else {
        if (useTop) throw;

        box.min() = box.max() = Vector::Zero();
    }

    for (auto& chain: chains) {
        int n = chain.end - chain.start;
        for (int kter = 0; kter < 9000; ++kter) {
            std::vector<double> phi(n), theta(n);

            auto pi = acos(-1.0);

            phi[0] = pi / 2.0;
            theta[0] = 0.0;

            phi[1] = 0.0;
            theta[1] = rand.uniform(0.0, M_PI/3.0);

            for (int i = 2; i < n - 1; ++i) {
                phi[i] = rand.uniform(-M_PI, M_PI);
                theta[i] = rand.uniform(0.0, M_PI/3.0);
            }

            std::vector<Eigen::Matrix3d> T(n);
            for (int i = 0; i < n - 1; ++i) {
                auto ct = cos(theta[i]), st = sin(theta[i]),
                     cp = cos(phi[i]), sp = sin(phi[i]);

                T[i] << ct,      st,       0.0,
                        st * cp, -ct * cp,  sp,
                        st * sp, -ct * sp, -cp;
            }

            auto theta0 = acos(-rand.uniform(-1.0, 1.0));
            auto phi0 = rand.uniform(0.0, 2.0 * pi);

            std::vector<Vector> R(n);
            for (int i = 0; i < n - 1; ++i) {
                Vector r;
                r << bond * sin(theta0) * cos(phi0),
                     bond * sin(theta0) * sin(phi0),
                     bond * cos(theta0);

                for (int j = i; j >= 0; --j) {
                    R[i] = T[j] * r;
                    r = R[i];
                }
            }

            Vector ran = sampleBox(box, rand);
            for (int i = 0; i < n; ++i) {
                Vector r = ran;
                for (int j = 0; j < i; ++j) {
                    r += R[j];
                }
                residues[i].r = r;
            }

            bool nonIntersecting = true;
            for (int i = 0; i < n - 3 && nonIntersecting; ++i) {
                auto r1 = residues[i].r;
                for (int j = i + 3; j < n && nonIntersecting; ++j) {
                    auto r2 = residues[j].r;
                    Vector dx = !useTop ? (r2 - r1) : top(r2 - r1);
                    if (dx.squaredNorm() < minDistSq) {
                        nonIntersecting = false;
                    }
                }
            }

            if (nonIntersecting) break;
        }
    }
}

void Model::initVelocity(Random& rand, double temperature, bool useMass) {
    int n = residues.size();
    vector<double> masses(n, 1.0 * f77mass);

    if (useMass) {
        //TODO
        assert(0);
    }

    for (int i = 0; i < n; ++i) {
        //TODO: do this only in LEGACY_MODE
        Vector& v = residues[i].v;
        for (int d = 0; d < 3; ++d) {
            v(d) = rand.uniform(-1.0, 1.0);
        }
        v.normalize();
    }

    // Shifting velocities so that total momentum is 0

    double totMass = 0.0;
    Vector totMomentum = Vector::Zero();

    for (int i = 0; i < n; ++i) {
        totMass += masses[i];
        totMomentum += masses[i] * residues[i].v;
    }

    for (int i = 0; i < n; ++i) {
        residues[i].v -= totMomentum * (1.0 / totMass);
    }

    // Scaling velocities to desired temperature

    double avgKinEnergy = 0.0;
    for (int i = 0; i < n; ++i) {
        avgKinEnergy += masses[i] * residues[i].v.squaredNorm() / 2.0;
    }
    avgKinEnergy /= n;

    double ratio = 1.5 * temperature / avgKinEnergy;
    for (int i = 0; i < n; ++i) {
        residues[i].v *= sqrt(ratio);
    }
}

Model::StructuredPart& Model::addContactMap(cmap::ContactMap const& contactMap) {
    auto& sp = addSP();
    sp.off = contactMap.offset;
    sp.len = contactMap.len;
    sp.angle = contactMap.angle;
    sp.dihedral = contactMap.dihedral;
    return sp;
}

void Model::addCMapContacts(cmap::ContactMap const& contactMap, Chain &chain) {
    for (auto const& cmapCont: contactMap.contacts) {
        auto& modelCont = addContact();
        modelCont.type = ContactType(ContactTypeIdx::NAT);
        modelCont.dist0 = cmapCont.dist0;
        for (int i = 0; i < 2; ++i) {
            auto resIdx = chain.start + contactMap.offset + cmapCont.res[i];
            modelCont.res[i] = resIdx;
        }
    }

    auto sp = addContactMap(contactMap);
    chain.structuredParts.push_back(sp.idx);
}

vector<pair<Model::Residue*, Model::Residue*>> Model::nonlocalPairs() {
    vector<pair<Residue*, Residue*>> pairs;
    for (auto& res1: residues) {
        for (auto& res2: residues) {
            if (res1.chainIdx == res2.chainIdx && res1.chainIdx >= 0)
                continue;
            if (abs(res1.idx - res2.idx) < 3)
                continue;

            pairs.emplace_back(&res1, &res2);
        }
    }

    return pairs;
}

void Model::useAverageMasses() {
    double sum = 0.0;
    for (auto& acid: AminoAcid::aminoAcids())
        sum += acid.info().mass;
    sum /= AminoAcid::N;

    for (auto& res: residues)
        res.mass = sum;
}


/**
 * @addtogroup DFNs
 * @{
 */

#include "ImageRectification.hpp"
#include <Eigen/Core>

namespace CDFF
{
namespace DFN
{
namespace ImageRectification
{

ImageRectification::ImageRectification()
:parameters(DEFAULT_PARAMETERS)
{
    parametersHelper.AddParameter<int>("ImageRectificationParams", "xratio", parameters.xratio, DEFAULT_PARAMETERS.xratio);
    parametersHelper.AddParameter<int>("ImageRectificationParams", "yratio", parameters.yratio, DEFAULT_PARAMETERS.yratio);
    parametersHelper.AddParameter<double>("ImageRectificationParams", "scaling", parameters.scaling, DEFAULT_PARAMETERS.scaling);
    parametersHelper.AddParameter<bool>("ImageRectificationParams", "centerPrincipalPoint", parameters.centerPrincipalPoint, DEFAULT_PARAMETERS.centerPrincipalPoint);
    parametersHelper.AddParameter<bool>("ImageRectificationParams", "fisheye", parameters.fisheye, DEFAULT_PARAMETERS.fisheye);

    configurationFilePath = "";
}

ImageRectification::~ImageRectification()
{
}

void ImageRectification::configure()
{
    if(configurationFilePath != ""){
        parametersHelper.ReadFile(configurationFilePath);
    }
    ValidateParameters();
}

void ImageRectification::process()
{
    cv::Mat in(static_cast<int>(inOriginalImage.data.rows), static_cast<int>(inOriginalImage.data.cols), CV_MAKETYPE(static_cast<int>(inOriginalImage.data.depth), static_cast<int>(inOriginalImage.data.channels)), inOriginalImage.data.data.arr, inOriginalImage.data.rowSize);
    cv::Mat out;

    // Generate correction maps if needed
    if( std::string(reinterpret_cast<char const *>(inOriginalImage.intrinsic.sensorId.arr)) != _sensorId ||
            parameters.xratio != _xratio ||
            parameters.yratio != _yratio ||
            parameters.scaling != _scaling ||
            parameters.centerPrincipalPoint != _centerPrincipalPoint ||
            parameters.fisheye != _fisheye){
        _sensorId = std::string(reinterpret_cast<char const *>(inOriginalImage.intrinsic.sensorId.arr));
        _xratio = parameters.xratio;
        _yratio = parameters.yratio;
        _scaling = parameters.scaling;
        _centerPrincipalPoint =parameters.centerPrincipalPoint;
        _fisheye = parameters.fisheye;

        cv::Mat1d cameraMatrix(3,3, Eigen::Map<Eigen::Matrix3d>(inOriginalImage.intrinsic.cameraMatrix.arr[0].arr, 3, 3).data());

        cv::Mat1d distCoeffs(inOriginalImage.intrinsic.distCoeffs.nCount, 1, Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, 1>>(inOriginalImage.intrinsic.distCoeffs.arr, inOriginalImage.intrinsic.distCoeffs.nCount, 1).data());

        if(_fisheye){
            cv::fisheye::estimateNewCameraMatrixForUndistortRectify(cameraMatrix, distCoeffs, cv::Size(in.cols, in.rows), cv::Matx33d::eye(), _newCameraMatrix, _scaling, cv::Size(in.cols / _xratio, in.rows / _yratio));
            cv::fisheye::initUndistortRectifyMap(cameraMatrix, distCoeffs, cv::Matx33d::eye(), _newCameraMatrix, cv::Size(in.cols / _xratio, in.rows / _yratio), CV_32F, _mapx, _mapy);
        }
        else{
            _newCameraMatrix = cv::getOptimalNewCameraMatrix(cameraMatrix, distCoeffs, cv::Size(in.cols, in.rows), _scaling, cv::Size(in.cols / _xratio, in.rows / _yratio), 0, _centerPrincipalPoint);
            cv::initUndistortRectifyMap(cameraMatrix, distCoeffs, cv::Mat(), _newCameraMatrix, cv::Size(in.cols / _xratio, in.rows / _yratio), CV_32F, _mapx, _mapy);
        }
    }

    cv::remap(in, out, _mapx, _mapy, cv::INTER_LINEAR);

    // Getting image
    {
        // init the structure
        outRectifiedImage.msgVersion = frame_Version;

        outRectifiedImage.intrinsic = inOriginalImage.intrinsic;

        Eigen::Map<Eigen::Matrix3d>(outRectifiedImage.intrinsic.cameraMatrix.arr[0].arr, 3, 3) = Eigen::Map<Eigen::Matrix3d>(reinterpret_cast<double *>(_newCameraMatrix.data), 3, 3);

        Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, 1>>(outRectifiedImage.intrinsic.distCoeffs.arr, inOriginalImage.intrinsic.distCoeffs.nCount, 1) = Eigen::Matrix<double, Eigen::Dynamic, 1>::Zero(inOriginalImage.intrinsic.distCoeffs.nCount, 1);
        outRectifiedImage.intrinsic.distCoeffs.nCount = inOriginalImage.intrinsic.distCoeffs.nCount;

        if(_fisheye){
            outRectifiedImage.intrinsic.cameraModel = asn1Scccam_FISHEYE;
        }
        else{
            outRectifiedImage.intrinsic.cameraModel = asn1Scccam_PINHOLE;
        }

        outRectifiedImage.extrinsic = inOriginalImage.extrinsic;
        outRectifiedImage.metadata = inOriginalImage.metadata;

        // Array3D
        {
            outRectifiedImage.data.msgVersion = array3D_Version;
            outRectifiedImage.data.rows = static_cast<asn1SccT_UInt32>(out.rows);
            outRectifiedImage.data.cols = static_cast<asn1SccT_UInt32>(out.cols);
            outRectifiedImage.data.channels = static_cast<asn1SccT_UInt32>(out.channels());
            outRectifiedImage.data.depth = static_cast<asn1SccArray3D_depth_t>(out.depth());
            outRectifiedImage.data.rowSize = static_cast<asn1SccT_UInt32>(out.step[0]);
            outRectifiedImage.data.data.nCount = static_cast<int>(outRectifiedImage.data.rows * outRectifiedImage.data.rowSize);
            memcpy(outRectifiedImage.data.data.arr, out.data, static_cast<size_t>(outRectifiedImage.data.data.nCount));
        }
    }

}

void ImageRectification::ValidateParameters()
{
    ASSERT(parameters.xratio >= 1 && parameters.xratio <= 25, "xratio has to be within [1..25]");
    ASSERT(parameters.yratio >= 1 && parameters.yratio <= 25, "yratio has to be within [1..25]");
    ASSERT((parameters.scaling >= 0 && parameters.scaling <= 1) || parameters.scaling == -1, "scaling has to be within [0..1] or equal to -1");
}

const ImageRectification::ImageRectificationParams ImageRectification::DEFAULT_PARAMETERS = {
    .xratio = 1,
    .yratio = 1,
    .scaling = -1,
    .centerPrincipalPoint = false,
    .fisheye = false
};

}
}
}

/** @} */


#define BOOST_TEST_MODULE C_EXEC_TEST

#include <dynet_c/init.h>
#include <dynet_c/dim.h>
#include <dynet_c/expr.h>
#include <dynet_c/graph.h>
#include <dynet_c/model.h>
#include <dynet_c/rnn-builder.h>
#include <dynet_c/tensor.h>
#include <vector>
#include <boost/test/unit_test.hpp>
#include "test_utils.h"

struct CExecTest {
  CExecTest() {
    const char *argv[] = {
        "ExecTest",
        "--dynet-seed", "10",
        "--dynet-mem", "10",
        "--dynet-autobatch", "1"
    };
    ::dynetDynetParams_t *params;
    ::dynetExtractDynetParams(7, const_cast<char**>(argv), false, &params);
    ::dynetInitialize(params);
    ::dynetDeleteDynetParams(params);
  }
};

BOOST_FIXTURE_TEST_SUITE(c_exec_test, CExecTest);

BOOST_AUTO_TEST_CASE(autobatch_lstm_gradient) {
  std::vector<float> results;
  ::dynetParameterCollection_t *model;
  ::dynetRNNBuilder_t *builder;
  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetCreateParameterCollection(&model));
  BOOST_CHECK_EQUAL(DYNET_C_OK,
                    ::dynetCreateVanillaLSTMBuilder(
                        2, 3, 10, model, false, 1.f, &builder));
  uint32_t dims[] = {3};
  ::dynetLookupParameter_t *lp;
  ::dynetDim_t *dim;
  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetCreateDimWithDimensions(dims, 1, &dim));
  BOOST_CHECK_EQUAL(DYNET_C_OK,
                    ::dynetAddLookupParametersToParameterCollection(
                        model, 10, dim, nullptr, nullptr, nullptr, &lp));

  for (size_t i = 0; i < 3; ++i) {
    ::dynetComputationGraph_t *cg;
    BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetCreateComputationGraph(&cg));
    BOOST_CHECK_EQUAL(DYNET_C_OK,
                      ::dynetResetRNNBuilderGraph(builder, cg, true));
    std::vector<::dynetExpression_t*> losses;
    for (size_t j = 0; j < 3; ++j) {
      ::dynetExpression_t *h_0[] = {};
      BOOST_CHECK_EQUAL(DYNET_C_OK,
                        ::dynetStartRNNBuilderNewSequence(
                            builder, h_0, 0));
      for (size_t k = 0; k < 3; ++k) {
        ::dynetExpression_t *x;
        BOOST_CHECK_EQUAL(DYNET_C_OK,
                          ::dynetApplyLookupOne(cg, lp, j * 3 + k, &x));
        ::dynetExpression_t *h;
        BOOST_CHECK_EQUAL(DYNET_C_OK,
                          ::dynetAddRNNBuilderInput(builder, x, &h));
        BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteExpression(h));
        BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteExpression(x));
      }
      std::size_t length = 0u;
      BOOST_CHECK_EQUAL(DYNET_C_OK,
                        ::dynetGetRNNBuilderFinalHiddenState(
                            builder, nullptr, &length));
      ::dynetExpression_t *hs[length];
      BOOST_CHECK_EQUAL(DYNET_C_OK,
                        ::dynetGetRNNBuilderFinalHiddenState(
                            builder, hs, &length));
      ::dynetExpression_t *loss;
      BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetApplySquaredNorm(hs[1], &loss));
      losses.emplace_back(loss);
      for (std::uint32_t t = 0; t < length; ++t) {
        BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteExpression(hs[t]));
      }
    }
    ::dynetExpression_t *loss;
    BOOST_CHECK_EQUAL(DYNET_C_OK,
                      ::dynetApplyAdd(losses[0], losses[2], &loss));
    losses.emplace_back(loss);
    ::dynetExpression_t *z;
    BOOST_CHECK_EQUAL(DYNET_C_OK,
                      ::dynetApplySum(&losses[0], losses.size(), &z));
    const ::dynetTensor_t *z_t;
    BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetGetExpressionValue(z, &z_t));
    float z_v = 0.f;
    BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetEvaluateTensorAsScalar(z_t, &z_v));
    results.emplace_back(z_v);
    BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteExpression(z));
    for (std::uint32_t t = 0; t < losses.size(); ++t) {
      BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteExpression(losses[t]));
    }
    BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteComputationGraph(cg));
  }
  for (size_t i = 1; i < results.size(); ++i) {
    BOOST_CHECK_CLOSE(results[0], results[i], 0.0001);
  }

  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteRNNBuilder(builder));
  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteDim(dim));
  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteLookupParameter(lp));
  BOOST_CHECK_EQUAL(DYNET_C_OK, ::dynetDeleteParameterCollection(model));
}

BOOST_AUTO_TEST_SUITE_END()


#include "similarity.h"
#include "types.h"
#include "random_pcg.h"

#include <iostream>
#include <cmath>
#include <armadillo>

#include <memory>

constexpr const static double HALF_PI = 1.57079632679490;
constexpr const static double PI =      3.14159265358979;
extern "C" {

double map_clusters( void *xi, void *xj,  py_int N,
                     py_int *ids, py_int *jds, py_int periodic,
                     void *xlo, void *xhi, py_int dims,
                     void *rotated, py_float *rot_axis )
{
	arr3f xxi(xi,N), xxj(xj,N);
	arr3f rot(rotated,N);

	test_rotations();

	py_float *xxlo = static_cast<py_float*>(xlo);
	py_float *xxhi = static_cast<py_float*>(xhi);

	return 0;
	/*
	return map_clusters_impl( xxi, xxj, N, ids, jds,
	                          periodic, xxlo, xxhi, dims,
	                          rot, rot_axis );
	*/
}

py_float rotate_to_template( void *xxi, void *xxt,  py_int N,
                             py_int periodic, void *xxlo, void *xxhi,
                             py_int dims, void *mmaj_ax, void *mmin_ax,
                             void *mmaj_ax_t, void *mmin_ax_t,
                             void *x_rotated )
{
	arr3f xi( xxi, N );
	arr3f xt( xxt, N );

	py_float *xlo      = static_cast<py_float*>( xxlo );
	py_float *xhi      = static_cast<py_float*>( xxhi );
	py_float *maj_ax_t = static_cast<py_float*>( mmaj_ax_t );
	py_float *min_ax_t = static_cast<py_float*>( mmin_ax_t );
	py_float *maj_ax   = static_cast<py_float*>( mmaj_ax );
	py_float *min_ax   = static_cast<py_float*>( mmin_ax );
	py_float *xrot     = nullptr;
	if( x_rotated ){
		// std::cerr << "Storing rotated values in array at " << x_rotated << ".\n";
		xrot = static_cast<py_float*>( x_rotated );
	}else{
		// std::cerr << "Not storing rotated values...\n";
	}

	return rotate_to_template_impl( xi, xt, N, periodic, xlo, xhi,
	                                dims, maj_ax, min_ax, maj_ax_t,
	                                min_ax_t, xrot );

}

} // extern "C"

template <typename T> inline
T sign( T& a )
{
	return (a > 0) - (a < 0);
}

void quat_prod( const double *q1, const double *q2, double *result )
{
	double a1 = q1[0];
	double b1 = q1[1];
	double c1 = q1[2];
	double d1 = q1[3];

	double a2 = q2[0];
	double b2 = q2[1];
	double c2 = q2[2];
	double d2 = q2[3];

	result[0] = a1*a2 - b1*b2 - c1*c2 - d1*d2;
	result[1] = a1*b2 + b1*a2 + c1*d2 - d1*c2;
	result[2] = a1*c2 - b1*d2 + c1*a2 + d1*b2;
	result[3] = a1*d2 + b1*c2 - c1*b2 + d1*a2;
}

double dot( const py_float *x, const py_float *y )
{
	return x[0]*y[0] + x[1]*y[1] + x[2]*y[2];
}

void cross( const py_float *x, const py_float *y, py_float *res )
{
	res[0] = x[1]*y[2] - x[2]*y[1];
	res[1] = x[2]*y[0] - x[0]*y[2];
	res[2] = x[0]*y[1] - x[1]*y[0];
}

double norm( const py_float *x, int N )
{
	double sum = 0;
	for( int i = 0; i < N; ++i ){
		sum += x[i]*x[i];
	}
	return sqrt(sum);
}

double norm3_2( const py_float *x )
{
	return x[0]*x[0] + x[1]*x[1] + x[2]*x[2];
}

double norm3( const py_float *x )
{
	return sqrt( norm3_2(x) );
}

double dist_2( const py_float *x, const py_float *y, py_int N )
{
	double r2 = 0.0;
	for( py_int i = 0; i < N; ++i ){
		double dx = x[i] - y[i];
		r2 += dx*dx;
	}
	return r2;
}

double dist3_2( const py_float *x, const py_float *y )
{
	py_float z[3];
	z[0] = x[0] - y[0];
	z[1] = x[1] - y[1];
	z[2] = x[2] - y[2];
	return norm3_2(z);
}

double dist3( const  py_float *x, const py_float *y )
{
	return sqrt( dist3_2(x,y) );
}

bool close_enough( const double *a, const double *b, int N, double tol = 1e-8 )
{
	double dr2 = dist3_2( a, b );

	double dx = a[0] - b[0];
	double dy = a[1] - b[1];
	double dz = a[2] - b[2];
	return dr2 < tol;
}

void rotate_vector( double *rotate, const double *rot_axis, double df )
{
	// Quaternion rotation:
	double q[4], p[4], qt[4];
	p[0] = 0;
	p[1] = rotate[0];
	p[2] = rotate[1];
	p[3] = rotate[2];

	double c = cos(0.5*df);
	double s = sin(0.5*df);
	q[0] = c;
	q[1] = rot_axis[0]*s;
	q[2] = rot_axis[1]*s;
	q[3] = rot_axis[2]*s;

	qt[0] =  q[0];
	qt[1] = -q[1];
	qt[2] = -q[2];
	qt[3] = -q[3];

	double result[4];
	quat_prod( q, p, result );
	quat_prod( result, qt, result );

	rotate[0] = result[1];
	rotate[1] = result[2];
	rotate[2] = result[3];
}

void gramm_schmid( const py_float *x1, const py_float *x2, py_float *res )
{
	double ab = x1[0]*x2[0] + x1[1]*x2[1] + x1[2]*x2[2];
	double b2 = norm3_2(x2);

	ab /= b2;

	res[0] = x1[0] - ab * x2[0]; // / b2;
	res[1] = x1[1] - ab * x2[1];
	res[2] = x1[2] - ab * x2[2];
}

// Tests the rotation and quaternion product functions.
void test_rotations()
{
	// Test quaternions:
	py_float q1[4], q2[4];
	q1[0] = 0;
	q1[1] = 1;
	q1[2] = 2;
	q1[3] = 0;

	q2[0] = 1;
	q2[1] = 1;
	q2[2] = q2[3] = 0;

	py_float q3[4];
	quat_prod( q1, q2, q3 );
	std::cerr << "q1 = ( " << q1[0] << ", " << q1[1] << ", " << q1[2] << ", " << q1[3] << " ).\n";
	std::cerr << "q2 = ( " << q2[0] << ", " << q2[1] << ", " << q2[2] << ", " << q2[3] << " ).\n";
	std::cerr << "q3 = ( " << q3[0] << ", " << q3[1] << ", " << q3[2] << ", " << q3[3] << " ).\n";
	// q3 should be
	py_float q3act[4];
	q3act[0] = -1;
	q3act[1] = 1;
	q3act[2] = 2;
	q3act[3] = -2;

	py_float x[3];
	py_float y[3];
	py_float z[3];

	x[0] = 1; x[1] = x[2] = 0;
	y[0] = y[2] = 0; y[1] = 1;
	z[0] = z[1] = 0; z[2] = 1;

	std::cerr << "\n\n";

	std::cerr << "x is now ( " << x[0] << ", " << x[1] << ", " << x[2] << " ).\n";
	rotate_vector( x, z, HALF_PI );
	std::cerr << "x is now ( " << x[0] << ", " << x[1] << ", " << x[2] << " ).\n";

	if( !close_enough( x, y, 3, 1e-10 ) ){
		std::cerr << "Error in rotations! Aborting.\n";
		std::terminate();
	}
}

void test_cross()
{
	py_float x[3], y[3], z[3];
	x[1] = x[2] = y[0] = y[2] = z[0] = z[1] = 0;
	x[0] = y[1] = z[2] = 1;
	py_float r[3];

	cross(x,y,r);
	std::cerr << "x cross y = ( " << r[0] << ", " << r[1] << ", " << r[2] << " ).\n";
	std::cerr << "should be   ( " << z[0] << ", " << z[1] << ", " << z[2] << " ).\n";

	cross(x,z,r);
	std::cerr << "x cross z = ( " << r[0] << ", " <<  r[1] << ", " << r[2] << " ).\n";
	std::cerr << "should be   ( " << y[0] << ", " << -y[1] << ", " << y[2] << " ).\n";

	cross(y,z,r);
	std::cerr << "y cross z = ( " << r[0] << ", " <<  r[1] << ", " << r[2] << " ).\n";
	std::cerr << "should be   ( " << x[0] << ", " <<  x[1] << ", " << x[2] << " ).\n";
	std::cerr << "\n\n\n";

}

double rotate_onto( const py_float *v, const py_float *target,
                    py_float *axis )
{
	// Finds axis and angle that rotate v __onto__ target.

	cross( v, target, axis );
	double n = 1.0 / norm3(axis);
	axis[0] *= n;
	axis[1] *= n;
	axis[2] *= n;

	double d = dot( v, target );
	d /= (norm3(v)*norm3(target));
	double angle = acos(d);

	return angle;
}

void copy_ensemble( const arr3f &x, arr3f &cp, py_int N )
{
	for( py_int i = 0; i < N; ++i ){
		cp[i][0] = x[i][0];
		cp[i][1] = x[i][1];
		cp[i][2] = x[i][2];

	}
}

void rotate_ensemble( arr3f &rot_data, const py_float *rot_ax,
                      py_int N, double angle )
{
	for( py_int i = 0; i < N; ++i ){
		rotate_vector( rot_data[i], rot_ax, angle );
	}
}

py_float rotate_to_template_impl( const arr3f &xi, const arr3f &xt, py_int N,
                                  py_int periodic, py_float *xlo, py_float *xhi,
                                  py_int dims, py_float *maj_ax, py_float *min_ax,
                                  py_float *maj_ax_t, py_float *min_ax_t,
                                  py_float *xrot )
{
	double *rot_data_;

	if( xrot ) rot_data_ = xrot;
	else       rot_data_ = new double[3*N];

	arr3f rot_data(rot_data_, N);

	double ang1, ang2;
	double rot_ax[3];

	// test_cross();

	// std::cerr << "Rotating " << N << " particles to template.\n";
	// std::cerr << "Source has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax[0] << ", " <<  maj_ax[1] << ", " << maj_ax[2] << " ).\n";
	// std::cerr << "     ( " << min_ax[0] << ", " <<  min_ax[1] << ", " << min_ax[2] << " ).\n";
	// std::cerr << "Template has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax_t[0] << ", " <<  maj_ax_t[1] << ", " << maj_ax_t[2] << " ).\n";
	// std::cerr << "     ( " << min_ax_t[0] << ", " <<  min_ax_t[1] << ", " << min_ax_t[2] << " ).\n";

	copy_ensemble( xi, rot_data, N );

	// Check if major axes are too much alike:

	double major_dot = dot(maj_ax_t, maj_ax);
	if( major_dot > 1 - 1e-4 ){
		// std::cerr << "Axes almost aligned...\n";

	}else if( major_dot < -1 + 1e-4 ){
		// std::cerr << "Axes almost anti-aligned...\n";

		// Idea: Just flip everything upside-down about the minor axis

		rotate_ensemble( rot_data, min_ax, N, PI );
		rotate_vector( maj_ax, min_ax, PI );
	}else{
		// std::cerr << "Axes not aligned...\n";

		ang1 = rotate_onto( maj_ax, maj_ax_t, rot_ax );
		rotate_ensemble( rot_data, rot_ax, N, ang1 );

		// Rotate axes as well:
		rotate_vector( maj_ax, rot_ax, ang1 );
		rotate_vector( min_ax, rot_ax, ang1 );
	}

	// std::cerr << "After rotation about major axes, this is the state:\n";
	// std::cerr << "Source has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax[0] << ", " <<  maj_ax[1] << ", " << maj_ax[2] << " ).\n";
	// std::cerr << "     ( " << min_ax[0] << ", " <<  min_ax[1] << ", " << min_ax[2] << " ).\n";
	// std::cerr << "Template has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax_t[0] << ", " <<  maj_ax_t[1] << ", " << maj_ax_t[2] << " ).\n";
	// std::cerr << "     ( " << min_ax_t[0] << ", " <<  min_ax_t[1] << ", " << min_ax_t[2] << " ).\n";

	double minor_dot = dot(min_ax_t, min_ax);
	if( minor_dot > 1 - 1e-4 ){
		// std::cerr << "Axes almost aligned...\n";

	}else if( minor_dot < -1 + 1e-4 ){
		// std::cerr << "Axes almost anti-aligned...\n";

		// Idea: Just flip everything upside-down about the major axis

		rotate_ensemble( rot_data, maj_ax, N, PI );
		rotate_vector( min_ax, maj_ax, PI );
	}else{
		// std::cerr << "Axes not aligned...\n";

		ang2 = rotate_onto( min_ax, min_ax_t, rot_ax );
		rotate_ensemble( rot_data, rot_ax, N, ang2 );

		// Rotate axes as well:
		rotate_vector( maj_ax, rot_ax, ang2 );
		rotate_vector( min_ax, rot_ax, ang2 );
	}

	// std::cerr << "After rotation about minor axes, this is the state:\n";
	// std::cerr << "Source has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax[0] << ", " <<  maj_ax[1] << ", " << maj_ax[2] << " ).\n";
	// std::cerr << "     ( " << min_ax[0] << ", " <<  min_ax[1] << ", " << min_ax[2] << " ).\n";
	// std::cerr << "Template has major and minor axes:\n";
	// std::cerr << "     ( " << maj_ax_t[0] << ", " <<  maj_ax_t[1] << ", " << maj_ax_t[2] << " ).\n";
	// std::cerr << "     ( " << min_ax_t[0] << ", " <<  min_ax_t[1] << ", " << min_ax_t[2] << " ).\n";

	// std::cerr << "\n\n";
	// std::cerr << "Dot of major axes: " << dot(maj_ax,maj_ax_t) << ".\n";
	// std::cerr << "Dot of minor axes: " << dot(min_ax,min_ax_t) << ".\n";

	// Now determine the rmsd:
	py_float rmsd = 0.0;
	for( int i = 0; i < N; ++i ){
		py_float dx, dy, dz;

		dx = rot_data[i][0] - xt[0][0];
		dy = rot_data[i][1] - xt[0][1];
		dz = rot_data[i][2] - xt[0][2];

		double r2 = dx*dx + dy*dy + dz*dz;
		double mr2 = r2;
		for( int j = 1; j < N; ++j ){
			dx = rot_data[i][0] - xt[j][0];
			dy = rot_data[i][1] - xt[j][1];
			dz = rot_data[i][2] - xt[j][2];

			r2 = dx*dx + dy*dy + dz*dz;
			if ( r2 < mr2 ){
				mr2 = r2;
			}
		}
		rmsd += mr2;
	}

	if ( !xrot ){
		delete [] rot_data_;
	}

	return sqrt(rmsd);
}

void test_rotate_template()
{

}