/* Copyright (C) 2011 Ion Torrent Systems, Inc. All Rights Reserved */

#include <armadillo>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <cinttypes>

#include "ByteSwapUtils.h"
#include "LSRowImageProcessor.h"

#define N_CORRECTION_GROUPS 8
#define SKIP_ROWS_AT_END    10

#define PINNED_HIGH_LIMIT   16373
#define PINNED_LOW_LIMIT    10

#define LSROWIMAGE_MAGIC_VALUE    0xFF115E3A

using namespace arma;
// generate an electrical cross-talk correction from the lsrow image
// if the file pointed to by lsimg_path does not exist, the method returns NULL
// If lsimg_path does exist, then a correction is generated and a pointer to
// a ChannelXTCorrection object with all the relevant information is returned
ChannelXTCorrection *LSRowImageProcessor::GenerateCorrection(const char *lsimg_path)
{
    uint32_t magic_value;
    uint32_t file_version;
    uint32_t rows = 0;
    uint32_t cols = 0;
    int nread;

    // the lsrowimage file contains:
    // uint32 magic value (0xFF115E3A)
    // uint32 file version number (0+)
    // uint32 rows
    // uint32 columns
    // uint16[rows*columns] high-speed and low-speed reference data, in row major order
    //  ..with the first row in the image being a high-speed collected row
    //    and the next row in the image being the low-speed reference for the previous high-speed row data

    FILE *lsrowfile;

    lsrowfile = fopen(lsimg_path,"rb");

    // if we have trouble opening the file, just return NULL
    if(lsrowfile == NULL)
        return NULL;

    nread = fread(&magic_value,sizeof(int32_t),1,lsrowfile);
    if(nread != 1)
    {
        printf("Ivalid lsrowfile detected\n");
        fclose(lsrowfile);
        return NULL;
    }
    magic_value = BYTE_SWAP_4(magic_value);

    if(magic_value != LSROWIMAGE_MAGIC_VALUE)
    {
        printf("Ivalid lsrowfile detected\n");
        fclose(lsrowfile);
        return NULL;
    }

    nread = fread(&file_version,sizeof(int32_t),1,lsrowfile);
    if(nread != 1)
    {
        printf("Ivalid lsrowfile detected\n");
        fclose(lsrowfile);
        return NULL;
    }
    file_version = BYTE_SWAP_4(file_version);

    if(file_version != 0)
    {
        printf("Unsupported lsrowimage file version\n");
        fclose(lsrowfile);
        return NULL;
    }

    nread = fread(&rows,sizeof(int32_t),1,lsrowfile);
    if(nread != 1)
    {
        printf("Ivalid lsrowfile detected\n");
        fclose(lsrowfile);
        return NULL;
    }
    rows = BYTE_SWAP_4(rows);

    nread = fread(&cols,sizeof(int32_t),1,lsrowfile);
    if(nread != 1)
    {
        printf("Ivalid lsrowfile detected\n");
        fclose(lsrowfile);
        return NULL;
    }
    cols = BYTE_SWAP_4(cols);

    int tot_pts = rows*cols;

    printf("reading lsrowfile with %d rows and %d columns\n",rows,cols);
    uint16_t *img = new uint16_t[tot_pts];

    nread = fread(img,sizeof(uint16_t),tot_pts,lsrowfile);
    if(nread != tot_pts)
    {
        printf("Ivalid lsrowfile detected\n");
        delete [] img;
        fclose(lsrowfile);
        return NULL;
    }

    // byte swap the image
    for(int i=0;i < tot_pts;i++)
        img[i] = BYTE_SWAP_2(img[i]);

    // this version of the code generates a correction for columns-modulo eight. (it generates a correction
    // for columns that belong to one of eight different groups, where membership is determined from
    // (column % 8))
    ChannelXTCorrection *xtptr = new ChannelXTCorrection();
    float *pvects = xtptr->AllocateVectorStorage(N_CORRECTION_GROUPS,nLen);
    float **vect_ptrs = xtptr->AllocateVectorPointerStorage(N_CORRECTION_GROUPS);
    xtptr->SetVectorIndicies(indicies,nLen);

    bool correction_valid = true;

    for(int i=0;i < N_CORRECTION_GROUPS;i++)
    {
        vect_ptrs[i] = pvects+nLen*i;
        if (GenerateGroupCorrection(i,vect_ptrs[i],rows,cols,img) == false)
        {
            correction_valid = false;
            break;
        }
    }

    delete [] img;
    fclose(lsrowfile);

    if(!correction_valid)
    {
        printf("Unable to compute valid correction\n");
        delete xtptr;
        return NULL;
    }

    return xtptr;
}

// generates a set of correction coefficients for a group of columns
// this basically solves the matrix equation Ax=B for x, where A are the measured high-speed
// pixel values and B are the low-speed measured pixels values.  The special image collected on the PGM
// contains pairs of high-speed and low-speed rows for the same pixels that can be used to populate
// these matricies
//
// As written, A is a nLen column matrix, with many rows in it (one for each example pixel we extract from the data)
// and B is an nLen row vector, again with one row per example pixel we extract from the data.
// When the equation is solved, both sides are multiplied by A-tranpose, which turns the left side into a square
// nLen by nLen matrix, and the rhs into an nLen row vector.  In order to facilitate ease of adding example pixels to
// the equations sequentially, A and B are not computed, and instead lhs = (A-transpose x A) and rhs = (A-transpose x B)
// are computed directly from the low-level data, and the solution is determined from lhs x = rhs
bool LSRowImageProcessor::GenerateGroupCorrection(int group_num,float *vect_output,int rows,int cols,uint16_t *img)
{
    double lhs[nLen*nLen];
    double rhs[nLen];

    memset(lhs,0,sizeof(lhs));
    memset(rhs,0,sizeof(rhs));

    // the image contains some invalid rows at the end that should be skipped
    int row_limit = rows - SKIP_ROWS_AT_END;

    int mcnt = 0;

    for(int column = group_num;column < cols;column += N_CORRECTION_GROUPS)
    {
        // make sure we have enough space on the left and right-hand sides to use this particular column
        if (((column + indicies[0]) < 0) || ((column + indicies[nLen-1]) >= cols))
            continue;

        // every other row is the start of a pair of rows, one high-speed and one reference
        for(int row = 0;row < row_limit;row += 2)
        {
            double amat[nLen];
            bool skip = false;

            uint16_t *hsrow = &img[row*cols];
            uint16_t *lsrow = &img[(row+1)*cols];

            // get the data points that are to be added to the matrix, making sure to filter out entries
            // that might reference pinned values
            for(int i=0;i < nLen;i++)
            {
                uint16_t temp = hsrow[indicies[i]+column];

                if((temp < PINNED_LOW_LIMIT) || (temp > PINNED_HIGH_LIMIT))
                {
                    skip = true;
                    break;
                }
                amat[i] = (double)temp;

                temp = lsrow[indicies[i]+column];
                if((temp < PINNED_LOW_LIMIT) || (temp > PINNED_HIGH_LIMIT))
                {
                    skip = true;
                    break;
                }
            }

            // if everything checks out, add this data into the matrix equation
            if(!skip)
            {
                mcnt++;
                AccumulateMatrixData(lhs,rhs,amat,(double)(lsrow[column]));
            }
        }
    }

    Mat<double> lhs_matrix(nLen,nLen);
    Col<double> rhs_vector(nLen);
    Col<double> coeffs(nLen);
    bool result_ok = true;

    for(int col=0;col < nLen;col++)
        for(int row=0;row <= col;row++)
        {
            lhs_matrix(row,col) = lhs[row*nLen+col];
            lhs_matrix(col,row) = lhs[row*nLen+col];
        }

    for(int row=0;row < nLen;row++)
        rhs_vector(row) = rhs[row];

    try {
      //LaSpdMatFactorize(lhs_matrix,lhs_matrix_fact);
      //LaLinearSolve(lhs_matrix_fact,coeffs,rhs_vector);
      coeffs = solve(lhs_matrix,rhs_vector);
    }
    catch (std::runtime_error& le) {
        result_ok = false;
        coeffs.zeros(nLen);
    }

    // make sure derived coefficients are valid
    for(int row=0;row < nLen;row++)
        if(std::isnan(coeffs(row)))
        {
            result_ok = false;
            break;
        }
        else
            vect_output[row] = coeffs(row);

    printf("group %d correction coefficients: ",group_num);
    for(int row=0;row < nLen;row++)
        printf("%11.8lf ",coeffs(row));

    printf("\n");

    return(result_ok);
}

