#pragma once
#pragma GCC diagnostic ignored "-Wsign-compare"

#include <vector>
#include <string>
#include <set>
#include <map>
#include <iostream>
#include <fstream>
#include <stdexcept>
#include <cmath>

#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/lexical_cast.hpp>

using std::vector;
using std::string;
using std::set;
using std::map;
using std::istream;
using std::ifstream;
using std::exception;
using std::invalid_argument;
using std::runtime_error;

class Variable{
private:
  string fName;
  vector<string> fDomainNames;

public: //formerly protected
  Variable(string name, vector<string> domainNames){
    fName = name;
    fDomainNames.assign(domainNames.begin(), domainNames.end());
  }

public:
  string name() const {
    return fName;
  }

  vector<string> domainNames() const {
    return fDomainNames;
  }

  //TODO overload << instead
  string toString() {
    return name();
  }

  // functions to be used by std::map

  bool operator< (const Variable& other) const {
    return this->fName < other.name();
  }

  Variable(){
  }

  Variable(const Variable& other){
    fName = other.name();
    fDomainNames = other.domainNames();
  }
};

class Potential {
private:
    string fName;
    int fNumPositions;

public: //formerly protected
    Potential(string name, int numPositions) {
        fName = name;
        fNumPositions = numPositions;
    }

public:
    string name() const{
        return fName;
    }

    int numPositions() const{
        return fNumPositions;
    }

    //TODO overload << instead
    string toString() {
        return "T (" + name() + ")";
    }

    // functions to be used by std::map

    bool operator<(const Potential& other) const {
      return this->fName < other.name();
    }

    Potential(){
    }

    Potential(const Potential& other){
      fName = other.name();
      fNumPositions = other.numPositions();
    }

};

class OnlineEngine;

class Evidence {

private:
  OnlineEngine& fEngine;

public: //formerly protected:
  vector<double> fVarToCurrentNegWeight;
  vector<double> fVarToCurrentPosWeight;

private:
  double defaultWeight (int l);
  void setCurrentWeight (int l, double w);
  void setCurrentWeightToDefault (int l);
  void setCurrentWeightsToDefaults (const vector<int>&);
  void setCurrentWeights (double w, const vector<int>&);

public:
  void retractAll ();
  Evidence (OnlineEngine& engine);
  void varCommit (Variable v, int u);
  void valCommit (Variable v, int u, bool allow);
  void varRetract (Variable v);
  void varSet (Variable v, double w);
  void parmCommit (Potential t, int p, double w);
  void parmRetract (Potential t, int p);

};

class OnlineEngine {
public: // formerly protected:
    static const char CONSTANT = 0;
    static const char LITERAL = 1;
    static const char MULTIPLY = 2;
    static const char ADD = 3;
    vector<char> fNodeToType;
    vector<int> fNodeToLastEdge;
    vector<int> fNodeToLit;
    vector<int> fEdgeToTailNode;
    vector<int> fVarToNegLitNode;
    vector<int> fVarToPosLitNode;
    string READ_DELIMITER;
    string DELIMITER;
    set<Variable> fVariables;
    set<Potential> fPotentials;
    map<string, Variable> fNameToSrcVar;
    map<string, Potential> fNameToSrcPot;
    vector<double> fLogicVarToDefaultNegWeight;
    vector<double> fLogicVarToDefaultPosWeight;
    map<Variable, vector<int>* > fSrcVarToSrcValToIndicator;
    map<Potential, vector<int>* > fSrcPotToSrcPosToParameter;
    vector<double> fNodeToValue;
    vector<double> fNodeToDerivative;
    vector<bool> fNodeToOneZero;
    bool fUpwardPassCompleted;
    bool fTwoPassesCompleted;
    vector<double> fAcVarToMostRecentNegWeight;
    vector<double> fAcVarToMostRecentPosWeight;

    vector<double> clone(vector<double> a);
    void readArithmeticCircuit(istream& r);
    void readLiteralMap(istream& r);
    void upwardPass (Evidence ev);
    void twoPasses (Evidence ev);
    double computedValue (int n);
    int first (int n);
    int rootNode ();
    int numAcNodes ();

public:
    OnlineEngine (string acFilename, string lmFilename);
    void initialize(istream& acReader, istream& lmReader);
    Variable varForName (string n);
    Potential potForName (string n);
    set<Variable> variables ();
    set<Potential> potentials ();
    void assertEvidence (Evidence e, bool secondPass);
    double probOfEvidence ();
    vector<double> varPartials (Variable v);
    map<Variable,vector<double> > varPartials (set<Variable> vs);
    vector<double> varMarginals (Variable v);
    map<Variable,vector<double> > varMarginals (set<Variable> vs);
    vector<double> varPosteriors (Variable v);
    map<Variable,vector<double> > varPosteriors (set<Variable> vs);
    vector<double> potPartials (Potential pot);
    map<Potential,vector<double> > potPartials (set<Potential> ps);
    vector<double> potMarginals (Potential p);
    map<Potential,vector<double> > potMarginals (set<Potential> ps);
    vector<double> potPosteriors (Potential p);
    map<Potential,vector<double> > potPosteriors (set<Potential> ps);
};

inline
double Evidence::defaultWeight (int l) {
    return l < 0 ? fEngine.fLogicVarToDefaultNegWeight[-l]:
           fEngine.fLogicVarToDefaultPosWeight[l];
}

inline
void Evidence::setCurrentWeight (int l, double w) {
    if (l < 0) {
        fVarToCurrentNegWeight[-l] = w;
    } else {
        fVarToCurrentPosWeight[l] = w;
    }
}

inline
void Evidence::setCurrentWeightToDefault (int l) {
    setCurrentWeight (l, defaultWeight (l));
}

inline
void Evidence::setCurrentWeightsToDefaults (const vector<int>& lits) {
    for (int l : lits) {
        setCurrentWeightToDefault (l);
    }
}

inline
void Evidence::setCurrentWeights (double w, const vector<int>& lits) {
    for (int l : lits) {
        setCurrentWeight (l, w);
    }
}

inline
void Evidence::retractAll () {
    fVarToCurrentNegWeight.assign(fEngine.fLogicVarToDefaultNegWeight.begin(),
                                  fEngine.fLogicVarToDefaultNegWeight.end());
    fVarToCurrentPosWeight.assign(fEngine.fLogicVarToDefaultPosWeight.begin(),
                                  fEngine.fLogicVarToDefaultPosWeight.end());
}

inline
Evidence::Evidence (OnlineEngine& engine) : fEngine(engine){
    fVarToCurrentNegWeight.resize(
        fEngine.fLogicVarToDefaultNegWeight.size());
    fVarToCurrentPosWeight.resize(
        fEngine.fLogicVarToDefaultPosWeight.size());
    retractAll ();
}

inline
void Evidence::varCommit (Variable v, int u) {
    varSet (v, 0.0);
    setCurrentWeightToDefault (
        fEngine.fSrcVarToSrcValToIndicator[v]->at(u));
}

inline
void Evidence::valCommit (Variable v, int u, bool allow) {
    int l = fEngine.fSrcVarToSrcValToIndicator[v]->at(u);
    if (allow) {
        setCurrentWeightToDefault (l);
    } else {
        setCurrentWeight (l, 0.0);
    }
}

inline
void Evidence::varRetract (Variable v) {
    setCurrentWeightsToDefaults (*(fEngine.fSrcVarToSrcValToIndicator[v]));
}

inline
void Evidence::varSet (Variable v, double w) {
    setCurrentWeights (w, *(fEngine.fSrcVarToSrcValToIndicator[v]));
}

inline
void Evidence::parmCommit (Potential t, int p, double w) {
    int l = fEngine.fSrcPotToSrcPosToParameter[t]->at(p);
    if (p == 0) {
        throw invalid_argument("Attempt to set value of parameter illegally!");
    }
    setCurrentWeight (l, w);
}

inline
void Evidence::parmRetract (Potential t, int p) {
    int l = fEngine.fSrcPotToSrcPosToParameter[t]->at(p);
    if (p == 0) {
        throw invalid_argument("Attempt to set value of parameter illegally!");
    }
    setCurrentWeightToDefault (l);
}

inline
vector<double> OnlineEngine::clone(vector<double> a) {
    vector<double> ans;
    ans.reserve(a.size());
    ans.assign(a.begin(), a.end());
    return ans;
}

inline
void OnlineEngine::upwardPass (Evidence ev) {
    vector<double>& negValues = ev.fVarToCurrentNegWeight;
    vector<double>& posValues = ev.fVarToCurrentPosWeight;
    int numNodes = numAcNodes ();
    for (int n = 0; n < numNodes; n++) {
        char type = fNodeToType[n];
        if (type == MULTIPLY) {
            double v = 1.0;
            int last = fNodeToLastEdge[n];
            for (int e = first (n); e < last; e++) {
                int ch = fEdgeToTailNode[e];
                double chVal = fNodeToValue[ch];
                if (chVal == 0.0) {
                    v = 0.0;
                    break;
                }
                v *= chVal;
                if (v == 0.0) {
		  throw runtime_error("Underflow");
                }
            }
            fNodeToValue[n] = v;
        } else if (type == ADD) {
            double v = 0.0;
            int last = fNodeToLastEdge[n];
            for (int e = first (n); e < last; e++) {
                int ch = fEdgeToTailNode[e];
                v += fNodeToValue[ch];
            }
            fNodeToValue[n] = v;
        } else if (type == LITERAL) {
            int l = fNodeToLit[n];
            fNodeToValue[n] = (l < 0 ? negValues[-l] : posValues[l]);
        }
        // do nothing for a constant
    }
}

inline
void OnlineEngine::twoPasses (Evidence ev) {

    // Upward pass.

    vector<double>& negValues = ev.fVarToCurrentNegWeight;
    vector<double>& posValues = ev.fVarToCurrentPosWeight;
    int numNodes = numAcNodes ();
    for (int n = 0; n < numNodes; n++) {
        fNodeToDerivative[n] = 0.0;
        char type = fNodeToType[n];
        if (type == MULTIPLY) {
            int numZeros = 0;
            double v = 1.0;
            int last = fNodeToLastEdge[n];
            for (int e = first (n); e < last; e++) {
                int ch = fEdgeToTailNode[e];
                double chVal = computedValue (ch);
                if (chVal == 0.0) {
                    if (++numZeros > 1) {
                        v = 0;
                        break;
                    }
                } else {
                    v *= chVal;
                    if (v == 0.0) {
                        throw runtime_error("Underflow");
                    }
                }
            }
            fNodeToValue[n] = v;
            fNodeToOneZero[n] = numZeros == 1;
        } else if (type == ADD) {
            double v = 0.0;
            int last = fNodeToLastEdge[n];
            for (int e = first (n); e < last; e++) {
                int ch = fEdgeToTailNode[e];
                double chVal = computedValue (ch);
                v += chVal;
            }
            fNodeToValue[n] = v;
        } else if (type == LITERAL) {
            int l = fNodeToLit[n];
            fNodeToValue[n] = (l < 0 ? negValues[-l] : posValues[l]);
        }
        // do nothing for a constant
    }

    // Downward pass.

    fNodeToDerivative[numNodes - 1] = 1.0;
    for (int n = numNodes - 1; n >= 0; n--) {
        char type = fNodeToType[n];
        if (type == LITERAL || type == CONSTANT) {
            continue;
        }
        int last = fNodeToLastEdge[n];
        if (type == MULTIPLY) {
            double value = fNodeToValue[n];
            if (value == 0.0) {
                continue;   // more than one zero
            }
            double x = fNodeToDerivative[n];
            if (x == 0.0) {
                continue;
            }
            x *= value;
            if (x == 0.0) {
                throw runtime_error("Underflow");
            }
            if (fNodeToOneZero[n]) { // exactly one zero
                for (int e = first (n); e < last; e++) {
                    int ch = fEdgeToTailNode[e];
                    double chVal = computedValue (ch);
                    if (chVal == 0.0) {
                        fNodeToDerivative[ch] += x;
                        break;
                    }
                }
            } else { // no zeros
                for (int e = first (n); e < last; e++) {
                    int ch = fEdgeToTailNode[e];
                    double chVal = computedValue (ch);
                    fNodeToDerivative[ch] += x / chVal;
                }
            }
        } else { /* PLUS NODE */
            double x = fNodeToDerivative[n];
            for (int e = first (n); e < last; e++) {
                int ch = fEdgeToTailNode[e];
                fNodeToDerivative[ch] += x;
            }
        }
    }

}

inline
double OnlineEngine::computedValue (int n) {
    return fNodeToOneZero[n] ? 0 : fNodeToValue[n];
}

inline
int OnlineEngine::first (int n) {
    return (n == 0) ? 0 : fNodeToLastEdge[n-1];
}

