#include <InverseFourierTransform.h>
#include <OpenCLContext.h>
#include <boost/log/trivial.hpp>

using namespace neneta;
using namespace neneta::net;

InverseFourierTransform::InverseFourierTransform(const std::string& layerId, const conf::ConfigurationReader& confReader, const gpu::OpenCLProgram& program, const gpu::OpenCLContext& oclContext)
    : gpu::OpenCLExecutionPlan(layerId, confReader, program)
    , m_layerParameters(conf::FourierConfiguration(layerId, confReader).getParamSet(layerId))
    , m_oclKernelParameters(m_layerParameters.m_size, m_layerParameters.m_size, getKernelConfiguration().getLWS(1), getKernelConfiguration().getLWS(2))
    , m_clContext(oclContext)
    , m_io()
    , m_reTmp(m_clContext.getContext(), CL_MEM_READ_WRITE, m_layerParameters.m_size*m_layerParameters.m_size*sizeof(cmn::GPUFLOAT))
    , m_imTmp(m_clContext.getContext(), CL_MEM_READ_WRITE, m_layerParameters.m_size*m_layerParameters.m_size*sizeof(cmn::GPUFLOAT))
{
}

InverseFourierTransform::~InverseFourierTransform()
{
}

void InverseFourierTransform::init()
{
    for(unsigned int channel = 0; channel < m_layerParameters.m_channels; ++channel)
    {
        // plan("fftShift", {m_layerParameters.m_size, m_layerParameters.m_size, m_layerParameters.m_size/2, m_layerParameters.m_size/2}, *reKernelIt, *imKernelIt, (cl_int)m_layerParameters.m_size);
        planFwd("rowIndexReverse", m_oclKernelParameters, m_io.m_reChannels[channel], m_io.m_imChannels[channel], (cl_int)m_layerParameters.m_size, (cl_int)std::log2(m_layerParameters.m_size));
        planFwd("columnIndexReverse", m_oclKernelParameters, m_io.m_reChannels[channel], m_io.m_imChannels[channel], (cl_int)m_layerParameters.m_size, (cl_int)std::log2(m_layerParameters.m_size));
        for(int stage = 1; stage <= std::log2(m_layerParameters.m_size); ++stage)
        {
           planFwd("dit_2x2radix_2dfft_1st", m_oclKernelParameters, m_io.m_reChannels[channel], m_io.m_imChannels[channel], m_reTmp, m_imTmp, (cl_int)m_layerParameters.m_size, (cl_int)stage, (cl_int)1);
           planFwd("dit_2x2radix_2dfft_2nd", m_oclKernelParameters, m_io.m_reChannels[channel], m_io.m_imChannels[channel], m_reTmp, m_imTmp, (cl_int)m_layerParameters.m_size, (cl_int)stage);
        }
        planFwd("fftScale", m_oclKernelParameters, m_io.m_reChannels[channel], m_io.m_imChannels[channel], (cl_int)m_layerParameters.m_size);
    }
}

void InverseFourierTransform::getRe(std::vector<std::vector<cmn::GPUFLOAT>>& re)
{
    cmn::GPUFLOAT* data = new cmn::GPUFLOAT[m_layerParameters.m_size*m_layerParameters.m_size];
    for(unsigned int channel = 0; channel < m_layerParameters.m_channels; ++channel)
    {
        re.emplace_back();
        readFromBuffer(m_clContext.getCommandQueue(), m_io.m_reChannels[channel], m_layerParameters.m_size*m_layerParameters.m_size*sizeof(cmn::GPUFLOAT), data);
        re.back().assign(data, data + m_layerParameters.m_size*m_layerParameters.m_size);
    }
    delete [] data;
}

void InverseFourierTransform::getIm(std::vector<std::vector<cmn::GPUFLOAT>>& im)
{
    cmn::GPUFLOAT* data = new cmn::GPUFLOAT[m_layerParameters.m_size*m_layerParameters.m_size];
    for(unsigned int channel = 0; channel < m_layerParameters.m_channels; ++channel)
    {
        im.emplace_back();
        readFromBuffer(m_clContext.getCommandQueue(), m_io.m_imChannels[channel], m_layerParameters.m_size*m_layerParameters.m_size*sizeof(cmn::GPUFLOAT), data);
        im.back().assign(data, data + m_layerParameters.m_size*m_layerParameters.m_size);
    }
    delete [] data;
}

void neneta::net::InverseFourierTransform::setInput(gpu::BufferIO input)
{
    m_io = input;
    init();
}

gpu::BufferIO neneta::net::InverseFourierTransform::getOutput()
{
    return m_io;
}

void InverseFourierTransform::setBkpInput(gpu::BufferIO input)
{

}

gpu::BufferIO InverseFourierTransform::getBkpOutput()
{
    return m_io;
}


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

#ifndef DIRECT_HH
#define DIRECT_HH

#include <cmath>
#include <memory>
#include <type_traits>

#include <boost/timer/timer.hpp>

#include "dune/istl/operators.hh"
#include <dune/istl/solvers.hh>

#include "fem/istlinterface.hh"
#include "linalg/factorization.hh"

namespace Kaskade
{
  /// \internal
  namespace DirectSolver_Detail
  {
    class NoScalar{};
    class NoFieldType{};

    template <class T>
    typename T::Scalar hasScalar(typename T::Scalar*);
    template <class T>
    NoScalar hasScalar(...);

    template <class T>
    typename T::field_type hasFieldType(typename T::field_type*);
    template <class T>
    NoFieldType hasFieldType(...);

    template <class T>
    struct HasScalar
    {
      typedef decltype(hasScalar<T>(nullptr)) type;
      static constexpr bool value = !std::is_same<type,NoScalar>::value;
    };

    template <class T>
    struct HasFieldType
    {
      typedef decltype(hasFieldType<T>(nullptr)) type;
      static constexpr bool value = !std::is_same<type,NoFieldType>::value;
    };

    template <class T>
    struct ScalarType
    {
      typedef typename std::conditional< HasScalar<T>::value, typename HasScalar<T>::type,
                                           typename std::conditional< HasFieldType<T>::value, typename HasFieldType<T>::type, void>::type >::type type;

    };
  }
  /// \endinternal

  /**
   * \ingroup direct
   * \brief Dune::InverseOperator and Dune::Preconditioner interface for direct solvers.
   *
   * This keeps a factorization during the lifetime of the object, however, due to
   * shared data, efficient copying is possible.
   */
  template <class Domain_, class Range_>
  class DirectSolver: public Dune::InverseOperator<Domain_,Range_>,
		      public Dune::Preconditioner<Domain_,Range_>
  {
  public:
    typedef typename DirectSolver_Detail::ScalarType<Domain_>::type Scalar;
    typedef Domain_ Domain;
    typedef Range_ Range;

    /**
     * \brief Default constructor.
     *
     * A default constructed DirectSolver implements a (pretty useless) zero operator.
     */
    DirectSolver() {}

    /**
     * \brief Constructs a direct solver from an assembled linear operator.
     *
     * A copy of the linear operator is held internally, thus the linear
     * operator object is better not too large.
     */
    template <class AssembledGOP>
    explicit DirectSolver(AssembledGOP const& A,
                          DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL):
        solver(getFactorization(directType,properties,A.template get<MatrixAsTriplet<typename AssembledGOP::Scalar> >()))
    {
      assert(solver.get()); // make sure the factorization has been successful - otherwise it should have raised an exception
    }

    /**
     * \brief Constructs a direct solver from a triplet matrix.
     */
    template <class FieldType>
    explicit DirectSolver(MatrixAsTriplet<FieldType> const& A,
                          DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL):
        solver(getFactorization(directType,properties,A))
    {
      assert(solver.get()); // make sure the factorization has been successful - otherwise it should have raised an exception
    }

    /**
     * \brief Constructs a direct solver from a BCRS matrix.
     * A copy of the linear operator is held internally, thus the linear
     * operator object is better not too large.
     */
    template <class FieldType>
    explicit DirectSolver(Dune::BCRSMatrix<Dune::FieldMatrix<FieldType,1,1>> const& A,
                          DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL):
        solver(getFactorization<FieldType,int>(directType,properties,A))
    {
      assert(solver.get()); // make sure the factorization has been successful - otherwise it should have raised an exception
    }

    /**
     * \brief Constructs a direct solver from a NumaBCRS matrix.
     * A copy of the linear operator is held internally, thus the linear
     * operator object is better not too large.
     */
    template <class FieldType, int n, int m, class Index>
    explicit DirectSolver(NumaBCRSMatrix<Dune::FieldMatrix<FieldType,n,m>,Index> const& A,
                          DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL):
        solver(getFactorization<FieldType,n,m,Index,int>(directType,properties,A))
    {
      assert(solver.get()); // make sure the factorization has been successful - otherwise it should have raised an exception
    }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * The provided right hand side \arg b is not modified (except possibly via aliasing
     * \arg x).
     */
    virtual void apply(Domain& x, Range& b, Dune::InverseOperatorResult& res) { apply(x,b,1e-10,res); }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * As an extension to the Dune::InverseOperator interface, we provide const versions of
     * the apply methods.
     */
    void apply(Domain& x, Range& b, Dune::InverseOperatorResult& res) const { apply(x,b,1e-10,res); }

    /**
     * Solves the system for the given right hand side \arg b, which is
     * guaranteed not to be overwritten (except possibly via aliasing
     * \arg x).
     */
    virtual void apply(Domain& x, Range& b, double reduction, Dune::InverseOperatorResult& res)
    {
      const_cast<DirectSolver<Domain,Range> const*>(this)->apply(x,b,res);
    }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * As an addition to the Dune::InverseOperator interface, we provide const versions of
     * the apply methods, since the factorization is not modified during the solution phase.
     */
    void apply(Domain& x, Range const& b, double reduction, Dune::InverseOperatorResult& res) const
    {
      boost::timer::cpu_timer timer;

      std::vector<Scalar> rhs(b.dim());
      vectorToSequence(b,begin(rhs));
      apply(rhs,reduction,res);
      vectorFromSequence(x,begin(rhs));
    }

    /**
     * \brief The input vector v contains the rhs. It will be overwritten by the solution.
     *
     * Factorization is not touched by solving. As an addition to the
     * Dune::InverseOperator interface, we provide const versions of
     * the apply methods.
     */
    void apply(std::vector<Scalar> &v, double reduction, Dune::InverseOperatorResult& res) const
    {
      boost::timer::cpu_timer timer;

      // Check for solver availability. Otherwise return zero solution.
      if (solver)
      {
        #ifndef NDEBUG
        size_t rhsNan = 0, solNan = 0;
        size_t rhsInf = 0;
        for (auto const& x: v)
        {
          if (std::isnan(x)) ++rhsNan;
          if (std::isinf(x)) ++rhsInf;
        }
        if (rhsNan>0) std::cerr << __FILE__ << ':' << __LINE__ << ": " << rhsNan << " rhs entries are nan\n";
        if (rhsInf>0) std::cerr << __FILE__ << ':' << __LINE__ << ": " << rhsInf << " rhs entries are inf\n";
        #endif
        solver->solve(v);
        #ifndef NDEBUG
        for (int i=0; i<v.size(); ++i)
          if (std::isnan(v[i])) ++solNan;
          if (solNan>0) std::cerr << __FILE__ << ':' << __LINE__ << ": " << solNan << " solution entries are nan\n";
          #endif
      } else
        std::fill(v.begin(),v.end(),0.0);

      // Write solver statistics. Currently dummy.
      res.clear();
      res.iterations = 1;
      res.reduction = 1e-10; // dummy!
      res.converged = true;
      res.conv_rate = 1e-10;
      res.elapsed = (double)(timer.elapsed().user)/1e9;
    }

    virtual void  apply (std::vector<Scalar> &v)
    {
      Dune::InverseOperatorResult dummy_res;
      apply(v,1e-10,dummy_res);
    }

    virtual void pre (Domain &x, Range &b) {}

    virtual void apply (Domain &v, const Range &d)
    {
      Dune::InverseOperatorResult dummy_res;
      apply(v,d, 1e-10,dummy_res);
    }

    virtual void post (Domain &x) {}

  private:
    std::shared_ptr<Factorization<Scalar>> solver;
  };