// Adds one example pixel's data into the lhs and rhs matrix and vector
void LSRowImageProcessor::AccumulateMatrixData(double *lhs,double *rhs,double *amat,double bval)
{
    for(int col=0;col < nLen;col++)
        for(int row=0;row <= col;row++)
            lhs[row*nLen+col] += amat[row]*amat[col];

    for(int row=0;row < nLen;row++)
        rhs[row] += amat[row]*bval;
}


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

#ifndef NMIII_APCG_HH
#define NMIII_APCG_HH

#include <numeric>
#include <string>

#include <boost/circular_buffer.hpp>
#include <boost/timer/timer.hpp>

#include "dune/common/timer.hh"
#include "dune/istl/istlexception.hh"
#include "dune/istl/operators.hh"
#include "dune/istl/preconditioners.hh"
#include "dune/istl/scalarproducts.hh"
#include "dune/istl/solvers.hh"

#include "linalg/symmetricOperators.hh"
#include "utilities/detailed_exception.hh"
#include "utilities/geometric_sequence.hh"
#include "utilities/scalar.hh"

namespace Kaskade
{
  /**
   * \ingroup iterative
   * \brief Whether a preconditioner needs zero initialized result vector or not.
   *
   * For a couple of known preconditioners, this function returns the actual information.
   * For unknown preconditioners, it errs on the conservative side, i.e. reports true.
   *
   * The function is moderately expensive, hence don't use it inside of loops.
   */
  template <class Domain, class Range>
  bool requiresInitializedInput(Dune::Preconditioner<Domain,Range> const* p)
  {
    if (dynamic_cast<Dune::Richardson<Domain,Range> const*>(p))
      return false;
    if (auto sp = dynamic_cast<SymmetricPreconditioner<Domain,Range> const*>(p))
      return sp->requiresInitializedInput();
    return true;
  }

  /**
   * \ingroup iterative
   * \brief Interface for IterateType::PCG termination criterion policy classes
   * \tparam R a floating point type for real numbers
   */
  template <class R>
  class PCGTerminationCriterion
  {
  public:
    /**
     * \brief real field type
     */
    typedef R Real;

    /**
     * \brief re-initializes the termination criterion for a new IterateType::CG run
     */
    virtual void clear() = 0;

    /**
     * \brief set requested tolerance
     *
     * \param tol the requested tolerance (nonnegative)
     */
    virtual PCGTerminationCriterion<R>& tolerance(Real tol) = 0;

    /**
     * \brief supplies energy of step to the termination criterion
     * \param gamma2 the energy of the step
     */
    virtual void step(Real gamma2) = 0;

    /**
     * \brief supplies the preconditioned residual to the termination criterion
     * \param sigma the preconditioned residual norm
     */
    virtual void residual(Real sigma) = 0;

    /**
     * \brief termination decision
     * \return true if the iteration has reached the required accuracy
     */
    virtual operator bool() const = 0;
  };

  /**
   * \ingroup iterative
   * \brief TerminationCriterion based on an absolute energy error estimate
   *
   * This termination criterion assumes the error reduction occurs in two phases: First a smoothing phase where
   * \f$ \epsilon_k^2 \approx a^2/(1+\sqrt{k})^2 + b^2 \f$, and a linear convergence phase with \f$ \epsilon_k \approx q^k \epsilon_0 \f$.
   * The termination criterion tries to fit the model parameters \f$ a, b, q \f$ and to decide where the crossover of both
   * phases occurs. Based on this decision, the energy error \f$ \epsilon_k \f$ is estimated as \f$ [\epsilon_k] \f$.
   *
   * For greater robustnes, the actual error estimate used is \f[ [\epsilon_{k-l}]^2 := [\epsilon_k]^2 + \sum_{i=k-l}^{k-1} \gamma_i^2. \f] with
   * \f$ \gamma_i = \|u_{i+1}-u_i\|_A \f$.The lookahead value \f$ l \f$ can be set using \ref lookahead.
   *
   * \tparam R the real field type. Currently instantiated for float and double in file apcg.cpp.
   */
  template <class R>
  class PCGEnergyErrorTerminationCriterion: public PCGTerminationCriterion<R> {
  public:
    typedef R Real;

    /**
     * \brief constructor
     *
     * The pcg iteration is terminated as soon as either \f$ [\epsilon] \le \mathrm{atol} \f$ or
     * the number of iterations exceeds the limit maxit.
     *
     * \param atol the absolute error tolerance for termination
     * \param maxit the maximum number of iterations
     */
    PCGEnergyErrorTerminationCriterion(Real atol, int maxit);

    /**
     * \brief re-initializes the termination criterion for a new IterateType::CG run
     */
    virtual void clear();

    /**
     * \brief set requested absolute tolerance
     *
     * \param atol the requested tolerance (nonnegative)
     */
    virtual PCGEnergyErrorTerminationCriterion<Real>& tolerance(Real atol);

    /**
     * \brief set requested lookahead value
     *
     * \param lah the requested lookahead (nonnegative)
     *
     * The default value is 6.
     */
     PCGEnergyErrorTerminationCriterion<Real>& lookahead(int lah);

    /**
     * \brief supplies energy of step to the termination criterion
     * \param gamma2 the energy \f$ \gamma^2 = \| u_{k+1}-u_k \|_A^2 \f$ of the step
     */
    virtual void step(Real gamma2);

    /**
     * \brief supplies the preconditioned residual to the termination criterion
     * \param sigma the preconditioned residual norm
     */
    virtual void residual(Real sigma) {}

    /**
     * \brief termination decision
     * \return true if the iteration has reached the required accuracy
     */
    virtual operator bool() const;

    /**
     * \brief returns the estimated absolute energy error
     */
    Real error() const;

    std::vector<Real> const& gamma2() const {return gammas2; }

  private:
    Real tol;
    int maxit;
    // squared gammas
    std::vector<Real> gammas2;
    int lookah;
  };

 /**
  * \ingroup iterative
  * \brief preconditioned conjugate gradient method
  *
  * This implements a preconditioned IterateType::CG iteration for an operator \f$ A: X\to X^* \f$, preconditioned by a
  * preconditioner \f$ B^{-1}: X^* \to X \f$. The termination is policy-based.
  *
  * The implementation follows Deuflhard/Weiser, Section 5.3.3.
  *
  * \tparam X the type of vectors from the primal space
  * \tparam Xstar the type of vectors from the dual space
  *
  * \see PCGEnergyErrorTerminationCriterion
  */
  template<class X, class Xstar>
  class Pcg : public Dune::InverseOperator<X,Xstar> {
  public:
    /// \brief The domain type of the operator to be inverted.
    typedef X domain_type;
    /// \brief The range type of the operator to be inverted.
    typedef Xstar range_type;
    /// \brief The field type of the operator to be inverted.
    typedef typename X::field_type field_type;

    /**
     * \brief the real field type corresponding to field_type
     */
    typedef typename ScalarTraits<field_type>::Real Real;

    /**
     * \brief Set up conjugate gradient solver.
     *
     * \param op the operator
     * \param prec the preconditioner
     * \param terminate the termination criterion. The object has to exist during the lifetime of the pcg object as it is referenced.
     * \param verbose controls the verbosity of logging to std::cout. 0 means no output at all (default). Values in 0,1,2 are valid.
     */
    Pcg(SymmetricLinearOperator<X,Xstar>& op_, SymmetricPreconditioner<X,Xstar>& prec_,
        PCGTerminationCriterion<Real>& terminate_, int verbose_=0) :
      proxyOp(op_,zeroDp), op(op_),
      proxyPrec(prec_,zeroDp), prec(prec_),
      terminate(terminate_), verbose(verbose_), requiresInit(requiresInitializedInput(&prec))
    {
// Do we need this in Kaskade7?
//         dune_static_assert( static_cast<int>(L::category) == static_cast<int>(P::category),
//                             "L and P must have the same category!");
//         dune_static_assert( static_cast<int>(L::category) == static_cast<int>(SolverCategory::sequential),
//                             "L must be sequential!");
    }

    /**
     * \brief Set up conjugate gradient solver.
     *
     * \param op the operator
     * \param prec the preconditioner
     * \param dp the dual pairing with respect to which the operator is symmetric
     * \param terminate the termination criterion. The object has to exist during the lifetime of the pcg object as it is referenced.
     * \param verbose controls the verbosity of logging to std::cout. 0 means no output at all (default). Values in 0,1,2 are valid.
     */
    Pcg(Dune::LinearOperator<X,Xstar>& op_, Dune::Preconditioner<X,Xstar>& prec_, DualPairing<X,Xstar> const& dp_,
        PCGTerminationCriterion<Real>& terminate_, int verbose_=0) :
      proxyOp(op_,dp_), op(proxyOp),
      proxyPrec(prec_,dp_), prec(proxyPrec),
      terminate(terminate_), verbose(verbose_), requiresInit(requiresInitializedInput(&prec))
    {    }

