/** This is free and unencumbered software released into the public domain.
The authors of ISIS do not claim copyright on the contents of this file.
For more details about the LICENSE terms and the AUTHORS, you will
find files of those names at the top level of this repository. **/

/* SPDX-License-Identifier: CC0-1.0 */
#include "InlineCalculator.h"

// std library
#include <cmath>

// Qt library
#include <QMap>
#include <QString>
#include <QStringList>
#include <QVariant>
#include <QVector>

// boost library
#include <boost/foreach.hpp>

// naif library
#include <SpiceUsr.h>
#include <SpiceZfc.h>
#include <SpiceZmc.h>

// other ISIS
#include "IException.h"
#include "InlineInfixToPostfix.h"
#include "IString.h"
#include "NaifStatus.h"

using namespace std;

namespace Isis {

  /**
   * Constructs an InlineCalculator object by initializing the operator lookup
   * list.
   */
  InlineCalculator::InlineCalculator() : Calculator() {
    initialize();
  }

  /**
   * Constructs an InlineCalculator object by initializing the operator lookup
   * list and compiles the given equation to prepare for evaluation.
   *
   * @param equation A string representing an equation in infix format.
   */
  InlineCalculator::InlineCalculator(const QString &equation) : Calculator() {
    initialize();
    compile(equation);
  }

  /**
   * Destroys the InlineCalculator object.
   */
  InlineCalculator::~InlineCalculator() {
    destruct();
  }

  /**
   * Accesses the number of functions, operators, variables, and scalars to be
   * executed.
   *
   * @return int The number of functions, operators, variables, and scalars
   *         stored in the Calculator.
   */
  int InlineCalculator::size() const {
    return (m_functions.size());
  }

  /**
   * Accesses the string representation of the current equation, in postfix
   * format. This equation is most commonly entered using infix, but it is
   * reformatted to postfix when set.
   *
   * @return QString A string representing the equation, in postfix.
   */
  QString InlineCalculator::equation() const {
    return (m_equation);
  }

  /**
   * @brief Compiles the given equation for evaluation.
   *
   * This method first converts the given infix equation into a postfix
   * equation for evaluation and saves the postfix formatted string. It then
   * ensures that the equation is ready for evaluation by parsing and
   * verifying that all tokens are recognized.
   *
   * @param equation A string representing an equation to be compiled, in
   *                 infix format.
   *
   * @return bool Indicates whether the compilation was successful.
   * @throw IException::User "Error parsing inline equation. Equation
   *        element invalid - token not recognized."
   *
   */
  bool InlineCalculator::compile(const QString &equation) {
    //  Transform equation to postfix order

    QString tokenOps = toPostfix(equation);
    tokenOps = tokenOps.simplified();

    int nerrors = 0;
    QString error =  "Errors parsing inline equation[" + equation + "].";
    IException errList(IException::User, error, _FILEINFO_);

    Clear();  // Clear the stack
    m_equation = equation;
    m_functions.clear();  // Clear function list

    QStringList tokenList = tokenOps.split(" ");
    while ( !tokenList.isEmpty() ) {
      QString token = tokenList.takeFirst();
      if ( !token.isEmpty() ) {

        // See if the function already exists.  Note that scalars and variables
        // that are already present can safely be reused.  New occuring ones are
        // created fresh!
        FxTypePtr fx = find(token);
        if ( 0 != fx ) {
          m_functions.push_back(fx);
        }
        // New scalars and variables will create new unique function objects if
        // they do not already exist (they would be found above then)
        else if ( isScalar(token)  ) {
          fx = addFunction(new ParameterFx(token, &InlineCalculator::scalar, this));
          m_functions.push_back(fx);
        }
        else if ( isVariable(token)  ) {
          // Will also get line, sample, band, etc...
          fx = addFunction(new ParameterFx(token, &InlineCalculator::variable, this));
          m_functions.push_back(fx);
        }
        else {
            //  Parameter not recognized during compile.  All unknown tokens are
            //  assumed to be variables until run time when they are searched for
            //  in the current state of the resource pool.
          try {
            if ( !orphanTokenHandler(token) ) {
              error = "Equation element (" + token + ") invalid - token not recognized.";
              errList.append(IException(IException::User, error, _FILEINFO_));
              nerrors++;
            }
          }
          // Catch all failures from orphaned tokens
          catch (IException &e) {
            errList.append(e);
            nerrors++;
          }
        }
      }
    }

    //  Might want to make this optional here
    if (nerrors > 0) {
      throw errList;
    }
    return (nerrors == 0);
  }

  /**
   * @brief Evaluate with a variable pool
   *
   * This method accepts a varible resource pool, stores it off for evaluation
   * and then executes the operands to provide the result of the equation.  It
   * stores the variable pool for subsequent runs that should yield identical
   * results when operator() is invoked.
   *
   * @param variablePool Keyword resource for variables substituted at run time
   *                     of the equation
   *
   * @return QVector \< double \> Result of the stored equation given the
   *         resource.
   * @throw IException::Programmer "Calculation with variable pool failed."
   *
   */
  QVector<double> InlineCalculator::evaluate(CalculatorVariablePool *variablePool) {
    QVector<double> value;
    try {
      pushVariables(variablePool);
      value = evaluate();
      popVariables();
    }
    catch (IException &ie) {
      popVariables();
      throw IException(ie, IException::Programmer,
                       "Calculation with variable pool failed.",
                       _FILEINFO_);
    }
    return (value);
  }

  /**
   * @brief Evaluate compiled equation with existing variable pool
   *
   * This executes the operands to provide the result of the equation.  The
   * currently stored variable pool, refernenced via variables(), is used to
   * provide the variables utilized in the equation.
   *
   * This is reentrant and as long as the variable pool is unchanged between
   * succesive calls, they all should return the same result.
   *
   * @return QVector \< double \> Result of the stored equation.
   * @throw IException::Unknown "Too many operands in the equation."
   *
   */
  QVector<double> InlineCalculator::evaluate() {

    BOOST_FOREACH (FxTypePtr function,  m_functions) {
      function->execute();
    }

    if (StackSize() != 1) {
      QString msg = "Too many operands in the equation [" + m_equation + "].";
      throw IException(IException::Unknown, msg, _FILEINFO_);
    }

    return (Pop(true));
  }

  /**
   * Converts the given string from infix to postfix format.
   *
   * @param equation A string representing an infix equation.
   *
   * @return QString A string representing the given equation converted to
   *         postfix format.
   */
  QString InlineCalculator::toPostfix(const QString &equation) const {
    InlineInfixToPostfix parser;
    return (parser.convert(equation));
  }

  /**
   * Determines whether the given string contains a scalar value (i.e. can be
   * converted to a double precision value).
   *
   * @param scalar A string to be checked.
   *
   * @return bool Indicates whether the given string is a scalar value.
   */
  bool InlineCalculator::isScalar(const QString &scalar) {
    if (scalar.isEmpty())  return (false);
    try {
      toDouble(scalar);
      return (true);
    }
    catch (IException &e) {
      return (false);
    }
  }

  /**
   * Determines whether the given string is a variable. If the string is empty
   * or scalar, this method returns false. Otherwise, it is assumed to be a
   * variable and this method returns true.
   *
   * @param str A string to be checked.
   * @return bool Indicates whether the given string is a variable.
   */
  bool InlineCalculator::isVariable(const QString &str) {
    if (str.isEmpty()) {
      return (false);
    }
    if (!isScalar(str)) {
      return (true);
    }
    return (false);
  }

  /**
   * @brief Pushes the given value onto the stack as a scalar.
   *
   * This method adds the given variant value to the stack as a scalar.
   *
   * @param scalar QVariant containing the scalar. Usually a QString.
   *
   */
  void InlineCalculator::scalar(const QVariant &scalar) {
    Push(toDouble(scalar.toString()));
  }

  /**
   * @brief Pushes the given value onto the stack as a variable.
   *
   * This method adds the given variant value to the stack as a variable.
   *
   * @param variable QVariant containing the name of the variable. Usually
   *                 a QString.
   *
   * @throw IException::User "Could not find variable in variable pool."
   */
  void InlineCalculator::variable(const QVariant &variable) {
    CalculatorVariablePool *variablePool = variables();
    QString key = variable.toString();
    if (variablePool->exists(key)) {
      QVector<double> values = variablePool->value(key);
      Push(values);
      return;
    }

    // Error!
    QString error = "Could not find variable [" + key + "] in variable pool.";
    throw IException(IException::User, error, _FILEINFO_);
  }

