#include "alchemist.h"
#include "data_stream.h"
#include <iostream>
#include <fstream>
#include <cstdlib>
#include <sys/stat.h>
#include <map>
#include <random>
#include <sstream>
#include <boost/asio.hpp>
#include <boost/chrono.hpp>
#include <boost/thread/thread.hpp>
#include <boost/tokenizer.hpp>
#include "arpackpp/arrssym.h"
#include "spdlog/spdlog.h"
#include <time.h>

using Eigen::MatrixXd;
using Eigen::VectorXd;

namespace alchemist {

using namespace El;

struct Driver {
  mpi::communicator world;
  DataInputStream input;
  DataOutputStream output;
  std::vector<WorkerInfo> workers;
  std::map<MatrixHandle, MatrixDescriptor> matrices;
  uint32_t nextMatrixId;
  std::shared_ptr<spdlog::logger> log;

  Driver(const mpi::communicator &world, std::istream &is, std::ostream &os, std::shared_ptr<spdlog::logger> log);
  void issue(const Command &cmd);
  MatrixHandle registerMatrix(size_t numRows, size_t numCols);
  void reshapeMatrix(MatrixHandle handle, size_t numRows, size_t numCols);
  int main();

  void handle_newMatrix();
  void handle_matrixMul();
  void handle_matrixDims();
  void handle_computeThinSVD();
  void handle_getMatrixRows();
  void handle_getTranspose();
  void handle_kmeansClustering();
  void handle_truncatedSVD();
  void handle_SkylarkKRR();
  void handle_SkylarkLSQR();
  void handle_FactorizedCGSolver();
  void handle_RandomFourierFeatures();
  void handle_ReadHDF5();
  void handle_normalizeMatInPlace();
};

Driver::Driver(const mpi::communicator &world, std::istream &is, std::ostream &os, std::shared_ptr<spdlog::logger> log) :
    world(world), input(is), output(os), log(log), nextMatrixId(42) {
}

void Driver::issue(const Command &cmd) {
  const Command *cmdptr = &cmd;
  mpi::broadcast(world, cmdptr, 0);
}

MatrixHandle Driver::registerMatrix(size_t numRows, size_t numCols) {
  MatrixHandle handle{nextMatrixId++};
  MatrixDescriptor info(handle, numRows, numCols);
  matrices.insert(std::make_pair(handle, info));
  return handle;
}

void Driver::reshapeMatrix(MatrixHandle handle, size_t numRows, size_t numCols) {
    matrices.erase(handle);
    MatrixDescriptor info(handle, numRows, numCols);
    matrices.insert(std::make_pair(handle, info));
}

int Driver::main() {

  // get WorkerInfo
  auto numWorkers = world.size() - 1;
  workers.resize(numWorkers);
  for(auto id = 0; id < numWorkers; ++id) {
    world.recv(id + 1, 0, workers[id]);
  }
  log->info("{} workers ready, sending hostnames and ports to Spark", numWorkers);

  // handshake
  ENSURE(input.readInt() == 0xABCD);
  ENSURE(input.readInt() == 0x1);
  output.writeInt(0xDCBA);
  output.writeInt(0x1);
  output.writeInt(numWorkers);
  for(auto id = 0; id < numWorkers; ++id) {
    output.writeString(workers[id].hostname);
    output.writeInt(workers[id].port);
  }
  output.flush();

  bool shouldExit = false;
  while(!shouldExit) {
    uint32_t typeCode = input.readInt();
    log->info("Received code {:#x}", typeCode);

    switch(typeCode) {
      // shutdown
      case 0xFFFFFFFF:
        shouldExit = true;
        issue(HaltCommand());
        output.writeInt(0x1);
        output.flush();
        break;

      // new matrix
      case 0x1:
        handle_newMatrix();
        break;

      // matrix multiplication
      case 0x2:
        handle_matrixMul();
        break;

      // get matrix dimensions
      case 0x3:
        handle_matrixDims();
        break;

      // return matrix to Spark
      case 0x4:
        handle_getMatrixRows();
        break;

      case 0x5:
        handle_computeThinSVD();
        break;

      case 0x6:
        handle_getTranspose();
        break;

      case 0x7:
        handle_kmeansClustering();
        break;

      case 0x8:
        handle_truncatedSVD();
        break;

      case 0x9:
        handle_SkylarkKRR();
        break;

      case 0x10:
        handle_SkylarkLSQR();
        break;

      case 0x11:
        handle_FactorizedCGSolver();
        break;

      case 0x12:
        handle_RandomFourierFeatures();
        break;

      case 0x13:
        handle_ReadHDF5();
        break;

      case 0x14:
        handle_normalizeMatInPlace();
        break;

      default:
        log->error("Unknown typeCode {#x}", typeCode);
        abort();
    }
    log->info("Waiting on next command");
  }

  // wait for workers to reach exit
  world.barrier();
  return EXIT_SUCCESS;
}

// A = USV
/*
void Driver::handle_computeLowrankSVD() {
  uint32_t inputMat = input.readInt();
  uint32_t kRank = input.readInt();
  uint32_t whichFactors = input.readInt(); // 0 = U, 1 = U and S, 2 = V and S, 3/default = U, S, and V

  MatrixHandle U;
  MatrixHandle S;
  MatrixHandle V;

  // TODO: register the U, S, V factors with false newmatrixcommands to track them
  switch(whichFactors) {
    case 0: U = MatrixHandle{nextMatrixId++};
            break;
    case 1: U = MatrixHandle{nextMatrixId++};
            S = MatrixHandle{nextMatrixId++};
            break;
    case 2: V = MatrixHandle{nextMatrixId++};
            S = MatrixHandle{nextMatrixId++};
            break;
    default: U = MatrixHandle{nextMatrixId++};
             S = MatrixHandle{nextMatrixId++};
             V = MatrixHandle{nextMatrixId++};
             break;
  }

  LowrankSVDCommand cmd(MatrixHandle{inputMat}, whichFactors, krank, U, S, V);
  issue(cmd);

  output.writeInt(0x1); // statusCode
  switch(whichFactors) {
    case 0: output.writeInt(U.id);
            break;
    case 1: output.writeInt(U.id);
            output.writeInt(S.id);
            break;
    case 2: output.writeInt(V.id);
            output.writeInt(S.id);
            break;
    default: output.writeInt(U.id);
             output.writeInt(S.id);
             output.writeInt(V.id);
             break;
  }
  output.flush();

  // wait for command to finish
  world.barrier();
  output.writeInt(0x1);
  output.flush();
}
*/

void Driver::handle_RandomFourierFeatures() {
    MatrixHandle A{input.readInt()};
    uint32_t numRandFeatures = input.readInt();
    double sigma = input.readDouble();
    uint32_t seed = input.readInt();

    auto numRows = matrices[A].numRows;
    uint32_t d = matrices[A].numCols;
    MatrixHandle X = registerMatrix(numRows, numRandFeatures);

    log->info("Computing Fourier Random Features on feature matrix {}");
    log->info("Input dim: {}, Output dim: {}, number of points: {}", d, numRandFeatures, numRows);

    RandomFourierFeaturesCommand cmd(A, X, numRandFeatures, sigma, seed);
    issue(cmd);

    world.barrier(); // wait for it to finish
    output.writeInt(0x1);
    output.writeInt(X.id);
    output.flush();
    log->info("Finished computing Fourier Random Features, stored matrix as {}", X);
}

void Driver::handle_normalizeMatInPlace() {
    MatrixHandle A{input.readInt()};
    log->info("Normalizing matrix {} in place by zero meaning the rows and setting the columns to have stdev 1", A);

    NormalizeMatInPlaceCommand cmd(A);
    issue(cmd);

    output.writeInt(0x1);
    log->info("Finished normalizing matrix {} in place", A);
}

void Driver::handle_ReadHDF5() {
    std::string fname = input.readString();
    std::string varname = input.readString();
    int colreplicas = input.readInt();
    MatrixHandle A = registerMatrix(0, 0);

    log->info("Reading variable {} from file {} and replicating columnwise {} times", varname, fname, colreplicas);
    ReadHDF5Command cmd(A, fname, varname, colreplicas);
    issue(cmd);

    El::Int numRows, numCols;
    world.recv(1, mpi::any_tag, numRows);
    world.recv(1, mpi::any_tag, numCols);
    world.barrier();

    reshapeMatrix(A, numRows, numCols);
    log->info("Loaded matrix is {}-by-{}", matrices[A].numRows, matrices[A].numCols);
    output.writeInt(0x1);
    output.writeInt(A.id);
    output.flush();
    log->info("Finished loading variable from HDF5 file");
}

void Driver::handle_FactorizedCGSolver() {
  MatrixHandle featureMat{input.readInt()};
  MatrixHandle targetMat{input.readInt()};
  double lambda = input.readDouble();
  uint32_t maxIters = input.readInt();

  auto numfeatures = matrices[featureMat].numCols;
  auto numtargets = matrices[targetMat].numCols;
  MatrixHandle coefs = registerMatrix(numfeatures, numtargets);

  log->info("Starting factorized CG solver on feature matrix {} and target matrix {}", featureMat, targetMat);
  log->info("lambda={}, maxIters={}", lambda, maxIters);

  FactorizedCGSolverCommand cmd(featureMat, targetMat, coefs, lambda, maxIters);
  issue(cmd);

  world.barrier(); // wait for it to finish
  output.writeInt(0x1);
  output.writeInt(coefs.id);
  output.flush();
  log->info("Finished CG solve, stored coefficient matrix as {}", coefs);
}

void Driver::handle_SkylarkKRR() {
    MatrixHandle featureMat{input.readInt()};
    MatrixHandle targetMat{input.readInt()};
    bool regression = input.readInt() > 0? true : false;
    uint32_t lossfunction = input.readInt();
    uint32_t regularizer = input.readInt();
    uint32_t kernel = input.readInt();
    double kernelparam = input.readDouble();
    double kernelparam2 = input.readDouble();
    double kernelparam3 = input.readDouble();
    double lambda = input.readDouble();
    uint32_t maxiter = input.readInt();
    double tolerance = input.readDouble();
    double rho = input.readDouble();
    uint32_t seed = input.readInt();
    uint32_t randomfeatures = input.readInt();
    uint32_t numfeaturepartitions = input.readInt();

    log->info("Starting Skylark's ADMM KRR solver on feature matrix {} and target matrix {}", featureMat, targetMat);
    log->info("<should list arguments here>");

    auto numfeatures = matrices[featureMat].numCols;
    auto numtargets = matrices[targetMat].numCols;
    MatrixHandle coefs = registerMatrix(numfeatures, numtargets);

    SkylarkKernelSolverCommand cmd(featureMat, targetMat, coefs, regression,
        lossfunction, regularizer, kernel, kernelparam, kernelparam2,
        kernelparam3, lambda, maxiter, tolerance, rho, seed, randomfeatures,
        numfeaturepartitions);
    issue(cmd);

    world.barrier(); // wait for it to finish
    log->info("Finished calling Skylark's ADMM Kernel solver for this KRR problem");
    output.writeInt(0x1);
    output.writeInt(coefs.id);
    output.flush();
    log->info("Finished KRR solve, stored coeffcient matrix as {}", coefs);
}

void Driver::handle_SkylarkLSQR() {
  MatrixHandle A{input.readInt()};
  MatrixHandle B{input.readInt()};
  double tolerance = input.readDouble();
  uint32_t iter_lim = input.readInt();

  auto p = matrices[A].numCols;
  auto m = matrices[B].numCols;
  MatrixHandle X = registerMatrix(p, m);

  log->info("Starting Skylark's LSQR solver on feature matrix {} and target matrix {}", A, B);
  log->info("Tolerance: {}", tolerance);
  log->info("Iteration Limit: {}", iter_lim);
  log->info("Result will be a {}-by-{} matrix", p, m);

  SkylarkLSQRSolverCommand cmd(A, B, X, tolerance, iter_lim);
  issue(cmd);

  world.barrier(); // wait for it to finish
  log->info("Finished call to Skylark's LSQR solver");
  output.writeInt(0x1);
  output.writeInt(X.id);
  output.flush();
  log->info("Finished LSQR computation");
}

// TODO: the cluster centers should be stored locally on driver and reduced/broadcasted. the current
// way of updating kmeans centers is ridiculous
// TODO: currently only implements kmeans||
void Driver::handle_kmeansClustering() {
  MatrixHandle inputMat{input.readInt()};
  uint32_t numCenters = input.readInt();
  uint32_t maxnumIters = input.readInt(); // how many iteration of Lloyd's algorithm to use
  uint32_t initSteps = input.readInt(); // number of initialization steps to use in kmeans||
  double changeThreshold = input.readDouble(); // if all the centers change by Euclidean distance less than changeThreshold, then we stop the iterations
  uint32_t method = input.readInt(); // which initialization method to use to choose initial cluster center guesses
  uint64_t seed = input.readLong(); // randomness seed used in driver and workers

  log->info("Starting K-means on matrix {}", inputMat);
  log->info("numCenters = {}, maxnumIters = {}, initSteps = {}, changeThreshold = {}, method = {}, seed = {}",
      numCenters, maxnumIters, initSteps, changeThreshold, method, seed);

  if (changeThreshold < 0.0 || changeThreshold > 1.0) {
    log->error("Unreasonable change threshold in k-means: {}", changeThreshold);
    abort();
  }
  if (method != 1) {
    log->warn("Sorry, only k-means|| initialization has been implemented, so ignoring your choice of method {}", method);
  }

  auto n = matrices[inputMat].numRows;
  auto d = matrices[inputMat].numCols;
  MatrixHandle centersHandle = this->registerMatrix(numCenters, d);
  MatrixHandle assignmentsHandle = this->registerMatrix(n, 1);
  KMeansCommand cmd(inputMat, numCenters, method, initSteps, changeThreshold, seed, centersHandle, assignmentsHandle);
  issue(cmd); // initial call initializes stuff and waits for next command

  /******** START of kmeans|| initialization ********/
  std::mt19937 gen(seed);
  std::uniform_int_distribution<unsigned long> dis(0, n-1);
  uint32_t rowidx = dis(gen);
  std::vector<double> initialCenter(d);

  mpi::broadcast(world, rowidx, 0); // tell the workers which row to use as initialization in kmeans||
  world.barrier(); // wait for workers to return oversampled cluster centers and sizes

  std::vector<uint32_t> clusterSizes;
  std::vector<MatrixXd> initClusterCenters;
  world.recv(1, mpi::any_tag, clusterSizes);
  world.recv(1, mpi::any_tag, initClusterCenters);
  world.barrier();

  log->info("Retrieved the k-means|| oversized set of potential cluster centers");
  log->debug("{}", initClusterCenters);

  // use kmeans++ locally to find the initial cluster centers
  std::vector<double> weights;
  weights.reserve(clusterSizes.size());
  std::for_each(clusterSizes.begin(), clusterSizes.end(), [&weights](const uint32_t & cnt){ weights.push_back((double) cnt); });
  MatrixXd clusterCenters(numCenters, d);

  kmeansPP(gen(), initClusterCenters, weights, clusterCenters, 30); // same number of maxIters as spark kmeans

  log->info("Ran local k-means on the driver to determine starting cluster centers");
  log->debug("{}", clusterCenters);

  mpi::broadcast(world, clusterCenters.data(), numCenters*d, 0);
  /******** END of kMeans|| initialization ********/

  /******** START of Lloyd's algorithm iterations ********/
  double percentAssignmentsChanged = 1.0;
  bool centersMovedQ = true;
  uint32_t numChanged = 0;
  uint32_t numIters = 0;
  std::vector<uint32_t> parClusterSizes(numCenters);
  std::vector<uint32_t> zerosVector(numCenters);

  for(uint32_t clusterIdx = 0; clusterIdx < numCenters; clusterIdx++)
    zerosVector[clusterIdx] = 0;

  uint32_t command = 1; // do another iteration
  while (centersMovedQ && numIters++ < maxnumIters)  {
    log->info("Starting iteration {} of Lloyd's algorithm, {} percentage changed in last iter",
        numIters, percentAssignmentsChanged*100);
    numChanged = 0;
    for(uint32_t clusterIdx = 0; clusterIdx < numCenters; clusterIdx++)
      parClusterSizes[clusterIdx] = 0;
    command = 1; // do a basic iteration
    mpi::broadcast(world, command, 0);
    mpi::reduce(world, (uint32_t) 0, numChanged, std::plus<int>(), 0);
    mpi::reduce(world, zerosVector.data(), numCenters, parClusterSizes.data(), std::plus<uint32_t>(), 0);
    world.recv(1, mpi::any_tag, centersMovedQ);
    percentAssignmentsChanged = ((double) numChanged)/n;

    for(uint32_t clusterIdx = 0; clusterIdx < numCenters; clusterIdx++) {
      if (parClusterSizes[clusterIdx] == 0) {
        // this is an empty cluster, so randomly pick a point in the dataset
        // as that cluster's centroid
        centersMovedQ = true;
        command = 2; // reinitialize this cluster center
        uint32_t rowIdx = dis(gen);
        mpi::broadcast(world, command, 0);
        mpi::broadcast(world, clusterIdx, 0);
        mpi::broadcast(world, rowIdx, 0);
        world.barrier();
      }
    }

  }
  command = 0xf; // terminate and finalize the k-means centers and assignments as distributed matrices
  mpi::broadcast(world, command, 0);
  double objVal = 0.0;
  mpi::reduce(world, 0.0, objVal, std::plus<double>(), 0);
  world.barrier();

  /******** END of Lloyd's iterations ********/

  log->info("Finished Lloyd's algorithm: took {} iterations, final objective value {}", numIters, objVal);
  output.writeInt(0x1);
  output.writeInt(assignmentsHandle.id);
  output.writeInt(centersHandle.id);
  output.writeInt(numIters);
  output.flush();
}

void Driver::handle_getTranspose() {
  MatrixHandle inputMat{input.readInt()};
  log->info("Constructing the transpose of matrix {}", inputMat);

  auto numRows = matrices[inputMat].numCols;
  auto numCols = matrices[inputMat].numRows;
  MatrixHandle transposeHandle = registerMatrix(numRows, numCols);
  TransposeCommand cmd(inputMat, transposeHandle);
  issue(cmd);

  world.barrier(); // wait for command to finish
  output.writeInt(0x1);
  output.writeInt(transposeHandle.id);
  log->info("Wrote handle for transpose");
  output.writeInt(0x1);
  output.flush();
}

// CAVEAT: Assumes tall-and-skinny for now, doesn't allow many options for controlling
// LIMITATIONS: assumes V small enough to fit on one machine (so can use ARPACK instead of PARPACK), but still distributes U,S,V and does distributed computations not needed
void Driver::handle_truncatedSVD() {
  log->info("Starting truncated SVD computation");
  MatrixHandle inputMat{input.readInt()};
  uint32_t k = input.readInt();
  uint32_t method = input.readInt();

  int LOCALEIGS = 0; // TODO: make these an enumeration, and global to Alchemist
  int LOCALEIGSPRECOMPUTE = 1;
  int DISTEIGS = 2;

  MatrixHandle UHandle{nextMatrixId++};
  MatrixHandle SHandle{nextMatrixId++};
  MatrixHandle VHandle{nextMatrixId++};

  auto m = matrices[inputMat].numRows;
  auto n = matrices[inputMat].numCols;
  TruncatedSVDCommand cmd(inputMat, UHandle, SHandle, VHandle, k, method);
  issue(cmd);

  if (method == DISTEIGS) {
      log->info("using distributed mat-vec prods against A, then A tranpose");
  }
  if (method == LOCALEIGS) {
      log->info("using local mat-vec prods computed on the fly against the local Gramians");
  }
  if (method == LOCALEIGSPRECOMPUTE) {
      log->info("using local mat-vec prods against the precomputed local Gramians");
  }

  ARrcSymStdEig<double> prob(n, k, "LM");
  uint32_t command;
  std::vector<double> zerosVector(n);
  for(uint32_t idx = 0; idx < n; idx++)
    zerosVector[idx] = 0;

  int iterNum = 0;

  while (!prob.ArnoldiBasisFound()) {
    prob.TakeStep();
    ++iterNum;
    if(iterNum % 20 == 0) {
        log->info("Computed {} mv products", iterNum);
    }
    if (prob.GetIdo() == 1 || prob.GetIdo() == -1) {
      command = 1;
      mpi::broadcast(world, command, 0);
      if (method == LOCALEIGS || method == LOCALEIGSPRECOMPUTE) {
          mpi::broadcast(world, prob.GetVector(), n, 0);
          mpi::reduce(world, zerosVector.data(), n, prob.PutVector(), std::plus<double>(), 0);
      }
      if (method == DISTEIGS) {
          world.send(1, 0, prob.GetVector(), n);
          world.recv(1, 0, prob.PutVector(), n);
      }
    }
  }

  prob.FindEigenvectors();
  uint32_t nconv = prob.ConvergedEigenvalues();
  uint32_t niters = prob.GetIter();
  log->info("Done after {} Arnoldi iterations, converged to {} eigenvectors of size {}", niters, nconv, n);

  //NB: it may be the case that n*nconv > 4 GB, then have to be careful!
  // assuming tall and skinny A for now
  MatrixXd rightVecs(n, nconv);
  log->info("Allocated matrix for right eigenvectors of A'*A");
  // Eigen uses column-major layout by default!
  for(uint32_t idx = 0; idx < nconv; idx++)
    std::memcpy(rightVecs.col(idx).data(), prob.RawEigenvector(idx), n*sizeof(double));
  log->info("Copied right eigenvectors into allocated storage");

  // Populate U, V, S
  command = 2;
  mpi::broadcast(world, command, 0);
  mpi::broadcast(world, nconv, 0);
  log->info("Broadcasted command and number of converged eigenvectors");
  mpi::broadcast(world, rightVecs.data(), n*nconv, 0);
  log->info("Broadcasted right eigenvectors");
  mpi::broadcast(world, prob.RawEigenvalues(), nconv, 0);
  log->info("Broadcasted eigenvalues");

  MatrixDescriptor Uinfo(UHandle, m, nconv);
  MatrixDescriptor Sinfo(SHandle, nconv, 1);
  MatrixDescriptor Vinfo(VHandle, n, nconv);
  ENSURE(matrices.insert(std::make_pair(UHandle, Uinfo)).second);
  ENSURE(matrices.insert(std::make_pair(SHandle, Sinfo)).second);
  ENSURE(matrices.insert(std::make_pair(VHandle, Vinfo)).second);

  log->info("Waiting on workers to store U,S,V");

  world.barrier();
  log->info("Writing ok status followed by U,S,V handles");
  output.writeInt(0x1);
  output.writeInt(UHandle.id);
  output.writeInt(SHandle.id);
  output.writeInt(VHandle.id);
  output.flush();
}

void Driver::handle_computeThinSVD() {
  MatrixHandle inputMat{input.readInt()};

  // this needs to be done automatically rather than hand-coded. e.g. what if
  // we switch to determining rank by sing-val thresholding instead of doing thin SVD?
  auto m = matrices[inputMat].numRows;
  auto n = matrices[inputMat].numCols;
  auto k = std::min(m,n);
  MatrixHandle Uhandle = registerMatrix(m, k);
  MatrixHandle Shandle = registerMatrix(k, 1);
  MatrixHandle Vhandle = registerMatrix(n, k);
  ThinSVDCommand cmd(inputMat, Uhandle, Shandle, Vhandle);
  issue(cmd);

  output.writeInt(0x1); // statusCode
  output.writeInt(Uhandle.id);
  output.writeInt(Shandle.id);
  output.writeInt(Vhandle.id);
  output.flush();

  // wait for command to finish
  world.barrier();
  log->info("Done with SVD computation");
  output.writeInt(0x1);
  output.flush();
}

void Driver::handle_matrixMul() {
  MatrixHandle matA{input.readInt()};
  MatrixHandle matB{input.readInt()};
  log->info("Multiplying matrices {} and {}", matA, matB);

  auto numRows = matrices[matA].numRows;
  auto numCols = matrices[matB].numCols;
  MatrixHandle destHandle = registerMatrix(numRows, numCols);
  MatrixMulCommand cmd(destHandle, matA, matB);
  issue(cmd);

  // tell spark id of resulting matrix
  output.writeInt(0x1); // statusCode
  output.writeInt(destHandle.id);
  output.flush();

  // wait for it to finish
  world.barrier();
  output.writeInt(0x1);
  output.flush();
}

void Driver::handle_matrixDims() {
  MatrixHandle matrixHandle{input.readInt()};
  log->info("Looking up dimensions for matrix {}", matrixHandle.id);
  auto info = matrices[matrixHandle];
  log->info("Returning dimensions for matrix {}: {}-by-{}", matrixHandle.id, info.numRows, info.numCols);
  output.writeInt(0x1);
  output.writeLong(info.numRows);
  output.writeLong(info.numCols);
  output.flush();

}

void Driver::handle_getMatrixRows() {
  MatrixHandle handle{input.readInt()};
  uint64_t layoutLen = input.readLong();
  std::vector<uint32_t> layout;
  layout.reserve(layoutLen);
  for(uint64_t part = 0; part < layoutLen; ++part) {
    layout.push_back(input.readInt());
  }
  log->info("Returning matrix {} to Spark", handle);

  MatrixGetRowsCommand cmd(handle, layout);
  issue(cmd);

  // tell Spark to start asking for rows
  output.writeInt(0x1);
  output.flush();

  // wait for it to finish
  world.barrier();
  output.writeInt(0x1);
  output.flush();
}

void Driver::handle_newMatrix() {
  // read args
  uint64_t numRows = input.readLong();
  uint64_t numCols = input.readLong();

  // assign id and notify workers
  MatrixHandle handle = registerMatrix(numRows, numCols);
  NewMatrixCommand cmd(matrices[handle]);
  log->info("Recieving new matrix {}, with dimensions {}x{}", handle, numRows, numCols);
  issue(cmd);

  output.writeInt(0x1);
  output.writeInt(handle.id);
  output.flush();

  // tell spark which worker expects each row
  std::vector<int> rowWorkerAssignments(numRows, 0);
  std::vector<uint64_t> rowsOnWorker;
  std::vector<uint32_t> numRowsOnEachWorker(world.size() -1, 0);
  for(int workerIdx = 1; workerIdx < world.size(); workerIdx++) {
    world.recv(workerIdx, 0, rowsOnWorker);
    numRowsOnEachWorker[workerIdx - 1] = rowsOnWorker.size();
    world.barrier();
    for(auto rowIdx: rowsOnWorker) {
      rowWorkerAssignments[rowIdx] = workerIdx;
    }
  }

  log->info("Sending list of which worker each row should go to");
  output.writeInt(0x1); // statusCode
  std::stringstream ss;
  for(auto workerIdx: rowWorkerAssignments) {
    output.writeInt(workerIdx);
    ALCHEMIST_TRACE(ss << workerIdx << ' ');
  }
  ALCHEMIST_TRACE(log->info(ss.str().c_str()));

  for(uint32_t workerIdx = 1; workerIdx < world.size(); workerIdx++)
    log->info("Worker with rank {} will receive {} rows", workerIdx, numRowsOnEachWorker[workerIdx-1]);

  output.flush();

  log->info("Waiting for spark to finish sending data to the workers");
  world.barrier();
  output.writeInt(0x1);  // statusCode
  output.flush();
  log->info("Entire matrix has been received");
}

inline bool exist_test (const std::string& name) {
    struct stat buffer;
    return (stat(name.c_str(), &buffer) == 0);
}

int driverMain(const mpi::communicator &world, int argc, char *argv[]) {
  //log to console as well as file (single-threaded logging)
  //TODO: allow to specify log directory, log level, etc.
  time_t rawtime = time(0);
  struct tm * timeinfo = localtime(&rawtime);
  char buf[80];
  strftime(buf, 80, "%F-%T", timeinfo);
  std::shared_ptr<spdlog::logger> log;
  std::vector<spdlog::sink_ptr> sinks;
  sinks.push_back(std::make_shared<spdlog::sinks::ansicolor_stderr_sink_st>());
  sinks.push_back(std::make_shared<spdlog::sinks::simple_file_sink_st>(str(format("driver-%s.log") % buf)));
  log = std::make_shared<spdlog::logger>("driver", std::begin(sinks), std::end(sinks));
  //log->flush_on(spdlog::level::warn); // flush whenever warning or more critical message is logged
  //log->set_level(spdlog::level::info); // only log stuff at or above info level, for production
  log->flush_on(spdlog::level::info); // flush always, for debugging
  log->info("Started Driver");
  log->info("Max number of OpenMP threads: {}", omp_get_max_threads());

  char machine[255];
  char port[255];

  if (argc == 3) { // we are on a non-NERSC system, so passed in Spark driver machine name and port
      log->info("Non-NERSC system assumed");
      log->info("Connecting to Spark executor at {}:{}", argv[1], argv[2]);
      std::strcpy(machine, argv[1]);
      std::strcpy(port, argv[2]);
  } else { // assume we are on NERSC, so look in a specific location for a file containing the machine name and port
      char const* tmp = std::getenv("SPARK_WORKER_DIR");
      std::string sockPath;
      if (tmp == NULL) {
          log->info("Couldn't find the SPARK_WORKER_DIR variable");
          world.abort(1);
      } else {
        sockPath = std::string(tmp) + "/connection.info";
      }
      log->info("NERSC system assumed");
      log->info("Searching for connection information in file {}", sockPath);

      while(!exist_test(sockPath)) {
          boost::this_thread::sleep_for(boost::chrono::milliseconds(50));
      }
      // now wait for a while for the connection file to be completely written, hopefully is enough time
      // TODO: need a more robust way of ensuring this is the case
      boost::this_thread::sleep_for(boost::chrono::milliseconds(500));

      std::string sockSpec;
      std::ifstream infile(sockPath);
      std::getline(infile, sockSpec);
      infile.close();
      boost::tokenizer<> tok(sockSpec);
      boost::tokenizer<>::iterator iter=tok.begin();
      std::string machineName = *iter;
      std::string portName = *(++iter);

      log->info("Connecting to Spark executor at {}:{}", machineName, portName);
      strcpy(machine, machineName.c_str());
      strcpy(port, portName.c_str());
  }

  using boost::asio::ip::tcp;
  boost::asio::ip::tcp::iostream stream(machine, port);
  ENSURE(stream);
  stream.rdbuf()->non_blocking(false);
  auto result = Driver(world, stream, stream, log).main();
  return result;
}

} // namespace alchemist