    /**
     * \brief Apply inverse operator by performing a number of IterateType::PCG iterations.
     *
     * \param u the initial value (starting iterate)
     * \param b the right hand side (which will not be modified)
     */
    virtual void apply (X& u, Xstar& b, Dune::InverseOperatorResult& res) {
      res.clear();                // clear solver statistics
      terminate.clear();          // clear termination criterion
      Dune::Timer watch;          // start a timer

      boost::timer::cpu_timer mvTimer;
      mvTimer.stop();
      boost::timer::cpu_timer apTimer;
      apTimer.stop();
      boost::timer::cpu_timer prTimer;
      prTimer.stop();

      Xstar r(b);
      op.applyscaleadd(-1,u,r);  // r = b-Au

      // Keeping rq and q in unique_ptr allows efficient computation of q = rq + beta*q below by swapping the pointers.
      // More elegant would be a swap method directly on X, but this is hard to implement as there is no swap either
      // for boost::fusion vectors nor for Dune::BlockVectors. Even more elegant would be a loop fusion expression
      // framework for BLAS level 1 operations.
      std::unique_ptr<X> rq(new X(u)); *rq = 0;
      prec.apply(*rq,r); // rq = B^{-1} r

      std::unique_ptr<X> q(new X(*rq));

      X Aq(b);

      // some local variables
      field_type alpha,beta,sigma,gamma;
      sigma = op.dp(*rq,r); // preconditioned residual norm squard
      terminate.residual(ScalarTraits<field_type>::real(sigma));
      double const sigma0 = std::abs(sigma);

      // Check for trivial case
      if (sigma0 == 0)
        return;

      if (verbose>0) {            // printing
        std::cout << "=== Kaskade7 IterateType::PCG" << std::endl;
        if (verbose>1) {
          this->printHeader(std::cout);
          this->printOutput(std::cout,0.0,sigma0);
        }
      }

      // the loop
      int i=1;
      for ( ; true; i++ ) {
        // minimize in given search direction p
        mvTimer.resume();
        field_type qAq = op.applyDp(*q,Aq);             // compute Aq = A*q and qAq = <q,A*q>
if (std::isnan(qAq)) std:: cerr << "qAq is nan\n";
        mvTimer.stop();

	if (std::abs(qAq) < 1e-28*sigma)
	  break; // oops.

        if (qAq <= 0)
          throw NonpositiveMatrixException("encountered nonpositive energy product " + std::to_string(qAq) + " vs sigma " + std::to_string(sigma) + " in IterateType::PCG.",__FILE__,__LINE__);

        alpha = sigma/qAq;
        apTimer.resume();
        u.axpy(alpha,*q);
        apTimer.stop();
        gamma = sigma*alpha;
	terminate.step(ScalarTraits<field_type>::real(gamma));
	resDebug.push_back(std::sqrt(sigma));

        // convergence test
        if (terminate)
	  break;

        apTimer.resume();
        r.axpy(-alpha,Aq); // r = r - alpha*A*q
        apTimer.stop();

        prTimer.resume();
        if (requiresInit)
          *rq = 0;
        double sigmaNew = prec.applyDp(*rq,r); // compute rq = B^{-1}r and <rq,r>
        prTimer.stop();

        // determine new search direction
	terminate.residual(ScalarTraits<field_type>::real(sigmaNew));
        if (verbose>1)             // print
          this->printOutput(std::cout,static_cast<double>(i),std::abs(sigmaNew),std::abs(sigma));

	// convergence check to prevent division by zero if we obtained an exact solution
	// (which may happen for low dimensional systems)
	if (std::abs(sigmaNew) < 1e-28*sigma0)
	  break;

        beta = sigmaNew/sigma;
        sigma = sigmaNew;
        apTimer.resume();
        rq->axpy(beta,*q); // compute q = rq + beta*q
        std::swap(rq,q);
        apTimer.stop();
      }

      if (verbose>0)                // printing for non verbose
        this->printOutput(std::cout,static_cast<double>(i),std::abs(sigma));

      res.iterations = i;               // fill statistics
      res.reduction = std::sqrt(std::abs(sigma)/sigma0);
      res.conv_rate  = pow(res.reduction,1.0/i);
      res.elapsed = watch.elapsed();

      if (verbose>0)                 // final print
      {
        std::cout << "=== rate=" << res.conv_rate
                  << ", T=" << res.elapsed
                  << ", TIT=" << res.elapsed/i
                  << ", IT=" << i << std::endl;

        std::cout << "Matrix-Vector time: " << mvTimer.format()
                  << "vector ops        : " << apTimer.format()
                  << "preconditioner    : " << prTimer.format() << '\n';
      }
  }

    /**
     * \brief Apply inverse operator with given tolerance.
     *
     * This method is equivalent to setting first the tolerance in the termination criterion, then calling apply().
     */
  virtual void apply (X& x, X& b, double tol, Dune::InverseOperatorResult& res) {
    terminate.tolerance(tol);
    (*this).apply(x,b,res);
  }

  private:
    // The solver may be constructed with either symmetric operator/preconditioner interface provided,
    // or with Dune interfaces provided. The implementation relies on the symmetric interface
    // in order to exploit performance gains with simultaneous evaluation of matrix-vector products
    // and dual pairings.
    // Therefore we hold "working" references to symmetric operator/preconditioners. Those reference
    // either the provided ones, or wrapper objects relaying the evaluation to provided Dune interfaces.
    // For the second case, we hold the (small) wrapper objects ourselves. In the first case, those
    // wrapper objects have to be initialized correctly (even though they are never accessed). For this
    // we need a dual pairing, hence the dummy zeroDp.
    ZeroDualPairing<X,Xstar>                zeroDp;  // just a dummy

    SymmetricLinearOperatorWrapper<X,Xstar> proxyOp; // proxy in case a Dune::LinearOperator is provided
    SymmetricLinearOperator<X,Xstar>& op;  // working reference to operator

    SymmetricPreconditionerWrapper<X,Xstar>  proxyPrec; // proxy in case a Dune::Preconditioner is provided
    SymmetricPreconditioner<X,Xstar>&        prec;      // working reference to preconditioner

    PCGTerminationCriterion<Real>& terminate;
    int verbose;
    bool requiresInit;

    std::vector<double> resDebug;
  };
} // namespace Kaskade
#endif


/**
 * @file origin.cpp
 * @brief Parse origin from xml string
 *
 * @author Levi Armstrong
 * @date September 1, 2019
 * @version TODO
 * @bug No known bugs
 *
 * @copyright Copyright (c) 2019, Southwest Research Institute
 *
 * @par License
 * Software License Agreement (Apache License)
 * @par
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 * @par
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <tesseract_common/macros.h>
TESSERACT_COMMON_IGNORE_WARNINGS_PUSH
#include <stdexcept>
#include <tesseract_common/utils.h>
#include <Eigen/Geometry>
#include <vector>
#include <boost/algorithm/string.hpp>
#include <tinyxml2.h>
TESSERACT_COMMON_IGNORE_WARNINGS_POP

#include <tesseract_urdf/origin.h>

Eigen::Isometry3d tesseract_urdf::parseOrigin(const tinyxml2::XMLElement* xml_element, int /*version*/)
{
  Eigen::Isometry3d origin = Eigen::Isometry3d::Identity();

  if (xml_element->Attribute("xyz") == nullptr && xml_element->Attribute("rpy") == nullptr &&
      xml_element->Attribute("wxyz") == nullptr)
    std::throw_with_nested(std::runtime_error("Origin: Error missing required attributes 'xyz' and 'rpy' and/or 'wxyz' "
                                              "for origin element!"));

  std::string xyz_string, rpy_string, wxyz_string;
  tinyxml2::XMLError status = tesseract_common::QueryStringAttribute(xml_element, "xyz", xyz_string);
  if (status != tinyxml2::XML_NO_ATTRIBUTE && status != tinyxml2::XML_SUCCESS)
    std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'xyz'!"));

  if (status != tinyxml2::XML_NO_ATTRIBUTE)
  {
    std::vector<std::string> tokens;
    boost::split(tokens, xyz_string, boost::is_any_of(" "), boost::token_compress_on);
    if (tokens.size() != 3 || !tesseract_common::isNumeric(tokens))
      std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'xyz' string!"));

    double x{ 0 }, y{ 0 }, z{ 0 };
    // No need to check return values because the tokens are verified above
    tesseract_common::toNumeric<double>(tokens[0], x);
    tesseract_common::toNumeric<double>(tokens[1], y);
    tesseract_common::toNumeric<double>(tokens[2], z);

    origin.translation() = Eigen::Vector3d(x, y, z);
  }

  if (xml_element->Attribute("wxyz") == nullptr)
  {
    status = tesseract_common::QueryStringAttribute(xml_element, "rpy", rpy_string);
    if (status != tinyxml2::XML_NO_ATTRIBUTE && status != tinyxml2::XML_SUCCESS)
      std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'rpy'!"));

    if (status != tinyxml2::XML_NO_ATTRIBUTE)
    {
      std::vector<std::string> tokens;
      boost::split(tokens, rpy_string, boost::is_any_of(" "), boost::token_compress_on);
      if (tokens.size() != 3 || !tesseract_common::isNumeric(tokens))
        std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'rpy' string!"));

      double r{ 0 }, p{ 0 }, y{ 0 };
      // No need to check return values because the tokens are verified above
      tesseract_common::toNumeric<double>(tokens[0], r);
      tesseract_common::toNumeric<double>(tokens[1], p);
      tesseract_common::toNumeric<double>(tokens[2], y);

      Eigen::AngleAxisd rollAngle(r, Eigen::Vector3d::UnitX());
      Eigen::AngleAxisd pitchAngle(p, Eigen::Vector3d::UnitY());
      Eigen::AngleAxisd yawAngle(y, Eigen::Vector3d::UnitZ());

      Eigen::Quaterniond rpy = yawAngle * pitchAngle * rollAngle;

      origin.linear() = rpy.toRotationMatrix();
    }
  }
  else
  {
    status = tesseract_common::QueryStringAttribute(xml_element, "wxyz", wxyz_string);
    if (status != tinyxml2::XML_NO_ATTRIBUTE && status != tinyxml2::XML_SUCCESS)
      std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'wxyz'!"));

    if (status != tinyxml2::XML_NO_ATTRIBUTE)
    {
      std::vector<std::string> tokens;
      boost::split(tokens, wxyz_string, boost::is_any_of(" "), boost::token_compress_on);
      if (tokens.size() != 4 || !tesseract_common::isNumeric(tokens))
        std::throw_with_nested(std::runtime_error("Origin: Failed to parse attribute 'wxyz' string!"));

      double qw{ 0 }, qx{ 0 }, qy{ 0 }, qz{ 0 };
      // No need to check return values because the tokens are verified above
      tesseract_common::toNumeric<double>(tokens[0], qw);
      tesseract_common::toNumeric<double>(tokens[1], qx);
      tesseract_common::toNumeric<double>(tokens[2], qy);
      tesseract_common::toNumeric<double>(tokens[3], qz);

      Eigen::Quaterniond q(qw, qx, qy, qz);
      q.normalize();

      origin.linear() = q.toRotationMatrix();
    }
  }
  return origin;
}


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

