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

#include <ros/console.h>
#include <limits.h>
#include <Eigen/LU>
#include <Eigen/Cholesky>
#include <math.h>
#include <stomp_core/utils.h>
#include <numeric>
#include "stomp_core/stomp.h"
#include <stomp_core/rosconsole_extras.h>

static const double DEFAULT_NOISY_COST_IMPORTANCE_WEIGHT = 1.0; /**< Default noisy cost importance weight */
static const double MIN_COST_DIFFERENCE = 1e-8; /**< Minimum cost difference allowed during probability calculation */
static const double MIN_CONTROL_COST_WEIGHT = 1e-8; /**< Minimum control cost weight allowed */

/**
 * @brief Compute a linear interpolated trajectory given a start and end state
 * @param first             The start position
 * @param last              The final position
 * @param num_timesteps     The number of timesteps
 * @param trajectory_joints The returned linear interpolated trajectory
 */
static void computeLinearInterpolation(const std::vector<double>& first,const std::vector<double>& last,
                         int num_timesteps,
                         Eigen::MatrixXd& trajectory_joints)
{
  trajectory_joints.setZero(first.size(),num_timesteps);
  for(int unsigned i = 0; i < first.size(); i++)
  {
    double dtheta = (last[i] - first[i])/(num_timesteps - 1);
    for(unsigned int j = 0; j < num_timesteps; j++)
    {
      trajectory_joints(i,j) = first[i] + j * dtheta;
    }
  }
}

/**
 * @brief Compute a cubic interpolated trajectory given a start and end state
 * @param first             The start position
 * @param last              The final position
 * @param num_points        The number of points in the trajectory
 * @param dt                The timestep in seconds
 * @param trajectory_joints The returned cubic interpolated trajectory
 */
static void computeCubicInterpolation(const std::vector<double>& first,const std::vector<double>& last,
                         int num_points,double dt,
                         Eigen::MatrixXd& trajectory_joints)
{
  std::vector<double> coeffs(4,0);
  double total_time = (num_points - 1) * dt;
  for(int unsigned i = 0; i < first.size(); i++)
  {
    coeffs[0] = first[i];
    coeffs[2] = (3/(pow(total_time,2))) * (last[i] - first[i]);
    coeffs[3] = (-2/(pow(total_time,3))) * (last[i] - first[i]);

    double t;
    for(unsigned j = 0; j < num_points; j++)
    {
      t = j*dt;
      trajectory_joints(i,j) = coeffs[0] + coeffs[2]*pow(t,2) + coeffs[3]*pow(t,3);
    }
  }
}

/**
 * @brief Compute a minimum cost trajectory given a start and end state
 * @param first                        The start position
 * @param last                         The final position
 * @param control_cost_matrix_R_padded The control cost matrix with padding
 * @param inv_control_cost_matrix_R    The inverse constrol cost matrix
 * @param trajectory_joints            The returned minimum cost trajectory
 * @return True if successful, otherwise false
 */
bool computeMinCostTrajectory(const std::vector<double>& first,
                              const std::vector<double>& last,
                              const Eigen::MatrixXd& control_cost_matrix_R_padded,
                              const Eigen::MatrixXd& inv_control_cost_matrix_R,
                              Eigen::MatrixXd& trajectory_joints)
{
  using namespace stomp_core;

  if(control_cost_matrix_R_padded.rows() != control_cost_matrix_R_padded.cols())
  {
    ROS_ERROR("Control Cost Matrix is not square");
    return false;
  }

  int timesteps = control_cost_matrix_R_padded.rows() - 2*(FINITE_DIFF_RULE_LENGTH - 1);
  int start_index_padded = FINITE_DIFF_RULE_LENGTH - 1;
  int end_index_padded = start_index_padded + timesteps-1;
  std::vector<Eigen::VectorXd> linear_control_cost(first.size(),Eigen::VectorXd::Zero(timesteps));
  trajectory_joints.setZero(first.size(),timesteps);

  for(unsigned int d = 0; d < first.size(); d++)
  {
    linear_control_cost[d].transpose() = first[d] * Eigen::VectorXd::Ones(FINITE_DIFF_RULE_LENGTH - 1).transpose() *
        control_cost_matrix_R_padded.block(0,start_index_padded,FINITE_DIFF_RULE_LENGTH - 1,timesteps);

    linear_control_cost[d].transpose() += last[d] * Eigen::VectorXd::Ones(FINITE_DIFF_RULE_LENGTH - 1).transpose() *
        control_cost_matrix_R_padded.block(end_index_padded + 1,start_index_padded,FINITE_DIFF_RULE_LENGTH - 1,
                                    timesteps);

    linear_control_cost[d] *=2;

    trajectory_joints.row(d) = -0.5*inv_control_cost_matrix_R*linear_control_cost[d];
    trajectory_joints(d,0) = first[d];
    trajectory_joints(d,timesteps - 1) = last[d];
  }

  return true;
}

/**
 * @brief Compute the parameters control costs
 * @param parameters            The parameters used to compute the control cost
 * @param dt                    The timestep in seconds
 * @param control_cost_weight   The control cost weight
 * @param control_cost_matrix_R The control cost matrix
 * @param control_costs returns The parameters control costs
 */
void computeParametersControlCosts(const Eigen::MatrixXd& parameters,
                                          double dt,
                                          double control_cost_weight,
                                          const Eigen::MatrixXd& control_cost_matrix_R,
                                          Eigen::MatrixXd& control_costs)
{
  std::size_t num_timesteps = parameters.cols();
  double cost = 0;
  for(auto d = 0u; d < parameters.rows(); d++)
  {
    cost = double(parameters.row(d)*(control_cost_matrix_R*parameters.row(d).transpose()));
    control_costs.row(d).setConstant( 0.5*(1/dt)*cost );
  }

  double max_coeff = control_costs.maxCoeff();
  control_costs /= (max_coeff > 1e-8) ? max_coeff : 1;
  control_costs *= control_cost_weight;
}

