#include <Eigen/Core>
#include <boost/make_shared.hpp>
#include <ceres/ceres.h>
#include <cmath>
#include <pcl/io/ply_io.h>
#include <pcl/kdtree/kdtree_flann.h>

#include "align/cost_func.h"
#include "align/nonrigid_pipe.h"
#include "face/model_cudahelper.h"
#include "type.h"
#include "util/convert_arr.h"
#include "util/cu_quaternion.h"
#include "util/transform.h"
#include "util/require.h"

#define EPS 0.005

using namespace telef::types;
using namespace telef::face;

namespace telef::align {

boost::shared_ptr<PCANonRigidFittingResult>
PCAToFittingResultPipe::_processData(boost::shared_ptr<PCANonRigidAlignmentSuite> in)
{
  auto result = boost::make_shared<PCANonRigidFittingResult>() ;

  result->image = in->image;
  result->pca_model = in->pca_model;
  result->cloud = in->rawCloud;
  result->landmark3d = in->fittingSuite->landmark3d;
  result->fx = in->fx;
  result->fy = in->fy;
  result->transformation = in->transformation;
  result->expressionCoeff = in->expressionCoeff;
  result->shapeCoeff = in->shapeCoeff;

  return result;
}
PCAGPUNonRigidFittingPipe::PCAGPUNonRigidFittingPipe()
    : isModelInitialized(false), geoWeight(0), geoMaxPoints(0),
      geoSearchRadius(0), addGeoTerm(false) {
  if (cublasCreate(&cublasHandle) != CUBLAS_STATUS_SUCCESS) {
    throw std::runtime_error("Cublas could not be initialized");
  }
}

PCAGPUNonRigidFittingPipe::PCAGPUNonRigidFittingPipe(
    const float geoWeight,
    const int geoMaxPoints,
    const float geoSearchRadius,
    const bool addGeoTerm,
    const bool usePrevFrame,
    const bool adjustCamera)
    : isModelInitialized(false), geoWeight(geoWeight),
      geoMaxPoints(geoMaxPoints), geoSearchRadius(geoSearchRadius),
      addGeoTerm(addGeoTerm), usePrevFrame(usePrevFrame),
      adjustCamera(adjustCamera) {
  if (cublasCreate(&cublasHandle) != CUBLAS_STATUS_SUCCESS) {
    throw std::runtime_error("Cublas could not be initialized");
  }
}

PCAGPUNonRigidFittingPipe::PCAGPUNonRigidFittingPipe(
    const PCAGPUNonRigidFittingPipe &that) {
  this->isModelInitialized = false;
  this->geoWeight = that.geoWeight;
  this->geoMaxPoints = that.geoMaxPoints;
  this->geoSearchRadius = that.geoSearchRadius;
  this->addGeoTerm = that.addGeoTerm;
  this->usePrevFrame = that.usePrevFrame;
  this->adjustCamera = that.adjustCamera;
  if (cublasCreate(&cublasHandle) != CUBLAS_STATUS_SUCCESS) {
    throw std::runtime_error("Cublas could not be initialized");
  }
}

PCAGPUNonRigidFittingPipe::PCAGPUNonRigidFittingPipe(
    PCAGPUNonRigidFittingPipe &&that) noexcept {
  this->isModelInitialized = that.isModelInitialized;
  this->c_deformModel = that.c_deformModel;
  this->cublasHandle = that.cublasHandle;
  this->geoWeight = that.geoWeight;
  this->geoMaxPoints = that.geoMaxPoints;
  this->geoSearchRadius = that.geoSearchRadius;
  this->addGeoTerm = that.addGeoTerm;
  this->usePrevFrame = that.usePrevFrame;
  this->adjustCamera = that.adjustCamera;
}

PCAGPUNonRigidFittingPipe &PCAGPUNonRigidFittingPipe::
operator=(const PCAGPUNonRigidFittingPipe &that) {
  if (&that != this) {
    if (isModelInitialized) {
      freeModelCUDA(c_deformModel);
      this->isModelInitialized = false;
    }
    cublasDestroy(cublasHandle);
    if (cublasCreate(&this->cublasHandle) != CUBLAS_STATUS_SUCCESS) {
      throw std::runtime_error("Cublas could not be initialized");
    }
  }

  this->geoWeight = that.geoWeight;
  this->geoMaxPoints = that.geoMaxPoints;
  this->geoSearchRadius = that.geoSearchRadius;
  this->addGeoTerm = that.addGeoTerm;
  this->usePrevFrame = that.usePrevFrame;
  return *this;
}

PCAGPUNonRigidFittingPipe &PCAGPUNonRigidFittingPipe::
operator=(PCAGPUNonRigidFittingPipe &&that) {
  if (&that != this) {
    if (isModelInitialized) {
      freeModelCUDA(c_deformModel);
      this->isModelInitialized = false;
    }
    cublasDestroy(cublasHandle);
    this->isModelInitialized = that.isModelInitialized;
    this->c_deformModel = that.c_deformModel;
    this->cublasHandle = that.cublasHandle;
    this->geoWeight = that.geoWeight;
    this->geoMaxPoints = that.geoMaxPoints;
    this->geoSearchRadius = that.geoSearchRadius;
    this->addGeoTerm = that.addGeoTerm;
    this->usePrevFrame = that.usePrevFrame;
    this->adjustCamera = that.adjustCamera;
  }

  return *this;
}

PCAGPUNonRigidFittingPipe::~PCAGPUNonRigidFittingPipe() {
  if (isModelInitialized) {
    freeModelCUDA(c_deformModel);
  }
  cublasDestroy(cublasHandle);
}

boost::shared_ptr<PCANonRigidFittingResult>
PCAGPUNonRigidFittingPipe::_processData(
    boost::shared_ptr<PCANonRigidAlignmentSuite> in) {
  /* Load data to cuda device */

  if (!isModelInitialized) {
    auto shapeBasis = in->pca_model->getShapeBasisMatrix();
    auto expressionBasis = in->pca_model->getExpressionBasisMatrix();
    auto ref = in->pca_model->getReferenceVector();
    auto meanShapeDeformation = in->pca_model->getShapeDeformationCenter();
    auto meanExpressionDeformation =
        in->pca_model->getExpressionDeformationCenter();
    auto landmarks = in->pca_model->getLandmarks();
    loadModelToCUDADevice(
        &this->c_deformModel,
        shapeBasis,
        expressionBasis,
        ref,
        meanShapeDeformation,
        meanExpressionDeformation,
        landmarks);
    isModelInitialized = true;
  }

  if(0 < in->shapeCoeff.size() && 0 < in->expressionCoeff.size())
    {
      shapeCoeff.resize(in->shapeCoeff.size());
      expressionCoeff.resize(in->expressionCoeff.size());
      std::copy_n(
          in->shapeCoeff.data(),
          in->shapeCoeff.size(),
          shapeCoeff.data());
      std::copy_n(
          in->expressionCoeff.data(),
          in->expressionCoeff.size(),
          expressionCoeff.data());
    }

  C_ScanPointCloud c_scanPointCloud;
  loadScanToCUDADevice(
      &c_scanPointCloud,
      in->rawCloud,
      in->fx,
      in->fy,
      (0 == fixedShapeCoeff.size()) ? landmarkSelection : landmarkSelection_exp,
      in->transformation,
      in->fittingSuite->landmark3d);

  /* Setup Optimizer */
  auto lmkCost = new PCAGPULandmarkDistanceFunctor(
      this->c_deformModel, c_scanPointCloud, cublasHandle);
  ceres::Problem problem;

  // Reset input parameters if we are first frame or not using prev frames
  if (!usePrevFrame || t.size() == 0)
    t.assign(3, 0.0);
  if (!usePrevFrame || u.size() == 0)
    u = {0.0, 0.0, 0.0};

  if(0 != fixedShapeCoeff.size())
    {
      problem.AddResidualBlock(
          lmkCost,
          nullptr,
          fixedShapeCoeff.data(),
          expressionCoeff.data(),
          t.data(),
          u.data());
      problem.SetParameterBlockConstant(fixedShapeCoeff.data());
    }
  else
    {
      problem.AddResidualBlock(
          lmkCost,
          nullptr,
          shapeCoeff.data(),
          expressionCoeff.data(),
          t.data(),
          u.data());
    }

  if (addGeoTerm && 0 == fixedShapeCoeff.size()) {
    std::vector<double> fsc(in->shapeCoeff.size());
    for(size_t i=0; i<fsc.size(); i++)
      {
        fsc[i] = in->shapeCoeff(i);
      }
    fixedShapeCoeff = fsc;
    auto geoCost = new PCAGPUGeometricDistanceFunctor(
        this->c_deformModel,
        c_scanPointCloud,
        cublasHandle,
        geoMaxPoints * 3,
        sqrtf(geoWeight),
        geoSearchRadius);
    problem.AddResidualBlock(
        geoCost,
        nullptr,
        fixedShapeCoeff.data(),
        expressionCoeff.data(),
        t.data(),
        u.data());
    }

  if(!adjustCamera)
    {
      problem.SetParameterBlockConstant(t.data());
      problem.SetParameterBlockConstant(u.data());
    }
  problem.SetParameterBlockConstant(expressionCoeff.data());

  /*
  problem.AddResidualBlock(
      new L2RegularizerFunctor(c_deformModel.shapeRank, 1e-7),
      NULL,
      shapeCoeff.data());
  */

  /*
  problem.AddResidualBlock(
      new L2RegularizerFunctor(c_deformModel.expressionRank, 1e-4),
      NULL,
      expressionCoeff.data());
  */

  ceres::Solver::Options options;
  options.minimizer_progress_to_stdout = false;
  //options.use_nonmonotonic_steps=true;
  options.max_num_iterations = 1000;
  //options.function_tolerance = 1e-15;
  //options.parameter_tolerance = 1e-12;
  options.linear_solver_type = ceres::LinearSolverType::DENSE_NORMAL_CHOLESKY;

  /* Run Optimization */
  auto summary = ceres::Solver::Summary();
  ceres::Solve(options, &problem, &summary);
  //std::cout << summary.FullReport() << std::endl;

  float fu[3];
  float ft[3];
  float r[9];
  float trans[16];
  convertArray(t.data(), ft, 3);
  convertArray(u.data(), fu, 3);
  calc_r_from_u(r, fu);
  create_trans_from_tu(trans, ft, r);
  Eigen::Map<Eigen::Matrix4f> eigenTrans(trans);

  /* Free Resources */
  freeScanCUDA(c_scanPointCloud);

  auto result = boost::make_shared<PCANonRigidFittingResult>();

  result->shapeCoeff =
      Eigen::Map<Eigen::VectorXd>(shapeCoeff.data(), c_deformModel.shapeRank)
          .cast<float>();
  result->expressionCoeff =
      Eigen::Map<Eigen::VectorXd>(
          expressionCoeff.data(), c_deformModel.expressionRank)
          .cast<float>();
  /*
  result->expressionCoeff = in->expressionCoeff;
  result->shapeCoeff =
    Eigen::Map<Eigen::VectorXd>(fixedShapeCoeff.data(), c_deformModel.shapeRank).cast<float>();

  */
 /*
  std::cout << "Fitted(Shape): " << std::endl;
  std::cout << result->shapeCoeff << std::endl;
  std::cout << "Fitted(Expression): " << std::endl;
  std::cout << result->expressionCoeff << std::endl;
  */
  result->image = in->image;
  result->pca_model = in->pca_model;
  result->cloud = in->rawCloud;
  result->landmark3d = in->fittingSuite->landmark3d;
  result->fx = in->fx;
  result->fy = in->fy;
  result->transformation = eigenTrans * in->transformation;

  return result;
}

const std::vector<int> PCAGPUNonRigidFittingPipe::landmarkSelection{
  0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
    11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
    21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
    31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
    51, 52, 53, 54, 55, 56, 57, 58, 59, 60,
    61, 62, 63, 64, 65, 66, 67,
    };

const std::vector<int> PCAGPUNonRigidFittingPipe::landmarkSelection_exp{
  0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
    11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
    21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
    31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
    51, 52, 53, 54, 55, 56, 57, 58, 59, 60,
    61, 62, 63, 64, 65, 66, 67,
};

} // namespace telef::align