#ifndef COARSENING_HH
#define COARSENING_HH

#include <boost/timer/timer.hpp>
#include <boost/type_traits/remove_pointer.hpp>
#include <boost/type_traits/remove_reference.hpp>

#include <boost/fusion/algorithm.hpp>
#include <boost/fusion/sequence.hpp>

#include "fem/errorest.hh"
#include "fem/embedded_errorest.hh"
#include "fem/fetransfer.hh"
#include "fem/firstless.hh"
#include "fem/fixfusion.hh"
#include "fem/iterate_grid.hh"
#include "utilities/power.hh"

#include "linalg/dynamicMatrix.hh"

namespace Kaskade
{
  namespace CoarseningDetail
  {

    using namespace boost::fusion;

    // A class representing a local projector for a given FE space
    template <class Sp>
    struct Projector
    {
      typedef Sp Space;

      Space const*                                                  space;  // pointer to FE space
      std::vector<size_t>                                           gidx;   // global indices
      DynamicMatrix<Dune::FieldMatrix<typename Space::Scalar,1,1>>  q;      // stores a factor of the projection matrix Q Q^T.

      // these are declared here to be used in GetLocalTransferProjection - avoid frequent reallocations
      DynamicMatrix<Dune::FieldMatrix<typename Space::Scalar,1,1>> pLocal;
      std::vector<int>                                             idx;
    };

    // A functor creating local projectors for a given FE space
    struct CreateProjection
    {
      template <class T> struct result {};

      template <class SpacePtr>
      struct result<CreateProjection(SpacePtr)> {
        typedef typename boost::remove_pointer<typename boost::remove_reference<SpacePtr>::type>::type Space;
        typedef Projector<Space> type;
      };

      template <class SpacePtr>
      typename result<CreateProjection(SpacePtr)>::type operator()(SpacePtr space) const
      {
        typename result<CreateProjection(SpacePtr)>::type p;
        p.space = space;
        return p;
      }
    };

    // A functor filling a local projector with actual data for a given cell
    template <class CellPointer>
    struct GetLocalTransferProjection
    {
      // father points to the father cell,
      // children contains pointer to its children
      GetLocalTransferProjection(CellPointer father_, std::vector<CellPointer> const& children_): father(father_), children(children_)
      { }

      template <class Projector>
      void operator()(Projector& p) const
      {
	typedef typename Projector::Space Space;
	// We could use LocalTransfer here, but this is conceptual and computational overkill. It supports different FE spaces (we need only one here)
	// and creates new shape function sets just in case the cell vanishes on coarsening. We do the projection here before coarsening...
	// Moreover we know here that all children of the father cell are leafs.

        // Following a direct implementation exploiting our a priori knowledge for simplicity and performance.
        // Obtain all global indices associated to leaf cells within this father cell
        p.gidx.clear();
        for (int i=0; i<children.size(); ++i)
        {
          auto gi = p.space->mapper().globalIndices(*children[i]);
          p.gidx.insert(p.gidx.end(),gi.begin(),gi.end());
        }
        // global indices can be duplicate - make unique
        std::sort(p.gidx.begin(),p.gidx.end());
        p.gidx.erase(std::unique(p.gidx.begin(),p.gidx.end()),p.gidx.end());

        // Step through all children and obtain the prolongation matrix. Scatter this to the
        // prolongation matrix P for all children's degrees of freedom.
        p.q.setSize(p.gidx.size(),p.space->mapper().shapefunctions(*father).size());
//        p.q = 0 // causes compilation error with dune-2.4.0 and clang++ on OS X (Darwin)
        p.q.fill(0);
        typedef typename Space::Grid Grid;
        typedef typename Space::Scalar   Scalar;
        int const sfComponents = Space::sfComponents;
        DynamicMatrix<Dune::FieldMatrix<Scalar,sfComponents,1>> afValues; // declare out of loop to prevent frequent reallocations
        std::vector<Dune::FieldVector<typename Grid::ctype,Grid::dimension> > iNodes; // declare out of loop to prevent frequent reallocations
        for (int i=0; i<children.size(); ++i)
        {
          // Compute the mapping of global indices on the child to the set of global indices on all children.
          auto gi = p.space->mapper().globalIndices(*children[i]);
          p.idx.resize(gi.size());
          for (int j=0; j<gi.size(); ++j)
            p.idx[j] = std::lower_bound(p.gidx.begin(),p.gidx.end(),gi[j]) - p.gidx.begin();

          // Compute local prolongation matrix p.pLocal. As father and child spaces are the same, we can use the same shape function set
          // (assuming that the type of cell doesn't change...).
          auto const& sfs = p.space->mapper().shapefunctions(*children[i]);
          localProlongationMatrix(*p.space,children[i],*p.space,father,p.pLocal,sfs,sfs);

          // Enter prolongation values on this child into the whole prolongation matrix for all children
          for(size_t k=0; k<p.idx.size(); ++k)
            for(int j=0; j<p.q.M(); ++j)
              p.q[p.idx[k]][j] = p.pLocal[k][j];
        }

        // Now the projection onto a coarser space is P P^+, where P^+, the pseudoinverse of P, is the restriction matrix.
        // Note that this formulation of the (not uniquely defined restriction) is different from the interpolation-based
        // restriction as obtained from LocalTransfer. Nevertheless, both make sense and both appear to work very similar.
        // In fact one might argue that the pseudoinverse formulation here makes even more sense, as it leads to a
        // least-squares approximation of the fine grid solution by the coarse grid restriction. "Least squares" is, however,
        // a norm-dependent concept, and the euclidean norm of FE coefficient vectors as implicitly used here might not
        // be the best suited.
        // With a QR decomposition of P = Q [R;0], the pseudoinverse is P^+ = [R^{-1} 0] Q^T and hence the projection
        // P P^+ = Q [I 0; 0 0] Q^T. If Q = [Q1 Q2], then P P^+ = Q1 Q1^T holds. The image space of Q1 is the same as that
        // of P. For computing Q1, we therefore do not need to compute a full QR decomposition, but only the first columns of
        // Q. The most simple method is Gram-Schmidt. Though it is numerically instable if the columns of P are almost linear
        // dependent, we use it here, as the columns of P are associated to the shape functions on the father cell and
        // should be comfortably linearly independent, though not orthogonal.
        // Instead of forming Q1 Q1^T, we store just Q1, and apply the product if needed. As Q1 is high and tall, this is
        // (slightly) more efficient.
        for (int j=0; j<p.q.M(); ++j) // step through all columns
        {
          Scalar tmp = 0;

          // Normalize column j.
          for (int i=0; i<p.q.N(); ++i)
            tmp += power<2>(p.q[i][j]);
          tmp = std::sqrt(tmp);
          for (int i=0; i<p.q.N(); ++i)
            p.q[i][j] /= tmp;

          // Orthogonalize remaining columns.
          for (int k=j+1; k<p.q.M(); ++k)
          {
            tmp = 0;
            for (int i=0; i<p.q.N(); ++i)
              tmp += p.q[i][j] * p.q[i][k];
            for (int i=0; i<p.q.N(); ++i)
              p.q[i][k] -= tmp*p.q[i][j];
          }
        }
      }

