/* Copyright (C) 2010 Ion Torrent Systems, Inc. All Rights Reserved */
#include <vector>
#include <algorithm>
#include <limits>
#include "BfMetric.h"
#include <Eigen/Dense>
#include <Eigen/LU>
#include "IonH5File.h"
#include "IonH5Eigen.h"
#include "Stats.h"
#include "ChipReduction.h"
#include "SampleStats.h"

#define BFMETRIC_REDUCE_STEP 5
#define BFMETRIC_MIN_GOOD_WELLS 15

void BfMetric::CalcBfSdMetric(const short *__restrict image,
                              Mask *mask,
                              const char *__restrict bad_wells,
                              const float *__restrict t0,
                              const std::string &dbg_name,
                              int t0_window_start,
                              int t0_window_end,
                              int row_step,
                              int col_step,
                              int num_row_neighbors,
                              int num_col_neighbors) {

  // Calculate the average trace in region around each well
  m_nn_avg.Init(m_num_rows, m_num_cols, m_num_frames);
  m_nn_avg.SetGainMinMult(m_gain_min_mult);
  m_nn_avg.CalcCumulativeSum(image, mask, bad_wells);
  m_nn_avg.CalcNNAvgAndMinFrame(row_step, col_step,
                                num_row_neighbors, num_col_neighbors,
                                m_trace_min, m_trace_min_frame,
                                image, true);

  short * __restrict well_nn_diff = (short *) memalign(VEC8F_SIZE_B, (sizeof(short) * m_num_rows) * (m_num_cols * m_num_frames));
  int * __restrict well_nn_sq_diff = (int *) memalign(VEC8F_SIZE_B, (sizeof(int) * m_num_rows) * (m_num_cols * m_num_frames));
  memset(well_nn_diff, 0, (sizeof(short) * m_num_rows) * (m_num_cols * m_num_frames));
  memset(well_nn_sq_diff, 0, (sizeof(int) * m_num_rows) * (m_num_cols * m_num_frames));
  // Subtract each well from the average around it, this is beadfind buffering estimate
  short *__restrict diff_start = well_nn_diff;
  short *__restrict diff_end = diff_start + (m_num_rows * m_num_cols * m_num_frames);
  const short *__restrict img = image;
  const float *__restrict nn_avg = m_nn_avg.GetNNAvgImagePtr();
  const float *__restrict nn_avg_start = m_nn_avg.GetNNAvgImagePtr();
  while (diff_start != diff_end) {
    *diff_start = (short) (*img - *nn_avg_start + .5);
    diff_start++;
    img++;
    nn_avg_start++;
  }

  // if (!dbg_name.empty()) {
  //   H5File h5(dbg_name);
  //   h5.Open(true);
  //   int num_wells = m_num_rows * m_num_cols;
  //   const short *__restrict diff = well_nn_diff;
  //   const float *__restrict avg = nn_avg;
  //   const short *__restrict img = image;
  //   Eigen::MatrixXf mat_diff(m_num_cols * m_num_rows, m_num_frames);
  //   Eigen::MatrixXf mat_avg(m_num_cols * m_num_rows, m_num_frames);
  //   Eigen::MatrixXf mat_data(m_num_cols * m_num_rows, m_num_frames);
  //   for (int fIx = 0; fIx < m_num_frames; fIx++) {
  //     for (int well_ix = 0; well_ix < num_wells; well_ix++) {
  //       mat_diff(well_ix,fIx) = *diff;
  //       mat_avg(well_ix, fIx) = *avg;
  //       mat_data(well_ix, fIx) = *img;
  //       diff++;
  //       avg++;
  //       img++;
  //     }
  //   }
  //   H5Eigen::WriteMatrix(h5, "/bf_diff", mat_diff);
  //   H5Eigen::WriteMatrix(h5, "/bf_avg", mat_avg);
  //   H5Eigen::WriteMatrix(h5, "/bf_img", mat_data);
  //   h5.Close();
  // }

  // Gotta cleanup memory asap as these things take a lot of mem for 318 series
  m_nn_avg.Cleanup();

  // Now calculate the sum of differences and squared differences to determine frame in region
  // that has the maximum variance over wells in region and use that as bf metric

  // Calculate the cumulative sum of difference in the region per frame
  m_diff_sum_avg.Init(m_num_rows, m_num_cols, m_num_frames);
  m_diff_sum_avg.FreeNNAvg();
  m_diff_sum_avg.CalcCumulativeSum(well_nn_diff, mask, bad_wells);

  // Calculate the squared cumulative sum of difference
  diff_start = well_nn_diff;
  diff_end = diff_start + (m_num_rows * m_num_cols * m_num_frames);
  int *__restrict diff_sq_start = well_nn_sq_diff;
  while (diff_start != diff_end) {
    *diff_sq_start = (*diff_start) * (*diff_start);
    diff_start++;
    diff_sq_start++;
  }
  m_diff_sum_sq_avg.Init(m_num_rows, m_num_cols, m_num_frames);
  m_diff_sum_sq_avg.FreeNNAvg();
  m_diff_sum_sq_avg.CalcCumulativeSum(well_nn_sq_diff, mask, bad_wells);
  m_diff_sum_sq_avg.FreeNumGoodWells();
  FREEZ(&well_nn_sq_diff);

  // Figure out the frame around which a well that has the largest deviation
  int cs_frame_stride = (m_num_cols + 1) * (m_num_rows + 1);
  size_t frame_stride = m_num_cols * m_num_rows;
  const int64_t *__restrict cum_sum = m_diff_sum_avg.GetNNCumSumPtr();
  const int64_t *__restrict cum_sum_sq = m_diff_sum_sq_avg.GetNNCumSumPtr();
  const int *__restrict num_good_wells = m_diff_sum_avg.GetNumGoodWellsPtr();

  // Allocate scratch data frames to keep track of the frame with maximum variance in region
  uint8_t *max_sd_frame = (uint8_t*) memalign(VEC8F_SIZE_B, sizeof(uint8_t) * m_num_cols * m_num_rows);
  memset(max_sd_frame, 0, sizeof(uint8_t) * m_num_cols * m_num_rows);
  float *max_sd_value = (float*) memalign(VEC8F_SIZE_B, sizeof(float) * m_num_cols * m_num_rows);
  memset(max_sd_value, 0, sizeof(float) * m_num_cols * m_num_rows);

  // Go through each frame and calculate the sd for each window for each well.
  for (int fIx = 0; fIx < m_num_frames; fIx++) {
    const int64_t *__restrict cum_sum_frame = cum_sum + (fIx * cs_frame_stride);
    const int64_t *__restrict cum_sum_sq_frame = cum_sum_sq + (fIx * cs_frame_stride);
    int cs_col_size = m_num_cols + 1;
    int q1, q2,q3,q4;
    for (int rowIx = 0; rowIx < m_num_rows; rowIx++) {
      for (int colIx = 0; colIx < m_num_cols; colIx++) {
        int well_ix = rowIx * m_num_cols + colIx;
        int row_below_clip = row_step * (rowIx / row_step) -1;
        int row_above_clip = std::min(row_below_clip + row_step, m_num_rows - 1);
        int col_below_clip = col_step * (colIx / col_step)  -1;
        int col_above_clip = std::min(col_below_clip + col_step, m_num_cols - 1);

        int r_start = std::max(row_below_clip,rowIx - num_row_neighbors - 1); // -1 ok for calc of edge conditions
        int r_end = std::min(rowIx+num_row_neighbors, row_above_clip);
        int c_start = std::max(col_below_clip, colIx - num_col_neighbors - 1); // -1 ok for calc of edge conditions
        int c_end = std::min(col_above_clip,colIx + num_col_neighbors);

        int64_t sum_sq = 0, sum = 0;
        int good_wells = 0;
        int bad_wells = 0;
        /* calculate sd stats for wells */
        q1 = (r_end + 1) * cs_col_size + c_end + 1;
        q2 = (r_start +1) * cs_col_size + c_end + 1;
        q3 = (r_end+1) * cs_col_size + c_start + 1;
        q4 = (r_start+1) * cs_col_size + c_start + 1;

        good_wells = num_good_wells[q1] - num_good_wells[q2] - num_good_wells[q3] + num_good_wells[q4];
        /* combine them if we got some good wells */

        if (good_wells > 0) {
          sum  = cum_sum_frame[q1] - cum_sum_frame[q2] - cum_sum_frame[q3] + cum_sum_frame[q4];
          sum_sq  = cum_sum_sq_frame[q1] - cum_sum_sq_frame[q2] - cum_sum_sq_frame[q3] + cum_sum_sq_frame[q4];
          float region_sd = (sum_sq - (float)(sum*sum/good_wells))/good_wells;
          if (region_sd > max_sd_value[well_ix] && fIx >= (t0[well_ix] + t0_window_start) && fIx <= (t0[well_ix] + t0_window_end)) {
            max_sd_value[well_ix] = region_sd;
            max_sd_frame[well_ix] = (uint8_t)fIx;
          }
        }
        else {
          // do nothing, leave at zero
        }

      }
    }
  }

  // Cleanup as fast as possible these things take up a lot of memory
  m_diff_sum_avg.Cleanup();
  m_diff_sum_sq_avg.Cleanup();
  FREEZ(&max_sd_value);
  // Loop through and grab the data from frame with the value at the maximum sd frame point
  int *__restrict bf_start = m_bf_metric;
  int *__restrict bf_end = bf_start + frame_stride;
  uint8_t *__restrict max_frame_start = max_sd_frame;
  int *__restrict trace_min_start = m_trace_min_frame;
  int well_index = 0;
  int frame_counts[m_num_frames];
  memset(frame_counts, 0, sizeof(int) * m_num_frames);
  while(bf_start != bf_end) {
    frame_counts[*max_frame_start]++;
    *trace_min_start = *max_frame_start;
    *bf_start += *(well_nn_diff + (*max_frame_start) * frame_stride + well_index);
    *bf_start += *(well_nn_diff + (std::max(*max_frame_start-1,0)) * frame_stride + well_index);
    *bf_start += *(well_nn_diff + (std::min(*max_frame_start+1,m_num_frames-1)) * frame_stride + well_index);
    trace_min_start++;
    well_index++;
    max_frame_start++;
    bf_start++;
  }

  FREEZ(&well_nn_diff);
  FREEZ(&max_sd_frame);
  FREEZ(&max_sd_value);
}