#define _USE_MATH_DEFINES
/// code to make agents
#include <vector>
#include <cassert>
#include <algorithm>
#include <iostream>
#include <random>

#include <boost/foreach.hpp>

#include <Rcpp.h>
#include <RcppParallel.h>

#include "network.h"
#include "landscape.h"
#include "agents.h"

// to shuffle pop id
void Population::shufflePop() {
    if (order[0] == order[nAgents - 1])
    {
        for (size_t i = 0; i < static_cast<size_t>(nAgents); i++)
        {
            order[i] = i;
        }
        std::random_shuffle ( order.begin(), order.end() );
    }
    else {
        std::random_shuffle ( order.begin(), order.end() );
    }

}

// to update agent Rtree
void Population::updateRtree () {
    // initialise rtree
    bgi::rtree< value, bgi::quadratic<16> > tmpRtree;
    for (int i = 0; i < nAgents; ++i)
    {
        point p = point(coordX[i], coordY[i]);
        tmpRtree.insert(std::make_pair(p, i));
    }
    std::swap(agentRtree, tmpRtree);
    tmpRtree.clear();
}

// uniform distribution for agent position
std::uniform_real_distribution<float> agent_ran_pos(0.0f, 1.f);

// function for initial positions
void Population::initPos(Resources food) {
    for (size_t i = 0; i < static_cast<size_t>(nAgents); i++) {
        coordX[i] = agent_ran_pos(rng) * food.dSize;
        initX[i] = coordX[i];
        coordY[i] = agent_ran_pos(rng) * food.dSize;
        initY[i] = coordY[i];
    }
    updateRtree();
}

// set agent trait
void Population::setTrait(const float mSize) {

    // create a cauchy distribution, mSize is the scale
    std::cauchy_distribution<float> agent_ran_trait(0.f, mSize);

    for(int i = 0; i < nAgents; i++) {
        sF[i] = agent_ran_trait(rng);
        sH[i] = agent_ran_trait(rng);
        sN[i] = agent_ran_trait(rng);
    }
}

float get_distance(float x1, float x2, float y1, float y2) {
    return std::sqrt(std::pow((x1 - x2), 2) + std::pow((y1 - y2), 2));
}

// general function for agents within distance
std::pair<int, int> Population::countAgents (
    const float xloc, const float yloc) {

    int handlers = 0;
    int nonhandlers = 0;
    std::vector<value> near_agents;
    // query for a simple box
    agentRtree.query(bgi::satisfies([&](value const& v) {
        return bg::distance(v.first, point(xloc, yloc)) < range_agents;}),
        std::back_inserter(near_agents));

    BOOST_FOREACH(value const& v, near_agents) {

        if(counter[v.second] > 0) handlers ++; else nonhandlers ++;
    }
    near_agents.clear();
    // first element is number of near entities
    // second is the identity of entities
    return std::pair<int, int> {handlers, nonhandlers};
}

// function for near agent ids
std::vector<int> Population::getNeighbourId (
    const float xloc, const float yloc) {

    std::vector<int> agent_id;
    std::vector<value> near_agents;
    // query for a simple box
    // neighbours for associations are counted over the MOVEMENT RANGE
    agentRtree.query(bgi::satisfies([&](value const& v) {
        return bg::distance(v.first, point(xloc, yloc)) < range_move;}),
        std::back_inserter(near_agents));

    BOOST_FOREACH(value const& v, near_agents) {
        agent_id.push_back(v.second);
    }
    near_agents.clear();
    // first element is number of near entities
    // second is the identity of entities
    return agent_id;
}

// general function for items within distance
int Population::countFood (
    const Resources &food,
    const float xloc, const float yloc) {

    int nFood = 0;
    std::vector<value> near_food;

    // check any available
    if (food.nAvailable > 0) {
        // query for a simple box
        food.rtree.query(bgi::satisfies([&](value const& v) {
            return bg::distance(v.first, point(xloc, yloc)) < range_food;}),
            std::back_inserter(near_food));

        BOOST_FOREACH(value const& v, near_food) {
            // count only which are available!
            if (food.available[v.second]) {
                nFood++;
            }
        }
        near_food.clear();
    }

    return nFood;
}

// function for the nearest available food item
std::vector<int> Population::getFoodId (
    const Resources &food,
    const float xloc, const float yloc) {

    std::vector<int> food_id;
    std::vector<value> near_food;
    // check any available
    if (food.nAvailable > 0) {
        // query for a simple box
        // food is accessed over the MOVEMENT RANGE
        food.rtree.query(bgi::satisfies([&](value const& v) {
            return bg::distance(v.first, point(xloc, yloc)) < range_move;}),
            std::back_inserter(near_food));

        BOOST_FOREACH(value const& v, near_food) {
            // count only which are available!
            if (food.available[v.second]) {
                food_id.push_back(v.second);
            }
        }
        near_food.clear();
    }

    // first element is number of near entities
    // second is the identity of entities
    return food_id;
}

/// rng for suitability
std::normal_distribution<float> noise(0.f, 0.01f);
std::cauchy_distribution<float> noise_cauchy(0.f, 0.001f);

/// population movement function
void Population::move(const Resources &food, const int nThreads) {

    float twopi = 2.f * M_PI;

    // what increment for 3 samples in a circle around the agent
    float increment = twopi / n_samples;
    float angle = 0.f;
    // for this increment what angles to sample at
    std::vector<float> sample_angles (static_cast<int>(n_samples), 0.f);
    for (int i_ = 0; i_ < static_cast<int>(n_samples); i_++)
    {
        sample_angles[i_] = angle;
        angle += increment;
    }

    // make random noise for each individual and each sample
    std::vector<std::vector<float> > noise_v (nAgents, std::vector<float>(static_cast<int>(n_samples), 0.f));
    for (size_t i_ = 0; i_ < noise_v.size(); i_++)
    {
        for (size_t j_ = 0; j_ < static_cast<size_t>(n_samples); j_++)
        {
            noise_v[i_][j_] = noise(rng);
        }
    }

    shufflePop();
    // loop over agents --- randomise
    if (nThreads > 1) {
        // any number above 1 will allow automatic n threads
        tbb::task_scheduler_init _tbb(tbb::task_scheduler_init::automatic); // automatic for now
        // try parallel
        tbb::parallel_for(
            tbb::blocked_range<unsigned>(1, order.size()),
            [&](const tbb::blocked_range<unsigned>& r) {
                for (unsigned i = r.begin(); i < r.end(); ++i) {
                    int id = order[i];
                    if (counter[id] > 0) {
                        counter[id] --;
                    }
                    else {
                        // first assess current location
                        float sampleX = coordX[id];
                        float sampleY = coordY[id];

                        float foodHere = 0.f;
                        // count local food only if items are available
                        if(food.nAvailable > 0) {
                            foodHere = static_cast<float>(countFood(
                                food, sampleX, sampleY
                            ));
                        }
                        // count local handlers and non-handlers
                        std::pair<int, int> agentCounts = countAgents(sampleX, sampleY);

                        // get suitability current
                        float suit_origin = (
                            (sF[id] * foodHere) + (sH[id] * agentCounts.first) +
                            (sN[id] * agentCounts.second)
                        );

                        float newX = sampleX;
                        float newY = sampleY;
                        // now sample at three locations around
                        for(size_t j = 0; j < sample_angles.size(); j++) {
                            float t1_ = static_cast<float>(cos(sample_angles[j]));
                            float t2_ = static_cast<float>(sin(sample_angles[j]));

                            // use range for agents to determine sample locs
                            sampleX = coordX[id] + (range_agents * t1_);
                            sampleY = coordY[id] + (range_agents * t2_);

                            // crudely wrap sampling location
                            if((sampleX > food.dSize) | (sampleX < 0.f)) {
                                sampleX = std::fabs(std::fmod(sampleX, food.dSize));
                            }
                            if((sampleY > food.dSize) | (sampleY < 0.f)) {
                                sampleY = std::fabs(std::fmod(sampleY, food.dSize));
                            }

                            // count food at sample locations if any available
                            if(food.nAvailable > 0) {
                                foodHere = static_cast<float>(countFood(
                                    food, sampleX, sampleY
                                ));
                            }

                            // count local handlers and non-handlers
                            std::pair<int, int> agentCounts = countAgents(sampleX, sampleY);

                            float suit_dest = (
                                (sF[id] * foodHere) + (sH[id] * agentCounts.first) +
                                (sN[id] * agentCounts.second) +
                                noise_v[id][j] // add same very very small noise to all
                            );

                            if (suit_dest > suit_origin) {
                                // where does the individual really go
                                newX = coordX[id] + (range_move * t1_);
                                newY = coordY[id] + (range_move * t2_);

                                // crudely wrap MOVEMENT location
                                if((newX > food.dSize) | (newX < 0.f)) {
                                    newX = std::fabs(std::fmod(newX, food.dSize));
                                }
                                if((newY > food.dSize) | (newY < 0.f)) {
                                    newY = std::fabs(std::fmod(newY, food.dSize));
                                }

                                assert(newX < food.dSize && newX > 0.f);
                                assert(newY < food.dSize && newY > 0.f);
                                suit_origin = suit_dest;
                            }
                        }
                        // distance to be moved
                        moved[id] += range_move;

                        // set locations
                        coordX[id] = newX; coordY[id] = newY;
                    }
                }
            }
        );
    } else if (nThreads == 1) {
        for (int i = 0; i < nAgents; ++i) {
            int id = order[i];
            if (counter[id] > 0) {
                counter[id] --;
            }
            else {
                // first assess current location
                float sampleX = coordX[id];
                float sampleY = coordY[id];

                float foodHere = 0.f;
                // count local food only if items are available
                if(food.nAvailable > 0) {
                    foodHere = static_cast<float>(countFood(
                        food, sampleX, sampleY
                    ));
                }
                // count local handlers and non-handlers
                std::pair<int, int> agentCounts = countAgents(sampleX, sampleY);

                // get suitability current
                float suit_origin = (
                    (sF[id] * foodHere) + (sH[id] * agentCounts.first) +
                    (sN[id] * agentCounts.second)
                );

                float newX = sampleX;
                float newY = sampleY;
                // now sample at three locations around
                for(size_t j = 0; j < sample_angles.size(); j++) {
                    float t1_ = static_cast<float>(cos(sample_angles[j]));
                    float t2_ = static_cast<float>(sin(sample_angles[j]));

                    // use range for agents to determine sample locs
                    sampleX = coordX[id] + (range_agents * t1_);
                    sampleY = coordY[id] + (range_agents * t2_);

                    // crudely wrap sampling location
                    if((sampleX > food.dSize) | (sampleX < 0.f)) {
                        sampleX = std::fabs(std::fmod(sampleX, food.dSize));
                    }
                    if((sampleY > food.dSize) | (sampleY < 0.f)) {
                        sampleY = std::fabs(std::fmod(sampleY, food.dSize));
                    }

                    // count food at sample locations if any available
                    if(food.nAvailable > 0) {
                        foodHere = static_cast<float>(countFood(
                            food, sampleX, sampleY
                        ));
                    }

                    // count local handlers and non-handlers
                    std::pair<int, int> agentCounts = countAgents(sampleX, sampleY);

                    float suit_dest = (
                        (sF[id] * foodHere) + (sH[id] * agentCounts.first) +
                        (sN[id] * agentCounts.second) +
                        noise_v[id][j] // add same very very small noise to all
                    );

                    if (suit_dest > suit_origin) {
                        // where does the individual really go
                        newX = coordX[id] + (range_move * t1_);
                        newY = coordY[id] + (range_move * t2_);

                        // crudely wrap MOVEMENT location
                        if((newX > food.dSize) | (newX < 0.f)) {
                            newX = std::fabs(std::fmod(newX, food.dSize));
                        }
                        if((newY > food.dSize) | (newY < 0.f)) {
                            newY = std::fabs(std::fmod(newY, food.dSize));
                        }

                        assert(newX < food.dSize && newX > 0.f);
                        assert(newY < food.dSize && newY > 0.f);
                        suit_origin = suit_dest;
                    }
                }
                // distance to be moved
                moved[id] += range_move;

                // set locations
                coordX[id] = newX; coordY[id] = newY;
            }
        }

    }

}