    private:
      CellPointer father;
      std::vector<CellPointer> const& children;
    };

    // A functor that applies the appropriate projector provided on construction to the
    // FE function that is given as argument.
    template <class Projectors>
    struct ProjectCoefficients
    {
      // Construct the functor, giving a list of (local) projectors indexed by the FE space index.
      ProjectCoefficients(Projectors const& projectors_): projectors(projectors_) {}

      // Apply the appropriate projector to the provided FE function. The pair consists of a variable description
      // (containing the space index) and the FE function itself.
      template <class Pair>
      void operator()(Pair const& pair) const
      {
        typedef typename boost::remove_reference<typename result_of::value_at_c<Pair,0>::type>::type VarDesc;
        typedef typename boost::remove_reference<typename result_of::value_at_c<Pair,1>::type>::type Function;

        int const sIdx = VarDesc::spaceIndex;
        auto const& proj = at_c<sIdx>(projectors);
        int const n = proj.gidx.size();

        // Multiplication of projector Q Q^T and coefficient vector. The
        // multiplication has to be done manually because of
        // type-incompatible entries. This results in the coefficients of
        // the projection residual.
        auto const& q = proj.q;
        int const m = q.M();
        typename Function::StorageType x(n);
        typename Function::StorageType y(m);

        // y <- Q^T v
        for (int i=0; i<m; ++i) {
          y[i] = 0;
          for (int j=0; j<n; ++j)
            y[i].axpy(q[j][i][0][0],(at_c<1>(pair).coefficients())[proj.gidx[j]]);
        }

        // x <- Q y
        for (int i=0; i<n; ++i) {
          typename Function::StorageValueType x(0);
          for (int j=0; j<m; ++j)
            x.axpy(q[i][j][0][0],y[j]);
          (at_c<1>(pair).coefficients())[proj.gidx[i]] = x;
        }
      }

    private:
      Projectors const& projectors;
    };

    // Convenience function for template type deduction.
    template <class Projectors>
    ProjectCoefficients<Projectors> getCoefficientProjectors(Projectors const& ps) { return ProjectCoefficients<Projectors>(ps); }

    class GroupByCell
    {
    public:
      GroupByCell(std::vector<size_t> const& groups_)
      : groups(&groups_), nGroups(*std::max_element(groups_.begin(),groups_.end())+1) {}

      size_t operator[](size_t idx) const { return (*groups)[idx]; }
      size_t nGroups;

    private:
      std::vector<size_t> const* groups;
    };

  } // End of namespace CoarseningDetail

  /** \ingroup adapt
   * In order to maintain compatibility with existing code,
   * this overload is needed, as non const reference parameters (coarsenedCells) can not
   * have default values. All it does is call the original method with a temporary
   * parameter.
   */
  template <class VariableSetDescription, class Scaling>
  void coarsening(VariableSetDescription const& varDesc,
                  typename VariableSetDescription::VariableSet const& sol,
                  Scaling const& scaling,
                  std::vector<std::pair<double,double> > const& tol,
                  GridManager<typename VariableSetDescription::Grid>& gridManager, int verbosity=1, int minRefLevel=0)
  {
    std::vector<bool> tmp;
    coarsening(varDesc, sol, scaling, tol, gridManager, tmp, verbosity, minRefLevel);
  }