inline
int OnlineEngine::rootNode () {
    return fNodeToType.size() - 1;
}

inline
int OnlineEngine::numAcNodes () {
    return fNodeToType.size();
}

inline
OnlineEngine::OnlineEngine (string acFilename, string lmFilename) {
	READ_DELIMITER  = "\\$";
	DELIMITER = "$";
	ifstream ac_fs (acFilename, ifstream::in);
        ifstream lm_fs (lmFilename, ifstream::in);
        initialize (ac_fs, lm_fs);
        ac_fs.close();
        lm_fs.close();
}

inline
void OnlineEngine::initialize (istream& acReader, istream& lmReader) {
    readArithmeticCircuit(acReader);
    readLiteralMap(lmReader);
    fNodeToValue.resize(numAcNodes());
    fNodeToDerivative.resize(numAcNodes());
    fNodeToOneZero.resize(numAcNodes());
    fUpwardPassCompleted = false;
    fTwoPassesCompleted = false;
}

inline
Variable OnlineEngine::varForName (string n) {
    return fNameToSrcVar[n];
}

inline
Potential OnlineEngine::potForName (string n) {
    return fNameToSrcPot[n];
}

inline
set<Variable> OnlineEngine::variables () {
    return fVariables;
}

inline
set<Potential> OnlineEngine::potentials () {
    return fPotentials;
}

inline
void OnlineEngine::assertEvidence (Evidence e, bool secondPass) {
    if (secondPass) {
        fAcVarToMostRecentNegWeight = clone (e.fVarToCurrentNegWeight);
        fAcVarToMostRecentPosWeight = clone (e.fVarToCurrentPosWeight);
        twoPasses (e);
    } else {
        fAcVarToMostRecentNegWeight.clear();
        fAcVarToMostRecentPosWeight.clear();
        upwardPass (e);
    }
    fUpwardPassCompleted = true;
    fTwoPassesCompleted = secondPass;
}

inline
double OnlineEngine::probOfEvidence () {
    if (!fUpwardPassCompleted) {
        throw runtime_error ("assertEvidence () must be called!");
    }
    int root = rootNode ();
    return fTwoPassesCompleted ? computedValue (root) : fNodeToValue[root];
}

inline
vector<double> OnlineEngine::varPartials (Variable v) {
    if (!fTwoPassesCompleted) {
        throw runtime_error (
            "assertEvidence () must be called with second pass flag set!");
    }
    vector<double> ans(v.domainNames().size());
    vector<int>& inds = *(fSrcVarToSrcValToIndicator[v]);
    for (int u = 0; u < ans.size(); u++) {
        int l = inds[u];
        ans[u] =
            (l < 0) ?
            fNodeToDerivative[fVarToNegLitNode[-l]] :
            fNodeToDerivative[fVarToPosLitNode[l]];
    }
    return ans;
}

inline
map<Variable,vector<double> > OnlineEngine::varPartials (set<Variable> vs) {
    map<Variable,vector<double> > ans;
    for (Variable v : vs) {
        ans[v] = varPartials (v);
    }
    return ans;
}

inline
vector<double> OnlineEngine::varMarginals (Variable v) {
    if (!fTwoPassesCompleted) {
        throw runtime_error (
            "assertEvidence () must be called with marginals flag set!");
    }
    vector<double> ans(v.domainNames ().size ());
    vector<int>& inds = *(fSrcVarToSrcValToIndicator[v]);
    for (int u = 0; u < ans.size(); u++) {
        int l = inds[u];
        ans[u] =
            (l < 0) ?
            fAcVarToMostRecentNegWeight[-l] *
            fNodeToDerivative[fVarToNegLitNode[-l]] :
            fAcVarToMostRecentPosWeight[l] *
            fNodeToDerivative[fVarToPosLitNode[l]];
    }
    return ans;
}

inline
map<Variable,vector<double> > OnlineEngine::varMarginals (set<Variable> vs) {
    map<Variable,vector<double> > ans;
    for (Variable v : vs) {
        ans[v] = varMarginals (v);
    }
    return ans;
}

inline
vector<double> OnlineEngine::varPosteriors (Variable v) {
    double pe = probOfEvidence ();
    vector<double> ans = varMarginals (v);
    for (int i = 0; i < ans.size(); i++) {
        ans[i] /= pe;
    }
    return ans;
}

inline
map<Variable,vector<double> > OnlineEngine::varPosteriors (set<Variable> vs) {
    map<Variable,vector<double> > ans;
    for (Variable v : vs) {
        ans[v] = varPosteriors (v);
    }
    return ans;
}

inline
vector<double> OnlineEngine::potPartials (Potential pot) {
    vector<int>& parms = *(fSrcPotToSrcPosToParameter[pot]);
    vector<double> ans(parms.size());
    for (int pos = 0; pos < ans.size(); pos++) {
        int l = parms[pos];
        ans[pos] =
            l == 0  ? NAN :
            l < 0 ? fNodeToDerivative[fVarToNegLitNode[-l]] :
            fNodeToDerivative[fVarToPosLitNode[l]];
    }
    return ans;
}

inline
map<Potential,vector<double> > OnlineEngine::potPartials (
    set<Potential> ps) {
    map<Potential,vector<double> > ans;
    for (Potential p : ps) {
        ans[p] = potPartials (p);
    }
    return ans;
}

inline
vector<double> OnlineEngine::potMarginals (Potential p) {
    vector<int>& parms = *(fSrcPotToSrcPosToParameter[p]);
    vector<double> ans = potPartials (p);
    for (int pos = 0; pos < ans.size(); pos++) {
        int l = parms[pos];
        if (!isnan (ans[pos])) {
            ans[pos] *=
                l < 0 ?
                fAcVarToMostRecentNegWeight[-l] : fAcVarToMostRecentPosWeight[l];
        }
    }
    return ans;
}

inline
map<Potential,vector<double> > OnlineEngine::potMarginals (set<Potential> ps) {
    map<Potential,vector<double> > ans;
    for (Potential p : ps) {
        ans[p] = potMarginals (p);
    }
    return ans;
}

inline
vector<double> OnlineEngine::potPosteriors (Potential p) {
    vector<double> ans = potMarginals (p);
    double pe = probOfEvidence ();
    for (int pos = 0; pos < ans.size(); pos++) {
        if (!isnan (ans[pos])) {
            ans[pos] /= pe;
        }
    }
    return ans;
}

inline
map<Potential,vector<double> > OnlineEngine::potPosteriors (
    set<Potential> ps) {
    map<Potential,vector<double> > ans;
    for (Potential p : ps) {
        ans[p] = potPosteriors (p);
    }
    return ans;
}

inline
void OnlineEngine::readArithmeticCircuit(istream& r) {

    int numNodes = __INT_MAX__;
    int nextEdge = 0;
    int nextNode = 0;

    // Process each line.

    while (nextNode < numNodes) {

      // Read the line.  Quit if eof.  Skip if comment or blank line.
      // Otherwise, split into tokens.

      string line;
      if (!getline(r, line)) {break;} // eof
      if (boost::starts_with(line, "c")) {continue;} // comment
      boost::trim(line);
      if (line.length () == 0) {continue;} // blank line
      vector<string> tokens;
      boost::split(tokens, line, boost::is_any_of("\t "), boost::token_compress_on);

      // A header line looks like: "nnf" numNodes numEdges numVars

      if (tokens[0] == "nnf") {
	numNodes = boost::lexical_cast<int>(tokens[1]);
        int numEdges = boost::lexical_cast<int>(tokens[2]);
        int numAcVars = boost::lexical_cast<int>(tokens[3]);
        fEdgeToTailNode.resize(numEdges);
        fNodeToLastEdge.resize(numNodes);
        fNodeToType.resize(numNodes);
        fNodeToLit.resize(numNodes, 0);
        fVarToNegLitNode.assign(numAcVars + 1, -1);
        fVarToPosLitNode.assign(numAcVars + 1, -1);
        continue;
      }

      // This is not a header line, so it must be a node line, which looks
      // like one of the following:
      //   "A" numChildren child+
      //   "O" splitVar numChildren child+
      //   "L" literal
      char ch = tokens[0].at(0);
      if (ch == 'A') {
        fNodeToType[nextNode] = MULTIPLY;
        for (int chIndex = 2; chIndex < tokens.size(); chIndex++) {
          fEdgeToTailNode[nextEdge++] = boost::lexical_cast<int> (tokens[chIndex]);
        }
      } else if (ch == 'O') {
        fNodeToType[nextNode] = ADD;
        for (int chIndex = 3; chIndex < tokens.size(); chIndex++) {
          fEdgeToTailNode[nextEdge++] = boost::lexical_cast<int> (tokens[chIndex]);
        }
      } else /* ch == 'L' */ {
        fNodeToType[nextNode] = LITERAL;
        int l = boost::lexical_cast<int> (tokens[1]);
        fNodeToLit[nextNode] = l;
        (l < 0 ? fVarToNegLitNode : fVarToPosLitNode)[abs (l)] =
          nextNode;
      }
      fNodeToLastEdge[nextNode] = nextEdge;
      nextNode++;

    }

}

inline
void OnlineEngine::readLiteralMap(istream& r) {

    // Prepare to parse.

    int numLits = __INT_MAX__;
    int litsFinished = 0;

    // Process each line.

    while (litsFinished < numLits) {

      // Read the line.  Quit if eof.  Skip if comment (including blank
      // lines).  Otherwise, split into tokens.

      string line;
      if (!getline(r, line)) {break;} // eof
      if (!boost::starts_with(line, ("cc" + DELIMITER))) {continue;} // comment
      boost::trim(line);
      vector<string> tokens;
      boost::split(tokens, line, boost::is_any_of(READ_DELIMITER), boost::token_compress_on);

      // If the line is a header, it is of the form: "cc" "N" numLogicVars.
      string type = tokens[1];
      if (type == "N") {
	int n = boost::lexical_cast<int>(tokens[2]);
        numLits = n * 2;
        fLogicVarToDefaultNegWeight.resize(n+1);
        fLogicVarToDefaultPosWeight.resize(n+1);
        continue;
      }

      // If the line is a variable line, then it is of the form:
      // "cc" "V" srcVarName numSrcVals (srcVal)+

      if (type == "V") {
        string vn = tokens[2];
        int valCount = boost::lexical_cast<int>(tokens[3]);
        vector<string> valNames(valCount, "");
        for (int i = 0; i < valCount; i++) {valNames[i] = tokens[4 + i];}
        Variable v(vn, valNames);
        fSrcVarToSrcValToIndicator[v] = new vector<int>(valCount);
        fNameToSrcVar[vn] = v;
        continue;
      }

      // If the line is a potential line, then it is of the form:
      // "cc" "T" srcPotName parameterCnt.

      if (type == "T") {
        string tn = tokens[2];
        int parmCount = boost::lexical_cast<int> (tokens[3]);
        Potential pot(tn, parmCount);
        fSrcPotToSrcPosToParameter[pot] = new vector<int>(parmCount);
        fNameToSrcPot[tn] = pot;
        continue;
      }

      // This is not a header line, a variable line, or potential line, so
      // it must be a literal description, which looks like one of the
      // following:
      //   "cc" "I" literal weight srcVarName srcValName srcVal
      //   "cc" "P" literal weight srcPotName pos+
      //   "cc" "A" literal weight

      int l = boost::lexical_cast<int>(tokens[2]);
      double w = boost::lexical_cast<double>(tokens[3]);
      (l < 0 ? fLogicVarToDefaultNegWeight : fLogicVarToDefaultPosWeight)
        [abs (l)] = w;
      if (type == "I") {
        string vn = tokens[4];
        //String un = tokens[5];
        int u = boost::lexical_cast<int>(tokens[6]);
        fSrcVarToSrcValToIndicator[fNameToSrcVar[vn]]->at(u) = l;
      } else if (type == "P") {
        string tn = tokens[4];
	vector<string> posStrings;
	boost::split(posStrings, tokens[5], boost::is_any_of(","));
        if (posStrings.size() == 1) {
	  int pos = boost::lexical_cast<int>(posStrings[0]);
          fSrcPotToSrcPosToParameter[fNameToSrcPot[tn]]->at(pos) = l;
        }
      } else if (type == "A") {
      } else {
        throw runtime_error (
          "\"cc\" must be followed by \"N\", \"V\", \"T\", \"I\", \"P\", or \"A\"");
      }
      ++litsFinished;

    }

    // Now create the variables, the map from variable name to variable, and
    // the map from variable to value to indicator.
    for (auto p : fNameToSrcVar)
      fVariables.insert(p.second);
    for (auto p : fNameToSrcPot)
      fPotentials.insert(p.second);

}


//!
//! Contains the implementation of Monte Carlo Markov Chain simulations.
//!
//! \file infer/mcmc.cpp
//! \author Lachlan McCalman
//! \author Darren Shen
//! \date 2014
//! \license Affero General Public License version 3 or later
//! \copyright (c) 2014, NICTA
//!