/*
 * This file is a part of
 *
 * ============================================
 * ###   Pteros molecular modeling library  ###
 * ============================================
 *
 * (C) 2009-2018, Semen Yesylevskyy
 *
 * All works, which use Pteros, should cite the following papers:
 *
 *  1.  Semen O. Yesylevskyy, "Pteros 2.0: Evolution of the fast parallel
 *      molecular analysis library for C++ and python",
 *      Journal of Computational Chemistry, 2015, 36(19), 1480–1488.
 *      doi: 10.1002/jcc.23943.
 *
 *  2.  Semen O. Yesylevskyy, "Pteros: Fast and easy to use open-source C++
 *      library for molecular analysis",
 *      Journal of Computational Chemistry, 2012, 33(19), 1632–1636.
 *      doi: 10.1002/jcc.22989.
 *
 * This is free software distributed under Artistic License:
 * http://www.opensource.org/licenses/artistic-license-2.0.php
 *
*/

#include "pteros/extras/membrane.h"
#include "pteros/core/pteros_error.h"
#include "pteros/core/distance_search.h"
#include "pteros/core/utilities.h"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <fstream>
#include "voro++.hh"

#ifndef M_PI
    #define M_PI 3.14159265358979323846
#endif

#ifndef M_PI_2
    #define M_PI_2 9.86960440108935712052951