  //---------------------------------------------------------------------

  // partial specialization for tiny fixed-size matrices
  template <class S, int n>
  class DirectSolver<Dune::FieldVector<S,n>,Dune::FieldVector<S,n>>: public Dune::InverseOperator<Dune::FieldVector<S,n>,Dune::FieldVector<S,n>>,
                                                                     public Dune::Preconditioner<Dune::FieldVector<S,n>,Dune::FieldVector<S,n>>
  {
  public:
    typedef S Scalar;
    typedef Dune::FieldVector<S,n> Domain;
    typedef Dune::FieldVector<S,n> Range;

    /**
     * \brief Default constructor.
     *
     * A default constructed DirectSolver implements a (pretty useless) zero operator.
     */
    DirectSolver() {}

    /**
     * \brief Constructs a direct solver from a FieldMatrix matrix.
     */
    template <class FieldType>
    explicit DirectSolver(Dune::FieldMatrix<FieldType,n,n> const& A):
        inverse(A)
    {
      inverse.invert();
    }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * The provided right hand side \arg b is not modified (except possibly via aliasing
     * \arg x).
     */
    virtual void apply(Domain& x, Range& b, Dune::InverseOperatorResult& res) { apply(x,b,1e-10,res); }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * As an extension to the Dune::InverseOperator interface, we provide const versions of
     * the apply methods.
     */
    void apply(Domain& x, Range& b, Dune::InverseOperatorResult& res) const { apply(x,b,1e-10,res); }

    /**
     * Solves the system for the given right hand side \arg b, which is
     * guaranteed not to be overwritten (except possibly via aliasing
     * \arg x).
     */
    virtual void apply(Domain& x, Range& b, double reduction, Dune::InverseOperatorResult& res)
    {
      const_cast<DirectSolver<Domain,Range> const*>(this)->apply(x,b,res);
    }

    /**
     * \brief Solves the system for the given right hand side \arg b.
     *
     * As an addition to the Dune::InverseOperator interface, we provide const versions of
     * the apply methods, since the factorization is not modified during the solution phase.
     */
    void apply(Domain& x, Range const& b, double reduction, Dune::InverseOperatorResult& res) const
    {
      x = inverse*b;

      // Write solver statistics. Currently dummy.
      res.clear();
      res.iterations = 1;
      res.reduction = 1e-10; // dummy!
      res.converged = true;
      res.conv_rate = 1e-10;
      res.elapsed = 0;
    }

    virtual void        pre (Domain &x, Range &b) {}

    virtual void        apply (Domain &v, const Range &d)
    {
      v = inverse*d;
    }

    virtual void        post (Domain &x) {}

  private:
    Dune::FieldMatrix<Scalar,n,n> inverse;
  };

  //---------------------------------------------------------------------

  /**
   * \ingroup linalgsolution
   * \brief Dune::LinearOperator interface for inverse operators.
   */
  template <class InverseOperator>
  class InverseLinearOperator: public Dune::LinearOperator<typename InverseOperator::Range, typename InverseOperator::Domain>
  {
  public:
    typedef typename InverseOperator::Domain Domain;
    typedef typename InverseOperator::Range Range;
    typedef typename InverseOperator::Scalar Scalar;

    InverseLinearOperator() = default;

    InverseLinearOperator(InverseOperator const& op_) : op(op_)
    {}

    virtual void apply(Domain const& x, Range& y) const
    {
      Dune::InverseOperatorResult result;
      Domain rhs(x);
      op.apply(y,rhs,result);
    }

    virtual void applyscaleadd(Scalar alpha, Domain const& x, Range& y) const
    {
      Range ynew(y);
      apply(x,ynew);
      y.axpy(alpha,ynew);
    }

  private:
    InverseOperator op;
  };

  //---------------------------------------------------------------------

  /**
   * \ingroup direct
   * \brief convenience function for constructing a DirectInverseOperator
   */
  template <class GOP, int firstRow, int lastRow, int firstCol, int lastCol>
  InverseLinearOperator<DirectSolver<typename AssembledGalerkinOperator<GOP,firstRow,lastRow,firstCol,lastCol>::Domain,
                                     typename AssembledGalerkinOperator<GOP,firstRow,lastRow,firstCol,lastCol>::Range> >
  directInverseOperator(AssembledGalerkinOperator<GOP,firstRow,lastRow,firstCol,lastCol> const& A,
                        DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL)
  {
    typedef typename AssembledGalerkinOperator<GOP,firstRow,lastRow,firstCol,lastCol>::Domain Domain;
    typedef typename AssembledGalerkinOperator<GOP,firstRow,lastRow,firstCol,lastCol>::Range Range;
    return InverseLinearOperator<DirectSolver<Domain,Range> >(DirectSolver<Domain,Range>(A,directType,properties));
  }

  template <class Matrix, class Domain, class Range>
  InverseLinearOperator<DirectSolver<Domain,Range> >
  directInverseOperator(MatrixRepresentedOperator<Matrix,Domain,Range> const& A,
                        DirectType directType=DirectType::UMFPACK, MatrixProperties properties=MatrixProperties::GENERAL)
  {
    return InverseLinearOperator<DirectSolver<Domain,Range> >(DirectSolver<Domain,Range>(A,directType,properties));
  }

} // namespace Kaskade

//---------------------------------------------------------------------

#endif


#ifndef CMAESAG_HPP
#define CMAESAG_HPP

#include <vector>
#include <string>
#include <type_traits>
#include <boost/serialization/list.hpp>
#include <boost/serialization/set.hpp>
#include <boost/serialization/vector.hpp>
#include <boost/functional/hash.hpp>

#include "arch/ARLAgent.hpp"
#include "bib/Seed.hpp"
#include "bib/Utils.hpp"
#include <bib/MetropolisHasting.hpp>
#include <bib/XMLEngine.hpp>
#include <bib/Combinaison.hpp>
#include "nn/MLP.hpp"
#include "nn/DevMLP.hpp"
#include "nn/DODevMLP.hpp"
#include "cmaes_interface.h"

template<typename NN = MLP>
class CMAESAg : public arch::ARLAgent<arch::AgentProgOptions> {
 public:
  CMAESAg(unsigned int _nb_motors, unsigned int _nb_sensors)
    : arch::ARLAgent<arch::AgentProgOptions>(_nb_motors, _nb_sensors), nb_sensors(_nb_sensors) {

  }

  virtual ~CMAESAg() {
    cmaes_exit(evo);
    delete evo;
    delete ann;
    delete hidden_unit_a;
  }

  const std::vector<double>& _run(double, const std::vector<double>& sensors,
                                  bool, bool, bool) override {

    vector<double>* next_action = ann->computeOut(sensors);

//  CMA-ES already implement exploration in parameter space
    last_action.reset(next_action);

    return *next_action;
  }