#pragma once

#include <limits>
#include <random>
#include <iomanip>
#include <chrono>
#include <Eigen/Dense>
#include <boost/circular_buffer.hpp>
#include <iostream>
#include <glog/logging.h>

#include "db/db.hpp"
#include "app/settings.hpp"
#include "infer/chainarray.hpp"
#include "infer/diagnostics.hpp"
#include "comms/transport.hpp"

namespace stateline
{
  namespace mcmc
  {
    //! Represents a Markov Chain Monte Carlo sampler that returns samples from
    //! a distribution.
    //!
    class Sampler
    {
    public:
      //! Create a new sampler.
      //!
      //! \param s Settings to tune various parameters of the MCMC.
      //! \param d Settings for configuring the database for storing samples.
      //! \param stateDim The number of dimensions in each sample.
      //! \param interrupted A flag used to monitor whether the sampler has been interrupted.
      //!
      Sampler(const MCMCSettings& s, const DBSettings& d, uint stateDim, volatile bool& interrupted)
          : db_(d),
            chains_(s.stacks, s.chains, s.initialTempFactor, s.proposalInitialSigma, s.initialSigmaFactor, db_, s.cacheLength, d.recover),
            lengths_(s.stacks * s.chains, 0),
            propStates_(s.stacks * s.chains, stateDim),
            locked_(s.stacks * s.chains, false),
            nextChainBeta_(s.stacks * s.chains),
            nAcceptsGlobal_(s.stacks * s.chains, 0),
            nSwapsGlobal_(s.stacks * (s.chains), 0),
            nSwapAttemptsGlobal_(s.stacks * s.chains, 0),
            sigmas_(s.stacks * s.chains),
            betas_(s.stacks * s.chains),
            acceptRates_(s.stacks * s.chains),
            swapRates_(s.stacks * s.chains),
            lowestEnergies_(s.stacks * s.chains),
            s_(s),
            recover_(d.recover),
            numOutstandingJobs_(0),
            interrupted_(interrupted),
            context_(1)
      {
        // Initialise for logging purposes
        for (uint i = 0; i < s.chains * s.stacks; i++)
        {
          lengths_[i] = chains_.length(i);
          sigmas_[i] = chains_.sigma(i);
          betas_[i] = chains_.beta(i);
          nextChainBeta_[i] = chains_.beta(i);
          acceptRates_[i] = 1;
          swapRates_[i] = 0;
          lowestEnergies_[i] = std::numeric_limits<double>::infinity();
          acceptBuffers_.push_back(boost::circular_buffer<bool>(s.adaptionLength));
          swapBuffers_.push_back(boost::circular_buffer<bool>(s.adaptionLength / s.swapInterval + 1));
          nAcceptsGlobal_[i] = 1;
          acceptBuffers_[i].push_back(true); // first state always accepts
          nSwapsGlobal_[i] = 0;
          nSwapAttemptsGlobal_[i] = 1;
          swapBuffers_[i].push_back(false); // gets rid of a nan, not really needed
        }
      }

      //! Run the sampler for a period of time.
      //!
      //! \param policy Async policy to evaluate states.
      //! \param initialStates Initial chain states. Ignored if recovering.
      //! \param propFn The proposal function.
      //! \param numSeconds The number of seconds to run the MCMC for.
      //!
      template<class AsyncPolicy, class PropFn>
      void run(AsyncPolicy &policy, const std::vector<Eigen::VectorXd>& initialStates, PropFn &propFn, uint numSeconds)
      {
        using namespace std::chrono;

        // Used for publishing statistics to visualisation server.
        zmq::socket_t publisher(context_, ZMQ_PUB);
        publisher.bind("tcp://*:5556");

        // Record the starting time of the MCMC
        steady_clock::time_point startTime = steady_clock::now();

        // Initialise the chains if we're not recovering
        if (!recover_)
        {
          initialise(policy, initialStates);
        }

        // Start all the chains from hottest to coldest
        for (uint i = 0; i < chains_.numTotalChains(); i++)
        {
          uint c = chains_.numTotalChains() - i - 1;
          propose(policy, c, propFn);
        }

        // Initialise the convergence criteria
        uint stateDim = initialStates[0].size();
        EpsrConvergenceCriteria cc(chains_.numStacks(), stateDim);

        // Listen for replies. As soon as a new state comes back,
        // add it to the corresponding chain, and submit a new proposed state
        auto lastLogTime = steady_clock::now();
        auto lastPrintTime = steady_clock::now();
        while (duration_cast<seconds>(steady_clock::now() - startTime).count() < numSeconds && !interrupted_)
        {
          std::pair<uint, double> result;

          // Wait a for reply
          try
          {
            result = policy.retrieve();
          }
          catch (...)
          {
            VLOG(3) << "Comms error -- probably shutting down";
          }

          // If we weree interrupted this state will be garbage
          if (interrupted_)
            break;

          numOutstandingJobs_--;
          uint id = result.first;
          double energy = result.second;

          // Check if this chain is either the coldest or the hottest
          bool isHottestChainInStack = id % chains_.numChains() == chains_.numChains() - 1;
          bool isColdestChainInStack = id % chains_.numChains() == 0;

          // Handle the new proposal and add a new state to the chain
          State propState { propStates_.row(id), energy, chains_.beta(id), false, SwapType::NoAttempt };
          bool propAccepted = chains_.append(id, propState);
          lengths_[id] += 1;
          updateAccepts(id, propAccepted);

          // Update the convergence test if this is the coldest chain in a stack
          if (isColdestChainInStack && chains_.numChains() > 1 && chains_.numStacks() > 1)
          {
            cc.update(id / chains_.numChains(), chains_.lastState(id).sample);
          }

          // Check if this chain was locked. If it was locked, it means that
          // the chain above (hotter) locked it so that it can swap with it
          if (locked_[id])
          {
            // Try swapping this chain with the one above it
            bool swapAccepted = chains_.swap(id, id + 1);
            updateSwaps(id + 1, swapAccepted);
            // Unlock this chain, propgating the lock downwards
            unlock(policy, id, propFn);
          }
          else if (isHottestChainInStack && lengths_[id] % s_.swapInterval == 0 && chains_.numChains() > 1)
          {
            // The hottest chain is ready to swap. Lock the next chain
            // to prevent it from proposing any more
            locked_[id - 1] = true;
          }
          else
          {
            // This chain is not locked, so we can propose
            try
            {
              propose(policy, id, propFn);
            }
            catch (...)
            {
              VLOG(3) << "Comms error -- probably shutting down";
            }
          }

          // Check again after a new interaction with comms
          if (interrupted_)
            break;

          // Log the best energy state so far
          lowestEnergies_[id] = std::min(lowestEnergies_[id], chains_.lastState(id).energy);

          // Check if we need to adapt the step size for this chain
          if (lengths_[id] % s_.proposalAdaptInterval == 0)
          {
            adaptSigma(id);
          }

          // Update the temperature which might have changed while waiting
          chains_.setBeta(id, nextChainBeta_[id]);
          betas_[id] = nextChainBeta_[id];
          // Check for adapting the temperatures of all the chains but the 1st
          if (lengths_[id] % s_.betaAdaptInterval == 0 && !isColdestChainInStack)
          {
            adaptBeta(id);
          }

          // Update the accept and swap rates
          if (duration_cast<milliseconds>(steady_clock::now() - lastLogTime).count() > 50)
          {
            lastLogTime = steady_clock::now();

            std::stringstream s;
            s << "\n\nChainID  Length  MinEngy  CurrEngy    Sigma      AcptRt    GlbAcptRt    Beta     SwapRt   GlbSwapRt\n";
            s << "-----------------------------------------------------------------------------------------------------\n";
            for (uint i = 0; i < chains_.numTotalChains(); i++)
            {
              if (i % chains_.numChains() == 0 && i != 0)
                s << '\n';
              s << std::setprecision(6) << std::showpoint << i << " " << std::setw(9) << lengths_[i] << " " << std::setw(10)
                  << lowestEnergies_[i] << " " << std::setw(10) << chains_.lastState(i).energy << " " << std::setw(10) << sigmas_[i] << " "
                  << std::setw(10) << acceptRates_[i] << " " << std::setw(10) << nAcceptsGlobal_[i] / (double) lengths_[i] << " "
                  << std::setw(10) << betas_[i] << " " << std::setw(10) << swapRates_[i] << " " << std::setw(10)
                  << nSwapsGlobal_[i] / (double) nSwapAttemptsGlobal_[i] << " \n";
            }

            // Quick and dirty way to get the data to the visualisation server
            comms::sendString(publisher, s.str());

            if (duration_cast<milliseconds>(steady_clock::now() - lastPrintTime).count() > 500)
            {
              lastPrintTime = steady_clock::now();

              LOG(INFO)<< s.str() << "\n";

              if (chains_.numStacks() > 1 && chains_.numChains() > 1)
              {
                if (cc.hasConverged())
                {
                  LOG(INFO)<< "converged: true ("<< cc.rHat().transpose() << " < 1.1)";
                }
                else
                {
                  LOG(INFO) << "converged: false ("<< cc.rHat().transpose() << " > 1.1)";
                }
              }
            }
          }
        }

        // Time limit reached. We need to now retrieve all outstanding job results.
        if (!interrupted_)
        {
          while (numOutstandingJobs_--)
          {
            auto result = policy.retrieve();
            uint id = result.first;
            double energy = result.second;
            State propState { propStates_.row(id), energy, chains_.beta(id), false, SwapType::NoAttempt };
            bool propAccepted = chains_.append(id, propState);
            lengths_[id] += 1;
            updateAccepts(id, propAccepted);
          }
        }
        else
        {
          LOG(INFO)<< "Inference interrupted by user: Exiting.";
        }

        // Manually flush any chain states that are in memory to disk
        for (uint i = 0; i < chains_.numTotalChains(); i++)
        {
          chains_.flushCache(i);
        }

        if (cc.hasConverged())
        {
          LOG(INFO)<< "MCMC has converged";
        }
        else
        {
          LOG(INFO) << "WARNING: MCMC has not converged";
        }

        LOG(INFO)<<"Length of chain 0: " << lengths_[0];
    }

    //! Get the MCMC chain array.
    //!
    //! \return A copy of the chain array.
    //!
    ChainArray chains()
    {
      return chains_;
    }

    //! Get the MCMC chain array.
    //!
    //! \return A const reference to the chain array.
    //!
    const ChainArray &chains() const
    {
      return chains_;
    }

  private:

    //! Initialise the sampler.
    //!
    //! \param policy Async policy to evaluate states.
    //! \param initialStates A list containing the initial states of the chains.
    //!
    template <class AsyncPolicy>
    void initialise(AsyncPolicy &policy, const std::vector<Eigen::VectorXd>& initialStates)
    {
      // Evaluate the initial states of the chains
      for (uint i = 0; i < chains_.numTotalChains(); i++)
      {
        propStates_.row(i) = initialStates[i];
        policy.submit(i, propStates_.row(i));
      }

      // Retrieve the energies and temperatures for the initial states
      for (uint i = 0; i < chains_.numTotalChains(); i++)
      {
        auto result = policy.retrieve();
        uint id = result.first;
        double energy = result.second;
        State s
        { initialStates[id], energy, chains_.beta(id), true, SwapType::NoAttempt};
        chains_.initialise(id, s);
      }
    }

    //! Propose a new state through an async policy.
    //!
    //! \param policy Async policy to evaluate proposals.
    //! \param id The id of the chain that is proposing.
    //! \param propFn The proposal function.
    //!
    template <class AsyncPolicy, class PropFn>
    void propose(AsyncPolicy &policy, uint id, PropFn &propFn)
    {
      propStates_.row(id) = propFn(chains_.lastState(id).sample, chains_.sigma(id));
      policy.submit(id, propStates_.row(id));
      numOutstandingJobs_++;
    }

    //! Unlock a chain and reactivate any chains that were waiting for it.
    //!
    //! \param policy Async policy to evulate proposals.
    //! \param id The id of the chain to unlock.
    //! \param propFn The proposal function.
    //!
    template <class AsyncPolicy, class PropFn>
    void unlock(AsyncPolicy &policy, uint id, PropFn &propFn)
    {
      // Unlock this chain
      locked_[id] = false;

      // The hotter chain no longer has to wait for this chain, so
      // it can propose new state
      propose(policy, id + 1, propFn);

      // Check if this was the coldest chain
      if (id % chains_.numChains() != 0)
      {
        // Lock the chain that is below (colder) than this.
        locked_[id - 1] = true;
      }
      else
      {
        // This is the coldest chain and there is no one to swap with
        propose(policy, id, propFn);
      }
    }