#endif

using namespace std;
using namespace pteros;
using namespace Eigen;

Membrane::Membrane(System *sys, const std::vector<Lipid_descr> &species): system(sys), lipid_species(species)
{
    log = create_logger("membrane");

    Selection all_mid_sel(*system);

    // Creating selections
    for(auto& sp: lipid_species){
        vector<Selection> res;
        system->select(sp.whole_sel_str).split_by_residue(res);
        log->info("Lipid {}: {}", sp.name,res.size());
        for(auto& lip: res){
            auto mol = Lipid(lip,sp);
            mol.set_markers();
            all_mid_sel.append(mol.mid_sel.index(0));
            lipids.push_back(mol);
            index_map[mol.mid_sel.index(0)] = lipids.size()-1;
        }
    }

    // Compute connectivity
    std::vector<Selection> leafs;
    all_mid_sel.split_by_connectivity(2.0,leafs,true);
    leaflets.resize(leafs.size());
    leaflets_sel.resize(leafs.size());
    // Set lipid leaflets
    for(int i=0;i<leafs.size();++i){
        leaflets_sel[i].set_system(*system);
        for(int a=0; a<leafs[i].size(); ++a){
            int ind = leafs[i].index(a);
            lipids[index_map[ind]].leaflet = i;
            leaflets[i].push_back(index_map[ind]);
            leaflets_sel[i].append(ind);
        }
    }

    // Print statictics
    log->info("Total number of lipids: {}",lipids.size());
    log->info("Number of leaflets: {}",leafs.size());
}

void fit_quad_surface(const MatrixXf& coord,
                      Matrix<float,6,1>& res,
                      Vector3f& smoothed_first_point,
                      float& rms){
    int N = coord.cols();
    // We fit with polynomial fit = A*x^2 + B*y^2 + C*xy + D*x + E*y + F
    // Thus we need a linear system of size 6
    Matrix<float,6,6> m;
    Matrix<float,6,1> rhs; // Right hand side and result

    // Functions returning powers after coefficients (x^2, y^2, x*y, x, y, 1)
    vector<std::function<float(int)> > coef(6);
    coef[0] = [&coord](int j){return coord.col(j)(0)*coord.col(j)(0);};
    coef[1] = [&coord](int j){return coord.col(j)(1)*coord.col(j)(1);};
    coef[2] = [&coord](int j){return coord.col(j)(0)*coord.col(j)(1);};
    coef[3] = [&coord](int j){return coord.col(j)(0);};
    coef[4] = [&coord](int j){return coord.col(j)(1);};
    coef[5] = [&coord](int j){return 1.0;};

    // Now form the matrix
    m.fill(0.0);
    rhs.fill(0.0);

    for(int r=0;r<6;++r){ //rows
        for(int c=0;c<6;++c){ //columns
            m(r,c) = 0.0;
            for(int j=0;j<N;++j){
                m(r,c) += coef[r](j)*coef[c](j);
            }
        }
        // Now rhs
        for(int j=0;j<N;++j){
            rhs(r) += coef[r](j)*coord.col(j)(2);
        }
    }

    // Now solve
    res = m.colPivHouseholderQr().solve(rhs);

    // Compute RMS of fitting
    rms = 0.0;
    for(int j=0;j<N;++j){
        float fit = 0.0;
        for(int r=0;r<6;++r) fit += coef[r](j)*res[r];
        rms += pow(coord.col(j)(2)-fit,2);
    }
    rms = sqrt(rms/float(N));

    // Get smoothed surface point
    smoothed_first_point.fill(0.0);
    for(int r=0;r<6;++r) smoothed_first_point(2) += coef[r](0)*res[r];
}