// function to paralellise choice of forage item
void Population::pickForageItem(const Resources &food, const int nThreads){
    shufflePop();
    // nearest food
    std::vector<int> idTargetFood (nAgents, -1);

    if (nThreads > 1)
    {
        // loop over agents --- no shuffling required here
        tbb::task_scheduler_init _tbb(tbb::task_scheduler_init::automatic); // automatic for now
        // try parallel foraging --- agents pick a target item
        tbb::parallel_for(
            tbb::blocked_range<unsigned>(1, order.size()),
                [&](const tbb::blocked_range<unsigned>& r) {
                for (unsigned i = r.begin(); i < r.end(); ++i) {
                    if ((counter[i] > 0) | (food.nAvailable == 0)) {
                        // nothing -- agent cannot forage or there is no food
                    }
                    else {
                        // find nearest item ids
                        std::vector<int> theseItems = getFoodId(food, coordX[i], coordY[i]);
                        int thisItem = -1;

                        // check near items count
                        if(theseItems.size() > 0) {
                            // take first item by default
                            thisItem = theseItems[0];
                            idTargetFood[i] = thisItem;
                        }
                    }
                }
            }
        );
    } else if (nThreads == 1)
    {
        for (int i = 0; i < nAgents; ++i) {
            if ((counter[i] > 0) | (food.nAvailable == 0)) {
                // nothing -- agent cannot forage or there is no food
            }
            else {
                // find nearest item ids
                std::vector<int> theseItems = getFoodId(food, coordX[i], coordY[i]);
                int thisItem = -1;

                // check near items count
                if(theseItems.size() > 0) {
                    // take first item by default
                    thisItem = theseItems[0];
                    idTargetFood[i] = thisItem;
                }
            }
        }
    }

    forageItem = idTargetFood;
}

// function to exploitatively forage on picked forage items
void Population::doForage(Resources &food) {
    // all agents have picked a food item if they can forage
    // now forage in a serial loop --- this cannot be parallelised
    // this order is randomised
    for (size_t i = 0; i < static_cast<size_t>(nAgents); i++)
    {
        int id = order[i];
        if ((counter[id] > 0) | (food.nAvailable == 0)) {
            // nothing
        } else {
            int thisItem = forageItem[id]; //the item picked by this agent
            // check selected item is available
            if (thisItem != -1)
            {
                counter[id] = handling_time;
                intake[id] += 1.0; // increased here --- not as described.

                // reset food availability
                food.available[thisItem] = false;
                food.counter[thisItem] = food.regen_time;
                food.nAvailable --;
            }
        }
    }
}

void Population::countAssoc(const int nThreads) {
    for (int i = 0; i < nAgents; ++i) {
        // count nearby agents and update raw associations
        std::vector<int> nearby_agents = getNeighbourId(coordX[i], coordY[i]);
        associations[i] += nearby_agents.size();

        // loop over nearby agents and update association matrix
        for (size_t j = 0; j < nearby_agents.size(); j++)
        {
            int target_agent = nearby_agents[j];
            pbsn.adjMat (i, target_agent) += 1;
        }
    }
}

/// minor function to normalise vector
std::vector<float> Population::handleFitness() {
    // sort vec fitness
    std::vector<float> vecFitness = energy;
    std::sort(vecFitness.begin(), vecFitness.end()); // sort to to get min-max
    // scale to max fitness
    float maxFitness = vecFitness[vecFitness.size()-1];
    float minFitness = vecFitness[0];

    // reset to energy
    vecFitness = energy;
    // rescale copied energy vector by min anx max fitness
    for(size_t i = 0; i < static_cast<size_t>(nAgents); i++) {
        vecFitness[i] = ((vecFitness[i]  - minFitness) / (maxFitness - minFitness)) +
         noise(rng);
    }

    return vecFitness;
}

// fun for replication
void Population::Reproduce(const Resources food, const bool infect_percent,
    const float dispersal, const float mProb, const float mSize)
{
    // std::bernoulli_distribution verticalInfect(0.01f);

    // mutation probability and size distribution --- inefficient but oh well
    std::bernoulli_distribution mutation_happens(mProb);
    std::cauchy_distribution<float> mutation_size(0.0, mSize);

    // choose the range over which individuals are dispersed
    std::normal_distribution<float> sprout(0.f, dispersal);
    std::vector<float> vecFitness;
    //normalise intake if percent infect is not true
    if (infect_percent) {
         vecFitness = energy;
    } else {
        vecFitness = handleFitness();
    }

    // set up weighted lottery
    std::discrete_distribution<> weightedLottery(vecFitness.begin(), vecFitness.end());

    // get parent trait based on weighted lottery
    std::vector<float> tmp_sF (nAgents, 0.f);
    std::vector<float> tmp_sH (nAgents, 0.f);
    std::vector<float> tmp_sN (nAgents, 0.f);

    // infected or not for vertical transmission
    std::vector<bool> infected_2 (nAgents, false);

    // reset infection source
    srcInfect = std::vector<int> (nAgents, 0);

    // reset associations
    associations = std::vector<int> (nAgents, 0);

    // reset distance moved
    moved = std::vector<float> (nAgents, 0.f);

    // reset adjacency matrix
    pbsn.adjMat = Rcpp::NumericMatrix(nAgents, nAgents);

    // positions
    std::vector<float> coord_x_2 (nAgents, 0.f);
    std::vector<float> coord_y_2 (nAgents, 0.f);

    for (int a = 0; a < nAgents; a++) {
        size_t parent_id = static_cast<size_t>(weightedLottery(rng));

        tmp_sF[a] = sF[parent_id];
        tmp_sH[a] = sH[parent_id];
        tmp_sN[a] = sN[parent_id];

        // inherit positions from parent
        coord_x_2[a] = coordX[parent_id] + sprout(rng);
        coord_y_2[a] = coordY[parent_id] + sprout(rng);

        // robustly wrap positions
        if(coord_x_2[a] < 0.f) coord_x_2[a] = food.dSize + coord_x_2[a];
        if(coord_x_2[a] > food.dSize) coord_x_2[a] = coord_x_2[a] - food.dSize;

        if(coord_y_2[a] < 0.f) coord_y_2[a] = food.dSize + coord_y_2[a];
        if(coord_y_2[a] > food.dSize) coord_y_2[a] = coord_y_2[a] - food.dSize;

        // // vertical transmission of infection.
        // if(infected[parent_id]) {
        //     if(verticalInfect(rng)) {
        //         infected_2[a] = true;
        //         srcInfect[a] = 1;
        //     }
        // }
    }

    // swap infected and infected_2
    std::swap(infected, infected_2);
    infected_2.clear();

    // swap coords --- this initialises individuals near their parent's position
    std::swap(coordX, coord_x_2);
    std::swap(coordY, coord_y_2);
    coord_x_2.clear(); coord_y_2.clear();

    // update initial positions!
    initX = coordX;
    initY = coordY;

    // reset counter and time infected
    counter = std::vector<int> (nAgents, 0);
    timeInfected = std::vector<int> (nAgents, 0);
    assert(static_cast<int>(counter.size()) == nAgents && "counter size wrong");

    // mutate trait: trait shifts up or down with an equal prob
    // trait mutation prob is mProb, in a two step process
    for (int a = 0; a < nAgents; a++) {
        if(mutation_happens(rng)) {
            tmp_sF[a] = tmp_sF[a] + mutation_size(rng);
        }
        if(mutation_happens(rng)) {
            tmp_sH[a] = tmp_sH[a] + mutation_size(rng);
        }
        if(mutation_happens(rng)) {
            tmp_sN[a] = tmp_sN[a] + mutation_size(rng);
        }
    }

    // reset nInfected and count natal infections
    // from vertical transmission
    countInfected();

    // swap trait matrices
    std::swap(sF, tmp_sF);
    std::swap(sH, tmp_sH);
    std::swap(sN, tmp_sN);

    tmp_sF.clear(); tmp_sH.clear(); tmp_sN.clear();

    // swap energy
    std::vector<float> tmpEnergy (nAgents, 0.001);
    std::swap(energy, tmpEnergy);
    tmpEnergy.clear();

    // swap intake
    std::vector<float> tmpIntake (nAgents, 0.001);
    std::swap(intake, tmpIntake);
    tmpIntake.clear();
}


#ifndef PARMCB_SVA_SIGNED_TBB_HPP_
#define PARMCB_SVA_SIGNED_TBB_HPP_

//    Copyright (C) Dimitrios Michail 2019 - 2021.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          https://www.boost.org/LICENSE_1_0.txt)

#include <cstddef>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <set>
#include <vector>

#include <boost/graph/graph_traits.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/tuple/detail/tuple_basic.hpp>
#include <boost/timer/timer.hpp>
#include <tbb/parallel_for.h>
#include <tbb/parallel_reduce.h>
#include <tbb/concurrent_vector.h>
#include <tbb/task_group.h>

#include <parmcb/detail/signed_dijkstra.hpp>
#include <parmcb/forestindex.hpp>
#include <parmcb/spvecgf2.hpp>
#include <parmcb/util.hpp>

namespace parmcb {

    namespace detail {

        template<class Graph, class WeightMap>
        struct OddCycleFinder {
            typedef typename boost::graph_traits<Graph>::vertex_descriptor Vertex;
            typedef typename boost::graph_traits<Graph>::edge_descriptor Edge;
            typedef typename boost::property_traits<WeightMap>::value_type WeightType;

            OddCycleFinder(const Graph &g, const WeightMap &weight_map, const ForestIndex<Graph> &forest_index,
                    const std::vector<Vertex> &vertices) :
                    g(g), weight_map(weight_map), forest_index(forest_index), vertices(vertices), compare(
                            std::less<WeightType>()) {
            }

            std::tuple<std::set<Edge>, WeightType, bool> find(const SpVecGF2<std::size_t> &support) {
                std::set<Edge> signed_edges;
                convert_edges(support, std::inserter(signed_edges, signed_edges.end()), forest_index);
                if (signed_edges.size() == 1) {
                    return find_single_edge(*signed_edges.begin());
                } else if (signed_edges.size() >= boost::num_vertices(g)) {
                    return find_all_vertices(signed_edges);
                } else {
                    return find_less_than_vertices(signed_edges);
                }
            }