    //! Return a new proposal step size for a chain.
    //!
    //! \param id The id of the chain to be adapted.
    //! \return The new proposal step size.
    //!
    double adaptSigma(uint id)
    {
      double acceptRate = acceptRates_[id];
      double oldSigma= chains_.sigma(id);
      double factor = std::pow(acceptRate / s_.proposalOptimalAccept, s_.proposalAdaptRate);
      double boundFactor = std::min(std::max(factor, s_.proposalMinFactor), s_.proposalMaxFactor);
      double gamma = s_.adaptionLength/(double)(s_.adaptionLength+lengths_[id]);

      double newSigma = oldSigma * std::pow(boundFactor, gamma);
      VLOG(2) << "Adapting Sigma" << id <<":" << oldSigma << "->" << newSigma << " @acceptrate:" << acceptRate;
      chains_.setSigma(id, newSigma);
      sigmas_[id] = newSigma;

      // Ensure higher temperature chains have larger sigmas than chains colder than it
      if (id % chains_.numChains() != 0 && newSigma < chains_.sigma(id - 1))
      {
        newSigma = chains_.sigma(id - 1);
      }

      return newSigma;
    }

    //! Adapt a new temperature for a chain.
    //!
    //! \param id The id of the chain to be adapted.
    //!
    void adaptBeta(uint id)
    {
      // Adapt the temperature
      double swapRate = swapRates_[id];
      double rawFactor = std::pow(swapRate/s_.betaOptimalSwapRate, s_.betaAdaptRate);
      double boundedFactor = std::min( std::max(rawFactor, s_.betaMinFactor), s_.betaMaxFactor);
      double beta = chains_.beta(id);
      double lowerBeta = chains_.beta(id-1);// temperature changes propogate UP
      double factor = 1.0/std::max(boundedFactor, 2*beta/(beta + lowerBeta));

      // Set the temperature for this chain (because it hasn't proposed yet)
      double gamma = s_.adaptionLength/(double)(s_.adaptionLength+lengths_[id]);
      double newbeta = chains_.beta(id) * std::pow(factor, gamma);
      chains_.setBeta(id,newbeta);
      nextChainBeta_[id] = newbeta;
      betas_[id] = newbeta;

      // Just for logging
      VLOG(2) << "Adapting Beta" << id << ":" << beta << "->" << newbeta << " @swaprate:" << swapRate;

      // Loop through the other temperatures
      uint coldestChainId = (uint)(id / (double)chains_.numChains()) * chains_.numChains();
      uint hottestChainId = coldestChainId + chains_.numChains()-1;

      // Set the next temperatures for the other chains (as they're busy)
      for (uint i = id+1; i <= hottestChainId; i++)
      {
        nextChainBeta_[i] = nextChainBeta_[i] * std::pow(factor, gamma);
      }
    }

    void updateAccepts(uint id, bool acc)
    {
      uint oldSize = acceptBuffers_[id].size();
      double oldRate = acceptRates_[id];
      bool isFull = acceptBuffers_[id].full();
      bool lastAcc = acceptBuffers_[id].front();

      // Now push on the new state
      acceptBuffers_[id].push_back(acc);

      // Compute the new rate
      uint newSize = acceptBuffers_[id].size();
      double delta = ((int)acc - (int)(lastAcc&&isFull))/(double)newSize;
      double scale = oldSize/(double)newSize;
      acceptRates_[id] = std::max(oldRate*scale + delta, 0.0);
      nAcceptsGlobal_[id] += acc;
      if (acceptRates_[id] > 1.0)
      {
        std::cout << "oldSize: " << oldSize << "\n"
        << "isFull:" << isFull << "\n"
        << "newSize:" << newSize << "\n"
        << "lastAcc:" << lastAcc << "\n"
        << "delta:" << delta << "\n"
        << "scale:" << scale << "\n"
        << "oldRate:" << oldRate << "\n"
        << "rate:" << acceptRates_[id] << "\n"
        << std::endl;
        exit(EXIT_FAILURE);
      }
    }

    void updateSwaps(uint id, bool sw)
    {
      uint oldSize = swapBuffers_[id].size();
      double oldRate = swapRates_[id];
      bool isFull = swapBuffers_[id].full();
      bool lastSw = swapBuffers_[id].front();

      // Now push back the new state
      swapBuffers_[id].push_back(sw);

      // Compute the new rate
      uint newSize = swapBuffers_[id].size();
      double delta = ((int)sw - (int)(lastSw&&isFull))/(double)newSize;
      double scale = oldSize/(double)newSize;
      swapRates_[id] = std::max(oldRate*scale + delta, 0.0);
      nSwapsGlobal_[id] += sw;// for global rate
      nSwapAttemptsGlobal_[id] += 1;
    }

    // The MCMC chain database
    db::Database db_;

    // The MCMC chain wrapper
    ChainArray chains_;

    // lengths of the chains
    std::vector<uint> lengths_;

    // Matrix of proposed states
    Eigen::MatrixXd propStates_;

    // Whether a chain is locked. A locked chain will wait for any outstanding
    // job results and propagate the lock.
    std::vector<bool> locked_;

    // Cache the next temperature as computed by lower chains in the stack
    std::vector<double> nextChainBeta_;

    // Keep track of the swaps and accepts for adaption
    std::vector<unsigned long long> nAcceptsGlobal_;
    std::vector<unsigned long long> nSwapsGlobal_;
    std::vector<unsigned long long> nSwapAttemptsGlobal_;

    std::vector<boost::circular_buffer<bool>> acceptBuffers_;
    std::vector<boost::circular_buffer<bool>> swapBuffers_;

    // For logging purposes
    std::vector<double> sigmas_;
    std::vector<double> betas_;
    std::vector<double> acceptRates_;
    std::vector<double> swapRates_;
    std::vector<double> lowestEnergies_;

    // The MCMC settings
    MCMCSettings s_;

    // Recovering?
    bool recover_;

    // How many jobs haven't been retrieved?
    uint numOutstandingJobs_;

    // Whether an interrupt signal has been sent to the sampler.
    volatile bool& interrupted_;

    zmq::context_t context_;
  };
}
// namespace mcmc
}// namespace obsidian


#include <float.h>
#include <math.h>
#include <time.h>
#include <functional>
#include <iostream>
#include <random>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include <boost/serialization/string.hpp>
#include <boost/serialization/unordered_map.hpp>
#include <boost/serialization/utility.hpp>

#include <clipper/datatypes.hpp>
#include <clipper/logging.hpp>
#include <clipper/selection_policies.hpp>
#include <clipper/util.hpp>