  /**
   * Determines the remainder of the quotient a/b whose sign is the same as that
   * of a. In other words, this method finds the value r = a - bq such that q is
   * the integer found by truncating a/b.
   *
   * @param a The dividend (numerator).
   * @param b The divisor (denominator).
   *
   * @return double The remainder of the quotient that has the same sign
   *         as the dividend.
   */
  double floatModulusOperator(double a, double b) {
    return (fmod(a, b));
  }

  /**
   * Pops the top two vectors off the current stack and performs the
   * floatModulusOperator() on the corresponding components of these vectors.
   * The result is then pushed back onto the stack.
   */
  void InlineCalculator::floatModulus() {
    QVector<double> y = Pop();
    QVector<double> x = Pop();
    QVector<double> result;
    PerformOperation(result, x.begin(), x.end(), y.begin(), y.end(),
                     floatModulusOperator);
    Push(result);
    return;
  }

  /**
   * Pops the top vector off the current stack and converts from degrees to
   * radians. The result is then pushed back onto the stack.
   */
  void InlineCalculator::radians() {
    QVector<double> degree = Pop();
    QVector<double> result;
    BOOST_FOREACH(double d, degree) {
      result.push_back(d * rpd_c());
    }
    Push(result);
    return;
  }

  /**
   * Pops the top vector off the current stack and converts from radians to
   * degrees. The result is then pushed back onto the stack.
   */
  void InlineCalculator::degrees() {
    QVector<double> radians = Pop();
    QVector<double> result;
    BOOST_FOREACH(double r, radians) {
      result.push_back(r * dpr_c());
    }
    Push(result);
    return;
  }

  /**
   * Pops the top two vectors off the current stack and performs a logical or
   * on each pair.
   *
   * @throws IException::Unknown "Failed performing logical or operation, "
   *                             "input vectors are of differnet lengths."
   */
  void InlineCalculator::logicalOr() {
    QVector<double> inputA = Pop();
    QVector<double> inputB = Pop();
    QVector<double> results;
    if ( inputA.size() != inputB.size() ) {
      QString msg = "Failed performing logical or operation, "
                    "input vectors are of differnet lengths.";
      throw IException(IException::Unknown, msg, _FILEINFO_);
    }
    for (int i = 0; i < inputA.size(); i++) {
      results.push_back( inputA[i] || inputB[i] );
    }
    Push(results);
    return;
  }

  /**
   * Pops the top two vectors off the current stack and performs a logical and
   * on each pair.
   *
   * @throws IException::Unknown "Failed performing logical and operation, "
   *                             "input vectors are of differnet lengths."
   */
  void InlineCalculator::logicalAnd() {
    QVector<double> inputA = Pop();
    QVector<double> inputB = Pop();
    QVector<double> results;
    if ( inputA.size() != inputB.size() ) {
      QString msg = "Failed performing logical and operation, "
                    "input vectors are of differnet lengths.";
      throw IException(IException::Unknown, msg, _FILEINFO_);
    }
    for (int i = 0; i < inputA.size(); i++) {
      results.push_back( inputA[i] && inputB[i] );
    }
    Push(results);
    return;
  }

  /**
   * Pushes the PI constant onto the current stack.
   */
  void InlineCalculator::pi() {
    Push(pi_c());
    return;
  }

  /**
   * Pushes the Euler constant (e) onto the current stack.
   */
  void InlineCalculator::eConstant() {
    Push(E);
    return;
  }

  /**
   * Determines whether the given function name exists in the current function
   * pool.
   *
   * @param fxname A string containing the name of the function we are looking
   *               for in the pool.
   *
   * @return bool Indicates whether the pool contains the given function.
   */
  bool InlineCalculator::fxExists(const QString &fxname) const {
    return (m_fxPool.contains(fxname));
  }

  /**
   * Adds a function to the function pool. Once in the pool, functions cannot
   * be overwritten. This method will throw an exception if we attempt to add
   * a function with the same name as one already in the pool.
   *
   * @param function The function type pointer to be added to the pool
   *
   * @return InlineCalculator::FxTypePtr The function type pointer that has been
   *         added to the pool.
   * @throw IException::Programmer "Function operator exists!  Cannot replace
   *                                existing functions in the pool"
   */
  InlineCalculator::FxTypePtr InlineCalculator::addFunction(InlineCalculator::FxTypePtr function) {
    FxTypePtr func = find(function->name());
    if (!func) {
      m_fxPool.insert(function->name(), function);
    }
    else {
      QString msg = "Function operator [" + function->name() +
                    "] exists!  Cannot replace existing functions in the pool :-(";
      throw IException(IException::Programmer, msg, _FILEINFO_);
    }
    return (function);
  }

  /**
   * @brief Default token handler if it is undefined during parsing/compilation
   *
   * This method provides the default handling of an undefined token found while
   * compiling an equation.  Users may code their own handler that, for example,
   * could provide additional further functionality.
   *
   * If the virtualized reimplementation is unable to process it, it can throw an
   * exception which will be caught, recorded and compiling will continue until
   * all tokens are processed.  If the optional implementation cannot handle it,
   * return false from your version and an exception will be thrown for you.
   *
   * @param token Character representation of the token.  This could be an
   *              undefined constant, function or soemthing else that occurs in
   *              the equation that is not resolved when this is called.
   *
   * @return bool True if it is handled, false if it cannot (which is an error
   *              and an exception will be invoked).
   *
   */
  bool InlineCalculator::orphanTokenHandler(const QString &token) {
    return (false);
  }

  /**
   * Push the given variable pool onto the current variable pool list.
   *
   * @param variablePool A pointer to the CalculatorVariablePool object to be
   *                     pushed onto the current list.
   */
  void InlineCalculator::pushVariables(CalculatorVariablePool *variablePool) {
    m_variablePoolList.push_back(variablePool);
    return;
  }

  /**
   * Accesses the last variable pool in the current pool list. If the list is
   * empty, an error will be thrown.
   *
   * @return CalculatorVariablePool The last variable pool in the current list.
   * @throw IException::Programmer "Request for nonexistent variable pool.",
   */
  CalculatorVariablePool *InlineCalculator::variables() {
    if ( !m_variablePoolList.isEmpty() ) {
      return (m_variablePoolList.back());
    }
    throw IException(IException::Programmer,
                     "Request for nonexistent variable pool.",
                     _FILEINFO_);
  }

  /**
   * Removes the last variable pool in the current variable pool list.
   */
  void InlineCalculator::popVariables() {
    Clear();
    if ( !m_variablePoolList.isEmpty() ) {
      m_variablePoolList.pop_back();
    }
    return;
  }

  /**
   * Gets a pointer to the function from the current pool that corresponds
   * to the given function name. This method returns null if no function
   * with the given name is found.
   *
   * @param fxname A string containing the name of the function to be
   *               found.
   * @return InlineCalculator::FxTypePtr A pointer to the corresponding
   *         function type, if found.
   */
  InlineCalculator::FxTypePtr InlineCalculator::find(const QString &fxname) {
    FxPoolType::iterator fx = m_fxPool.find(fxname);
    if ( m_fxPool.end() == fx ) return NULL;
    return (*fx);
  }