            std::tuple<std::set<Edge>, WeightType, bool> find_single_edge(const Edge &se) {
                std::tuple<std::set<Edge>, WeightType, bool> best;

                auto se_v = boost::source(se, g);
                auto se_u = boost::target(se, g);
                auto res = bidirectional_signed_dijkstra(g, weight_map, std::set<Edge> { }, std::set<Edge> { se }, true,
                        se_v, true, se_u, true, std::get<2>(best), std::get<1>(best));
                if (std::get<2>(res) && std::get<0>(res).find(se) == std::get<0>(res).end()) {
                    std::get<1>(res) += boost::get(weight_map, se);
                    if (!std::get<2>(best) || compare(std::get<1>(res), std::get<1>(best))) {
                        std::get<0>(res).insert(se);
                        best = res;
                    }
                }
                return best;
            }

            std::tuple<std::set<Edge>, WeightType, bool> find_all_vertices(const std::set<Edge> &signed_edges) {
                typedef std::tuple<std::set<Edge>, WeightType, bool> cycle_t;
                auto cycle_min = [&](const cycle_t &c1, const cycle_t &c2) {
                    if (!std::get<2>(c1) || !std::get<2>(c2)) {
                        if (std::get<2>(c1)) {
                            return c1;
                        } else {
                            return c2;
                        }
                    }
                    // both valid, compare
                    if (!compare(std::get<1>(c2), std::get<1>(c1))) {
                        return c1;
                    }
                    return c2;
                };

                return tbb::parallel_reduce(tbb::blocked_range<std::size_t>(0, boost::num_vertices(g)),
                        std::make_tuple(std::set<Edge>(), (std::numeric_limits<WeightType>::max)(), false),
                        [&](tbb::blocked_range<std::size_t> r, auto running_min) {
                            for (std::size_t i = r.begin(); i < r.end(); i++) {
                                auto v = vertices[i];
                                const bool use_hidden_edges = false;
                                auto res = bidirectional_signed_dijkstra(g, weight_map, signed_edges,
                                        std::set<Edge> { }, use_hidden_edges, v, true, v, false,
                                        std::get<2>(running_min), std::get<1>(running_min));
                                if (std::get<2>(res)
                                        && (!std::get<2>(running_min)
                                                || compare(std::get<1>(res), std::get<1>(running_min)))) {
                                    running_min = res;
                                }
                            }
                            return running_min;
                        },
                        cycle_min);

            }

            std::tuple<std::set<Edge>, WeightType, bool> find_less_than_vertices(const std::set<Edge> &signed_edges) {
                /*
                 * Heuristic in case number of signed edges is small compared to the number of vertices.
                 */
                typedef std::tuple<std::set<Edge>, WeightType, bool> cycle_t;
                auto cycle_min = [&](const cycle_t &c1, const cycle_t &c2) {
                    if (!std::get<2>(c1) || !std::get<2>(c2)) {
                        if (std::get<2>(c1)) {
                            return c1;
                        } else {
                            return c2;
                        }
                    }
                    // both valid, compare
                    if (!compare(std::get<1>(c2), std::get<1>(c1))) {
                        return c1;
                    }
                    return c2;
                };

                std::map<Edge, std::set<Edge>> hidden_edges_per_edge;
                std::vector<Edge> signed_edges_as_vector;
                std::set<Edge> tmp_signed_edges = signed_edges;
                while (!tmp_signed_edges.empty()) {
                    auto bit = tmp_signed_edges.begin();
                    hidden_edges_per_edge.insert(std::make_pair(*bit, tmp_signed_edges));
                    signed_edges_as_vector.push_back(*bit);
                    tmp_signed_edges.erase(bit);
                }
                return tbb::parallel_reduce(tbb::blocked_range<std::size_t>(0, signed_edges_as_vector.size()),
                        std::make_tuple(std::set<Edge>(), (std::numeric_limits<WeightType>::max)(), false),
                        [&](tbb::blocked_range<std::size_t> r, auto running_min) {
                            for (std::size_t i = r.begin(); i < r.end(); i++) {
                                auto se = signed_edges_as_vector.at(i);
                                auto se_v = boost::source(se, g);
                                auto se_u = boost::target(se, g);
                                auto hidden_edges = hidden_edges_per_edge.at(se);
                                auto res = bidirectional_signed_dijkstra(g, weight_map, signed_edges, hidden_edges,
                                        true, se_v, true, se_u, true, std::get<2>(running_min),
                                        std::get<1>(running_min));
                                if (std::get<2>(res) && std::get<0>(res).find(se) == std::get<0>(res).end()) {
                                    std::get<1>(res) += boost::get(weight_map, se);
                                    if (!std::get<2>(running_min)
                                            || compare(std::get<1>(res), std::get<1>(running_min))) {
                                        std::get<0>(res).insert(se);
                                        running_min = res;
                                    }
                                }
                            }
                            return running_min;
                        },
                        cycle_min);
            }

            const Graph &g;
            const WeightMap &weight_map;
            const ForestIndex<Graph> &forest_index;
            const std::vector<Vertex> &vertices;
            const std::less<WeightType> compare;
        };

    }

    template<class Graph, class WeightMap, class CycleOutputIterator>
    typename boost::property_traits<WeightMap>::value_type mcb_sva_signed_tbb(const Graph &g, WeightMap weight_map,
            CycleOutputIterator out, const std::size_t hardware_concurrency_hint = 0) {

        typedef typename boost::graph_traits<Graph>::vertex_descriptor Vertex;
        typedef typename boost::graph_traits<Graph>::vertex_iterator VertexIt;
        typedef typename boost::graph_traits<Graph>::edge_descriptor Edge;
        typedef typename boost::property_traits<WeightMap>::value_type WeightType;

        /*
         * Index the graph
         */
        ForestIndex<Graph> forest_index(g);
        auto csd = forest_index.cycle_space_dimension();
        std::cout << "Cycle dimension " << csd << std::endl;
        std::vector<Vertex> vertices;
        {
            VertexIt vi, viend;
            for (boost::tie(vi, viend) = boost::vertices(g); vi != viend; ++vi) {
                vertices.push_back(*vi);
            }
        }

        /*
         * Initialize support vectors
         */
        tbb::concurrent_vector<SpVecGF2<std::size_t>> support;
        tbb::parallel_for(tbb::blocked_range<std::size_t>(0, csd), [&](const tbb::blocked_range<std::size_t> &r) {
            for (std::size_t i = r.begin(); i != r.end(); ++i) {
                support.push_back(SpVecGF2<std::size_t> { i });
            }
        });

        boost::timer::cpu_timer cycle_timer;
        cycle_timer.stop();
        boost::timer::cpu_timer support_timer;
        support_timer.stop();

        /*
         * Main loop
         */
        WeightType mcb_weight = WeightType();
        parmcb::detail::OddCycleFinder<Graph, WeightMap> odd_cycle_finder(g, weight_map, forest_index, vertices);
        for (std::size_t k = 0; k < csd; k++) {
            if (k % 250 == 0) {
                std::cout << "iteration = " << k << std::endl;
            }

            /*
             * Choose the sparsest support heuristic
             */
            auto min_support = k;
            for (auto r = k + 1; r < csd; ++r) {
                if (support[r].size() < support[min_support].size())
                    min_support = r;
            }
            if (min_support != k) {  // swap
                std::swap(support[k], support[min_support]);
            }

            /*
             * Compute shortest odd cycle
             */
            cycle_timer.resume();
            auto cycle = odd_cycle_finder.find(support[k]);
            cycle_timer.stop();

            /*
             * Update support vectors
             */
            support_timer.resume();
            std::set<std::size_t> cyclek;
            convert_edges(std::get<0>(cycle), std::inserter(cyclek, cyclek.end()), forest_index);
            tbb::parallel_for(tbb::blocked_range<std::size_t>(k + 1, csd),
                    [&](const tbb::blocked_range<std::size_t> &r) {
                        auto e = r.end();
                        for (std::size_t i = r.begin(); i != e; ++i) {
                            if (support[i] * cyclek == 1) {
                                support[i] += support[k];
                            }
                        }
                    });
            support_timer.stop();

            /*
             * Report cycle
             */
            std::list<Edge> cyclek_edgelist;
            std::copy(std::get<0>(cycle).begin(), std::get<0>(cycle).end(), std::back_inserter(cyclek_edgelist));
            *out++ = cyclek_edgelist;
            mcb_weight += std::get<1>(cycle);
        }

        std::cout << "cycle   timer" << cycle_timer.format();
        std::cout << "support timer" << support_timer.format();

        return mcb_weight;
    }

} // namespace parmcb

#endif


/***************************************************************************
 * Copyright 1998-2015 by authors (see AUTHORS.txt)                        *
 *                                                                         *
 *   This file is part of LuxRender.                                       *
 *                                                                         *
 * 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.                                          *
 ***************************************************************************/

#include <stdexcept>
#include <boost/foreach.hpp>
#include <boost/regex.hpp>

#include "slg/film/film.h"
#include "slg/film/imagepipeline/plugins/gaussianblur3x3.h"

using namespace std;
using namespace luxrays;
using namespace slg;

//------------------------------------------------------------------------------
// GaussianBlur filter plugin
//------------------------------------------------------------------------------

BOOST_CLASS_EXPORT_IMPLEMENT(slg::GaussianBlur3x3FilterPlugin)

ImagePipelinePlugin *GaussianBlur3x3FilterPlugin::Copy() const {
	return new GaussianBlur3x3FilterPlugin(weight);
}

void GaussianBlur3x3FilterPlugin::ApplyBlurFilterXR1(
		const u_int filmWidth, const u_int filmHeight,
		const Spectrum *src,
		Spectrum *dst,
		const float aF,
		const float bF,
		const float cF) const {
	// Do left edge
	Spectrum a;
	Spectrum b = src[0];
	Spectrum c = src[1];

	const float leftTotF = bF + cF;
	const Spectrum bLeftK = bF / leftTotF;
	const Spectrum cLeftK = cF / leftTotF;
	dst[0] = bLeftK  * b + cLeftK * c;

    // Main loop
	const float totF = aF + bF + cF;
	const Spectrum aK = aF / totF;
	const Spectrum bK = bF / totF;
	const Spectrum cK = cF / totF;

	for (unsigned int x = 1; x < filmWidth - 1; ++x) {
		a = b;
		b = c;
		c = src[x + 1];

		dst[x] = aK * a + bK  * b + cK * c;
    }

    // Do right edge
	const float rightTotF = aF + bF;
	const Spectrum aRightK = aF / rightTotF;
	const Spectrum bRightK = bF / rightTotF;
	a = b;
	b = c;
	dst[filmWidth - 1] = aRightK  * a + bRightK * b;
}