  void _unique_invoke(boost::property_tree::ptree* pt, boost::program_options::variables_map* vm) override {
    hidden_unit_a               = bib::to_array<uint>(pt->get<std::string>("agent.hidden_unit_a"));
    actor_hidden_layer_type     = pt->get<uint>("agent.actor_hidden_layer_type");
    actor_output_layer_type     = pt->get<uint>("agent.actor_output_layer_type");
    batch_norm                  = pt->get<uint>("agent.batch_norm");
    population                  = pt->get<uint>("agent.population");
    initial_deviation           = pt->get<double>("agent.initial_deviation");

    check_feasible = true;
    ignore_null_lr = true;
    racing = false;
    error_count = 0;
    xbestever_score = std::numeric_limits<double>::max();

    try {
      check_feasible = pt->get<bool>("agent.check_feasible");
    } catch(boost::exception const& ) {
    }

    try {
      ignore_null_lr = pt->get<bool>("agent.ignore_null_lr");
    } catch(boost::exception const& ) {
    }

    try {
      racing = pt->get<bool>("agent.racing");
    } catch(boost::exception const& ) {
    }

    episode = 0;

    ann = new NN(nb_sensors, *hidden_unit_a, nb_motors, 0.1, 1, actor_hidden_layer_type, actor_output_layer_type,
                 batch_norm);
    if(std::is_same<NN, DevMLP>::value)
      ann->exploit(pt, static_cast<CMAESAg *>(old_ag)->ann);
    else if(std::is_same<NN, DODevMLP>::value)
      ann->exploit(pt, nullptr);

//     const uint dimension = (nb_sensors+1)*hidden_unit_a->at(0) + (hidden_unit_a->at(0)+1)*nb_motors;
    const uint dimension = ann->number_of_parameters(ignore_null_lr);
    double* startx  = new double[dimension];
    double* deviation  = new double[dimension];
    for(uint j=0; j< dimension; j++) {
      deviation[j] = initial_deviation;
    }
    ann->copyWeightsTo(startx, ignore_null_lr);

    evo = new cmaes_t;
    arFunvals = cmaes_init(evo, dimension, startx, deviation, 0, population, NULL/*"config.cmaes.ini"*/);
    delete[] startx;
    delete[] deviation;
//     evo->sp.stopTolFun = 1e-150;
//     evo->sp.stopTolFunHist = 1e-150;
//     evo->sp.stopTolUpXFactor = 1e50;

    printf("%s\n", cmaes_SayHello(evo));
    new_population();
    LOG_DEBUG(cmaes_Get(evo, "lambda") << " " << dimension << " " << population << " " << cmaes_Get(evo, "N"));
    if (population < 2)
      LOG_DEBUG("population too small, changed to : " << (4+(int)(3*log((double)dimension))));

    if(vm->count("continue") > 0) {
      uint continue_save_each          = DEFAULT_AGENT_SAVE_EACH_CONTINUE;
      try {
        continue_save_each            = pt->get<uint>("simulation.continue_save_each");
      } catch(boost::exception const& ) {
      }

      if(continue_save_each % (int) cmaes_Get(evo, "lambda") != 0) {
        LOG_ERROR("continue_save_each must be a multiple of the population size !");
        exit(1);
      }
    }

    if(std::is_same<NN, DODevMLP>::value){
      try {
        if(pt->get<double>("devnn.ewc") >= 0.){
          LOG_ERROR("CMA-ES doesn't rely on gradient (difficult to compute Fisher Matrix)");
          exit(1);
        }
      } catch(boost::exception const& ) {
      }
    }
  }

  bool is_feasible(const double* parameters) {
    for(uint i=0; i < (uint) cmaes_Get(evo, "dim"); i++)
      if(fabs(parameters[i]) >= 500.f) {
        return false;
      }

    return true;
  }

  void new_population() {
    const char * terminate =  cmaes_TestForTermination(evo);
    if(terminate) {
      LOG_INFO("mismatch "<< terminate);
      error_count++;

      if(error_count > 20 && racing) {
        LOG_FILE(DEFAULT_END_FILE, "-1");
        exit(0);
      }
    }
    //ASSERT(!cmaes_TestForTermination(evo), "mismatch "<< cmaes_TestForTermination(evo));

    current_individual = 0;
    pop = cmaes_SamplePopulation(evo);

    if(check_feasible) {
      //check that the population is feasible
      bool allfeasible = true;
      for (int i = 0; i < cmaes_Get(evo, "popsize"); ++i)
        while (!is_feasible(pop[i])) {
          cmaes_ReSampleSingle(evo, i);
          allfeasible = false;
        }

      if(!allfeasible)
        LOG_INFO("non feasible solution produced");
    }
  }

  void start_instance(bool learning) override {
    last_action = nullptr;
    scores.clear();

    if(std::is_same<NN, DODevMLP>::value && learning) {
      auto dodevmlp = static_cast<DODevMLP *>(ann);
      bool reset_operator, changed;
      std::tie(reset_operator, changed) = dodevmlp->inform(episode, last_sum_weighted_reward);
      if(reset_operator && changed) {
        LOG_INFO("reset learning catched");

        if(!dodevmlp->ewc_enabled()){
          const double* parameters = nullptr;
          parameters = getBestSolution();
          loadPolicyParameters(parameters);
        }

        cmaes_exit(evo);
        delete evo;

        const uint dimension = ann->number_of_parameters(ignore_null_lr);
        double* startx  = new double[dimension];
        double* deviation  = new double[dimension];
        for(uint j=0; j< dimension; j++) {
          deviation[j] = initial_deviation;
        }
        ann->copyWeightsTo(startx, ignore_null_lr);

        xbestever_score = std::numeric_limits<double>::max();
        cmaes_UpdateDistribution_done_once = false;
        evo = new cmaes_t;
        arFunvals = cmaes_init(evo, dimension, startx, deviation, 0, population, NULL/*"config.cmaes.ini"*/);
        delete[] startx;
        delete[] deviation;
        new_population();
      }
    }

    if(!justLoaded) {
      //put individual into NN
      const double* parameters = nullptr;
      if(learning || !cmaes_UpdateDistribution_done_once)
        parameters = pop[current_individual];
      else
        parameters = getBestSolution();

      loadPolicyParameters(parameters);
    }

    if(learning)
      episode++;
    //LOG_FILE("policy_exploration", ann->hash());
  }

  void restoreBest() override {
    const double* parameters = getBestSolution();
    loadPolicyParameters(parameters);
  }

  void end_episode(bool) override {
    scores.push_back(-sum_weighted_reward + ann->ewc_cost());
  }

  void end_instance(bool learning) override {
    if(learning) {
      arFunvals[current_individual] = std::accumulate(scores.begin(), scores.end(), 0.f) / scores.size();

      //TODO with instance
      ann->update_best_param_previous_task(sum_weighted_reward);

      current_individual++;
      if(current_individual >= cmaes_Get(evo, "lambda")) {
        cmaes_UpdateDistribution(evo, arFunvals);
        cmaes_UpdateDistribution_done_once=true;
        new_population();
      }
    }

    justLoaded = false;
  }

  void save(const std::string& path, bool save_best, bool) override {
    if(!save_best || !cmaes_UpdateDistribution_done_once) {
      ann->save(path+".actor");
    } else if(save_best && -sum_weighted_reward < xbestever_score ) {
      xbestever_score = -sum_weighted_reward;
      ann->save(path+".actor");
    }
//     else {
//       //TODO : check this part, apparently it modify the learning
//       LOG_WARNING("be careful it might be a problem here");
//       NN* to_be_restaured = new NN(*ann, false);
//       const double* parameters = getBestSolution();
//       loadPolicyParameters(parameters);
//       ann->save(path+".actor");
//       delete ann;
//       ann = to_be_restaured;
//     }
  }

  void load(const std::string& path) override {
    justLoaded = true;
    ann->load(path+".actor");
  }

  void save_run() override {
    if(current_individual == 1) {
      const double* xbptr_ = getBestSolution();
      std::vector<double> xbestever_((int)cmaes_Get(evo, "N"));
      xbestever_.assign(xbptr_, xbptr_ + (int)cmaes_Get(evo, "N"));
      double bs = getBestScore();

      struct algo_state st = {scores, justLoaded, cmaes_UpdateDistribution_done_once,
               current_individual, episode, error_count, bs, xbestever_
      };
      bib::XMLEngine::save(st, "algo_state", "continue.algo_state.data");
      cmaes_WriteToFile(evo, "resume", "continue.cmaes.data");
    }
  }

  void load_previous_run() override {
    auto algo_state_ = bib::XMLEngine::load<struct algo_state>("algo_state", "continue.algo_state.data");
    scores = algo_state_->scores;
    justLoaded = algo_state_->justLoaded;
    cmaes_UpdateDistribution_done_once = algo_state_->cmaes_UpdateDistribution_done_once;
    current_individual = algo_state_->current_individual;
    episode = algo_state_->episode;
    error_count = algo_state_->error_count;
    xbestever_score = algo_state_->xbestever_score;
    xbestever_ptr = algo_state_->xbestever_ptr;
    delete algo_state_;
    char file_[] = "continue.cmaes.data";
    cmaes_resume_distribution(evo, file_);
    new_population();
  }

  MLP* getNN() {
    return ann;
  }

 protected:
  void _display(std::ostream& out) const override {
    out << std::setw(8) << std::fixed << std::setprecision(5) << sum_weighted_reward << " " << ann->ewc_cost();
  }

  void _dump(std::ostream& out) const override {
    out << std::setw(8) << std::fixed << std::setprecision(5) << sum_weighted_reward;
  }

 private:
  void loadPolicyParameters(const double* parameters) {
    ann->copyWeightsFrom(parameters, ignore_null_lr);
  }

  const double* getBestSolution() {
    //TODO:
//     LOG_DEBUG(cmaes_Get(evo, "fbestever")  << " " << xbestever_score);
    if(cmaes_Get(evo, "fbestever") < xbestever_score)
      return cmaes_GetPtr(evo, "xbestever");
    else
      return xbestever_ptr.data();
  }

  double getBestScore() {
    return std::min(cmaes_Get(evo, "fbestever"), xbestever_score);
  }

 private:
  //initilized by constructor
  uint nb_sensors;

  //initialized by invoke
  std::vector<uint>* hidden_unit_a;
  uint population, actor_hidden_layer_type, actor_output_layer_type, batch_norm;
  double initial_deviation;
  bool check_feasible;
  bool ignore_null_lr;
  bool racing;
  MLP* ann;
  cmaes_t* evo;
  double *arFunvals;

  //internal mecanisms
  std::shared_ptr<std::vector<double>> last_action;
  std::list<double> scores;
  double *const *pop;
  bool justLoaded = false;
  bool cmaes_UpdateDistribution_done_once = false;
  uint current_individual;
  uint episode;
  uint error_count;
  double xbestever_score;
  std::vector<double> xbestever_ptr;

  struct algo_state {
    std::list<double> scores;
    bool justLoaded;
    bool cmaes_UpdateDistribution_done_once;
    uint current_individual;
    uint episode;
    uint error_count;
    double xbestever_score;
    std::vector<double> xbestever_ptr;