  /**
   * Adds the recognized functions to the function pool.
   */
  void InlineCalculator::initialize() {
    //  Set up calculator function lookup list
    addFunction(new VoidFx("^", &Calculator::Exponent, this));
    addFunction(new VoidFx("/", &Calculator::Divide, this));
    addFunction(new VoidFx("*", &Calculator::Multiply, this));
    addFunction(new VoidFx("<<", &Calculator::LeftShift, this));
    addFunction(new VoidFx(">>", &Calculator::RightShift, this));
    addFunction(new VoidFx("+", &Calculator::Add, this));
    addFunction(new VoidFx("-", &Calculator::Subtract, this));
    addFunction(new VoidFx(">", &Calculator::GreaterThan, this));
    addFunction(new VoidFx("<", &Calculator::LessThan, this));
    addFunction(new VoidFx(">=", &Calculator::GreaterThanOrEqual, this));
    addFunction(new VoidFx("<=", &Calculator::LessThanOrEqual, this));
    addFunction(new VoidFx("==", &Calculator::Equal, this));
    addFunction(new VoidFx("!=", &Calculator::NotEqual, this));

    //  These are not part of the Calculator class because InfixToPostfix didn't
    //  add/recognize them in the tokenizer.  See InlineInfixToPostfix class which
    //  is part of this calculator.
    addFunction(new VoidFx("&", &Calculator::And, this));
    addFunction(new VoidFx("and", &Calculator::And, this));
    addFunction(new VoidFx("|", &Calculator::Or, this));
    addFunction(new VoidFx("or", &Calculator::Or, this));
    addFunction(new InlineVoidFx("%", &InlineCalculator::floatModulus, this));
    addFunction(new VoidFx("mod", &Calculator::Modulus, this));
    addFunction(new InlineVoidFx("fmod", &InlineCalculator::floatModulus, this));

    addFunction(new VoidFx("--", &Calculator::Negative, this));
    addFunction(new VoidFx("neg", &Calculator::Negative, this));

    addFunction(new VoidFx("min", &Calculator::MinimumPixel, this));
    addFunction(new VoidFx("max", &Calculator::MaximumPixel, this));
    addFunction(new VoidFx("abs", &Calculator::AbsoluteValue, this));
    addFunction(new VoidFx("sqrt", &Calculator::SquareRoot, this));
    addFunction(new VoidFx("log", &Calculator::Log, this));
    addFunction(new VoidFx("ln", &Calculator::Log, this));
    addFunction(new VoidFx("log10", &Calculator::Log10, this));
    addFunction(new InlineVoidFx("pi", &InlineCalculator::pi, this));

    addFunction(new VoidFx("sin", &Calculator::Sine, this));
    addFunction(new VoidFx("cos", &Calculator::Cosine, this));
    addFunction(new VoidFx("tan", &Calculator::Tangent, this));
    addFunction(new VoidFx("sec", &Calculator::Secant, this));
    addFunction(new VoidFx("csc", &Calculator::Cosecant, this));
    addFunction(new VoidFx("cot", &Calculator::Cotangent, this));
    addFunction(new VoidFx("asin", &Calculator::Arcsine, this));
    addFunction(new VoidFx("acos", &Calculator::Arccosine, this));
    addFunction(new VoidFx("atan", &Calculator::Arctangent, this));
    addFunction(new VoidFx("atan2", &Calculator::Arctangent2, this));

    addFunction(new InlineVoidFx("degs", &InlineCalculator::degrees, this));
    addFunction(new InlineVoidFx("rads", &InlineCalculator::radians, this));
    addFunction(new InlineVoidFx("e", &InlineCalculator::eConstant, this));
    addFunction(new InlineVoidFx("||", &InlineCalculator::logicalOr, this));
    addFunction(new InlineVoidFx("&&", &InlineCalculator::logicalAnd, this));

    // Add new functions available for inlining
  //  m_variablePoolList = defaultVariables();

    return;
  }

  /**
   * Discard of all the function pool and class resources.
   *
   */
  void InlineCalculator::destruct() {
    BOOST_FOREACH (FxTypePtr function,  m_fxPool) {
      delete function;
    }

    m_fxPool.clear();
    m_functions.clear();
    return;
  }

  /**
   * Constructs a CalculatorVariablePool object.
   */
  CalculatorVariablePool::CalculatorVariablePool() {
  }

  /**
   * Destroys the CalculatorVariablePool object.
   */
  CalculatorVariablePool::~CalculatorVariablePool() {
  }

  /**
   * Returns true so the real error can be reported.
   *
   * @param variable A string containing the variable we are looking for.
   * @return bool True
   */
  bool CalculatorVariablePool::exists(const QString &variable) const {
    return (true);
  }

  /**
   * Return vector of doubles for Calculator functions.
   *
   * @param variable A string containing the variable.
   * @param index The location in the pool.
   *
   * @return QVector \< double \> A vector of calculator functions.
   * @throw IException::Programmer "No implementation in Calculator variable pool to provide
   *                                value for variable."
   */
  QVector<double> CalculatorVariablePool::value(const QString &variable,
                                                const int &index) const {
    QString mess = "No implementation in Calculator variable pool to provide "
                   " value for variable [" + variable + "].";
    throw IException(IException::Programmer, mess, _FILEINFO_);
  }

  /**
   * Add a parameter to the variable pool.  Some implementations can take
   * advantage of this if desired but it is not standard.
   *
   * @param key A string containing the name of the parameter to be
   *            added.
   * @param values A vector of double precision values to be added to the
   *               variable pool.
   *
   * @throw IException::Programmer "No implementation in Calculator variable pool to add
   *                                value for variable."
   */
  void CalculatorVariablePool::add(const QString &key, QVector<double> &values) {
    QString mess = "No implementation in Calculator variable pool to add "
                   " value for variable [" + key + "].";
    throw IException(IException::Programmer, mess, _FILEINFO_);
  }

  /**
   * Constructs a function binder given a name.
   *
   * @param name A string containing a name for this function.
   */
  FxBinder::FxBinder(const QString &name) : m_name(name) {
  }

  /**
   * Destroys the FxBinder object.
   */
  FxBinder::~FxBinder() {
  }

  /**
   * The name assigned to this function binder.
   *
   * @return QString A string containing the name of this function.
   */
  QString FxBinder::name() const {
    return (m_name);
  }

  /**
   * Executes the function. This method is a wrapper for the virtual dispatch
   * method.
   */
  void FxBinder::execute() {
    dispatch();
  }

  /**
   * Executes the function. This method is a wrapper for the virtual dispatch
   * method.
   */
  void FxBinder::operator()() {
    dispatch();
  }

  /**
   * Accesses the arguments for this function. For scalars and variables, the
   * argument is also the function name.
   *
   * @return QVariant The parameters of this function, as a QVariant.
   */
  QVariant FxBinder::args()  {
    return (QVariant(m_name));
  }

  /**
   * Constructs an InlineVoid function from the given name, InlineCalculator
   * operator, and InlineCalculator.
   *
   * @param name A string containing a name for this function.
   * @param function An InlineCalculator operator that takes no arguments.
   * @param calculator The InlineCalculator used to evaluate this function.
   */
  InlineVoidFx::InlineVoidFx(const QString &name, calcOp function,
                InlineCalculator *calculator) : FxBinder(name),
                                                m_func(function),
                                                m_calc(calculator)  {
  }

  /**
   * Destroys the InlineVoidFx object.
   */
  InlineVoidFx::~InlineVoidFx() {
  }

  /**
   * Calls the function corresponding to this object using its
   * stored InlineCalculator and InlineCalculator operator.
   */
  void InlineVoidFx::dispatch() {
    CALL_MEMBER_FN(*m_calc, m_func)();
  }

  /**
   * Constructs a Parameter function from the given name (containing the
   * appropriate parameters), InlineCalculator operator, and
   * InlineCalculator.
   *
   * @param name A string containing a name for this function. Note: The
   *             name given should include the parameters to be passed
   *             into this function.
   * @param function An InlineCalculator operator that takes parameters.
   * @param calculator The InlineCalculator used to evaluate this function.
   */
  ParameterFx::ParameterFx(const QString &name, calcOp function,
              InlineCalculator *calculator) : FxBinder(name),
                                              m_func(function),
                                              m_calc(calculator){
  }

  /**
   * Destroys the ParameterFx object.
   */
  ParameterFx::~ParameterFx() {
  }

  /**
   * Calls the function corresponding to this object using its
   * stored InlineCalculator, InlineCalculator operator, and
   * arguments.
   */
  void ParameterFx::dispatch() {
    CALL_MEMBER_FN(*m_calc, m_func)(args());
  }

  /**
   * Constructs a Void function from the given name, Calculator operator, and
   * Calculator.
   *
   * @param name A string containing a name for this function.
   * @param function A Calculator operator that takes no arguments.
   * @param calculator The Calculator used to evaluate this function.
   */
  VoidFx::VoidFx(const QString &name, calcOp function,
         InlineCalculator *calculator) : FxBinder(name),
                                         m_func(function),
                                         m_calc(calculator)  {
  }