namespace clipper {

// *********
// * State *
// *********

void BanditPolicyState::set_model_map(Map map) { model_map_ = map; }
void BanditPolicyState::add_model(VersionedModelId id, ModelInfo model) {
  model_map_.insert({id, model});
}
void BanditPolicyState::set_weight_sum(double sum) { weight_sum_ = sum; }

std::string BanditPolicyState::serialize() const {
  std::stringstream ss;
  boost::archive::binary_oarchive oa(ss);
  oa << weight_sum_ << model_map_.size()
     << model_map_;  // save weight_sum, map size and map
  return ss.str();
}

BanditPolicyState BanditPolicyState::deserialize(const std::string& bytes) {
  std::stringstream ss;
  ss.str(bytes);
  boost::archive::binary_iarchive ia(ss);
  BanditPolicyState state;
  double sum;
  size_t size;
  ia >> sum >> size;  // load weight_sum and map size
  Map map(size, &versioned_model_hash);
  ia >> map;
  state.set_model_map(map);
  state.set_weight_sum(sum);

  return state;
}

std::string BanditPolicyState::debug_string() const {
  /** State string representation:
    *  For each model: Model Name, Model ID, Model Property Value 1, Model
   * Property Value 2
    *  Different models are separated by semi-colon
    *  e.g. "3.0;classification,00001,1.0,0.4;regression,203422,1.0,0.4;......."
    */

  std::string string_state = "Exp3State;";
  if (model_map_.empty()) {
    std::cout << "State is empty" << std::endl;
    return string_state;
  }
  string_state += std::to_string(weight_sum_) + ";";  // Weight Sum
  for (auto it : model_map_) {
    string_state +=
        it.first.first + std::to_string(it.first.second);  // Model Name
    for (auto model_info_it : it.second) {                 // Model information
      string_state += "," + std::to_string(model_info_it.second);
    }
    string_state += ";";
  }
  return string_state;
}

// ********
// * EXP3 *
// ********

BanditPolicyState Exp3Policy::initialize(
    const std::vector<VersionedModelId>& candidate_models_) {
  Map map(candidate_models_.size(), &versioned_model_hash);
  for (VersionedModelId id : candidate_models_) {
    ModelInfo info = {{"weight", 1.0}};
    map.insert({id, info});
  }
  BanditPolicyState state;
  state.set_model_map(map);
  state.set_weight_sum(map.size() * 1.0);
  return state;
}

BanditPolicyState Exp3Policy::add_models(
    BanditPolicyState state, const std::vector<VersionedModelId>& new_models) {
  double avg;
  if (state.model_map_.empty()) {  // State hasn't been initiated or no models
    avg = 1.0;
  } else {
    avg = state.weight_sum_ / state.model_map_.size();
  }

  for (VersionedModelId id : new_models) {
    ModelInfo info = {{"weight", avg}};
    state.add_model(id, info);
    state.set_weight_sum(state.weight_sum_ + avg);
  }
  return state;
}

VersionedModelId Exp3Policy::select(BanditPolicyState& state) {
  // Helper function for randomly drawing an arm based on its normalized weight
  VersionedModelId selected_model;
  if (state.model_map_.empty()) {
    std::cout << "No models to select from" << std::endl;
    return selected_model;
  }
  double rand_num =
      (double)rand() / (RAND_MAX);  // Pick random number between [0, 1]
  for (auto it = state.model_map_.begin();
       it != state.model_map_.end() && rand_num >= 0; ++it) {
    rand_num -= it->second["weight"] / state.weight_sum_;
    selected_model = it->first;
  }
  return selected_model;
}

std::vector<PredictTask> Exp3Policy::select_predict_tasks(
    BanditPolicyState state, Query query, long query_id) {
  auto selected_model = select(state);
  auto task = PredictTask(query.input_, selected_model, 1.0, query_id,
                          query.latency_micros_);
  std::vector<PredictTask> tasks{task};
  return tasks;
}

Output Exp3Policy::combine_predictions(BanditPolicyState /*state*/,
                                       Query /*query*/,
                                       std::vector<Output> predictions) {
  if (predictions.empty()) {
    std::cout << "No predictions to combine" << std::endl;
    Output output;
    return output;
  }
  return predictions.front();
}

std::pair<std::vector<PredictTask>, std::vector<FeedbackTask>>
Exp3Policy::select_feedback_tasks(BanditPolicyState& state,
                                  FeedbackQuery feedback, long query_id) {
  // Predict Task
  auto selected_model = select(state);
  auto predict_task =
      PredictTask(feedback.feedback_.input_, selected_model, -1, query_id, -1);
  std::vector<PredictTask> predict_tasks{predict_task};
  // Feedback Task
  std::vector<FeedbackTask> feedback_tasks;

  return make_pair(predict_tasks, feedback_tasks);
}

BanditPolicyState Exp3Policy::process_feedback(
    BanditPolicyState state, Feedback feedback,
    std::vector<Output> predictions) {
  if (predictions.empty()) {  // Edge case
    std::cout << "No predictions, so can't update state." << std::endl;
    return state;
  }

  // Compute loss and find which model to update
  auto loss = std::abs(predictions.front().y_hat_ - feedback.y_);
  auto model_id = predictions.front().models_used_.front();
  // Update arm weight and weight_sum
  auto s_i = state.model_map_[model_id]["weight"];
  if (s_i != 0) {
    auto update = exp(-eta * loss / (s_i / state.weight_sum_));
    state.model_map_[model_id]["weight"] = s_i * update;
    state.set_weight_sum(state.weight_sum_ - s_i +
                         state.model_map_[model_id]["weight"]);
  }
  return state;
}

std::string Exp3Policy::serialize_state(BanditPolicyState state) {
  return state.serialize();
}

BanditPolicyState Exp3Policy::deserialize_state(const std::string& bytes) {
  return BanditPolicyState::deserialize(bytes);
}

std::string Exp3Policy::state_debug_string(const BanditPolicyState& state) {
  return state.debug_string();
}

//// ********
//// * EXP4 *
//// ********

BanditPolicyState Exp4Policy::initialize(
    const std::vector<VersionedModelId>& candidate_models_) {
  return Exp3Policy::initialize(candidate_models_);
}

BanditPolicyState Exp4Policy::add_models(
    BanditPolicyState state, const std::vector<VersionedModelId>& new_models) {
  return Exp3Policy::add_models(state, new_models);
}

std::vector<PredictTask> Exp4Policy::select_predict_tasks(
    BanditPolicyState& /*state*/, Query query, long query_id) {
  // Pass along all models selected
  std::vector<PredictTask> tasks;
  for (VersionedModelId id : query.candidate_models_) {
    auto task =
        PredictTask(query.input_, id, 1.0, query_id, query.latency_micros_);
    tasks.push_back(task);
  }
  return tasks;
}

Output Exp4Policy::combine_predictions(BanditPolicyState state, Query /*query*/,
                                       std::vector<Output> predictions) {
  // Weighted Combination of All predictions
  auto y_hat = 0;
  std::vector<VersionedModelId> models;
  for (auto p : predictions) {
    auto model_id = (p.models_used_).front();
    y_hat +=
        (state.model_map_[model_id]["weight"] / state.weight_sum_) * p.y_hat_;
    models.push_back(model_id);
  }
  // Turn y_hat into either 0 or 1
  if (y_hat < 0.5) {
    y_hat = 0;
  } else {
    y_hat = 1;
  }

  auto output = Output(y_hat, models);
  return output;
}

std::pair<std::vector<PredictTask>, std::vector<FeedbackTask>>
Exp4Policy::select_feedback_tasks(BanditPolicyState& /*state*/,
                                  FeedbackQuery feedback, long query_id) {
  std::vector<PredictTask> predict_tasks;
  std::vector<FeedbackTask> feedback_tasks;
  for (VersionedModelId id : feedback.candidate_models_) {
    auto predict_task =
        PredictTask(feedback.feedback_.input_, id, -1, query_id, -1);
    predict_tasks.push_back(predict_task);
  }
  return std::make_pair(predict_tasks, feedback_tasks);
}

BanditPolicyState Exp4Policy::process_feedback(
    BanditPolicyState state, Feedback feedback,
    std::vector<Output> predictions) {
  if (predictions.empty()) {  // Edge case
    std::cout << "No predictions, so can't update state." << std::endl;
    return state;
  }
  // Update every individual model's distribution
  for (auto p : predictions) {
    // Compute loss and find which model to update
    auto loss = std::abs(feedback.y_ - p.y_hat_);
    auto model_id = p.models_used_.front();

    // Update arm weight and weight_sum
    auto s_i = state.model_map_[model_id]["weight"];
    if (s_i != 0) {
      double update = exp(-eta * loss / (s_i / state.weight_sum_));
      state.model_map_[model_id]["weight"] *= update;
      state.set_weight_sum(state.weight_sum_ - s_i +
                           state.model_map_[model_id]["weight"]);
    }
  }
  return state;
}

std::string Exp4Policy::serialize_state(BanditPolicyState state) {
  return state.serialize();
}

BanditPolicyState Exp4Policy::deserialize_state(const std::string& bytes) {
  return BanditPolicyState::deserialize(bytes);
}

std::string Exp4Policy::state_debug_string(const BanditPolicyState& state) {
  return state.debug_string();
}

// ******************
// * Epsilon Greedy *
// ******************

BanditPolicyState EpsilonGreedyPolicy::initialize(
    const std::vector<VersionedModelId>& candidate_models_) {
  BanditPolicyState state;
  Map map(candidate_models_.size(), &versioned_model_hash);
  for (VersionedModelId id : candidate_models_) {
    ModelInfo info = {{"expected_loss", 0.0}, {"times_selected", 0.0}};
    map.insert({id, info});
  }
  state.set_model_map(map);
  return state;
}

BanditPolicyState EpsilonGreedyPolicy::add_models(
    BanditPolicyState state, const std::vector<VersionedModelId>& new_models) {
  // Calculate expected loss from old models
  auto sum = 0.0;
  for (auto model : state.model_map_) {
    sum += model.second.at("expected_loss");
  }
  auto avg = sum / state.model_map_.size();
  // Add new model with average reward
  for (auto id : new_models) {
    ModelInfo info = {{"expected_loss", avg}, {"times_selected", 0.0}};
    state.add_model(id, info);
  }
  return state;
}

VersionedModelId EpsilonGreedyPolicy::select(BanditPolicyState& state) {
  // Helper function for selecting an arm based on lowest expected loss
  VersionedModelId selected_model;
  if (state.model_map_.empty()) {  // Edge case
    std::cout << "No models to select from." << std::endl;
    return selected_model;
  }
  double rand_num = (double)rand() / RAND_MAX;
  if (rand_num >= epsilon) {  // Select best model
    auto min_loss = DBL_MAX;
    for (auto id = state.model_map_.begin(); id != state.model_map_.end();
         ++id) {
      auto model_loss = id->second["expected_loss"];
      if (model_loss < min_loss) {
        min_loss = model_loss;
        selected_model = id->first;
      }
    }
  } else {  // Randomly select
    int rand_draw = rand() % state.model_map_.size();
    auto random_it = next(begin(state.model_map_), rand_draw);
    selected_model = random_it->first;
  }

  return selected_model;
}

std::vector<PredictTask> EpsilonGreedyPolicy::select_predict_tasks(
    BanditPolicyState& state, Query query, long query_id) {
  auto selected_model = select(state);
  auto task = PredictTask(query.input_, selected_model, 1.0, query_id,
                          query.latency_micros_);
  std::vector<PredictTask> tasks{task};
  return tasks;
}

Output EpsilonGreedyPolicy::combine_predictions(
    BanditPolicyState state, Query query, std::vector<Output> predictions) {
  return Exp3Policy::combine_predictions(state, query, predictions);
}

std::pair<std::vector<PredictTask>, std::vector<FeedbackTask>>
EpsilonGreedyPolicy::select_feedback_tasks(BanditPolicyState& state,
                                           FeedbackQuery feedback,
                                           long query_id) {
  return Exp3Policy::select_feedback_tasks(state, feedback, query_id);
}

BanditPolicyState EpsilonGreedyPolicy::process_feedback(
    BanditPolicyState state, Feedback feedback,
    std::vector<Output> predictions) {
  // Edge case
  if (predictions.empty()) {
    return state;
  }
  auto model_id = predictions.front().models_used_.front();
  auto new_loss = std::abs(feedback.y_ - predictions.front().y_hat_);
  // Update expected loss
  int times = state.model_map_[model_id]["times_selected"];
  state.model_map_[model_id]["expected_loss"] =
      (state.model_map_[model_id]["expected_loss"] * times + new_loss) /
      (times + 1);
  // Update times selected
  state.model_map_[model_id]["times_selected"] = times + 1;

  return state;
}

std::string EpsilonGreedyPolicy::serialize_state(BanditPolicyState state) {
  return state.serialize();
}

BanditPolicyState EpsilonGreedyPolicy::deserialize_state(
    const std::string& bytes) {
  return BanditPolicyState::deserialize(bytes);
}

std::string EpsilonGreedyPolicy::state_debug_string(
    const BanditPolicyState& state) {
  return state.debug_string();
}

// ********
// * UCB1 *
// ********

BanditPolicyState UCBPolicy::initialize(
    const std::vector<VersionedModelId>& candidate_models_) {
  return EpsilonGreedyPolicy::initialize(candidate_models_);
}

BanditPolicyState UCBPolicy::add_models(
    BanditPolicyState state, const std::vector<VersionedModelId>& new_models) {
  return EpsilonGreedyPolicy::add_models(state, new_models);
}

VersionedModelId UCBPolicy::select(BanditPolicyState& state) {
  // Helper function for selecting an arm based on lowest upper confidence bound
  VersionedModelId selected_model;
  if (state.model_map_.empty()) {  // Edge case
    std::cout << "No models to select from." << std::endl;
  } else {
    auto min_upper_bound = DBL_MAX;
    for (auto id = state.model_map_.begin(); id != state.model_map_.end();
         ++id) {
      auto model_loss = id->second["expected_loss"];
      auto bound =
          sqrt(2 * log(state.model_map_.size()) / id->second["times_selected"]);
      if (model_loss + bound < min_upper_bound) {
        min_upper_bound = model_loss + bound;
        selected_model = id->first;
      }
    }
  }

  return selected_model;
}

std::vector<PredictTask> UCBPolicy::select_predict_tasks(
    BanditPolicyState& state, Query query, long query_id) {
  auto selected_model = select(state);
  auto task = PredictTask(query.input_, selected_model, 1.0, query_id,
                          query.latency_micros_);
  std::vector<PredictTask> tasks{task};
  return tasks;
}

Output UCBPolicy::combine_predictions(BanditPolicyState state, Query query,
                                      std::vector<Output> predictions) {
  return Exp3Policy::combine_predictions(state, query, predictions);
}

std::pair<std::vector<PredictTask>, std::vector<FeedbackTask>>
UCBPolicy::select_feedback_tasks(BanditPolicyState& state,
                                 FeedbackQuery feedback, long query_id) {
  return Exp3Policy::select_feedback_tasks(state, feedback, query_id);
}

BanditPolicyState UCBPolicy::process_feedback(BanditPolicyState state,
                                              Feedback feedback,
                                              std::vector<Output> predictions) {
  return EpsilonGreedyPolicy::process_feedback(state, feedback, predictions);
}

std::string UCBPolicy::serialize_state(BanditPolicyState state) {
  return state.serialize();
}

BanditPolicyState UCBPolicy::deserialize_state(const std::string& bytes) {
  return BanditPolicyState::deserialize(bytes);
}

std::string UCBPolicy::state_debug_string(const BanditPolicyState& state) {
  return state.debug_string();
}

}  // namespace clipper


/*
 * Software License Agreement (Apache License)
 *
 * Copyright (c) 2021, Southwest Research Institute
 *
 * 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.
 */
#pragma once

#include <descartes_light/descartes_macros.h>
DESCARTES_IGNORE_WARNINGS_PUSH
#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <console_bridge/console.h>
#include <omp.h>
DESCARTES_IGNORE_WARNINGS_POP

#include <descartes_light/solvers/bgl/bgl_ladder_graph_solver.h>
#include <descartes_light/types.h>

#define UNUSED(x) (void)(x)

static void reportFailedEdges(const std::vector<std::size_t>& indices)
{
  if (indices.empty())
    CONSOLE_BRIDGE_logInform("No failed edges");
  else
  {
    std::stringstream ss;
    ss << "Failed edges:\n";
    for (const auto& i : indices)
      ss << "\t" << i << "\n";

    CONSOLE_BRIDGE_logWarn(ss.str().c_str());
  }
}

static void reportFailedVertices(const std::vector<std::size_t>& indices)
{
  if (indices.empty())
    CONSOLE_BRIDGE_logInform("No failed vertices");
  else
  {
    std::stringstream ss;
    ss << "Failed vertices:\n";
    for (const auto& i : indices)
      ss << "\t" << i << "\n";

    CONSOLE_BRIDGE_logWarn(ss.str().c_str());
  }
}