void GaussianBlur3x3FilterPlugin::ApplyBlurFilterYR1(
		const u_int filmWidth, const u_int filmHeight,
		const Spectrum *src,
		Spectrum *dst,
		const float aF,
		const float bF,
		const float cF) const {
	// Do left edge
	Spectrum a;
	Spectrum b = src[0];
	Spectrum c = src[filmWidth];

	const float leftTotF = bF + cF;
	const Spectrum bLeftK = bF / leftTotF;
	const Spectrum cLeftK = cF / leftTotF;
	dst[0] = bLeftK  * b + cLeftK * c;

    // Main loop
	const float totF = aF + bF + cF;
	const Spectrum aK = aF / totF;
	const Spectrum bK = bF / totF;
	const Spectrum cK = cF / totF;

    for (unsigned int y = 1; y < filmHeight - 1; ++y) {
		const unsigned index = y * filmWidth;

		a = b;
		b = c;
		c = src[index + filmWidth];

		dst[index] = aK * a + bK  * b + cK * c;
    }

    // Do right edge
	const float rightTotF = aF + bF;
	const Spectrum aRightK = aF / rightTotF;
	const Spectrum bRightK = bF / rightTotF;
	a = b;
	b = c;
	dst[(filmHeight - 1) * filmWidth] = aRightK  * a + bRightK * b;
}

void GaussianBlur3x3FilterPlugin::ApplyGaussianBlurFilterXR1(
		const u_int filmWidth, const u_int filmHeight,
		const Spectrum *src, Spectrum *dst) const {
	const float aF = weight;
	const float bF = 1.f;
	const float cF = weight;

	ApplyBlurFilterXR1(filmWidth, filmHeight, src, dst, aF, bF, cF);
}

void GaussianBlur3x3FilterPlugin::ApplyGaussianBlurFilterYR1(
		const u_int filmWidth, const u_int filmHeight,
		const Spectrum *src, Spectrum *dst) const {
	const float aF = weight;
	const float bF = 1.f;
	const float cF = weight;

	ApplyBlurFilterYR1(filmWidth, filmHeight, src, dst, aF, bF, cF);
}

void GaussianBlur3x3FilterPlugin::Apply(const Film &film, Spectrum *pixels, vector<bool> &pixelsMask) const {
	const u_int width = film.GetWidth();
	const u_int height = film.GetHeight();

	// Allocate the temporary buffer if required
	if ((!tmpBuffer) || (width * height != tmpBufferSize)) {
		delete tmpBuffer;

		tmpBufferSize = width * height;
		tmpBuffer = new Spectrum[tmpBufferSize];
	}

	for (u_int i = 0; i < 3; ++i) {
		for (u_int y = 0; y < height; ++y) {
			const u_int index = y * width;
			ApplyGaussianBlurFilterXR1(width, height, &pixels[index], &tmpBuffer[index]);
		}

		for (u_int x = 0; x < width; ++x)
			ApplyGaussianBlurFilterYR1(width, height, &tmpBuffer[x], &pixels[x]);
	}
}


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

#ifdef USE_GUI

#include <Eigen/Dense>

#include <cstring>
#include <cmath>

#include <NVX/Application.hpp>
#include <NVX/ProfilerRange.hpp>

#include "Render/CUDA-OpenGL/BaseRender3DImpl.hpp"

// row-major storage order for compatibility with vx_matrix
typedef Eigen::Matrix<float, 4, 4, Eigen::RowMajor> Matrix4f_rm;

//============================================================
// Utility stuff
//============================================================

namespace
{

float toRadians(float degrees)
{
    return degrees * 0.017453292519f;
}

void multiplyMatrix(const nvidiaio::matrix4x4f_t & m1,
                    const nvidiaio::matrix4x4f_t & m2,
                    nvidiaio::matrix4x4f_t & res)
{
    std::memset((void *)res.ptr, 0, sizeof(float) * 16);

    for (int32_t i = 0; i < 4; ++i)
        for (int32_t j = 0; j < 4; ++j)
            for (int32_t k = 0; k < 4; ++k)
                res.ptr[4 *i + j] += m1.ptr[4 * i + k] * m2.ptr[4 * k + j];
}

void calcProjectionMatrix(float fovY, float aspect, float zNear, float zFar, nvidiaio::matrix4x4f_t & projection)
{
    std::memset((void *)projection.ptr, 0, sizeof(float) * 16);

    float ctg = 1.f / tan( fovY/2 );

    projection.ptr[0] = ctg / aspect;
    projection.ptr[5] = ctg;
    projection.ptr[10] = - (zFar + zNear) / (zFar - zNear);
    projection.ptr[11] = - 1.f;
    projection.ptr[14] = - 2.f * zFar * zNear / (zFar - zNear);
}

Matrix4f_rm yawPitchRoll(float yaw, float pitch, float roll)
{
    float tmp_ch = std::cos(yaw);
    float tmp_sh = std::sin(yaw);
    float tmp_cp = std::cos(pitch);
    float tmp_sp = std::sin(pitch);
    float tmp_cb = std::cos(roll);
    float tmp_sb = std::sin(roll);

    Matrix4f_rm Result;
    Result(0,0) = tmp_ch * tmp_cb + tmp_sh * tmp_sp * tmp_sb;
    Result(0,1) = tmp_sb * tmp_cp;
    Result(0,2) = -tmp_sh * tmp_cb + tmp_ch * tmp_sp * tmp_sb;
    Result(0,3) = 0.0f;
    Result(1,0) = -tmp_ch * tmp_sb + tmp_sh * tmp_sp * tmp_cb;
    Result(1,1) = tmp_cb * tmp_cp;
    Result(1,2) = tmp_sb * tmp_sh + tmp_ch * tmp_sp * tmp_cb;
    Result(1,3) = 0.0f;
    Result(2,0) = tmp_sh * tmp_cp;
    Result(2,1) = -tmp_sp;
    Result(2,2) = tmp_ch * tmp_cp;
    Result(2,3) = 0.0f;
    Result(3,0) = 0.0f;
    Result(3,1) = 0.0f;
    Result(3,2) = 0.0f;
    Result(3,3) = 1.0f;

    return Result;
}

void lookAt(nvidiaio::matrix4x4f_t & view, const Eigen::Vector3f& eye, const Eigen::Vector3f& center, const Eigen::Vector3f& up)
{
    Eigen::Vector3f f = Eigen::Vector3f(center - eye).normalized();
    Eigen::Vector3f s = f.cross(up).normalized();
    Eigen::Vector3f u = s.cross(f);

    Matrix4f_rm result = Matrix4f_rm::Identity();
    result(0,0) = s(0);
    result(1,0) = s(1);
    result(2,0) = s(2);
    result(0,1) = u(0);
    result(1,1) = u(1);
    result(2,1) = u(2);
    result(0,2) = -f(0);
    result(1,2) = -f(1);
    result(2,2) = -f(2);
    result(3,0) = -s.dot(eye);
    result(3,1) = -u.dot(eye);
    result(3,2) = f.dot(eye);

    std::memcpy(view.ptr, result.data(), sizeof(float) * 16);
}

void updateOrbitCamera(nvidiaio::matrix4x4f_t & view, float xr, float yr, float distance, const Eigen::Vector3f & target)
{
    Matrix4f_rm R = yawPitchRoll(xr, yr, 0.0f);

    Eigen::Vector3f T(0.0f, 0.0f, -distance);

    Eigen::Vector4f T_ = R * Eigen::Vector4f(T(0), T(1), T(2), 0.0f);

    T = Eigen::Vector3f(T_(0), T_(1), T_(2));

    Eigen::Vector3f position = target + T;

    Eigen::Vector4f up_ = R*Eigen::Vector4f(0.0f, -1.0f, 0.0f, 0.0f);

    Eigen::Vector3f up(up_[0], up_[1], up_[2]);

    lookAt(view, position, target, up);
}

void matrixSetEye(nvidiaio::matrix4x4f_t & m)
{
    static const float data[4*4] = {1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1};
    std::memcpy(m.ptr, data, sizeof(data));
}

} // namespace

//============================================================
// Callbacks and events
//============================================================

void nvidiaio::BaseRender3DImpl::enableDefaultKeyboardEventCallback()
{
    if (!useDefaultCallback_)
    {
        useDefaultCallback_ = true;

        fov_ = defaultFOV_;
        initMVP();
    }
}

void nvidiaio::BaseRender3DImpl::disableDefaultKeyboardEventCallback()
{
    if(useDefaultCallback_)
    {
        useDefaultCallback_ = false;

        fov_ = defaultFOV_;
        orbitCameraParams_.setDefault();
        initMVP();
    }
}

bool nvidiaio::BaseRender3DImpl::useDefaultKeyboardEventCallback()
{
    return useDefaultCallback_;
}

void nvidiaio::BaseRender3DImpl::setOnKeyboardEventCallback(nvidiaio::Render3D::OnKeyboardEventCallback callback, void * context)
{
    keyboardCallback_ = callback;
    keyboardCallbackContext_ = context;

    glfwSetKeyCallback(window_, keyboardCallback);
}

void nvidiaio::BaseRender3DImpl::keyboardCallbackDefault(GLFWwindow* window, int key, int /*scancode*/, int action, int /*mods*/)
{
    nvidiaio::BaseRender3DImpl* impl = static_cast<nvidiaio::BaseRender3DImpl*>(glfwGetWindowUserPointer(window));

    if (action == GLFW_PRESS)
    {
        static const float stepAngle = toRadians(4); // in radians
        static const float stepR = 1;

        switch (key)
        {
            case GLFW_KEY_ESCAPE:
            {
                glfwSetWindowShouldClose(window, 1);
                break;
            }
            case GLFW_KEY_MINUS:
            {
                impl->orbitCameraParams_.R += stepR;
                break;
            }
            case GLFW_KEY_EQUAL:
            {
                impl->orbitCameraParams_.R -= stepR;
                break;
            }

            case GLFW_KEY_A:
            {
                impl->orbitCameraParams_.xr -= stepAngle;
                break;
            }
            case GLFW_KEY_D:
            {
                impl->orbitCameraParams_.xr += stepAngle;
                break;
            }
            case GLFW_KEY_W:
            {
                impl->orbitCameraParams_.yr += stepAngle;
                break;
            }
            case GLFW_KEY_S:
            {
                impl->orbitCameraParams_.yr -= stepAngle;
                break;
            }
        }
        impl->updateView();
    }
}

void nvidiaio::BaseRender3DImpl::keyboardCallback(GLFWwindow* window, int key, int scancode, int action, int mods)
{
    nvidiaio::BaseRender3DImpl* impl = static_cast<nvidiaio::BaseRender3DImpl*>(glfwGetWindowUserPointer(window));

    if (action == GLFW_PRESS)
    {
        double x, y;
        glfwGetCursorPos(window, &x, &y);

        if (key == GLFW_KEY_ESCAPE)
            key = 27;

        if(impl->keyboardCallback_)
            (impl->keyboardCallback_)(impl->keyboardCallbackContext_, tolower(key),
                                      static_cast<uint32_t>(x),
                                      static_cast<uint32_t>(y));

        if (impl->useDefaultKeyboardEventCallback())
        {
            keyboardCallbackDefault(window, key, scancode, action, mods);
        }
    }
}

void nvidiaio::BaseRender3DImpl::setOnMouseEventCallback(OnMouseEventCallback callback, void * context)
{
    mouseCallback_ = callback;
    mouseCallbackContext_ = context;

    glfwSetMouseButtonCallback(window_, mouse_button);
    glfwSetCursorPosCallback(window_, cursor_pos);
}