    friend class boost::serialization::access;
    template <typename Archive>
    void serialize(Archive& ar, const unsigned int) {
      ar& BOOST_SERIALIZATION_NVP(scores);
      ar& BOOST_SERIALIZATION_NVP(justLoaded);
      ar& BOOST_SERIALIZATION_NVP(cmaes_UpdateDistribution_done_once);
      ar& BOOST_SERIALIZATION_NVP(current_individual);
      ar& BOOST_SERIALIZATION_NVP(episode);
      ar& BOOST_SERIALIZATION_NVP(error_count);
      ar& BOOST_SERIALIZATION_NVP(xbestever_score);
      ar& BOOST_SERIALIZATION_NVP(xbestever_ptr);
    }
  };
};

#endif


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

// Currently needs -DMTL_DEEP_COPY_CONSTRUCTOR !!!

#include <iostream>
#include <boost/numeric/mtl/mtl.hpp>

template <typename T>
struct as {};

template <>
struct as<int>
{
	typedef mtl::ashape::scal  type;
};

template <typename V>
struct as<mtl::compressed2D<V> >
{
	typedef mtl::ashape::mat<typename as<V>::type>   type;
};

template <typename V>
struct as<mtl::dense2D<V> >
{
	typedef mtl::ashape::mat<typename as<V>::type>   type;
};

int main(int, char**)
{
    using namespace std;

#if 0
    cout << typeid(as<int>::type).name() << endl;
    cout << typeid(as<mtl::compressed2D<int> >::type).name() << endl;
    cout << typeid(as<mtl::compressed2D<mtl::dense2D<int> > >::type).name() << endl;

    cout << typeid(mtl::ashape::ashape<int>::type).name() << endl;
    cout << typeid(mtl::ashape::ashape<mtl::compressed2D<int> >::type).name() << endl;
    cout << typeid(mtl::ashape::ashape<typename mtl::compressed2D<mtl::dense2D<int> >::value_type >::type).name() << endl;
#endif

    // Define a 6x5 sparse matrix in a 3x3 block-sparse
    typedef mtl::dense2D<double>    m_t;
    typedef mtl::compressed2D<m_t>  matrix_t;
    matrix_t                        A(3, 3);
    {
	mtl::mat::inserter<matrix_t> ins(A);

	// First block
	m_t  b1(1, 1);
	b1(0, 0)= 1.0;
	ins(0, 2) << b1;

	// Second block
	m_t  b2(2, 3);
	b2=       0.0;
	b2[0][1]= 2.0;
	b2[1][2]= 3.0;
	ins(1, 0) << b2;

	m_t b3(3, 1);
	b3= 0.0;
	b3[1][0]= 4.0;
	ins(2, 1) << b3;
    }
    // cout << "A is " << A << endl; // doesn't works and before it was completely unreadable anyway

    /* Should be something like this:

       [                   [ 1]] // b1
       [[  0   2   0]          ] // b2
       [[  0   0   3]          ]
       [             [  0]     ] // b3
       [             [  4]     ]
       [             [  0]     ]

    */

    // Access blocks (they are read-only) for sparse matrices
    cout << "The block A(1, 0) is \n" << A(1, 0) << endl;
    cout << "The block A[1][0] is \n" << A[1][0] << endl;

    // Access elements in blocks
    cout << "In block A(1, 0), the element (0, 1) is " << A(1, 0)(0, 1) << endl;
    cout << "In block A(1, 0), the element [0][1] is " << A(1, 0)[0][1] << endl;
    cout << "In block A[1][0], the element [0][1] is " << A[1][0][0][1] << endl << endl;

    typedef mtl::dense_vector<double> v_t;
    typedef mtl::dense_vector<v_t>    vector_t;
    vector_t                          x(3), y(3);

    // x= [[0, 5, 3], [1], [8]]^T
    x[0]= v_t(3, 0.0); x[0][1]= 5.0; x[0][2]= 3.0; // first block of x = [0, 5, 3]^T
    x[1]= v_t(1, 1.0);                             // second block of x
    x[2]= v_t(1, 8.0);                             // third block

    cout << "x is " << x << endl;

    // For y we would only need the vector sizes [[?], [?, ?], [?, ?, ?]]^T
    // To avoid valgrind complains we set to 0
    y[0]= v_t(1, 0.0); y[1]= v_t(2, 0.0); y[2]= v_t(3, 0.0);

    cout << "y is " << y << endl;

    // Block-sparse matrix * blocked vector !!!
    y= A*x;

    cout << "y after multiplication is " << y << endl
	 << "Should be [[8], [10, 9], [0, 4, 0]]^T." << endl;

    return 0;
}


#include <iostream>
#include <cmath>
#include <typeinfo>
#include <type_traits>
#include <boost/numeric/mtl/mtl.hpp>

// namespace to avoid ambiguities with mtl
namespace tst {

    template <typename T>
    struct is_const
    {
	static const bool value= false;
    };

    template <typename T>
    struct is_const<const T>
    {
	static const bool value= true;
    };

    template <bool Condition, typename ThenType, typename ElseType>
    struct conditional
    {
	using type= ThenType;
    };

    template <typename ThenType, typename ElseType>
    struct conditional<false, ThenType, ElseType>
    {
	using type= ElseType;
    };

    template <bool Condition, typename ThenType, typename ElseType>
    using conditional_t= typename conditional<Condition, ThenType, ElseType>::type;

    template <typename T>
    struct is_matrix
      : std::false_type
    {};

    template <typename Value, typename Para>
    struct is_matrix<mtl::dense2D<Value, Para> >
      : std::true_type
    {};

#if 1
    template <typename T>
    struct is_matrix<const T>
      : is_matrix<T> {};
#endif

#if 0
    template <typename Matrix>
    class transposed_view
    {
      public:
	using value_type= typename Matrix::value_type;
	using size_type=  typename Matrix::size_type;

	// typedef typename mtl::Collection<Matrix>::value_type  value_type;
	// typedef typename mtl::Collection<Matrix>::size_type   size_type;

	explicit transposed_view(Matrix& A) : ref(A) {}

	value_type& operator()(size_type r, size_type c) { return ref(c, r); }
	const value_type& operator()(size_type r, size_type c) const { return ref(c, r); }

      private:
	Matrix& ref;
    };

    template <typename Matrix>
    class const_transposed_view
    {
      public:
	using value_type= typename Matrix::value_type;
	using size_type=  typename Matrix::size_type;

	explicit const_transposed_view(const Matrix& A) : ref(A) {}

	const value_type& operator()(size_type r, size_type c) const { return ref(c, r); }

      private:
	const Matrix& ref;
    };

    template <typename Matrix>
    inline const_transposed_view<Matrix> trans(const Matrix& A)
    {
	return const_transposed_view<Matrix>(A);
    }

#else
    template <typename Matrix>
    class transposed_view
    {
	static_assert(is_matrix<Matrix>::value, "template argument is not a Matrix");
      public:
	using value_type= typename Matrix::value_type;
	using size_type=  typename Matrix::size_type;

      private:
	using vref_type= conditional_t<is_const<Matrix>::value,
				       const value_type&,
				       value_type&>;
      public:
	explicit transposed_view(Matrix& A) : ref(A) {}

	vref_type operator()(size_type r, size_type c) { return ref(c, r); }
	const value_type& operator()(size_type r, size_type c) const { return ref(c, r); }

      private:
	Matrix& ref;
    };
#endif

    template <typename Matrix>
    transposed_view<Matrix> inline trans(Matrix& A)
    {
	return transposed_view<Matrix>(A);
    }

    template <typename Matrix>
    struct is_matrix<transposed_view<Matrix> >
      : is_matrix<Matrix>
    {};

}

template <int M>
void f()
{
    typedef typename tst::conditional<M < 100, double, float>::type & value_type;
    typedef const typename tst::conditional<M < 100, const double, const float>::type const_value_type;
    std::cout << "typeid = " << typeid(value_type).name() << '\n';
}

int main (int argc, char* argv[])
{
    // const double eps= 0.00000001;
    const int n = 10;
    typedef tst::conditional<n < 100, double, float>::type& value_type;
    typedef const tst::conditional<n < 100, double, float>::type const_value_type;
    std::cout << "typeid = " << typeid(value_type).name() << '\n';
    f<17>();

    mtl::dense2D<float> A= {{2, 3, 4},
			    {5, 6, 7},
			    {8, 9, 10}};
    std::cout << "trans(A) is\n" << trans(A);

    trans(A)[2][0]= 4.5;
    std::cout << "trans(A) is\n" << trans(A);

    const mtl::dense2D<float> B(A);
    std::cout << "trans(B) is\n" << trans(B);

    tst::transposed_view<mtl::dense2D<float> >  At(A);
    At(2, 0)= 4.5;

    std::cout << "tst::trans(A)(2, 0) = " << tst::trans(A)(2, 0) << '\n';
    tst::trans(A)(2, 0)= 4.6;
    std::cout << "trans(A) after modification is\n" << trans(A);

    std::cout << "typeid of trans(A) = " << typeid(tst::trans(A)).name() << '\n';
    std::cout << "typeid of trans(B) = " << typeid(tst::trans(B)).name() << '\n';
    // int ta= trans(A);
    // int tb= trans(B);
    // int tar= trans(A).ref;
    // int tbr= trans(B).ref;

    std::cout << "tst::trans(B)(2, 0) = " << tst::trans(B)(2, 0) << '\n';
    // tst::trans(B)(2, 0)= 4.6;

    const tst::transposed_view<const mtl::dense2D<float> >  Bt(B);
    std::cout << "Bt(2, 0) = " << Bt(2, 0) << '\n';

    return 0 ;
}


#ifndef __Rocket_Flight_DM_HH__
#define __Rocket_Flight_DM_HH__
/********************************* TRICK HEADER *******************************
PURPOSE:
      (Describe the Rocket Flgiht Dynamics Module Variables and Algorithm)
LIBRARY DEPENDENCY:
      ((../src/Rocket_Flight_DM.cpp)
       (../../math/src/nrutil.cpp))
PROGRAMMERS:
      (((Lai Jun Xu) () () () ))
*******************************************************************************/
#include <armadillo>
#include "Module.hh"
#include "Time_management.hh"
#include "aux.hh"
#include "global_constants.hh"
#include "icf_trx_ctrl.h"
#include "simgen_remote.h"