void BfMetric::CalcBfSdMetricReduce(const short *__restrict image,
                                    Mask *mask,
                                    const char *__restrict bad_wells,
                                    const float *__restrict t0,
                                    const std::string &dbg_name,
                                    int t0_window_start,
                                    int t0_window_end,
                                    int row_step,
                                    int col_step,
                                    int num_row_neighbors,
                                    int num_col_neighbors,
                                    int frame_window) {

  // Calculate the average trace in region around each well
  ChipReduction reduce_avg(m_num_rows, m_num_cols, m_num_frames,
                       BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                       row_step, col_step, BFMETRIC_MIN_GOOD_WELLS);
  reduce_avg.Reduce(image, bad_wells);
  reduce_avg.SmoothBlocks(num_row_neighbors, num_col_neighbors);

  short * __restrict well_nn_diff = (short *) memalign(VEC8F_SIZE_B, (sizeof(short) * m_num_rows) * (m_num_cols * m_num_frames));
  int * __restrict well_nn_sq_diff = (int *) memalign(VEC8F_SIZE_B, (sizeof(int) * m_num_rows) * (m_num_cols * m_num_frames));
  memset(well_nn_diff, 0, (sizeof(short) * m_num_rows) * (m_num_cols * m_num_frames));
  memset(well_nn_sq_diff, 0, (sizeof(int) * m_num_rows) * (m_num_cols * m_num_frames));

  // Subtract each well from the average around it, this is beadfind buffering estimate
  short *__restrict diff_start = well_nn_diff;
  int *__restrict diff_sq_start = well_nn_sq_diff;
  const short *__restrict img = image;

  for (int frame_ix = 0; frame_ix < m_num_frames; frame_ix++) {
    for (int row_ix = 0; row_ix < m_num_rows; row_ix++) {
      for (int col_ix = 0; col_ix < m_num_cols; col_ix++) {
        *diff_start = (short)(*img - reduce_avg.GetSmoothEst(row_ix, col_ix, frame_ix) + .5f);
        *diff_sq_start = (int)(*diff_start) * (*diff_start);
        diff_start++;
        diff_sq_start++;
        img++;
      }
    }
  }

  if (!dbg_name.empty()) {
    H5File h5(dbg_name);
    h5.Open(true);
    int num_wells = m_num_rows * m_num_cols;
    const short *__restrict diff = well_nn_diff;
    const short *__restrict img = image;
    Eigen::MatrixXf mat_diff(m_num_cols * m_num_rows, m_num_frames);
    Eigen::MatrixXf mat_avg(m_num_cols * m_num_rows, m_num_frames);
    Eigen::MatrixXf mat_data(m_num_cols * m_num_rows, m_num_frames);
    for (int fIx = 0; fIx < m_num_frames; fIx++) {
      for (int well_ix = 0; well_ix < num_wells; well_ix++) {
        mat_diff(well_ix,fIx) = *diff;
        mat_avg(well_ix, fIx) = reduce_avg.GetSmoothEst(well_ix / m_num_cols, well_ix % m_num_cols, fIx);
        mat_data(well_ix, fIx) = *img;
        diff++;
        img++;
      }
    }
    H5Eigen::WriteMatrix(h5, "/bf_diff", mat_diff);
    H5Eigen::WriteMatrix(h5, "/bf_avg", mat_avg);
    H5Eigen::WriteMatrix(h5, "/bf_img", mat_data);
    h5.Close();
  }

  ChipReduction diff_avg(m_num_rows, m_num_cols, m_num_frames,
                         BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                         row_step, col_step, BFMETRIC_MIN_GOOD_WELLS);
  diff_avg.SetAvgReduce(false);
  diff_avg.Reduce(well_nn_diff, bad_wells);
  diff_avg.SmoothBlocks(row_step, col_step);

  ChipReduction diff_sq_avg(m_num_rows, m_num_cols, m_num_frames,
                            BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                            row_step, col_step, BFMETRIC_MIN_GOOD_WELLS);
  diff_sq_avg.SetAvgReduce(false);
  diff_sq_avg.Reduce(well_nn_sq_diff, bad_wells);
  diff_sq_avg.SmoothBlocks(row_step, col_step);

  NNAvg good_wells(diff_avg.GetBlockHeight(), diff_avg.GetBlockWidth(), 1);
  good_wells.CalcCumulativeSum(diff_avg.GetGoodWells(), diff_avg.GetBadWells());

  // Now calculate the sum of differences and squared differences to determine frame in region
  // that has the maximum variance over wells in region and use that as bf metric

  // Allocate scratch data frames to keep track of the frame with maximum variance in region
  uint8_t *max_sd_frame = (uint8_t*) memalign(VEC8F_SIZE_B, sizeof(uint8_t) * m_num_cols * m_num_rows);
  memset(max_sd_frame, 0, sizeof(uint8_t) * m_num_cols * m_num_rows);
  float *max_sd_value = (float*) memalign(VEC8F_SIZE_B, sizeof(float) * m_num_cols * m_num_rows);
  memset(max_sd_value, 0, sizeof(float) * m_num_cols * m_num_rows);

  // Go through each frame and calculate the sd for each window for each well.
  for (int fIx = 0; fIx < m_num_frames; fIx++) {
    int well_ix = 0;
    for (int rowIx = 0; rowIx < m_num_rows; rowIx++) {
      for (int colIx = 0; colIx < m_num_cols; colIx++) {
        float sum_sq = 0;
        float sum = 0;
        float total_sum = 0;
        int good_count = 0;
        diff_avg.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, fIx,
                                       row_step/BFMETRIC_REDUCE_STEP, col_step/BFMETRIC_REDUCE_STEP,
                                       num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                       sum, good_count);
        diff_sq_avg.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, fIx,
                                          row_step/BFMETRIC_REDUCE_STEP, col_step/BFMETRIC_REDUCE_STEP,
                                          num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                          sum_sq, good_count);
        good_wells.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                                row_step/BFMETRIC_REDUCE_STEP, col_step/BFMETRIC_REDUCE_STEP,
                                num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                total_sum, good_count);
        // if (rowIx == 100 && colIx == 100) {
        //   fprintf(stdout, "Here we go\n");
        // }
        if (total_sum > 0) {
          float region_sd = (sum_sq - (float)(sum * sum/total_sum))/total_sum;
          if (region_sd > max_sd_value[well_ix] && fIx >= (t0[well_ix] + t0_window_start) && fIx <= (t0[well_ix] + t0_window_end)) {
            max_sd_value[well_ix] = region_sd;
            max_sd_frame[well_ix] = (uint8_t)fIx;
          }
        }
        else {
          // do nothing, leave at zero
        }
        well_ix++;
      }
    }
  }

  FREEZ(&max_sd_value);
  // Loop through and grab the data from frame with the value at the maximum sd frame point
  int *__restrict bf_start = m_bf_metric;
  int frame_stride = (m_num_rows * m_num_cols);
  int *__restrict bf_end = bf_start + frame_stride;
  uint8_t *__restrict max_frame_start = max_sd_frame;
  int *__restrict trace_min_start = m_trace_min_frame;
  int well_index = 0;
  int frame_counts[m_num_frames];
  memset(frame_counts, 0, sizeof(int) * m_num_frames);
  while(bf_start != bf_end) {
    frame_counts[*max_frame_start]++;
    *trace_min_start = *max_frame_start;
    // *bf_start += *(well_nn_diff + (*max_frame_start) * frame_stride + well_index);
    // *bf_start += *(well_nn_diff + (std::max(*max_frame_start-1,0)) * frame_stride + well_index);
    // *bf_start += *(well_nn_diff + (std::min(*max_frame_start+1,m_num_frames-1)) * frame_stride + well_index);
    int f_start = max(0, *max_frame_start - frame_window);
    int f_end = min(m_num_frames - 1, *max_frame_start + frame_window);
    for (int f_ix = f_start; f_ix <= f_end; f_ix++) {
      *bf_start += *(well_nn_diff + f_ix * frame_stride + well_index);
    }
    trace_min_start++;
    well_index++;
    max_frame_start++;
    bf_start++;
  }

  FREEZ(&well_nn_diff);
  FREEZ(&well_nn_sq_diff);
  FREEZ(&max_sd_frame);
}

void BfMetric::NormalizeMetric(const char *bad_wells,
                               int row_clip, int col_clip,
                               int num_row_neighbors,
                               int num_col_neighbors) {

  /* Get some filteres to avoid really bad outliers... */
  std::vector<char> norm_bad_wells(m_num_cols * m_num_cols);
  std::vector<float> bf_metric_sorted(norm_bad_wells.size());
  std::copy(m_bf_metric, m_bf_metric + norm_bad_wells.size(), bf_metric_sorted.begin());
  std::sort(bf_metric_sorted.begin(), bf_metric_sorted.end());
  float high_threshold = std::numeric_limits<float>::max();
  float low_threshold = -1.0f * high_threshold;
  if (bf_metric_sorted.size() > BFMETRIC_MIN_GOOD_WELLS) {
    float q25 = ionStats::quantile_sorted(bf_metric_sorted, .25);
    float q75 = ionStats::quantile_sorted(bf_metric_sorted, .75);
    float q02 = ionStats::quantile_sorted(bf_metric_sorted, .02);
    float q98 = ionStats::quantile_sorted(bf_metric_sorted, .98);
    float iqr = q75-q25;
    float hx3 = q75 + 3 * iqr;
    float lx3 = q25 - 3 * iqr;
    high_threshold = min(q98, hx3);
    low_threshold = max(lx3, q02);
  }
  /* If we fail our filters don't use it, otherwise default to regular mask. */
  for (size_t i = 0; i < norm_bad_wells.size(); i++) {
    if (m_bf_metric[i] <= low_threshold || m_bf_metric[i] >= high_threshold) {
      norm_bad_wells[i] = 1;
    }
    else {
      norm_bad_wells[i] = bad_wells[i];
    }
  }

  ChipReduction sum(m_num_rows, m_num_cols, 1,
                    BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                    row_clip, col_clip,
                    BFMETRIC_MIN_GOOD_WELLS);
  sum.SetAvgReduce(false);
  sum.Reduce(m_bf_metric, &norm_bad_wells[0]);
  sum.SmoothBlocks(num_row_neighbors, num_col_neighbors);
  std::vector<int> sq_bf_metric(m_num_cols * m_num_rows);
  int *__restrict bf_start = m_bf_metric;
  int *__restrict bf_end = m_bf_metric + m_num_rows * m_num_cols;
  int *__restrict sq_start = &sq_bf_metric[0];
  while(bf_start != bf_end) {
    *sq_start = *bf_start * *bf_start;
    sq_start++;
    bf_start++;
  }
  ChipReduction sum_sq(m_num_rows, m_num_cols, 1,
                       BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                       row_clip, col_clip,
                       BFMETRIC_MIN_GOOD_WELLS);
  sum_sq.SetAvgReduce(false);
  sum_sq.Reduce(&sq_bf_metric[0], &norm_bad_wells[0]);
  sum_sq.SmoothBlocks(num_row_neighbors, num_col_neighbors);

  std::vector<float> norm_ref(m_num_cols * m_num_rows);
  NNAvg good_wells(sum.GetBlockHeight(), sum.GetBlockWidth(), 1);
  good_wells.CalcCumulativeSum(sum.GetGoodWells(), sum.GetBadWells());
  SampleStats<float> bf_stats;
  SampleStats<float> ref_bf_stats;
  bf_start = m_bf_metric;
  bf_end = bf_start + norm_ref.size();
  const char *bad_start = &norm_bad_wells[0];
  while(bf_start != bf_end) {
    if (*bad_start == 0) {
      bf_stats.AddValue(*bf_start);
    }
    bad_start++;
    bf_start++;
  }

  int well_ix = 0;
  for (int rowIx = 0; rowIx < m_num_rows; rowIx++) {
    for (int colIx = 0; colIx < m_num_cols; colIx++) {
      float s = 0;
      float sq = 0;
      float total_sum = 0;
      int good_count = 0;
      sum.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                                row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                                num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                s, good_count);
      sum_sq.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                                   row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                                   num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                   sq, good_count);
      good_wells.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                              row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                              num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                              total_sum, good_count);
      float ref_avg = sum.GetSmoothEst(rowIx, colIx, 0) / total_sum;
      float var = (sq - (float)(s * s)/total_sum)/total_sum;
      //        norm_ref[well_ix] = (m_bf_metric[well_ix] - ref_avg)/var
      if (good_count > 0 && var > 0) {
        m_norm_bf_metric[well_ix] = (m_bf_metric[well_ix] - ref_avg)/sqrt(var);
      }
      else {
        m_norm_bf_metric[well_ix] = (m_bf_metric[well_ix] - bf_stats.GetMean())/bf_stats.GetSD();
      }
      well_ix++;
    }
  }
}

void BfMetric::NormalizeMetricRef(const char *ref_wells, const char *bad_wells,
                                  int row_clip, int col_clip,
                                  int num_row_neighbors,
                                  int num_col_neighbors) {

  /* Get some filteres to avoid really bad outliers... */
  std::vector<char> norm_bad_wells(m_num_rows * m_num_cols);
  std::vector<char> ref_flip(norm_bad_wells.size());
  std::vector<float> bf_metric_sorted(norm_bad_wells.size());
  std::copy(m_bf_metric, m_bf_metric + norm_bad_wells.size(), bf_metric_sorted.begin());
  std::sort(bf_metric_sorted.begin(), bf_metric_sorted.end());
  float high_threshold = std::numeric_limits<float>::max();
  float low_threshold = -1.0f * high_threshold;
  if (bf_metric_sorted.size() > BFMETRIC_MIN_GOOD_WELLS) {
    float q25 = ionStats::quantile_sorted(bf_metric_sorted, .25);
    float q75 = ionStats::quantile_sorted(bf_metric_sorted, .75);
    float q02 = ionStats::quantile_sorted(bf_metric_sorted, .02);
    float q98 = ionStats::quantile_sorted(bf_metric_sorted, .98);
    float iqr = q75-q25;
    float hx3 = q75 + 3 * iqr;
    float lx3 = q25 - 3 * iqr;
    high_threshold = min(q98, hx3);
    low_threshold = max(lx3, q02);
  }
  /* If we fail our filters don't use it, otherwise default to regular mask. */
  for (size_t i = 0; i < norm_bad_wells.size(); i++) {
    if (ref_wells[i] == 0) {
      ref_flip[i] = 1;
    }
    else {
      ref_flip[i] = 0;
    }
    if (m_bf_metric[i] <= low_threshold || m_bf_metric[i] >= high_threshold) {
      norm_bad_wells[i] = 1;
    }
    else {
      norm_bad_wells[i] = bad_wells[i];
    }
  }

  ChipReduction sum(m_num_rows, m_num_cols, 1,
                    BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                    row_clip, col_clip,
                    BFMETRIC_MIN_GOOD_WELLS);
  sum.SetAvgReduce(false);
  sum.Reduce(m_bf_metric, &norm_bad_wells[0]);
  sum.SmoothBlocks(num_row_neighbors, num_col_neighbors);
  std::vector<int> sq_bf_metric(m_num_cols * m_num_rows);
  int *__restrict bf_start = m_bf_metric;
  int *__restrict bf_end = m_bf_metric + m_num_rows * m_num_cols;
  int *__restrict sq_start = &sq_bf_metric[0];
  while(bf_start != bf_end) {
    *sq_start = *bf_start * *bf_start;
    sq_start++;
    bf_start++;
  }
  ChipReduction sum_sq(m_num_rows, m_num_cols, 1,
                       BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                       row_clip, col_clip,
                       BFMETRIC_MIN_GOOD_WELLS);
  sum_sq.SetAvgReduce(false);
  sum_sq.Reduce(&sq_bf_metric[0], &norm_bad_wells[0]);
  sum_sq.SmoothBlocks(num_row_neighbors, num_col_neighbors);

  std::vector<float> norm_ref(m_num_cols * m_num_rows);
  NNAvg good_wells(sum.GetBlockHeight(), sum.GetBlockWidth(), 1);
  good_wells.CalcCumulativeSum(sum.GetGoodWells(), sum.GetBadWells());
  SampleStats<float> bf_stats;
  SampleStats<float> ref_bf_stats;
  bf_start = m_bf_metric;
  bf_end = bf_start + norm_ref.size();
  const char *bad_start = &norm_bad_wells[0];
  const char *ref_start = ref_wells;
  while(bf_start != bf_end) {
    if (*bad_start == 0 && *ref_start == 1) {
      bf_stats.AddValue(*bf_start);
    }
    bad_start++;
    bf_start++;
    ref_start++;
  }

  ChipReduction ref_avg(m_num_rows, m_num_cols, 1,
                       BFMETRIC_REDUCE_STEP, BFMETRIC_REDUCE_STEP,
                       row_clip, col_clip,
                       BFMETRIC_MIN_GOOD_WELLS);
  ref_avg.Reduce(m_bf_metric, &ref_flip[0]);
  ref_avg.SmoothBlocks(num_row_neighbors, num_col_neighbors);

  int well_ix = 0;
  for (int rowIx = 0; rowIx < m_num_rows; rowIx++) {
    for (int colIx = 0; colIx < m_num_cols; colIx++) {
      float s = 0;
      float sq = 0;
      float total_sum = 0;
      int good_count = 0;
      sum.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                                row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                                num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                s, good_count);
      sum_sq.m_nn_avg.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                                   row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                                   num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                                   sq, good_count);
      good_wells.GetRegionSum(rowIx/BFMETRIC_REDUCE_STEP, colIx/BFMETRIC_REDUCE_STEP, 0,
                              row_clip/BFMETRIC_REDUCE_STEP, col_clip/BFMETRIC_REDUCE_STEP,
                              num_row_neighbors/BFMETRIC_REDUCE_STEP, num_col_neighbors/BFMETRIC_REDUCE_STEP,
                              total_sum, good_count);
      float all_mean = sum.GetSmoothEst(rowIx, colIx, 0) / total_sum;
      float ref_mean = ref_avg.GetSmoothEst(rowIx, colIx, 0);
      float var = (sq - (float)(s * s)/total_sum)/total_sum;
      if (good_count > 0 && var > 0) {
        float sd = sqrt(var);
        float ref_center = ref_mean / sd;
        m_norm_bf_metric[well_ix] = (m_bf_metric[well_ix] - all_mean)/sd - ref_center;
      }
      else {
        m_norm_bf_metric[well_ix] = (m_bf_metric[well_ix] - bf_stats.GetMean())/bf_stats.GetSD();
      }
      well_ix++;
    }
  }
}