void nvidiaio::BaseRender3DImpl::mouse_button(GLFWwindow* window, int button, int action, int /*mods*/)
{
    BaseRender3DImpl* impl = static_cast<BaseRender3DImpl*>(glfwGetWindowUserPointer(window));

    if (impl->mouseCallback_)
    {
        Render3D::MouseButtonEvent event = nvxio::Render3D::MouseMove;

        if (button == GLFW_MOUSE_BUTTON_LEFT)
        {
            if (action == GLFW_RELEASE)
                event = nvxio::Render3D::LeftButtonUp;
            else
                event = nvxio::Render3D::LeftButtonDown;
        }
        if (button == GLFW_MOUSE_BUTTON_RIGHT)
        {
            if (action == GLFW_RELEASE)
                event = nvxio::Render3D::RightButtonUp;
            else
                event = nvxio::Render3D::RightButtonDown;
        }
        if (button == GLFW_MOUSE_BUTTON_MIDDLE)
        {
            if (action == GLFW_RELEASE)
                event = nvxio::Render3D::MiddleButtonUp;
            else
                event = nvxio::Render3D::MiddleButtonDown;
        }

        double x = 0.0, y = 0.0;
        glfwGetCursorPos(window, &x, &y);
        (impl->mouseCallback_)(impl->mouseCallbackContext_, event,
                               static_cast<uint32_t>(x),
                               static_cast<uint32_t>(y));
    }
}

void nvidiaio::BaseRender3DImpl::cursor_pos(GLFWwindow* window, double x, double y)
{
    BaseRender3DImpl* impl = static_cast<BaseRender3DImpl*>(glfwGetWindowUserPointer(window));

    if (impl->mouseCallback_)
        (impl->mouseCallback_)(impl->mouseCallbackContext_, nvxio::Render3D::MouseMove,
                               static_cast<uint32_t>(x),
                               static_cast<uint32_t>(y));
}

//============================================================
// Orbit Camera Params
//============================================================

nvidiaio::BaseRender3DImpl::OrbitCameraParams::OrbitCameraParams(float R_min_, float R_max_) :
    R_min(R_min_), R_max(R_max_)
{
    setDefault();
}

void nvidiaio::BaseRender3DImpl::OrbitCameraParams::applyConstraints()
{
    if (R < R_min) R = R_min;
    else if (R > R_max) R = R_max;

    if (yr > 2 * ovxio::PI_F) yr -= 2 * ovxio::PI_F;
    else if (yr < 0) yr += 2 * ovxio::PI_F;

    if (xr > 2 * ovxio::PI_F) xr -= 2 * ovxio::PI_F;
    else if (xr < 0) xr += 2 * ovxio::PI_F;
}

void nvidiaio::BaseRender3DImpl::OrbitCameraParams::setDefault()
{
    R = R_min;
    xr = 0.0f;
    yr = 0.0f;
}

//============================================================
// Setters and getters
//============================================================

void nvidiaio::BaseRender3DImpl::setDefaultFOV(float fov)
{
    defaultFOV_ = std::max(std::min(fov, 180.0f), 0.0f);
}

void nvidiaio::BaseRender3DImpl::setViewMatrix(const matrix4x4f_t & view)
{
    view_ = view;
}

void nvidiaio::BaseRender3DImpl::setProjectionMatrix(const matrix4x4f_t & projection)
{
    projection_ = projection;
}

void nvidiaio::BaseRender3DImpl::getViewMatrix(matrix4x4f_t & view) const
{
    view = view_;
}

void nvidiaio::BaseRender3DImpl::getProjectionMatrix(matrix4x4f_t & projection) const
{
    projection = projection_;
}

void nvidiaio::BaseRender3DImpl::setModelMatrix(const matrix4x4f_t & model)
{
    model_ = model;
}

//============================================================
// Put objects to OpenGL framebuffer
//============================================================

void nvidiaio::BaseRender3DImpl::putPointCloud(const array_t & points, const matrix4x4f_t & model, const PointCloudStyle& style)
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::putPointCloud (NVXIO)");

    setModelMatrix(model);

    // MVP = model * view * projection
    matrix4x4f_t MVP;

    multiplyMatrix(model_, view_, MVP);
    multiplyMatrix(matrix4x4f_t(MVP), projection_,  MVP);

    OpenGLContextSafeSetter setter(holder_);

    // invoke
    pointCloudRender_.render(points, MVP, style);
}

void nvidiaio::BaseRender3DImpl::putPlanes(const array_t & planes, const matrix4x4f_t & model, const nvidiaio::Render3D::PlaneStyle& style)
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::putPlanes (NVXIO)");

    setModelMatrix(model);

    // MVP = model * view * projection
    matrix4x4f_t MVP;

    multiplyMatrix(model_, view_, MVP);
    multiplyMatrix(matrix4x4f_t(MVP), projection_,  MVP);

    OpenGLContextSafeSetter setter(holder_);

    // Invoke
    fencePlaneRender_.render(planes, MVP, style);
}

void nvidiaio::BaseRender3DImpl::putImage(const image_t & image)
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::putImage (NVXIO)");

    // calculate actual ScaleRatio that will be applied to other primitives like lines, circles, etc.

    uint32_t imageWidth = image.width, imageHeight = image.height;

    float widthRatio = static_cast<float>(windowWidth_) / imageWidth;
    float heightRatio = static_cast<float>(windowHeight_) / imageHeight;
    scaleRatio_ = std::min(widthRatio, heightRatio);

    textureWidth_ = static_cast<uint32_t>(scaleRatio_ * imageWidth);
    textureHeight_ = static_cast<uint32_t>(scaleRatio_ * imageHeight);

    OpenGLContextSafeSetter setter(holder_);

    if (image.format == NVXCU_DF_IMAGE_NV12)
        nv12ImageRender_.render(image, textureWidth_, textureHeight_);
    else
        imageRender_.render(image, textureWidth_, textureHeight_);
}

void nvidiaio::BaseRender3DImpl::putText(const std::string& text, const nvidiaio::Render::TextBoxStyle& style)
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::putText (NVXIO)");

    OpenGLContextSafeSetter setter(holder_);
    textRender_.render(text, *(const nvidiaio::Render::TextBoxStyle *)&style, textureWidth_, textureHeight_, scaleRatio_);
}

//============================================================
// Initialize and deinitialize
//============================================================

nvidiaio::BaseRender3DImpl::BaseRender3DImpl():
    nvidiaio::Render3D(nvxio::Render3D::BASE_RENDER_3D, "BaseOpenGlRender3D"),
    gl_(nullptr),
    model_(),
    view_(),
    projection_(),
    window_(nullptr),
    holder_(nullptr),
    keyboardCallback_(nullptr),
    keyboardCallbackContext_(nullptr),
    useDefaultCallback_(true),
    windowWidth_(0u),
    windowHeight_(0u),
    textureWidth_(0u),
    textureHeight_(0u),
    defaultFOV_(70),// in degrees
    Z_NEAR_(0.01f),
    Z_FAR_(500.0f),
    fov_(defaultFOV_),
    orbitCameraParams_(1e-6f, 100.f)
{
}

void nvidiaio::BaseRender3DImpl::initMVP()
{
    matrixSetEye(model_);

    const static float viewData[4*4] = {1, 0, 0, 0,
                                        0, -1, 0, 0,
                                        0, 0, -1, 0,
                                        0, 0, 0, 1};

    std::memcpy(view_.ptr, viewData, sizeof(viewData));

    calcProjectionMatrix(toRadians(fov_), (float)windowWidth_ / windowHeight_, Z_NEAR_, Z_FAR_, projection_);
}

void nvidiaio::BaseRender3DImpl::updateView()
{
    orbitCameraParams_.applyConstraints();
    updateOrbitCamera(view_, orbitCameraParams_.xr, orbitCameraParams_.yr, orbitCameraParams_.R, Eigen::Vector3f(0, 0, 0));
}

namespace {

class GLFWContextHolderImpl :
        public nvidiaio::OpenGLContextHolder
{
public:
    explicit GLFWContextHolderImpl(GLFWwindow * currentWindow_) :
        prevWindow(nullptr), currentWindow(currentWindow_)
    {
        if (!currentWindow)
            NVXIO_THROW_EXCEPTION("The render is closed, you must open it before");
    }

    virtual void set()
    {
        // save current context
        prevWindow = glfwGetCurrentContext();
        // attach our OpenGL context
        glfwMakeContextCurrent(currentWindow);
    }

    virtual void unset()
    {
        // attach previous context
        glfwMakeContextCurrent(prevWindow);
    }

private:
    GLFWwindow * prevWindow, * currentWindow;
};

}

void nvidiaio::BaseRender3DImpl::createOpenGLContextHolder()
{
    holder_ = std::make_shared<GLFWContextHolderImpl>(window_);
}

bool nvidiaio::BaseRender3DImpl::open(int32_t xPos, int32_t yPos, uint32_t windowWidth, uint32_t windowHeight, const std::string& windowTitle)
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::open (NVXIO)");

    windowWidth_ = windowWidth;
    windowHeight_ = windowHeight;

    if (!nvxio::Application::get().initGui())
    {
        NVXIO_PRINT("Error: Failed to init GUI");
        return false;
    }

    bool result = initWindow(xPos, yPos, windowWidth_, windowHeight_, windowTitle.c_str());

    fov_ = defaultFOV_;
    initMVP();

    return result;
}

bool nvidiaio::BaseRender3DImpl::initWindow(int32_t xpos, int32_t ypos, uint32_t width, uint32_t height, const std::string& wintitle)
{
    int count = 0;
    GLFWmonitor ** monitors = glfwGetMonitors(&count);
    if (count == 0)
    {
        NVXIO_THROW_EXCEPTION("GLFW: no monitors found");
    }

    int maxPixels = 0;
    const GLFWvidmode* mode = nullptr;

    for (int i = 0; i < count; ++i)
    {
        const GLFWvidmode* currentMode = glfwGetVideoMode(monitors[i]);
        int currentPixels = currentMode->width * currentMode->height;

        if (maxPixels < currentPixels)
        {
            mode = currentMode;
            maxPixels = currentPixels;
        }
    }

    uint32_t cur_width = 0u, cur_height = 0u;
    if ((width <= (uint32_t)mode->width) && (height <= (uint32_t)mode->height))
    {
        cur_width = width;
        cur_height = height;
    }
    else
    {
        float widthRatio = static_cast<float>(mode->width) / width;
        float heightRatio = static_cast<float>(mode->height) / height;
        float scaleRatio = std::min(widthRatio, heightRatio);
        cur_width = static_cast<uint32_t>(scaleRatio * width);
        cur_height = static_cast<uint32_t>(scaleRatio * height);
    }

    glfwWindowHint(GLFW_RESIZABLE, 0);
#ifdef USE_GLES
    glfwWindowHint(GLFW_CLIENT_API, GLFW_OPENGL_ES_API);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
#else
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
#endif
    window_ = glfwCreateWindow(cur_width, cur_height,
                               wintitle.c_str(),
                               nullptr, nullptr);
    if (!window_)
    {
        NVXIO_PRINT("Error: Failed to create GLFW window");
        return false;
    }

    glfwSetWindowUserPointer(window_, this);
    glfwSetWindowPos(window_, xpos, ypos);
    glfwSetInputMode(window_, GLFW_STICKY_KEYS, GL_TRUE);

    // Create OpenGL context holder

    createOpenGLContextHolder();
    OpenGLContextSafeSetter setter(holder_);

    // Initialize OpenGL renders

    if (!gl_)
    {
        gl_ = std::make_shared<nvidiaio::GLFunctions>();
        nvidiaio::loadGLFunctions(gl_.get());
    }

    if (!imageRender_.init(gl_, width, height))
        return false;

    if (!nv12ImageRender_.init(gl_, width, height))
        return false;

    if (!textRender_.init(gl_))
        return false;

    if (!pointCloudRender_.init(gl_))
        return false;

    if (!fencePlaneRender_.init(gl_))
        return false;

    return true;
}