  /** \ingroup adapt
   * \brief coarsening routine
   *
   * Perform a projection from a fine grid onto a coarse grid. If the error is small, locally
   * coarsen the grid there.
   *
   * \param varDesc
   * \param sol
   * \param scaling
   * \param tol
   * \param gridManager
   * \param coarsenedCells
   * \param verbosity
   * \param minRefLevel do not coarse cells that are on this level or below.
   */
  template <class VariableSetDescription, class Scaling>
  void coarsening(VariableSetDescription const& varDesc,
                  typename VariableSetDescription::VariableSet const& sol,
                  Scaling const& scaling,
                  std::vector<std::pair<double,double> > const& tol,
                  GridManager<typename VariableSetDescription::Grid>& gridManager,
                  std::vector<bool>& coarsenedCells, int verbosity=1, int minRefLevel=0)
  {
    using namespace boost::fusion;
    using namespace CoarseningDetail;

// boost::timer::cpu_timer timer;
// boost::timer::cpu_timer timerA, timerB;
// timerA.stop(); timerB.stop();

    typedef typename VariableSetDescription::Grid           Grid;
    typedef typename Grid::LocalIdSet    IdSet;
    typedef typename Grid::LeafIndexSet  IndexSet;
    typedef typename Grid::template Codim<0>::Entity        Cell;
    typedef typename Grid::template Codim<0>::EntityPointer CellPointer;
    typedef typename Cell::HierarchicIterator               HierarchicIterator;

    IndexSet const& indexSet = varDesc.indexSet;

    // A set of FE functions that will contain the projection.
    typename VariableSetDescription::VariableSet pSol(sol);

    // A container of local projectors - defined here to avoid multiple allocations inside the cell loop
    auto projectors = as_vector(transform(varDesc.spaces,CreateProjection()));

    // a container for caching computed cell pointers to children. Prevents doubled hierarchical grid traversal
    // and frequent reallocation
    std::vector<CellPointer> children;

    // A mapping from leaf cells to their coarsening group number. A
    // coarsening group is defined as follows. If all the direct
    // children of a non-leaf cell are leafs (i.e. none of the children
    // has been refined), these children form a coarsening group. Either
    // all or none cells of a coarsening group are marked for
    // coarsening. The special group number 0 collects all cells which
    // are not element of a regular coarsening group (and hence cannot be
    // removed during refinement).
    std::vector<size_t> coarseningGroup(indexSet.size(0),0);
    size_t coarseningGroupCount = 0;

    // Step through all leaf cells and process their fathers.
    typedef typename Grid::template Codim<0>::template Partition<Dune::All_Partition>::LeafIterator CellIterator;
    CellIterator end = gridManager.grid().template leafend<0>();
    for (CellIterator ci=gridManager.grid().template leafbegin<0>(); ci!=end; ++ci) // TODO: parallelize this loop
      // Father has already been processed if cell belongs to a regular
      // cell group. No father exists if the level is 0.
      if (coarseningGroup[indexSet.index(*ci)]==0 && ci->level()>0) { // not yet processed, not on coarse grid
        // Only process fathers which are suitable for coarsening.
        CellPointer father = ci->father();
        bool canBeCoarsened = true;
        children.clear();
        HierarchicIterator first = father->hbegin(father->level()+1), last = father->hend(father->level()+1);
        for (HierarchicIterator hi=first; canBeCoarsened && hi!=last; ++hi) // early termination in case we find a non-leaf child
        {
          canBeCoarsened = canBeCoarsened && hi->isLeaf();
          children.push_back(CellPointer(hi));
        }

        if (canBeCoarsened) {
// timerA.resume();
          // Gather all the children into a new coarsening group.
          ++coarseningGroupCount;
          for (auto const& c: children)
            coarseningGroup[indexSet.index(*c)] = coarseningGroupCount; // not coarseningGroupCount-1: group number 0 is for not coarsenable cells
          // Compute the projections to the locally coarser subspace.
          for_each(projectors,GetLocalTransferProjection<CellPointer>(father,children));
// timerA.stop();
// timerB.resume();

          // Apply the projectors to the FE functions.
          for_each2(typename VariableSetDescription::Variables(),pSol.data,getCoefficientProjectors(projectors));
// timerB.stop();
        }
      }

    // Create the projection residual (i.e. the error introduced by coarsening).
    pSol -= sol;

//     std::cerr << "Coarsening part Ia time: " << timer.format()
//               << "       A               : " << timerA.format()
//               << "       B               : " << timerB.format();
//     timer.start();

    // Compute scaled L2 norms of the function and difference,
    // attributing the cell contributions to its coarsening group.
    ErrorestDetail::GroupedSummationCollector<GroupByCell> sum{GroupByCell(coarseningGroup)};
    scaledTwoNormSquared(join(typename VariableSetDescription::Variables(),typename VariableSetDescription::Variables()),
                         join(pSol.data,sol.data),varDesc.spaces,scaling,sum);
//     std::cerr << "Coarsening part Ib time: " << timer.format();
//     timer.start();

    // For each variable, compute the (squared) relative error
    // contribution, that is e_i^2 = |p_i|^2/(atol^2+|f|^2*rtol^2),
    // where f is the function and p the difference to its hierarchic
    // projection.
    std::vector<double> norm2(varDesc.noOfVariables,0);
    for (int i=0; i<varDesc.noOfVariables; ++i)
      for (int j=0; j<=coarseningGroupCount; ++j)
        norm2[i] += sum.sums[j][i+varDesc.noOfVariables];
    if ( verbosity>0 )
    {
      std::cout << "coarsening solution norms2: ";
      std::copy(norm2.begin(),norm2.end(),
      std::ostream_iterator<double>(std::cout," ")); std::cout  << '\n';
    };

    for (int i=0; i<varDesc.noOfVariables; ++i) {
      double toli = power(tol[i].first,2) + norm2[i]*power(tol[i].second,2);
      for (int j=0; j<=coarseningGroupCount; ++j)
        sum.sums[j][i] /= toli;
    }

    // Select a maximal subset C of the coarsening groups, such that for
    // each variable the sum of C's relative errors is below 1.
    //
    // A greedy algorithm may work as follows: Maintain a vector E_i (of
    // size the number of variables) that contains the remaining allowed
    // relative error. Of the remaining coarsening groups j with
    // relative errors e_ji select the one for which max_i e_ji / E_i is
    // minimal.
    //
    // However, this has quadratic complexity. Thus we resort to the
    // following simplification.  First we assign to each coarsening
    // group j the maximal relative error e_j = max_i e_ji encountered
    // in any variable. Subsequently, these are sorted in ascending
    // order. Then the first k groups for which sum_{j=0}^k e_j <= 1
    // hold are selected. Non-coarsenable cells (in the irregular
    // coarsening group 0) are ignored.
    //
    // WARNING: Despite the fact that in principle we can compute the
    // projection error exactly, the value we obtain is just an
    // estimate. This is not only due to the heuristic suboptimal
    // selection of C, but also due to the fact that the coarsening
    // errors are modeled strictly locally. In conforming meshes, mesh
    // adaptation can lead to nonlocal effects, e.g. the removal of no
    // longer needed green closures (the introduced error is not taken
    // into account), the introduction of green closures on newly
    // coarsened cells (in which case the error could be below our
    // estimate, depending on the actual transfer - currently the
    // transfer is suboptimal and our error estimate should be exact),
    // or a cell group can actually be retained for mesh topology
    // reasons even if marked for coarsening (in which case the actual
    // local error would be zero).
    std::vector<std::pair<double,size_t> > e(coarseningGroupCount);
    for (int j=0; j<coarseningGroupCount; ++j) {
      e[j].second = j+1;                            // omit the irregular coarsening group 0
      for (int i=0; i<varDesc.noOfVariables; ++i)
        e[j].first = std::max(e[j].first,sum.sums[j+1][i]);
    }

    std::sort(e.begin(),e.end(),FirstLess());

    std::vector<size_t> selectedCoarseningGroups;
    double totalError = 0;
    for (int i=0; i<e.size(); ++i) {
      totalError += e[i].first;
      if (totalError<=1)
        selectedCoarseningGroups.push_back(e[i].second);
      else
        break;
    }

    // for compression: need to keep track of coarsening history
    coarsenedCells.clear() ;
    coarsenedCells.resize( indexSet.size(0), false ) ;
    // no coarsening: just stop here
    //   return ;

    // If no cell can be coarsened, we don't need to touch the grid.
    if (selectedCoarseningGroups.empty())
      return;

    // Now we sweep over the whole grid and mark all cells for
    // coarsening which belong to a selected coarsening group.
    std::sort(selectedCoarseningGroups.begin(),selectedCoarseningGroups.end());
    for (CellIterator ci=gridManager.grid().template leafbegin<0>(); ci!=gridManager.grid().template leafend<0>(); ++ci)
    {
      auto idx = indexSet.index(*ci);
      if (ci->level()>minRefLevel && std::binary_search(selectedCoarseningGroups.begin(),selectedCoarseningGroups.end(),
                                                        coarseningGroup[idx])) // for uniform refinement: just delete level
      {
        coarsenedCells[idx] = true ; // for compression: coarsening history
        gridManager.mark(-1,*ci);
      }
    }

//     std::cerr << "Coarsening part II time: " << timer.format();
//     timer.start();

    gridManager.adaptAtOnce();

//     std::cerr << "Coarsening part III time: " << timer.format();
  }
} /* end of namespace Kaskade */

#endif


#include <cstdio>
#include <cstdlib>
#include <cmath>

#include <boost/numeric/odeint.hpp>
#include <boost/numeric/odeint/stepper/dense_output_runge_kutta.hpp>
#include <boost/numeric/odeint/stepper/runge_kutta_dopri5.hpp>
#include "ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver.h"

#include "Utilities.hpp"

using namespace boost::numeric::odeint;

typedef boost::numeric::ublas::vector<double> vector_type;
typedef vector_type state_type;
typedef result_of::make_dense_output<
    runge_kutta_dopri5< state_type > >::type dense_dopri5_stepper_type;

//typedef int (*CALLBACK_F_PTR)(double t, N_Vector y, N_Vector ydot, void* user_data);
//typedef int (*CALLBACK_G_PTR)(double t, N_Vector y, double* gout, void* user_data);

struct ode_system;

struct boost_solver_data {
//	dense_output_runge_kutta< runge_kutta_dopri5< state_type > > stepper;
//	runge_kutta_dopri5< state_type > stepper;
	ode_system* system;
	dense_dopri5_stepper_type stepper;
    JNIEnv* env;
    jobject ode;
    jmethodID vectorFieldMethodID;
    jdouble stepSize;
    jdouble t;
    jdouble t1;
    jdouble* x;
    jdouble* xTmp;
    jdouble* xDot;
    state_type y;
    state_type yDot;
};

struct ode_system {

    boost_solver_data* data;

	void operator()(const state_type &x, state_type &dxdt, double t )
    {
	    JNIEnv *env = data->env;
//	    for (state_type::const_iterator it=x.begin(); it != x.end(); ++it) {
		for (int i=0; i < x.size(); i++)
	    	data->xTmp[i] = x[i];
	    env->CallVoidMethod(data->ode, data->vectorFieldMethodID, t);
	    if (env->ExceptionCheck() == JNI_TRUE) {
	        fprintf(stderr, "f_direct_bridge: Failed to call computeVectorField!\n");
	        return;
	    }
//	    for (state_type::const_iterator it=x.begin(); it != x.end(); ++it)
		for (int i=0; i < dxdt.size(); i++)
	    	dxdt[i] = data->xDot[i];
    }

};

struct ode_system_wrapper {

	ode_system* system;

	ode_system_wrapper(ode_system* system)
		: system(system) { }
	void operator()(const state_type &x, state_type &dxdt, double t ) {
		system->operator ()(x, dxdt, t);
	}

};

/*int jac_bridge(long int N, double t, N Vector y, N Vector fy, DlsMat Jac,
               void *user data, N Vector tmp1, N Vector tmp2, N Vector tmp3) {
    return 0;
}*/