#ifndef TEMPLATE_MATCH
#define TEMPLATE_MATCH
#include <pcl/console/time.h>
#include <std_msgs/Int8.h>
//#include <pcl/features/fpfh.h>
#include <pcl/features/normal_3d.h>
#include <pcl/filters/crop_box.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/io/pcd_io.h>
#include <pcl/kdtree/kdtree_flann.h>
// #include <pcl/memory.h>
#include <pcl/features/fpfh.h>
#include <pcl/features/normal_3d.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/io/pcd_io.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/pcl_macros.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/registration/ia_ransac.h>
#include <pcl/segmentation/extract_clusters.h>
#include <pcl/visualization/cloud_viewer.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl_conversions/pcl_conversions.h>
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
typedef pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ>
    PCLHandler;
#include <math.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/registration/icp.h>

#include <Eigen/Core>
#include <fstream>
#include <iostream>
#include <limits>
#include <map>
#include <vector>

#include "bpmsg/arm_state.h"
#include "bpmsg/pose.h"

// #include "template_match/pose.h"
typedef pcl::PointXYZ PointT;
typedef pcl::PointCloud<pcl::PointXYZ> PointCloud;
typedef pcl::PointCloud<pcl::Normal> SurfaceNormals;
typedef pcl::PointCloud<pcl::FPFHSignature33> LocalFeatures;
typedef pcl::search::KdTree<pcl::PointXYZ> SearchMethod;
using namespace std;

class FeatureCloud {
 public:
  // A bit of shorthand

  FeatureCloud()
      : search_method_xyz_(new SearchMethod),
        normal_radius_(0.05f),
        feature_radius_(0.05f) {}

  ~FeatureCloud() {}

  // Process the given cloud
  void setInputCloud(PointCloud::Ptr xyz) {
    xyz_ = xyz;
    processInput();
  }

  // Load and process the cloud in the given PCD file
  void loadInputCloud(const std::string &pcd_file) {
    xyz_ = PointCloud::Ptr(new PointCloud);
    pcl::io::loadPCDFile(pcd_file, *xyz_);
    processInput();
  }

  // Get a pointer to the cloud 3D points
  PointCloud::Ptr getPointCloud() const { return (xyz_); }

  // Get a pointer to the cloud of 3D surface normals
  SurfaceNormals::Ptr getSurfaceNormals() const { return (normals_); }

  // Get a pointer to the cloud of feature descriptors
  LocalFeatures::Ptr getLocalFeatures() const { return (features_); }

 protected:
  // Compute the surface normals and local features
  void processInput() {
    computeSurfaceNormals();
    computeLocalFeatures();
  }

  // Compute the surface normals
  void computeSurfaceNormals() {
    normals_ = SurfaceNormals::Ptr(new SurfaceNormals);

    pcl::NormalEstimation<pcl::PointXYZ, pcl::Normal> norm_est;
    norm_est.setInputCloud(xyz_);
    norm_est.setSearchMethod(search_method_xyz_);
    norm_est.setRadiusSearch(normal_radius_);
    norm_est.compute(*normals_);
  }

  // Compute the local feature descriptors
  void computeLocalFeatures() {
    features_ = LocalFeatures::Ptr(new LocalFeatures);

    pcl::FPFHEstimation<pcl::PointXYZ, pcl::Normal, pcl::FPFHSignature33>
        fpfh_est;
    fpfh_est.setInputCloud(xyz_);
    fpfh_est.setInputNormals(normals_);
    fpfh_est.setSearchMethod(search_method_xyz_);
    fpfh_est.setRadiusSearch(feature_radius_);
    fpfh_est.compute(*features_);
  }

 private:
  // Point cloud data
  PointCloud::Ptr xyz_;
  SurfaceNormals::Ptr normals_;
  LocalFeatures::Ptr features_;
  SearchMethod::Ptr search_method_xyz_;

  // Parameters
  float normal_radius_;
  float feature_radius_;
};

class TemplateAlignment {
 public:
  // A struct for storing alignment results
  struct Result {
    float fitness_score;
    Eigen::Matrix4f final_transformation;
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  };

  TemplateAlignment()
      : min_sample_distance_(0.005f),               // 0.005f
        max_correspondence_distance_(0.1f * 0.1f),  // 0.1f * 0.1f
        nr_iterations_(1000) {                      // 1000
    // Initialize the parameters in the Sample Consensus Initial Alignment
    // (SAC-IA) algorithm
    sac_ia_.setMinSampleDistance(min_sample_distance_);
    sac_ia_.setMaxCorrespondenceDistance(max_correspondence_distance_);
    sac_ia_.setMaximumIterations(nr_iterations_);
  }

  ~TemplateAlignment() {}

  // Set the given cloud as the target to which the templates will be aligned
  void setTargetCloud(FeatureCloud &target_cloud) {
    target_ = target_cloud;
    sac_ia_.setInputTarget(target_cloud.getPointCloud());
    sac_ia_.setTargetFeatures(target_cloud.getLocalFeatures());
  }

  // Add the given cloud to the list of template clouds
  void addTemplateCloud(FeatureCloud &template_cloud) {
    templates_.push_back(template_cloud);
  }

  // Align the given template cloud to the target specified by setTargetCloud ()
  void align(FeatureCloud &template_cloud, TemplateAlignment::Result &result) {
    sac_ia_.setInputCloud(template_cloud.getPointCloud());
    sac_ia_.setSourceFeatures(template_cloud.getLocalFeatures());

    pcl::PointCloud<pcl::PointXYZ> registration_output;
    sac_ia_.align(registration_output);

    result.fitness_score =
        (float)sac_ia_.getFitnessScore(max_correspondence_distance_);
    result.final_transformation = sac_ia_.getFinalTransformation();
  }

  // Align all of template clouds set by addTemplateCloud to the target
  // specified by setTargetCloud ()
  void alignAll(std::vector<TemplateAlignment::Result,
                            Eigen::aligned_allocator<Result> > &results) {
    results.resize(templates_.size());
    for (size_t i = 0; i < templates_.size(); ++i) {
      align(templates_[i], results[i]);
    }
  }

  // Align all of template clouds to the target cloud to find the one with best
  // alignment score
  int findBestAlignment(TemplateAlignment::Result &result) {
    // Align all of the templates to the target cloud
    std::vector<Result, Eigen::aligned_allocator<Result> > results;
    alignAll(results);

    // Find the template with the best (lowest) fitness score
    float lowest_score = std::numeric_limits<float>::infinity();
    int best_template = 0;
    for (size_t i = 0; i < results.size(); ++i) {
      const Result &r = results[i];
      if (r.fitness_score < lowest_score) {
        lowest_score = r.fitness_score;
        best_template = (int)i;
      }
    }

    // Output the best alignment
    result = results[best_template];
    return (best_template);
  }

 private:
  // A list of template clouds and the target to which they will be aligned
  std::vector<FeatureCloud> templates_;
  FeatureCloud target_;

  // The Sample Consensus Initial Alignment (SAC-IA) registration routine and
  // its parameters
  pcl::SampleConsensusInitialAlignment<pcl::PointXYZ, pcl::PointXYZ,
                                       pcl::FPFHSignature33>
      sac_ia_;
  float min_sample_distance_;
  float max_correspondence_distance_;
  int nr_iterations_;
};

class template_match {
 public:
  double min_x, min_y, min_z, max_x, max_y, max_z;
  double model_size;
  int object_index = 0;
  double theta, dx, dy, dz;
  int arm_state;
  int max_num, min_num;
  int object_num = 0;
  int com_flag = 0;
  int pick_index;
  bool view_on;
  Eigen::Vector3f euler_angles;
  Eigen::Matrix<float, 4, 1> origin_pos;
  Eigen::Matrix<float, 3, 1> origin_angle;
  std::vector<Eigen::Matrix<float, 3, 1> > target_angle;
  Eigen::Matrix<float, 3, 1> target_vector;
  std::vector<Eigen::Matrix<float, 4, 1> > target_pos;
  friend class TemplateAlignment;
  friend class FeatureCloud;
  ros::Subscriber bat_sub;
  ros::Subscriber arm_sub;
  ros::Publisher trans_pub;
  bpmsg::pose result_pose;
  PointCloud::Ptr model_;
  // pcl::PointCloud<PointT>::Ptr cloud_;
  map<int, double> height_map;
  vector<pcl::PointCloud<PointT>::Ptr> goals;
  // vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> t_cloud;
  std::vector<FeatureCloud> object_templates;
  vector<Eigen::Matrix4f> final_trans;
  template_match();
  void showCloud(pcl::PointCloud<PointT>::Ptr cloud1,
                 pcl::PointCloud<PointT>::Ptr cloud2);
  void cloudCB(const sensor_msgs::PointCloud2 &input);
  void armCB(const bpmsg::arm_state &msg);
  void init();
  void mainloop();
  void match(pcl::PointCloud<pcl::PointXYZ>::Ptr goal,
             pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_2, int index);
  void cluster(pcl::PointCloud<PointT>::Ptr cloud_,
               pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_2);
  void commandCB(const std_msgs::Int8 &msg);
  pcl::PointCloud<pcl::PointXYZ>::Ptr transclouds(
      pcl::PointCloud<pcl::PointXYZ>::Ptr source_cloud, double theta, double dx,
      double dy, double dz);
  Eigen::Matrix4f icp(pcl::PointCloud<PointT>::Ptr cloud_in,
                      pcl::PointCloud<PointT>::Ptr cloud_out);
};
#endif


//////////////////////////////////////////////////////////////////////////////
//  Boost.Assign v2                                                         //
//                                                                          //
//  Copyright (C) 2003-2004 Thorsten Ottosen                                //
//  Copyright (C) 2011 Erwann Rogard                                        //
//  Use, modification and distribution are subject to the                   //
//  Boost Software License, Version 1.0. (See accompanying file             //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)        //
//////////////////////////////////////////////////////////////////////////////
#include <cmath>
#include <deque>
#include <list>
#include <map>
#include <vector>
#include <boost/array.hpp>
#include <boost/assign/v2/deque.hpp>
#include <boost/assign/v2/option/data.hpp>
#include <boost/assign/v2/interpreter/csv/csv_basic.hpp>
#include <boost/assign/v2/include/csv_deque_basic.hpp>
#include <boost/assign/v2/include/csv_put_basic.hpp>
#include <boost/assign/v2/include/put.hpp>
#include <boost/assign/v2/include/deque.hpp>
#include <boost/assign/v2/support/config/check.hpp>

#include <boost/function.hpp>
#include <boost/lambda/bind.hpp>
#include <boost/lambda/lambda.hpp>
#include <boost/numeric/conversion/bounds.hpp>
#include <boost/range/algorithm/equal.hpp>
#include <boost/typeof/typeof.hpp>
#include <libs/assign/v2/test/option/data.h>

namespace test_assign_v2{
namespace xxx_option{
namespace xxx_data{

    void test()
    {
        namespace as2 = boost::assign::v2;
        {
            // (*fp) resolves error C2440 using MSVC
            //[test_data1
            std::vector<double> exponent;
            /*<-*/typedef double(*fp)(double);/*->*/
            typedef boost::function<double(double)> f_;

            BOOST_AUTO(
                _data,
                ( as2::_data = f_( /*<-*/fp(/*->*/ log10 /*<-*/)/*->*/ ) )
            );

            csv(
                as2::put( exponent ) % _data
                , 1.0, 10.0, 100.0, 1000.0, 10000.0
            );

            double eps = boost::numeric::bounds<double>::smallest();
            BOOST_ASSIGN_V2_CHECK( fabs( exponent.front() - 0.0 ) < eps );
            BOOST_ASSIGN_V2_CHECK( fabs( exponent.back() - 4.0 ) < eps );
            //]
        }
        {
            //[test_data2
            int k = 1;

            BOOST_AUTO(
                _data,
                (
                    as2::_data = (
                        boost::lambda::var( k ) *= boost::lambda::_1
                    )
                )
            );

            BOOST_ASSIGN_V2_CHECK(
                boost::range::equal(
                    csv(
                        as2::deque<int>( as2::_nil ) % _data,
                        1, 2, 3, 4, 5
                    ),
                    as2::csv_deque( 1, 2, 6, 24, 120 )
                )
            );
            //]
        }
        {
            //[test_data3
            typedef std::string month_; typedef int days_;
            typedef std::map<month_, days_> map_;
            typedef map_::value_type pair_;

            map_ map;
            as2::csv_put<as2::value_>(
                map,
                pair_( "jan", 31 ),
                pair_( "feb", 28 ),
                pair_( "mar", 31 )
            );

            BOOST_ASSIGN_V2_CHECK( map["jan"] == 31 );
            BOOST_ASSIGN_V2_CHECK( map["feb"] == 28 );
            BOOST_ASSIGN_V2_CHECK( map["mar"] == 31 );
            //]
        }
    }

}// xxx_data
}// xxx_option
}// test_assign_v2