class Rocket_Flight_DM : public Dynamics {
  TRICK_INTERFACE(Rocket_Flight_DM);

 public:
  Rocket_Flight_DM(Data_exchang &input);
  Rocket_Flight_DM(const Rocket_Flight_DM &other);
  Rocket_Flight_DM &operator=(const Rocket_Flight_DM &other);
  struct icf_ctrlblk_t *dm_icf_info_hook;
  int enqueue_to_simgen_buffer(struct icf_ctrlblk_t *C, double ext_porlation);
  int stand_still_motion_data(struct icf_ctrlblk_t *C, double ext_porlation);

  virtual void init();
  virtual void algorithm(double int_step);
  void load_location(double lonx_in, double latx_in, double alt_in);
  void load_geodetic_velocity(double alpha0x, double beta0x, double dvbe);
  void load_angle(double yaw, double roll, double pitch);
  void load_angular_velocity(double ppx_in, double qqx_in, double rrx_in);
  void update_diagnostic_attributes(double int_step);
  void Interpolation_Extrapolation(double T, double int_step,
                                   double ext_porlation);
  void set_reference_point(double rp);

  arma::vec3 get_VBED();
  void set_liftoff(int in);

  struct TX_data {
    double SBEE[3];
    double VBEE[3];
    double ABEE[3];
    double JBEE[3];
    double psibd;
    double thtbd;
    double phibd;
    double WBEB[3];
  } TX_data_forward;

  void collect_forces_and_propagate();
  void set_DOF(int ndof);
  void set_aero_flag(unsigned int in);

 private:
  double get_ppx();
  double get_qqx();
  double get_rrx();
  double get_psibdx();
  double get_thtbdx();
  double get_phibdx();
  double get_dvbe();
  double get_dbi();
  double get_dvbi();
  double get_thtvdx();
  double get_psivdx();
  double get_thtbdx_in(double &cthtbd);

  void propagate_position_speed_acceleration(double int_step);
  void propagate_aeroloss(double int_step);
  void propagate_gravityloss(double int_step);
  void propagate_control_loss(double int_step);
  void vibration(double int_step);
  void propagate_TBI(double int_step, arma::vec3 WBIB_in);
  void propagate_TBI_Q(double int_step, arma::vec3 WBIB_in);
  void propagate_WBIB(double int_step, arma::vec3 FMB_in, arma::mat33 IBBB);
  void orbital(arma::vec3 SBII_in, arma::vec3 VBII_in, double dbi_in);
  void build_WEII();
  void aux_calulate(arma::mat33 TEI, arma::mat33 TBI_in);
  void RK4F(std::vector<arma::vec> Var_in, std::vector<arma::vec> &Var_out);
  void Send();

  double calculate_alphaix(arma::vec3 VBIB);
  double calculate_betaix(arma::vec3 VBIB);
  double calculate_alppx(arma::vec3 VBAB_in, double dvba);
  double calculate_phipx(arma::vec3 VBAB_in);
  double calculate_alphax(arma::vec3 VBAB_in);
  double calculate_betax(arma::vec3 VBAB, double dvba);

  arma::vec build_VBEB(double _alpha0x, double _beta0x, double dvbe);
  arma::mat calculate_TBD(double lonx_in, double latx_in, double alt_in);
  arma::vec3 calculate_WBII(arma::mat33 TBI_in);
  arma::vec3 calculate_fspb(arma::vec3 FAPB_in, double vmass);
  arma::vec3 calculate_WBEB(arma::mat33 TBI_in);
  arma::vec3 euler_angle(arma::mat33 TBD_in);

  void gamma_beta();
  void Gravity_Q();
  void AeroDynamics_Q();
  void calculate_I1();
  void funcv(int n, double *x, double *ff);
  void broydn(double x[], int n, int *check);
  void rsolv(double **a, int n, double d[], double b[]);
  void fdjac(int n, double x[], double fvec_in[], double **df);
  double f_min(double x[]);
  void lnsrch(int n, double xold[], double fold, double g[], double p[],
              double x[], double *f_in, double stpmax, int *check);
  void qrdcmp(double **a, int n, double *c, double *d, int *sing);
  void qrupdt(double **r, double **qt, int n, double u[], double v[]);
  void rotate(double **r, double **qt, int n, int i, double a, double b);

  MATRIX(TBI, 3, 3); /* *o (--)    Transformation Matrix of body coord wrt
                    inertia coord */

  MATRIX(TBID, 3, 3); /* *o (--)    Transformation Matrix of body coord wrt
                     inertia coord derivative */

  VECTOR(TBI_Q, 4); /* *o (--)    Transformation Matrix of body coord wrt
                      inertia coord (Quaternion) */

  VECTOR(TBID_Q, 4); /* *o (--)    Transformation Matrix of body coord wrt
                       inertia coord derivative (Quaternion) */

  MATRIX(TBD, 3, 3); /* *o (--)    Transformation Matrix of body coord wrt
                    geodetic coord */
  VECTOR(TBDQ, 4);   /* *o  (--)    Quaternion from geodetic to body */

  VECTOR(VBAB, 3); /* *o  (m/s)    Air speed in body frame */

  MATRIX(WEII_skew, 3,
         3); /* *o  (r/s)    Earth's angular velocity (skew-sym) */

  VECTOR(SBIIP, 3); /* *o  (m)      Vehicle position in inertia coord */

  VECTOR(VBIIP, 3); /* *o  (m/s)    Vehicle inertia velocity */

  VECTOR(SBII, 3); /* *o  (m)      Vehicle position in inertia coord */

  VECTOR(VBII, 3); /* *o  (m/s)    Vehicle inertia velocity */

  VECTOR(ABII, 3); /* *o  (m/s2)   Vehicle inertia acceleration */

  VECTOR(ABIB,
         3); /* *o  (m/s2)   Vehicle inertia acceleration on body coordinate */

  VECTOR(SBEE, 3); /* *o  (m)     Vehicle position in earth coord  */

  VECTOR(VBEE, 3); /* *o  (m/s)     Vehicle speed in earth coord  */

  VECTOR(ABEE, 3); /* *o  (m/s2)   Vehicle acceleration in ECEF */

  VECTOR(SBEE_old, 3); /* *o  (m)     Vehicle position in earth coord  */

  VECTOR(VBEE_old, 3); /* *o  (m/s)     Vehicle speed in earth coord  */

  VECTOR(ABEE_old, 3); /* *o  (m/s2)   Vehicle acceleration in ECEF */

  VECTOR(JBII, 3); /* *o (m/s3)    Vehicle Jerk in ECI*/

  VECTOR(JBEE, 3); /* *o (m/s3)    Vehicle Jerk in ECEF */

  MATRIX(TDI, 3, 3); /* **  (--)     Transformation Matrix of geodetic wrt
                    inertial coordinates */

  MATRIX(
      TGI, 3,
      3); /* **  (--)     Transformation Matrix geocentric wrt inertia coord */

  VECTOR(VBED, 3); /* *o (m/s)   NED velocity */

  VECTOR(FSPB, 3); /* *o  (m/s2)   Specific force in body coord */

  VECTOR(NEXT_ACC, 3); /* *o (m/s2)   New Inertial acceleration */

  MATRIX(TDE, 3, 3); /* ** (--)  T.M. from ECEF to geodetic */

  VECTOR(VBII_old, 3); /* ** (m/s)  Prior body inertia velocity */

  VECTOR(WEII, 3); /* ** (--)  Earth rate in inertia coordinate */

  VECTOR(WBII, 3); /* *o (r/s)        Vehicle's inertia angular velocity in
                     inertia coord */

  VECTOR(WBEB, 3); /* *o (r/s)        Angular velocity of vehicle wrt earth in
                     body coord */

  VECTOR(WBIB, 3); /* *o (r/s)        Augular velocity of vehicle wrt inertia in
                     body coord */

  VECTOR(WBIBD, 3); /* *o (r/s2)       Angular velocity of vehicle wrt inertia
                      in body coord - derivative */

  MATRIX(TVD, 3, 3); /* **  (--)     Transformation Matrix of geographic
                    velocity wrt geodetic coord */

  VECTOR(SBEE_test, 3); /* *o  (m)     Vehicle position in earth coord  */

  VECTOR(VBEE_test, 3); /* *o  (m/s)     Vehicle speed in earth coord  */

  VECTOR(ABEE_test, 3); /* *o  (m/s2)   Vehicle acceleration in ECEF */

  MATRIX(TLI, 3, 3); /* **  (--)  T.M. from inertia to launch site */

  VECTOR(LT_euler, 3); /* ** (rad)  Launch site coordinate Euler angle */

  VECTOR(TBLQ, 4); /* ** (--) Quaternion from launch site to body */

  VECTOR(ABID, 3); /* ** (--) Vehicle inertia body acceleration in geodetic */

  VECTOR(FAPB,
         3); /* *o (N)      Aerodynamic and propulsion forces in body axes */

  VECTOR(FAP, 3); /* *o (N)      Aerodynamic force in body axes */

  VECTOR(FMB,
         3); /* *o (N*m)    Aerodynamic and propulsion moment in body axes */

  VECTOR(FMAB,
         3); /* *o (N*m)    Aerodynamic and propulsion moment in body axes */

  VECTOR(Q_G, 6); /* *o (--)     External force generated by gravity */

  VECTOR(Q_Aero, 6); /* *o (--)  External force generated by aerodynamics */

  VECTOR(rhoC_1, 3); /* *o  (m)  Level arm from reference ponit to CG */

  MATRIX(I1, 3, 3); /* *o (kg*m2)  MOI of vehicle */

  VECTOR(ddrP_1, 3); /* *o  (m/s2)  Vehicle acceleration */

  VECTOR(ddang_1, 3); /* *o  (r/s2)  Vehicle angular acceleration */

  VECTOR(dang_1, 3); /* *o  (r/s)   Vehicle angular rate */

  VECTOR(ddrhoC_1, 3); /* *o  (m/s2)  Centrifugal acceleration and tangential
                          acceleration term */