  /**
   * Destroys the VoidFx object.
   */
  VoidFx::~VoidFx() {
  }

  /**
   * Calls the function corresponding to this object using its stored Calculator
   * and Calculator operator.
   */
  void VoidFx::dispatch() {
    CALL_MEMBER_FN(*m_calc, m_func)();
  }
} // namespace Isis


/// \file   fund.cpp
///
/// \brief
///
/// \authors    maarten
/// \date       2020-01-20
/// \copyright  Copyright 2017-2020 The Institute for New Economic Thinking, Oxford Martin School, University of Oxford
///
///             Licensed under the Apache License, Version 2.0 (the "License");
///             you may not use this file except in compliance with the License.
///             You may obtain a copy of the License at
///
///                 http://www.apache.org/licenses/LICENSE-2.0
///
///             Unless required by applicable law or agreed to in writing, software
///             distributed under the License is distributed on an "AS IS" BASIS,
///             WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
///             See the License for the specific language governing permissions and
///             limitations under the License.
///
///             You may obtain instructions to fulfill the attribution requirements in CITATION.cff
///
#include "fund.hpp"

#include <esl/economics/markets/walras/price_setter.hpp>

#include <esl/economics/markets/walras/quote_message.hpp>
using esl::economics::markets::walras::quote_message;

using esl::law::owner;
using namespace esl::economics;
using namespace esl::economics::accounting;
using namespace esl::economics::finance;

fund::fund(const identity<fund> &i, const jurisdiction &j)
: agent(i)
, law::owner<cash>(i)
, owner<stock>(i)
, identifiable_as<fund>()
, company(i,j)
, lookup_(primary_jurisdiction.tender)
, reset_amount(1.00, currencies::USD)
{
    output_net_asset_value  = create_output<price>("net_asset_value");
    output_signal           = create_output<double>("signal");
    output_cash             = create_output<price>("cash");
    output_stocks           = create_output<price>("stocks");
    output_loans            = create_output<price>("loans");
    output_lending          = create_output<price>("lending");
    output_pnl              = create_output<price>("pnl");

    auto invest_ = [this](auto msg, simulation::time_interval ti, std::seed_seq &seed) {
        return invest(msg, ti, seed);
    };

    ESL_REGISTER_CALLBACK(quote_message, -100, invest_, "make investment decisions");

    auto process_dividends_ = [this](std::shared_ptr<dividend_announcement_message> m,
            simulation::time_interval step,
            std::seed_seq &) {

        for(const auto &[share_, sharedetails_] : m->policy.dividend_per_share){
            std::tuple<identity<company>, share_class> key_ =
                    {esl::reinterpret_identity_cast<company>(m->sender), share_};

            auto iterator_ = stocks.find(key_);
            if (stocks.end() == iterator_){
                LOG(warning) << "fund.cpp(" << __LINE__ << ") fund receives info about stock it is not tracking" << std::endl;
            }else{
                identity<property> stock_ = iterator_->second;
                market_data.shares_outstanding[stock_] = std::get<0>(sharedetails_);

                uint64_t stocks_ = 0;
                for(const auto& s: owner<stock>::inventory){
                    if(s.first->identifier == stock_){
                        stocks_ += s.second.amount;
                    }
                }

                auto dps_ = std::get<1>(sharedetails_);
                auto usd_ = std::make_shared<cash>(currencies::USD);
                auto [i, b] = owner<cash>::inventory.emplace(usd_, quantity(0));

                /// TODO: this needs to be moved someplace else
                auto dividend_received_ = uint64_t((stocks_ * dps_.value * (1.+risk_free_rate)) /  std::get<0>(sharedetails_));
                i->second.amount += dividend_received_;

                for(auto [p, q]: owner<securities_lending_contract>::properties.items){
                    if(p->security == stock_){
                        auto pay_on_short_ = uint64_t((q.amount * dps_.value * (1.+risk_free_rate)) / std::get<0>(sharedetails_));

                        if(i->second.amount < pay_on_short_){
                            LOG(trace) << "pay " << pay_on_short_ << " to fund short position " << std::endl;
                            throw std::invalid_argument("out of cash " + identifier.representation());
                        }
                        i->second.amount -= pay_on_short_;
                    }
                }

                if(market_data.dividend_per_share.end() == market_data.dividend_per_share.find(stock_)){
                    market_data.dividend_per_share.insert(std::make_pair(stock_,std::get<1>(sharedetails_)));
                }else{
                    market_data.dividend_per_share[stock_].value = std::get<1>(sharedetails_).value;
                }
            }
        }
        return step.upper;
    };

    ESL_REGISTER_CALLBACK(quote_message, 16, [&](std::shared_ptr<quote_message> m, simulation::time_interval ti, std::seed_seq &s){

        for(auto [property_, quote_]: m->proposed) {
            auto price_ = std::get<price>(quote_.type);
            lookup_.mark_to_market.erase(property_->identifier);
            auto i = lookup_.mark_to_market.emplace(property_->identifier, price_);
        }

        return ti.upper;
        }, "update prices");

    ESL_REGISTER_CALLBACK(dividend_announcement_message, 0, process_dividends_, "process dividend announcement and store dividends")
}

template<typename T>
price valuate(const esl::law::property_map<quantity> &inventory, const standard &s)
{
    price value_ = price::approximate(0.00, s.reporting_currency);
    for(const auto &[p,q]: inventory){
        auto cast_ = std::dynamic_pointer_cast<T>(p);
        if(cast_){
            value_ += s.value(*cast_, q);
        }
    }
    return value_;
}

///
/// If we are re-setting wealth in an experiment then
/// `target_net_asset_value` is set to the value that we reset wealth to.
/// `target_date` determines the last date on which we reset, for example
/// in experiments where we always reset it is set to 200 years
///
void fund::reset_wealth(price &net_asset_value_, simulation::time_interval ti)
{
    // if the target NAV rule does not apply, exit
    // if we are no longer resetting, exit
    if(!target_net_asset_value.has_value() || double(target_net_asset_value.value()) <= 0. || ti.lower > target_date){
        if(! output_net_asset_value->values.empty() ){
            auto pnl = net_asset_value_ - std::get<1>( output_net_asset_value->values.back() );
            output_pnl->put(ti.lower, pnl);
        }
        return;
    }

    // bailout as a fraction of assets and liabilities
    double bailout_ratio = double(target_net_asset_value.value()) / double(net_asset_value_);

    //std::cout << this->describe() << " bailout_ratio " << std::setprecision(5) << bailout_ratio << " (" << target_net_asset_value.value() << " / " << net_asset_value_ << ")" <<std::endl;
    // bailout monetary amount
    auto bailout_ = price::approximate((double(target_net_asset_value.value()) - double(net_asset_value_)), currencies::USD);
    auto remainder_ = bailout_.value;

    output_pnl->put(ti.lower, -bailout_);

    for(auto &[p, q]: inventory) {
        auto cast_ = std::dynamic_pointer_cast<loan>(p);
        if(cast_) {
            if(remainder_ > 0) {
                auto dec = std::min<std::uint64_t>(inventory[p].amount, remainder_ / 100);
                inventory[p].amount -= dec;  // std::uint64_t(inventory[p].amount * bailout_ratio);
                remainder_ -= dec;
            } else if(remainder_ < 0) {
                inventory[p].amount += (-remainder_/ 100);  // std::uint64_t(inventory[p].amount * bailout_ratio);
                remainder_ = 0;
            }
        }
    }

    if(remainder_){

        for(auto &[p, q]: inventory) {
            auto cast_ = std::dynamic_pointer_cast<cash>(p);
            if(cast_){
                if(remainder_ < 0){
                    auto dec = std::min<std::uint64_t>(inventory[p].amount, -remainder_);
                    inventory[p].amount -= dec;//std::uint64_t(inventory[p].amount * bailout_ratio);
                    remainder_ += dec;
                }else if(remainder_ > 0){
                    inventory[p].amount += remainder_;//std::uint64_t(inventory[p].amount * bailout_ratio);
                    remainder_ = 0;
                }
            }
        }

    }

    if(remainder_){
        for(auto &[p, q]: inventory) {
            auto cast_ = std::dynamic_pointer_cast<stock>(p);
            if(cast_){
                auto i = lookup_.mark_to_market.find(p->identifier);
                if(lookup_.mark_to_market.end() == i){
                    continue;
                }
                auto change_ = remainder_ / std::get<price>(i->second.type).value;
                if(change_ < 0){
                    auto dec = std::min<std::uint64_t>(inventory[p].amount, -change_);
                    inventory[p].amount -= dec;//std::uint64_t(inventory[p].amount * bailout_ratio);
                    remainder_ += change_ * std::get<price>(i->second.type).value;
                }else if(change_ > 0){
                    inventory[p].amount += change_;//std::uint64_t(inventory[p].amount * bailout_ratio);
                    remainder_ = 0;
                }
            }

        }
    }
}