namespace stomp_core {

Stomp::Stomp(const StompConfiguration& config,TaskPtr task):
    config_(config),
    task_(task)
{

  resetVariables();

}

bool Stomp::clear()
{
  return resetVariables();
}

void Stomp::setConfig(const StompConfiguration& config)
{
  config_ = config;
  resetVariables();
}

bool Stomp::solve(const std::vector<double>& first,const std::vector<double>& last,
                  Eigen::MatrixXd& parameters_optimized)
{
  // initialize trajectory
  if(!computeInitialTrajectory(first,last))
  {
    ROS_ERROR("Unable to generate initial trajectory");
  }

  return solve(parameters_optimized_,parameters_optimized);
}

bool Stomp::solve(const Eigen::VectorXd& first,const Eigen::VectorXd& last,
                  Eigen::MatrixXd& parameters_optimized)
{
  // converting to std vectors
  std::vector<double> start(first.size());
  std::vector<double> end(last.size());

  Eigen::VectorXd::Map(&start[0],first.size()) = first;
  Eigen::VectorXd::Map(&end[0],last.size()) = last;

  return solve(start,end,parameters_optimized);
}

bool Stomp::solve(const Eigen::MatrixXd& initial_parameters,
                  Eigen::MatrixXd& parameters_optimized)
{
  if(parameters_optimized_.isZero())
  {
    parameters_optimized_ = initial_parameters;
  }

  // check initial trajectory size
  if(initial_parameters.rows() != config_.num_dimensions || initial_parameters.cols() != config_.num_timesteps)
  {
    ROS_ERROR("Initial trajectory dimensions is incorrect");
    return false;
  }
  else
  {
    if(initial_parameters.cols() != config_.num_timesteps)
    {
      ROS_ERROR("Initial trajectory number of time steps is incorrect");
      return false;
    }
  }

  if(not config_.optimization_enabled)
  {
    ROS_YELLOW_STREAM("OPTIMIZATION NOT ENABLED, returning initial trajectory as result.");
    parameters_optimized = initial_parameters;
    return true;
  }

  current_iteration_ = 1;
  unsigned int valid_iterations = 0;
  current_lowest_cost_ = std::numeric_limits<double>::max();

  {
    MeasureTime timer("computeOptimizedCost");

    // computing initialial trajectory cost
    if(!computeOptimizedCost())
    {
      ROS_ERROR("Failed to calculate initial trajectory cost");
      return false;
    }
  }

  while(current_iteration_ <= config_.num_iterations && runSingleIteration())
  {
    ROS_WARN("STOMP completed iteration %i with cost %f",current_iteration_,current_lowest_cost_);

    if(parameters_valid_)
    {
      ROS_INFO("Found valid solution, will iterate %i more time(s) ",
               config_.num_iterations_after_valid - valid_iterations);

      valid_iterations++;
    }
    else
    {
      valid_iterations = 0;
    }

    if(valid_iterations > config_.num_iterations_after_valid)
    {
      break;
    }

    current_iteration_++;
  }

  if(parameters_valid_)
  {
    ROS_INFO("STOMP found a valid solution with cost %f after %i iterations",
             current_lowest_cost_,current_iteration_);
  }
  else
  {
    if (proceed_)
      ROS_ERROR("STOMP failed to find a valid solution after %i iterations",current_iteration_);
    else
      ROS_ERROR_STREAM("Stomp was terminated");
  }

  parameters_optimized = parameters_optimized_;

  // notifying task
  task_->done(parameters_valid_,current_iteration_,current_lowest_cost_,parameters_optimized);

  return parameters_valid_;
}

bool Stomp::resetVariables()
{
  proceed_= true;
  parameters_total_cost_ = 0;
  parameters_valid_ = false;
  num_active_rollouts_ = 0;
  current_iteration_ = 0;

  // verifying configuration
  if(config_.max_rollouts <= config_.num_rollouts)
  {
    ROS_DEBUG_STREAM("'max_rollouts' must be greater than 'num_rollouts_per_iteration'.");
    config_.max_rollouts = config_.num_rollouts + 1; // one more to accommodate optimized trajectory
  }

  // noisy rollouts allocation
  int d = config_.num_dimensions;
  num_active_rollouts_ = 0;
  noisy_rollouts_.resize(config_.max_rollouts);
  reused_rollouts_.resize(config_.max_rollouts);

  // initializing rollout
  Rollout rollout;
  rollout.noise.resize(d, config_.num_timesteps);
  rollout.noise.setZero();

  rollout.parameters_noise.resize(d, config_.num_timesteps);
  rollout.parameters_noise.setZero();

  rollout.probabilities.resize(d, config_.num_timesteps);
  rollout.probabilities.setZero();

  rollout.full_probabilities.clear();
  rollout.full_probabilities.resize(d);

  rollout.full_costs.clear();
  rollout.full_costs.resize(d);

  rollout.control_costs.resize(d, config_.num_timesteps);
  rollout.control_costs.setZero();

  rollout.total_costs.resize(d, config_.num_timesteps);
  rollout.total_costs.setZero();

  rollout.state_costs.resize(config_.num_timesteps);
  rollout.state_costs.setZero();

  rollout.importance_weight = DEFAULT_NOISY_COST_IMPORTANCE_WEIGHT;

  for(unsigned int r = 0; r < config_.max_rollouts ; r++)
  {
    noisy_rollouts_[r] = rollout;
    reused_rollouts_[r] = rollout;
  }

  // parameter updates
  parameters_updates_.resize(d, config_.num_timesteps);
  parameters_updates_.setZero();

  parameters_control_costs_.resize(d, config_.num_timesteps);
  parameters_control_costs_.setZero();

  parameters_state_costs_.resize(config_.num_timesteps);
  parameters_state_costs_.setZero();

  parameters_optimized_.resize(config_.num_dimensions,config_.num_timesteps);
  parameters_optimized_.setZero();

  // generate finite difference matrix
  start_index_padded_ = FINITE_DIFF_RULE_LENGTH-1;
  num_timesteps_padded_ = config_.num_timesteps + 2*(FINITE_DIFF_RULE_LENGTH-1);
  generateFiniteDifferenceMatrix(num_timesteps_padded_,DerivativeOrders::STOMP_ACCELERATION,
                                 config_.delta_t,finite_diff_matrix_A_padded_);

  /* control cost matrix (R = A_transpose * A):
   * Note: Original code multiplies the A product by the time interval.  However this is not
   * what was described in the literature
   */
  control_cost_matrix_R_padded_ = config_.delta_t*finite_diff_matrix_A_padded_.transpose() * finite_diff_matrix_A_padded_;
  control_cost_matrix_R_ = control_cost_matrix_R_padded_.block(
      start_index_padded_,start_index_padded_,config_.num_timesteps,config_.num_timesteps);
  inv_control_cost_matrix_R_ = control_cost_matrix_R_.fullPivLu().inverse();

  /*
   * Applying scale factor to ensure that max(R^-1)==1
   */
  double maxVal = std::abs(inv_control_cost_matrix_R_.maxCoeff());
  control_cost_matrix_R_padded_ *= maxVal;
  control_cost_matrix_R_ *= maxVal;
  inv_control_cost_matrix_R_ /= maxVal; // used in computing the minimum control cost initial trajectory

  return true;
}

bool Stomp::computeInitialTrajectory(const std::vector<double>& first,const std::vector<double>& last)
{
  bool valid = true;

  switch(config_.initialization_method)
  {
    case TrajectoryInitializations::CUBIC_POLYNOMIAL_INTERPOLATION:

      computeCubicInterpolation(first,last,config_.num_timesteps,config_.delta_t,parameters_optimized_);
      break;
    case TrajectoryInitializations::LINEAR_INTERPOLATION:

      computeLinearInterpolation(first,last,config_.num_timesteps,parameters_optimized_);
      break;
    case TrajectoryInitializations::MININUM_CONTROL_COST:

      valid = computeMinCostTrajectory(first,last,control_cost_matrix_R_padded_,inv_control_cost_matrix_R_,parameters_optimized_);
      break;
  }

  return valid;
}

bool Stomp::cancel()
{
  ROS_WARN("Interrupting STOMP");
  proceed_ = false;
  return !proceed_;
}

bool Stomp::runSingleIteration()
{
  MeasureTime timer("runSingleIteration()");

  if(!proceed_)
  {
    return false;
  }

  bool proceed = generateNoisyRollouts() &&
      computeNoisyRolloutsCosts() &&
      filterNoisyRollouts() &&
      computeProbabilities() &&
      updateParameters() &&
      computeOptimizedCost();

  // notifying end of iteration
  task_->postIteration(0,config_.num_timesteps,current_iteration_,current_lowest_cost_,parameters_optimized_);

  return proceed;
}

bool Stomp::generateNoisyRollouts()
{
  MeasureTime timer("generateNoisyRollouts");
  // calculating number of rollouts to reuse from previous iteration
  std::vector< std::pair<double,int> > rollout_cost_sorter; // Used to sort noisy trajectories in ascending order wrt their total cost
  double h = config_.exponentiated_cost_sensitivity;
  int rollouts_stored = num_active_rollouts_-1; // don't take the optimized rollout into account
  rollouts_stored = rollouts_stored < 0 ? 0 : rollouts_stored;
  int rollouts_generate = config_.num_rollouts;
  int rollouts_total = rollouts_generate + rollouts_stored +1;
  int rollouts_reuse =  rollouts_total < config_.max_rollouts  ? rollouts_stored :  config_.max_rollouts - (rollouts_generate + 1) ; // +1 for optimized params

  // selecting least costly rollouts from previous iteration
  if(rollouts_reuse > 0)
  {
    // find min and max cost for exponential cost scaling
    double min_cost = std::numeric_limits<double>::max();
    double max_cost = std::numeric_limits<double>::min();
    for (int r=1; r<rollouts_stored; ++r)
    {
      double c = noisy_rollouts_[r].total_cost;
      if (c < min_cost)
        min_cost = c;
      if (c > max_cost)
        max_cost = c;
    }

    double cost_denom = max_cost - min_cost;
    if (cost_denom < 1e-8)
      cost_denom = 1e-8;

    // compute weighted cost on all rollouts
    double cost_prob;
    double weighted_prob;
    for (auto r = 0u; r<rollouts_stored; ++r)
    {

      // Apply noise generated on the previous iteration onto the current trajectory
      noisy_rollouts_[r].noise = noisy_rollouts_[r].parameters_noise
          - parameters_optimized_;

      cost_prob = exp(-h*(noisy_rollouts_[r].total_cost - min_cost)/cost_denom);
      weighted_prob = cost_prob * noisy_rollouts_[r].importance_weight;
      rollout_cost_sorter.push_back(std::make_pair(-weighted_prob,r));
    }

    std::sort(rollout_cost_sorter.begin(), rollout_cost_sorter.end());

    // use the best ones: (copy them into reused_rollouts)
    for (auto r = 0u; r<rollouts_stored; ++r)
    {
      int reuse_index = rollout_cost_sorter[r].second;
      reused_rollouts_[r] = noisy_rollouts_[reuse_index];
    }

    // copy them back from reused_rollouts_ into rollouts_
    for (auto r = 0u; r<rollouts_reuse; ++r)
    {
      noisy_rollouts_[rollouts_generate + r ] = reused_rollouts_[r];
    }
  }

  // adding optimized trajectory as the last rollout
  noisy_rollouts_[rollouts_generate + rollouts_reuse].parameters_noise = parameters_optimized_;
  noisy_rollouts_[rollouts_generate + rollouts_reuse].noise.setZero();
  noisy_rollouts_[rollouts_generate + rollouts_reuse].state_costs = parameters_state_costs_;
  noisy_rollouts_[rollouts_generate + rollouts_reuse].control_costs = parameters_control_costs_;

  // generate new noisy rollouts
  for(auto r = 0u; r < rollouts_generate; r++)
  {
    if(!task_->generateNoisyParameters(parameters_optimized_,
                                      0,config_.num_timesteps,
                                      current_iteration_,r,
                                      noisy_rollouts_[r].parameters_noise,
                                      noisy_rollouts_[r].noise))
    {
      ROS_ERROR("Failed to generate noisy parameters at iteration %i",current_iteration_);
      return false;
    }

  }

  // update total active rollouts
  num_active_rollouts_ = rollouts_reuse + rollouts_generate + 1;

  return true;
}

bool Stomp::filterNoisyRollouts()
{
  MeasureTime timer("filterNoisyRollouts");
  // apply post noise generation filters
  bool filtered = false;
  for(auto r = 0u ; r < config_.num_rollouts; r++)
  {
    if(!task_->filterNoisyParameters(0,config_.num_timesteps,current_iteration_,r,noisy_rollouts_[r].parameters_noise,filtered))
    {
      ROS_ERROR_STREAM("Failed to filter noisy parameters");
      return false;
    }

    if(filtered)
    {
      noisy_rollouts_[r].noise = noisy_rollouts_[r].parameters_noise - parameters_optimized_;
    }
  }

  return true;
}

bool Stomp::computeNoisyRolloutsCosts()
{
  MeasureTime timer("computeNoisyRollouts");
  // computing state and control costs
  bool valid = computeRolloutsStateCosts() && computeRolloutsControlCosts();

  if(valid)
  {
    // compute total costs
    double total_state_cost ;
    double total_control_cost;

    for(auto r = 0u ; r < num_active_rollouts_;r++)
    {
      Rollout& rollout = noisy_rollouts_[r];
      total_state_cost = rollout.state_costs.sum();

      // Compute control + state cost for each joint
      total_control_cost = 0;
      double ccost = 0;
      for(auto d = 0u; d < config_.num_dimensions; d++)
      {
        ccost = rollout.control_costs.row(d).sum();
        total_control_cost += ccost;
        rollout.full_costs[d] = ccost + total_state_cost;
      }
      rollout.total_cost = total_state_cost + total_control_cost;

      // Compute total cost for each time step
      for(auto d = 0u; d < config_.num_dimensions; d++)
      {
        rollout.total_costs.row(d) = rollout.state_costs.transpose() + rollout.control_costs.row(d);
      }
    }
  }

  return valid;
}

bool Stomp::computeRolloutsStateCosts()
{
  MeasureTime timer("computeRolloutsStateCosts");
  bool all_valid = true;
  bool proceed = true;
  for(auto r = 0u ; r < config_.num_rollouts; r++)
  {
    if(!proceed_)
    {
      proceed = false;
      break;
    }

    Rollout& rollout = noisy_rollouts_[r];
    if(!task_->computeNoisyCosts(rollout.parameters_noise,0,
                            config_.num_timesteps,
                            current_iteration_,r,
                            rollout.state_costs,all_valid))
    {
      ROS_ERROR("Trajectory cost computation failed for rollout %i.",r);
      proceed = false;
      break;
    }
  }

  return proceed;
}
bool Stomp::computeRolloutsControlCosts()
{
  MeasureTime timer("computeRolloutsControlCosts");
  Eigen::ArrayXXd Ax; // accelerations
  for(auto r = 0u ; r < num_active_rollouts_; r++)
  {
    Rollout& rollout = noisy_rollouts_[r];

    if(config_.control_cost_weight < MIN_CONTROL_COST_WEIGHT)
    {
      for(auto d = 0u; d < config_.num_dimensions; d++)
      {
        rollout.control_costs.row(d).setConstant(0);
      }
    }
    else
    {
      computeParametersControlCosts(rollout.parameters_noise,
                                    config_.delta_t,
                                    config_.control_cost_weight,
                                    control_cost_matrix_R_,rollout.control_costs);
    }
  }
  return true;
}

bool Stomp::computeProbabilities()
{
  MeasureTime timer("computeProbabilities");
  double cost;
  double min_cost;
  double max_cost;
  double denom;
  double numerator;
  double probl_sum = 0.0; // total probability sum of all rollouts for each joint
  const double h = config_.exponentiated_cost_sensitivity;
  double exponent = 0;

  for (auto d = 0u; d<config_.num_dimensions; ++d)
  {

    for (auto t = 0u; t<config_.num_timesteps; t++)
    {

      // find min and max cost over all rollouts at timestep 't':
      min_cost = noisy_rollouts_[0].total_costs(d,t);
      max_cost = min_cost;
      for (auto r=0u; r<num_active_rollouts_; ++r)
      {
          cost = noisy_rollouts_[r].total_costs(d,t);
          if (cost < min_cost)
              min_cost = cost;
          if (cost > max_cost)
              max_cost = cost;
      }

      denom = max_cost - min_cost;

      // prevent division by zero:
      if (denom < MIN_COST_DIFFERENCE)
      {
        denom = MIN_COST_DIFFERENCE;
      }

      probl_sum = 0.0;
      for (auto r = 0u; r<num_active_rollouts_; ++r)
      {
        // this is the exponential term in the probability calculation described in the literature
        exponent = -h*(noisy_rollouts_[r].total_costs(d,t) - min_cost)/denom;
        noisy_rollouts_[r].probabilities(d,t) = noisy_rollouts_[r].importance_weight *
            exp(exponent);

        probl_sum += noisy_rollouts_[r].probabilities(d,t);
      }

      // scaling each probability value by the sum of all probabilities corresponding to all rollouts at time "t"
      for (auto r = 0u; r<num_active_rollouts_; ++r)
      {
        noisy_rollouts_[r].probabilities(d,t) /= probl_sum;
      }
    }

    // computing full probabilities
    min_cost = noisy_rollouts_[0].full_costs[d];
    max_cost = min_cost;
    double c = 0.0;
    for (int r=1; r<num_active_rollouts_; ++r)
    {
      c = noisy_rollouts_[r].full_costs[d];
      if (c < min_cost)
        min_cost = c;
      if (c > max_cost)
        max_cost = c;
    }

    denom = max_cost - min_cost;
    denom = denom < MIN_COST_DIFFERENCE ? MIN_COST_DIFFERENCE : denom;

    probl_sum = 0.0;
    for (int r=0; r<num_active_rollouts_; ++r)
    {
      noisy_rollouts_[r].full_probabilities[d] = noisy_rollouts_[r].importance_weight *
          exp(-h*(noisy_rollouts_[r].full_costs[d] - min_cost)/denom);
      probl_sum += noisy_rollouts_[r].full_probabilities[d];
    }
    for (int r=0; r<num_active_rollouts_; ++r)
    {
        noisy_rollouts_[r].full_probabilities[d] /= probl_sum;
    }
  }

  return true;
}

bool Stomp::updateParameters()
{
  MeasureTime timer("updateParameters");
  // computing updates from probabilities using convex combination
  parameters_updates_.setZero();
  for(auto d = 0u; d < config_.num_dimensions ; d++)
  {

    for(auto r = 0u; r < num_active_rollouts_; r++)
    {
      auto& rollout = noisy_rollouts_[r];
      parameters_updates_.row(d) +=  (rollout.noise.row(d).array() * rollout.probabilities.row(d).array()).matrix();
    }

  }

  // filtering updates
  if(!task_->filterParameterUpdates(0,config_.num_timesteps,current_iteration_,parameters_optimized_,parameters_updates_))
  {
    ROS_ERROR("Updates filtering step failed");
    return false;
  }

  // updating parameters
  parameters_optimized_ += parameters_updates_;

  return true;
}

bool Stomp::computeOptimizedCost()
{
  MeasureTime timer("computeOptimizedCosts");
  // control costs
  parameters_total_cost_ = 0;
  if(config_.control_cost_weight > MIN_CONTROL_COST_WEIGHT)
  {
    computeParametersControlCosts(parameters_optimized_,
                                  config_.delta_t,
                                  config_.control_cost_weight,
                                  control_cost_matrix_R_,
                                  parameters_control_costs_);

    // adding all costs
    parameters_total_cost_ = parameters_control_costs_.rowwise().sum().sum();

  }

  // state costs
  if(task_->computeCosts(parameters_optimized_,
                         0,config_.num_timesteps,current_iteration_,
                         parameters_state_costs_,parameters_valid_))
  {
    parameters_total_cost_ += parameters_state_costs_.sum();
  }
  else
  {
    return false;
  }

  if(current_lowest_cost_ > parameters_total_cost_)
  {
    current_lowest_cost_ = parameters_total_cost_;
  }
  else
  {
    // reverting updates as no improvement was made
    parameters_optimized_ -= parameters_updates_;
  }

  return true;
}

} /* namespace stomp */