  VECTOR(p_b1_ga,
         3); /* *o (--)   General dynamics equations 1st DoF to 3rd DoF */

  VECTOR(p_b1_be,
         3); /* *o (--)   General dynamics equations 4th DoF to 6th DoF */

  VECTOR(f, 6); /* *o  (--)  Summation of external force & internal force */

  VECTOR(gamma_b1_q1, 3); /* *o (--)  Vehicle's 1st DoF velocity coefficient */

  VECTOR(gamma_b1_q2, 3); /* *o (--)  Vehicle's 2nd DoF velocity coefficient */

  VECTOR(gamma_b1_q3, 3); /* *o (--)  Vehicle's 3rd DoF velocity coefficient */

  VECTOR(beta_b1_q4,
         3); /* *o (--)  Vehicle's 4th DoF angular velocity coefficient */

  VECTOR(beta_b1_q5,
         3); /* *o (--)  Vehicle's 5th DoF angular velocity coefficient */

  VECTOR(beta_b1_q6,
         3); /* *o (--)  Vehicle's 6th DoF angular velocity coefficient */

  /* Generating Outputs */
  double
      ortho_error; /* *o (--)    Direction cosine matrix orthogonality error*/
  double alphax;   /* *o (d)     Angle of attack */
  double betax;    /* *o (d)     Sideslip angle */
  double alppx;    /* *o (d)     Total angle of attack */
  double phipx;    /* *o (d)     Aerodynamic roll angle*/
  double alphaix;  /* *o (d)     Angle of attack, inertia velocity*/
  double betaix;   /* *o (d)     Sideslip angle, inertia velocity*/
  double psibdx; /* *o (d)     Yaw angle of Vehicle wrt geodetic coord - deg */
  double
      thtbdx; /* *o (d)     Pitch angle of Vehicle wrt geodetic coord - deg */
  double phibdx; /* *o (d)     Roll angle of Vehicle wrt geodetic coord - deg */
  double psibd;  /* *o (r)     Yaw angle of Vehicle wrt geodetic coord - rad */
  double thtbd; /* *o (r)     Pitch angle of Vehicle wrt geodetic coord - rad */
  double phibd; /* *o (r)     Roll angle of Vehicle wrt geodetic coord - rad */
  double alt;   /* *o  (m)      Vehicle altitude */
  double lonx;  /* *o  (d)      Vehicle longitude */
  double latx;  /* *o  (d)      Vehicle latitude */
  double _aero_loss; /* *o  (m/s)    Velocity loss caused by aerodynamic drag */
  double gravity_loss; /* *o  (m/s)    Velocity loss caused by gravity */
  // double t;            /* *o (s)       timer */
  double _grndtrck; /* *o  (m)     [DIAG] Vehicle ground track on earth */
  double _gndtrkmx; /* *o  (km)    [DIAG] Ground track - km */
  double _gndtrnmx; /* **  (nm)    [DIAG] Ground track - nm */
  double _ayx;      /* *o  (m/s2)  [DIAG] Achieved side acceleration */
  double _anx;      /* *o  (m/s2)  [DIAG] Achieved normal acceleration */
  double _dbi;    /* *o  (m)     [DIAG] Vehicle distance from center of earth */
  double _dvbi;   /* *o  (m/s)   [DIAG] Vehicle inertia speed */
  double _dvbe;   /* *o  (m/s)   [DIAG] Vehicle geographic speed */
  double _thtvdx; /* *o  (d)     [DIAG] Vehicle's flight path angle */
  double _psivdx; /* *o  (d)     [DIAG] Vehicle's heading angle */
  int liftoff;    /* *i  (--)     To check wether the rocket liftoff or
                              not: liftoff = 1, not liftoff = 0 */
  int cadorbin_flag; /* Orbit calculation status flag */
  double ppx; /* *o (d/s)        Body roll angular velocity wrt earth in body
                 axes */
  double qqx; /* *o (d/s)        Body pitch angular velocity wrt earth in body
                 axes */
  double rrx; /* *o (d/s)        Body yaw angular velocity wrt earth in body
                 axes */
  double control_loss; /* *o (--) Velocity loss due to control effect */

  /* Orbital Logging */
  double _inclination;  /* *o  (deg)   [DIAG] Orbital inclination is the minimun
                        angle between reference plane and the orbital plane or
                        direction of an object in orbit around another object */
  double _eccentricity; /* *o  (--)    [DIAG] Determines the amount by which its
                           orbit around another body deviates from a perfect
                           circle */
  double _semi_major; /* *o  (m)     [DIAG] the major axis of an ellipse is its
                         longest diameter */
  double _ha;         /* *o  (m)     [DIAG] Orbital Apogee */
  double _hp;         /* *o  (m)     [DIAG] Orbital Perigee */
  double _lon_anodex; /* *o  (deg)   [DIAG] The longitude of the ascending node
                         (☊ or Ω) is one of the orbital elements used to specify
                         the orbit of an object in space. It is the angle from a
                         reference direction, called the origin of longitude, to
                         the direction of the ascending node, measured in a
                         reference plane */
  double _arg_perix;  /* *o  (deg)   [DIAG] The argument of periapsis (also
                         called argument of perifocus or argument of pericenter),
                         symbolized as ω, is one of the orbital elements of an
                         orbiting body. Parametrically, ω is the angle from the
                         body's ascending node to its periapsis, measured in the
                         direction of motion */
  double _true_anomx; /* *o  (deg)   [DIAG] In celestial mechanics, true anomaly
                      is an angular parameter that defines the position of a
                      body moving along a Keplerian orbit. It is the angle
                      between the direction of periapsis and the current
                      position of the body, as seen from the main focus of the
                      ellipse (the point around which the object orbits) */
  double _ref_alt;    /* *o  (m)     [DIAG] */
  double reference_point; /* *o (m)    Multibody dynamics reference point */
  double Roll;
  double Pitch;
  double Yaw;
  unsigned int Interpolation_Extrapolation_flag;
  // double xp; /* *o (m) Reference point  */
  int its;                /* *o (--) Number of iterations */
  int DOF;                /* *o (--)  Number of Degree of Freedom */
  unsigned int Aero_flag; /* *o (-)  Aerodynamics flag */
};
#endif


/* Copyright (C) 2010 Ion Torrent Systems, Inc. All Rights Reserved */
#include <armadillo>
#include <iostream>
#include <vector>
#include <string>
#include "IonErr.h"
#include "Image.h"
#include "Mask.h"
#include "OptArgs.h"
#include "SampleQuantiles.h"

using namespace arma;
using namespace std;

#define D '\t'
#define NUM_SAMPLE 10000
#define MAX_MILL_SEC 80
class CompressReporter {
public:
  virtual void Init() = 0;
  virtual void Report(int rowStart, int rowEnd, int colStart, int colEnd,
                      int flow, const std::vector<int> &wellMapping,
                      Mask &mask,
                      const Mat<float> &raw, const Mat<float> &compressed) = 0;
  virtual void Finish() = 0;
};

class WellReporter : public CompressReporter {
public:
  WellReporter(const std::string &fileOut) {
    mFileName = fileOut;
    mMad.Init(10000);
    mAutoCor.Init(10000);
  }
  void Init() {
    mOut.open(mFileName.c_str());
    mOut << "well\trow\tcol\tmad\tabs25\tabs75\trms\tacor" << endl;
  }

  void Report(int rowStart, int rowEnd, int colStart, int colEnd,
              int flow, const std::vector<int> &wellMapping,
              Mask &mask,
              const Mat<float> &raw, const Mat<float> &compressed) {
    Mat<float> diff = raw - compressed;
    for (size_t row = 0; row < raw.n_rows; row++) {
      SampleQuantiles<float> q(diff.n_cols);
      SampleStats<float>  s;
      SampleStats<float>  m;
      for (size_t col = 0; col < diff.n_cols; col++) {
        float val = diff(row, col);
        m.AddValue(val);
        q.AddValue(fabs(val));
        s.AddValue(val*val);
      }
      SampleStats<float> a;
      double mean = m.GetMean();
      for (size_t col = 1; col < diff.n_cols; col++) {
        float val1 = diff(row, col);
        float val2 = diff(row, col-1);
        a.AddValue((val1 - mean) * (val2 - mean));
      }
      double acor = a.GetMean() / m.GetVar();
      mAutoCor.AddValue(acor);
      mMad.AddValue(q.GetQuantile(.5));
      mOut << wellMapping[row] << D << wellMapping[row] / mask.W() << D << wellMapping[row] % mask.W() << D
           << q.GetQuantile(.5) << D << q.GetQuantile(.25) << D << q.GetQuantile(.75) << D << sqrt(s.GetMean()) << D << acor << endl;
    }
  }

  virtual void Finish() {
    cout << "Quantiles\tmad\tautocor" << endl;
    for (int i = 0; i <= 10; i++) {
      cout << i * 10 << D << mMad.GetQuantile(i / 10.0) << D << mAutoCor.GetQuantile(i / 10.0) << endl;
    }
    mOut.close();
  }
  SampleQuantiles<float> mMad;
  SampleQuantiles<float> mAutoCor;
  std::string mFileName;
  std::ofstream mOut;
};

class RegionIndividualReporter : public CompressReporter {
public:
  RegionIndividualReporter(const std::string &fileOut, int rowStart, int rowEnd, int colStart, int colEnd) {
    mFileName = fileOut;
    mRowStart = rowStart;
    mRowEnd = rowEnd;
    mColStart = colStart;
    mColEnd = colEnd;
    mFirst = true;
  }

  void Init() {
    mOut.open(mFileName.c_str());
    mOut << "well\ttype";
  }