void fund::apply_reinvestment(price &net_asset_value_, simulation::time_interval ti)
{
    ///
    /// \brief  This section relates to the reinvestment rate experiments,
    ///         where we introduce (unseen) retail investors that
    ///         invest based on the performance of the fund
    ///
    if(previous_net_asset_value.has_value() && reinvestment_rate != 1.){
        double returns_ =  double(net_asset_value_) / previous_net_asset_value.value() - 1;
        double compound_returns_f_ = std::pow(1 + returns_, reinvestment_rate - 1) - 1;

        //std::cout <<"returns: " << returns_ << " compounded " << compound_returns_f_ << std::endl;

        // OLD rule matching papers
        //int64_t change_ = static_cast<int64_t>(round(  (reinvestment_rate -1.) * (double(net_asset_value_) - previous_net_asset_value.value()) ));
        // new rule

        double change_ = compound_returns_f_ * double(net_asset_value_);

        for(auto &[p, q]: inventory){
            auto cast_ = std::dynamic_pointer_cast<cash>(p);
            if(cast_){
                net_asset_value_ += price::approximate(change_, net_asset_value_.valuation);
                auto qchange_ = std::max<int64_t>(-((int64_t)inventory[p].amount), int64_t(change_ * 100) );
                inventory[p].amount += qchange_;
            }
        }
    }

}

price fund::net_asset_value(esl::simulation::time_interval ti)
{
    for(auto p: shareholder::prices){
        lookup_.mark_to_market.emplace(p.first->identifier, p.second);
    }

    for(auto p: bond_prices){
        lookup_.mark_to_market.emplace(p.first->identifier, p.second);
    }

    for(auto p: owner<stock>::properties.items){
        if(lookup_.mark_to_market.end() == lookup_.mark_to_market.find(p.first->identifier)){
            lookup_.mark_to_market.emplace(p.first->identifier, price::approximate(100.00, currencies::USD));
        }
    }

    auto cash1_ = valuate<cash>(inventory, lookup_);
    auto stocks_ = valuate<stock>(inventory, lookup_);
    output_stocks->put(ti.lower, stocks_);

    auto loans_value_ = valuate<loan>(inventory, lookup_);
    auto loans_ = price::approximate(double(loans_value_) * (1.+risk_free_rate), loans_value_.valuation);
    output_loans->put(ti.lower, loans_);

    auto lending_ = valuate<securities_lending_contract>(inventory, lookup_);
    output_lending->put(ti.lower, lending_);
    auto cash_ = price::approximate(double(cash1_) * (1.+risk_free_rate), cash1_.valuation);
    auto net_asset_value_ = cash_ + stocks_ + lending_ + loans_;

    output_cash->put(ti.lower, cash_);
    net_asset_value_ = cash_ + stocks_ + lending_ + loans_;

    if(ti.lower > target_date) {
        apply_reinvestment(net_asset_value_, ti);
    }

    reset_wealth(net_asset_value_, ti);
    previous_net_asset_value = double(net_asset_value_);

    output_net_asset_value->put(ti.lower, net_asset_value_);
    return net_asset_value_;
}

time_point fund::act(time_interval interval, std::seed_seq &s)
{
    (void)s;
    return interval.upper;
}

#include <boost/serialization/export.hpp>
BOOST_CLASS_EXPORT(esl::data::output<esl::economics::price>);


/**
BSD 3-Clause License

This file is part of the RootBA project.
https://github.com/NikolausDemmel/rootba

Copyright (c) 2021, Nikolaus Demmel.
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 the copyright holder 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 HOLDER 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.

*/
#pragma once

#include <fstream>

#include <Eigen/Dense>

#include "rootba/bal/bal_bundle_adjustment_helper.hpp"
#include "rootba/bal/bal_problem.hpp"
#include "rootba/bal/bal_residual_options.hpp"
#include "rootba/cg/block_sparse_matrix.hpp"
#include "rootba/util/assert.hpp"
#include "rootba/util/format.hpp"

namespace rootba {

template <typename Scalar, int POSE_SIZE>
class LandmarkBlockSC {
 public:
  struct Options {
    // use Householder instead of Givens for marginalization
    bool use_householder = true;

    // use_valid_projections_only: if true, set invalid projection's
    // residual and jacobian to 0; invalid means z <= 0
    bool use_valid_projections_only = true;

    // huber norm with given threshold, else squared norm
    BalResidualOptions residual_options;

    // ceres uses 1.0 / (1.0 + sqrt(SquaredColumnNorm))
    // we use 1.0 / (eps + sqrt(SquaredColumnNorm))
    Scalar jacobi_scaling_eps = 1.0;
  };

  enum State { UNINITIALIZED = 0, ALLOCATED, NUMERICAL_FAILURE, LINEARIZED };

  inline bool is_numerical_failure() const {
    return state_ == NUMERICAL_FAILURE;
  }

  using Vec2 = Eigen::Matrix<Scalar, 2, 1>;
  using Vec3 = Eigen::Matrix<Scalar, 3, 1>;
  using Vec4 = Eigen::Matrix<Scalar, 4, 1>;

  using VecX = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;

  using Mat3 = Eigen::Matrix<Scalar, 3, 3>;
  using Mat36 = Eigen::Matrix<Scalar, 3, 6>;

  using MatX = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
  using RowMatX =
      Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;

  using Landmark = typename BalProblem<Scalar>::Landmark;
  using Camera = typename BalProblem<Scalar>::Camera;
  using Landmarks = typename BalProblem<Scalar>::Landmarks;
  using Cameras = typename BalProblem<Scalar>::Cameras;

  inline void allocate_landmark(Landmark& lm, const Options& options) {
    options_ = options;

    pose_idx_.clear();
    pose_idx_.reserve(lm.obs.size());
    for (const auto& [cam_idx, obs] : lm.obs) {
      pose_idx_.push_back(cam_idx);
    }

    num_rows_ = pose_idx_.size() * 2;

    lm_idx_ = POSE_SIZE;
    res_idx_ = lm_idx_ + 3;
    num_cols_ = res_idx_ + 1;

    storage_.resize(num_rows_, num_cols_);

    state_ = ALLOCATED;

    lm_ptr_ = &lm;
  }

  // may set state to NumericalFailure --> linearization at this state is
  // unusable. Numeric check is only performed for residuals that were
  // considered to be used (valid), which depends on use_valid_projections_only
  // setting.
  inline void linearize_landmark(const Cameras& cameras) {
    ROOTBA_ASSERT(state_ == ALLOCATED || state_ == NUMERICAL_FAILURE ||
                  state_ == LINEARIZED);

    storage_.setZero(num_rows_, num_cols_);

    bool numerically_valid = true;

    for (size_t i = 0; i < pose_idx_.size(); i++) {
      size_t cam_idx = pose_idx_[i];
      size_t obs_idx = i * 2;
      size_t pose_idx = 0;

      const auto& obs = lm_ptr_->obs.at(cam_idx);
      const auto& cam = cameras.at(cam_idx);

      typename BalBundleAdjustmentHelper<Scalar>::MatRP Jp;
      typename BalBundleAdjustmentHelper<Scalar>::MatRI Ji;
      typename BalBundleAdjustmentHelper<Scalar>::MatRL Jl;

      Vec2 res;
      const bool valid = BalBundleAdjustmentHelper<Scalar>::linearize_point(
          obs.pos, lm_ptr_->p_w, cam.T_c_w, cam.intrinsics, true, res, &Jp, &Ji,
          &Jl);

      if (!options_.use_valid_projections_only || valid) {
        numerically_valid = numerically_valid && Jl.array().isFinite().all() &&
                            Jp.array().isFinite().all() &&
                            Ji.array().isFinite().all() &&
                            res.array().isFinite().all();

        const Scalar res_squared = res.squaredNorm();
        const auto [weighted_error, weight] =
            BalBundleAdjustmentHelper<Scalar>::compute_error_weight(
                options_.residual_options, res_squared);
        const Scalar sqrt_weight = std::sqrt(weight);

        storage_.template block<2, 6>(obs_idx, pose_idx) = sqrt_weight * Jp;
        storage_.template block<2, 3>(obs_idx, pose_idx + 6) = sqrt_weight * Ji;
        storage_.template block<2, 3>(obs_idx, lm_idx_) = sqrt_weight * Jl;
        storage_.template block<2, 1>(obs_idx, res_idx_) = sqrt_weight * res;
      }
    }

    if (numerically_valid) {
      state_ = LINEARIZED;
    } else {
      state_ = NUMERICAL_FAILURE;
    }
  }