namespace descartes_light
{
template <typename FloatType>
BGLLadderGraphSolver<FloatType>::BGLLadderGraphSolver(unsigned num_threads) : num_threads_{ num_threads } {};

template <typename FloatType>
BuildStatus BGLLadderGraphSolver<FloatType>::buildImpl(
    const std::vector<typename WaypointSampler<FloatType>::ConstPtr>& trajectory,
    const std::vector<typename EdgeEvaluator<FloatType>::ConstPtr>& edge_evaluators,
    const std::vector<typename StateEvaluator<FloatType>::ConstPtr>& state_evaluators)
{
  BuildStatus status;

  // Build Vertices
  ladder_rungs_.resize(trajectory.size());
  long cnt = 0;

  using Clock = std::chrono::high_resolution_clock;
  std::chrono::time_point<Clock> start_time = Clock::now();
#pragma omp parallel for num_threads(num_threads_)
  for (long i = 0; i < static_cast<long>(trajectory.size()); ++i)
  {
    std::vector<StateSample<FloatType>> samples = trajectory[static_cast<size_t>(i)]->sample();
    if (!samples.empty())
    {
      for (StateSample<FloatType>& sample : samples)
      {
        VertexDesc<FloatType> vd;
        if (state_evaluators.empty())
        {
          vd = add_vertex(sample, graph_);
          ladder_rungs_[static_cast<size_t>(i)].push_back(vd);
        }
        else
        {
          std::pair<bool, FloatType> results = state_evaluators[static_cast<size_t>(i)]->evaluate(*sample.state);
          if (results.first)
          {
            sample.cost += results.second;
            vd = add_vertex(sample, graph_);
            ladder_rungs_[static_cast<size_t>(i)].push_back(vd);
          }
        }
      }
    }
    else
    {
#pragma omp critical
      {
        status.failed_vertices.push_back(static_cast<size_t>(i));
      }
    }
#ifndef NDEBUG
#pragma omp critical
    {
      ++cnt;
      std::stringstream ss;
      ss << "Descartes Processed: " << cnt << " of " << trajectory.size() << " vertices";
      CONSOLE_BRIDGE_logInform(ss.str().c_str());
    }
#endif
  }
  double duration = std::chrono::duration<double>(Clock::now() - start_time).count();
  CONSOLE_BRIDGE_logDebug("Descartes took %0.4f seconds to build vertices.", duration);

  // Build Edges
  cnt = 0;
  start_time = Clock::now();
#pragma omp parallel for num_threads(num_threads_)
  for (long i = 1; i < static_cast<long>(trajectory.size()); ++i)
  {
    auto& from = ladder_rungs_[static_cast<size_t>(i) - 1];
    const auto& to = ladder_rungs_[static_cast<size_t>(i)];

    bool found = false;
    for (long j = 0; j < static_cast<long>(from.size()); ++j)
    {
      StateSample<FloatType> from_sample = graph_[from[static_cast<size_t>(j)]];
      for (long k = 0; k < static_cast<long>(to.size()); ++k)
      {
        // Consider the edge:
        StateSample<FloatType> to_sample = graph_[to[static_cast<size_t>(k)]];
        std::pair<bool, FloatType> results =
            edge_evaluators[static_cast<size_t>(i - 1)]->evaluate(*from_sample.state, *to_sample.state);
        if (results.first)
        {
          found = true;
          if (i == 1)
          {
            // first edge captures first rung weights
            boost::add_edge(from[static_cast<size_t>(j)],
                            to[static_cast<size_t>(k)],
                            from_sample.cost + results.second + to_sample.cost,
                            graph_);
          }
          else
          {
            boost::add_edge(
                from[static_cast<size_t>(j)], to[static_cast<size_t>(k)], results.second + to_sample.cost, graph_);
          }
        }
      }
    }  // node loop

    if (!found)
    {
#pragma omp critical
      {
        status.failed_edges.push_back(static_cast<size_t>(i) - 1);
      }
    }
#ifndef NDEBUG
#pragma omp critical
    {
      ++cnt;
      std::stringstream ss;
      ss << "Descartes Processed: " << cnt << " of " << (trajectory.size() - 1) << " edges";
      CONSOLE_BRIDGE_logInform(ss.str().c_str());
    }
#endif
  }  // rung loop
  UNUSED(cnt);
  duration = std::chrono::duration<double>(Clock::now() - start_time).count();
  CONSOLE_BRIDGE_logDebug("Descartes took %0.4f seconds to build edges.", duration);

  // Create a zero-value, zero-cost start node and connect it with a zero-cost edge to each node in the first rung
  {
    auto arr = std::make_shared<State<FloatType>>();
    StateSample<FloatType> start_sample = StateSample<FloatType>{ arr, static_cast<FloatType>(0.0) };
    source_ = add_vertex(start_sample, graph_);
    for (const VertexDesc<FloatType>& target : ladder_rungs_[0])
    {
      boost::add_edge(source_, target, static_cast<FloatType>(0.0), graph_);
    }
  }

  std::sort(status.failed_vertices.begin(), status.failed_vertices.end());
  std::sort(status.failed_edges.begin(), status.failed_edges.end());

  // todo: move these to a utilites function
  reportFailedVertices(status.failed_vertices);
  reportFailedEdges(status.failed_edges);

  if (!status)
  {
    CONSOLE_BRIDGE_logError("LadderGraphSolver failed to build graph.");
  }

  return status;
}

template <typename FloatType>
std::vector<typename State<FloatType>::ConstPtr> BGLLadderGraphSolver<FloatType>::reconstructPath(
    const VertexDesc<FloatType>& source,
    const VertexDesc<FloatType>& target,
    const std::map<VertexDesc<FloatType>, VertexDesc<FloatType>>& predecessor_map) const
{
  // Reconstruct the path from predecessors
  std::vector<typename State<FloatType>::ConstPtr> path;

  VertexDesc<FloatType> v = target;
  path.push_back(graph_[v].state);

  for (VertexDesc<FloatType> u = predecessor_map.at(v); u != v; v = u, u = predecessor_map.at(v))
  {
    path.push_back(graph_[u].state);
  }
  std::reverse(path.begin(), path.end());

  // Check that the last traversed vertex is the source vertex
  if (v != source)
    throw std::runtime_error("Failed to find path through the graph");

  return path;
}

template <typename FloatType>
SearchResult<FloatType> BGLLadderGraphSolver<FloatType>::search()
{
  SearchResult<FloatType> result;

  // Internal properties
  auto index_prop_map = boost::get(boost::vertex_index, graph_);
  auto weight_prop_map = boost::get(boost::edge_weight, graph_);

  result.cost = std::numeric_limits<FloatType>::max();
  result.trajectory = {};

  std::map<VertexDesc<FloatType>, VertexDesc<FloatType>> predecessor_map;
  boost::associative_property_map<std::map<VertexDesc<FloatType>, VertexDesc<FloatType>>> predecessor_prop_map(
      predecessor_map);

  std::map<VertexDesc<FloatType>, FloatType> distance_map;
  boost::associative_property_map<std::map<VertexDesc<FloatType>, FloatType>> distance_prop_map(distance_map);
  boost::dijkstra_shortest_paths(graph_,
                                 source_,
                                 predecessor_prop_map,
                                 distance_prop_map,
                                 weight_prop_map,
                                 index_prop_map,
                                 std::less<>(),
                                 std::plus<>(),
                                 std::numeric_limits<FloatType>::max(),
                                 static_cast<FloatType>(0.0),
                                 boost::default_dijkstra_visitor());

  // Find lowest cost node in last rung
  auto target = std::min_element(ladder_rungs_.back().begin(),
                                 ladder_rungs_.back().end(),
                                 [&distance_map](const VertexDesc<FloatType>& a, const VertexDesc<FloatType>& b) {
                                   return distance_map.at(a) < distance_map.at(b);
                                 });

  // Reconstruct the path from the predecesor map; remove the artificial start state
  result.trajectory = reconstructPath(source_, *target, predecessor_map);
  result.trajectory.erase(result.trajectory.begin());

  result.cost = distance_map.at(*target);

  return result;
}

}  // namespace descartes_light


// Copyright (C) 2019  Rhys Mainwaring
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#include "asv_wave_sim_gazebo_plugins/Wavefield.hh"
#include "asv_wave_sim_gazebo_plugins/CGALTypes.hh"
#include "asv_wave_sim_gazebo_plugins/Convert.hh"
#include "asv_wave_sim_gazebo_plugins/Geometry.hh"
#include "asv_wave_sim_gazebo_plugins/Grid.hh"
#include "asv_wave_sim_gazebo_plugins/Physics.hh"
#include "asv_wave_sim_gazebo_plugins/Utilities.hh"

#include <CGAL/Aff_transformation_3.h>
#include <CGAL/number_utils.h>
#include <CGAL/Simple_cartesian.h>
#include <CGAL/Surface_mesh.h>
#include <CGAL/Timer.h>

#include <Eigen/Dense>

#include <gazebo/gazebo.hh>
#include <gazebo/common/common.hh>
#include <gazebo/msgs/msgs.hh>

#include <ignition/math/Pose3.hh>
#include <ignition/math/Vector2.hh>
#include <ignition/math/Vector3.hh>

#include <tbb/tbb.h>

#include <array>
#include <iostream>
#include <cmath>
#include <string>

namespace asv
{
  typedef CGAL::Aff_transformation_2<Kernel> TransformMatrix;

///////////////////////////////////////////////////////////////////////////////
// Utilities

  std::ostream& operator<<(std::ostream& os, const std::vector<double>& _vec)
  {
    for (auto&& v : _vec )
      os << v << ", ";
    return os;
  }

///////////////////////////////////////////////////////////////////////////////
// WaveParametersPrivate

  /// \internal
  /// \brief Private data for the WavefieldParameters.
  class WaveParametersPrivate
  {
    /// \brief Constructor.
    public: WaveParametersPrivate():
      number(1),
      scale(2.0),
      angle(2.0*M_PI/10.0),
      steepness(1.0),
      amplitude(0.0),
      period(1.0),
      phase(0.0),
      direction(1, 0),
      angularFrequency(2.0*M_PI),
      wavelength(2*M_PI/Physics::DeepWaterDispersionToWavenumber(2.0*M_PI)),
      wavenumber(Physics::DeepWaterDispersionToWavenumber(2.0*M_PI))
    {
    }

    /// \brief The number of component waves.
    public: size_t number;

    /// \brief Set the scale of the largest and smallest waves.
    public: double scale;

    /// \brief Set the angle between component waves and the mean direction.
    public: double angle;

    /// \brief Control the wave steepness. 0 is sine waves, 1 is Gerstner waves.
    public: double steepness;

    /// \brief The mean wave amplitude [m].
    public: double amplitude;

    /// \brief The mean wave period [s]
    public: double period;

    /// \brief The mean wve phase (not currently enabled).
    public: double phase;

    /// \brief The mean wave direction.
    public: Vector2 direction;

    /// \brief The mean wave angular frequency (derived).
    public: double angularFrequency;

    /// \brief The mean wavelength (derived).
    public: double wavelength;

    /// \brief The mean wavenumber (derived).
    public: double wavenumber;

    /// \brief The component wave angular frequencies (derived).
    public: std::vector<double> angularFrequencies;

    /// \brief The component wave amplitudes (derived).
    public: std::vector<double> amplitudes;

    /// \brief The component wave phases (derived).
    public: std::vector<double> phases;

    /// \brief The component wave steepness factors (derived).
    public: std::vector<double> steepnesses;

    /// \brief The component wavenumbers (derived).
    public: std::vector<double> wavenumbers;

    /// \brief The component wave dirctions (derived).
    public: std::vector<Vector2> directions;

    /// \brief Recalculate all derived quantities from inputs.
    public: void Recalculate()
    {
      // Normalize direction
      this->direction = Geometry::Normalize(this->direction);

      // Derived mean values
      this->angularFrequency = 2.0 * M_PI / this->period;
      this->wavenumber = Physics::DeepWaterDispersionToWavenumber(this->angularFrequency);
      this->wavelength = 2.0 * M_PI / this->wavenumber;

      // Update components
      this->angularFrequencies.clear();
      this->amplitudes.clear();
      this->phases.clear();
      this->wavenumbers.clear();
      this->steepnesses.clear();
      this->directions.clear();

      for (size_t i=0; i<this->number; ++i)
      {
        const int n = i - this->number/2;
        const double scaleFactor = std::pow(this->scale, n);
        const double a = scaleFactor * this->amplitude;
        const double k = this->wavenumber / scaleFactor;
        const double omega = Physics::DeepWaterDispersionToOmega(k);
        const double phi = this->phase;
        double q = 0.0;
        if (a != 0)
        {
          q = std::min(1.0, this->steepness / (a * k * this->number));
        }

        this->amplitudes.push_back(a);
        this->angularFrequencies.push_back(omega);
        this->phases.push_back(phi);
        this->steepnesses.push_back(q);
        this->wavenumbers.push_back(k);

        // Direction
        const double c = std::cos(n * this->angle);
        const double s = std::sin(n * this->angle);
        const TransformMatrix T(
          c, -s,
          s,  c
        );
        const Vector2 d = T(this->direction);
        directions.push_back(d);
      }
    }
  };

///////////////////////////////////////////////////////////////////////////////
// WaveParameters

  WaveParameters::~WaveParameters()
  {
  }

  WaveParameters::WaveParameters()
    : data(new WaveParametersPrivate())
  {
    this->data->Recalculate();
  }