/**
 * @file bcsr_matrix_tools.cpp
 * @brief
 * @date 2009-11-24
 */
#ifdef BSPY_EXPORTING_PLUGIN
#include <boost/python.hpp>
#endif
#include "bcsr_matrix_tools.h"

#include <time.h>
#include <stdlib.h>
#include <map>
#include <boost/format.hpp>

//#include "strategies.h"

using namespace std;
using namespace boost::python;

namespace blue_sky
{
  bcsr_matrix_tools::bcsr_matrix_tools (bs_type_ctor_param)
        : bcsr_matrix_tools_iface ()
    {
    }
  bcsr_matrix_tools::bcsr_matrix_tools (const bcsr_matrix_tools & /*src*/) : bs_refcounter ()
     {
     }

/*!
  \brief read matrix from ascii file
  \param matrix         -- reference to the BCSR matrix interface
  \param file_name -- name of the file
  \return 0 if success
*/
int
bcsr_matrix_tools::ascii_read_from_csr_format (sp_bcsr_t matrix, const std::string &file_name) const
{
  //locale_keeper lkeeper ("C", LC_NUMERIC);

  FILE *fp = 0;
  char buf[4096];
  int state = 0;
  t_long nc, nr, nnz = 0, nbs, b_sqr = 0;
  t_long row_ind = 0;
  t_long j_ind = 0, j = 0;
  t_long *rows_ptr;
  t_long *cols_ind;
  t_float *values;
  char *start_ptr, *end_ptr;

  fp = fopen (file_name.c_str (), "r");
  //BS_ASSERT (fp) (file_name);
  if (!fp)
    {
      //bs_throw_exception ("Can't open file");
      return -1;
    }

  while (fgets(buf, 4096, fp))
    {
      // skip comments
      if (!strncmp (buf, "//", 2))
        continue;
      if (state == 0) // read n_rows, n_cols, n_non_zeros
        {
          start_ptr = buf;
          nr = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;
          start_ptr = end_ptr;
          nc = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;
          start_ptr = end_ptr;
          nnz = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;
          start_ptr = end_ptr;
          nbs = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;

          if (matrix->init (nr, nc, nbs, nnz))
            {
              //bs_throw_exception ("Can't init matrix");
              return -45;
            }

          //BOSOUT << "nr = " << nr << ", nc = " << nc << ", nbs = " << nbs << ", nnz = " << nnz << bs_end;

          state = 1;

          rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
          cols_ind = &(*(matrix->get_cols_ind ()))[0];
          values = &(*(matrix->get_values ()))[0];

          rows_ptr[0] = 0;
          b_sqr = nbs * nbs;
        }
      else if (state == 1) // read rows_ptr
        {
          t_long n_rows = matrix->get_n_rows ();

          start_ptr = buf;
          rows_ptr[row_ind + 1] = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;
          ++row_ind;
          if (row_ind == n_rows)
            state = 2;
        }
      else if (state == 2) // read values
        {
          start_ptr = buf;
          cols_ind[j_ind] = strtol (start_ptr, &end_ptr, 10);
          if (start_ptr == end_ptr)
            continue;
          for (j = 0; j < b_sqr; ++j)
            {
              start_ptr = end_ptr;
              values[j_ind * b_sqr + j] = strtod (start_ptr, &end_ptr);
              if (start_ptr == end_ptr)
                continue;
            }
          j_ind++;
          if (j_ind == nnz)
            break;
        }
    }

  //BOSOUT << "nnz = " << nnz << bs_end;

  fclose (fp);
  return 0;
}

/*!
  \brief write matrix to ascii file
  \param matrix         -- reference to the BCSR matrix interface
  \param file_name -- name of the file
  \return 0 if success
*/
int
bcsr_matrix_tools::ascii_write_to_csr_format (const sp_bcsr_t matrix,
                                              const std::string &file_name,
                                              const bool sort_cols) const
{
  FILE *fp = 0;
  t_long i, j, j1, j2, jj, jj1, jj2, counter, b_sqr = 0;
  t_long n_row_cols = 0;
  spv_long sp_sort_index = BS_KERNEL.create_object (v_long::bs_type ());
  BS_ASSERT (sp_sort_index);
  t_long *sort_index = 0;

  t_long *rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
  t_long *cols_ind = &(*(matrix->get_cols_ind ()))[0];
  t_float *values = &(*(matrix->get_values ()))[0];
  t_long n_rows = matrix->get_n_rows ();
  t_long n_cols = matrix->get_n_cols ();
  t_long n_block_size = matrix->get_n_block_size ();
  t_long n_nnz = matrix->get_n_non_zeros ();
  b_sqr = n_block_size * n_block_size;

  fp = fopen (file_name.c_str (), "w");
  if (!fp)
    //TODO: write error message
    return -1;
  fprintf (fp, "// N_ROWS\tN_COLS\tN_NON_ZEROS\tN_BLOCK_SIZE\n");
  fprintf (fp,
      (boost::format("%ld\t%ld\t%ld\t%ld\n") %
       n_rows % n_cols % n_nnz % n_block_size).str().c_str());

  fprintf (fp, "// Rows indexes[1..n_rows] (with out 0)\n");
  for (i = 1; i <= n_rows; ++i)
    fprintf (fp, (boost::format("%ld\n") % rows_ptr[i]).str().c_str());

  fprintf (fp, "// END of Rows indexes\n");

  fprintf (fp, "// Values n_non_zeros elements\n");
  fprintf (fp, "//COLUMN\tVALUE\n");
  for (i = 0, counter = 0; i < n_rows; ++i)
    {
      fprintf (fp, (boost::format("// ROW %ld\n") % i).str().c_str());
      j1 = rows_ptr[i];
      j2 = rows_ptr[i + 1];

      if (sort_cols)
        {
          sp_sort_index->resize (j2 - j1);
          sort_index = &(*(sp_sort_index))[0];
          for (j = j1, jj = 0; j < j2; ++j, ++jj)
            {
              sort_index[jj] = j;
            }

          n_row_cols = j2 - j1;
          for (j = 0; j < n_row_cols - 1; ++j)
            {
              for (jj = j; jj < n_row_cols; ++jj)
                {
                  if (cols_ind[sort_index[j]] > cols_ind[sort_index[jj]])
                    {
                      jj1 = sort_index[j];
                      sort_index[j] = sort_index[jj];
                      sort_index[jj] = jj1;
                    }
                }
            }
          for (j = 0; j < n_row_cols; ++j, ++counter)
            {
              fprintf (fp, (boost::format("%ld") % cols_ind[sort_index[j]]).str().c_str());
              jj1 = sort_index[j] * b_sqr;
              jj2 = jj1 + b_sqr;
                for (jj = jj1; jj < jj2; ++jj)
                  fprintf (fp, "\t%le", values[jj]);
              fprintf (fp, "\n");
            }
        }
      else
        {
          for (j = j1; j < j2; ++j, ++counter)
            {
              fprintf (fp, (boost::format("%ld") % cols_ind[j]).str().c_str());
              jj1 = j * b_sqr;
              jj2 = jj1 + b_sqr;
                for (jj = jj1; jj < jj2; ++jj)
                  fprintf (fp, "\t%le", values[jj]);
              fprintf (fp, "\n");
            }
        }
    }
  if (counter != n_nnz)
    return -1;
  fprintf (fp, "// END OF FILE\n");
  fclose (fp);

  return 0;
}

int
bcsr_matrix_tools::random_init (sp_bcsr_t matrix,
                                const t_long new_n_rows,
                                const t_long new_n_block_size,
                                const t_double rand_value_dispersion,
                                const t_long elems_in_row
                               ) const
{
  int r_code = 0;
  t_long i, j, j1, j2, b_sqr;
  t_long n_non_zeros;

  // check input data
  if (elems_in_row <= 0)
    // internal error
    return -1;

  // setup integer vars
  n_non_zeros = new_n_rows * elems_in_row; // approx value, final nonzeros value can be equal or less
  b_sqr = new_n_block_size * new_n_block_size;

  r_code = matrix->init (new_n_rows, new_n_rows, new_n_block_size, n_non_zeros);

  if (r_code)
    //TODO: print error message
    return -2;

  t_float *values = &(*(matrix->get_values ()))[0];
  t_long *rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
  t_long *cols_ind = &(*(matrix->get_cols_ind ()))[0];

  srand ((unsigned)time( NULL ));

  for (i = 0; i < n_non_zeros * b_sqr; ++i)
    values[i] = (t_float)rand () / (t_float)RAND_MAX * (t_float)rand_value_dispersion;

  rows_ptr[0] = 0;
  for (i = 0; i < new_n_rows; ++i)
    {
      rows_ptr[i + 1] = rows_ptr[i] +  elems_in_row;
      j1 = rows_ptr[i];
      j2 = rows_ptr[i + 1];
      cols_ind[j1++] = i; // diagonal should be the first

      for (j = j1; j < j2; ++j)
        {

          int cl = rand () % (int)((double)new_n_rows / (double)(j2 - j1)) + (j - j1) * new_n_rows / (j2 - j1);
          if (cl == i)
            {
              --j;
              continue;
            }
          cols_ind[j] = cl;
        }
    }

  //n_non_zeros = rows_ptr[new_n_rows];
  return 0;
}

int
bcsr_matrix_tools::gen_2d_laplas (sp_bcsr_t matrix, const t_long n) const
{
  int r_code = 0;
  t_long i, j;

  t_long elems_in_row = 5;
  t_long new_n_block_size = 1;

  // setup integer vars
  t_long N = n * n;
  t_long n_non_zeros = N * elems_in_row; // approx value, final nonzeros value can be equal or less

  r_code = matrix->init (N, N, new_n_block_size, n_non_zeros);

  if (r_code)
    //TODO: print error message
    return -2;

  t_float *values = &(*(matrix->get_values ()))[0];
  t_long *rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
  t_long *cols_ind = &(*(matrix->get_cols_ind ()))[0];

  rows_ptr[0] = 0;
  for (i = 0, j = 0; i < N; ++i)
    {
      values[j] = 4;
      cols_ind[j] = i;
      ++j;

      if (i - n >= 0)
        {
          values[j] = -1;
          cols_ind[j] = i - n;
          ++j;
        }

      if (i % n)
        {
          values[j] = -1;
          cols_ind[j] = i - 1;
          ++j;
        }

      if ((i + 1) % n)
        {
          values[j] = -1;
          cols_ind[j] = i + 1;
          ++j;
        }

      if (i + n < N)
        {
          values[j] = -1;
          cols_ind[j] = i + n;
          ++j;
        }

      rows_ptr[i + 1] = j;
    }

  //n_non_zeros = rows_ptr[new_n_rows];
  return 0;
}

int
bcsr_matrix_tools::gen_diagonal (sp_bcsr_t matrix, const t_long n,
                                 const t_long nb, const t_double val) const
{
  int r_code = 0;
  t_long i, k;
  t_long b_sqr = nb * nb;

  r_code = matrix->init (n, n, nb, n * b_sqr);
  if (r_code)
    //TODO: print error message
    return -2;

  t_float *values = &(*(matrix->get_values ()))[0];
  t_long *rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
  t_long *cols_ind = &(*(matrix->get_cols_ind ()))[0];

  rows_ptr[0] = 0;
  for (i = 0; i < n; ++i)
    {
      for (k = 0; k < b_sqr; ++k)
        {
          if (k % (nb + 1) == 0)
            values[i * b_sqr + k] = val;
          else
            values[i * b_sqr + k] = 0.0;
        }
      cols_ind[i] = i;
      rows_ptr[i + 1] = i + 1;
    }

  return 0;
}

int
bcsr_matrix_tools::dense_init (sp_bcsr_t matrix,
                               const t_long n_rows,
                               const t_long block_size,
                               const t_double rand_value_dispersion) const
{
  int r_code = 0;
  t_long b_sqr = block_size * block_size;

  srand ((unsigned)time( NULL ));

  r_code = matrix->init (n_rows, n_rows, block_size, n_rows * n_rows);
  if (r_code)
    return -2;

  t_float *values = &(*(matrix->get_values ()))[0];
  t_long *rows_ptr = &(*(matrix->get_rows_ptr ()))[0];
  t_long *cols_ind = &(*(matrix->get_cols_ind ()))[0];
  rows_ptr[0] = 0;

  for (t_long i = 0; i < n_rows * n_rows * b_sqr; ++i)
    {
      values[i] = (t_float)rand () / (t_float)RAND_MAX * (t_float)rand_value_dispersion;
    }
  for (t_long i = 0; i < n_rows; ++i)
    {
      t_long cl;
      t_long j1 = rows_ptr[i];

      rows_ptr[i + 1] = j1 + n_rows;

      // diagonal should be the first
      cols_ind[j1] = i;
      cl = j1 + 1;
      for (t_long j = 0; j < n_rows; ++j)
        {
          if (j != i)
            {
              cols_ind[cl] = j;
              ++cl;
            }
        }
    }
  return 0;
}
/////////////////////////////////BS Register
/////////////////////////////////Stuff//////////////////////////