  inline void add_Jp_diag2(VecX& res) const {
    ROOTBA_ASSERT(state_ == LINEARIZED);

    for (size_t i = 0; i < pose_idx_.size(); i++) {
      size_t cam_idx = pose_idx_[i];
      res.template segment<POSE_SIZE>(POSE_SIZE * cam_idx) +=
          storage_.template block<2, POSE_SIZE>(2 * i, 0)
              .colwise()
              .squaredNorm();
    }
  }

  inline void scale_Jl_cols() {
    ROOTBA_ASSERT(state_ == LINEARIZED);

    // ceres uses 1.0 / (1.0 + sqrt(SquaredColumnNorm))
    // we use 1.0 / (eps + sqrt(SquaredColumnNorm))
    Jl_col_scale =
        (options_.jacobi_scaling_eps +
         storage_.block(0, lm_idx_, num_rows_, 3).colwise().norm().array())
            .inverse();

    storage_.block(0, lm_idx_, num_rows_, 3) *= Jl_col_scale.asDiagonal();
  }

  inline void scale_Jp_cols(const VecX& jacobian_scaling) {
    ROOTBA_ASSERT(state_ == LINEARIZED);

    for (size_t i = 0; i < pose_idx_.size(); i++) {
      size_t cam_idx = pose_idx_[i];

      storage_.template block<2, POSE_SIZE>(2 * i, 0) *=
          jacobian_scaling.template segment<POSE_SIZE>(POSE_SIZE * cam_idx)
              .asDiagonal();
    }
  }

  inline State get_state() const { return state_; }

  inline void set_landmark_damping(Scalar lambda) { lambda_ = lambda; }

  // Fill the reduced H, b linear system.
  inline void add_Hb(BlockSparseMatrix<Scalar>& accu, VecX& b) const {
    ROOTBA_ASSERT(state_ == LINEARIZED);

    // Compute landmark-landmark block plus inverse
    Mat3 H_ll;
    Mat3 H_ll_inv;
    Vec3 H_ll_inv_bl;
    {
      auto Jl = storage_.block(0, lm_idx_, num_rows_, 3);
      H_ll = Jl.transpose() * Jl;
      H_ll.diagonal().array() += lambda_;
      H_ll_inv = H_ll.inverse();

      H_ll_inv_bl = H_ll_inv * Jl.transpose() * storage_.col(res_idx_);
    }

    // Add pose-pose blocks and
    for (size_t i = 0; i < pose_idx_.size(); i++) {
      size_t cam_idx_i = pose_idx_[i];

      auto Jp_i = storage_.template block<2, POSE_SIZE>(2 * i, 0);
      auto Jl_i = storage_.template block<2, 3>(2 * i, lm_idx_);
      auto r_i = storage_.template block<2, 1>(2 * i, res_idx_);

      MatX H_pp = Jp_i.transpose() * Jp_i;
      accu.add(cam_idx_i, cam_idx_i, std::move(H_pp));

      // Schur complement blocks
      for (size_t j = 0; j < pose_idx_.size(); j++) {
        size_t cam_idx_j = pose_idx_[j];
        auto Jp_j = storage_.template block<2, POSE_SIZE>(2 * j, 0);
        auto Jl_j = storage_.template block<2, 3>(2 * j, lm_idx_);

        MatX H_pl_H_ll_inv_H_lp =
            -Jp_i.transpose() * Jl_i * H_ll_inv * Jl_j.transpose() * Jp_j;

        accu.add(cam_idx_i, cam_idx_j, std::move(H_pl_H_ll_inv_H_lp));
      }

      // add blocks to b
      b.template segment<POSE_SIZE>(cam_idx_i * POSE_SIZE) +=
          Jp_i.transpose() * (r_i - Jl_i * H_ll_inv_bl);
    }
  }

  void back_substitute(const VecX& pose_inc, Scalar& l_diff) {
    ROOTBA_ASSERT(state_ == LINEARIZED);

    Mat3 H_ll = Mat3::Zero();
    Vec3 tmp = Vec3::Zero();
    VecX J_inc;
    J_inc.setZero(num_rows_);

    // Add pose-pose blocks and
    for (size_t i = 0; i < pose_idx_.size(); i++) {
      size_t cam_idx_i = pose_idx_[i];

      auto Jp_i = storage_.template block<2, POSE_SIZE>(2 * i, 0);
      auto Jl_i = storage_.template block<2, 3>(2 * i, lm_idx_);
      auto r_i = storage_.template block<2, 1>(2 * i, res_idx_);

      H_ll += Jl_i.transpose() * Jl_i;

      auto p_inc = pose_inc.template segment<POSE_SIZE>(cam_idx_i * POSE_SIZE);

      tmp += Jl_i.transpose() * (r_i + Jp_i * p_inc);
      J_inc.template segment<2>(2 * i) += Jp_i * p_inc;
    }

    // TODO: store additionally "Hllinv" (inverted with lambda), so we don't
    // need lambda in the interface
    H_ll.diagonal().array() += lambda_;
    Vec3 inc = -H_ll.inverse() * tmp;

    // Add landmark jacobian cost change
    J_inc += storage_.block(0, lm_idx_, num_rows_, 3) * inc;

    l_diff -= J_inc.transpose() * (0.5 * J_inc + storage_.col(res_idx_));

    // Note: scale only after computing model cost change
    inc.array() *= Jl_col_scale.array();
    lm_ptr_->p_w += inc;
  }

  void print_storage(const std::string& filename) const {
    std::ofstream f(filename);

    Eigen::IOFormat clean_fmt(4, 0, " ", "\n", "", "");

    f << "Storage (state: " << state_
      << " Jl_col_scale: " << Jl_col_scale.transpose() << "):\n"
      << storage_.format(clean_fmt) << std::endl;

    f.close();
  }

 protected:
  // Dense storage for pose Jacobians, padding, landmark Jacobians and
  // residuals [J_p (all jacobians in one column) | J_l | res]
  Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>
      storage_;

  Vec3 Jl_col_scale;
  Scalar lambda_ = 0;

  std::vector<size_t> pose_idx_;
  size_t lm_idx_ = 0;
  size_t res_idx_ = 0;

  size_t num_cols_ = 0;
  size_t num_rows_ = 0;

  Options options_;

  State state_ = UNINITIALIZED;

  Landmark* lm_ptr_ = nullptr;
};

}  // namespace rootba


/*

Copyright (c) 2010--2011, Stephane Magnenat, ASL, ETHZ, Switzerland
You can contact the author at <stephane at magnenat dot net>

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 the <organization> 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 ETH-ASL 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.

*/
#include "stdafx.h"
#include "nabo_private.h"
#include "index_heap.h"
#include <iostream>
#include <stdexcept>
#include <limits>
#include <queue>
#include <algorithm>
#include <utility>
#include <boost/numeric/conversion/bounds.hpp>
#include <boost/limits.hpp>
#include <boost/format.hpp>
#ifdef HAVE_OPENMP
#include <omp.h>
#endif

/*!	\file kdtree_cpu.cpp
	\brief kd-tree search, cpu implementation
	\ingroup private
*/

namespace Nabo
{
	//! \ingroup private
	//@{

	using namespace std;

	//! Return the number of bit required to store a value
	/** \param v value to store
	 * \return number of bits required
	 */
	template<typename T>
	T getStorageBitCount(T v)
	{
		for (T i = 0; i < 64; ++i)
		{
			if (v == 0)
				return i;
			v >>= 1;
		}
		return 64;
	}