float compute_area(const Eigen::MatrixXf& coord, double dist, std::vector<int>& neib){
    // Perform Voronoi computation in the tangent plane
    using namespace voro;

    voronoicell_neighbor c;

    // Z dimension of cotainer is 1, so volume=area
    container con(-dist,dist,
                  -dist,dist,
                  -0.5,0.5,
                  10,10,10,false,false,false,8);

    particle_order po;
    // Cycle over particles, first particle is our target
    for(int i=0;i<coord.cols();++i){
        con.put(po, i, coord.col(i)(0), coord.col(i)(1), 0.0); // Z is zero
    }
    c_loop_order clo(con,po);
    if(! clo.start()) return -1.0; // Get first point, which is our target
    if(! con.compute_cell(c,clo) ) return -1.0;

    // Check if any of vertices is on outer bounding box and skip cell if they are
    vector<double> vert;
    c.vertices(0,0,0,vert);
    for(int i=0; i<vert.size(); i+=3){
        //     X                         Y
        if(abs(abs(vert[i])-dist)<=1.0e-5 || abs(abs(vert[i+1])-dist)<=1.0e-5){
            //cout << "| " << vert[i] << " " << vert[i+1] << " " << vert[i+2] << " " << abs(vert[i])-dist << endl;
            return -1.0;
        }
    }

    // Extract neighbors of this particle
    std::vector<int> n;
    c.neighbors(n);
    // Filter out negatives
    neib.clear();
    for(int i: n)
        if(i>0) neib.push_back(i);

/*
    cout << "Npoints: " << coord.cols() << endl;
    // Output the particle positions in gnuplot format
    con.draw_particles("points.gnu");

    // Output the Voronoi cells in gnuplot format
    con.draw_cells_gnuplot("points_v.gnu");
*/
    return c.volume();
}

void get_curvature(const Matrix<float,6,1>& coeffs, float& gaussian_curvature, float& mean_curvature){
    /* Compute the curvatures

        First fundamental form:  I = E du^2 + 2F du dv + G dv^2
        E= r_u dot r_u, F= r_u dot r_v, G= r_v dot r_v

        For us parametric variables (u,v) are just (x,y) in local space.
        Derivatives:
            r_u = {1, 0, 2Ax+Cy+D}
            r_v = {0, 1, 2By+Cx+E}

        In central point x=0, y=0 so:
            r_u={1,0,D}
            r_v={0,1,E}

        Thus: E_ =1+D^2; F_ = D*E; G_ = 1+E^2;

        Second fundamental form: II = L du2 + 2M du dv + N dv2
        L = r_uu dot n, M = r_uv dot n, N = r_vv dot n

        Normal is just  n = {0, 0, 1}
        Derivatives:
            r_uu = {0, 0, 2A}
            r_uv = {0 ,0, C}
            r_vv = {0, 0, 2B}

        Thus: L_ = 2A; M_ = C; N_ = 2B;
        */

    float E_ = 1.0+coeffs[3]*coeffs[3];
    float F_ = coeffs[3]*coeffs[4];
    float G_ = 1.0+coeffs[4]*coeffs[4];

    float L_ = 2.0*coeffs[0];
    float M_ = coeffs[2];
    float N_ = 2.0*coeffs[1];

    //Curvatures:
    gaussian_curvature = (L_*N_-M_*M_)/(E_*G_-F_*F_);
    mean_curvature = 0.5*(E_*N_-2.0*F_*M_+G_*L_)/(E_*G_-F_*F_);
}

void Membrane::compute_properties(float d, bool use_external_normal, Vector3f_const_ref external_pivot, Vector3i_const_ref external_dist_dim)
{
    // Clear everything
    neighbor_pairs.clear();

    Eigen::Matrix3f tr, tr_inv;

    // Set markers for all lipids
    // This unwraps each lipid
    for(auto& l: lipids) l.set_markers();

    // Compute per leaflet properties
    for(int l=0;l<leaflets.size();++l){
        // Get connectivity in this leaflet
        vector<Vector2i> bon;
        search_contacts(d,leaflets_sel[l],bon,false,true);

        // Convert the list of bonds to convenient form
        // atom ==> 1 2 3...
        vector<vector<int> > conn(leaflets_sel[l].size());
        for(int i=0;i<bon.size();++i){
            conn[bon[i](0)].push_back(bon[i](1));
            conn[bon[i](1)].push_back(bon[i](0));
        }

        // Process lipids from this leaflet
        for(int i=0;i<leaflets[l].size();++i){
            // Find current lipid
            Lipid& lip = lipids[leaflets[l][i]];

            // Save local selection for this lipid
            lip.local_sel = leaflets_sel[l].select(conn[i]);

            if(lip.local_sel.size()==0){
                log->warn("Empty locality of lipid {} in leaflet {}! Skipped.",i,l);
                continue;
            }

            // Local coordinate axes
            Matrix3f axes;

            //-----------------------------------
            // Compute normal
            //-----------------------------------

            Vector3f normal;

            if(use_external_normal){
                // Compute external normal as a vector from pivot to COM of mid_sel (set as marker currently)
                // over specified dimensions
                axes.col(2) = (lip.mid_marker-external_pivot).array()*external_dist_dim.cast<float>().array();

                // Find two vectors perpendicular to normal
                if( axes.col(2).dot(Vector3f(1,0,0)) != 1.0 ){
                    axes.col(0) = axes.col(2).cross(Vector3f(1,0,0));
                } else {
                    axes.col(0) = axes.col(2).cross(Vector3f(0,1,0));
                }
                axes.col(1) = axes.col(2).cross(axes.col(0));
            } else {
                // Compute normals from inertia axes
                // Create selection for locality of this lipid including itself
                Selection local_self(lip.local_sel);

                local_self.append(leaflets_sel[l].index(i)); // Add central atom

                // Get inertial axes
                Vector3f moments;
                local_self.inertia(moments,axes,fullPBC); // Have to use periodic variant
                // axes.col(2) will be a normal
            }

            normal = axes.col(2);

            // transformation matrix to local basis
            for(int j=0;j<3;++j) tr.col(j) = axes.col(j).normalized();
            tr_inv = tr.inverse();

            // Need to check direction of the normal
            float ang = angle_between_vectors(normal, lip.head_marker-lip.tail_marker);
            if(ang < M_PI_2){
                lip.normal = normal;
                lip.tilt = ang;
            } else {
                lip.normal = -normal;
                lip.tilt = M_PI-ang;
            }
            lip.normal.normalized();

            //-----------------------------------
            // Smooth and find local curvatures
            //-----------------------------------

            // Create array of local points in local basis
            MatrixXf coord(3,lip.local_sel.size()+1);
            Vector3f c0 = lip.mid_marker; // Real coord of central point - the marker
            coord.col(0) = Vector3f::Zero(); // Local coord of central point is zero
            for(int j=0; j<lip.local_sel.size(); ++j)
                coord.col(j+1) = tr_inv * system->box(0).shortest_vector(c0,lip.local_sel.xyz(j));

            // Fit a quad surface
            Matrix<float,6,1> res;
            Vector3f sm; // smoothed coords of central point
            fit_quad_surface(coord,res,sm,lip.quad_fit_rms);
            // Compute curvatures using quad fit coeefs
            get_curvature(res, lip.gaussian_curvature, lip.mean_curvature);

            // Get smoothed surface point in lab coords
            lip.smoothed_mid_xyz = tr*sm + lip.mid_marker;

            //-----------------------------------
            // Area and neighbours
            //-----------------------------------

            vector<int> neib;
            lip.area = compute_area(coord,10.0,neib);
            // Save coordination number of the lipid
            lip.coord_number = neib.size();

            // Add neighbor pairs. Only use nearest neighbor lipids from area computation
            // use only i<j to avoid adding duplicates
            // If area is -1 skip since this lipid has weird surrounding
            if(lip.area>0){
                int cur_ind = index_map[leaflets_sel[l].index(i)];
                for(int j=0; j<neib.size(); ++j){
                    int n = index_map[lip.local_sel.index(neib[j]-1)];
                    if(cur_ind<n) neighbor_pairs.emplace_back(cur_ind,n);
                }
            }

        } // Over lipids in leaflet

    } // over leaflets

    //-----------------------------------
    // Splay and triangilation
    //-----------------------------------

    // Go over neighbor pairs and compute mean splay
    // Also form neigbhor array as i ==> 1,2,3...
    splay.resize(neighbor_pairs.size());
    neighbors.resize(lipids.size());
    for(int i=0;i<neighbor_pairs.size();++i){
        neighbors[neighbor_pairs[i](0)].push_back(neighbor_pairs[i](1));
        neighbors[neighbor_pairs[i](1)].push_back(neighbor_pairs[i](0));
        // We have to use periodic distance since atoms could be from different images
        Lipid& lip1 = lipids[neighbor_pairs[i](0)];
        Lipid& lip2 = lipids[neighbor_pairs[i](1)];
        Vector3f& n1 = lip1.normal;
        Vector3f& n2 = lip2.normal;

        if(n1.dot(n2)<0) continue;

        auto x= system->box(0).shortest_vector(lip1.mid_marker, lip2.mid_marker);
        float d = x.norm();
        x = x-x.dot(n1)*n1;
        x.normalize();
        Vector3f v1 = (lip1.head_marker-lip1.tail_marker).normalized();
        Vector3f v2 = (lip2.head_marker-lip2.tail_marker).normalized();

        splay[i] = {
                    neighbor_pairs[i](0),
                    neighbor_pairs[i](1),
                    (v2.dot(x)-v1.dot(x)-n2.dot(x)+n1.dot(x))/d
                   };
    }

    //-----------------------------------
    // Order parameter
    //-----------------------------------

    for(auto& lip: lipids){
        // Compute Sz order parameter if the tails are provided
        for(int t=0; t<lip.tail_carbon_indexes.size(); ++t){
            // Go over atoms in tail t
            for(int at=1; at<lip.tail_carbon_indexes[t].size()-1; ++at){
                // Vector from at+1 to at-1
                auto coord1 = system->xyz(lip.tail_carbon_indexes[t][at+1]);
                auto coord2 = system->xyz(lip.tail_carbon_indexes[t][at-1]);
                float ang = angle_between_vectors(coord1-coord2,lip.normal);
                lip.order[t][at-1] = 1.5*pow(cos(ang),2)-0.5;
            }
        }
    }

}

string tcl_arrow(Vector3f_const_ref p1, Vector3f_const_ref p2, float r, string color){
    stringstream ss;
    Vector3f p = (p2-p1)*0.8+p1;
    ss << "draw color " << color << endl;

    ss << "draw cylinder \"" << p1.transpose()*10.0 << "\" ";
    ss << "\"" << p.transpose()*10.0 << "\" radius " << r << endl;

    ss << "draw cone \"" << p.transpose()*10.0 << "\" ";
    ss << "\"" << p2.transpose()*10.0 << "\" radius " << r*3.0 << endl;

    return ss.str();
}