  BLUE_SKY_TYPE_STD_CREATE (bcsr_matrix_tools);
  BLUE_SKY_TYPE_STD_COPY (bcsr_matrix_tools);

  BLUE_SKY_TYPE_IMPL(bcsr_matrix_tools,  bcsr_matrix_tools_iface, "bcsr_matrix_tools", "Tools for Block CSR Matrix class", "Tools realization of Block CSR Matricies");
}  // blue_sky namespace


// This file is distributed under the MIT license.
// See the LICENSE file for details.

#include <common/config.h>

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstring>
#include <fstream>
#include <iostream>
#include <memory>
#include <ostream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include <boost/iostreams/device/mapped_file.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/utility/string_ref.hpp>
#include <boost/assign.hpp>
#include <boost/bimap.hpp>
#include <boost/filesystem.hpp>

#include <rapidjson/document.h>
#include <rapidjson/filereadstream.h>
#include <rapidjson/filewritestream.h>
#include <rapidjson/prettywriter.h>

#include <visionaray/math/constants.h>
#include <visionaray/math/forward.h>
#include <visionaray/math/unorm.h>
#include <visionaray/math/vector.h>
#include <visionaray/texture/texture.h>

#include "cfile.h"
#include "model.h"
#include "sg.h"
#include "vsnray_loader.h"

using namespace visionaray;

namespace data_file
{

//-------------------------------------------------------------------------------------------------
// (included) data file meta data
//

struct meta_data
{
    enum encoding_t
    {
        Ascii,
        Binary
    };

    // VecN are binary compatible w/ visionaray::vecN
    enum data_type_t
    {
        U8,
        Float,
        Vec2u8,
        Vec2f,
        Vec3u8,
        Vec3f,
        Vec4u8,
        Vec4f,
    };

    static boost::bimap<data_type_t, std::string> data_type_map;

    enum compression_t
    {
        Raw
    };

    std::string   path;
    encoding_t    encoding    = Binary;
    data_type_t   data_type   = U8;
    int           num_items   = 0;
    compression_t compression = Raw;
    char          separator   = ' ';
};

boost::bimap<meta_data::data_type_t, std::string> meta_data::data_type_map
    = boost::assign::list_of<typename boost::bimap<meta_data::data_type_t, std::string>::relation>
        ( U8,     "u8" )
        ( Float,  "float" )
        ( Vec2u8, "vec2u8" )
        ( Vec2f,  "vec2f" )
        ( Vec3u8, "vec3u8" )
        ( Vec3f,  "vec3f" )
        ( Vec4u8, "vec4u8" )
        ( Vec4f,  "vec4f" );

} // data_file

//-------------------------------------------------------------------------------------------------
// Floating point number parser
//

template <typename It, typename Vector>
bool parse_floats(It first, It last, Vector& vec, char separator = ' ')
{
    namespace qi = boost::spirit::qi;

    return qi::phrase_parse(
            first,
            last,
            qi::float_ % *qi::char_(separator),
            qi::ascii::space,
            vec
            );
}