//---------------------------------------------------------------------------//
// Copyright (c) 2013 Kyle Lutz <kyle.r.lutz@gmail.com>
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//
// See http://kylelutz.github.com/compute for more information.
//---------------------------------------------------------------------------//

#ifndef BOOST_COMPUTE_ALGORITHM_DETAIL_COUNT_IF_WITH_THREADS_HPP
#define BOOST_COMPUTE_ALGORITHM_DETAIL_COUNT_IF_WITH_THREADS_HPP

#include <numeric>

#include <boost/compute/detail/meta_kernel.hpp>
#include <boost/compute/container/vector.hpp>

namespace boost {
namespace compute {
namespace detail {

template<class InputIterator, class Predicate>
class count_if_with_threads_kernel : meta_kernel
{
public:
    typedef typename
        std::iterator_traits<InputIterator>::value_type
        value_type;

    count_if_with_threads_kernel()
        : meta_kernel("count_if_with_threads")
    {
    }

    void set_args(InputIterator first,
                  InputIterator last,
                  Predicate predicate)

    {
        typedef typename std::iterator_traits<InputIterator>::value_type T;

        m_size = detail::iterator_range_size(first, last);

        m_size_arg = add_arg<const ulong_>("size");
        m_counts_arg = add_arg<ulong_ *>(memory_object::global_memory, "counts");

        *this <<
            // thread parameters
            "const uint gid = get_global_id(0);\n" <<
            "const uint block_size = size / get_global_size(0);\n" <<
            "const uint start = block_size * gid;\n" <<
            "uint end = 0;\n" <<
            "if(gid == get_global_size(0) - 1)\n" <<
            "    end = size;\n" <<
            "else\n" <<
            "    end = block_size * gid + block_size;\n" <<

            // count values
            "uint count = 0;\n" <<
            "for(uint i = start; i < end; i++){\n" <<
                decl<const T>("value") << "="
                    << first[expr<uint_>("i")] << ";\n" <<
                if_(predicate(var<const T>("value"))) << "{\n" <<
                    "count++;\n" <<
                "}\n" <<
            "}\n" <<

            // write count
            "counts[gid] = count;\n";
    }

    size_t exec(command_queue &queue)
    {
        const device &device = queue.get_device();
        const context &context = queue.get_context();

        size_t threads = device.compute_units();

        const size_t minimum_block_size = 2048;
        if(m_size / threads < minimum_block_size){
            threads = static_cast<size_t>(
                          (std::max)(
                              std::ceil(float(m_size) / minimum_block_size),
                              1.0f
                          )
                      );
        }

        // storage for counts
        ::boost::compute::vector<ulong_> counts(threads, context);

        // exec kernel
        set_arg(m_size_arg, static_cast<ulong_>(m_size));
        set_arg(m_counts_arg, counts.get_buffer());
        exec_1d(queue, 0, threads, 1);

        // copy counts to the host
        std::vector<ulong_> host_counts(threads);
        ::boost::compute::copy(counts.begin(), counts.end(), host_counts.begin(), queue);

        // return sum of counts
        return std::accumulate(host_counts.begin(), host_counts.end(), size_t(0));
    }

private:
    size_t m_size;
    size_t m_size_arg;
    size_t m_counts_arg;
};

// counts values that match the predicate using one thread per block. this is
// optimized for cpu-type devices with a small number of compute units.
template<class InputIterator, class Predicate>
inline size_t count_if_with_threads(InputIterator first,
                                    InputIterator last,
                                    Predicate predicate,
                                    command_queue &queue)
{
    count_if_with_threads_kernel<InputIterator, Predicate> kernel;
    kernel.set_args(first, last, predicate);
    return kernel.exec(queue);
}

} // end detail namespace
} // end compute namespace
} // end boost namespace

#endif // BOOST_COMPUTE_ALGORITHM_DETAIL_COUNT_IF_WITH_THREADS_HPP


#include "Common/Common.h"
#include "Simulation/CubicSDFCollisionDetection.h"
#include "Simulation/DistanceFieldCollisionDetection.h"
#include "Simulation/TimeManager.h"
#include "Simulation/SimulationModel.h"
#include "Simulation/TimeStepController.h"
#include "Utils/FileSystem.h"
#include "Utils/Logger.h"
#include "Utils/OBJLoader.h"
#include "Utils/TetGenLoader.h"
#include "Utils/Timing.h"
#include "Demos/Visualization/MiniGL.h"
#include "Demos/Visualization/Selection.h"
#include "Demos/Visualization/Visualization.h"
#include "Demos/Common/DemoBase.h"
#include "Demos/Common/TweakBarParameters.h"
#include "Simulation/Simulation.h"
#include "GL/glut.h"
#include <Eigen/Dense>
#include <iostream>
#include <set>

#include "StiffRodsSceneLoader.h"

#define _USE_MATH_DEFINES
#include "math.h"

// Enable memory leak detection
#if defined(_DEBUG) && !defined(EIGEN_ALIGN)
#define new DEBUG_NEW
#endif

using namespace PBD;
using namespace Eigen;
using namespace std;
using namespace Utilities;

void initParameters();
void timeStep();
void buildModel();
void readScene();
void render();
void reset();

void TW_CALL setContactTolerance(const void *value, void *clientData);
void TW_CALL getContactTolerance(void *value, void *clientData);
void TW_CALL setContactStiffnessRigidBody(const void *value, void *clientData);
void TW_CALL getContactStiffnessRigidBody(void *value, void *clientData);
void TW_CALL setContactStiffnessParticleRigidBody(const void *value, void *clientData);
void TW_CALL getContactStiffnessParticleRigidBody(void *value, void *clientData);
void TW_CALL setBendingMethod(const void *value, void *clientData);
void TW_CALL getBendingMethod(void *value, void *clientData);
void TW_CALL setClothSimulationMethod(const void *value, void *clientData);
void TW_CALL getClothSimulationMethod(void *value, void *clientData);
void TW_CALL setSolidSimulationMethod(const void *value, void *clientData);
void TW_CALL getSolidSimulationMethod(void *value, void *clientData);

void singleStep();

DemoBase *base;
CubicSDFCollisionDetection cd;

short clothSimulationMethod = 2;
short solidSimulationMethod = 2;
short bendingMethod = 2;

string sceneFileName = "/scenes/Wilberforce_scene.json";
string sceneName;
Vector3r camPos;
Vector3r camLookat;

// main
int main(int argc, char **argv)
{
	REPORT_MEMORY_LEAKS;

	std::string exePath = FileSystem::getProgramPath();
	std::string dataPath = exePath + "/" + std::string(PBD_DATA_PATH);

	if (argc > 1)
		sceneFileName = string(argv[1]);
	else
		sceneFileName = FileSystem::normalizePath(dataPath + sceneFileName);

	base = new DemoBase();
	base->init(argc, argv, sceneFileName.c_str());
	base->setSceneLoader(new StiffRodsSceneLoader());

	SimulationModel *model = new SimulationModel();
	model->init();
	Simulation::getCurrent()->setModel(model);

	buildModel();

	initParameters();
	base->readParameters();

	Simulation::getCurrent()->setSimulationMethodChangedCallback([&]() { reset(); initParameters(); base->getSceneLoader()->readParameterObject(Simulation::getCurrent()->getTimeStep()); });

	// OpenGL
	MiniGL::setClientIdleFunc(50, timeStep);
	MiniGL::setKeyFunc(0, 'r', reset);
	MiniGL::setKeyFunc(1, 's', singleStep);
	MiniGL::setClientSceneFunc(render);
	MiniGL::setViewport(40.0f, 0.1f, 500.0f, camPos, camLookat);

	SimulationModel::TriangleModelVector &triModels = model->getTriangleModels();
	SimulationModel::TetModelVector &tetModels = model->getTetModels();
	TwAddVarCB(MiniGL::getTweakBar(), "ContactTolerance", TW_TYPE_REAL, setContactTolerance, getContactTolerance, &cd, " label='Contact tolerance'  min=0.0 step=0.001 precision=3 group=Simulation ");
	TwAddVarCB(MiniGL::getTweakBar(), "ContactStiffnessRigidBody", TW_TYPE_REAL, setContactStiffnessRigidBody, getContactStiffnessRigidBody, &model, " label='Contact stiffness RB'  min=0.0 step=0.1 precision=2 group=Simulation ");
	TwAddVarCB(MiniGL::getTweakBar(), "ContactStiffnessParticleRigidBody", TW_TYPE_REAL, setContactStiffnessParticleRigidBody, getContactStiffnessParticleRigidBody, &model, " label='Contact stiffness Particle-RB'  min=0.0 step=0.1 precision=2 group=Simulation ");
	if (triModels.size() > 0)
	{
		TwType enumType2 = TwDefineEnum("ClothSimulationMethodType", NULL, 0);
		TwAddVarCB(MiniGL::getTweakBar(), "ClothSimulationMethod", enumType2, setClothSimulationMethod, getClothSimulationMethod, &clothSimulationMethod, " label='Cloth sim. method' enum='0 {None}, 1 {Distance constraints}, 2 {FEM based PBD}, 3 {Strain based dynamics}' group=Simulation");
	}
	if (tetModels.size() > 0)
	{
	TwType enumType3 = TwDefineEnum("SolidSimulationMethodType", NULL, 0);
	TwAddVarCB(MiniGL::getTweakBar(), "SolidSimulationMethod", enumType3, setSolidSimulationMethod, getSolidSimulationMethod, &solidSimulationMethod,
		" label='Solid sim. method' enum='0 {None}, 1 {Volume constraints}, 2 {FEM based PBD}, 3 {Strain based dynamics (no inversion handling)}, 4 {Shape matching (no inversion handling)}' group=Simulation");
	}
	if (triModels.size() > 0)
	{
		TwType enumType4 = TwDefineEnum("BendingMethodType", NULL, 0);
		TwAddVarCB(MiniGL::getTweakBar(), "BendingMethod", enumType4, setBendingMethod, getBendingMethod, &bendingMethod, " label='Bending method' enum='0 {None}, 1 {Dihedral angle}, 2 {Isometric bending}' group=Bending");
	}

	glutMainLoop();

	base->cleanup();

	Utilities::Timing::printAverageTimes();
	Utilities::Timing::printTimeSums();

	delete Simulation::getCurrent();
	delete base;
	delete model;

	return 0;
}

void initParameters()
{
	TwRemoveAllVars(MiniGL::getTweakBar());
	TweakBarParameters::cleanup();

	MiniGL::initTweakBarParameters();

	TweakBarParameters::createParameterGUI();
	TweakBarParameters::createParameterObjectGUI(base);
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent());
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getModel());
	TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getTimeStep());
}

void reset()
{
	Utilities::Timing::printAverageTimes();
	Utilities::Timing::reset();

	Simulation::getCurrent()->reset();
	base->getSelectedParticles().clear();

	Simulation::getCurrent()->getModel()->cleanup();
	Simulation::getCurrent()->getTimeStep()->getCollisionDetection()->cleanup();

	readScene();
}

void singleStep()
{
	// Simulation code
	SimulationModel *model = Simulation::getCurrent()->getModel();
	Simulation::getCurrent()->getTimeStep()->step(*model);

	// Update visualization models
	const ParticleData &pd = model->getParticles();
	for (unsigned int i = 0; i < model->getTetModels().size(); i++)
	{
		model->getTetModels()[i]->updateMeshNormals(pd);
		model->getTetModels()[i]->updateVisMesh(pd);
	}
	for (unsigned int i = 0; i < model->getTriangleModels().size(); i++)
	{
		model->getTriangleModels()[i]->updateMeshNormals(pd);
	}
}

void timeStep()
{
	const Real pauseAt = base->getValue<Real>(DemoBase::PAUSE_AT);
	if ((pauseAt > 0.0) && (pauseAt < TimeManager::getCurrent()->getTime()))
		base->setValue(DemoBase::PAUSE, true);

	if (base->getValue<bool>(DemoBase::PAUSE))
		return;

	// Simulation code
	SimulationModel *model = Simulation::getCurrent()->getModel();
	const unsigned int numSteps = base->getValue<unsigned int>(DemoBase::NUM_STEPS_PER_RENDER);
	for (unsigned int i = 0; i < numSteps; i++)
	{
		START_TIMING("SimStep");
		Simulation::getCurrent()->getTimeStep()->step(*model);
		STOP_TIMING_AVG;
	}

	// Update visualization models
	const ParticleData &pd = model->getParticles();
	for (unsigned int i = 0; i < model->getTetModels().size(); i++)
	{
		model->getTetModels()[i]->updateMeshNormals(pd);
		model->getTetModels()[i]->updateVisMesh(pd);
	}
	for (unsigned int i = 0; i < model->getTriangleModels().size(); i++)
	{
		model->getTriangleModels()[i]->updateMeshNormals(pd);
	}
}