void Membrane::write_vmd_arrows(const string &fname)
{
    ofstream f(fname);
    for(auto& lip: lipids){
        f << tcl_arrow(lip.mid_sel.center(true),lip.mid_sel.center(true)+lip.normal,0.2,"green");
        f << tcl_arrow(lip.tail_sel.center(true),lip.head_sel.center(true),0.2,"red");
    }
    f.close();

}

void Membrane::write_smoothed(const string& fname){
    System out;
    for(auto& l: lipids){
        auto s = out.append(l.mid_sel(0,0));
        s.name(0) = "M";
        s = out.append(l.head_sel(0,0));
        s.xyz(0) = l.smoothed_mid_xyz;
        s.name(0) = "S";
    }
    out().write(fname);
}

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

Lipid::Lipid(const Selection &sel, const Lipid_descr &descr){
    name = descr.name;
    whole_sel = sel;
    head_sel = whole_sel(descr.head_sel_str);
    tail_sel = whole_sel(descr.tail_sel_str);
    mid_sel = whole_sel(descr.mid_sel_str);
    // Fill tail indexes if any
    tail_carbon_indexes.resize(descr.tail_carbon_sels.size());
    order.resize(descr.tail_carbon_sels.size());
    for(int t=0; t<descr.tail_carbon_sels.size(); ++t){
        tail_carbon_indexes[t] = whole_sel(descr.tail_carbon_sels[t]).get_index();
        // Allocate array for order
        order[t].resize(tail_carbon_indexes[t].size()-2);
    }
}

void Lipid::set_markers()
{
    // Unwrap this lipid with leading index of position[0]
    whole_sel.unwrap(fullPBC, mid_sel.index(0)-whole_sel.index(0));

    // Set markers to COM
    head_marker = head_sel.center(true);
    tail_marker = tail_sel.center(true);
    mid_marker = mid_sel.center(true);
}


/********************************************************************************
 * Copyright 2009 The Robotics Group, The Maersk Mc-Kinney Moller Institute,
 * Faculty of Engineering, University of Southern Denmark
 *
 * 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.
 ********************************************************************************/

#ifndef RWLIBS_ALGORTIHMS_KDTREE_KDTREE_HPP_
#define RWLIBS_ALGORTIHMS_KDTREE_KDTREE_HPP_

#include <rw/core/Ptr.hpp>
#include <rw/core/macros.hpp>
#include <rw/math/Math.hpp>
#include <rw/math/Metric.hpp>

#include <algorithm>
#include <boost/any.hpp>
#include <float.h>
#include <list>
#include <queue>
#include <vector>

namespace rwlibs { namespace algorithms {

    // We can't define this as a static variable within the class, so we put it
    // here for now.
    const double kdtree_epsilon = 0.000000001;

    /** \addtogroup algorithms */
    /*@{*/