bool nvidiaio::BaseRender3DImpl::flush()
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::flush (NVXIO)");

    if (!window_)
        NVXIO_THROW_EXCEPTION("The render is closed, you must open it before");

    if (glfwWindowShouldClose(window_))
    {
        close();
        return false;
    }

    // GLFW says that we don't need current OpenGL context, but
    // it's wrong for EGL (OpenGL ES).
    // See EGL 1.4 spec. 3.9.3. Posting Semantics;
    // See EGL 1.5 spec. 3.10.3. Posting Semantics;
    {
        OpenGLContextSafeSetter setter(holder_);
        glfwSwapBuffers(window_);
    }

    glfwPollEvents();

    clearGlBuffer();

    return true;
}

void nvidiaio::BaseRender3DImpl::clearGlBuffer()
{
    OpenGLContextSafeSetter setter(holder_);

    gl_->ClearColor(0.0f, 0.0f, 0.0f, 1.0f);
    NVXIO_CHECK_GL_ERROR();
    gl_->Clear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    NVXIO_CHECK_GL_ERROR();
}

void nvidiaio::BaseRender3DImpl::close()
{
    nvxio::ProfilerRange range(nvxio::COLOR_ARGB_FUSCHIA, "Render3D::close (NVXIO)");

    if (window_)
    {
        OpenGLContextSafeSetter setter(holder_);

        textRender_.release();
        imageRender_.release();
        nv12ImageRender_.release();
        pointCloudRender_.release();
        fencePlaneRender_.release();

        glfwDestroyWindow(window_);
        window_ = nullptr;
    }
}

nvidiaio::BaseRender3DImpl::~BaseRender3DImpl()
{
    close();
}

#endif


//MIT License
//
//Copyright(c) 2020 Zheng Jiaqi @NUSComputing
//
//Permission is hereby granted, free of charge, to any person obtaining a copy
//of this software and associated documentation files(the "Software"), to deal
//in the Software without restriction, including without limitation the rights
//to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//copies of the Software, and to permit persons to whom the Software is
//furnished to do so, subject to the following conditions :
//
//The above copyright notice and this permission notice shall be included in all
//copies or substantial portions of the Software.
//
//THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.IN NO EVENT SHALL THE
//AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//SOFTWARE.

#include <CGAL/Surface_mesh_parameterization/Discrete_conformal_map_parameterizer_3.h>
#include <CGAL/Surface_mesh_parameterization/Discrete_authalic_parameterizer_3.h>
#include <CGAL/Surface_mesh_parameterization/Square_border_parameterizer_3.h>
#include <CGAL/Surface_mesh_parameterization/Two_vertices_parameterizer_3.h>
#include <CGAL/Surface_mesh_parameterization/internal/Containers_filler.h>
#include <CGAL/Surface_mesh_parameterization/ARAP_parameterizer_3.h>
#include <CGAL/Surface_mesh_parameterization/LSCM_parameterizer_3.h>
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>
#include <CGAL/Polygon_mesh_processing/connected_components.h>
#include <CGAL/Surface_mesh_parameterization/parameterize.h>
#include <CGAL/Surface_mesh_parameterization/IO/File_off.h>
#include <CGAL/Polygon_mesh_processing/detect_features.h>
#include <CGAL/Polygon_mesh_processing/smooth_mesh.h>
#include <CGAL/AABB_face_graph_triangle_primitive.h>
#include <CGAL/Polygon_mesh_processing/measure.h>
#include <CGAL/Polygon_mesh_processing/locate.h>
#include <CGAL/Polygon_mesh_processing/remesh.h>
#include <boost/property_map/property_map.hpp>
#include <boost/function_output_iterator.hpp>
#include <CGAL/boost/graph/Seam_mesh.h>
#include <CGAL/boost/graph/iterator.h>
#include <CGAL/disable_warnings.h>
#include <CGAL/Simple_cartesian.h>
#include <boost/unordered_map.hpp>
#include <boost/unordered_set.hpp>
#include <CGAL/Unique_hash_map.h>
#include <CGAL/Polyhedron_3.h>
#include <CGAL/Surface_mesh.h>
#include <CGAL/AABB_traits.h>
#include <CGAL/circulator.h>
#include <CGAL/AABB_tree.h>
#include <iostream>
#include <fstream>
#include <cstddef>
#include <vector>

#include "gDel2D/gDel2D/GpuDelaunay.h"
#include "gDel2D/gDel2D/PerfTimer.h"
#include "gDel2D/gDel2D/CPU/PredWrapper.h"
#include "gDel2D/DelaunayChecker.h"
#include "delaunay.h"

#if defined(_WIN32)
#include <Windows.h>
#endif

//#define OUTPUT_PARAMETERIZATION
//#define VISUALIZE_TRIANGULATION
#ifdef VISUALIZE_TRIANGULATION
#include "gDel2D/Visualizer.h"
#endif

typedef CGAL::Simple_cartesian<double>	Kernel;

#include "split.h"
#include "weighting.h"
#include "triangulation.h"
#include "discretization.h"
#include "constrains.h"
#include "gcvt.h"
#include "recover.h"

typedef Kernel::Point_2					Point_2;
typedef Kernel::Point_3					Point_3;
typedef CGAL::Surface_mesh<Point_3>		SurfMesh;
typedef CGAL::Surface_mesh<Point_2>		Mesh2D;

typedef boost::graph_traits<SurfMesh>::edge_descriptor			SM_edge_descriptor;
typedef boost::graph_traits<SurfMesh>::halfedge_descriptor		SM_halfedge_descriptor;
typedef boost::graph_traits<SurfMesh>::vertex_descriptor		SM_vertex_descriptor;

typedef boost::graph_traits<Mesh2D>::edge_descriptor			edge_2d;
typedef boost::graph_traits<Mesh2D>::halfedge_descriptor		halfedge_2d;
typedef boost::graph_traits<Mesh2D>::vertex_descriptor			vertex_2d;

typedef CGAL::Unique_hash_map<SM_halfedge_descriptor, Point_2>	UV_uhm;
typedef CGAL::Unique_hash_map<SM_edge_descriptor, bool>			Seam_edge_uhm;
typedef CGAL::Unique_hash_map<SM_vertex_descriptor, bool>		Seam_vertex_uhm;

typedef boost::associative_property_map<UV_uhm>					UV_pmap;
typedef boost::associative_property_map<Seam_edge_uhm>			Seam_edge_pmap;
typedef boost::associative_property_map<Seam_vertex_uhm>		Seam_vertex_pmap;

typedef CGAL::Seam_mesh<SurfMesh, Seam_edge_pmap, Seam_vertex_pmap>	SeamMesh;

typedef boost::graph_traits<SeamMesh>::vertex_descriptor		vertex_descriptor;
typedef boost::graph_traits<SeamMesh>::edge_descriptor			edge_descriptor;
typedef boost::graph_traits<SeamMesh>::halfedge_descriptor		halfedge_descriptor;
typedef boost::graph_traits<SeamMesh>::face_descriptor			face_descriptor;

typedef boost::unordered_map<vertex_descriptor, vertex_2d>		_3d_to_2d_uhm;
typedef boost::unordered_map<vertex_2d, vertex_descriptor>		_2d_to_3d_uhm;
typedef boost::unordered_map<vertex_2d, double>					weight_uhm;

typedef boost::associative_property_map<_3d_to_2d_uhm>			_3d_to_2d_pmap;
typedef boost::associative_property_map<_2d_to_3d_uhm>			_2d_to_3d_pmap;
typedef boost::associative_property_map<weight_uhm>				weight_pmap;

namespace PMP = CGAL::Polygon_mesh_processing;
namespace SMP = CGAL::Surface_mesh_parameterization;

typedef PMP::Face_location<SurfMesh, Kernel::FT>				Face_Location;

int imageSize	= 2048;
int depth		= 1;
int maxIter		= 1000;
int vertices    = 20000;
double scale, leftbound, lowerbound;
float *densityMap;
short *Voronoi;
bool *constrainMask;

GDel2DInput		DelInput;
GDel2DOutput	DelOutput;