	//! Return the index of the maximum value of a vector of positive values
	/** \param v vector of positive values
	 * \return index of maximum value, 0 if the vector is empty
	 */
	template<typename T, typename CloudType>
	size_t argMax(const typename NearestNeighbourSearch<T, CloudType>::Vector& v)
	{
		T maxVal(0);
		size_t maxIdx(0);
		for (int i = 0; i < v.size(); ++i)
		{
			if (v[i] > maxVal)
			{
				maxVal = v[i];
				maxIdx = i;
			}
		}
		return maxIdx;
	}

	// OPT
	template<typename T, typename Heap, typename CloudType>
	pair<T,T> KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::getBounds(const BuildPointsIt first, const BuildPointsIt last, const unsigned dim)
	{
		T minVal(boost::numeric::bounds<T>::highest());
		T maxVal(boost::numeric::bounds<T>::lowest());

		for (BuildPointsCstIt it(first); it != last; ++it)
		{
			const T val(cloud.coeff(dim, *it));
			minVal = min(val, minVal);
			maxVal = max(val, maxVal);
		}

		return make_pair(minVal, maxVal);
	}

	template<typename T, typename Heap, typename CloudType>
	unsigned KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::buildNodes(const BuildPointsIt first, const BuildPointsIt last, const Vector minValues, const Vector maxValues)
	{
		const int count(last - first);
		assert(count >= 1);
		const unsigned pos(nodes.size());

		//cerr << count << endl;
		if (count <= int(bucketSize))
		{
			const uint32_t initBucketsSize(buckets.size());
			//cerr << "creating bucket with " << count << " values" << endl;
			for (int i = 0; i < count; ++i)
			{
				const Index index(*(first+i));
				assert(index < cloud.cols());
				buckets.push_back(BucketEntry(&cloud.coeff(0, index), index));
				//cerr << "  " << &cloud.coeff(0, index) << ", " << index << endl;
			}
			//cerr << "at address " << bucketStart << endl;
			nodes.push_back(Node(createDimChildBucketSize(dim, count),initBucketsSize));
			return pos;
		}

		// find the largest dimension of the box
		const unsigned cutDim = argMax<T, CloudType>(maxValues - minValues);
		const T idealCutVal((maxValues(cutDim) + minValues(cutDim))/2);

		// get bounds from actual points
		const pair<T,T> minMaxVals(getBounds(first, last, cutDim));

		// correct cut following bounds
		T cutVal;
		if (idealCutVal < minMaxVals.first)
			cutVal = minMaxVals.first;
		else if (idealCutVal > minMaxVals.second)
			cutVal = minMaxVals.second;
		else
			cutVal = idealCutVal;

		int l(0);
		int r(count-1);
		// partition points around cutVal
		while (1)
		{
			while (l < count && cloud.coeff(cutDim, *(first+l)) < cutVal)
				++l;
			while (r >= 0 && cloud.coeff(cutDim, *(first+r)) >= cutVal)
				--r;
			if (l > r)
				break;
			swap(*(first+l), *(first+r));
			++l; --r;
		}
		const int br1 = l;	// now: points[0..br1-1] < cutVal <= points[br1..count-1]
		r = count-1;
		// partition points[br1..count-1] around cutVal
		while (1)
		{
			while (l < count && cloud.coeff(cutDim, *(first+l)) <= cutVal)
				++l;
			while (r >= br1 && cloud.coeff(cutDim, *(first+r)) > cutVal)
				--r;
			if (l > r)
				break;
			swap(*(first+l), *(first+r));
			++l; --r;
		}
		const int br2 = l; // now: points[br1..br2-1] == cutVal < points[br2..count-1]

		// find best split index
		int leftCount;
		if (idealCutVal < minMaxVals.first)
			leftCount = 1;
		else if (idealCutVal > minMaxVals.second)
			leftCount = count-1;
		else if (br1 > count / 2)
			leftCount = br1;
		else if (br2 < count / 2)
			leftCount = br2;
		else
			leftCount = count / 2;
		assert(leftCount > 0);
		/*if (leftCount >= count)
		{
			cerr << "Error found in kdtree:" << endl;
			cerr << "cloud size: " << cloud.cols() << endl;
			cerr << "count:" << count << endl;
			cerr << "leftCount: " << leftCount << endl;
			cerr << "br1: " << br1 << endl;
			cerr << "br2: " << br2 << endl;
			cerr << "idealCutVal: " << idealCutVal << endl;
			cerr << "cutVal: " << cutVal << endl;
			cerr << "minMaxVals.first: " << minMaxVals.first << endl;
			cerr << "minMaxVals.second: " << minMaxVals.second << endl;
		}*/
		assert(leftCount < count);

		// update bounds for left
		Vector leftMaxValues(maxValues);
		leftMaxValues[cutDim] = cutVal;
		// update bounds for right
		Vector rightMinValues(minValues);
		rightMinValues[cutDim] = cutVal;

		// add this
		nodes.push_back(Node(0, cutVal));

		// recurse
		const unsigned _UNUSED leftChild = buildNodes(first, first + leftCount, minValues, leftMaxValues);
		assert(leftChild == pos + 1);
		const unsigned rightChild = buildNodes(first + leftCount, last, rightMinValues, maxValues);

		// write right child index and return
		nodes[pos].dimChildBucketSize = createDimChildBucketSize(cutDim, rightChild);
		return pos;
	}

	template<typename T, typename Heap, typename CloudType>
	KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt(const CloudType& cloud, const Index dim, const unsigned creationOptionFlags, const Parameters& additionalParameters):
		NearestNeighbourSearch<T, CloudType>::NearestNeighbourSearch(cloud, dim, creationOptionFlags),
		bucketSize(additionalParameters.get<unsigned>("bucketSize", 8)),
		dimBitCount(getStorageBitCount<uint32_t>(this->dim)),
		dimMask((1<<dimBitCount)-1)
	{
		if (bucketSize < 2)
			throw runtime_error((boost::format("Requested bucket size %1%, but must be larger than 2") % bucketSize).str());
		if (cloud.cols() <= bucketSize)
		{
			// make a single-bucket tree
			for (int i = 0; i < cloud.cols(); ++i)
				buckets.push_back(BucketEntry(&cloud.coeff(0, i), i));
			nodes.push_back(Node(createDimChildBucketSize(this->dim, cloud.cols()),uint32_t(0)));
			return;
		}

		const uint64_t maxNodeCount((0x1ULL << (32-dimBitCount)) - 1);
		const uint64_t estimatedNodeCount(cloud.cols() / (bucketSize / 2));
		if (estimatedNodeCount > maxNodeCount)
		{
			throw runtime_error((boost::format("Cloud has a risk to have more nodes (%1%) than the kd-tree allows (%2%). The kd-tree has %3% bits for dimensions and %4% bits for node indices") % estimatedNodeCount % maxNodeCount % dimBitCount % (32-dimBitCount)).str());
		}

		// build point vector and compute bounds
		BuildPoints buildPoints;
		buildPoints.reserve(cloud.cols());
		for (int i = 0; i < cloud.cols(); ++i)
		{
			const Vector& v(cloud.block(0,i,this->dim,1));
			buildPoints.push_back(i);
#ifdef EIGEN3_API
			const_cast<Vector&>(minBound) = minBound.array().min(v.array());
			const_cast<Vector&>(maxBound) = maxBound.array().max(v.array());
#else // EIGEN3_API
			const_cast<Vector&>(minBound) = minBound.cwise().min(v);
			const_cast<Vector&>(maxBound) = maxBound.cwise().max(v);
#endif // EIGEN3_API
		}

		// create nodes
		buildNodes(buildPoints.begin(), buildPoints.end(), minBound, maxBound);
		buildPoints.clear();
	}