    /**
     * @brief a space partitioning structure for organizing points in k-dimensional space.
     * Used for searches involving multi.dimensional search keys, including nearest
     * neighbor and range search.
     *
     * KEY must implement:
     * copyable
     * operator[]
     * operator()
     * TODO: remove one of the operators (change code in KDTREE)
     *
     */
    template< class KEY, size_t DIM > class KDTree
    {
      private:
        struct TreeNode;

        size_t _nrOfNodes;
        TreeNode* _root;
        std::vector< TreeNode >* _nodes;
        rw::core::Ptr< rw::math::Metric< KEY > > _metric;

      public:
        //! a struct for the node in the tree
        struct KDNode
        {
            KDNode (const KEY& k, boost::any val) : key (k), value (val) {}
            KEY key;
            boost::any value;
        };

        struct KDResult
        {
            KDResult (KDNode* node, double d) : n (node), dist (d) {}
            KDNode* n;
            double dist;
        };

        /**
         * @brief Constructor
         * @cond
         * @param dim [in] the dimension of the keys in the KDTree
         * @endcond
         * @param metric documentation missing !
         */
        KDTree (rw::core::Ptr< rw::math::Metric< KEY > > metric) :
            _nrOfNodes (0), _root (NULL), _nodes (new std::vector< TreeNode > ())
        {}

        /**
         * @brief destructor
         */
        virtual ~KDTree ()
        {
            if (_root != NULL)
                delete _root;
            if (_nodes != NULL)
                delete _nodes;
        };

        rw::core::Ptr< rw::math::Metric< KEY > > getMetric () { return _metric; }

        /**
         * @brief Builds a KDTree from a list of key values and nodes. This method is more efficient
         * than creating an empty KDTree and then inserting nodes
         * @param nodes [in] a list of KDNode's
         * @param metric documentation missing !
         * @return if build succesfull then a pointer to a KD-tree is returned else NULL
         */
        static KDTree< KEY, DIM >* buildTree (std::vector< KDNode >& nodes,
                                              rw::core::Ptr< rw::math::Metric< KEY > > metric)
        {
            if (nodes.size () == 0)
                return NULL;

            // create all tree nodes in a list
            std::vector< TreeNode >* tNodes = new std::vector< TreeNode > (nodes.size ());

            // copy the KDNodes into the tree nodes
            int i = 0;
            for (KDNode& n : nodes) {
                (*tNodes)[i]._kdnode = &n;
                i++;
            }

            // create a simple median balanced tree
            TreeNode* root = buildBalancedRec (*tNodes, 0, (int) tNodes->size (), 0, DIM);

            return new KDTree< KEY, DIM > (*root, tNodes, metric);
        }

        /**
         * @brief Builds a KDTree from a list of key values and nodes. This method is more efficient
         * than creating an empty KDTree and then inserting nodes
         * @param nodes [in] a list of KDNode's
         * @param metric documentation missing !
         * @return if build succesfull then a pointer to a KD-tree is returned else NULL
         */
        static KDTree< KEY, DIM >* buildTree (const std::vector< KDNode* >& nodes,
                                              rw::core::Ptr< rw::math::Metric< KEY > > metric)
        {
            if (nodes.size () == 0)
                return NULL;

            // create all tree nodes in a list
            std::vector< TreeNode >* tNodes = new std::vector< TreeNode > (nodes.size ());

            // copy the KDNodes into the tree nodes
            int i = 0;
            for (KDNode* n : nodes) {
                (*tNodes)[i]._kdnode = n;
                i++;
            }

            // create a simple median balanced tree
            TreeNode* root = buildBalancedRec (*tNodes, 0, (int) tNodes->size (), 0, DIM);

            return new KDTree< KEY, DIM > (*root, tNodes, metric);
        }

        /**
         * @brief gets the number of dimensions that this KDTree supports
         * @return the nr of dimensions of this KD-Tree
         */
        size_t getDimensions () const { return DIM; };

        /**
         * @brief adds a key value pair to the KDTree.
         * @param key [in] must be the same length as the dimensionality of the KDTree
         * @param val [in] value that is to be stored at the keys position
         */
        void addNode (const KEY& key, boost::any val)
        {
            RW_THROW ("There is no implementation yet!");
        }

        /**
         * @brief finds the KDNode with key equal to nnkey
         * @param nnkey [in] the key that is to be found
         * @return KDNode with key equal to nnkey if existing, else NULL
         */
        KDNode* search (const KEY& nnkey)
        {
            TreeNode* tmpNode = _root;
            for (size_t lev = 0; tmpNode != NULL; lev = (lev + 1) % DIM) {
                KEY& key = tmpNode->_kdnode->key;
                if (nnkey[lev] == key[lev] && !(tmpNode->_deleted) &&
                    (_metric->distance (nnkey, key) < 0.000001)) {
                    return tmpNode->_kdnode;
                }
                else if (nnkey (lev) > key (lev)) {
                    tmpNode = tmpNode->_right;
                }
                else {
                    tmpNode = tmpNode->_left;
                }
            }
            return NULL;
        };

        /**
         * @brief finds the KDNode with the key closest too nnkey
         * @param nnkey [in] the key to which the nearest neighbor is found
         * @return the nearest neighbor to nnkey
         */
        KDNode& nnSearch (const KEY& nnkey)
        {
            // std::cout << "nnSearch " << DIM << std::endl;

            // RW_ASSERT(_metric->size(nnkey)==DIM);
            if (_root == NULL)
                RW_THROW ("KDTree has no data!");

            KEY min = nnkey, max = nnkey;
            KDResult result (NULL, DBL_MAX);
            for (size_t i = 0; i < DIM; i++) {
                min[i] = -DBL_MAX;
                max[i] = DBL_MAX;
            }
            // std::cout << "nnSearchRec" << std::endl;
            nnSearchRec (nnkey, _root, min, max, result);
            if (result.n == NULL)
                RW_THROW ("KDTree has no data!");
            return *result.n;
        }

        /**
         * @brief finds all neighbors in the hyperelipse with radius radi and center in nnkey.
         * @param nnkey [in] the center of the hyperelipse
         * @param radi [in] the radius of the hyperelipse in euclidean 2-norm
         * @param nodes [out] a container for all nodes that is found within the hyperelipse
         */
        void nnSearchElipse (const KEY& nnkey, const KEY& radi, std::list< const KDNode* >& nodes)
        {
            // typedef std::pair<TreeNode*,size_t> QElem;
            using namespace rw::math;
            std::queue< TreeNode* > unhandled;
            unhandled.push (_root);
            double distRadi = _metric->distance (radi);
            KEY low = nnkey, upp = nnkey;
            for (size_t i = 0; i < DIM; i++) {
                low[i] = nnkey[i] - distRadi;
                upp[i] = nnkey[i] + distRadi;
            }

            nnSearchElipseRec (nnkey, _root, low, upp, distRadi, nodes);
        }

        /**
         * @brief finds all neighbors in the hyperelipse with radius radi and center in nnkey.
         * @param nnkey [in] the center of the hyperelipse
         * @param radi [in] the radius of the hyperelipse in euclidean 2-norm
         * @param nodes [out] a container for all nodes that is found within the hyperelipse
         */
        void nnSearchElipseRect (const KEY& nnkey, const KEY& radi,
                                 std::list< const KDNode* >& nodes)
        {
            // typedef std::pair<TreeNode*,size_t> QElem;
            using namespace rw::math;
            std::queue< TreeNode* > unhandled;
            unhandled.push (_root);

            double distRadi = _metric->distance (radi);
            KEY low = nnkey, upp = nnkey;
            for (size_t i = 0; i < DIM; i++) {
                low (i) = nnkey (i) - distRadi;
                upp (i) = nnkey (i) + distRadi;
            }

            while (!unhandled.empty ()) {
                // std::cout << "unhandled size: " << unhandled.size() << std::endl;
                TreeNode* n = unhandled.front ();
                unhandled.pop ();

                unsigned char axis = n->_axis;
                KEY& key           = n->_kdnode->key;

                // std::cout << "Axis: " << axis << std::endl;

                // if the key is in range then add it to the result
                size_t j;
                for (j = 0; j < DIM && low[j] <= key[j] && upp[j] >= key[j]; j++)
                    ;
                // std::cout << j << "==" << DIM << " k:" << key << std::endl;
                if (j == DIM) {    // this is in range if
                    double dist = _metric->distance (nnkey, key);
                    // std::cout << "Dist: " << dist << " < " << distSqr << std::endl;
                    if (dist < distRadi)
                        nodes.push_back (n->_kdnode);
                }

                // add the children to the unhandled queue if the current dimension
                if ((low (axis) <= key (axis)) && (n->_left != NULL))
                    unhandled.push (n->_left);
                if ((upp (axis) > key (axis)) && (n->_right != NULL))
                    unhandled.push (n->_right);
            }
        }

        /**
         * @brief finds all neighbors in the hyperrectangle defined by the lower bound and the
         * upper bound
         */
        void nnSearchRect (const KEY& low, const KEY& upp, std::list< const KDNode* >& nodes)
        {
            // typedef std::pair<TreeNode*,size_t> QElem;
            std::queue< TreeNode* > unhandled;
            unhandled.push (_root);

            // std::cout << "nnSearchRect: "<< std::endl;
            // std::cout << "- low bound: "<< low << std::endl;
            // std::cout << "- upp bound: "<< upp << std::endl;

            while (!unhandled.empty ()) {
                // std::cout << "unhandled size: " << unhandled.size() << std::endl;
                TreeNode* n = unhandled.front ();
                unhandled.pop ();

                unsigned char axis = n->_axis;
                KEY& key           = n->_kdnode->key;

                // std::cout << "Axis: " << axis << std::endl;

                // if the key   is in range then add it to the result
                size_t j;
                for (j = 0; j < DIM && low[j] <= key[j] && key[j] <= upp[j]; j++)
                    ;
                // std::cout << j << "==" << DIM << " k:" << key << std::endl;
                if (j == DIM)    // this is in range
                    nodes.push_back (n->_kdnode);

                // add the children to the unhandled queue if the current dimension
                if ((low (axis) <= key (axis)) && (n->_left != NULL))
                    unhandled.push (n->_left);
                if ((upp (axis) > key (axis)) && (n->_right != NULL))
                    unhandled.push (n->_right);
            }
        };

      private:
        KDTree (){};

        KDTree (TreeNode& root, std::vector< TreeNode >* nodes,
                rw::core::Ptr< rw::math::Metric< KEY > > metric) :
            _root (&root),
            _nodes (nodes), _metric (metric){};

        /**
         * @brief Internal representation of a KD Tree Node. To save processing time when deleting
         * TreeNodes, a boolean is kept that say if the node is deleted or not. If deleted all
         * rutines kan skip the node and forward the call to its children.
         */
        struct TreeNode
        {
          public:
            TreeNode () :
                _left (NULL), _right (NULL), _kdnode (NULL), _deleted (false), _axis (0){};

            TreeNode (KDNode* node) :
                _left (NULL), _right (NULL), _kdnode (node), _deleted (false), _axis (0){};

            TreeNode (TreeNode* left, TreeNode* right, KDNode* node) :
                _left (left), _right (right), _kdnode (node), _deleted (false), _axis (0){};

            static void swap (TreeNode& n1, TreeNode& n2)
            {
                std::swap (n1._left, n2._left);
                std::swap (n1._right, n2._right);
                std::swap (n1._kdnode, n2._kdnode);
            }

            TreeNode *_left, *_right;
            KDNode* _kdnode;
            bool _deleted;          //
            unsigned char _axis;    // the splitting axis
        };

        struct SimpleCompare
        {
          private:
            size_t _dim;

          public:
            SimpleCompare (size_t dim) : _dim (dim){};

            bool operator() (const TreeNode& e1, const TreeNode& e2)
            {
                return e1._kdnode->key[_dim] < e2._kdnode->key[_dim];
            }
        };

        static TreeNode* buildBalancedRec (std::vector< TreeNode >& tNodes, int startIdx,
                                           int endIdx, size_t depth, size_t nrOfDims)
        {
            if (endIdx <= startIdx)
                return NULL;

            // std::cout << "RecBuild(" << startIdx << "," << endIdx << ")" << std::endl;
            size_t len = endIdx - startIdx;

            size_t dim = depth % nrOfDims;
            // the compare func can
            std::sort (&tNodes[startIdx], &tNodes[endIdx - 1], SimpleCompare (dim));
            size_t medianIdx = startIdx + len / 2;

            TreeNode& mNode = tNodes[medianIdx];
            mNode._axis     = 0xFF & dim;
            mNode._left = buildBalancedRec (tNodes, startIdx, (int) medianIdx, depth + 1, nrOfDims);
            mNode._right =
                buildBalancedRec (tNodes, (int) medianIdx + 1, endIdx, depth + 1, nrOfDims);
            return &mNode;
        }

        template< class T > T clamp (const T& val, const T& min, const T& max)
        {
            T result = val;
            for (size_t i = 0; i < DIM; i++) {
                result[i] = rw::math::Math::clamp (val[i], min[i], max[i]);
            }
            return result;
        }

        void nnSearchRec (const KEY& nnkey, TreeNode* node, KEY& min, KEY& max, KDResult& out)
        {
            using namespace rw::math;
            if (node == NULL)
                return;
            size_t axis = node->_axis;
            // std::cout << "nnSearchRec("<< axis << ")" << std::endl;

            KEY& key       = node->_kdnode->key;
            double distSqr = _metric->distance (nnkey, key);
            // std::cout << "le" << std::endl;

            // if this node is closer than any other then update out
            if (distSqr < out.dist && !node->_deleted) {
                out.dist = distSqr;
                out.n    = node->_kdnode;
            }
            // stop if the distance is very small
            if (distSqr < kdtree_epsilon)
                return;
            // std::cout << "1" << std::endl;
            // call nnSearch recursively with closerNode,
            // closestNode and closestDistSqr is updated
            bool isLeftClosest = nnkey (axis) < key (axis);
            if (isLeftClosest) {
                // std::cout << "left" << std::endl;
                // left is closest, backup split value and make the recursive call
                double maxTmp = max (axis);
                max (axis)    = key (axis);
                nnSearchRec (nnkey, node->_left, min, max, out);
                // undo the change of max
                max (axis) = maxTmp;
            }
            else {
                // std::cout << "right" << std::endl;
                // right is closest, backup split value and make the recursive call
                double minTmp = min (axis);
                min (axis)    = key (axis);
                nnSearchRec (nnkey, node->_right, min, max, out);
                // undo the change of max
                min (axis) = minTmp;
            }
            // std::cout << "2" << std::endl;

            // next check if fartherNode split plane lies closer than closestDistSqr
            if (Math::sqr (nnkey (axis) - key (axis)) >= out.dist)
                return;
            // std::cout << "3" << std::endl;

            bool isLeftFarthest = !isLeftClosest;
            // if closest point in hyperrect of farther node is closer than closest
            // then call nnSearch recursively with farther node
            if (isLeftFarthest) {
                double maxTmp = max (axis);
                max (axis)    = key (axis);
                KEY closest   = clamp (nnkey, min, max);
                if (_metric->distance (nnkey, closest) < out.dist)
                    nnSearchRec (nnkey, node->_left, min, max, out);
                // undo the change of max
                max (axis) = maxTmp;
            }
            else {
                double minTmp = min (axis);
                min (axis)    = key (axis);
                KEY closest   = clamp (nnkey, min, max);
                if (_metric->distance (nnkey, closest) < out.dist)
                    nnSearchRec (nnkey, node->_right, min, max, out);
                // undo the change of max
                min (axis) = minTmp;
            }
        };

        void nnSearchElipseRec (const KEY& nnkey, TreeNode* node, KEY& min, KEY& max,
                                double maxRadiSqr, std::list< const KDNode* >& nodes)
        {
            using namespace rw::math;
            if (node == NULL)
                return;
            size_t axis = node->_axis;
            // std::cout << "nnSearchRec("<< axis << ")" << std::endl;

            KEY& key       = node->_kdnode->key;
            double distSqr = _metric->distance (nnkey, key);

            // if this node is closer than any other then update out
            if (distSqr < maxRadiSqr && !node->_deleted) {
                nodes.push_back (node->_kdnode);
            }
            // stop if the distance is very small
            if (distSqr < kdtree_epsilon)
                return;

            // call nnSearch recursively with closerNode,
            // closestNode and closestDistSqr is updated
            bool isLeftClosest = nnkey (axis) < key (axis);
            if (isLeftClosest) {
                // left is closest, backup split value and make the recursive call
                double maxTmp = max (axis);
                max (axis)    = key (axis);
                nnSearchElipseRec (nnkey, node->_left, min, max, maxRadiSqr, nodes);
                // undo the change of max
                max (axis) = maxTmp;
            }
            else {
                // right is closest, backup split value and make the recursive call
                double minTmp = min (axis);
                min (axis)    = key (axis);
                nnSearchElipseRec (nnkey, node->_right, min, max, maxRadiSqr, nodes);
                // undo the change of max
                min (axis) = minTmp;
            }

            // next check if fartherNode split plane lies closer than closestDistSqr
            if (Math::sqr (nnkey (axis) - key (axis)) >= maxRadiSqr)
                return;

            bool isLeftFarthest = !isLeftClosest;
            // if closest point in hyperrect of farther node is closer than closest
            // then call nnSearch recursively with farther node
            if (isLeftFarthest) {
                double maxTmp = max (axis);
                max (axis)    = key (axis);
                KEY closest   = clamp (nnkey, min, max);
                if (_metric->distance (nnkey, closest) < maxRadiSqr)
                    nnSearchElipseRec (nnkey, node->_left, min, max, maxRadiSqr, nodes);
                // undo the change of max
                max (axis) = maxTmp;
            }
            else {
                double minTmp = min (axis);
                min (axis)    = key (axis);
                KEY closest   = clamp (nnkey, min, max);
                if (_metric->distance (nnkey, closest) < maxRadiSqr)
                    nnSearchElipseRec (nnkey, node->_right, min, max, maxRadiSqr, nodes);
                // undo the change of max
                min (axis) = minTmp;
            }
        };
    };

    extern template class KDTree< rw::math::Vector3D<>, 3 >;

    /**@}*/

}}    // namespace rwlibs::algorithms

#endif /*KDTREE_HPP_*/


#include <boost/config/warning_disable.hpp>
#include <boost/bind.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/qi_real.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_stl.hpp>
#include <boost/spirit/include/classic_core.hpp>
#include <boost/spirit/include/classic_increment_actor.hpp>
#include <string>
#include <map>
#include <iostream>
#include <fstream>
#include <stdexcept>
#include <equation_system/clientp.hpp>