JNIEXPORT jlong JNICALL Java_ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver_jni_1initialize
	(JNIEnv *env, jobject obj, jobject ode, jobject xBuffer, jobject xTmpBuffer, jobject xDotBuffer,
			jdouble stepSize, jdouble relTol, jdouble absTol, jint stepperType) {
    jclass odeCls = env->GetObjectClass(ode);
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to get class of ode object\n");
        return 0;
    }
    jmethodID methodID = env->GetMethodID(odeCls, "getDimensionOfVectorField", "()I");
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to get method ID of getDimensionOfVectorField\n");
        return 0;
    }
    jint jdimension = env->CallIntMethod(ode, methodID);
    state_type::size_type dimension = jdimension;
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to call getDimensionOfVectorField\n");
        return 0;
    }

    // Allocate and initialize boost_solver_data
    boost_solver_data* data = new boost_solver_data();
    // FIXME
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to allocate boost_solver_data\n");
        return 0;
    }
    data->stepSize = stepSize;

    data->env = env;
    // Acquire method IDs of callbacks and put them into boost_solver_data
    data->ode = env->NewGlobalRef(ode);
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to create global reference to ode object\n");
        delete data;
        return 0;
    }
    data->vectorFieldMethodID = env->GetMethodID(odeCls, "computeVectorField", "(D)V");
    if (env->ExceptionCheck() == JNI_TRUE) {
        fprintf(stderr, "jni_initialize: Failed to get method ID of computeVectorField\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }

    // Check capacities of direct buffer objects
    if (env->GetDirectBufferCapacity(xBuffer) < dimension) {
        throw_java_exception(env, "jni_initialize: Direct buffer xBuffer is not big enough\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }
    if (env->GetDirectBufferCapacity(xTmpBuffer) < dimension) {
        throw_java_exception(env, "jni_initialize: Direct buffer xTmpBuffer is not big enough\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }
    if (env->GetDirectBufferCapacity(xDotBuffer) < dimension) {
        throw_java_exception(env, "jni_initialize: Direct buffer xDotBuffer is not big enough\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }
    // Get addresses of direct buffer objects
    data->x = static_cast<jdouble*>(env->GetDirectBufferAddress(xBuffer));
    if (data->x == NULL) {
        fprintf(stderr, "jni_initialize: Failed to get direct address of xBuffer\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }
    data->xTmp = static_cast<jdouble*>(env->GetDirectBufferAddress(xTmpBuffer));
    if (data->xTmp == NULL) {
        fprintf(stderr, "jni_initialize: Failed to get direct address of xTmpBuffer\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }
    data->xDot = static_cast<jdouble*>(env->GetDirectBufferAddress(xDotBuffer));
    if (data->xDot == NULL) {
        fprintf(stderr, "jni_initialize: Failed to get direct address of xDotBuffer\n");
        // Delete global java references
        env->DeleteGlobalRef(data->ode);
        delete data;
        return 0;
    }

    data->y.resize(dimension, false);
    data->yDot.resize(dimension, false);

    data->system = new ode_system;
    data->system->data = data;

    // Pass pointer-address of boost_solver_data back to Java
    jlong jni_pointer = (jlong)data;
    return jni_pointer;
}

JNIEXPORT void JNICALL Java_ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver_jni_1dispose
  (JNIEnv *env, jobject obj, jlong jni_pointer) {
    boost_solver_data* data = (boost_solver_data*)jni_pointer;
    if (data != NULL) {
        // Delete global java references
        if (data->ode != NULL)
            env->DeleteGlobalRef(data->ode);
        // Free data structure
        delete data;
    }
}

JNIEXPORT void JNICALL Java_ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver_jni_1prepareStep
	(JNIEnv *env, jobject obj, jlong jni_pointer, jdouble t0, jdouble t1) {
    boost_solver_data* data = reinterpret_cast< boost_solver_data* >(jni_pointer);
    data->t = t0;
    data->t1 = t1;
    for (int i=0; i < data->y.size(); i++)
    	data->y[i] = data->x[i];
    data->stepper.initialize(data->y, data->t, data->stepSize);
}

/*
 * Class:     ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver
 * Method:    jni_integrateOneStep
 * Signature: (J)D
 */
JNIEXPORT jdouble JNICALL Java_ch_ethz_bhepp_ode_boost_FixedBoostOdeintSolver_jni_1integrateOneStep
(JNIEnv *env, jobject obj, jlong jni_pointer) {
    boost_solver_data* data = reinterpret_cast< boost_solver_data* >(jni_pointer);
    std::pair< double , double > times = data->stepper.do_step(ode_system_wrapper(data->system));
    const dense_dopri5_stepper_type::state_type& current_state = data->stepper.current_state();
    for (int i=0; i < data->y.size(); i++)
    	data->x[i] = current_state[i];
    return times.second;
}


// __BEGIN_LICENSE__
// Copyright (C) 2006-2011 United States Government as represented by
// the Administrator of the National Aeronautics and Space Administration.
// All Rights Reserved.
// __END_LICENSE__

/// \file StereoSessionIsis.cc
///

// Vision Workbench
#include <vw/FileIO.h>
#include <vw/Math/Functors.h>
#include <vw/Math/Geometry.h>
#include <vw/Math/RANSAC.h>
#include <vw/InterestPoint.h>
#include <vw/Stereo/DisparityMap.h>
#include <vw/Cartography.h>

// Stereo Pipeline
#include <asp/Sessions/ISIS/StereoSessionIsis.h>
#include <asp/IsisIO/IsisCameraModel.h>
#include <asp/Core/StereoSettings.h>
#include <asp/IsisIO/IsisAdjustCameraModel.h>
#include <asp/IsisIO/DiskImageResourceIsis.h>
#include <asp/Sessions/ISIS/PhotometricOutlier.h>

// Boost
#include <boost/filesystem/operations.hpp>
#include <boost/shared_ptr.hpp>
namespace fs = boost::filesystem;

#include <algorithm>

using namespace vw;
using namespace vw::camera;
using namespace asp;

// Allows FileIO to correctly read/write these pixel types
namespace vw {
  template<> struct PixelFormatID<Vector3>   { static const PixelFormatEnum value = VW_PIXEL_GENERIC_3_CHANNEL; };
}

// Process a single ISIS image to find an ideal min max. The reason we
// need to do this, is that ASP is to get image intensity values in
// the range of 0-1. To some extent we are compressing the dynamic
// range, but we try to minimize that.
void find_ideal_isis_range( std::string const& in_file,
                            std::string const& tag,
                            float & isis_lo, float & isis_hi ) {

  boost::shared_ptr<DiskImageResourceIsis> isis_rsrc( new DiskImageResourceIsis(in_file) );
  DiskImageView<PixelGray<float> > disk_image(isis_rsrc);

  float isis_mean, isis_std;

  // Calculating statistics. We subsample the images so statistics
  // only does about a million samples.
  {
    vw_out(InfoMessage) << "\t--> Computing statistics for the "+tag+" image\n";
    int left_stat_scale = int(ceil(sqrt(float(disk_image.cols())*float(disk_image.rows()) / 1000000)));
    ChannelAccumulator<math::CDFAccumulator<float> > accumulator;
    for_each_pixel(
      subsample(create_mask( edge_extend(disk_image, ConstantEdgeExtension()),
                             isis_rsrc->valid_minimum(),
                             isis_rsrc->valid_maximum() ),
                left_stat_scale ),
      accumulator );
    isis_lo = accumulator.quantile(0);
    isis_hi = accumulator.quantile(1);
    isis_mean = accumulator.approximate_mean();
    isis_std  = accumulator.approximate_stddev();

    vw_out(InfoMessage) << "\t  "+tag+": [ lo:" << isis_lo << " hi:" << isis_hi
                        << " m: " << isis_mean << " s: " << isis_std <<  "]\n";
  }

  // Normalizing to -+2 sigmas around mean
  if ( stereo_settings().force_max_min == 0 ) {
    vw_out(InfoMessage) << "\t--> Adjusting hi and lo to -+2 sigmas around mean.\n";

    if ( isis_lo < isis_mean - 2*isis_std )
      isis_lo = isis_mean - 2*isis_std;
    if ( isis_hi > isis_mean + 2*isis_std )
      isis_hi = isis_mean + 2*isis_std;

    vw_out(InfoMessage) << "\t    "+tag+" changed: [ lo:"
                        << isis_lo << " hi:" << isis_hi << "]\n";
  }
}

// This actually modifies and writes the pre-processed image.
void write_preprocessed_isis_image( BaseOptions const& opt,
                                    std::string const& in_file,
                                    std::string const& out_file,
                                    std::string const& tag,
                                    float isis_lo, float isis_hi,
                                    float out_lo, float out_hi,
                                    Matrix<double> const& matrix,
                                    Vector2i const& crop_size ) {
  DiskImageView<PixelGray<float> > disk_image(in_file);
  ImageViewRef<PixelGray<float> > applied_image;
  if ( matrix == math::identity_matrix<3>() ) {
    applied_image =
      crop(edge_extend(clamp(normalize(remove_isis_special_pixels(disk_image,
                                                                  isis_lo, isis_hi, out_lo),
                           out_lo, out_hi, 0.0, 1.0)),
                       ZeroEdgeExtension()),
           0, 0, crop_size[0], crop_size[1]);
  } else {
    applied_image =
      transform(clamp(normalize(remove_isis_special_pixels(disk_image,
                                                           isis_lo, isis_hi,
                                                           out_lo),
                                out_lo, out_hi, 0.0, 1.0)),
                HomographyTransform(matrix),
                crop_size[0], crop_size[1]);
  }

  // Write the results to disk.
  vw_out() << "\t--> Writing normalized images.\n";
  block_write_gdal_image( out_file, applied_image, opt,
                          TerminalProgressCallback("asp", "\t  "+tag+":  ") );
}

void
asp::StereoSessionIsis::pre_preprocessing_hook(std::string const& input_file1,
                                               std::string const& input_file2,
                                               std::string & output_file1,
                                               std::string & output_file2) {
  output_file1 = m_out_prefix + "-L.tif";
  output_file2 = m_out_prefix + "-R.tif";

  if ( fs::exists(output_file1) && fs::exists(output_file2) ) {
    try {
      vw_log().console_log().rule_set().add_rule(-1,"fileio");
      DiskImageView<PixelGray<float32> > out1(output_file1);
      DiskImageView<PixelGray<float32> > out2(output_file2);
      vw_out(InfoMessage) << "\t--> Using cached normalized input images.\n";
      vw_settings().reload_config();
      return;
    } catch (vw::ArgumentErr const& e) {
      // This throws on a corrupted file.
      vw_settings().reload_config();
    } catch (vw::IOErr const& e) {
      vw_settings().reload_config();
    }
  }

  float left_lo, left_hi, right_lo, right_hi;
  find_ideal_isis_range( input_file1, "left",
                         left_lo, left_hi );
  find_ideal_isis_range( input_file2, "right",
                         right_lo, right_hi );

  // Working out alignment
  float lo = std::min (left_lo, right_lo);  // Finding global
  float hi = std::max (left_hi, right_hi);
  Matrix<double> align_matrix(3,3);
  align_matrix.set_identity();
  if ( stereo_settings().keypoint_alignment) {
    DiskImageView<PixelGray<float> > left_disk_image(input_file1);
    DiskImageView<PixelGray<float> > right_disk_image(input_file2);
    ImageViewRef<PixelGray<float> > left_view =
      normalize(remove_isis_special_pixels(left_disk_image, lo),
                lo, hi, 0, 1.0);
    ImageViewRef<PixelGray<float> > right_view =
      normalize(remove_isis_special_pixels(right_disk_image, lo),
                lo, hi, 0, 1.0);
    align_matrix = determine_image_align(input_file1, input_file2,
                                         left_view, right_view );
  }
  write_matrix( m_out_prefix + "-align.exr", align_matrix );

  // Getting left image size
  Vector2i left_size;
  {
    DiskImageView<PixelGray<float> > left_image(input_file1);
    left_size = Vector2i(left_image.cols(), left_image.rows());
  }

  // Apply alignment and normalization
  if (stereo_settings().individually_normalize == 0 ) {
    vw_out() << "\t--> Normalizing globally to: ["<<lo<<" "<<hi<<"]\n";
    write_preprocessed_isis_image( m_options, input_file1, output_file1, "left",
                                   left_lo, left_hi, lo, hi,
                                   math::identity_matrix<3>(), left_size );
    write_preprocessed_isis_image( m_options, input_file2, output_file2, "right",
                                   right_lo, right_hi, lo, hi,
                                   align_matrix, left_size );
  } else {
    vw_out() << "\t--> Individually normalizing.\n";
    write_preprocessed_isis_image( m_options, input_file1, output_file1, "left",
                                   left_lo, left_hi, left_lo, left_hi,
                                   math::identity_matrix<3>(), left_size );
    write_preprocessed_isis_image( m_options, input_file2, output_file2, "right",
                                   right_lo, right_hi, right_lo, right_hi,
                                   align_matrix, left_size );
  }
}

inline std::string write_shadow_mask( BaseOptions const& opt,
                                      std::string const& output_prefix,
                                      std::string const& input_image,
                                      std::string const& mask_postfix ) {
  // This thresholds at -25000 as the input sub4s for Apollo that I've
  // processed have a range somewhere between -32000 and +32000. -ZMM
  DiskImageView<PixelGray<float> > disk_image( input_image );
  DiskImageView<uint8> disk_mask( output_prefix + mask_postfix );
  ImageViewRef<uint8> mask =
    apply_mask(intersect_mask(create_mask(disk_mask),
                              create_mask(threshold(disk_image,-25000,0,1.0))));
  std::string output_mask =
    output_prefix+mask_postfix.substr(0,mask_postfix.size()-4)+"Debug.tif";

  block_write_gdal_image( output_mask, mask, opt,
                          TerminalProgressCallback("asp","\t  Shadow:") );
  return output_mask;
}

// Stage 2: Correlation
//
// Pre file is a pair of grayscale images.  ( ImageView<PixelGray<float> > )
// Post file is a disparity map.            ( ImageView<PixelMask<Vector2f> > )
void
asp::StereoSessionIsis::pre_filtering_hook(std::string const& input_file,
                                           std::string & output_file) {
  output_file = input_file;

  // ****************************************************
  // The following code is for Apollo Metric Camera ONLY!
  // (use at your own risk)
  // ****************************************************
  if (stereo_settings().mask_flatfield) {
    vw_out() << "\t--> Masking pixels that are less than 0.0.  (NOTE: Use this option with Apollo Metric Camera frames only!)\n";
    output_file = m_out_prefix + "-R-masked.exr";

    std::string shadowLmask_name =
      write_shadow_mask( m_options, m_out_prefix, m_left_image_file,
                         "-lMask.tif" );
    std::string shadowRmask_name =
      write_shadow_mask( m_options, m_out_prefix, m_right_image_file,
                         "-rMask.tif" );

    DiskImageView<uint8> shadowLmask( shadowLmask_name );
    DiskImageView<uint8> shadowRmask( shadowRmask_name );

    DiskImageView<PixelMask<Vector2f> > disparity_disk_image(input_file);
    ImageViewRef<PixelMask<Vector2f> > disparity_map =
      stereo::disparity_mask(disparity_disk_image,
                             shadowLmask, shadowRmask );

    DiskImageResourceOpenEXR disparity_map_rsrc(output_file, disparity_map.format() );
    Vector2i block_size(std::min<size_t>(vw_settings().default_tile_size(),
                                         disparity_map.cols()),
                        std::min<size_t>(vw_settings().default_tile_size(),
                                         disparity_map.rows()));
    disparity_map_rsrc.set_block_write_size(block_size);
    block_write_image( disparity_map_rsrc, disparity_map,
                       TerminalProgressCallback( "asp", "\t--> Saving Mask :") );
  }
}

// Reverse any pre-alignment that was done to the images.
void
asp::StereoSessionIsis::pre_pointcloud_hook(std::string const& input_file,
                                            std::string & output_file) {

  // ****************************************************
  // The following code is for Apollo Metric Camera ONLY!
  // (use at your own risk)
  // ****************************************************
  std::string dust_result = input_file;
  if ( stereo_settings().mask_flatfield ) {
    vw_out() << "\t--> Masking pixels that appear to be dust.  (NOTE: Use this option with Apollo Metric Camera frames only!)\n";
    photometric_outlier_rejection( m_out_prefix, input_file,
                                   dust_result, stereo_settings().h_kern );
  }

  DiskImageView<PixelMask<Vector2f> > disparity_map(dust_result);
  output_file = m_out_prefix + "-F-corrected.tif";

  // We used a homography to line up the images, we may want
  // to generate pre-alignment disparities before passing this information
  // onto the camera model in the next stage of the stereo pipeline.
  Matrix<double> align_matrix;
  try {
    read_matrix(align_matrix, m_out_prefix + "-align.exr");
    vw_out(DebugMessage) << "Alignment Matrix: " << align_matrix << "\n";
  } catch (vw::IOErr const& e) {
    vw_out() << "\nCould not read in aligment matrix: " << m_out_prefix
             << "-align.exr.  Exiting. \n\n";
    exit(1);
  }

  // Remove pixels that are outside the bounds of the secondary image.
  DiskImageView<PixelGray<float> > right_disk_image(m_right_image_file);
  ImageViewRef<PixelMask<Vector2f> > result =
    stereo::disparity_range_mask(stereo::transform_disparities(disparity_map,
                                          HomographyTransform(align_matrix)),
                                 Vector2f(0,0),
                                 Vector2f( right_disk_image.cols(),
                                           right_disk_image.rows() ) );

  block_write_gdal_image( output_file, result, m_options,
                          TerminalProgressCallback("asp", "\t    Processing:") );
}

boost::shared_ptr<vw::camera::CameraModel>
asp::StereoSessionIsis::camera_model(std::string const& image_file,
                                     std::string const& camera_file) {

  if (boost::ends_with(boost::to_lower_copy(camera_file), ".isis_adjust")){
    // Creating Equations for the files
    std::ifstream input( camera_file.c_str() );
    boost::shared_ptr<asp::BaseEquation> posF = read_equation( input );
    boost::shared_ptr<asp::BaseEquation> poseF = read_equation( input );
    input.close();

    // Finally creating camera model
    return boost::shared_ptr<camera::CameraModel>(new IsisAdjustCameraModel( image_file, posF, poseF ));

  } else {
    return boost::shared_ptr<camera::CameraModel>(new IsisCameraModel(image_file));
  }

}