template <size_t N, typename Container>
bool parse_as_vecN(data_file::meta_data md, Container& vecNs)
{
    boost::iostreams::mapped_file_source file(md.path);

    if (md.data_type == data_file::meta_data::Float)
    {
        if (md.num_items % N != 0)
        {
            return false;
        }

        std::vector<float> floats;

        if (md.encoding == data_file::meta_data::Ascii)
        {
            boost::string_ref text(file.data(), file.size());

            parse_floats(text.cbegin(), text.cend(), floats, md.separator);

            if (static_cast<int>(floats.size()) != md.num_items)
            {
                return false;
            }
        }
        else // Binary
        {
            floats.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(floats.data())
                );
        }

        vecNs.resize(md.num_items / N);
        for (size_t i = 0; i < vecNs.size(); ++i)
        {
            for (size_t j = 0; j < N; ++j)
            {
                vecNs[i][j] = floats[i * N + j];
            }
        }
    }
    else if (md.data_type == data_file::meta_data::Vec2u8)
    {
        if (N != 2)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }
    else if (md.data_type == data_file::meta_data::Vec2f)
    {
        if (N != 2)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }
    else if (md.data_type == data_file::meta_data::Vec3u8)
    {
        if (N != 3)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }
    else if (md.data_type == data_file::meta_data::Vec3f)
    {
        if (N != 3)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }
    else if (md.data_type == data_file::meta_data::Vec4u8)
    {
        if (N != 4)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }
    else if (md.data_type == data_file::meta_data::Vec4f)
    {
        if (N != 4)
        {
            throw std::runtime_error("");
        }

        if (md.encoding == data_file::meta_data::Ascii)
        {
            // Not implemented yet
            return false;
        }
        else // Binary
        {
            vecNs.resize(md.num_items);
            std::copy(
                file.data(),
                file.data() + file.size(),
                reinterpret_cast<char*>(vecNs.data())
                );
        }
    }

    return true;
}

template <typename Container>
bool parse_as_vec2f(data_file::meta_data md, Container& vec2fs)
{
    return parse_as_vecN<2>(md, vec2fs);
}

template <typename Container>
bool parse_as_vec3f(data_file::meta_data md, Container& vec3fs)
{
    return parse_as_vecN<3>(md, vec3fs);
}

//-------------------------------------------------------------------------------------------------
// Parse tiny objects from json object
//

template <typename Object>
void parse_optional(float& f, Object const& obj, char const* member)
{
    if (obj.HasMember(member))
    {
        auto const& val = obj[member];

        if (!val.IsFloat())
        {
            throw std::runtime_error("");
        }

        f = val.GetFloat();
    }
}

template <size_t N, typename Object>
void parse_optional(vector<N, float>& v, Object const& obj, char const* member)
{
    if (obj.HasMember(member))
    {
        auto const& val = obj[member];

        if (!val.IsArray())
        {
            throw std::runtime_error("");
        }

        if (val.Capacity() != N)
        {
            throw std::runtime_error("");
        }

        for (size_t i = 0; i < N; ++i)
        {
            v[i] = val[i].GetFloat();
        }
    }
}

template <typename Object>
void parse_optional(recti& r, Object const& obj, char const* member)
{
    if (obj.HasMember(member))
    {
        auto const& val = obj[member];

        if (!val.IsArray())
        {
            throw std::runtime_error("");
        }

        if (val.Capacity() != 4)
        {
            throw std::runtime_error("");
        }

        for (size_t i = 0; i < 4; ++i)
        {
            r.data()[i] = val[i].GetFloat();
        }
    }
}

//-------------------------------------------------------------------------------------------------
// .vsnray parser
//

class vsnray_parser
{
public:

    vsnray_parser(std::string filename)
        : filename_(filename)
    {
    }

    void parse_children(std::shared_ptr<sg::node> parent, rapidjson::Value const& entries);