namespace System {
using namespace std;
using namespace BOOST_SPIRIT_CLASSIC_NS;
namespace clientp
{

namespace {
	void sad_little_message(int line_no, string input){
		cerr << "Input line : " << line_no << endl;
		cerr << "Received : " << input << endl;
		cerr << "Please express equations using arithmetic combinations +/- of one or more terms, " << endl;
		cerr << "terms consisting of products expressed with *, powers with ^ and powers/indices" << endl;
		cerr << "expressed within parenthesis for constants k and variables y (1-based index) eg.," << endl;
		cerr << " " << endl;
		cerr << "1.0*k(717)*y(516)-1.0*k(416)*y(1)*y(392)-1.0*k(718)*y(1)^(1/2)*y(517)" << endl;
		cerr << " " << endl;
		throw std::invalid_argument( "Received bad equation format on stdin" );
	}
}

template <typename T>
void update(std::map<T,T> &x, const T& key, const T& val){

	x[key]+=val;
}

namespace qi = boost::spirit::qi;
namespace ascii = boost::spirit::ascii;
namespace phoenix = boost::phoenix;

using qi::float_;
using qi::double_;
using qi::phrase_parse;
using qi::_1;
using ascii::space;
using phoenix::push_back;

template <typename Iterator>
	bool parse_csv(Iterator first, Iterator last, std::vector<float>& k)
	{
		bool r = phrase_parse(first, last,

				//  grammar for csv files
				(
				 //*(float_ >> ',' | float_ [push_back(phoenix::ref(k),_1)])
				 *(float_ [push_back(phoenix::ref(k),_1)])
				)
				,
				space);

		return r;

	}

template <typename Iterator>
	bool parse_k_vals(Iterator first, Iterator last, std::vector<float>& w)
	{
		bool r = phrase_parse(first, last,

				//  grammar for k constants
				(
				 ("k(" >> float_[push_back(phoenix::ref(w), _1)] >> ')')
				)
				,
				//skip
				space | '*' | float_ >> '*' | '-'>>float_ >> '*'| '+' >> float_ >> '*'| "y(" >> float_ >> ')');

		return r;
	}

template <typename Iterator>
	bool parse_constants(Iterator first, Iterator last, std::vector<float>& v)
	{
		bool r = phrase_parse(first, last,

				//  grammar for numerical constants & sign
				(
				 (float_[push_back(phoenix::ref(v), _1)] >> '*')
				)
				,
				//skip
				space);
		return r;
	}
template <typename Iterator>
	bool parse_y_vals(Iterator first, Iterator last, std::map<float,float>& x)
	{
		float base =1.0;
		float exp = 0.0;
		bool r = phrase_parse(first, last,

				//  grammar for y variables
				//  this looks like hell
				(
				 *("y(" >> float_[phoenix::ref(base)=_1] >> ")^(" >> float_[phoenix::ref(exp)=_1] >> '/' >> float_[phoenix::ref(exp)/=_1]\
					 [boost::bind(&update<float>,boost::ref(x),boost::ref(base),boost::ref(exp))] >> ')' |
					 "y(" >> float_[phoenix::ref(base)=_1, boost::bind(&update<float>,boost::ref(x),boost::ref(base),1.0f)] >> ')')
				)
				,
				//skip
				space | '*' | float_ >> '*' | '-'>>float_ >> '*'| '+' >> float_ >> '*'| "k(" >> float_ >> ')');

		return r;
	}

template <typename Iterator>
	bool parse_csv(Iterator first, Iterator last, std::vector<double>& k)
	{
		bool r = phrase_parse(first, last,

				//  grammar for csv files
				(
				 //*(double_ >> ',' | double_ [push_back(phoenix::ref(k),_1)])
				 *(double_ [push_back(phoenix::ref(k),_1)])
				)
				,
				space);

		return r;

	}

template <typename Iterator>
	bool parse_k_vals(Iterator first, Iterator last, std::vector<double>& w)
	{
		bool r = phrase_parse(first, last,

				//  grammar for k constants
				(
				 ("k(" >> double_[push_back(phoenix::ref(w), _1)] >> ')')
				)
				,
				//skip
				space | '*' | double_ >> '*' | '-'>>double_ >> '*'| '+' >> double_ >> '*'| "y(" >> double_ >> ')');

		return r;
	}

template <typename Iterator>
	bool parse_constants(Iterator first, Iterator last, std::vector<double>& v)
	{
		bool r = phrase_parse(first, last,

				//  grammar for numerical constants & sign
				(
				 (double_[push_back(phoenix::ref(v), _1)] >> '*')
				)
				,
				//skip
				space);
		return r;
	}
template <typename Iterator>
	bool parse_y_vals(Iterator first, Iterator last, std::map<double,double>& x)
	{
		double base =1.0;
		double exp = 0.0;
		bool r = phrase_parse(first, last,

				//  grammar for y variables
				//  this looks like hell
				(
				 *("y(" >> double_[phoenix::ref(base)=_1] >> ")^(" >> double_[phoenix::ref(exp)=_1] >> '/' >> double_[phoenix::ref(exp)/=_1]\
					 [boost::bind(&update<double>,boost::ref(x),boost::ref(base),boost::ref(exp))] >> ')' |
					 "y(" >> double_[phoenix::ref(base)=_1, boost::bind(&update<double>,boost::ref(x),boost::ref(base),1.0f)] >> ')')
				)
				,
				//skip
				space | '*' | double_ >> '*' | '-'>>double_ >> '*'| '+' >> double_ >> '*'| "k(" >> double_ >> ')');

		return r;
	}
}
}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2014, Oracle and/or its affiliates.

// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Licensed under the Boost Software License version 1.0.
// http://www.boost.org/users/license.html

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_CLOSEST_FEATURE_POINT_TO_RANGE_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_CLOSEST_FEATURE_POINT_TO_RANGE_HPP

#include <utility>

#include <boost/range.hpp>

#include <boost/geometry/core/assert.hpp>
#include <boost/geometry/core/closure.hpp>
#include <boost/geometry/strategies/distance.hpp>
#include <boost/geometry/util/math.hpp>

namespace geofeatures_boost {} namespace boost = geofeatures_boost; namespace geofeatures_boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace closest_feature
{

// returns the segment (pair of iterators) that realizes the closest
// distance of the point to the range
template
<
    typename Point,
    typename Range,
    closure_selector Closure,
    typename Strategy
>
class point_to_point_range
{
protected:
    typedef typename geofeatures_boost::range_iterator<Range const>::type iterator_type;

    template <typename Distance>
    static inline void apply(Point const& point,
                             iterator_type first,
                             iterator_type last,
                             Strategy const& strategy,
                             iterator_type& it_min1,
                             iterator_type& it_min2,
                             Distance& dist_min)
    {
        BOOST_GEOMETRY_ASSERT( first != last );

        Distance const zero = Distance(0);

        iterator_type it = first;
        iterator_type prev = it++;
        if (it == last)
        {
            it_min1 = it_min2 = first;
            dist_min = strategy.apply(point, *first, *first);
            return;
        }

        // start with first segment distance
        dist_min = strategy.apply(point, *prev, *it);
        iterator_type prev_min_dist = prev;

        // check if other segments are closer
        for (++prev, ++it; it != last; ++prev, ++it)
        {
            Distance dist = strategy.apply(point, *prev, *it);
            if (geometry::math::equals(dist, zero))
            {
                dist_min = zero;
                it_min1 = prev;
                it_min2 = it;
                return;
            }
            else if (dist < dist_min)
            {
                dist_min = dist;
                prev_min_dist = prev;
            }
        }

        it_min1 = it_min2 = prev_min_dist;
        ++it_min2;
    }

public:
    typedef typename std::pair<iterator_type, iterator_type> return_type;

    template <typename Distance>
    static inline return_type apply(Point const& point,
                                    iterator_type first,
                                    iterator_type last,
                                    Strategy const& strategy,
                                    Distance& dist_min)
    {
        iterator_type it_min1, it_min2;
        apply(point, first, last, strategy, it_min1, it_min2, dist_min);

        return std::make_pair(it_min1, it_min2);
    }

    static inline return_type apply(Point const& point,
                                    iterator_type first,
                                    iterator_type last,
                                    Strategy const& strategy)
    {
        typename strategy::distance::services::return_type
            <
                Strategy,
                Point,
                typename geofeatures_boost::range_value<Range>::type
            >::type dist_min;

        return apply(point, first, last, strategy, dist_min);
    }

    template <typename Distance>
    static inline return_type apply(Point const& point,
                                    Range const& range,
                                    Strategy const& strategy,
                                    Distance& dist_min)
    {
        return apply(point,
                     geofeatures_boost::begin(range),
                     geofeatures_boost::end(range),
                     strategy,
                     dist_min);
    }

    static inline return_type apply(Point const& point,
                                    Range const& range,
                                    Strategy const& strategy)
    {
        return apply(point, geofeatures_boost::begin(range), geofeatures_boost::end(range), strategy);
    }
};

// specialization for open ranges
template <typename Point, typename Range, typename Strategy>
class point_to_point_range<Point, Range, open, Strategy>
    : point_to_point_range<Point, Range, closed, Strategy>
{
private:
    typedef point_to_point_range<Point, Range, closed, Strategy> base_type;
    typedef typename base_type::iterator_type iterator_type;

    template <typename Distance>
    static inline void apply(Point const& point,
                             iterator_type first,
                             iterator_type last,
                             Strategy const& strategy,
                             iterator_type& it_min1,
                             iterator_type& it_min2,
                             Distance& dist_min)
    {
        BOOST_GEOMETRY_ASSERT( first != last );

        base_type::apply(point, first, last, strategy,
                         it_min1, it_min2, dist_min);

        iterator_type it_back = --last;
        Distance const zero = Distance(0);
        Distance dist = strategy.apply(point, *it_back, *first);

        if (geometry::math::equals(dist, zero))
        {
            dist_min = zero;
            it_min1 = it_back;
            it_min2 = first;
        }
        else if (dist < dist_min)
        {
            dist_min = dist;
            it_min1 = it_back;
            it_min2 = first;
        }
    }

public:
    typedef typename std::pair<iterator_type, iterator_type> return_type;

    template <typename Distance>
    static inline return_type apply(Point const& point,
                                    iterator_type first,
                                    iterator_type last,
                                    Strategy const& strategy,
                                    Distance& dist_min)
    {
        iterator_type it_min1, it_min2;

        apply(point, first, last, strategy, it_min1, it_min2, dist_min);

        return std::make_pair(it_min1, it_min2);
    }

    static inline return_type apply(Point const& point,
                                    iterator_type first,
                                    iterator_type last,
                                    Strategy const& strategy)
    {
        typedef typename strategy::distance::services::return_type
            <
                Strategy,
                Point,
                typename geofeatures_boost::range_value<Range>::type
            >::type distance_return_type;

        distance_return_type dist_min;

        return apply(point, first, last, strategy, dist_min);
    }

    template <typename Distance>
    static inline return_type apply(Point const& point,
                                    Range const& range,
                                    Strategy const& strategy,
                                    Distance& dist_min)
    {
        return apply(point,
                     geofeatures_boost::begin(range),
                     geofeatures_boost::end(range),
                     strategy,
                     dist_min);
    }

    static inline return_type apply(Point const& point,
                                    Range const& range,
                                    Strategy const& strategy)
    {
        return apply(point, geofeatures_boost::begin(range), geofeatures_boost::end(range), strategy);
    }
};

}} // namespace detail::closest_feature
#endif // DOXYGEN_NO_DETAIL

}} // namespace geofeatures_boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_CLOSEST_FEATURE_POINT_TO_RANGE_HPP


/*
  Author(s):      Robert I A Patterson
  Project:        sweepc (population balance solver)
  Sourceforge:    http://sourceforge.net/projects/mopssuite

  Copyright (C) 2011 Robert I A Patterson.

  File purpose:
    Implementation of constant additive coagulation kernel
	for the hybrid particle-number and particle model

  Licence:
    This file is part of "sweepc".

    sweepc 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 2
    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, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

  Contact:
    Dr Markus Kraft
    Dept of Chemical Engineering
    University of Cambridge
    New Museums Site
    Pembroke Street
    Cambridge
    CB2 3RA
    UK

    Email:       mk306@cam.ac.uk
    Website:     http://como.cheng.cam.ac.uk
*/

#include "swp_hybrid_constcoag.h"

#include "swp_params.h"
#include "swp_cell.h"
#include "swp_mechanism.h"
#include <boost/random/uniform_01.hpp>

using namespace Sweep::Processes;

const double Sweep::Processes::HybridConstantCoagulation::s_MajorantFactor = 1.5;

/**
 * Main way of building a new coagulation object
 * @param[in] mech      Mechanism to which coagulation will belong
 *
 */
Sweep::Processes::HybridConstantCoagulation::HybridConstantCoagulation(const Sweep::Mechanism &mech)
: Coagulation(mech)
{
    m_name = "ConstantCoagulation";
}

// Stream-reading constructor.
Sweep::Processes::HybridConstantCoagulation::HybridConstantCoagulation(std::istream &in, const Sweep::Mechanism &mech)
: Coagulation(mech)
{
    m_name = "ConstantCoagulation";
    Deserialize(in, mech);
}