  void WaveParameters::FillMsg(gazebo::msgs::Param_V& _msg) const
  {
    // Clear
    _msg.mutable_param()->Clear();

    // "number"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("number");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::INT32);
      nextParam->mutable_value()->set_int_value(this->data->number);
    }
    // "scale"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("scale");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::DOUBLE);
      nextParam->mutable_value()->set_double_value(this->data->scale);
    }
    // "angle"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("angle");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::DOUBLE);
      nextParam->mutable_value()->set_double_value(this->data->angle);
    }
    // "steepness"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("steepness");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::DOUBLE);
      nextParam->mutable_value()->set_double_value(this->data->steepness);
    }
    // "amplitude"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("amplitude");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::DOUBLE);
      nextParam->mutable_value()->set_double_value(this->data->amplitude);
    }
    // "period"
    {
      auto nextParam = _msg.add_param();
      nextParam->set_name("period");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::DOUBLE);
      nextParam->mutable_value()->set_double_value(this->data->period);
    }
    // "direction"
    {
      const auto& direction = this->data->direction;
      auto nextParam = _msg.add_param();
      nextParam->set_name("direction");
      nextParam->mutable_value()->set_type(gazebo::msgs::Any::VECTOR3D);
      nextParam->mutable_value()->mutable_vector3d_value()->set_x(direction.x());
      nextParam->mutable_value()->mutable_vector3d_value()->set_y(direction.y());
      nextParam->mutable_value()->mutable_vector3d_value()->set_z(0);
    }
  }

  void WaveParameters::SetFromMsg(const gazebo::msgs::Param_V& _msg)
  {
    this->data->number    = Utilities::MsgParamSizeT(_msg,    "number",     this->data->number);
    this->data->amplitude = Utilities::MsgParamDouble(_msg,   "amplitude",  this->data->amplitude);
    this->data->period    = Utilities::MsgParamDouble(_msg,   "period",     this->data->period);
    this->data->phase     = Utilities::MsgParamDouble(_msg,   "phase",      this->data->phase);
    this->data->direction = Utilities::MsgParamVector2(_msg,  "direction",  this->data->direction);
    this->data->scale     = Utilities::MsgParamDouble(_msg,   "scale",      this->data->scale);
    this->data->angle     = Utilities::MsgParamDouble(_msg,   "angle",      this->data->angle);
    this->data->steepness = Utilities::MsgParamDouble(_msg,   "steepness",  this->data->steepness);

    this->data->Recalculate();
  }

  void WaveParameters::SetFromSDF(sdf::Element& _sdf)
  {
    this->data->number    = Utilities::SdfParamSizeT(_sdf,    "number",     this->data->number);
    this->data->amplitude = Utilities::SdfParamDouble(_sdf,   "amplitude",  this->data->amplitude);
    this->data->period    = Utilities::SdfParamDouble(_sdf,   "period",     this->data->period);
    this->data->phase     = Utilities::SdfParamDouble(_sdf,   "phase",      this->data->phase);
    this->data->direction = Utilities::SdfParamVector2(_sdf,  "direction",  this->data->direction);
    this->data->scale     = Utilities::SdfParamDouble(_sdf,   "scale",      this->data->scale);
    this->data->angle     = Utilities::SdfParamDouble(_sdf,   "angle",      this->data->angle);
    this->data->steepness = Utilities::SdfParamDouble(_sdf,   "steepness",  this->data->steepness);

    this->data->Recalculate();
  }

  size_t WaveParameters::Number() const
  {
    return this->data->number;
  }

  double WaveParameters::Angle() const
  {
    return this->data->angle;
  }

  double WaveParameters::Scale() const
  {
    return this->data->scale;
  }

  double WaveParameters::Steepness() const
  {
    return this->data->steepness;
  }

  double WaveParameters::AngularFrequency() const
  {
    return this->data->angularFrequency;
  }

  double WaveParameters::Amplitude() const
  {
    return this->data->amplitude;
  }

  double WaveParameters::Period() const
  {
    return this->data->period;
  }

  double WaveParameters::Phase() const
  {
    return this->data->phase;
  }

  double WaveParameters::Wavelength() const
  {
    return this->data->wavelength;
  }

  double WaveParameters::Wavenumber() const
  {
    return this->data->wavenumber;
  }

  Vector2 WaveParameters::Direction() const
  {
    return this->data->direction;
  }

  void WaveParameters::SetNumber(size_t _number)
  {
    this->data->number = _number;
    this->data->Recalculate();
  }

  void WaveParameters::SetAngle(double _angle)
  {
    this->data->angle = _angle;
    this->data->Recalculate();
  }

  void WaveParameters::SetScale(double _scale)
  {
    this->data->scale = _scale;
    this->data->Recalculate();
  }

  void WaveParameters::SetSteepness(double _steepness)
  {
    this->data->steepness = _steepness;
    this->data->Recalculate();
  }

  void WaveParameters::SetAmplitude(double _amplitude)
  {
    this->data->amplitude = _amplitude;
    this->data->Recalculate();
  }

  void WaveParameters::SetPeriod(double _period)
  {
    this->data->period = _period;
    this->data->Recalculate();
  }

  void WaveParameters::SetPhase(double _phase)
  {
    this->data->phase = _phase;
    this->data->Recalculate();
  }

  void WaveParameters::SetDirection(const Vector2& _direction)
  {
    this->data->direction = _direction;
    this->data->Recalculate();
  }

  const std::vector<double>& WaveParameters::AngularFrequency_V() const
  {
    return this->data->angularFrequencies;
  }

  const std::vector<double>& WaveParameters::Amplitude_V() const
  {
    return this->data->amplitudes;
  }

  const std::vector<double>& WaveParameters::Phase_V() const
  {
    return this->data->phases;
  }

  const std::vector<double>& WaveParameters::Steepness_V() const
  {
    return this->data->steepnesses;
  }

  const std::vector<double>& WaveParameters::Wavenumber_V() const
  {
    return this->data->wavenumbers;
  }

  const std::vector<Vector2>& WaveParameters::Direction_V() const
  {
    return this->data->directions;
  }

  void WaveParameters::DebugPrint() const
  {
    gzmsg << "number:     " << this->data->number << std::endl;
    gzmsg << "scale:      " << this->data->scale << std::endl;
    gzmsg << "angle:      " << this->data->angle << std::endl;
    gzmsg << "period:     " << this->data->period << std::endl;
    gzmsg << "amplitude:  " << this->data->amplitudes << std::endl;
    gzmsg << "wavenumber: " << this->data->wavenumbers << std::endl;
    gzmsg << "omega:      " << this->data->angularFrequencies << std::endl;
    gzmsg << "phase:      " << this->data->phases << std::endl;
    gzmsg << "steepness:  " << this->data->steepnesses << std::endl;
    for (auto&& d : this->data->directions)
    {
      gzmsg << "direction:  " << d << std::endl;
    }
  }

///////////////////////////////////////////////////////////////////////////////
// WavefieldPrivate

  /// \internal
  /// \brief Private data for the WavefieldParameters.
  class WavefieldPrivate
  {
    /// \brief Constructor.
    public: WavefieldPrivate() :
      params(new WaveParameters()),
      size({ 1000, 1000 }),
      cellCount({ 50, 50 })
      // isDirty(false)
    {
    }

    /// \brief Constructor.
    ///
    /// \brief param[in] _size      The dimensions of the wave field [m].
    /// \brief param[in] _cellCount The number of cells in each direction.
    public: WavefieldPrivate(
      const std::array<double, 2>& _size,
      const std::array<size_t, 2>& _cellCount
    ) :
      params(new WaveParameters()),
      size(_size),
      cellCount(_cellCount)
      // isDirty(false)
    {
    }

    /// \brief Wave parameters
    public: std::shared_ptr<WaveParameters> params;

    /// brief The dimensions of the wave field in each direction [m].
    public: std::array<double, 2> size;

    /// \brief The number of grid cells in each direction.
    public: std::array<size_t, 2> cellCount;

    /// \brief The initial wave field mesh. This is a triangulated regular grid.
    public: std::shared_ptr<const Grid> initialGrid;

    /// \brief The current position of the wave field.
    public: std::shared_ptr<Grid> grid;

    /// \brief Retain the Gazebo mesh as it's required to update visuals.
    // public: std::string name;
    // public: std::shared_ptr<gazebo::common::Mesh> gzMesh;
    // public: gazebo::common::SubMesh* gzSubMesh;
    // public: bool isDirty;
  };

///////////////////////////////////////////////////////////////////////////////
// Wavefield

  Wavefield::~Wavefield()
  {
  }

  Wavefield::Wavefield(
    const std::string& _name) :
    data(new WavefieldPrivate())
  {
    // Grid
    this->data->initialGrid.reset(new Grid(
      this->data->size, this->data->cellCount));
    this->data->grid.reset(new Grid(
      this->data->size, this->data->cellCount));

    // GzMesh
    // this->data->name = _name;
    // this->data->gzMesh = std::make_shared<gazebo::common::Mesh>();
    // this->InitGzMesh();

    // Update
    this->Update(0.0);
  }

  Wavefield::Wavefield(
    const std::string& _name,
    const std::array<double, 2>& _size,
    const std::array<size_t, 2>& _cellCount) :
    data(new WavefieldPrivate(_size, _cellCount))
  {
    // Grid
    this->data->initialGrid.reset(new Grid(
      this->data->size, this->data->cellCount));
    this->data->grid.reset(new Grid(
      this->data->size, this->data->cellCount));

    // GzMesh
    // this->data->name = _name;
    // this->data->gzMesh = std::make_shared<gazebo::common::Mesh>();
    // this->InitGzMesh();

    // Update
    this->Update(0.0);
  }

  std::shared_ptr<const Mesh> Wavefield::GetMesh() const
  {
    return this->data->grid->GetMesh();
  }

  std::shared_ptr<const Grid> Wavefield::GetGrid() const
  {
    return this->data->grid;
  }

  // std::shared_ptr<const gazebo::common::Mesh> Wavefield::GetGzMesh() const
  // {
  //   this->LazyUpdateGzMesh();
  //   return this->data->gzMesh;
  // }

  std::shared_ptr<const WaveParameters> Wavefield::GetParameters() const
  {
    return this->data->params;
  }

  void Wavefield::SetParameters(std::shared_ptr<WaveParameters> _params) const
  {
    GZ_ASSERT(_params != nullptr, "Invalid parameter _params");
    this->data->params = _params;
  }

  void Wavefield::Update(double _time)
  {
    this->UpdateGerstnerWave(_time);

    // this->data->isDirty = true;
  }

  void Wavefield::UpdateGerstnerWave(double _time)
  {
    // Single wave params
    // auto amplitude  = this->data->params->Amplitude();
    // auto wavenumber = this->data->params->Wavenumber();
    // auto omega      = this->data->params->AngularFrequency();
    // auto phase      = this->data->params->Phase();
    // auto q          = this->data->params->Steepness();
    // auto direction  = this->data->params->Direction();

    // Multiple wave params
    const auto  number     = this->data->params->Number();
    const auto& amplitude  = this->data->params->Amplitude_V();
    const auto& wavenumber = this->data->params->Wavenumber_V();
    const auto& omega      = this->data->params->AngularFrequency_V();
    const auto& phase      = this->data->params->Phase_V();
    const auto& q          = this->data->params->Steepness_V();
    const auto& direction  = this->data->params->Direction_V();

    const auto& initMesh = *this->data->initialGrid->GetMesh();
    auto& mesh = *this->data->grid->GetMesh();

    // Reset points to original positions
    for (
      auto&& it = std::make_pair(std::begin(initMesh.vertices()), std::begin(mesh.vertices()));
      it.first != std::end(initMesh.vertices()) && it.second != std::end(mesh.vertices());
      ++it.first, ++it.second)
    {
      auto& vtx0 = *it.first;
      auto& vtx1 = *it.second;
      mesh.point(vtx1) = initMesh.point(vtx0);
    }

    // Multiple wave update
    for (size_t i=0; i<number; ++i)
    // size_t i = 2;
    {
      const auto& amplitude_i = amplitude[i];
      const auto& wavenumber_i = wavenumber[i];
      const auto& omega_i = omega[i];
      const auto& phase_i = phase[i];
      const auto& direction_i = direction[i];
      const auto& q_i = q[i];

      for (
        auto&& it = std::make_pair(std::begin(initMesh.vertices()), std::begin(mesh.vertices()));
        it.first != std::end(initMesh.vertices()) && it.second != std::end(mesh.vertices());
        ++it.first, ++it.second)
      {
        auto& vtx0 = *it.first;
        auto& vtx1 = *it.second;

        const Point3& p0 = initMesh.point(vtx0);
        Vector2 v0(p0.x(), p0.y()); // - CGAL::ORIGIN;

        // Multiple waves
        const double angle  = CGAL::to_double(direction_i * v0) * wavenumber_i - omega_i * _time + phase_i;
        const double s = std::sin(angle);
        const double c = std::cos(angle);
        Vector3 v1(
          - direction_i.x() * q_i * amplitude_i * s,
          - direction_i.y() * q_i * amplitude_i * s,
          + amplitude_i * c
        );

        mesh.point(vtx1) += v1;
      }
    }

    // Single wave update
    // for (
    //   auto&& it = std::make_pair(std::begin(initMesh.vertices()), std::begin(mesh.vertices()));
    //   it.first != std::end(initMesh.vertices()) && it.second != std::end(mesh.vertices());
    //   ++it.first, ++it.second)
    // {
    //   auto& vtx0 = *it.first;
    //   auto& vtx1 = *it.second;

    //   const Point3& p0 = initMesh.point(vtx0);
    //   Vector3 v0 = p0 - CGAL::ORIGIN;
    //   Vector3 v1 = CGAL::NULL_VECTOR;

    //   auto prj    = direction * v0;
    //   double ang  = CGAL::to_double(prj) * wavenumber - omega * _time + phase;
    //   double sang = std::sin(ang);
    //   double cang = std::cos(ang);
    //   v1 += Vector3(
    //     - direction.x() * q * amplitude * sang,
    //     - direction.y() * q * amplitude * sang,
    //     + amplitude * cang
    //   );

    //   mesh.point(vtx1) = p0 + v1;
    // }

  }

  // void Wavefield::InitGzMesh()
  // {
  //   auto& grid = *this->data->grid;
  //   auto& mesh = *grid.GetMesh();

  //   // Update Normals
  //   grid.RecalculateNormals();

  //   // Make GzMesh
  //   auto& name = this->data->name;
  //   this->data->gzMesh.reset(new gazebo::common::Mesh());
  //   this->data->gzMesh->SetName(name);

  //   const double Lx = this->data->size[0];
  //   const double Ly = this->data->size[1];

  //   // Create the submesh and retain a pointer
  //   std::unique_ptr<gazebo::common::SubMesh> gzSubMesh(new gazebo::common::SubMesh());
  //   this->data->gzSubMesh = gzSubMesh.get();

  //   size_t iv = 0;
  //   size_t it = 0;
  //   for(auto&& face : mesh.faces())
  //   {
  //     Triangle tri  = Geometry::MakeTriangle(mesh, face);
  //     auto& normal = grid.GetNormal(it++);

  //     ignition::math::Vector3d ignP0(ToIgn(tri[0]));
  //     ignition::math::Vector3d ignP1(ToIgn(tri[1]));
  //     ignition::math::Vector3d ignP2(ToIgn(tri[2]));
  //     ignition::math::Vector3d ignNormal(ToIgn(normal));

  //     // Vertices
  //     gzSubMesh->AddVertex(ignP0);
  //     gzSubMesh->AddVertex(ignP1);
  //     gzSubMesh->AddVertex(ignP2);

  //     // Vertex normals
  //     gzSubMesh->AddNormal(ignNormal);
  //     gzSubMesh->AddNormal(ignNormal);
  //     gzSubMesh->AddNormal(ignNormal);

  //     // Face
  //     gzSubMesh->AddIndex(iv++);
  //     gzSubMesh->AddIndex(iv++);
  //     gzSubMesh->AddIndex(iv++);

  //     // Vertex texture coordinates
  //     gzSubMesh->AddTexCoord(0.5 + ignP0.X()/Lx, 0.5 - ignP0.Y()/Ly);
  //     gzSubMesh->AddTexCoord(0.5 + ignP1.X()/Lx, 0.5 - ignP1.Y()/Ly);
  //     gzSubMesh->AddTexCoord(0.5 + ignP2.X()/Lx, 0.5 - ignP2.Y()/Ly);
  //   }

  //   this->data->gzMesh->AddSubMesh(gzSubMesh.release());
  // }

  // void Wavefield::LazyUpdateGzMesh() const
  // {
  //   // Update the Gazebo mesh vertices
  //   if (this->data->isDirty)
  //   {
  //     auto& grid = *this->data->grid;
  //     auto& mesh = *grid.GetMesh();
  //     auto& gzSubMesh = *this->data->gzSubMesh;

  //     // Update Normals
  //     grid.RecalculateNormals();

  //     // Update GzMesh
  //     size_t iv = 0;
  //     size_t it = 0;
  //     for(auto&& face : mesh.faces())
  //     {
  //       Triangle tri  = Geometry::MakeTriangle(mesh, face);
  //       auto& normal = grid.GetNormal(it++);

  //       ignition::math::Vector3d ignP0(ToIgn(tri[0]));
  //       ignition::math::Vector3d ignP1(ToIgn(tri[1]));
  //       ignition::math::Vector3d ignP2(ToIgn(tri[2]));
  //       ignition::math::Vector3d ignNormal(ToIgn(normal));

  //       // Update Vertices and Normals
  //       gzSubMesh.SetVertex(iv, ignP0);
  //       gzSubMesh.SetNormal(iv++, ignNormal);

  //       gzSubMesh.SetVertex(iv, ignP1);
  //       gzSubMesh.SetNormal(iv++, ignNormal);

  //       gzSubMesh.SetVertex(iv, ignP2);
  //       gzSubMesh.SetNormal(iv++, ignNormal);
  //     }
  //     this->data->isDirty = false;
  //   }
  // }