void render()
{
	base->render();
}

void buildModel()
{
	TimeManager::getCurrent()->setTimeStepSize(static_cast<Real>(0.005));

	SimulationModel *model = Simulation::getCurrent()->getModel();
	Simulation::getCurrent()->getTimeStep()->setCollisionDetection(*model, &cd);

	readScene();
}

void initTriangleModelConstraints()
{
	// init constraints
	SimulationModel *model = Simulation::getCurrent()->getModel();
	for (unsigned int cm = 0; cm < model->getTriangleModels().size(); cm++)
	{
		const unsigned int offset = model->getTriangleModels()[cm]->getIndexOffset();
		if (clothSimulationMethod == 1)
		{
			const unsigned int nEdges = model->getTriangleModels()[cm]->getParticleMesh().numEdges();
			const IndexedFaceMesh::Edge *edges = model->getTriangleModels()[cm]->getParticleMesh().getEdges().data();
			for (unsigned int i = 0; i < nEdges; i++)
			{
				const unsigned int v1 = edges[i].m_vert[0] + offset;
				const unsigned int v2 = edges[i].m_vert[1] + offset;

				model->addDistanceConstraint(v1, v2);
			}
		}
		else if (clothSimulationMethod == 2)
		{

			TriangleModel::ParticleMesh &mesh = model->getTriangleModels()[cm]->getParticleMesh();
			const unsigned int *tris = mesh.getFaces().data();
			const unsigned int nFaces = mesh.numFaces();
			for (unsigned int i = 0; i < nFaces; i++)
			{
				const unsigned int v1 = tris[3 * i] + offset;
				const unsigned int v2 = tris[3 * i + 1] + offset;
				const unsigned int v3 = tris[3 * i + 2] + offset;
				model->addFEMTriangleConstraint(v1, v2, v3);
			}
		}
		else if (clothSimulationMethod == 3)
		{
			TriangleModel::ParticleMesh &mesh = model->getTriangleModels()[cm]->getParticleMesh();
			const unsigned int *tris = mesh.getFaces().data();
			const unsigned int nFaces = mesh.numFaces();
			for (unsigned int i = 0; i < nFaces; i++)
			{
				const unsigned int v1 = tris[3 * i] + offset;
				const unsigned int v2 = tris[3 * i + 1] + offset;
				const unsigned int v3 = tris[3 * i + 2] + offset;
				model->addStrainTriangleConstraint(v1, v2, v3);
			}
		}
		if (bendingMethod != 0)
		{
			TriangleModel::ParticleMesh &mesh = model->getTriangleModels()[cm]->getParticleMesh();
			unsigned int nEdges = mesh.numEdges();
			const TriangleModel::ParticleMesh::Edge *edges = mesh.getEdges().data();
			const unsigned int *tris = mesh.getFaces().data();
			for (unsigned int i = 0; i < nEdges; i++)
			{
				const int tri1 = edges[i].m_face[0];
				const int tri2 = edges[i].m_face[1];
				if ((tri1 != 0xffffffff) && (tri2 != 0xffffffff))
				{
					// Find the triangle points which do not lie on the axis
					const int axisPoint1 = edges[i].m_vert[0];
					const int axisPoint2 = edges[i].m_vert[1];
					int point1 = -1;
					int point2 = -1;
					for (int j = 0; j < 3; j++)
					{
						if ((tris[3 * tri1 + j] != axisPoint1) && (tris[3 * tri1 + j] != axisPoint2))
						{
							point1 = tris[3 * tri1 + j];
							break;
						}
					}
					for (int j = 0; j < 3; j++)
					{
						if ((tris[3 * tri2 + j] != axisPoint1) && (tris[3 * tri2 + j] != axisPoint2))
						{
							point2 = tris[3 * tri2 + j];
							break;
						}
					}
					if ((point1 != -1) && (point2 != -1))
					{
						const unsigned int vertex1 = point1 + offset;
						const unsigned int vertex2 = point2 + offset;
						const unsigned int vertex3 = edges[i].m_vert[0] + offset;
						const unsigned int vertex4 = edges[i].m_vert[1] + offset;
						if (bendingMethod == 1)
							model->addDihedralConstraint(vertex1, vertex2, vertex3, vertex4);
						else if (bendingMethod == 2)
							model->addIsometricBendingConstraint(vertex1, vertex2, vertex3, vertex4);
					}
				}
			}
		}
	}
}

void initTetModelConstraints()
{
	// init constraints
	SimulationModel *model = Simulation::getCurrent()->getModel();
	for (unsigned int cm = 0; cm < model->getTetModels().size(); cm++)
	{
		const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
		const unsigned int nTets = model->getTetModels()[cm]->getParticleMesh().numTets();
		const unsigned int *tets = model->getTetModels()[cm]->getParticleMesh().getTets().data();
		const IndexedTetMesh::VertexTets *vTets = model->getTetModels()[cm]->getParticleMesh().getVertexTets().data();
		if (solidSimulationMethod == 1)
		{
			const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
			const unsigned int nEdges = model->getTetModels()[cm]->getParticleMesh().numEdges();
			const IndexedTetMesh::Edge *edges = model->getTetModels()[cm]->getParticleMesh().getEdges().data();
			for (unsigned int i = 0; i < nEdges; i++)
			{
				const unsigned int v1 = edges[i].m_vert[0] + offset;
				const unsigned int v2 = edges[i].m_vert[1] + offset;

				model->addDistanceConstraint(v1, v2);
			}

			for (unsigned int i = 0; i < nTets; i++)
			{
				const unsigned int v1 = tets[4 * i] + offset;
				const unsigned int v2 = tets[4 * i + 1] + offset;
				const unsigned int v3 = tets[4 * i + 2] + offset;
				const unsigned int v4 = tets[4 * i + 3] + offset;

				model->addVolumeConstraint(v1, v2, v3, v4);
			}
		}
		else if (solidSimulationMethod == 2)
		{
			TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
			for (unsigned int i = 0; i < nTets; i++)
			{
				const unsigned int v1 = tets[4 * i] + offset;
				const unsigned int v2 = tets[4 * i + 1] + offset;
				const unsigned int v3 = tets[4 * i + 2] + offset;
				const unsigned int v4 = tets[4 * i + 3] + offset;

				model->addFEMTetConstraint(v1, v2, v3, v4);
			}
		}
		else if (solidSimulationMethod == 3)
		{
			TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
			for (unsigned int i = 0; i < nTets; i++)
			{
				const unsigned int v1 = tets[4 * i] + offset;
				const unsigned int v2 = tets[4 * i + 1] + offset;
				const unsigned int v3 = tets[4 * i + 2] + offset;
				const unsigned int v4 = tets[4 * i + 3] + offset;

				model->addStrainTetConstraint(v1, v2, v3, v4);
			}
		}
		else if (solidSimulationMethod == 4)
		{
			TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
			for (unsigned int i = 0; i < nTets; i++)
			{
				const unsigned int v[4] = { tets[4 * i] + offset,
					tets[4 * i + 1] + offset,
					tets[4 * i + 2] + offset,
					tets[4 * i + 3] + offset };
				// Important: Divide position correction by the number of clusters
				// which contain the vertex.
				const unsigned int nc[4] = { vTets[v[0]].m_numTets, vTets[v[1]].m_numTets, vTets[v[2]].m_numTets, vTets[v[3]].m_numTets };
				model->addShapeMatchingConstraint(4, v, nc);
			}
		}
	}
}

void loadObj(const std::string &filename, VertexData &vd, IndexedFaceMesh &mesh, const Vector3r &scale)
{
	std::vector<OBJLoader::Vec3f> x;
	std::vector<OBJLoader::Vec3f> normals;
	std::vector<OBJLoader::Vec2f> texCoords;
	std::vector<MeshFaceIndices> faces;
	OBJLoader::Vec3f s = { (float)scale[0], (float)scale[1], (float)scale[2] };
	OBJLoader::loadObj(filename, &x, &faces, &normals, &texCoords, s);

	mesh.release();
	const unsigned int nPoints = (unsigned int)x.size();
	const unsigned int nFaces = (unsigned int)faces.size();
	const unsigned int nTexCoords = (unsigned int)texCoords.size();
	mesh.initMesh(nPoints, nFaces * 2, nFaces);
	vd.reserve(nPoints);
	for (unsigned int i = 0; i < nPoints; i++)
	{
		vd.addVertex(Vector3r(x[i][0], x[i][1], x[i][2]));
	}
	for (unsigned int i = 0; i < nTexCoords; i++)
	{
		mesh.addUV(texCoords[i][0], texCoords[i][1]);
	}
	for (unsigned int i = 0; i < nFaces; i++)
	{
		// Reduce the indices by one
		int posIndices[3];
		int texIndices[3];
		for (int j = 0; j < 3; j++)
		{
			posIndices[j] = faces[i].posIndices[j] - 1;
			if (nTexCoords > 0)
			{
				texIndices[j] = faces[i].texIndices[j] - 1;
				mesh.addUVIndex(texIndices[j]);
			}
		}

		mesh.addFace(&posIndices[0]);
	}
	mesh.buildNeighbors();

	mesh.updateNormals(vd, 0);
	mesh.updateVertexNormals(vd);

	LOG_INFO << "Number of triangles: " << nFaces;
	LOG_INFO << "Number of vertices: " << nPoints;
}

CubicSDFCollisionDetection::GridPtr generateSDF(const std::string &modelFile, const std::string &collisionObjectFileName, const Eigen::Matrix<unsigned int, 3, 1> &resolutionSDF,
	VertexData &vd, IndexedFaceMesh &mesh)
{
	const std::string basePath = FileSystem::getFilePath(base->getSceneFile());
	const string cachePath = basePath + "/Cache";
	const std::string modelFileName = FileSystem::getFileNameWithExt(modelFile);
	CubicSDFCollisionDetection::GridPtr distanceField;

	if (collisionObjectFileName == "")
	{
		std::string md5FileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + ".md5");
		string md5Str = FileSystem::getFileMD5(modelFile);
		bool md5 = false;
		if (FileSystem::fileExists(md5FileName))
			md5 = FileSystem::checkMD5(md5Str, md5FileName);

		// check MD5 if cache file is available
		const string resStr = to_string(resolutionSDF[0]) + "_" + to_string(resolutionSDF[1]) + "_" + to_string(resolutionSDF[2]);
		const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");
		bool foundCacheFile = FileSystem::fileExists(sdfFileName);

		if (foundCacheFile && md5)
		{
			LOG_INFO << "Load cached SDF: " << sdfFileName;
			distanceField = std::make_shared<CubicSDFCollisionDetection::Grid>(sdfFileName);
		}
		else
		{
			std::vector<unsigned int> &faces = mesh.getFaces();
			const unsigned int nFaces = mesh.numFaces();
#ifdef USE_DOUBLE
			Discregrid::TriangleMesh sdfMesh(&vd.getPosition(0)[0], faces.data(), vd.size(), nFaces);
#else
			// if type is float, copy vector to double vector
			std::vector<double> doubleVec;
			doubleVec.resize(3 * vd.size());
			for (unsigned int i = 0; i < vd.size(); i++)
				for (unsigned int j = 0; j < 3; j++)
					doubleVec[3 * i + j] = vd.getPosition(i)[j];
			Discregrid::TriangleMesh sdfMesh(&doubleVec[0], faces.data(), vd.size(), nFaces);
#endif
			Discregrid::MeshDistance md(sdfMesh);
			Eigen::AlignedBox3d domain;
			for (auto const& x : sdfMesh.vertices())
			{
				domain.extend(x);
			}
			domain.max() += 0.1 * Eigen::Vector3d::Ones();
			domain.min() -= 0.1 * Eigen::Vector3d::Ones();

			LOG_INFO << "Set SDF resolution: " << resolutionSDF[0] << ", " << resolutionSDF[1] << ", " << resolutionSDF[2];
			distanceField = std::make_shared<CubicSDFCollisionDetection::Grid>(domain, std::array<unsigned int, 3>({ resolutionSDF[0], resolutionSDF[1], resolutionSDF[2] }));
			auto func = Discregrid::DiscreteGrid::ContinuousFunction{};
			func = [&md](Eigen::Vector3d const& xi) {return md.signedDistanceCached(xi); };
			LOG_INFO << "Generate SDF for " << modelFile;
			distanceField->addFunction(func, true);
			if (FileSystem::makeDir(cachePath) == 0)
			{
				LOG_INFO << "Save SDF: " << sdfFileName;
				distanceField->save(sdfFileName);
				FileSystem::writeMD5File(modelFile, md5FileName);
			}
		}
	}
	else
	{
		std::string fileName = collisionObjectFileName;
		if (FileSystem::isRelativePath(fileName))
		{
			fileName = FileSystem::normalizePath(basePath + "/" + fileName);
		}
		LOG_INFO << "Load SDF: " << fileName;
		distanceField = std::make_shared<CubicSDFCollisionDetection::Grid>(fileName);
	}
	return distanceField;
}