  void Report(int rowStart, int rowEnd, int colStart, int colEnd,
              int flow, const std::vector<int> &wellMapping,
              Mask &mask,
              const Mat<float> &raw, const Mat<float> &compressed) {
    if (rowStart == mRowStart && rowEnd == mRowEnd && colStart == mColStart && colEnd == mColEnd) {
      if (mFirst) {
        for (size_t col = 0; col < raw.n_cols; col++) {
          mOut << D << "col." << col;
        }
        mOut << endl;
        mFirst = false;
      }
      for (size_t row = 0; row < raw.n_rows; row++) {
        mOut << wellMapping[row] << D << "raw";
        for (size_t col = 0; col < raw.n_cols; col++) {
          mOut << D << raw(row,col);
        }
        mOut << endl;
      }
      for (size_t row = 0; row < raw.n_rows; row++) {
        mOut << wellMapping[row] << D << "compressed";
        for (size_t col = 0; col < raw.n_cols; col++) {
          mOut << D << compressed(row,col);
        }
        mOut << endl;
      }
    }
  }

  virtual void Finish() { mOut.close(); }
  bool mFirst;
  int mRowStart, mRowEnd, mColStart, mColEnd;
  std::string mFileName;
  std::ofstream mOut;
};

class DeviationSummaryReporter : public CompressReporter {
public:
  DeviationSummaryReporter(const std::string &fileOut, const std::string &runName) {
    mFileName = fileOut;
    mRunName = runName;
  }

  void Init() {
    mQuantiles.Init(10000);
    mSample.Clear();
    mSqSample.Clear();
  }

  void Report(int rowStart, int rowEnd, int colStart, int colEnd,
              int flow, const std::vector<int> &wellMapping,
              Mask &mask,
              const Mat<float> &raw, const Mat<float> &compressed) {
    Mat<float> diff = raw - compressed;
    for (size_t row = 0; row < raw.n_rows; row++) {
      for (size_t col = 0; col < diff.n_cols; col++) {
        float val = fabs(diff(row, col));
        mQuantiles.AddValue(val);
        mSqSample.AddValue(val*val);
        mSample.AddValue(val);
      }
    }
  }

  virtual void Finish() {
    mOut.open(mFileName.c_str());
    mOut << "{" << endl;;
    mOut << "  run_name : " << mRunName << "," << endl;
    mOut << "  fabs_mean : " << mSample.GetMean() << ", "  << endl;
    mOut << "  fabs_sd : " << mSample.GetSD() << ", " << endl;
    mOut << "  rms : " << sqrt(mSqSample.GetMean()) << ", "  << endl;
    mOut << "  fabs_median : " << mQuantiles.GetQuantile(.5) << ", " << endl;
    mOut << "  fabs_iqr : " << mQuantiles.GetIQR() << ", "  << endl;
    mOut << "  fabs_quantiles: [";
    for (size_t i = 0; i < 10; i++) {
      float q = i / 10.0f;
      mOut << mQuantiles.GetQuantile(q) << ",";
    }
    mOut << mQuantiles.GetQuantile(1.0) <<  "], "  << endl;
    mOut << "  fabs_quantiles_100: [";
    for (size_t i = 0; i < 100; i++) {
      float q = i / 100.0f;
      mOut << mQuantiles.GetQuantile(q) << ",";
    }
    mOut << mQuantiles.GetQuantile(1.0) <<  "]"  << endl;
    mOut << "}" << endl;
    mOut.close();
  }

  std::string mFileName;
  SampleQuantiles<float> mQuantiles;
  SampleStats<float> mSqSample;
  SampleStats<float> mSample;
  std::string mRunName;
  std::ofstream mOut;
};

class FrameReporter : public CompressReporter {
public:
  FrameReporter(const std::string &fileOut) {
    mFileName = fileOut;
    mStarted = false;
  }
  void Init() {
    mOut.open(mFileName.c_str());
    mOut << "rowStart\trowEnd\tcolStart\tcolEnd\tnumWells\tflow\tstat";
    mStarted = false;
  }

  void Report(int rowStart, int rowEnd, int colStart, int colEnd,
              int flow, const std::vector<int> &wellMapping,
              Mask &mask,
              const Mat<float> &raw, const Mat<float> &compressed) {
    if (!mStarted) {
      for (size_t i = 0; i < raw.n_cols; i++) {
        mOut << D << "frame." << i;
      }
      mOut << endl;
      mQuantiles.resize(raw.n_cols);
      mStats.resize(raw.n_cols);
      mStarted = true;
    }
    for (size_t i = 0; i < raw.n_cols; i++) {
      mQuantiles[i].Clear();
      mQuantiles[i].Init(raw.n_rows);
      mStats[i].Clear();
    }
    Mat<float> diff = raw - compressed;
    for (size_t row = 0; row < diff.n_rows; row++) {
      for (size_t col = 0; col < diff.n_cols; col++) {
        float val = fabs(diff(row, col));
        mQuantiles[col].AddValue(val);
        mStats[col].AddValue(val);
      }
    }
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.Q25";
    for (size_t i = 0; i < mQuantiles.size(); i++) { mOut << D << mQuantiles[i].GetQuantile(.25); }
    mOut << std::endl;
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.Q50";
    for (size_t i = 0; i < mQuantiles.size(); i++) { mOut << D << mQuantiles[i].GetQuantile(.5); }
    mOut << std::endl;
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.Q75";
    for (size_t i = 0; i < mQuantiles.size(); i++) { mOut << D << mQuantiles[i].GetQuantile(.75); }
    mOut << std::endl;
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.IQR";
    for (size_t i = 0; i < mQuantiles.size(); i++) { mOut << D << mQuantiles[i].GetQuantile(.75) - mQuantiles[i].GetQuantile(.25); }
    mOut << std::endl;
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.Mean";
    for (size_t i = 0; i < mStats.size(); i++) { mOut << D << mStats[i].GetMean(); }
    mOut << std::endl;
    mOut << rowStart << D << rowEnd << D << colStart << D << raw.n_rows << D << flow << D << "Fabs.SD";
    for (size_t i = 0; i < mStats.size(); i++) { mOut << D << mStats[i].GetSD(); }
    mOut << std::endl;
  }

  virtual void Finish() { mOut.close(); }

  std::vector<SampleQuantiles<float> > mQuantiles;
  std::vector<SampleStats<float>  > mStats;
  bool mStarted;
  std::string mFileName;
  std::ofstream mOut;
};

class CompressOpts {
public:
  int rowStart, rowEnd, colStart, colEnd;
  int rowStep, colStep;
  double minPercent;
  std::string datIn;
  std::string maskFile;
  int flowNum;
  std::string prefixOut;
  std::string wellOut;
  std::string frameOut;
  std::string runName;
  int numBasis;
};

void AddMatVector(Mat<float> &M, const Col<float> &vec) {
  for (size_t i = 0; i < M.n_rows; i++) {
    for (size_t j = 0; j < M.n_cols; j++) {
      M(i,j) = M(i,j) + vec(i);
    }
  }
}

void Compress(const Mat<float> &raw, int numBasis, Mat<float> &compressed) {
  Col<float> dc = mean(raw, 1);
  Mat<float> X = raw;
  dc = dc * -1;
  AddMatVector(X, dc);
  Mat<float> Cov = X.t() * X;    // L x L
  Mat<float> EVec;
  Col<float> EVal;
  eig_sym(EVal, EVec, Cov);
  Mat<float> V(EVec.n_rows, numBasis );  // L x numBasis
  int count = 0;
  for(size_t v = V.n_rows - 1; v >= V.n_rows - numBasis; v--) {
    copy(EVec.begin_col(v), EVec.end_col(v), V.begin_col(count++));
  }
  Mat<float> A = X * V;  // N x numBasis
  Mat<float> P = A * V.t();  // Prediction N x L
  dc = dc * -1;
  AddMatVector(P, dc);
  compressed = P;
}

int LoadData(int rowStart, int rowEnd, int colStart, int colEnd,
              Mask &mask, Image &img, vector<int> &mapping, Mat<float> &raw) {
  int numOk = 0;
  for (int row = rowStart; row < rowEnd; row++) {
    for (int col = colStart; col < colEnd; col++) {
      if (!mask.Match(col, row, MaskPinned) && !mask.Match(col, row, MaskExclude)) {
        numOk++;
        mapping.push_back(mask.ToIndex(row, col));
      }
    }
  }
  const RawImage *rawImg = img.GetImage();
  cout << "Timestamps:" << endl;
  vector<bool> colOk(img.GetFrames(), false);
  int count = 0;
  if (rawImg->timestamps[0] < MAX_MILL_SEC) {
    count++;
    colOk[0] = true;
  }
  for (int i = 1; i < img.GetFrames(); i++) {
    cout << i << ":\t" << rawImg->timestamps[i] << "\t" << rawImg->timestamps[i] - rawImg->timestamps[i-1] << endl;
    if (rawImg->timestamps[i] - rawImg->timestamps[i-1] < MAX_MILL_SEC) {
      colOk[i] = true;
      count++;
    }
  }
  cout << "Got: " << count << " frames with no vfr" << endl;
  raw.set_size(numOk, count);
  for (size_t wIx = 0; wIx < mapping.size(); wIx++) {
    int idx = mapping[wIx];
    int cIx = 0;
    for (int fIx = 0; fIx < img.GetFrames(); fIx++) {
      if (colOk[fIx]) {
        raw(wIx,cIx++) = img.At(idx, fIx);
      }
    }
  }
  return numOk;
}