	template<typename T, typename Heap, typename CloudType>
	unsigned long KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::knn(const Matrix& query, IndexMatrix& indices, Matrix& dists2, const Index k, const T epsilon, const unsigned optionFlags, const T maxRadius) const
	{
		checkSizesKnn(query, indices, dists2, k, optionFlags);

		const bool allowSelfMatch(optionFlags & NearestNeighbourSearch<T>::ALLOW_SELF_MATCH);
		const bool sortResults(optionFlags & NearestNeighbourSearch<T>::SORT_RESULTS);
		const bool collectStatistics(creationOptionFlags & NearestNeighbourSearch<T>::TOUCH_STATISTICS);
		const T maxRadius2(maxRadius * maxRadius);
		const T maxError2((1+epsilon)*(1+epsilon));
		const int colCount(query.cols());

		assert(nodes.size() > 0);

		IndexMatrix result(k, query.cols());
		unsigned long leafTouchedCount(0);

#pragma omp parallel
		{

		Heap heap(k);
		std::vector<T> off(dim, 0);

#pragma omp for reduction(+:leafTouchedCount) schedule(guided,32)
		for (int i = 0; i < colCount; ++i)
		{
			leafTouchedCount += onePointKnn(query, indices, dists2, i, heap, off, maxError2, maxRadius2, allowSelfMatch, collectStatistics, sortResults);
		}
		}
		return leafTouchedCount;
	}

	template<typename T, typename Heap, typename CloudType>
	unsigned long KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::knn(const Matrix& query, IndexMatrix& indices, Matrix& dists2, const Vector& maxRadii, const Index k, const T epsilon, const unsigned optionFlags) const
	{
		checkSizesKnn(query, indices, dists2, k, optionFlags, &maxRadii);

		const bool allowSelfMatch(optionFlags & NearestNeighbourSearch<T>::ALLOW_SELF_MATCH);
		const bool sortResults(optionFlags & NearestNeighbourSearch<T>::SORT_RESULTS);
		const bool collectStatistics(creationOptionFlags & NearestNeighbourSearch<T>::TOUCH_STATISTICS);
		const T maxError2((1+epsilon)*(1+epsilon));
		const int colCount(query.cols());

		assert(nodes.size() > 0);
		IndexMatrix result(k, query.cols());
		unsigned long leafTouchedCount(0);

#pragma omp parallel
		{

		Heap heap(k);
		std::vector<T> off(dim, 0);

#pragma omp for reduction(+:leafTouchedCount) schedule(guided,32)
		for (int i = 0; i < colCount; ++i)
		{
			const T maxRadius(maxRadii[i]);
			const T maxRadius2(maxRadius * maxRadius);
			leafTouchedCount += onePointKnn(query, indices, dists2, i, heap, off, maxError2, maxRadius2, allowSelfMatch, collectStatistics, sortResults);
		}
		}
		return leafTouchedCount;
	}

	template<typename T, typename Heap, typename CloudType>
	unsigned long KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::onePointKnn(const Matrix& query, IndexMatrix& indices, Matrix& dists2, int i, Heap& heap, std::vector<T>& off, const T maxError2, const T maxRadius2, const bool allowSelfMatch, const bool collectStatistics, const bool sortResults) const
	{
		fill(off.begin(), off.end(), 0);
		heap.reset();
		unsigned long leafTouchedCount(0);

		if (allowSelfMatch)
		{
			if (collectStatistics)
				leafTouchedCount += recurseKnn<true, true>(&query.coeff(0, i), 0, 0, heap, off, maxError2, maxRadius2);
			else
				recurseKnn<true, false>(&query.coeff(0, i), 0, 0, heap, off, maxError2, maxRadius2);
		}
		else
		{
			if (collectStatistics)
				leafTouchedCount += recurseKnn<false, true>(&query.coeff(0, i), 0, 0, heap, off, maxError2, maxRadius2);
			else
				recurseKnn<false, false>(&query.coeff(0, i), 0, 0, heap, off, maxError2, maxRadius2);
		}

		if (sortResults)
			heap.sort();

		heap.getData(indices.col(i), dists2.col(i));
		return leafTouchedCount;
	}

	template<typename T, typename Heap, typename CloudType> template<bool allowSelfMatch, bool collectStatistics>
	unsigned long KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<T, Heap, CloudType>::recurseKnn(const T* query, const unsigned n, T rd, Heap& heap, std::vector<T>& off, const T maxError2, const T maxRadius2) const
	{
		const Node& node(nodes[n]);
		const uint32_t cd(getDim(node.dimChildBucketSize));

		if (cd == uint32_t(dim))
		{
			//cerr << "entering bucket " << node.bucket << endl;
			const BucketEntry* bucket(&buckets[node.bucketIndex]);
			const uint32_t bucketSize(getChildBucketSize(node.dimChildBucketSize));
			for (uint32_t i = 0; i < bucketSize; ++i)
			{
				//cerr << "  " << bucket-> pt << endl;
				//const T dist(dist2<T>(query, cloud.col(index)));
				//const T dist((query - cloud.col(index)).squaredNorm());
				T dist(0);
				const T* qPtr(query);
				const T* dPtr(bucket->pt);
				for (int i = 0; i < this->dim; ++i)
				{
					const T diff(*qPtr - *dPtr);
					dist += diff*diff;
					qPtr++; dPtr++;
				}
				if ((dist <= maxRadius2) &&
					(dist < heap.headValue()) &&
					(allowSelfMatch || (dist > numeric_limits<T>::epsilon()))
				)
					heap.replaceHead(bucket->index, dist);
				++bucket;
			}
			return (unsigned long)(bucketSize);
		}
		else
		{
			const unsigned rightChild(getChildBucketSize(node.dimChildBucketSize));
			unsigned long leafVisitedCount(0);
			T& offcd(off[cd]);
			//const T old_off(off.coeff(cd));
			const T old_off(offcd);
			const T new_off(query[cd] - node.cutVal);
			if (new_off > 0)
			{
				if (collectStatistics)
					leafVisitedCount += recurseKnn<allowSelfMatch, true>(query, rightChild, rd, heap, off, maxError2, maxRadius2);
				else
					recurseKnn<allowSelfMatch, false>(query, rightChild, rd, heap, off, maxError2, maxRadius2);
				rd += - old_off*old_off + new_off*new_off;
				if ((rd <= maxRadius2) &&
					(rd * maxError2 < heap.headValue()))
				{
					offcd = new_off;
					if (collectStatistics)
						leafVisitedCount += recurseKnn<allowSelfMatch, true>(query, n + 1, rd, heap, off, maxError2, maxRadius2);
					else
						recurseKnn<allowSelfMatch, false>(query, n + 1, rd, heap, off, maxError2, maxRadius2);
					offcd = old_off;
				}
			}
			else
			{
				if (collectStatistics)
					leafVisitedCount += recurseKnn<allowSelfMatch, true>(query, n+1, rd, heap, off, maxError2, maxRadius2);
				else
					recurseKnn<allowSelfMatch, false>(query, n+1, rd, heap, off, maxError2, maxRadius2);
				rd += - old_off*old_off + new_off*new_off;
				if ((rd <= maxRadius2) &&
					(rd * maxError2 < heap.headValue()))
				{
					offcd = new_off;
					if (collectStatistics)
						leafVisitedCount += recurseKnn<allowSelfMatch, true>(query, rightChild, rd, heap, off, maxError2, maxRadius2);
					else
						recurseKnn<allowSelfMatch, false>(query, rightChild, rd, heap, off, maxError2, maxRadius2);
					offcd = old_off;
				}
			}
			return leafVisitedCount;
		}
	}

	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapSTL<int,float> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapBruteForceVector<int,float> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapSTL<int,double> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapBruteForceVector<int,double> >;

	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapSTL<int,float>,Eigen::Matrix3Xf>;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapBruteForceVector<int,float>,Eigen::Matrix3Xf>;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapSTL<int,double>,Eigen::Matrix3Xd>;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapBruteForceVector<int,double>,Eigen::Matrix3Xd>;

	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapSTL<int,float>,Eigen::Map<const Eigen::Matrix3Xf, Eigen::Aligned> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<float,IndexHeapBruteForceVector<int,float>,Eigen::Map<const Eigen::Matrix3Xf, Eigen::Aligned> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapSTL<int,double>,Eigen::Map<const Eigen::Matrix3Xd, Eigen::Aligned> >;
	template struct KDTreeUnbalancedPtInLeavesImplicitBoundsStackOpt<double,IndexHeapBruteForceVector<int,double>,Eigen::Map<const Eigen::Matrix3Xd, Eigen::Aligned> >;

	//@}
}