///////////////////////////////////////////////////////////////////////////////
// WavefieldSamplerPrivate

  /// \internal
  /// \brief Private data for the WavefieldSampler.
  class WavefieldSamplerPrivate
  {
    /// \brief The wavefield grid.
    public: std::shared_ptr<const Wavefield> wavefield;

    /// \brief A sample of the wavefield. This copy is located at the initial pose.
    public: std::shared_ptr<const Grid> initWaterPatch;

    /// \brief A sample of the wavefield. This copy is updated.
    public: std::shared_ptr<Grid> waterPatch;
  };

///////////////////////////////////////////////////////////////////////////////
// WavefieldSampler

  WavefieldSampler::~WavefieldSampler()
  {
  }

  WavefieldSampler::WavefieldSampler(
    std::shared_ptr<const Wavefield> _wavefield,
    std::shared_ptr<const Grid> _waterPatch
  ) : data(new WavefieldSamplerPrivate())
  {
    this->data->wavefield = _wavefield;
    this->data->initWaterPatch = _waterPatch;
    this->data->waterPatch.reset(new Grid(*_waterPatch));
  }

  std::shared_ptr<const Grid> WavefieldSampler::GetWaterPatch() const
  {
    return this->data->waterPatch;
  }

  void WavefieldSampler::ApplyPose(const ignition::math::Pose3d& _pose)
  {
    // @TODO_FRAGILE - Move to Grid as changing internal state
    // Apply pose to center
    const Point3& c0 = this->data->initWaterPatch->GetCenter();
    Point3 c1(c0.x() + _pose.Pos().X(), c0.y() + _pose.Pos().Y(), c0.z());
    this->data->waterPatch->SetCenter(c1);

    // Iterate over vertices
    auto& source = *this->data->initWaterPatch->GetMesh();
    auto& target = *this->data->waterPatch->GetMesh();
    for (
      auto&& it = std::make_pair(std::begin(source.vertices()), std::begin(target.vertices()));
      it.first != std::end(source.vertices()) && it.second != std::end(target.vertices());
      ++it.first, ++it.second)
    {
      auto& v0 = *it.first;
      auto& v1 = *it.second;
      const Point3& p0 = source.point(v0);

      // Transformation: slide the patch in the xy - plane only
      Point3 p1(p0.x() + _pose.Pos().X(), p0.y() + _pose.Pos().Y(), p0.z());
      target.point(v1) = p1;
    }
  }

  void WavefieldSampler::UpdatePatch()
  {
    // Direction of the line search (i.e. positive z-axis)
    Direction3 direction(0, 0, 1);

    // Update the water patch Mesh
    const auto& target = this->data->waterPatch->GetMesh();
    for (
      auto&& vb = std::begin(target->vertices());
      vb != std::end(target->vertices());
      ++vb
    )
    {
      auto& v1 = *vb;

      Point3& origin = target->point(v1);
      Point3 point = CGAL::ORIGIN;

      auto& wavefieldGrid = *this->data->wavefield->GetGrid();
      std::array<size_t, 3> cellIndex = { 0, 0, 0 };
      bool isFound = GridTools::FindIntersectionIndex(
        wavefieldGrid, origin.x(), origin.y(), cellIndex);
      if (!isFound)
      {
        // @DEBUG_INFO
        gzmsg << "origin:   " << origin << std::endl;
        gzerr << "Wavefield is too small" << std::endl;
        return;
      }
      isFound = GridTools::FindIntersectionGrid(
        wavefieldGrid, origin, direction, cellIndex, point);

      if (!isFound)
      {
        // @DEBUG_INFO
        gzmsg << "origin:   " << origin << std::endl;
        gzerr << "Wavefield is too small" << std::endl;
        return;
      }

      target->point(v1) = point;
    }
  }

  double WavefieldSampler::ComputeDepth(const Point3& _point) const
  {
    auto& grid = *this->data->waterPatch;
    return WavefieldSampler::ComputeDepth(grid, _point);

    // @TODO_EXPERIMENTAL - direct calculation (currently static only)
    // auto& waveParams = *this->data->wavefield->GetParameters();
    // return WavefieldSampler::ComputeDepthDirectly(waveParams, _point, 0.0);
  }

  double WavefieldSampler::ComputeDepth(
    const Grid& _patch,
    const Point3& _point
  )
  {
    // Calculate the depth
    Direction3 direction(0, 0, 1);
    Point3 wavePoint = CGAL::ORIGIN;
    std::array<size_t, 3> index;
    bool isFound = GridTools::FindIntersectionIndex(
      _patch, _point.x(), _point.y(), index);
    if (!isFound)
    {
      // @DEBUG_INFO
      gzerr << "point:  " << _point << std::endl;
      // _patch.DebugPrint();
      gzerr << "Water patch is too small" << std::endl;
      return 0;
    }
    isFound = GridTools::FindIntersectionGrid(
      _patch, _point, direction, index, wavePoint);
    if (!isFound)
    {
      // @DEBUG_INFO
      gzerr << "point:  " << _point << std::endl;
      // _patch.DebugPrint();
      gzerr << "Water patch is too small" << std::endl;
      return 0;
    }
    double h = wavePoint.z() - _point.z();
    return h;
  }

  double WavefieldSampler::ComputeDepthDirectly(
    const WaveParameters& _waveParams,
    const Point3& _point,
    double time
  )
  {
    // Struture for passing wave parameters to lambdas
    struct WaveParams
    {
      WaveParams(
        const std::vector<double>& _a,
        const std::vector<double>& _k,
        const std::vector<double>& _omega,
        const std::vector<double>& _phi,
        const std::vector<double>& _q,
        const std::vector<Vector2>& _dir) :
        a(_a), k(_k), omega(_omega), phi(_phi), q(_q), dir(_dir) {}

      const std::vector<double>& a;
      const std::vector<double>& k;
      const std::vector<double>& omega;
      const std::vector<double>& phi;
      const std::vector<double>& q;
      const std::vector<Vector2>& dir;
    };

    // Compute the target function and Jacobian. Also calculate pz,
    // the z-componen of the Gerstner wave, which we essentially get for free.
    auto wave_fdf = [=](auto x, auto p, auto t, auto& wp, auto& F, auto& J)
    {
      double pz = 0;
      F(0) = p.x() - x.x();
      F(1) = p.y() - x.y();
      J(0, 0) = -1;
      J(0, 1) =  0;
      J(1, 0) =  0;
      J(1, 1) = -1;
      const size_t n = wp.a.size();
      for (auto&& i=0; i<n; ++i)
      {
        const double dx = wp.dir[i].x();
        const double dy = wp.dir[i].y();
        const double q = wp.q[i];
        const double a = wp.a[i];
        const double k = wp.k[i];
        const double dot = x.x() * dx + x.y() * dy;
        const double theta = k * dot - wp.omega[i] * t;
        const double s = std::sin(theta);
        const double c = std::cos(theta);
        const double qakc = q * a * k * c;
        const double df1x = qakc * dx * dx;
        const double df1y = qakc * dx * dy;
        const double df2x = df1y;
        const double df2y = qakc * dy * dy;
        pz += a * c;
        F(0) += a * dx * s;
        F(1) += a * dy * s;
        J(0, 0) += df1x;
        J(0, 1) += df1y;
        J(1, 0) += df2x;
        J(1, 1) += df2y;
      }
      return pz;
    };

    // Simple multi-variate Newton solver - this version returns the z-component of the
    // wave field at the desired point p.
    auto solver = [=](auto& fdfunc, auto x0, auto p, auto t, auto& wp, auto tol, auto nmax)
    {
      int n = 0;
      double err = 1;
      double pz = 0;
      auto xn = x0;
      Eigen::Vector2d F;
      Eigen::Matrix2d J;
      while (std::abs(err) > tol && n < nmax)
      {
        pz = fdfunc(x0, p, t, wp, F, J);
        xn = x0 - J.inverse() * F;
        x0 = xn;
        err = F.norm();
        n++;
      }
      return pz;
    };

    // Set up parameter references
    WaveParams wp(
      _waveParams.Amplitude_V(),
      _waveParams.Wavenumber_V(),
      _waveParams.AngularFrequency_V(),
      _waveParams.Phase_V(),
      _waveParams.Steepness_V(),
      _waveParams.Direction_V()
    );

    // Tolerances etc.
    const double tol = 1.0E-10;
    const double nmax = 30;

    // Use the target point as the initial guess (this is within sum{amplitudes} of the solution)
    Eigen::Vector2d p2(_point.x(), _point.y());
    const double pz = solver(wave_fdf, p2, p2, time, wp, tol, nmax);
    const double h = pz - _point.z();
    return h;
  }

///////////////////////////////////////////////////////////////////////////////

} // namespace asv