// Returns the rate of the process for the given system.
double Sweep::Processes::HybridConstantCoagulation::Rate(double t, const Cell &sys,
                                                        const Geometry::LocalGeometry1d &local_geom) const
{
    // Get the number of particles in the system.
    unsigned int n = sys.ParticleCount() + sys.Particles().GetTotalParticleNumber();

    return A() * n * (n - 1) * s_MajorantFactor / sys.SampleVolume() / 2;
}

/**
 * Number of terms in the expression for the sum of the majorant
 * kernel over all particle pairs.
 */
unsigned int Sweep::Processes::HybridConstantCoagulation::TermCount() const { return TYPE_COUNT; }

/**
 * Calculate the terms in the sum of the majorant kernel over all particle
 * pairs, placing each term in successive positions of the sequence
 * beginning at iterm and return the sum of the terms added to that
 * vector.
 *
 * @param[in]     t            Time for which rates are requested
 * @param[in]     sys          Details of the particle population and environment
 * @param[in]     local_geom   Position information
 * @param[in,out] iterm        Pointer to start of sequence to hold the rate terms, returned as one past the end.
 */
double Sweep::Processes::HybridConstantCoagulation::RateTerms(double t, const Cell &sys,
                                                             const Geometry::LocalGeometry1d &local_geom,
                                                             fvector::iterator &iterm) const
{
    return *(iterm++) = Rate(t, sys, local_geom);
}

/*!
 *
 *
 * \param[in]       t           Time
 * \param[in,out]   sys         System to update
 * \param[in]       local_geom  Details of local phsyical layout
 * \param[in]       iterm       Process term responsible for this event
 * \param[in,out]   rng         Random number generator
 *
 * \return      0 on success, otherwise negative.
 */
int HybridConstantCoagulation::Perform(double t, Sweep::Cell &sys,
                             const Geometry::LocalGeometry1d& local_geom,
                             unsigned int iterm,
                             Sweep::rng_type &rng) const
{
	PartPtrVector dummy;
    // Select properties by which to choose particles.
    // Note we need to choose 2 particles.  One particle must be chosen
    // uniformly and one with probability proportional
    // to particle mass.

    int ip1 = -1, ip2 = -1;
    unsigned int index1 = 0, index2 = 0;

    // Hybrid particle model flags
    bool hybrid_flag = m_mech->IsHybrid();
    bool ip1_flag = false;
    bool ip2_flag = false;
    bool coag_in_place = false;

    double n_incep = sys.Particles().GetTotalParticleNumber();
    double n_other = sys.ParticleCount();
    double n_total = n_incep + n_other;

    boost::uniform_01<rng_type&, double> unifDistrib(rng);
    double alpha1 = unifDistrib() * n_total;
    double alpha2 = unifDistrib() * n_total;
    if (n_total <= 0)
        return -1;

    if (n_total < 2) // if there are < 2 SPs but incepting class has weight >= 2, we can still act
        return 1;
    else
    {
        if (hybrid_flag)
        {
            // Particle 1 is picked uniformly. Here, the
            // incepting class has multiple weight 1 particles
            // Account for this by selecting this by default and
            // switching with probability n_other/n_total
            ip1 = -2;
            ip2 = -2;
            if (n_other >= alpha1)
            {
                ip1 = sys.Particles().Select_usingGivenRand(iUniform, alpha1, rng);
            }
        }
        else
        {
            ip1 = sys.Particles().Select(rng);
            ip2 = sys.Particles().Select(rng);
        }
    }

    // Choose and get first particle, then update it.
    Particle *sp1 = NULL;
    double dsp1 = 0.0;
    bool must_switch = false;

    // Is this an incepting class particle?
    if (hybrid_flag && ip1 == -2)
    {
        // Note don't need to add it to the ensemble unless coagulation is successful
        index1 = m_mech->SetRandomParticle(sys.Particles(), t, alpha1 - n_other, iUniform, rng);
		if (index1 == 0)
			return -1;
        sp1 = sys.Particles().GetPNParticleAt(index1)->Clone();
        sp1->SetTime(t);
        ip1_flag = true;                                                                // Flag sp1 as an incepting class particle
    }
    else
    {
        if (ip1 >= 0) {
            sp1 = sys.Particles().At(ip1);
        }
        else {
            // Failed to choose a particle.
            return -1;
        }
    }

    // Choose and get unique second particle, then update it.  Note, we are allowed to do
    // this even if the first particle was invalidated.
    unsigned int guard = 0;
    if (!hybrid_flag)
    {
        while ((ip2 == ip1) && (++guard < 1000))
        {
            ip2 = sys.Particles().Select(rng);
        }
    }
    else
    {
        ip2 = ip1;
        bool unsuitableChoice = true;
        unsigned int n_index1 = sys.Particles().NumberAtIndex(index1);
        while (unsuitableChoice && (++guard < 1000))
        {
            alpha2 = unifDistrib() * n_total;
            if (alpha2 <= n_incep)
            {
                index2 = m_mech->SetRandomParticle(sys.Particles(), t, alpha2, iUniform, rng);
				if (index2 == 0)
				{
					if (ip1_flag)
					{
						delete sp1;
						sp1 = NULL;
					}
					return -1;
				}
                if (!((index2 == index1) && (n_index1 == 1)))
                {
                    unsuitableChoice = false;
                    ip2 = -2;
                }
            }
            else
            {
                ip2 = sys.Particles().Select_usingGivenRand(iUniform, alpha2 - n_incep, rng);
                if (!(ip2 == ip1))
                    unsuitableChoice = false;
            }
        }
    }

    // Choose and get second particle, then update it.
    Particle *sp2 = NULL;
    double dsp2 = 0.0;

    // Is this an incepting class particle?
    if (hybrid_flag && ip2 == -2)
    {
        // Note don't need to add it to the ensemble unless coagulation is successful
        sp2 = sys.Particles().GetPNParticleAt(index2)->Clone();
        sp2->SetTime(t);
        ip2_flag = true;                                                             // Flag sp2 as an incepting class particle
    }
    else
    {
        if ((ip2 >= 0) && (ip2 != ip1)) {
            sp2 = sys.Particles().At(ip2);
        }
        else {
            // Failed to select a unique particle.
            return -1;
        }
    }

    //Calculate the majorant rate before updating the particles
    const double majk = MajorantKernel(*sp1, *sp2, sys, Default);

    //Update the particles
    if (t > sp1->LastUpdateTime())
        m_mech->UpdateParticle(*sp1, sys, t, ip1, rng, dummy);

    // Check that particle is still valid.  If not,
    // remove it and cease coagulating.
    if (!sp1->IsValid()) {
        if (!ip1_flag)
        {
            // Must remove first particle now.
            sys.Particles().Remove(ip1);
        }
        else
        {
            // Particle sp1 is not in the ensemble, must manually delete it
            delete sp1;
            sp1 = NULL;
        }
        // Invalidating the index tells this routine not to perform coagulation.
        ip1 = -1;
        return 0;
    }

    if (t > sp2->LastUpdateTime())
        m_mech->UpdateParticle(*sp2, sys, t, ip2, rng, dummy);

    // Check validity of particles after update.
    if (!sp2->IsValid()) {
        // Tell the ensemble to update particle one before we confuse things
        // by removing particle 2
        if (!ip1_flag)
            sys.Particles().Update(ip1);

        if (!ip2_flag)
        {
            // Must remove second particle now.
            sys.Particles().Remove(ip2);
        }
        else
        {
            // Particle sp2 is not in the ensemble, must manually delete it
            delete sp2;
            sp2 = NULL;
        }

        // Invalidating the index tells this routine not to perform coagulation.
        ip2 = -1;

        return 0;
    }

    // Check that both the particles are still valid.
    if ((ip1 != -1) && (ip2 != -1)) {
        // Must check for ficticious event now by comparing the original
        // majorant rate and the current (after updates) true rate.

        double truek = CoagKernel(*sp1, *sp2, sys);

        if (!Fictitious(majk, truek, rng)) {
            if (ip1_flag)
            {
                sys.Particles().UpdateTotalsWithIndex(index1, -1.0);
                sys.Particles().UpdateNumberAtIndex(index1, -1);
                sys.Particles().UpdateTotalParticleNumber(-1);
                unsigned int index12 = index1 + index2;
                // Allow for coagulation in place if the combined particle is small enough
				if ((m_mech->CoagulateInList()) && ip2_flag && (index12 < sys.Particles().GetHybridThreshold()))
                {
                    coag_in_place = true;
                    sys.Particles().UpdateTotalsWithIndex(index12, 1.0);
                    sys.Particles().UpdateNumberAtIndex(index12, 1);
                    sys.Particles().UpdateTotalParticleNumber(1);
                    if (sp1 != NULL)
                    {
                        delete sp1;
                        sp1 = NULL;
                    }
                }
                else
                {
                    // otherwise add the particle to the ensemble
                    ip1 = sys.Particles().Add(*sp1, rng, ip2, true);
                }
            }
            if (ip2_flag)
            {
                sys.Particles().UpdateTotalsWithIndex(index2, -1.0);
                sys.Particles().UpdateNumberAtIndex(index2, -1);
                sys.Particles().UpdateTotalParticleNumber(-1);
            }
            if (!coag_in_place)
                JoinParticles(t, ip1, sp1, ip2, sp2, sys, rng);
            if (ip2_flag && sp2 != NULL)
            {
                delete sp2;
                sp2 = NULL;
            }
        } else {
            if (!ip1_flag)
                sys.Particles().Update(ip1);
            else if (sp1 != NULL)
            {
                delete sp1;
                sp1 = NULL;
            }
            if (!ip2_flag)
                sys.Particles().Update(ip2);
            else if (sp2 != NULL)
            {
                delete sp2;
                sp2 = NULL;
            }
            return 1; // Ficticious event.
        }
    } else {
        // One or both particles were invalidated on update,
        // but that's not a problem.  Information on the update
        // of valid particles must be propagated into the binary
        // tree
        if (ip1 != -1)
        {
            if (!ip1_flag)
                sys.Particles().Update(ip1);
        }
        if (ip2 != -1 && !ip2_flag)
            sys.Particles().Update(ip2);

        if (ip1_flag && sp1 != NULL)
        {
            delete sp1;
            sp1 = NULL;
        }
        if (ip2_flag && sp2 != NULL)
        {
            delete sp2;
            sp2 = NULL;
        }
    }

    if (ip1_flag && sp1 != NULL && ip1 == -2)
    {
        delete sp1;
        sp1 = NULL;
    }

    if (ip2_flag && sp2 != NULL)
    {
        delete sp2;
        sp2 = NULL;
    }

    return 0;
}

/**
 * Calculate the coagulation kernel between two particles in a given environment.
 * Note that the details of the environment are not currently used.
 *
 *@param[in]    sp1         First particle
 *@param[in]    sp2         Second particle
 *@param[in]    sys         Details of the environment
 *
 *@return       Value of kernel
 */
double Sweep::Processes::HybridConstantCoagulation::CoagKernel(const Particle &sp1,
                                                              const Particle &sp2,
                                                              const Cell& sys) const
{
    return A();
}

/**
 * Calculate the majorant kernel between two particles in a given environment.
 * Note that the details of the environment are not currently used.
 *
 *@param[in]    sp1         First particle
 *@param[in]    sp2         Second particle
 *@param[in]    sys         Details of the environment
 *@param[in]    maj         Unused flag to indicate which majorant kernel is required
 *
 *@return       Value of majorant kernel
 */
double Sweep::Processes::HybridConstantCoagulation::MajorantKernel(const Particle &sp1,
                                                                  const Particle &sp2,
                                                                  const Cell& sys,
                                                                  const MajorantType maj) const
{
    return CoagKernel(sp1, sp2, sys) * s_MajorantFactor;
}