    template <typename Object>
    std::shared_ptr<sg::node> parse_node(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_camera(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_include(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_point_light(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_spot_light(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_reference(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_transform(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_surface_properties(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_triangle_mesh(Object const& obj);

    template <typename Object>
    std::shared_ptr<sg::node> parse_indexed_triangle_mesh(Object const& obj);

    template <typename Object>
    data_file::meta_data parse_file_meta_data(Object const& obj);

private:

    std::string filename_;

};

//-------------------------------------------------------------------------------------------------
// Parse nodes
//

void vsnray_parser::parse_children(std::shared_ptr<sg::node> parent, rapidjson::Value const& entries)
{
    if (!entries.IsArray())
    {
        throw std::runtime_error("");
    }

    parent->children().resize(entries.Capacity());

    size_t i = 0;
    for (auto const& c : entries.GetArray())
    {
        auto const& obj = c.GetObject();

        parent->children().at(i++) = parse_node(obj);
    }

    if (i != entries.Capacity())
    {
        throw std::runtime_error("");
    }
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_node(Object const& obj)
{
    std::shared_ptr<sg::node> result = nullptr;

    if (obj.HasMember("type"))
    {
        // Parse individual node types
        auto const& type_string = obj["type"];
        if (strncmp(type_string.GetString(), "node", 4) == 0)
        {
            // Empty node, (may still contain children, e.g. root)
            result = std::make_shared<sg::node>();
        }
        else if (strncmp(type_string.GetString(), "camera", 6) == 0)
        {
            result = parse_camera(obj);
        }
        else if (strncmp(type_string.GetString(), "include", 6) == 0)
        {
            result = parse_include(obj);
        }
        else if (strncmp(type_string.GetString(), "point_light", 11) == 0)
        {
            result = parse_point_light(obj);
        }
        else if (strncmp(type_string.GetString(), "spot_light", 10) == 0)
        {
            result = parse_spot_light(obj);
        }
        else if (strncmp(type_string.GetString(), "reference", 9) == 0)
        {
            result = parse_reference(obj);
        }
        else if (strncmp(type_string.GetString(), "transform", 9) == 0)
        {
            result = parse_transform(obj);
        }
        else if (strncmp(type_string.GetString(), "surface_properties", 18) == 0)
        {
            result = parse_surface_properties(obj);
        }
        else if (strncmp(type_string.GetString(), "triangle_mesh", 13) == 0)
        {
            result = parse_triangle_mesh(obj);
        }
        else if (strncmp(type_string.GetString(), "indexed_triangle_mesh", 21) == 0)
        {
            result = parse_indexed_triangle_mesh(obj);
        }
        else
        {
            throw std::runtime_error("");
        }

        // Parse common node properties
        if (obj.HasMember("name"))
        {
            assert(result != nullptr);

            rapidjson::Value const& name = obj["name"];
            result->name() = name.GetString();
        }

        if (obj.HasMember("children"))
        {
            assert(result != nullptr);

            rapidjson::Value const& children = obj["children"];
            parse_children(result, children);
        }
    }
    else
    {
        throw std::runtime_error("");
    }

    return result;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_camera(Object const& obj)
{
    auto cam = std::make_shared<sg::camera>();

    vec3 eye(0.0f);
    parse_optional(eye, obj, "eye");

    vec3 center(0.0f);
    parse_optional(center, obj, "center");

    vec3 up(0.0f);
    parse_optional(up, obj, "up");

    float fovy = 45.0f;
    parse_optional(fovy, obj, "fovy");

    float znear = 0.001f;
    parse_optional(znear, obj, "znear");

    float zfar = 1000.0f;
    parse_optional(zfar, obj, "zfar");

    recti viewport(0, 0, 0, 0);
    parse_optional(viewport, obj, "viewport");

    float lens_radius = 0.1f;
    parse_optional(lens_radius, obj, "lens_radius");

    float focal_distance = 10.0f;
    parse_optional(focal_distance, obj, "focal_distance");

    float aspect = viewport.w > 0 && viewport.h > 0
                 ? viewport.w / static_cast<float>(viewport.h)
                 : 1;

    cam->perspective(fovy * constants::degrees_to_radians<float>(), aspect, znear, zfar);
    if (viewport.w > 0 && viewport.h > 0)
    {
        cam->set_viewport(viewport);
    }
    cam->set_lens_radius(lens_radius);
    cam->set_focal_distance(focal_distance);
    cam->look_at(eye, center, up);

    return cam;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_include(Object const& obj)
{
    auto inc = std::make_shared<sg::node>();

    if (obj.HasMember("path"))
    {
        std::string path_string(obj["path"].GetString());

        boost::filesystem::path p(path_string);

        if (!p.is_absolute())
        {
            // Extract base path
            boost::filesystem::path bp(filename_);
            bp = bp.parent_path();

            // Append path to base path
            p = bp / p;

            path_string = p.string();
        }

        model mod;
        if (mod.load(path_string))
        {
            if (mod.scene_graph == nullptr)
            {
                std::unordered_map<std::string, std::shared_ptr<sg::texture2d<vector<4, unorm<8>>>>> texture_map;

                for (auto it = mod.texture_map.begin(); it != mod.texture_map.end(); ++it)
                {
                    auto tex = std::make_shared<sg::texture2d<vector<4, unorm<8>>>>();
                    tex->name() = it->first;
                    tex->resize(it->second.width(), it->second.height());
                    tex->reset(it->second.data());
                    tex->set_filter_mode(it->second.get_filter_mode());
                    tex->set_address_mode(it->second.get_address_mode());

                    texture_map.insert(std::make_pair(it->first, tex));
                }

                if (mod.primitives.size() > 0)
                {
                    // Vertices (disassemble triangles..)
                    for (auto tri : mod.primitives)
                    {
                        if (tri.geom_id >= inc->children().size())
                        {
                            auto props = std::make_shared<sg::surface_properties>();

                            // Add material
                            auto obj = std::make_shared<sg::obj_material>();
                            obj->ca = mod.materials[tri.geom_id].ca;
                            obj->cd = mod.materials[tri.geom_id].cd;
                            obj->cs = mod.materials[tri.geom_id].cs;
                            obj->ce = mod.materials[tri.geom_id].ce;
                            obj->cr = mod.materials[tri.geom_id].cr;
                            obj->ior = mod.materials[tri.geom_id].ior;
                            obj->absorption = mod.materials[tri.geom_id].absorption;
                            obj->transmission = mod.materials[tri.geom_id].transmission;
                            obj->specular_exp = mod.materials[tri.geom_id].specular_exp;
                            obj->illum = mod.materials[tri.geom_id].illum;
                            props->material() = obj;

                            bool insert_dummy = false;

                            if (tri.geom_id < mod.textures.size())
                            {
                                // Find texture in texture_map
                                bool found = false;
                                for (auto it = mod.texture_map.begin(); it != mod.texture_map.end(); ++it)
                                {
                                    auto ref = texture_ref<vector<4, unorm<8>>, 2>(it->second);

                                    if (ref.data() == mod.textures[tri.geom_id].data())
                                    {
                                        std::string name = it->first;
                                        // Find in local texture map
                                        auto res = texture_map.find(name);
                                        if (res != texture_map.end())
                                        {
                                            props->add_texture(res->second);
                                            found = true;
                                            break;
                                        }
                                    }
                                }

                                if (!found)
                                {
                                    insert_dummy = true;
                                }
                            }
                            else
                            {
                                insert_dummy = true;
                            }

                            if (insert_dummy)
                            {
                                // Add a dummy texture
                                vector<4, unorm<8>> dummy_texel(1.0f, 1.0f, 1.0f, 1.0f);
                                auto tex = std::make_shared<sg::texture2d<vector<4, unorm<8>>>>();
                                tex->resize(1, 1);
                                tex->set_address_mode(Wrap);
                                tex->set_filter_mode(Nearest);
                                tex->reset(&dummy_texel);
                                props->add_texture(tex);
                            }

                            // Add to scene graph
                            props->add_child(std::make_shared<sg::triangle_mesh>());
                            inc->add_child(props);
                        }

                        auto mesh = std::dynamic_pointer_cast<sg::triangle_mesh>(
                                inc->children()[tri.geom_id]->children()[0]
                                );

                        vec3 verts[3] = {
                            tri.v1,
                            tri.v1 + tri.e1,
                            tri.v1 + tri.e2
                            };
                        mesh->vertices.insert(mesh->vertices.end(), verts, verts + 3);

                        if (mod.shading_normals.size() > tri.prim_id * 3 + 3)
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->normals.push_back(mod.shading_normals[tri.prim_id * 3 + i]);
                            }
                        }
                        else
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->normals.push_back(normalize(cross(tri.e1, tri.e2)));
                            }
                        }

                        if (mod.tex_coords.size() >= tri.prim_id * 3 + 3)
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->tex_coords.push_back(mod.tex_coords[tri.prim_id * 3 + i]);
                            }
                        }
                        else
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->tex_coords.push_back(vec2(0.0f, 0.0f));
                            }
                        }

                        if (mod.colors.size() >= tri.prim_id * 3 + 3)
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->colors.push_back(vector<3, unorm<8>>(mod.colors[tri.prim_id * 3 + i]));
                            }
                        }
                        else
                        {
                            for (int i = 0; i < 3; ++i)
                            {
                                mesh->colors.push_back(vector<3, unorm<8>>(1.0f, 1.0f, 1.0f));
                            }
                        }
                    }
                }
                else
                {
                    throw std::runtime_error("");
                }
            }
            else
            {
                // TODO: don't allow circular references..
                inc = mod.scene_graph;
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        throw std::runtime_error("");
    }

    return inc;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_point_light(Object const& obj)
{
    auto light = std::make_shared<sg::point_light>();

    vec3 cl(1.0f);
    parse_optional(cl, obj, "cl");

    float kl = 1.0f;
    parse_optional(kl, obj, "kl");

    vec3 position(0.0f);
    parse_optional(position, obj, "position");

    float constant_attenuation = 1.0f;
    parse_optional(constant_attenuation, obj, "constant_attenuation");

    float linear_attenuation = 0.0f;
    parse_optional(linear_attenuation, obj, "linear_attenuation");

    float quadratic_attenuation = 0.0f;
    parse_optional(quadratic_attenuation, obj, "quadratic_attenuation");

    light->set_cl(cl);
    light->set_kl(kl);
    light->set_position(position);
    light->set_constant_attenuation(constant_attenuation);
    light->set_linear_attenuation(linear_attenuation);
    light->set_quadratic_attenuation(quadratic_attenuation);

    return light;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_spot_light(Object const& obj)
{
    auto light = std::make_shared<sg::spot_light>();

    vec3 cl(1.0f);
    parse_optional(cl, obj, "cl");

    float kl = 1.0f;
    parse_optional(kl, obj, "kl");

    vec3 position(0.0f);
    parse_optional(position, obj, "position");

    vec3 spot_direction(0.0f, 0.0f, -1.0f);
    parse_optional(spot_direction, obj, "spot_direction");
    assert(length(spot_direction) == 1.0f);

    float spot_cutoff = 180.0f * constants::degrees_to_radians<float>();
    parse_optional(spot_cutoff, obj, "spot_cutoff");

    float spot_exponent = 0.0f;
    parse_optional(spot_exponent, obj, "spot_exponent");

    float constant_attenuation = 1.0f;
    parse_optional(constant_attenuation, obj, "constant_attenuation");

    float linear_attenuation = 0.0f;
    parse_optional(linear_attenuation, obj, "linear_attenuation");

    float quadratic_attenuation = 0.0f;
    parse_optional(quadratic_attenuation, obj, "quadratic_attenuation");

    light->set_cl(cl);
    light->set_kl(kl);
    light->set_position(position);
    light->set_spot_direction(spot_direction);
    light->set_spot_cutoff(spot_cutoff);
    light->set_spot_exponent(spot_exponent);
    light->set_constant_attenuation(constant_attenuation);
    light->set_linear_attenuation(linear_attenuation);
    light->set_quadratic_attenuation(quadratic_attenuation);

    return light;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_reference(Object const& obj)
{
    return std::make_shared<sg::node>();
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_transform(Object const& obj)
{
    auto transform = std::make_shared<sg::transform>();

    if (obj.HasMember("matrix"))
    {
        auto const& mat = obj["matrix"];

        if (mat.Capacity() != 16)
        {
            throw std::runtime_error("");
        }

        for (rapidjson::SizeType i = 0; i < mat.Capacity(); ++i)
        {
            transform->matrix().data()[i] = mat[i].GetFloat();
        }
    }

    return transform;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_surface_properties(Object const& obj)
{
    auto props = std::make_shared<sg::surface_properties>();

    if (obj.HasMember("material"))
    {
        auto const& mat = obj["material"];

        if (mat.HasMember("type"))
        {
            auto const& type_string = mat["type"];
            if (strncmp(type_string.GetString(), "obj", 3) == 0)
            {
                auto obj = std::make_shared<sg::obj_material>();

                parse_optional(obj->ca, mat, "ca");
                parse_optional(obj->cd, mat, "cd");
                parse_optional(obj->cs, mat, "cs");
                parse_optional(obj->ce, mat, "ce");

                props->material() = obj;
            }
            else if (strncmp(type_string.GetString(), "glass", 5) == 0)
            {
                auto glass = std::make_shared<sg::glass_material>();

                parse_optional(glass->ct, mat, "ct");
                parse_optional(glass->cr, mat, "cr");
                parse_optional(glass->ior, mat, "ior");

                props->material() = glass;
            }
            else if (strncmp(type_string.GetString(), "disney", 6) == 0)
            {
                auto disney = std::make_shared<sg::disney_material>();

                parse_optional(disney->base_color, mat, "base_color");
                parse_optional(disney->spec_trans, mat, "spec_trans");
                parse_optional(disney->sheen, mat, "sheen");
                parse_optional(disney->sheen_tint, mat, "sheen_tint");
                parse_optional(disney->ior, mat, "ior");
                parse_optional(disney->refractive, mat, "refractive");
                parse_optional(disney->roughness, mat, "roughness");

                props->material() = disney;
            }
            else
            {
                throw std::runtime_error("");
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        // Set default material (wavefront obj)
        auto obj = std::make_shared<sg::obj_material>();
        props->material() = obj;
    }

    if (obj.HasMember("diffuse"))
    {
        // TODO: load from file
#if 1
        vector<4, unorm<8>> dummy_texel(1.0f, 1.0f, 1.0f, 1.0f);
        auto tex = std::make_shared<sg::texture2d<vector<4, unorm<8>>>>();
        tex->resize(1, 1);
        tex->set_address_mode(Wrap);
        tex->set_filter_mode(Nearest);
        tex->reset(&dummy_texel);

        props->add_texture(tex);
#endif
    }
    else
    {
        // Set a dummy texture
        vector<4, unorm<8>> dummy_texel(1.0f, 1.0f, 1.0f, 1.0f);
        auto tex = std::make_shared<sg::texture2d<vector<4, unorm<8>>>>();
        tex->resize(1, 1);
        tex->set_address_mode(Wrap);
        tex->set_filter_mode(Nearest);
        tex->reset(&dummy_texel);

        props->add_texture(tex);
    }

    return props;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_triangle_mesh(Object const& obj)
{
    auto mesh = std::make_shared<sg::triangle_mesh>();

    if (obj.HasMember("vertices"))
    {
        auto const& verts = obj["vertices"];

        if (verts.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < verts.Capacity(); i += 3)
            {
                mesh->vertices.emplace_back(
                    verts[i].GetFloat(),
                    verts[i + 1].GetFloat(),
                    verts[i + 2].GetFloat()
                    );
            }
        }
        else if (verts.IsObject())
        {
            auto const& type_string = verts["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(verts);

                if (!parse_as_vec3f(md, mesh->vertices))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    if (obj.HasMember("normals"))
    {
        auto const& normals = obj["normals"];

        if (normals.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < normals.Capacity(); i += 3)
            {
                mesh->normals.emplace_back(
                    normals[i].GetFloat(),
                    normals[i + 1].GetFloat(),
                    normals[i + 2].GetFloat()
                    );
            }
        }
        else if (normals.IsObject())
        {
            auto const& type_string = normals["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(normals);

                if (!parse_as_vec3f(md, mesh->normals))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        for (size_t i = 0; i < mesh->vertices.size(); i += 3)
        {
            vec3 v1 = mesh->vertices[i];
            vec3 v2 = mesh->vertices[i + 1];
            vec3 v3 = mesh->vertices[i + 2];

            vec3 e1 = v2 - v1;
            vec3 e2 = v3 - v1;

            vec3 gn = normalize(cross(e1, e2));

            mesh->normals.emplace_back(gn);
            mesh->normals.emplace_back(gn);
            mesh->normals.emplace_back(gn);
        }
    }

    if (obj.HasMember("tex_coords"))
    {
        auto const& tex_coords = obj["tex_coords"];

        if (tex_coords.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < tex_coords.Capacity(); i += 2)
            {
                mesh->tex_coords.emplace_back(
                    tex_coords[i].GetFloat(),
                    tex_coords[i + 1].GetFloat()
                    );
            }
        }
        else if (tex_coords.IsObject())
        {
            auto const& type_string = tex_coords["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(tex_coords);

                if (!parse_as_vec2f(md, mesh->tex_coords))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        for (size_t i = 0; i < mesh->vertices.size(); i += 3)
        {
            mesh->tex_coords.emplace_back(0.0f, 0.0f);
            mesh->tex_coords.emplace_back(0.0f, 0.0f);
            mesh->tex_coords.emplace_back(0.0f, 0.0f);
        }
    }

    if (obj.HasMember("colors"))
    {
        auto const& colors = obj["colors"];

        if (colors.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < colors.Capacity(); i += 3)
            {
                mesh->colors.emplace_back(
                    colors[i].GetFloat(),
                    colors[i + 1].GetFloat(),
                    colors[i + 2].GetFloat()
                    );
            }
        }
        else if (colors.IsObject())
        {
            auto const& type_string = colors["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(colors);

                if (!parse_as_vec3f(md, mesh->colors))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        for (size_t i = 0; i < mesh->vertices.size(); i += 3)
        {
            mesh->colors.emplace_back(1.0f, 1.0f, 1.0f);
            mesh->colors.emplace_back(1.0f, 1.0f, 1.0f);
            mesh->colors.emplace_back(1.0f, 1.0f, 1.0f);
        }
    }

    return mesh;
}

template <typename Object>
std::shared_ptr<sg::node> vsnray_parser::parse_indexed_triangle_mesh(Object const& obj)
{
    auto mesh = std::make_shared<sg::indexed_triangle_mesh>();

    if (obj.HasMember("vertex_indices"))
    {
        auto const& vertex_indices = obj["vertex_indices"];

        for (auto const& item : vertex_indices.GetArray())
        {
            mesh->vertex_indices.push_back(item.GetInt());
        }
    }

    if (obj.HasMember("normal_indices"))
    {
        auto const& normal_indices = obj["normal_indices"];

        for (auto const& item : normal_indices.GetArray())
        {
            mesh->normal_indices.push_back(item.GetInt());
        }
    }

    if (obj.HasMember("tex_coord_indices"))
    {
        auto const& tex_coord_indices = obj["tex_coord_indices"];

        for (auto const& item : tex_coord_indices.GetArray())
        {
            mesh->tex_coord_indices.push_back(item.GetInt());
        }
    }

    if (obj.HasMember("color_indices"))
    {
        auto const& color_indices = obj["color_indices"];

        for (auto const& item : color_indices.GetArray())
        {
            mesh->color_indices.push_back(item.GetInt());
        }
    }

    if (obj.HasMember("vertices"))
    {
        mesh->vertices = std::make_shared<aligned_vector<vec3>>();

        auto const& verts = obj["vertices"];

        if (verts.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < verts.Capacity(); i += 3)
            {
                mesh->vertices->emplace_back(
                    verts[i].GetFloat(),
                    verts[i + 1].GetFloat(),
                    verts[i + 2].GetFloat()
                    );
            }
        }
        else if (verts.IsObject())
        {
            auto const& type_string = verts["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(verts);

                if (!parse_as_vec3f(md, *mesh->vertices))
                {
                    throw std::runtime_error("");
                }
            }
        }
    }

    if (obj.HasMember("normals"))
    {
        mesh->normals = std::make_shared<aligned_vector<vec3>>();

        auto const& normals = obj["normals"];

        if (normals.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < normals.Capacity(); i += 3)
            {
                mesh->normals->emplace_back(
                    normals[i].GetFloat(),
                    normals[i + 1].GetFloat(),
                    normals[i + 2].GetFloat()
                    );
            }
        }
        else if (normals.IsObject())
        {
            auto const& type_string = normals["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(normals);

                if (!parse_as_vec3f(md, *mesh->normals))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    if (obj.HasMember("tex_coords"))
    {
        mesh->tex_coords = std::make_shared<aligned_vector<vec2>>();

        auto const& tex_coords = obj["tex_coords"];

        if (tex_coords.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < tex_coords.Capacity(); i += 2)
            {
                mesh->tex_coords->emplace_back(
                    tex_coords[i].GetFloat(),
                    tex_coords[i + 1].GetFloat()
                    );
            }
        }
        else if (tex_coords.IsObject())
        {
            auto const& type_string = tex_coords["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(tex_coords);

                if (!parse_as_vec2f(md, *mesh->tex_coords))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    if (obj.HasMember("colors"))
    {
        mesh->colors = std::make_shared<aligned_vector<vector<3, unorm<8>>>>();

        auto const& colors = obj["colors"];

        if (colors.IsArray())
        {
            for (rapidjson::SizeType i = 0; i < colors.Capacity(); i += 3)
            {
                mesh->colors->emplace_back(
                    colors[i].GetFloat(),
                    colors[i + 1].GetFloat(),
                    colors[i + 2].GetFloat()
                    );
            }
        }
        else if (colors.IsObject())
        {
            auto const& type_string = colors["type"];
            if (strncmp(type_string.GetString(), "file", 4) == 0)
            {
                auto md = parse_file_meta_data(colors);

                if (!parse_as_vec3f(md, *mesh->colors))
                {
                    throw std::runtime_error("");
                }
            }
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    return mesh;
}

template <typename Object>
data_file::meta_data vsnray_parser::parse_file_meta_data(Object const& obj)
{
    data_file::meta_data result;

    if (obj.HasMember("path"))
    {
        result.path = obj["path"].GetString();
    }
    else
    {
        throw std::runtime_error("");
    }

    if (obj.HasMember("encoding"))
    {
        std::string encoding = obj["encoding"].GetString();
        if (encoding == "ascii")
        {
            result.encoding = data_file::meta_data::Ascii;
        }
        else if (encoding == "binary")
        {
            result.encoding = data_file::meta_data::Binary;
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        throw std::runtime_error("");
    }

    if (obj.HasMember("data_type"))
    {
        std::string data_type = obj["data_type"].GetString();

        auto& data_type_map = data_file::meta_data::data_type_map;

        auto it = data_type_map.right.find(data_type);

        if (it != data_type_map.right.end())
        {
            result.data_type = it->second;
        }
        else
        {
            throw std::runtime_error("");
        }
    }
    else
    {
        throw std::runtime_error("");
    }

    if (obj.HasMember("num_items"))
    {
        result.num_items = obj["num_items"].GetInt();
    }
    else
    {
        throw std::runtime_error("");
    }

    if (obj.HasMember("compression"))
    {
        std::string compression = obj["compression"].GetString();
        if (compression == "none" || compression == "raw")
        {
            result.compression = data_file::meta_data::Raw;
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    if (obj.HasMember("separator"))
    {
        std::string separator = obj["separator"].GetString();
        result.separator = separator[0];
    }

    return result;
}

//-------------------------------------------------------------------------------------------------
// .vsnray writer
//

class vsnray_writer
{
public:

    vsnray_writer(rapidjson::Document& doc, std::string filename)
        : document_(doc)
        , filename_(filename)
    {
    }

    template <typename Object>
    void write_node(Object obj, std::shared_ptr<sg::node> const& n);

    template <typename Object>
    void write_transform(Object obj, std::shared_ptr<sg::transform> const& tr);

    template <typename Object>
    void write_surface_properties(Object obj, std::shared_ptr<sg::surface_properties> const& sp);

    template <typename Object>
    void write_triangle_mesh(Object obj, std::shared_ptr<sg::triangle_mesh> const& tm);

    template <typename Object>
    void write_indexed_triangle_mesh(Object obj, std::shared_ptr<sg::indexed_triangle_mesh> const& tm);

    template <typename Object, typename Container>
    void write_data_file(Object obj, data_file::meta_data md, Container const& cont);

private:

    std::string make_inline_filename(std::string node_name, std::string suffix);

    rapidjson::Document& document_;

    std::string filename_;

};

//-------------------------------------------------------------------------------------------------
// Write nodes
//

template <typename Object>
void vsnray_writer::write_node(Object obj, std::shared_ptr<sg::node> const& n)
{
    if (n == nullptr)
    {
        return;
    }

    auto& allocator = document_.GetAllocator();

    rapidjson::Value name(n->name().c_str(), allocator);
    obj.AddMember(
        rapidjson::StringRef("name"),
        name,
        allocator
        );

    if (auto tr = std::dynamic_pointer_cast<sg::transform>(n))
    {
        obj.AddMember(
            rapidjson::StringRef("type"),
            rapidjson::StringRef("transform"),
            allocator
            );

        write_transform(obj, tr);
    }
    else if (auto sp = std::dynamic_pointer_cast<sg::surface_properties>(n))
    {
        obj.AddMember(
            rapidjson::StringRef("type"),
            rapidjson::StringRef("surface_properties"),
            allocator
            );

        write_surface_properties(obj, sp);
    }
    else if (auto tm = std::dynamic_pointer_cast<sg::triangle_mesh>(n))
    {
        if (tm->flags() == 0)
        {
            obj.AddMember(
                rapidjson::StringRef("type"),
                rapidjson::StringRef("triangle_mesh"),
                allocator
                );

            write_triangle_mesh(obj, tm);

            tm->flags() = 1;
        }
    }
    else if (auto itm = std::dynamic_pointer_cast<sg::indexed_triangle_mesh>(n))
    {
        obj.AddMember(
            rapidjson::StringRef("type"),
            rapidjson::StringRef("indexed_triangle_mesh"),
            allocator
            );

        write_indexed_triangle_mesh(obj, itm);
    }
    else
    {
        obj.AddMember(
            rapidjson::StringRef("type"),
            rapidjson::StringRef("node"),
            allocator
            );
    }

    if (n->children().size() > 0)
    {
        rapidjson::Value arr(rapidjson::kArrayType);

        for (auto& c : n->children())
        {
            rapidjson::Value child;
            child.SetObject();
            write_node(child.GetObject(), c);

            arr.PushBack(child, allocator);
        }

        obj.AddMember("children", arr, allocator);
    }
}

template <typename Object>
void vsnray_writer::write_transform(Object obj, std::shared_ptr<sg::transform> const& tr)
{
    auto& allocator = document_.GetAllocator();

    rapidjson::Value mat(rapidjson::kArrayType);

    for (int i = 0; i < 16; ++i)
    {
        rapidjson::Value val;
        val.SetDouble(tr->matrix().data()[i]);
        mat.PushBack(val, allocator);
    }

    obj.AddMember("matrix", mat, allocator);
}

template <typename Object>
void vsnray_writer::write_surface_properties(Object obj, std::shared_ptr<sg::surface_properties> const& sp)
{
    auto& allocator = document_.GetAllocator();

    if (sp->material() != nullptr)
    {
        rapidjson::Value name(sp->material()->name().c_str(), allocator);
        obj.AddMember(
            rapidjson::StringRef("name"),
            name,
            allocator
            );

        rapidjson::Value jmat;
        jmat.SetObject();

        if (auto mat = std::dynamic_pointer_cast<sg::obj_material>(sp->material()))
        {
            jmat.AddMember(
                rapidjson::StringRef("type"),
                rapidjson::StringRef("obj"),
                allocator
                );

            rapidjson::Value ca(rapidjson::kArrayType);
            rapidjson::Value cd(rapidjson::kArrayType);
            rapidjson::Value cs(rapidjson::kArrayType);
            rapidjson::Value ce(rapidjson::kArrayType);

            for (int i = 0; i < 3; ++i)
            {
                ca.PushBack(rapidjson::Value().SetFloat(mat->ca[i]), allocator);
                cd.PushBack(rapidjson::Value().SetFloat(mat->cd[i]), allocator);
                cs.PushBack(rapidjson::Value().SetFloat(mat->cs[i]), allocator);
                ce.PushBack(rapidjson::Value().SetFloat(mat->ce[i]), allocator);
            }

            jmat.AddMember(rapidjson::StringRef("ca"), ca, allocator);
            jmat.AddMember(rapidjson::StringRef("cd"), cd, allocator);
            jmat.AddMember(rapidjson::StringRef("cs"), cs, allocator);
            jmat.AddMember(rapidjson::StringRef("ce"), ce, allocator);
        }
        else if (auto mat = std::dynamic_pointer_cast<sg::glass_material>(sp->material()))
        {
            jmat.AddMember(
                rapidjson::StringRef("type"),
                rapidjson::StringRef("glass"),
                allocator
                );

            rapidjson::Value ct(rapidjson::kArrayType);
            rapidjson::Value cr(rapidjson::kArrayType);
            rapidjson::Value ior(rapidjson::kArrayType);

            for (int i = 0; i < 3; ++i)
            {
                ct.PushBack(rapidjson::Value().SetFloat(mat->ct[i]), allocator);
                cr.PushBack(rapidjson::Value().SetFloat(mat->cr[i]), allocator);
                ior.PushBack(rapidjson::Value().SetFloat(mat->ior[i]), allocator);
            }

            jmat.AddMember(rapidjson::StringRef("ct"), ct, allocator);
            jmat.AddMember(rapidjson::StringRef("cr"), cr, allocator);
            jmat.AddMember(rapidjson::StringRef("ior"), ior, allocator);
        }
        else if (auto mat = std::dynamic_pointer_cast<sg::disney_material>(sp->material()))
        {
            jmat.AddMember(
                rapidjson::StringRef("type"),
                rapidjson::StringRef("disney"),
                allocator
                );

            rapidjson::Value base_color(rapidjson::kArrayType);

            for (int i = 0; i < 4; ++i)
            {
                base_color.PushBack(rapidjson::Value().SetFloat(mat->base_color[i]), allocator);
            }

            jmat.AddMember(rapidjson::StringRef("base_color"), base_color, allocator);
            jmat.AddMember(
                rapidjson::StringRef("spec_trans"),
                rapidjson::Value().SetFloat(mat->spec_trans),
                allocator
                );
            jmat.AddMember(
                rapidjson::StringRef("sheen"),
                rapidjson::Value().SetFloat(mat->sheen),
                allocator
                );
            jmat.AddMember(
                rapidjson::StringRef("sheen_tint"),
                rapidjson::Value().SetFloat(mat->sheen_tint),
                allocator
                );
            jmat.AddMember(
                rapidjson::StringRef("ior"),
                rapidjson::Value().SetFloat(mat->ior),
                allocator
                );
            jmat.AddMember(
                rapidjson::StringRef("refractive"),
                rapidjson::Value().SetFloat(mat->refractive),
                allocator
                );
            jmat.AddMember(
                rapidjson::StringRef("roughness"),
                rapidjson::Value().SetFloat(mat->roughness),
                allocator
                );
        }
        else
        {
            throw std::runtime_error("");
        }

        obj.AddMember(
            rapidjson::StringRef("material"),
            jmat,
            allocator
            );
    }

    for (auto& t : sp->textures())
    {
        std::string channel_name = t.first;

        rapidjson::Value jtex;
        jtex.SetObject();

        if (auto tex = std::dynamic_pointer_cast<sg::texture2d<vector<4, unorm<8>>>>(t.second))
        {
            // Not implemented yet
        }
#if VSNRAY_COMMON_HAVE_PTEX
        else if (auto tex = std::dynamic_pointer_cast<sg::ptex_texture>(t.second))
        {
            jtex.AddMember(
                rapidjson::StringRef("type"),
                rapidjson::StringRef("ptex"),
                allocator
                );

            rapidjson::Value filename(tex->filename().c_str(), allocator);
            jtex.AddMember(
                rapidjson::StringRef("filename"),
                filename,
                allocator
                );
        }
#endif // VSNRAY_COMMON_HAVE_PTEX
        else
        {
            throw std::runtime_error("");
        }
    }
}

template <typename Object>
void vsnray_writer::write_triangle_mesh(Object obj, std::shared_ptr<sg::triangle_mesh> const& tm)
{
    auto& allocator = document_.GetAllocator();

    // Write binary files

    if (!tm->vertices.empty())
    {
        data_file::meta_data md;
        md.path = make_inline_filename(tm->name(), "vert");
        md.encoding = data_file::meta_data::Binary;
        md.data_type = data_file::meta_data::Vec3f;
        md.num_items = tm->vertices.size();

        rapidjson::Value verts;
        verts.SetObject();

        write_data_file(verts.GetObject(), md, tm->vertices);

        obj.AddMember("vertices", verts, allocator);
    }

    if (!tm->normals.empty())
    {
        data_file::meta_data md;
        md.path = make_inline_filename(tm->name(), "norm");
        md.encoding = data_file::meta_data::Binary;
        md.data_type = data_file::meta_data::Vec3f;
        md.num_items = tm->normals.size();

        rapidjson::Value norms;
        norms.SetObject();

        write_data_file(norms.GetObject(), md, tm->normals);

        obj.AddMember("normals", norms, allocator);
    }

    if (!tm->tex_coords.empty())
    {
        data_file::meta_data md;
        md.path = make_inline_filename(tm->name(), "texc");
        md.encoding = data_file::meta_data::Binary;
        md.data_type = data_file::meta_data::Vec2f;
        md.num_items = tm->tex_coords.size();

        rapidjson::Value tcs;
        tcs.SetObject();

        write_data_file(tcs.GetObject(), md, tm->tex_coords);

        obj.AddMember("tex_coords", tcs, allocator);
    }

    if (!tm->colors.empty())
    {
        data_file::meta_data md;
        md.path = make_inline_filename(tm->name(), "colo");
        md.encoding = data_file::meta_data::Binary;
        md.data_type = data_file::meta_data::Vec3u8;
        md.num_items = tm->colors.size();

        rapidjson::Value cols;
        cols.SetObject();

        write_data_file(cols.GetObject(), md, tm->colors);

        obj.AddMember("colors", cols, allocator);
    }
}

template <typename Object>
void vsnray_writer::write_indexed_triangle_mesh(Object obj, std::shared_ptr<sg::indexed_triangle_mesh> const& itm)
{
}

template <typename Object, typename Container>
void vsnray_writer::write_data_file(Object obj, data_file::meta_data md, Container const& cont)
{
    // First try to write the actual data file

    // Don't overwrite
    if (boost::filesystem::exists(md.path))
    {
        throw std::runtime_error("");
    }

    // Check for consistency
    if (static_cast<int>(cont.size()) != md.num_items)
    {
        throw std::runtime_error("");
    }

    std::ofstream file(md.path, std::ios::binary);

    if (!file.good())
    {
        throw std::runtime_error("");
    }

    // Write data
    try
    {
        file.write(reinterpret_cast<char const*>(cont.data()), cont.size() * sizeof(typename Container::value_type));
    }
    catch (std::ios_base::failure)
    {
        throw std::runtime_error("");
    }

    assert(boost::filesystem::exists(md.path));

    // Now store a JSON node containing meta data to the document
    auto& allocator = document_.GetAllocator();

    std::string data_type = "";

    auto& data_type_map = data_file::meta_data::data_type_map;

    auto it = data_type_map.left.find(md.data_type);

    if (it != data_type_map.left.end())
    {
        data_type = it->second;
    }
    else
    {
        throw std::runtime_error("");
    }

    assert(!data_type.empty());

    if (md.encoding == data_file::meta_data::Binary)
    {
        obj.AddMember(
            rapidjson::StringRef("type"),
            rapidjson::StringRef("file"),
            allocator
            );

        rapidjson::Value path(md.path.c_str(), allocator);
        obj.AddMember(
            rapidjson::StringRef("path"),
            path,
            allocator
            );

        obj.AddMember(
            rapidjson::StringRef("encoding"),
            rapidjson::StringRef("binary"),
            allocator
            );

        rapidjson::Value dt(data_type.c_str(), allocator);
        obj.AddMember(
            rapidjson::StringRef("data_type"),
            dt,
            allocator
            );

        rapidjson::Value num_items;
        num_items.SetInt(md.num_items);
        obj.AddMember("num_items", num_items, allocator);

        obj.AddMember(
            rapidjson::StringRef("compression"),
            rapidjson::StringRef("none"),
            allocator
            );
    }
    else
    {
        // Not implemented yet
        throw std::runtime_error("");
    }
}

std::string vsnray_writer::make_inline_filename(std::string node_name, std::string suffix)
{
    std::string result;

    std::string insert = node_name;

    int inc = 0;

    for (;;)
    {
        std::string fn = filename_;

        if (!node_name.empty())
        {
            fn.append(std::string(".") + node_name);
        }
        else
        {
            std::string inc_str = std::to_string(inc);

            while (inc_str.length() < 8)
            {
                inc_str = std::string("0") + inc_str;
            }

            fn.append(std::string(".") + inc_str);

            ++inc;
        }

        if (!suffix.empty())
        {
            if (suffix[0] != '.')
            {
                fn.append(".");
            }

            fn.append(suffix);
        }

        if (!boost::filesystem::exists(fn))
        {
            result = fn;
            break;
        }
        else
        {
            if (!node_name.empty())
            {
                throw std::runtime_error("");
            }
        }
    }

    return result;
}

namespace visionaray
{

//-------------------------------------------------------------------------------------------------
// Interface
//

void load_vsnray(std::string const& filename, model& mod)
{
    std::vector<std::string> filenames(1);

    filenames[0] = filename;

    load_vsnray(filenames, mod);
}

void save_vsnray(std::string const& filename, model const& mod, file_base::save_options const& options)
{
    cfile file(filename, "w+");
    if (!file.good())
    {
        std::cerr << "Cannot open " << filename << '\n';
        return;
    }

    rapidjson::Document doc;
    doc.SetObject();

    vsnray_writer writer(doc, filename);
    writer.write_node(doc.GetObject(), mod.scene_graph);

    char buffer[65536];
    rapidjson::FileWriteStream fws(file.get(), buffer, sizeof(buffer));

    rapidjson::PrettyWriter<rapidjson::FileWriteStream> w(fws);
    doc.Accept(w);
}

void load_vsnray(std::vector<std::string> const& filenames, model& mod)
{
    auto root = std::make_shared<sg::node>();

    for (auto filename : filenames)
    {
        cfile file(filename, "r");
        if (!file.good())
        {
            std::cerr << "Cannot open " << filename << '\n';
            return;
        }

        char buffer[65536];
        rapidjson::FileReadStream frs(file.get(), buffer, sizeof(buffer));
        rapidjson::Document doc;
        doc.ParseStream(frs);

        if (doc.IsObject())
        {
            vsnray_parser parser(filename);
            root = parser.parse_node(doc.GetObject());
        }
        else
        {
            throw std::runtime_error("");
        }
    }

    if (mod.scene_graph == nullptr)
    {
        mod.scene_graph = root;
    }
    else
    {
        mod.scene_graph->add_child(root);
    }
}

} // visionaray