/** Create the rigid body model
*/
void readScene()
{
	SimulationModel *model = Simulation::getCurrent()->getModel();
	SimulationModel::RigidBodyVector &rb = model->getRigidBodies();
	SimulationModel::TriangleModelVector &triModels = model->getTriangleModels();
	SimulationModel::TetModelVector &tetModels = model->getTetModels();
	SimulationModel::ConstraintVector &constraints = model->getConstraints();

	StiffRodsSceneLoader::RodSceneData data;
	StiffRodsSceneLoader *loader = static_cast<StiffRodsSceneLoader*>(base->getSceneLoader());
	loader->readStiffRodsScene(sceneFileName, data);
	LOG_INFO << "Scene: " << sceneFileName;

	camPos = data.m_camPosition;
	camLookat = data.m_camLookat;

	TimeManager::getCurrent()->setTimeStepSize(data.m_timeStepSize);

	sceneName = data.m_sceneName;

	if (data.m_triangleModelSimulationMethod != -1)
		clothSimulationMethod = data.m_triangleModelSimulationMethod;
	if (data.m_tetModelSimulationMethod != -1)
		solidSimulationMethod = data.m_tetModelSimulationMethod;
	if (data.m_triangleModelBendingMethod != -1)
		bendingMethod = data.m_triangleModelBendingMethod;
	cd.setTolerance(data.m_contactTolerance);
	model->setContactStiffnessRigidBody(data.m_contactStiffnessRigidBody);
	model->setContactStiffnessParticleRigidBody(data.m_contactStiffnessParticleRigidBody);

	model->setValue<Real>(SimulationModel::CLOTH_BENDING_STIFFNESS, data.m_cloth_bendingStiffness);
	model->setValue<Real>(SimulationModel::CLOTH_STIFFNESS_XX, data.m_cloth_xxStiffness);
	model->setValue<Real>(SimulationModel::CLOTH_STIFFNESS_YY, data.m_cloth_yyStiffness);
	model->setValue<Real>(SimulationModel::CLOTH_STIFFNESS_XY, data.m_cloth_xyStiffness);
	model->setValue<Real>(SimulationModel::CLOTH_POISSON_RATIO_XY, data.m_cloth_xyPoissonRatio);
	model->setValue<Real>(SimulationModel::CLOTH_POISSON_RATIO_YX, data.m_cloth_yxPoissonRatio);
	model->setValue<bool>(SimulationModel::CLOTH_NORMALIZE_STRETCH, data.m_cloth_normalizeStretch);
	model->setValue<bool>(SimulationModel::CLOTH_NORMALIZE_SHEAR, data.m_cloth_normalizeShear);

	model->setValue<Real>(SimulationModel::SOLID_STIFFNESS, data.m_solid_stiffness);
	model->setValue<Real>(SimulationModel::SOLID_POISSON_RATIO, data.m_solid_poissonRatio);
	model->setValue<bool>(SimulationModel::SOLID_NORMALIZE_STRETCH, data.m_solid_normalizeStretch);
	model->setValue<bool>(SimulationModel::SOLID_NORMALIZE_SHEAR, data.m_solid_normalizeShear);

	//////////////////////////////////////////////////////////////////////////
	// rigid bodies
	//////////////////////////////////////////////////////////////////////////

	// map file names to loaded geometry to prevent multiple imports of same files
	std::map<std::string, pair<VertexData, IndexedFaceMesh>> objFiles;
	std::map<std::string, CubicSDFCollisionDetection::GridPtr> distanceFields;
	for (unsigned int rbIndex = 0; rbIndex < data.m_rigidBodyData.size(); rbIndex++)
	{
		SceneLoader::RigidBodyData &rbd = data.m_rigidBodyData[rbIndex];

		// Check if already loaded
		rbd.m_modelFile = FileSystem::normalizePath(rbd.m_modelFile);
		if (objFiles.find(rbd.m_modelFile) == objFiles.end())
		{
			IndexedFaceMesh mesh;
			VertexData vd;
			loadObj(rbd.m_modelFile, vd, mesh, Vector3r::Ones());
			objFiles[rbd.m_modelFile] = { vd, mesh };
		}

		const std::string basePath = FileSystem::getFilePath(sceneFileName);
		const string cachePath = basePath + "/Cache";
		const string resStr = to_string(rbd.m_resolutionSDF[0]) + "_" + to_string(rbd.m_resolutionSDF[1]) + "_" + to_string(rbd.m_resolutionSDF[2]);
		const std::string modelFileName = FileSystem::getFileNameWithExt(rbd.m_modelFile);
		const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");

		std::string sdfKey = rbd.m_collisionObjectFileName;
		if (sdfKey == "")
		{
			sdfKey = sdfFileName;
		}
		if (distanceFields.find(sdfKey) == distanceFields.end())
		{
			// Generate SDF
			if (rbd.m_collisionObjectType == SceneLoader::SDF)
			{
				if (rbd.m_collisionObjectFileName == "")
				{
					std::string md5FileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + ".md5");
					string md5Str = FileSystem::getFileMD5(rbd.m_modelFile);
					bool md5 = false;
					if (FileSystem::fileExists(md5FileName))
						md5 = FileSystem::checkMD5(md5Str, md5FileName);

					// check MD5 if cache file is available
					const string resStr = to_string(rbd.m_resolutionSDF[0]) + "_" + to_string(rbd.m_resolutionSDF[1]) + "_" + to_string(rbd.m_resolutionSDF[2]);
					const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");
					bool foundCacheFile = FileSystem::fileExists(sdfFileName);

					if (foundCacheFile && md5)
					{
						LOG_INFO << "Load cached SDF: " << sdfFileName;
						distanceFields[sdfFileName] = std::make_shared<CubicSDFCollisionDetection::Grid>(sdfFileName);
					}
					else
					{
						VertexData &vd = objFiles[rbd.m_modelFile].first;
						IndexedFaceMesh &mesh = objFiles[rbd.m_modelFile].second;

						std::vector<unsigned int> &faces = mesh.getFaces();
						const unsigned int nFaces = mesh.numFaces();

#ifdef USE_DOUBLE
						Discregrid::TriangleMesh sdfMesh(&vd.getPosition(0)[0], faces.data(), vd.size(), nFaces);
#else
						// if type is float, copy vector to double vector
						std::vector<double> doubleVec;
						doubleVec.resize(3 * vd.size());
						for (unsigned int i = 0; i < vd.size(); i++)
							for (unsigned int j = 0; j < 3; j++)
								doubleVec[3 * i + j] = vd.getPosition(i)[j];
						Discregrid::TriangleMesh sdfMesh(&doubleVec[0], faces.data(), vd.size(), nFaces);
#endif
						Discregrid::MeshDistance md(sdfMesh);
						Eigen::AlignedBox3d domain;
						for (auto const& x : sdfMesh.vertices())
						{
							domain.extend(x);
						}
						domain.max() += 0.1 * Eigen::Vector3d::Ones();
						domain.min() -= 0.1 * Eigen::Vector3d::Ones();

						LOG_INFO << "Set SDF resolution: " << rbd.m_resolutionSDF[0] << ", " << rbd.m_resolutionSDF[1] << ", " << rbd.m_resolutionSDF[2];
						distanceFields[sdfFileName] = std::make_shared<CubicSDFCollisionDetection::Grid>(domain, std::array<unsigned int, 3>({ rbd.m_resolutionSDF[0], rbd.m_resolutionSDF[1], rbd.m_resolutionSDF[2] }));
						auto func = Discregrid::DiscreteGrid::ContinuousFunction{};
						func = [&md](Eigen::Vector3d const& xi) {return md.signedDistanceCached(xi); };
						LOG_INFO << "Generate SDF for " << rbd.m_modelFile;
						distanceFields[sdfFileName]->addFunction(func, true);
						if (FileSystem::makeDir(cachePath) == 0)
						{
							LOG_INFO << "Save SDF: " << sdfFileName;
							distanceFields[sdfFileName]->save(sdfFileName);
							FileSystem::writeMD5File(rbd.m_modelFile, md5FileName);
						}
					}
				}
				else
				{
					std::string fileName = rbd.m_collisionObjectFileName;
					if (FileSystem::isRelativePath(fileName))
					{
						fileName = FileSystem::normalizePath(basePath + "/" + fileName);
					}
					LOG_INFO << "Load SDF: " << fileName;
					distanceFields[rbd.m_collisionObjectFileName] = std::make_shared<CubicSDFCollisionDetection::Grid>(fileName);
				}
			}
		}
	}

	for (unsigned int rbIndex = 0; rbIndex < data.m_tetModelData.size(); rbIndex++)
	{
		const SceneLoader::TetModelData &tmd = data.m_tetModelData[rbIndex];

		// Check if already loaded
		if ((tmd.m_modelFileVis != "") &&
			(objFiles.find(tmd.m_modelFileVis) == objFiles.end()))
		{
			IndexedFaceMesh mesh;
			VertexData vd;
			loadObj(FileSystem::normalizePath(tmd.m_modelFileVis), vd, mesh, Vector3r::Ones());
			objFiles[tmd.m_modelFileVis] = { vd, mesh };
		}
	}

	// Stiff rods: gather rigid body id of bodies belonging to a tree
	std::set<unsigned int> treeSegments;
	for (auto &tree : data.m_TreeModels)
	{
		treeSegments.insert(tree.m_SegmentIds.begin(), tree.m_SegmentIds.end());
	}
	// end stiff rods

	for (unsigned int i = 0; i < data.m_tetModelData.size(); i++)
	{
		const SceneLoader::TetModelData &tmd = data.m_tetModelData[i];

		// Check if already loaded
		if ((tmd.m_modelFileVis != "") &&
			(objFiles.find(tmd.m_modelFileVis) == objFiles.end()))
		{
			IndexedFaceMesh mesh;
			VertexData vd;
			loadObj(FileSystem::normalizePath(tmd.m_modelFileVis), vd, mesh, Vector3r::Ones());
			objFiles[tmd.m_modelFileVis] = { vd, mesh };
		}

		const std::string basePath = FileSystem::getFilePath(base->getSceneFile());
		const string cachePath = basePath + "/Cache";
		const string resStr = to_string(tmd.m_resolutionSDF[0]) + "_" + to_string(tmd.m_resolutionSDF[1]) + "_" + to_string(tmd.m_resolutionSDF[2]);
		const std::string modelFileName = FileSystem::getFileNameWithExt(tmd.m_modelFileVis);
		const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");

		std::string sdfKey = tmd.m_collisionObjectFileName;
		if (sdfKey == "")
		{
			sdfKey = sdfFileName;
		}
		if (distanceFields.find(sdfKey) == distanceFields.end())
		{
			// Generate SDF
			if (tmd.m_collisionObjectType == SceneLoader::SDF)
			{
				VertexData &vd = objFiles[tmd.m_modelFileVis].first;
				IndexedFaceMesh &mesh = objFiles[tmd.m_modelFileVis].second;

				CubicSDFCollisionDetection::GridPtr distanceField = generateSDF(tmd.m_modelFileVis, tmd.m_collisionObjectFileName, tmd.m_resolutionSDF, vd, mesh);

				if (tmd.m_collisionObjectFileName == "")
					distanceFields[sdfFileName] = distanceField;
				else
					distanceFields[tmd.m_collisionObjectFileName] = distanceField;
			}
		}
	}

	rb.resize(data.m_rigidBodyData.size());
	std::map<unsigned int, unsigned int> id_index;
	unsigned int rbIndex = 0;
	for (unsigned int idx(0); idx < data.m_rigidBodyData.size(); ++idx)
	{
		const SceneLoader::RigidBodyData &rbd = data.m_rigidBodyData[idx];

		// Stiff rods: do not handle segments of rods here
		if (treeSegments.find(rbd.m_id) != treeSegments.end())
			continue;
		// end stiff rods

		if (objFiles.find(rbd.m_modelFile) == objFiles.end())
			continue;

		id_index[rbd.m_id] = rbIndex;

		VertexData &vd = objFiles[rbd.m_modelFile].first;
		IndexedFaceMesh &mesh = objFiles[rbd.m_modelFile].second;

		rb[rbIndex] = new RigidBody();
		rb[rbIndex]->initBody(rbd.m_density,
			rbd.m_x,
			rbd.m_q,
			vd, mesh,
			rbd.m_scale);

		if (!rbd.m_isDynamic)
			rb[rbIndex]->setMass(0.0);
		else
		{
			rb[rbIndex]->setVelocity(rbd.m_v);
			rb[rbIndex]->setAngularVelocity(rbd.m_omega);
		}
		rb[rbIndex]->setRestitutionCoeff(rbd.m_restitutionCoeff);
		rb[rbIndex]->setFrictionCoeff(rbd.m_frictionCoeff);

		const std::vector<Vector3r> *vertices = rb[rbIndex]->getGeometry().getVertexDataLocal().getVertices();
		const unsigned int nVert = static_cast<unsigned int>(vertices->size());

		switch (rbd.m_collisionObjectType)
		{
		case SceneLoader::No_Collision_Object: break;
		case SceneLoader::Sphere:
			cd.addCollisionSphere(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale[0], rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::Box:
			cd.addCollisionBox(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale, rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::Cylinder:
			cd.addCollisionCylinder(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale.head<2>(), rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::Torus:
			cd.addCollisionTorus(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale.head<2>(), rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::HollowSphere:
			cd.addCollisionHollowSphere(rbIndex, PBD::CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale[0], rbd.m_thicknessSDF, rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::HollowBox:
			cd.addCollisionHollowBox(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, rbd.m_collisionObjectScale, rbd.m_thicknessSDF, rbd.m_testMesh, rbd.m_invertSDF);
			break;
		case SceneLoader::SDF:
		{
			if (rbd.m_collisionObjectFileName == "")
			{
				const std::string basePath = FileSystem::getFilePath(sceneFileName);
				const string cachePath = basePath + "/Cache";
				const string resStr = to_string(rbd.m_resolutionSDF[0]) + "_" + to_string(rbd.m_resolutionSDF[1]) + "_" + to_string(rbd.m_resolutionSDF[2]);
				const std::string modelFileName = FileSystem::getFileNameWithExt(rbd.m_modelFile);
				const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");
				cd.addCubicSDFCollisionObject(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, distanceFields[sdfFileName], rbd.m_collisionObjectScale, rbd.m_testMesh, rbd.m_invertSDF);
			}
			else
				cd.addCubicSDFCollisionObject(rbIndex, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices)[0], nVert, distanceFields[rbd.m_collisionObjectFileName], rbd.m_collisionObjectScale, rbd.m_testMesh, rbd.m_invertSDF);
			break;
		}
		}
		// stiff rods
		++rbIndex;
		// end stiff rods
	}

	//////////////////////////////////////////////////////////////////////////
	// stiff rod models
	//////////////////////////////////////////////////////////////////////////

	if (data.m_TreeModels.size())
		LOG_INFO << "Loading " << data.m_TreeModels.size() << " stiff rods" ;

	std::map<unsigned int, unsigned int> index_id;

	std::map<unsigned int, StiffRodsSceneLoader::StretchBendingTwistingConstraintData*> cIDTocData;
	for (auto &constraint : data.m_StretchBendingTwistingConstraints)
	{
		cIDTocData[constraint.m_Id] = &constraint;
	}

	unsigned int lengthIdx(1), widthIdx(0), heightIdx(2); // y-axis is used as the length axis to be compatible to Blender's armatures (Blender 2.78)
	std::map<unsigned int, Vector3r> segmentScale;
	std::map<unsigned int, Real> segmentDensity;

	for (auto &treeModel : data.m_TreeModels)
	{
		START_TIMING("Load tree");

		std::vector<RigidBody*> treeSegments;
		treeSegments.reserve(treeModel.m_SegmentIds.size());

		const Vector3r &modelPosition(treeModel.m_X);
		const Quaternionr &modelRotationQuat(treeModel.m_Q);
		Matrix3r modelRotationMat(modelRotationQuat.toRotationMatrix());

		std::cout << "Loading " << treeModel.m_SegmentIds.size() << " stiff rod segments" << std::endl;
		for (size_t idx(0); idx < data.m_rigidBodyData.size(); ++idx)
		{
			auto &segmentData(data.m_rigidBodyData[static_cast<unsigned int>(idx)]);
			if (treeModel.m_SegmentIds.find(segmentData.m_id) == treeModel.m_SegmentIds.end())
				continue;

			if (objFiles.find(segmentData.m_modelFile) == objFiles.end())
				continue;

			// continue if segment has already been created because it is also a member of another tree
			if (id_index.find(segmentData.m_id) != id_index.end())
				continue;

			VertexData &vd = objFiles[segmentData.m_modelFile].first;
			IndexedFaceMesh &mesh = objFiles[segmentData.m_modelFile].second;

			//compute mass and inertia of a cylinder
			Vector3r & scale(segmentData.m_scale);
			Real length(scale[lengthIdx]), width(scale[widthIdx]), height(scale[heightIdx]);
			Real density(segmentData.m_density);
			segmentDensity[segmentData.m_id] = density;

			Real mass(static_cast<Real>(M_PI_4) * width * height * length * density);

			const Real Iy = static_cast<Real>(1. / 12.) * mass*(static_cast<Real>(3.)*(static_cast<Real>(0.25)*height*height) + length*length);
			const Real Iz = static_cast<Real>(1. / 12.) * mass*(static_cast<Real>(3.)*(static_cast<Real>(0.25)*width*width) + length*length);
			const Real Ix = static_cast<Real>(0.25) * mass*(static_cast<Real>(0.25)*(height*height + width*width));

			Vector3r inertiaTensor;
			inertiaTensor(lengthIdx) = Ix;
			inertiaTensor(widthIdx) = Iy;
			inertiaTensor(heightIdx) = Iz;

			rb[rbIndex] = new RigidBody();
			rb[rbIndex]->initBody(mass,
				modelPosition + modelRotationMat * segmentData.m_x,
				inertiaTensor,
				modelRotationQuat * segmentData.m_q,
				vd, mesh, scale);

			if (!segmentData.m_isDynamic)
				rb[rbIndex]->setMass(0.0);
			else
			{
				rb[rbIndex]->setVelocity(segmentData.m_v);
				rb[rbIndex]->setAngularVelocity(segmentData.m_omega);
			}
			rb[rbIndex]->setRestitutionCoeff(segmentData.m_restitutionCoeff);
			rb[rbIndex]->setFrictionCoeff(segmentData.m_frictionCoeff);

			segmentScale[segmentData.m_id] = scale;

			id_index[segmentData.m_id] = rbIndex;
			index_id[rbIndex] = segmentData.m_id;
			treeSegments.push_back(rb[rbIndex]);

			++rbIndex;
		}

		// Make roots static;
		for (auto id : treeModel.m_StaticSegments)
		{
			std::map<unsigned int, unsigned int>::iterator it(id_index.find(id));
			if (it != id_index.end())
			{
				rb[it->second]->setMass(0.);
			}
		}

		// create zero-stretch, bending and twisting constraints

		std::vector<std::pair<unsigned int, unsigned int>>  constraintSegmentIndices;
		std::vector<Vector3r> constraintPositions;
		std::vector<Real> averageRadii;
		std::vector<Real> averageSegmentLengths;
		std::vector<Real> youngsModuli(treeModel.m_StretchBendingTwistingConstraintIds.size(), treeModel.m_YoungsModulus);
		std::vector<Real> torsionModuli(treeModel.m_StretchBendingTwistingConstraintIds.size(), treeModel.m_TorsionModulus);

		constraintSegmentIndices.reserve(treeModel.m_StretchBendingTwistingConstraintIds.size());
		constraintPositions.reserve(treeModel.m_StretchBendingTwistingConstraintIds.size());
		averageRadii.reserve(treeModel.m_StretchBendingTwistingConstraintIds.size());
		averageSegmentLengths.reserve(treeModel.m_StretchBendingTwistingConstraintIds.size());

		for (unsigned int cID : treeModel.m_StretchBendingTwistingConstraintIds)
		{
			if (cIDTocData.find(cID) == cIDTocData.end())
				continue;

			auto &sbtConstraint = *cIDTocData[cID];

			std::map<unsigned int, unsigned int>::iterator itFirstIdx(id_index.find(sbtConstraint.m_SegmentID[0])),
				itSecondIdx(id_index.find(sbtConstraint.m_SegmentID[1]));

			if (itFirstIdx == id_index.end() || itSecondIdx == id_index.end())
				continue;

			constraintSegmentIndices.push_back(std::pair<unsigned int, unsigned int>(itFirstIdx->second, itSecondIdx->second));

			const Vector3r &localPosition(sbtConstraint.m_Position);
			Vector3r x(modelPosition + modelRotationMat * localPosition);
			constraintPositions.push_back(x);

			// compute average length
			Real avgLength(static_cast<Real>(0.5)*(segmentScale[sbtConstraint.m_SegmentID[0]][lengthIdx] + segmentScale[sbtConstraint.m_SegmentID[1]][lengthIdx]));
			averageSegmentLengths.push_back(avgLength);

			// compute average radius
			averageRadii.push_back(static_cast<Real>(0.125)*(segmentScale[sbtConstraint.m_SegmentID[0]][widthIdx]
				+ segmentScale[sbtConstraint.m_SegmentID[0]][heightIdx]
				+ segmentScale[sbtConstraint.m_SegmentID[1]][widthIdx] + segmentScale[sbtConstraint.m_SegmentID[1]][heightIdx]));
		}

		model->addDirectPositionBasedSolverForStiffRodsConstraint(constraintSegmentIndices,
			constraintPositions, averageRadii, averageSegmentLengths, youngsModuli, torsionModuli);

		STOP_TIMING_PRINT;// load tree
	}

	//////////////////////////////////////////////////////////////////////////
	// triangle models
	//////////////////////////////////////////////////////////////////////////

	// map file names to loaded geometry to prevent multiple imports of same files
	std::map<std::string, pair<VertexData, IndexedFaceMesh>> triFiles;
	for (unsigned int i = 0; i < data.m_triangleModelData.size(); i++)
	{
		const SceneLoader::TriangleModelData &tmd = data.m_triangleModelData[i];

		// Check if already loaded
		if (triFiles.find(tmd.m_modelFile) == triFiles.end())
		{
			IndexedFaceMesh mesh;
			VertexData vd;
			loadObj(FileSystem::normalizePath(tmd.m_modelFile), vd, mesh, Vector3r::Ones());
			triFiles[tmd.m_modelFile] = { vd, mesh };
		}
	}

	triModels.reserve(data.m_triangleModelData.size());
	std::map<unsigned int, unsigned int> tm_id_index;
	for (unsigned int i = 0; i < data.m_triangleModelData.size(); i++)
	{
		const SceneLoader::TriangleModelData &tmd = data.m_triangleModelData[i];

		if (triFiles.find(tmd.m_modelFile) == triFiles.end())
			continue;

		tm_id_index[tmd.m_id] = i;

		VertexData vd = triFiles[tmd.m_modelFile].first;
		IndexedFaceMesh &mesh = triFiles[tmd.m_modelFile].second;

		const Matrix3r R = tmd.m_q.matrix();
		for (unsigned int j = 0; j < vd.size(); j++)
		{
			vd.getPosition(j) = R * (vd.getPosition(j).cwiseProduct(tmd.m_scale)) + tmd.m_x;
		}

		model->addTriangleModel(vd.size(), mesh.numFaces(), &vd.getPosition(0), mesh.getFaces().data(), mesh.getUVIndices(), mesh.getUVs());

		TriangleModel *tm = triModels[triModels.size() - 1];
		ParticleData &pd = model->getParticles();
		unsigned int offset = tm->getIndexOffset();

		for (unsigned int j = 0; j < tmd.m_staticParticles.size(); j++)
		{
			const unsigned int index = tmd.m_staticParticles[j] + offset;
			pd.setMass(index, 0.0);
		}

		tm->setRestitutionCoeff(tmd.m_restitutionCoeff);
		tm->setFrictionCoeff(tmd.m_frictionCoeff);
	}

	initTriangleModelConstraints();

	//////////////////////////////////////////////////////////////////////////
	// tet models
	//////////////////////////////////////////////////////////////////////////

	// map file names to loaded geometry to prevent multiple imports of same files
	std::map<pair<string, string>, pair<vector<Vector3r>, vector<unsigned int>>> tetFiles;
	for (unsigned int i = 0; i < data.m_tetModelData.size(); i++)
	{
		const SceneLoader::TetModelData &tmd = data.m_tetModelData[i];

		// Check if already loaded
		pair<string, string> fileNames = { tmd.m_modelFileNodes, tmd.m_modelFileElements };
		if (tetFiles.find(fileNames) == tetFiles.end())
		{
			vector<Vector3r> vertices;
			vector<unsigned int> tets;
			TetGenLoader::loadTetgenModel(FileSystem::normalizePath(tmd.m_modelFileNodes), FileSystem::normalizePath(tmd.m_modelFileElements), vertices, tets);
			tetFiles[fileNames] = { vertices, tets };
		}
	}

	tetModels.reserve(data.m_tetModelData.size());
	std::map<unsigned int, unsigned int> tm_id_index2;
	for (unsigned int i = 0; i < data.m_tetModelData.size(); i++)
	{
		const SceneLoader::TetModelData &tmd = data.m_tetModelData[i];

		pair<string, string> fileNames = { tmd.m_modelFileNodes, tmd.m_modelFileElements };
		auto geo = tetFiles.find(fileNames);
		if (geo == tetFiles.end())
			continue;

		tm_id_index2[tmd.m_id] = i;

		vector<Vector3r> vertices = geo->second.first;
		vector<unsigned int> &tets = geo->second.second;

		const Matrix3r R = tmd.m_q.matrix();
		for (unsigned int j = 0; j < vertices.size(); j++)
		{
			vertices[j] = R * (vertices[j].cwiseProduct(tmd.m_scale)) + tmd.m_x;
		}

		model->addTetModel((unsigned int)vertices.size(), (unsigned int)tets.size() / 4, vertices.data(), tets.data());

		TetModel *tm = tetModels[tetModels.size() - 1];
		ParticleData &pd = model->getParticles();
		unsigned int offset = tm->getIndexOffset();

		tm->setInitialX(tmd.m_x);
		tm->setInitialR(R);
		tm->setInitialScale(tmd.m_scale);

		for (unsigned int j = 0; j < tmd.m_staticParticles.size(); j++)
		{
			const unsigned int index = tmd.m_staticParticles[j] + offset;
			pd.setMass(index, 0.0);
		}

		// read visualization mesh
		if (tmd.m_modelFileVis != "")
		{
			if (objFiles.find(tmd.m_modelFileVis) != objFiles.end())
			{
				IndexedFaceMesh &visMesh = tm->getVisMesh();
				VertexData &vdVis = tm->getVisVertices();
				vdVis = objFiles[tmd.m_modelFileVis].first;
				visMesh = objFiles[tmd.m_modelFileVis].second;

				for (unsigned int j = 0; j < vdVis.size(); j++)
					vdVis.getPosition(j) = R * (vdVis.getPosition(j).cwiseProduct(tmd.m_scale)) + tmd.m_x;

				tm->updateMeshNormals(pd);
				tm->attachVisMesh(pd);
				tm->updateVisMesh(pd);
			}
		}

		tm->setRestitutionCoeff(tmd.m_restitutionCoeff);
		tm->setFrictionCoeff(tmd.m_frictionCoeff);

		tm->updateMeshNormals(pd);
	}

	initTetModelConstraints();

	// init collision objects for deformable models
	ParticleData &pd = model->getParticles();
	for (unsigned int i = 0; i < data.m_triangleModelData.size(); i++)
	{
		TriangleModel *tm = triModels[i];
		unsigned int offset = tm->getIndexOffset();
		const unsigned int nVert = tm->getParticleMesh().numVertices();
		cd.addCollisionObjectWithoutGeometry(i, CollisionDetection::CollisionObject::TriangleModelCollisionObjectType, &pd.getPosition(offset), nVert, true);

	}
	for (unsigned int i = 0; i < data.m_tetModelData.size(); i++)
	{
		TetModel *tm = tetModels[i];
		unsigned int offset = tm->getIndexOffset();
		const unsigned int nVert = tm->getParticleMesh().numVertices();
		const IndexedTetMesh &tetMesh = tm->getParticleMesh();

		const SceneLoader::TetModelData &tmd = data.m_tetModelData[i];

		switch (tmd.m_collisionObjectType)
		{
		case SceneLoader::No_Collision_Object:
			cd.addCollisionObjectWithoutGeometry(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, true);
			break;
		case SceneLoader::Sphere:
			cd.addCollisionSphere(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale[0], tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::Box:
			cd.addCollisionBox(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale, tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::Cylinder:
			cd.addCollisionCylinder(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale.head<2>(), tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::Torus:
			cd.addCollisionTorus(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale.head<2>(), tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::HollowSphere:
			cd.addCollisionHollowSphere(i, PBD::CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale[0], tmd.m_thicknessSDF, tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::HollowBox:
			cd.addCollisionHollowBox(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, tmd.m_collisionObjectScale, tmd.m_thicknessSDF, tmd.m_testMesh, tmd.m_invertSDF);
			break;
		case SceneLoader::SDF:
		{
			if (tmd.m_collisionObjectFileName == "")
			{
				const std::string basePath = FileSystem::getFilePath(base->getSceneFile());
				const string cachePath = basePath + "/Cache";
				const string resStr = to_string(tmd.m_resolutionSDF[0]) + "_" + to_string(tmd.m_resolutionSDF[1]) + "_" + to_string(tmd.m_resolutionSDF[2]);
				const std::string modelFileName = FileSystem::getFileNameWithExt(tmd.m_modelFileVis);
				const string sdfFileName = FileSystem::normalizePath(cachePath + "/" + modelFileName + "_" + resStr + ".csdf");
				cd.addCubicSDFCollisionObject(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, distanceFields[sdfFileName], tmd.m_collisionObjectScale, tmd.m_testMesh, tmd.m_invertSDF);
			}
			else
				cd.addCubicSDFCollisionObject(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, distanceFields[tmd.m_collisionObjectFileName], tmd.m_collisionObjectScale, tmd.m_testMesh, tmd.m_invertSDF);
			break;
		}
		}
	}

	//////////////////////////////////////////////////////////////////////////
	// joints
	//////////////////////////////////////////////////////////////////////////

	for (unsigned int i = 0; i < data.m_ballJointData.size(); i++)
	{
		const SceneLoader::BallJointData &jd = data.m_ballJointData[i];
		model->addBallJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position);
	}

	for (unsigned int i = 0; i < data.m_ballOnLineJointData.size(); i++)
	{
		const SceneLoader::BallOnLineJointData &jd = data.m_ballOnLineJointData[i];
		model->addBallOnLineJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position, jd.m_axis);
	}

	for (unsigned int i = 0; i < data.m_hingeJointData.size(); i++)
	{
		const SceneLoader::HingeJointData &jd = data.m_hingeJointData[i];
		model->addHingeJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position, jd.m_axis);
	}

	for (unsigned int i = 0; i < data.m_universalJointData.size(); i++)
	{
		const SceneLoader::UniversalJointData &jd = data.m_universalJointData[i];
		model->addUniversalJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position, jd.m_axis[0], jd.m_axis[1]);
	}

	for (unsigned int i = 0; i < data.m_sliderJointData.size(); i++)
	{
		const SceneLoader::SliderJointData &jd = data.m_sliderJointData[i];
		model->addSliderJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_axis);
	}

	for (unsigned int i = 0; i < data.m_rigidBodyParticleBallJointData.size(); i++)
	{
		const SceneLoader::RigidBodyParticleBallJointData &jd = data.m_rigidBodyParticleBallJointData[i];
		model->addRigidBodyParticleBallJoint(id_index[jd.m_bodyID[0]], jd.m_bodyID[1]);
	}

	for (unsigned int i = 0; i < data.m_targetAngleMotorHingeJointData.size(); i++)
	{
		const SceneLoader::TargetAngleMotorHingeJointData &jd = data.m_targetAngleMotorHingeJointData[i];
		model->addTargetAngleMotorHingeJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position, jd.m_axis);
		((MotorJoint*)constraints[constraints.size() - 1])->setTarget(jd.m_target);
		((MotorJoint*)constraints[constraints.size() - 1])->setTargetSequence(jd.m_targetSequence);
		((MotorJoint*)constraints[constraints.size() - 1])->setRepeatSequence(jd.m_repeat);
	}

	for (unsigned int i = 0; i < data.m_targetVelocityMotorHingeJointData.size(); i++)
	{
		const SceneLoader::TargetVelocityMotorHingeJointData &jd = data.m_targetVelocityMotorHingeJointData[i];
		model->addTargetVelocityMotorHingeJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position, jd.m_axis);
		((MotorJoint*)constraints[constraints.size() - 1])->setTarget(jd.m_target);
		((MotorJoint*)constraints[constraints.size() - 1])->setTargetSequence(jd.m_targetSequence);
		((MotorJoint*)constraints[constraints.size() - 1])->setRepeatSequence(jd.m_repeat);
	}

	for (unsigned int i = 0; i < data.m_targetPositionMotorSliderJointData.size(); i++)
	{
		const SceneLoader::TargetPositionMotorSliderJointData &jd = data.m_targetPositionMotorSliderJointData[i];
		model->addTargetPositionMotorSliderJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_axis);
		((MotorJoint*)constraints[constraints.size() - 1])->setTarget(jd.m_target);
		((MotorJoint*)constraints[constraints.size() - 1])->setTargetSequence(jd.m_targetSequence);
		((MotorJoint*)constraints[constraints.size() - 1])->setRepeatSequence(jd.m_repeat);
	}

	for (unsigned int i = 0; i < data.m_targetVelocityMotorSliderJointData.size(); i++)
	{
		const SceneLoader::TargetVelocityMotorSliderJointData &jd = data.m_targetVelocityMotorSliderJointData[i];
		model->addTargetVelocityMotorSliderJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_axis);
		((MotorJoint*)constraints[constraints.size() - 1])->setTarget(jd.m_target);
		((MotorJoint*)constraints[constraints.size() - 1])->setTargetSequence(jd.m_targetSequence);
		((MotorJoint*)constraints[constraints.size() - 1])->setRepeatSequence(jd.m_repeat);
	}

	for (unsigned int i = 0; i < data.m_damperJointData.size(); i++)
	{
		const SceneLoader::DamperJointData &jd = data.m_damperJointData[i];
		model->addDamperJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_axis, jd.m_stiffness);
	}

	for (unsigned int i = 0; i < data.m_rigidBodySpringData.size(); i++)
	{
		const SceneLoader::RigidBodySpringData &jd = data.m_rigidBodySpringData[i];
		model->addRigidBodySpring(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position1, jd.m_position2, jd.m_stiffness);
	}

	for (unsigned int i = 0; i < data.m_distanceJointData.size(); i++)
	{
		const SceneLoader::DistanceJointData &jd = data.m_distanceJointData[i];
		model->addDistanceJoint(id_index[jd.m_bodyID[0]], id_index[jd.m_bodyID[1]], jd.m_position1, jd.m_position2);
	}
}