int main(int argc, const char * argv[]) {

  OptArgs opts;
  opts.ParseCmdLine(argc, argv);
  CompressOpts o;
  opts.GetOption(o.maskFile, "", '-', "mask");
  Mask mask(o.maskFile.c_str());
  opts.GetOption(o.rowStart, "0", '-', "row-start");
  opts.GetOption(o.rowEnd, "-1", '-', "row-end");
  if(o.rowEnd < 0) { o.rowEnd = mask.H(); }

  opts.GetOption(o.colStart, "0", '-', "col-start");
  opts.GetOption(o.colEnd, "-1", '-', "col-end");
  if(o.colEnd < 0) { o.colEnd = mask.W(); }

  opts.GetOption(o.datIn, "", '-', "dat-file");
  opts.GetOption(o.prefixOut, "noise", '-', "out-prefix");
  opts.GetOption(o.numBasis, "6", '-', "num-basis");
  opts.GetOption(o.rowStep, "100", '-', "row-step");
  opts.GetOption(o.colStep, "100", '-', "col-step");
  opts.GetOption(o.minPercent, ".5", '-', "min-percent");
  opts.GetOption(o.runName, "", '-', "run-name");
  // Load file
  Image img;
  img.SetImgLoadImmediate (false);
  img.SetIgnoreChecksumErrors (false);
  ION_ASSERT(img.LoadRaw (o.datIn.c_str()), "Couldn't load: " + o.datIn);

  // Create reporters
  vector<CompressReporter *> reporters;
  WellReporter wRep(o.prefixOut + ".wells.txt");
  reporters.push_back(&wRep);
  FrameReporter frameRep(o.prefixOut + ".frames.txt");
  reporters.push_back(&frameRep);
  DeviationSummaryReporter devRep(o.prefixOut + ".summary.json", o.runName);
  reporters.push_back(&devRep);
  int regRowStart = 0;
  if (o.rowStep * 5 < mask.H()) {
    regRowStart = o.rowStep * 4;
  }
  int regColStart = 0;
  if (o.colStep * 5 < mask.H()) {
    regColStart = o.colStep * 4;
  }
  RegionIndividualReporter regionRep(o.prefixOut + ".region.txt", regRowStart, regRowStart + o.rowStep, regColStart, regColStart + o.colStep);
  reporters.push_back(&regionRep);
  for (size_t rIx = 0; rIx < reporters.size(); rIx++) {
    reporters[rIx]->Init();
  }
  // For each region do compression and run report.
  for (int row = o.rowStart; row < o.rowEnd; row += o.rowStep) {
    for (int col = o.colStart; col < o.colEnd; col += o.colStep) {
      int rowStop = min(row + o.rowStep, o.rowEnd);
      int colStop = min(col + o.colStep, o.colEnd);
      Mat<float> raw;
      vector<int> mapping;
      int numOk = LoadData(row, rowStop, col, colStop, mask, img,  mapping, raw);
      if (numOk > o.minPercent * (rowStop - row) * (colStop - col)) {
        Mat<float> compressed(raw.n_rows, raw.n_cols);
        Compress(raw, o.numBasis, compressed);
        for (size_t rIx = 0; rIx < reporters.size(); rIx++) {
          reporters[rIx]->Report(row, rowStop, col, colStop, 0,
                                 mapping, mask, raw, compressed);
        }
      }
    }
  }
  for (size_t rIx = 0; rIx < reporters.size(); rIx++) {
    reporters[rIx]->Finish();
  }

  // Spin things down.

}


// Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
//       its contributors may be used to endorse or promote products derived
//       from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Author: Johannes L. Schoenberger (jsch at inf.ethz.ch)

#include "base/polynomial.h"

#include <Eigen/Eigenvalues>

#include "util/logging.h"

namespace colmap {
namespace {

// Remove leading zero coefficients.
Eigen::VectorXd RemoveLeadingZeros(const Eigen::VectorXd& coeffs) {
  Eigen::VectorXd::Index num_zeros = 0;
  for (; num_zeros < coeffs.size(); ++num_zeros) {
    if (coeffs(num_zeros) != 0) {
      break;
    }
  }
  return coeffs.tail(coeffs.size() - num_zeros);
}

// Remove trailing zero coefficients.
Eigen::VectorXd RemoveTrailingZeros(const Eigen::VectorXd& coeffs) {
  Eigen::VectorXd::Index num_zeros = 0;
  for (; num_zeros < coeffs.size(); ++num_zeros) {
    if (coeffs(coeffs.size() - 1 - num_zeros) != 0) {
      break;
    }
  }
  return coeffs.head(coeffs.size() - num_zeros);
}

}  // namespace

bool FindLinearPolynomialRoots(const Eigen::VectorXd& coeffs,
                               Eigen::VectorXd* real, Eigen::VectorXd* imag) {
  CHECK_EQ(coeffs.size(), 2);

  if (coeffs(0) == 0) {
    return false;
  }

  if (real != nullptr) {
    real->resize(1);
    (*real)(0) = -coeffs(1) / coeffs(0);
  }

  if (imag != nullptr) {
    imag->resize(1);
    (*imag)(0) = 0;
  }

  return true;
}

bool FindQuadraticPolynomialRoots(const Eigen::VectorXd& coeffs,
                                  Eigen::VectorXd* real,
                                  Eigen::VectorXd* imag) {
  CHECK_EQ(coeffs.size(), 3);

  const double a = coeffs(0);
  if (a == 0) {
    return FindLinearPolynomialRoots(coeffs.tail(2), real, imag);
  }

  const double b = coeffs(1);
  const double c = coeffs(2);
  if (b == 0 && c == 0) {
    if (real != nullptr) {
      real->resize(1);
      (*real)(0) = 0;
    }
    if (imag != nullptr) {
      imag->resize(1);
      (*imag)(0) = 0;
    }
    return true;
  }

  const double d = b * b - 4 * a * c;

  if (d >= 0) {
    const double sqrt_d = std::sqrt(d);
    if (real != nullptr) {
      real->resize(2);
      if (b >= 0) {
        (*real)(0) = (-b - sqrt_d) / (2 * a);
        (*real)(1) = (2 * c) / (-b - sqrt_d);
      } else {
        (*real)(0) = (2 * c) / (-b + sqrt_d);
        (*real)(1) = (-b + sqrt_d) / (2 * a);
      }
    }
    if (imag != nullptr) {
      imag->resize(2);
      imag->setZero();
    }
  } else {
    if (real != nullptr) {
      real->resize(2);
      real->setConstant(-b / (2 * a));
    }
    if (imag != nullptr) {
      imag->resize(2);
      (*imag)(0) = std::sqrt(-d) / (2 * a);
      (*imag)(1) = -(*imag)(0);
    }
  }

  return true;
}

bool FindPolynomialRootsDurandKerner(const Eigen::VectorXd& coeffs_all,
                                     Eigen::VectorXd* real,
                                     Eigen::VectorXd* imag) {
  CHECK_GE(coeffs_all.size(), 2);

  const Eigen::VectorXd coeffs = RemoveLeadingZeros(coeffs_all);

  const int degree = coeffs.size() - 1;

  if (degree <= 0) {
    return false;
  } else if (degree == 1) {
    return FindLinearPolynomialRoots(coeffs, real, imag);
  } else if (degree == 2) {
    return FindQuadraticPolynomialRoots(coeffs, real, imag);
  }

  // Initialize roots.
  Eigen::VectorXcd roots(degree);
  roots(degree - 1) = std::complex<double>(1, 0);
  for (int i = degree - 2; i >= 0; --i) {
    roots(i) = roots(i + 1) * std::complex<double>(1, 1);
  }

  // Iterative solver.
  const int kMaxNumIterations = 100;
  const double kMaxRootChange = 1e-10;
  for (int iter = 0; iter < kMaxNumIterations; ++iter) {
    double max_root_change = 0.0;
    for (int i = 0; i < degree; ++i) {
      const std::complex<double> root_i = roots(i);
      std::complex<double> numerator = coeffs[0];
      std::complex<double> denominator = coeffs[0];
      for (int j = 0; j < degree; ++j) {
        numerator = numerator * root_i + coeffs[j + 1];
        if (i != j) {
          denominator = denominator * (root_i - roots(j));
        }
      }
      const std::complex<double> root_i_change = numerator / denominator;
      roots(i) = root_i - root_i_change;
      max_root_change =
          std::max(max_root_change, std::abs(root_i_change.real()));
      max_root_change =
          std::max(max_root_change, std::abs(root_i_change.imag()));
    }

    // Break, if roots do not change anymore.
    if (max_root_change < kMaxRootChange) {
      break;
    }
  }

  if (real != nullptr) {
    real->resize(degree);
    *real = roots.real();
  }
  if (imag != nullptr) {
    imag->resize(degree);
    *imag = roots.imag();
  }

  return true;
}

bool FindPolynomialRootsCompanionMatrix(const Eigen::VectorXd& coeffs_all,
                                        Eigen::VectorXd* real,
                                        Eigen::VectorXd* imag) {
  CHECK_GE(coeffs_all.size(), 2);

  Eigen::VectorXd coeffs = RemoveLeadingZeros(coeffs_all);

  const int degree = coeffs.size() - 1;

  if (degree <= 0) {
    return false;
  } else if (degree == 1) {
    return FindLinearPolynomialRoots(coeffs, real, imag);
  } else if (degree == 2) {
    return FindQuadraticPolynomialRoots(coeffs, real, imag);
  }

  // Remove the coefficients where zero is a solution.
  coeffs = RemoveTrailingZeros(coeffs);

  // Check if only zero is a solution.
  if (coeffs.size() == 1) {
    if (real != nullptr) {
      real->resize(1);
      (*real)(0) = 0;
    }
    if (imag != nullptr) {
      imag->resize(1);
      (*imag)(0) = 0;
    }
    return true;
  }

  // Fill the companion matrix.
  Eigen::MatrixXd C(coeffs.size() - 1, coeffs.size() - 1);
  C.setZero();
  for (Eigen::MatrixXd::Index i = 1; i < C.rows(); ++i) {
    C(i, i - 1) = 1;
  }
  C.row(0) = -coeffs.tail(coeffs.size() - 1) / coeffs(0);

  // Solve for the roots of the polynomial.
  Eigen::EigenSolver<Eigen::MatrixXd> solver(C, false);
  if (solver.info() != Eigen::Success) {
    return false;
  }

  // If there are trailing zeros, we must add zero as a solution.
  const int effective_degree =
      coeffs.size() - 1 < degree ? coeffs.size() : coeffs.size() - 1;

  if (real != nullptr) {
    real->resize(effective_degree);
    real->head(coeffs.size() - 1) = solver.eigenvalues().real();
    if (effective_degree > coeffs.size() - 1) {
      (*real)(real->size() - 1) = 0;
    }
  }
  if (imag != nullptr) {
    imag->resize(effective_degree);
    imag->head(coeffs.size() - 1) = solver.eigenvalues().imag();
    if (effective_degree > coeffs.size() - 1) {
      (*imag)(imag->size() - 1) = 0;
    }
  }

  return true;
}

}  // namespace colmap