int main(int argc, char** argv)
{
	// input surface mesh
	const char* filename = argc > 1 ? argv[1] : "data/horse.off";
	std::ifstream input(filename);
	SurfMesh surface;
	if (!input || !(input >> surface) || surface.is_empty()) {
		std::cerr << "Not a valid .off file." << std::endl;
		return EXIT_FAILURE;
	}
	input.close();

	double edge_length_limit = 0.72;

	// split long edges
	std::vector<SM_edge_descriptor> constrain_vec;
	detectConstrains(surface, constrain_vec);
	PMP::split_long_edges(constrain_vec, edge_length_limit, surface);

	// create seam mesh
	const char* seamfile = "data/horse.selection.txt";
	Seam_edge_uhm seam_edge_uhm(false);
	Seam_edge_pmap seam_edges(seam_edge_uhm);
	Seam_vertex_uhm seam_vertex_uhm(false);
	Seam_vertex_pmap seam_vertices(seam_vertex_uhm);
	std::vector<SM_edge_descriptor> seam_vec;
	addSeams(surface, seam_edges, seam_vec, seamfile);
	PMP::split_long_edges(seam_vec, edge_length_limit, surface, PMP::parameters::edge_is_constrained_map(seam_edges));
	getSeamVertices(surface, seam_edges, seam_vertices);

	SeamMesh seam_mesh(surface, seam_edges, seam_vertices);

	// the property map stores the 2D points in the uv parameter space
	UV_uhm uv_uhm;
	UV_pmap uv_coord(uv_uhm);

	// a halfedge on the border
	halfedge_descriptor border_halfedge = PMP::longest_border(seam_mesh).first;

	// planar parameterization
	//typedef SMP::Square_border_uniform_parameterizer_3<SeamMesh> Border_parameterizer;
	typedef SMP::ARAP_parameterizer_3<SeamMesh> Parameterizer;
	SMP::parameterize(seam_mesh, Parameterizer(), border_halfedge, uv_coord);
	printf("Parameterization done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);
#ifdef OUTPUT_PARAMETERIZATION
	std::ofstream out("parameterization.off");
	SMP::IO::output_uvmap_to_off(seam_mesh, border_halfedge, uv_coord, out);
#endif

	// create 2d triangulation
	Mesh2D mesh_2d;
	_3d_to_2d_uhm _32_uhm;
	_3d_to_2d_pmap vertex_3d_to_2d(_32_uhm);
	_2d_to_3d_uhm _23_uhm;
	_2d_to_3d_pmap vertex_2d_to_3d(_23_uhm);
	build_2d_triangulation(seam_mesh, border_halfedge, uv_coord, mesh_2d, vertex_3d_to_2d, vertex_2d_to_3d);
	printf("Triangulation done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// weighting vertices
	weight_uhm wgt_uhm;
	weight_pmap vertex_weight(wgt_uhm);
	weighting_vertices_with_area(seam_mesh, mesh_2d, vertex_3d_to_2d, vertex_weight);
	printf("Weighting done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// discretization
	densityMap = (float *)malloc(imageSize * imageSize * sizeof(float));
	discretization(mesh_2d, vertex_weight, densityMap, imageSize, scale, leftbound, lowerbound);
	printf("Discretization done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// detect constrains
	std::vector<std::pair<vertex_2d, vertex_2d> > constrain_edge;
	std::unordered_set<vertex_2d> constrain_point;
	detect_constrains(seam_mesh, vertex_3d_to_2d, constrain_edge, constrain_point);
	printf("Detect constrains done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// construct centroidal Voronoi diagram
	Voronoi = (short *)malloc(sizeof(short) * imageSize * imageSize * 2);
	constrainMask = (bool *)malloc(sizeof(bool) * imageSize * imageSize);
	generateMask(mesh_2d, constrain_point, constrainMask, imageSize, scale, leftbound, lowerbound);
	centroidalVoronoi(Voronoi, densityMap, constrainMask, vertices, imageSize, depth, maxIter);
	printf("Centroidal Voronoi tessellation done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// construct constraint Delaunay triangulation
#ifdef VISUALIZE_TRIANGULATION
	Visualizer *vis = Visualizer::instance();
	if (vis->isEnable()) {
		vis->init(argc, argv, "gDel2D");
		vis->printHelp();
	}
	Visualizer::instance()->pause();
#endif
	GpuDel gpuDel;
	Point2HVec().swap(DelInput.pointVec);
	SegmentHVec().swap(DelInput.constraintVec);
	TriHVec().swap(DelOutput.triVec);
	TriOppHVec().swap(DelOutput.triOppVec);
	delaunayInput(mesh_2d, Voronoi, constrainMask, constrain_point, constrain_edge, DelInput.pointVec, DelInput.constraintVec, imageSize, scale, leftbound, lowerbound);
	//printf("%d points, %d segments", DelInput.pointVec.size(), DelInput.constraintVec.size());
	gpuDel.compute(DelInput, &DelOutput);
	printf("Constraint Delaunay triangulation done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// Parameter space -> surface
	SurfMesh resultMesh;
	recover(mesh_2d, seam_mesh, vertex_2d_to_3d, DelInput.pointVec, DelOutput.triVec, constrain_point, resultMesh);
	printf("Recover done. %f s\n", clock()*1.0 / CLOCKS_PER_SEC);

	// compute average edge length
	double edge_sum = 0, edge_avg = 0;
	for (auto e : resultMesh.edges()) {
		edge_sum += PMP::edge_length(e, resultMesh);
	}
	edge_avg = edge_sum / resultMesh.number_of_edges();
	printf("Average edge length: %f\n", edge_sum / resultMesh.number_of_halfedges());

	// Output to off file
	std::ofstream ostream("result.off");
	ostream << resultMesh;

	free(densityMap);
	free(Voronoi);
	free(constrainMask);

#ifdef VISUALIZE_TRIANGULATION
	// Visualize the constraint Delaunay triangualtion
	Visualizer::instance()->resume();
	Visualizer::instance()->addFrame(DelInput.pointVec, DelInput.constraintVec, DelOutput.triVec);
	vis->run();
#endif

	return EXIT_SUCCESS;
}


#include <stan/callbacks/stream_logger.hpp>
#include <stan/mcmc/hmc/xhmc/unit_e_xhmc.hpp>
#include <boost/random/additive_combine.hpp>
#include <test/test-models/good/mcmc/hmc/common/gauss3D.hpp>
#include <stan/io/dump.hpp>
#include <fstream>

#include <gtest/gtest.h>

typedef boost::ecuyer1988 rng_t;

TEST(McmcUnitEXHMC, build_tree) {
  rng_t base_rng(4839294);

  stan::mcmc::unit_e_point z_init(3);
  z_init.q(0) = 1;
  z_init.q(1) = -1;
  z_init.q(2) = 1;
  z_init.p(0) = -1;
  z_init.p(1) = 1;
  z_init.p(2) = -1;

  std::stringstream debug, info, warn, error, fatal;
  stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);

  std::fstream empty_stream("", std::fstream::in);
  stan::io::dump data_var_context(empty_stream);
  gauss3D_model_namespace::gauss3D_model model(data_var_context);

  stan::mcmc::unit_e_xhmc<gauss3D_model_namespace::gauss3D_model, rng_t>
    sampler(model, base_rng);

  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.set_nominal_stepsize(0.1);
  sampler.set_stepsize_jitter(0);
  sampler.sample_stepsize();

  stan::mcmc::ps_point z_propose = z_init;

  double ave = 0;
  double log_sum_weight = -std::numeric_limits<double>::infinity();

  double H0 = -0.1;
  int n_leapfrog = 0;
  double sum_metro_prob = 0;

  bool valid_subtree = sampler.build_tree(3, z_propose, ave, log_sum_weight,
                                          H0, 1, n_leapfrog,
                                          sum_metro_prob,
                                          logger);

  EXPECT_EQ(0.1, sampler.get_nominal_stepsize());

  EXPECT_TRUE(valid_subtree);

  EXPECT_FLOAT_EQ(-0.022019938, sampler.z().q(0));
  EXPECT_FLOAT_EQ(0.022019938, sampler.z().q(1));
  EXPECT_FLOAT_EQ(-0.022019938, sampler.z().q(2));

  EXPECT_FLOAT_EQ(-1.4131583, sampler.z().p(0));
  EXPECT_FLOAT_EQ(1.4131583, sampler.z().p(1));
  EXPECT_FLOAT_EQ(-1.4131583, sampler.z().p(2));

  EXPECT_FLOAT_EQ(0.78105003, z_propose.q(0));
  EXPECT_FLOAT_EQ(-0.78105003, z_propose.q(1));
  EXPECT_FLOAT_EQ(0.78105003, z_propose.q(2));

  EXPECT_FLOAT_EQ(-1.1785525, z_propose.p(0));
  EXPECT_FLOAT_EQ(1.1785525, z_propose.p(1));
  EXPECT_FLOAT_EQ(-1.1785525, z_propose.p(2));

  EXPECT_EQ(8, n_leapfrog);
  EXPECT_FLOAT_EQ(4.2207355, ave);
  EXPECT_FLOAT_EQ(std::log(0.36134657), log_sum_weight);
  EXPECT_FLOAT_EQ(0.36134657, sum_metro_prob);

  EXPECT_EQ("", debug.str());
  EXPECT_EQ("", info.str());
  EXPECT_EQ("", warn.str());
  EXPECT_EQ("", error.str());
  EXPECT_EQ("", fatal.str());
}

TEST(McmcUnitEXHMC, transition) {
  rng_t base_rng(4839294);

  stan::mcmc::unit_e_point z_init(3);
  z_init.q(0) = 1;
  z_init.q(1) = -1;
  z_init.q(2) = 1;
  z_init.p(0) = -1;
  z_init.p(1) = 1;
  z_init.p(2) = -1;

  std::stringstream debug, info, warn, error, fatal;
  stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);

  std::fstream empty_stream("", std::fstream::in);
  stan::io::dump data_var_context(empty_stream);
  gauss3D_model_namespace::gauss3D_model model(data_var_context);

  stan::mcmc::unit_e_xhmc<gauss3D_model_namespace::gauss3D_model, rng_t>
    sampler(model, base_rng);

  sampler.z() = z_init;
  sampler.init_hamiltonian(logger);
  sampler.set_nominal_stepsize(0.1);
  sampler.set_stepsize_jitter(0);
  sampler.sample_stepsize();

  stan::mcmc::sample init_sample(z_init.q, 0, 0);

  stan::mcmc::sample s = sampler.transition(init_sample, logger);

  EXPECT_FLOAT_EQ(1, s.cont_params()(0));
  EXPECT_FLOAT_EQ(-1, s.cont_params()(1));
  EXPECT_FLOAT_EQ(1, s.cont_params()(2));
  EXPECT_FLOAT_EQ(-1.5, s.log_prob());
  EXPECT_FLOAT_EQ(0.99805242, s.accept_stat());
  EXPECT_EQ("", debug.str());
  EXPECT_EQ("", info.str());
  EXPECT_EQ("", warn.str());
  EXPECT_EQ("", error.str());
  EXPECT_EQ("", fatal.str());
}


#ifndef OPENGM_PYTHON_INTERFACE
#define OPENGM_PYTHON_INTERFACE 1
#endif

#include <stdexcept>
#include <stddef.h>
#include <string>
#include <boost/python.hpp>
#include <opengm/graphicalmodel/graphicalmodel.hxx>
#include <opengm/inference/inference.hxx>
#include <opengm/inference/gibbs.hxx>
#include "nifty_iterator.hxx"
#include "inferencehelpers.hxx"
#include"../export_typedes.hxx"
using namespace boost::python;

namespace pygibbs{
   template<class PARAM>
   inline void set
   (
      PARAM & p,
      const size_t steps,
      const size_t burnInSteps,
      const pyenums::GibbsVariableProposal  proposal
   ){
      p.maxNumberOfSamplingSteps_=steps;
      p.numberOfBurnInSteps_=burnInSteps;
      if(proposal==pyenums::RANDOM)
         p.variableProposal_=PARAM::RANDOM;
      else
         p.variableProposal_=PARAM::CYCLIC;
   }
   template<class PARAM>
   inline void setProposal
   (
      PARAM & p,
      const pyenums::GibbsVariableProposal  proposal
   ){
      if(proposal==pyenums::RANDOM)
         p.variableProposal_=PARAM::RANDOM;
      else
         p.variableProposal_=PARAM::CYCLIC;
   }
   template<class PARAM>
   inline pyenums::GibbsVariableProposal getProposal
   (
      PARAM & p,
      const pyenums::GibbsVariableProposal  proposal
   ){
      if(p.variableProposal_==PARAM::RANDOM)
         return pyenums::RANDOM;
      else
         return pyenums::CYCLIC;
   }
}

template<class GM,class ACC>
void export_gibbs(){
   import_array();
// Py Inference Types
   typedef opengm::Gibbs<GM, ACC>  PyGibbs;
   typedef typename PyGibbs::Parameter PyGibbsParameter;
   typedef typename PyGibbs::VerboseVisitorType PyGibbsVerboseVisitor;

//   enum_<typename PyGibbsParameter::VariableProposal > ("GibbsVariableProposal")
//           .value("RANDOM", PyGibbsParameter::RANDOM)
//           .value("CYCLIC", PyGibbsParameter::CYCLIC)
           ;
   class_<PyGibbsParameter > ("GibbsParameter", init<const size_t ,const size_t,const typename PyGibbsParameter::VariableProposal> (args("numberOfSamplingSteps,numberOfburnInSteps,variableProposal")))
   .def(init<>())
   .def_readwrite("numberOfSamplingSteps", &PyGibbsParameter::maxNumberOfSamplingSteps_)
   .def_readwrite("numberOfburnInSteps", &PyGibbsParameter::numberOfBurnInSteps_)
   .add_property("variableProposal",&pygibbs::getProposal<PyGibbsParameter>, pygibbs::setProposal<PyGibbsParameter>)
   .def ("set", &pygibbs::set<PyGibbsParameter>,
         (
         arg("steps")=1e5,
         arg("burnInSteps")= 0,
         arg("variableProposal")=pyenums::RANDOM
         )
   )
   ;

   OPENGM_PYTHON_VERBOSE_VISITOR_EXPORTER(PyGibbsVerboseVisitor,"GibbsVerboseVisitor" );
   OPENGM_PYTHON_INFERENCE_EXPORTER(PyGibbs,"Gibbs");

}

template void export_gibbs<GmAdder,opengm::Minimizer>();
//template void export_gibbs<GmAdder,opengm::Maximizer>();
//template void export_gibbs<GmMultiplier,opengm::Minimizer>();
template void export_gibbs<GmMultiplier,opengm::Maximizer>();