void TW_CALL setContactTolerance(const void *value, void *clientData)
{
	const Real val = *(const Real *)(value);
	((DistanceFieldCollisionDetection*)clientData)->setTolerance(val);
}

void TW_CALL getContactTolerance(void *value, void *clientData)
{
	*(Real *)(value) = ((DistanceFieldCollisionDetection*)clientData)->getTolerance();
}

void TW_CALL setContactStiffnessRigidBody(const void *value, void *clientData)
{
	const Real val = *(const Real *)(value);
	((SimulationModel*)clientData)->setContactStiffnessRigidBody(val);
}

void TW_CALL getContactStiffnessRigidBody(void *value, void *clientData)
{
	*(Real *)(value) = ((SimulationModel*)clientData)->getContactStiffnessRigidBody();
}

void TW_CALL setContactStiffnessParticleRigidBody(const void *value, void *clientData)
{
	const Real val = *(const Real *)(value);
	((SimulationModel*)clientData)->setContactStiffnessParticleRigidBody(val);
}

void TW_CALL getContactStiffnessParticleRigidBody(void *value, void *clientData)
{
	*(Real *)(value) = ((SimulationModel*)clientData)->getContactStiffnessParticleRigidBody();
}

void TW_CALL setBendingMethod(const void *value, void *clientData)
{
	const short val = *(const short *)(value);
	*((short*)clientData) = val;
	reset();
}

void TW_CALL getBendingMethod(void *value, void *clientData)
{
	*(short *)(value) = *((short*)clientData);
}

void TW_CALL setClothSimulationMethod(const void *value, void *clientData)
{
	const short val = *(const short *)(value);
	*((short*)clientData) = val;
	reset();
}

void TW_CALL getClothSimulationMethod(void *value, void *clientData)
{
	*(short *)(value) = *((short*)clientData);
}

void TW_CALL setSolidSimulationMethod(const void *value, void *clientData)
{
	const short val = *(const short *)(value);
	*((short*)clientData) = val;
	reset();
}

void TW_CALL getSolidSimulationMethod(void *value, void *clientData)
{
	*(short *)(value) = *((short*)clientData);
}


#define BOOST_TEST_MODULE "Test Groups class"

#include <boost/test/unit_test.hpp>

#include "distance/Distance.hpp"
#include "Exception.hpp"
#include "Group.hpp"

#define TOLERANCE 1e-9

using namespace konex;

struct MockData
{
  data_t dat_1[10] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

  data_t dat_2[10] = {1, 1, 2, 2, 3, 3, 4, 4, 5, 5};

  data_t centroid_1[5] = {1, 1.5, 2.5, 3, 4};
  data_t centroid_2[10] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.5};

  data_t dat_3[10] = {2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
  data_t dat_4[10] = {10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  data_t dat_5[10] = {9, 9, 8, 7, 6, 5, 4, 3, 2, 1};

  data_t dat_sim_1[5] = {1, 2, 3, 4, 5};
  data_t dat_sim_2[5] = {2, 3, 4, 5, 6};
  data_t dat_diff[5]  = {12, 13, 14, 15};

  std::string test_5_10_space = "datasets/test/test_5_10_space.txt";
  std::string test_3_10_space = "datasets/test/test_3_10_space.txt";
};

BOOST_AUTO_TEST_CASE( basic_groups, *boost::unit_test::tolerance(TOLERANCE) )
{
  MockData data;

  int timeSeriesCount = 5;
  int timeSeriesLengths = 10;
  int memberLength = 5;
  int subTimeSeriesCount = timeSeriesLengths - memberLength + 1;

  TimeSeriesSet tsSet;
  tsSet.loadData(data.test_5_10_space, timeSeriesCount, 0, " ");

  BOOST_CHECK_EQUAL( tsSet.getItemLength(), timeSeriesLengths );
  BOOST_CHECK_EQUAL( tsSet.getItemCount(), timeSeriesCount );
  BOOST_CHECK( tsSet.getFilePath() == data.test_5_10_space );

  std::vector<group_membership_t> memberMap =
    std::vector<group_membership_t>(tsSet.getItemCount() * tsSet.getItemLength());
  Group g(0, memberLength, subTimeSeriesCount, tsSet, memberMap);

  const TimeSeries& c = g.getCentroid();

  // test initial centroid is all zeros
  for (int i = 0; i < memberLength; i++)
  {
    BOOST_TEST( c[i] == 0 );
  }

  g.addMember(0, 0);
  g.setCentroid(0, 0);

  //checking if centroid is updated
  BOOST_CHECK_EQUAL( g.getCount(), 1 );
  for (int i = 0; i < memberLength; i++)
  {
    BOOST_TEST( c[i] == data.dat_1[i] );
  }

  g.addMember(1, 0);

  BOOST_CHECK_EQUAL( g.getCount(), 2 );
}

BOOST_AUTO_TEST_CASE( group_get_best_match, *boost::unit_test::tolerance(TOLERANCE) )
{
  MockData data;
  dist_t distance = warpedDistance;

  int timeSeriesCount = 3;
  int timeSeriesLengths = 10;
  int memberLength = 10;
  int subTimeSeriesCount = timeSeriesLengths - memberLength + 1;

  TimeSeriesSet tsSet;
  tsSet.loadData(data.test_3_10_space, timeSeriesCount, 0, " ");

  std::vector<group_membership_t> memberMap =
    std::vector<group_membership_t>(tsSet.getItemCount() * tsSet.getItemLength());

  Group g(0, memberLength, subTimeSeriesCount, tsSet, memberMap);
  g.addMember(2, 0);
  g.addMember(0, 0);
  TimeSeries t = tsSet.getTimeSeries(1,0,memberLength);
  BOOST_TEST(t[0] == 1.0);
  candidate_time_series_t best = g.getBestMatch(t, distance);
  BOOST_TEST(best.dist == sqrt(1.0)/(2 * 10.0));
}