/******************************************************************************
 * Copyright 2017 Baidu Robotic Vision Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
#include "feature_utils.h"
#include <Eigen/Dense>

#define SUBPIX_VERBOSE 0

namespace XP {
namespace align {

using Eigen::Vector2f;
using Eigen::Vector3f;
using Eigen::Matrix2f;
using Eigen::Matrix3f;
using Eigen::Matrix3i;
const int W_BITS = 14;
const int half_patch_size = 4;
const int patch_size = 8;
const int patch_area = 64;
bool align1D(const cv::Mat& cur_img,
             const Vector2f& dir,  // direction in which the patch is allowed to move
             uint8_t* ref_patch_with_border,
             uint8_t* ref_patch,
             const int n_iter,
             Vector2f* cur_px_estimate,
             float* h_inv) {
  bool converged = false;

  // compute derivative of template and prepare inverse compositional
  float __attribute__((__aligned__(16))) ref_patch_dv[patch_area];
  Matrix2f H; H.setZero();

  // compute gradient and hessian
  const int ref_step = patch_size + 2;
  float* it_dv = ref_patch_dv;
  for (int y=0; y < patch_size; ++y) {
    uint8_t* it = ref_patch_with_border + (y + 1) * ref_step + 1;
    for (int x = 0; x < patch_size; ++x, ++it, ++it_dv) {
      Vector2f J;
      J[0] = 0.5*(dir[0] * (it[1] - it[-1]) + dir[1] * (it[ref_step] - it[-ref_step]));
      J[1] = 1;
      *it_dv = J[0];
      H += J * J.transpose();
    }
  }
  *h_inv = 1.0 / H(0, 0) * patch_size * patch_size;
  Matrix2f Hinv = H.inverse();
  float mean_diff = 0;

  // Compute pixel location in new image:
  float u = (*cur_px_estimate)(0);
  float v = (*cur_px_estimate)(1);

  // termination condition
  const float min_update_squared = 0.03 * 0.03;
  const int cur_step = cur_img.step.p[0];
  float chi2 = 0;
  Vector2f update; update.setZero();
  for (int iter = 0; iter < n_iter; ++iter) {
    int u_r = floor(u);
    int v_r = floor(v);
    if (u_r < half_patch_size || v_r < half_patch_size ||
        u_r >= cur_img.cols-half_patch_size || v_r >= cur_img.rows - half_patch_size) {
      break;
    }
    if (std::isnan(u) || std::isnan(v)) {
      // TODO(SVO): very rarely this can happen, maybe H is singular? should not be at corner.
      return false;
    }

    // compute interpolation weights
    float subpix_x = u - u_r;
    float subpix_y = v - v_r;
    float wTL = (1.0 - subpix_x) * (1.0 - subpix_y);
    float wTR = subpix_x * (1.0 - subpix_y);
    float wBL = (1.0 - subpix_x) * subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interpolate
    uint8_t* it_ref = ref_patch;
    float* it_ref_dv = ref_patch_dv;
    float new_chi2 = 0.0;
    Vector2f Jres; Jres.setZero();
    for (int y = 0; y < patch_size; ++y) {
      uint8_t* it = static_cast<uint8_t*>(cur_img.data) +
          (v_r + y - half_patch_size) * cur_step + u_r - half_patch_size;
      for (int x = 0; x < patch_size; ++x, ++it, ++it_ref, ++it_ref_dv) {
        float search_pixel = wTL * it[0] + wTR * it[1] + wBL * it[cur_step] + wBR * it[cur_step+1];
        float res = search_pixel - *it_ref + mean_diff;
        Jres[0] -= res * (*it_ref_dv);
        Jres[1] -= res;
        new_chi2 += res * res;
      }
    }

    if (iter > 0 && new_chi2 > chi2) {
#if SUBPIX_VERBOSE
      cout << "error increased." << endl;
#endif
      u -= update[0];
      v -= update[1];
      break;
    }

    chi2 = new_chi2;
    update = Hinv * Jres;
    u += update[0] * dir[0];
    v += update[0] * dir[1];
    mean_diff += update[1];

#if SUBPIX_VERBOSE
    VLOG(2) << "Iter " << iter << ":"
            << "\t u=" << u << ", v=" << v
            << "\t update = " << update[0] << ", " << update[1]
            << "\t new chi2 = " << new_chi2;
#endif

    if (update[0] * update[0] + update[1] * update[1] < min_update_squared) {
#if SUBPIX_VERBOSE
      VLOG(2) << "converged.";
#endif
      converged = true;
      break;
    }
  }

  *cur_px_estimate << u, v;
  return converged;
}

 //逆向组合法进行2D图像对齐,这个好像是在LK十年那个论文里有说过
 //输入右相机的图像,从左相机变换到右相机上的patch_border,从左相机变换到右相机上的patch,最大迭代次数
 //输出最终块匹配残差最小的右目中特征点的像素坐标
 //残差: r = I(p_cur) - I(p_ref) + m ,其中I(*)表示在某个像素位置的光度, m是均值差.cur是右相机,ref是左相机
 //优化变量是p_ref,m，这里为什么是p_ref而不是p_cur,是因为如果用p_cur，每次增量以后还需要再计算一次这个的J,用p_ref只用算一次,均值差也做为优化变量是防止噪声
bool align2D(const cv::Mat& cur_img,
             uint8_t* ref_patch_with_border,
             uint8_t* ref_patch,
             const int n_iter,
             Vector2f* cur_px_estimate,
             bool no_simd)
             {
/*
#ifdef __ARM_NEON__
  Vector2f cur_px_estimate_neon = *cur_px_estimate;
  if (!no_simd) {
    align2D_NEON(cur_img, ref_patch_with_border, ref_patch, n_iter, &cur_px_estimate_neon);
  }
#endif
*/
  bool converged = false;
// 模板图像对像素 x,y 坐标进行求导
  // compute derivative of template and prepare inverse compositional
  float __attribute__((__aligned__(16))) ref_patch_dx[patch_area];
  float __attribute__((__aligned__(16))) ref_patch_dy[patch_area];
  Matrix3f H; H.setZero();

  // compute gradient and hessian
  const int ref_step = patch_size+2;//patch框边长
     // 图像导数的指针
  float* it_dx = ref_patch_dx;
  float* it_dy = ref_patch_dy;
  //构造H矩阵
  for (int y = 0; y < patch_size; ++y) {
    uint8_t* it = ref_patch_with_border + (y + 1) * ref_step + 1;
    for (int x = 0; x < patch_size; ++x, ++it, ++it_dx, ++it_dy) {
      Vector3f J;
      //p_ref像素值关于位置，其实就是梯度
      J[0] = 0.5 * (it[1] - it[-1]);
      J[1] = 0.5 * (it[ref_step] - it[-ref_step]);
      J[2] = 1;//均值差的导数
      *it_dx = J[0];
      *it_dy = J[1];
      H += J*J.transpose();
    }
  }
  Matrix3f Hinv = H.inverse();
  float mean_diff = 0;

  // Compute pixel location in new image:
     //初值
  float u = (*cur_px_estimate)(0);
  float v = (*cur_px_estimate)(1);

  // termination condition
  const float min_update_squared = 0.03 * 0.03;//收敛条件
  const int cur_step = cur_img.step.p[0];
  float chi2 = std::numeric_limits<int>::max();
  Vector3f update; update.setZero();//优化变量的增量
  for (int iter = 0; iter < n_iter; ++iter) //优化开始
  {
    int u_r = floor(u);
    int v_r = floor(v);
    //边缘处跳过
    if (u_r < half_patch_size || v_r < half_patch_size ||
        u_r >= cur_img.cols-half_patch_size || v_r >= cur_img.rows-half_patch_size) {
      break;
    }
    if (std::isnan(u) || std::isnan(v)) {
      // TODO(SVO): very rarely this can happen, maybe H is singular? should not be at corner
      return false;
    }

    // compute interpolation weights
    //双线性插值I(i+u,j+v) = (1-u)(1-v)I(i,j) + (1-u)vI(i,j+1) + u(1-v)I(i+1,j) + uvI(i+1,j+1) 以算出这个特征点在右目图像中的像素值
    float subpix_x = u - u_r;
    float subpix_y = v - v_r;
    float wTL = (1.0 - subpix_x) * (1.0 - subpix_y);
    float wTR = subpix_x * (1.0 - subpix_y);
    float wBL = (1.0 - subpix_x) * subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interpolate
    uint8_t* it_ref = ref_patch;// 左图像 patch指针
    float* it_ref_dx = ref_patch_dx; // 左图像 patch x方向导数指针
    float* it_ref_dy = ref_patch_dy;// 左图像 patch y方向导数指针
    float new_chi2 = 0.0;
    Vector3f Jres; Jres.setZero();
    for (int y = 0; y < patch_size; ++y)
    {
      uint8_t* it = static_cast<uint8_t*>(cur_img.data) +
          (v_r + y - half_patch_size) * cur_step + u_r - half_patch_size;
      for (int x = 0; x < patch_size; ++x, ++it, ++it_ref, ++it_ref_dx, ++it_ref_dy)
      {
          //残差r = I(p_cur) - I(p_ref) + m
        float search_pixel = wTL*it[0] + wTR*it[1] + wBL*it[cur_step] + wBR*it[cur_step+1];//双线性插值得到curimg中的像素值
        float res = search_pixel - *it_ref + mean_diff;
        //计算 H * deltax = -Jr中的Jres
        Jres[0] -= res*(*it_ref_dx);
        Jres[1] -= res*(*it_ref_dy);
        Jres[2] -= res;
        new_chi2 += res*res;// 卡方
      }
    }

    /*
    if(iter > 0 && new_chi2 > chi2) {
#if SUBPIX_VERBOSE
      VLOG(2) << "error increased.";
#endif
      u -= update[0];
      v -= update[1];
      break;
    }
    */
    chi2 = new_chi2;

    //更新变量
    //deltax = H.inv * -Jr
    update = Hinv * Jres;
    u += update[0];
    v += update[1];
    mean_diff += update[2];

#if SUBPIX_VERBOSE
    VLOG(2) << "Iter " << iter << ":"
            << "\t u=" << u << ", v=" << v
            << "\t update = " << update[0] << ", " << update[1]
            << "\t chi2 = " << chi2;
#endif

    if (update[0] * update[0] + update[1] * update[1] < min_update_squared)
    {
        //满足收敛条件
#if SUBPIX_VERBOSE
      VLOG(2) << "converged.";
#endif
      converged = true;
      break;
    }
  }
  *cur_px_estimate << u, v;
  return converged;
}

// TODO(mingyu): Add the SSE and NEON version
bool align2D_SSE2(const cv::Mat& cur_img,
                  uint8_t* ref_patch_with_border,
                  uint8_t* ref_patch,
                  const int n_iter,
                  Vector2f* cur_px_estimate) {
  LOG(FATAL) << "align2D_SSE2 is not implemented";
  return false;
}
#ifdef __ARM_NEON__
bool align2D_NEON(const cv::Mat& cur_img,
                  const uint8_t* ref_patch_with_border,
                  const uint8_t* ref_patch,
                  const int n_iter,
                  Vector2f* cur_px_estimate) {
  bool converged = false;

  // compute derivative of template and prepare inverse compositional
  /*************************************
   * memory layout
   * dx dy 1 dx dy 1 dx dy 1 ...
   */
  float __attribute__((__aligned__(16))) ref_patch_dx_dy[patch_area * 3];
  Matrix3f H = Matrix3f::Zero();

  // compute gradient
  const int ref_step = patch_size + 2;
  for (int n = 0; n < patch_size; ++n) {
    const uint8_t* it = ref_patch_with_border + (n + 1) * ref_step + 1;
    int16x8_t horizontal_diff, vertical_diff;
    {
      int16x8_t raw_data_mid_left = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it - 1)));
      int16x8_t raw_data_mid_right = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it + 1)));
      int16x8_t raw_data_top = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it - ref_step)));
      int16x8_t raw_data_bot = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it + ref_step)));
      horizontal_diff = vsubq_s16(raw_data_mid_right, raw_data_mid_left);
      vertical_diff = vsubq_s16(raw_data_bot, raw_data_top);
    }
    float32x4x3_t left4 = {
      vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(horizontal_diff))), 0.5),
      vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(vertical_diff))), 0.5),
      vdupq_n_f32(1.f)
    };
    vst3q_f32(ref_patch_dx_dy + n * 24, left4);
    float32x4x3_t right4 = {
      vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(horizontal_diff))), 0.5),
      vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(vertical_diff))), 0.5),
      vdupq_n_f32(1.f)
    };
    vst3q_f32(ref_patch_dx_dy + n * 24 + 12, right4);
  }
  // compute hessian
  float* ptr_dxy = ref_patch_dx_dy;
  for (int cnt = patch_size * patch_size; cnt > 0; --cnt) {
    Eigen::Map<Vector3f> J(ptr_dxy);
    H += J * J.transpose();
    ptr_dxy += 3;
  }

  Matrix3f Hinv = H.inverse();
  float mean_diff = 0.f;

  // Compute pixel location in new image:
  float u = (*cur_px_estimate)(0);
  float v = (*cur_px_estimate)(1);

  // termination condition
  const float min_update_squared = 0.03 * 0.03;
  const int cur_step = cur_img.step.p[0];
  float chi2 = std::numeric_limits<int>::max();
  Vector3f update = Vector3f::Zero();
  for (int iter = 0; iter < n_iter; ++iter) {
    int u_r = floor(u);
    int v_r = floor(v);
    if (u_r < half_patch_size ||
        v_r < half_patch_size ||
        u_r >= cur_img.cols - half_patch_size ||
        v_r >= cur_img.rows - half_patch_size) {
      break;
    }
    if (std::isnan(u) || std::isnan(v)) {
      // TODO(SVO): very rarely this can happen, maybe H is singular? should not be at corner
      return false;
    }

    // compute interpolation weights
    float subpix_x = u - u_r;
    float subpix_y = v - v_r;

    int iw00 = cvRound((1.f - subpix_x) * (1.f - subpix_y) * (1 << W_BITS));
    int iw01 = cvRound(subpix_x * (1.f - subpix_y) * (1 << W_BITS));
    int iw10 = cvRound((1.f - subpix_x) * subpix_y * (1 << W_BITS));
    int iw11 = (1 << W_BITS) - iw00 - iw01 - iw10;
    const int16x4_t viw00 = vdup_n_s16(static_cast<int16_t>(iw00));
    const int16x4_t viw01 = vdup_n_s16(static_cast<int16_t>(iw01));
    const int16x4_t viw10 = vdup_n_s16(static_cast<int16_t>(iw10));
    const int16x4_t viw11 = vdup_n_s16(static_cast<int16_t>(iw11));
    const int32x4_t shift = vdupq_n_s32(-W_BITS);
    const float32x4_t vmean_diff = vdupq_n_f32(mean_diff);

    // loop through search_patch, interpolate
    const uint8_t* it_ref = ref_patch;
    float new_chi2 = 0.f;
    Vector3f Jres = Vector3f::Zero();
    for (int n = 0; n < patch_size; ++n, it_ref += 8) {
      uint8_t* it = static_cast<uint8_t*>(cur_img.data) +
          (v_r + n - half_patch_size) * cur_step + u_r - half_patch_size;
      int16x8_t topleft = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it)));
      int16x8_t topright = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it + 1)));
      int16x8_t botleft = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it + cur_step)));
      int16x8_t botright = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it + cur_step + 1)));

      int32x4_t left_half1 = vaddq_s32(
                  vmull_s16(vget_low_s16(topleft), viw00),
                  vmull_s16(vget_low_s16(topright), viw01));
      int32x4_t left_half2 = vaddq_s32(
                  vmull_s16(vget_low_s16(botleft), viw10),
                  vmull_s16(vget_low_s16(botright), viw11));

      int32x4_t right_half1 = vaddq_s32(
                  vmull_s16(vget_high_s16(topleft), viw00),
                  vmull_s16(vget_high_s16(topright), viw01));
      int32x4_t right_half2 = vaddq_s32(
                  vmull_s16(vget_high_s16(botleft), viw10),
                  vmull_s16(vget_high_s16(botright), viw11));

      int16x8_t v_it_ref8 = vreinterpretq_s16_u16(vmovl_u8(vld1_u8(it_ref)));
      float32x4x2_t search_pixel = {
        vcvtq_f32_s32(vqrshlq_s32(vaddq_s32(left_half1, left_half2), shift)),
        vcvtq_f32_s32(vqrshlq_s32(vaddq_s32(right_half1, right_half2), shift))
      };

      search_pixel.val[0] = vsubq_f32(search_pixel.val[0],
                           vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_it_ref8))));
      search_pixel.val[1] = vsubq_f32(search_pixel.val[1],
                           vcvtq_f32_s32(vmovl_s16(vget_high_s16(v_it_ref8))));
      search_pixel.val[0] = vaddq_f32(search_pixel.val[0], vmean_diff);
      search_pixel.val[1] = vaddq_f32(search_pixel.val[1], vmean_diff);

      float32x4x3_t dxy_left = vld3q_f32(ref_patch_dx_dy + n * 24);
      float32x4x3_t dxy_right = vld3q_f32(ref_patch_dx_dy + n * 24 + 12);

      // dx
      dxy_left.val[0]  = vmulq_f32(dxy_left.val[0], search_pixel.val[0]);
      dxy_right.val[0] = vmulq_f32(dxy_right.val[0], search_pixel.val[1]);
      // dy
      dxy_left.val[1]  = vmulq_f32(dxy_left.val[1], search_pixel.val[0]);
      dxy_right.val[1] = vmulq_f32(dxy_right.val[1], search_pixel.val[1]);

      Jres[2] -= (vgetq_lane_f32(search_pixel.val[0], 0) + vgetq_lane_f32(search_pixel.val[0], 1) +
                  vgetq_lane_f32(search_pixel.val[0], 2) + vgetq_lane_f32(search_pixel.val[0], 3) +
                  vgetq_lane_f32(search_pixel.val[1], 0) + vgetq_lane_f32(search_pixel.val[1], 1) +
                  vgetq_lane_f32(search_pixel.val[1], 2) + vgetq_lane_f32(search_pixel.val[1], 3));

      search_pixel.val[0] = vmulq_f32(search_pixel.val[0], search_pixel.val[0]);
      search_pixel.val[1] = vmulq_f32(search_pixel.val[1], search_pixel.val[1]);
      dxy_left.val[0] = vaddq_f32(dxy_left.val[0], dxy_right.val[0]);
      dxy_left.val[1] = vaddq_f32(dxy_left.val[1], dxy_right.val[1]);
      Jres[0] -= (vgetq_lane_f32(dxy_left.val[0], 0)  + vgetq_lane_f32(dxy_left.val[0], 1) +
                  vgetq_lane_f32(dxy_left.val[0], 2)  + vgetq_lane_f32(dxy_left.val[0], 3));
      Jres[1] -= (vgetq_lane_f32(dxy_left.val[1], 0)  + vgetq_lane_f32(dxy_left.val[1], 1) +
                  vgetq_lane_f32(dxy_left.val[1], 2)  + vgetq_lane_f32(dxy_left.val[1], 3));
      search_pixel.val[0] = vaddq_f32(search_pixel.val[0], search_pixel.val[1]);
      new_chi2 += vgetq_lane_f32(search_pixel.val[0], 0) + vgetq_lane_f32(search_pixel.val[0], 1) +
                  vgetq_lane_f32(search_pixel.val[0], 2) + vgetq_lane_f32(search_pixel.val[0], 3);
    }
    /*
    if(iter > 0 && new_chi2 > chi2) {
#if SUBPIX_VERBOSE
      VLOG(2) << "error increased.";
#endif
      u -= update[0];
      v -= update[1];
      break;
    }
    */
    chi2 = new_chi2;

    update = Hinv * Jres;
    u += update[0];
    v += update[1];
    mean_diff += update[2];

#if SUBPIX_VERBOSE
    VLOG(2) << "Iter " << iter << ":"
            << "\t u=" << u << ", v=" << v
            << "\t update = " << update[0] << ", " << update[1]
            << "\t chi2 = " << chi2;
#endif

    if (update[0] * update[0] + update[1] * update[1] < min_update_squared) {
#if SUBPIX_VERBOSE
      VLOG(2) << "converged.";
#endif
      converged = true;
      break;
    }
  }

  *cur_px_estimate << u, v;
  return converged;
}
#endif
}  // namespace align
}  // namespace XP


///////////////////////////////////////////////////////////////////////////////
// statistics::survival::data::algorithm::logit_log.hpp                      //
//                                                                           //
//  Copyright 2009 Erwann Rogard. 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)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_SURVIVAL_DATA_ALGORITHM_LOGIT_LOG_HPP_ER_2009
#define BOOST_STATISTICS_SURVIVAL_DATA_ALGORITHM_LOGIT_LOG_HPP_ER_2009
#include <stdexcept>
#include <iterator>
#include <boost/format.hpp>
#include <boost/statistics/survival/data/algorithm/detail/log_shift.hpp>

namespace boost{
namespace statistics{
namespace survival{
namespace data{

// Transforms an alternating sequence of (usually mean) failures
// and event times to their logit and log, with an offset that prevents nan.
//
// The rationale for a logit is that 0<= mean failure <=1
template<typename It,typename ItO,typename T>
ItO logit_log(
    It b,    // Alternating sequence, f[0],t[0],f[1],t[1],...
    It e,
    ItO o,   // Output iterator
    T t0,
    T t1
){
    static const char* str = "logit_log : distance(b,e) = %1% mod 2 !=2";

   if( std::distance(b,e)%2 != 0 ){
        format f(str); f%distance(b,e);
        throw std::runtime_error(f.str());
   }
   while(b!=e){
        (*o++) = detail::logit_shift((*b++),t0);
        (*o++) = detail::log_shift((*b++),t1);
    }
    return o;
}

}// data
}// survival
}// statistics
}// boost

#endif


/*
Copyright 2014 Alberto Crivellaro, Ecole Polytechnique Federale de Lausanne (EPFL), Switzerland.
alberto.crivellaro@epfl.ch

terms of the GNU General Public License as published by the Free Software
Foundation; either version 2 of the License, or (at your option) any later
version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY
WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with
this program; if not, write to the Free Software Foundation, Inc., 51 Franklin
Street, Fifth Floor, Boston, MA 02110-1301, USA
 */

#ifndef ITERATIVEOPTIMIZATION_HPP_
#define ITERATIVEOPTIMIZATION_HPP_

#include "Typedefs.hpp"
#include "Utilities.hpp"
#include "Homography.hpp"

#include <Eigen/LU>

#include <math.h>
#include <vector>
using namespace std;

class IterativeOptimization{
public:
	AlignmentResults SSDCalibration(const StructOfArray2di & pixelsOnTemplate,
		cv::Mat &grayscaleFloatTemplate, cv::Mat &grayscaleFloatImage, vector<float> &parameters,
		OptimizationParameters & optimizationParameters);
	AlignmentResults DescriptorFieldsCalibration(const StructOfArray2di & pixelsOnTemplate,
		cv::Mat &grayscaleFloatTemplate, cv::Mat &grayscaleFloatImage, vector<float> &paramters,
		OptimizationParameters & optimizationParameters);
	AlignmentResults GradientModuleCalibration(const StructOfArray2di & pixelsOnTemplate,
		cv::Mat &grayscaleFloatTemplate, cv::Mat &grayscaleFloatImage, vector<float> &paramters,
		OptimizationParameters & optimizationParameters);

	//these should be protected but there are tests on them.
	void ComputeWarpedPixels(const StructOfArray2di & pixelsOnTemplate, const vector<float>&  parameters, StructOfArray2di & warpedPixels);
    void AssembleSDImages(const vector<float>&  parameters, const cv::Mat &imageDx, const cv::Mat &imageDy,
		const StructOfArray2di & warpedPixels, const vector<Eigen::Matrix<float, 2, N_PARAM>,
		Eigen::aligned_allocator<Eigen::Matrix<float, 2, N_PARAM> > > & warpJacobians,
		Eigen::MatrixXf & sdImages);
protected:
	virtual AlignmentResults GaussNewtonMinimization(const StructOfArray2di & pixelsOnTemplate, const vector<cv::Mat> & images, const vector<cv::Mat> & templates, const OptimizationParameters optParam, vector<float> & parameters) = 0;
	AlignmentResults PyramidMultiLevelCalibration(const StructOfArray2di & pixelsOnTemplate, vector<cv::Mat> &templateDescriptorFields, vector<cv::Mat> &imageDescriptorFields, vector<float> & parameters, OptimizationParameters & optimizationParameters);
	void ComputeResiduals(const cv::Mat &image,  vector<float> & templatePixelIntensities, const StructOfArray2di & warpedPixels, vector<float> & errorImage);
	float ComputeResidualNorm(vector<float> &errorImage);
	int CheckConvergenceOptimization(float deltaPoseNorm, int nIter, float residualNormIncrement, OptimizationParameters optParam);
};

class LucasKanade: public IterativeOptimization{
private:
	AlignmentResults GaussNewtonMinimization(const StructOfArray2di & pixelsOnTemplate, const vector<cv::Mat> & images, const vector<cv::Mat> & templates, const OptimizationParameters optParam, vector<float> & parameters);
};

// TODO: implement this
//class ICA: public IterativeOptimization{
//private:
//	AlignmentResults GaussNewtonMinimization(const StructOfArray2di & pixelsOnTemplate, const vector<Mat> & images, const vector<Mat> & templates, const OptimizationParameters optParam, vector<float> & parameters);
//};

#endif /* ITERATIVEOPTIMIZATION_HPP_*/


#include <cstdint>
#include <fstream>
#include <iostream>
#include <random>
#include <vector>

#include <boost/program_options.hpp>

#include "hash.hpp"
#include "sampler.hpp"

int main(int argc, char * argv[])
{
    boost::program_options::options_description cli("Options");
    cli.add_options()
    ("help,h", "help")
    ("epsilon,e", boost::program_options::value<double>(), "epsilon value used in l0 sampling, must be in [0, 1], defaults to 1e-8")
    ("probability,p", boost::program_options::value<double>(), "success probability for binomial distribution, must be in [0, 1], defaults to 0.5")
    ("range,r", boost::program_options::value<uint16_t>(), "data range, must be positive, defaults to the stream size / 100")
    ("size,s", boost::program_options::value<uint16_t>(), "stream size, must be positive, defaults to 10000")
    ("trial,t", boost::program_options::value<uint16_t>(), "number of trials, must be positive, defaults to 10 * range");

    boost::program_options::variables_map options;
    boost::program_options::store(boost::program_options::parse_command_line(argc, argv, cli), options);
    boost::program_options::notify(options);

    if (options.count("help"))
    {
        std::cout << cli << std::endl;
        return 0;
    }

    double e = 1e-8;
    if (options.count("epsilon"))
    {
        e = options["epsilon"].as<double>();
        if (e <= 0 || e >= 1)
        {
            std::cerr << "-e requires a value in range (0, 1)" << std::endl;
            return 1;
        }
    }

    double p = 0.5;
    if (options.count("probability"))
    {
        p = options["probability"].as<double>();
        if (p <= 0 || p >= 1)
        {
            std::cerr << "-p requires a value in range (0, 1)" << std::endl;
            return 1;
        }
    }

    uint16_t s = 10000;
    if (options.count("size"))
        s = options["size"].as<uint16_t>();

    uint16_t r = s / 100;
    if (options.count("range"))
        r = options["range"].as<uint16_t>();

    uint16_t t = 10 * r;
    if (options.count("trial"))
        t = options["trial"].as<uint16_t>();

    // used to record distribution
    std::vector<uint16_t> vs(r), vlk(r);
    std::mt19937_64 g { std::random_device()() };
    uint8_t const k = log2(floor(1 / e));
    hash::k_universal_family<uint16_t> kuf(k);

    std::uniform_int_distribution<uint16_t> ud(0, r - 1);
    std::vector<sampler::l0_insertion> l0ks;
    l0ks.reserve(t);
    for (uint64_t i = 0; i < t; ++i)
        l0ks.emplace_back(r, kuf());
    for (uint64_t i = 0; i < s; ++i)
    {
        uint64_t v = ud(g);
        ++vs[v];
        for (auto & l0k : l0ks)
            l0k += v;
    }
    for (auto const & l0k : l0ks)
        ++vlk[l0k];
    for (uint64_t i = 0; i < r; ++i)
        std::cout << i << " " << vs[i] << " " << vlk[i] << std::endl;

    vs = std::vector<uint16_t>(r);
    vlk = std::vector<uint16_t>(r);

    std::binomial_distribution<uint16_t> ub(r - 1, p);
    l0ks.clear();
    for (uint64_t i = 0; i < t; ++i)
        l0ks.emplace_back(r, kuf());
    for (uint64_t i = 0; i < s; ++i)
    {
        uint64_t v = ub(g);
        ++vs[v];
        for (auto & l0k : l0ks)
            l0k += v;
    }
    for (auto const & l0k : l0ks)
        ++vlk[l0k];
    for (uint64_t i = 0; i < r; ++i)
        std::cout << i << " " << vs[i] << " " << vlk[i] << std::endl;

    return 0;
}


/*
 * StrainInvariants.cpp
 *
 *  Created on: Dec 23, 2018
 *      Author: lorenzo
 */

#include "StrainInvariants.h"

#define QUICKHULL_IMPLEMENTATION
#include <quickhull/quickhull.h>
#include <Eigen/Eigenvalues>

#include <iostream>
#include <sstream>

namespace ashell {

using std::endl;

StrainInvariants::StrainInvariants(std::string n_name) :
		Observable(n_name) {

}

StrainInvariants::~StrainInvariants() {

}

std::string StrainInvariants::as_string() {
	return _line;
}

void StrainInvariants::_observe(ullint step) {
	std::stringstream ss;

	std::vector<qh_vertex_t> vertices(_particles->N());

	auto poss = _particles->positions();
	vec3 com(0., 0., 0.);
	for(uint i = 0; i < _particles->N(); i++) {
		com += poss[i];
	}
	com /= _particles->N();

	for(uint i = 0; i < _particles->N(); i++) {
		vec3 v = poss[i] - com;
		vertices[i].x = v[0];
		vertices[i].y = v[1];
		vertices[i].z = v[2];
	}

	qh_mesh_t mesh = qh_quickhull3d(vertices.data(), _particles->N());

	vec3 ch_com(0., 0., 0.);
	double ch_volume = 0.;
	for(uint i = 0; i < mesh.nindices; i += 3) {
		vec3 p1(mesh.vertices[mesh.indices[i + 0]].x, mesh.vertices[mesh.indices[i + 0]].y, mesh.vertices[mesh.indices[i + 0]].z);
		vec3 p2(mesh.vertices[mesh.indices[i + 1]].x, mesh.vertices[mesh.indices[i + 1]].y, mesh.vertices[mesh.indices[i + 1]].z);
		vec3 p3(mesh.vertices[mesh.indices[i + 2]].x, mesh.vertices[mesh.indices[i + 2]].y, mesh.vertices[mesh.indices[i + 2]].z);

		ch_volume += (p1.dot(p2.cross(p3))) / 6.;
		ch_com += (p1 + p2 + p3);
	}

	ch_com /= mesh.nindices;

	// Gyration tensor
	mat3 gyration_tensor = mat3::Zero();
	for(int i = 0, j = 0; i < (int)mesh.nindices; i += 3, j++) {
		vec3 p1(mesh.vertices[mesh.indices[i + 0]].x, mesh.vertices[mesh.indices[i + 0]].y, mesh.vertices[mesh.indices[i + 0]].z);
		vec3 p2(mesh.vertices[mesh.indices[i + 1]].x, mesh.vertices[mesh.indices[i + 1]].y, mesh.vertices[mesh.indices[i + 1]].z);
		vec3 p3(mesh.vertices[mesh.indices[i + 2]].x, mesh.vertices[mesh.indices[i + 2]].y, mesh.vertices[mesh.indices[i + 2]].z);
		vec3 triangle_com = (p1 + p2 + p3) / 3. - ch_com;

		gyration_tensor(0, 0) += SQR(triangle_com[0]);
		gyration_tensor(0, 1) += triangle_com[0] * triangle_com[1];
		gyration_tensor(0, 2) += triangle_com[0] * triangle_com[2];

		gyration_tensor(1, 1) += SQR(triangle_com[1]);
		gyration_tensor(1, 2) += triangle_com[1] * triangle_com[2];

		gyration_tensor(2, 2) += SQR(triangle_com[2]);
	}
	gyration_tensor(1, 0) = gyration_tensor(0, 1);
	gyration_tensor(2, 0) = gyration_tensor(0, 2);
	gyration_tensor(2, 1) = gyration_tensor(1, 2);
	gyration_tensor /= mesh.nindices / 3.;

	vec3 eigenvalues = gyration_tensor.eigenvalues().real();

	double eigen_volume = 4 * M_PI * sqrt(3) * sqrt(eigenvalues[0]) * sqrt(eigenvalues[1]) * sqrt(eigenvalues[2]);

	ss << " " << ch_volume;
	ss << " " << eigen_volume;
	ss << " " << sqrt(eigenvalues[0]);
	ss << " " << sqrt(eigenvalues[1]);
	ss << " " << sqrt(eigenvalues[2]);

	qh_free_mesh(mesh);

	_line = ss.str();
}

} /* namespace ashell */


// Local
#include "visualise/ar_example.h"

// Eigen
#include <Eigen/Geometry>

// OpenCV
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

// Standard Library
#include <iomanip>

// Define a few BGR-colors for convenience.
namespace color
{
const cv::Scalar blue(255, 0, 0);
const cv::Scalar green(0, 255, 0);
const cv::Scalar red(0, 0, 255);
}

// Define font parameters.
namespace font
{
constexpr auto face = cv::FONT_HERSHEY_PLAIN;
constexpr auto scale_small = 1.0;
constexpr auto scale_large = 2.0;
constexpr auto thickness_bold = 3;
}

namespace visualise
{

ARExample::ARExample(double axes_length)
    : axes_length_{axes_length}
    , origin_{0, 0, 0, 1}
    , X_{axes_length_, 0, 0, 1}
    , Y_{0, axes_length_, 0, 1}
    , Z_{0, 0, axes_length_, 1}
{}

void ARExample::update(const cv::Mat& image,
                       const estimators::PoseEstimate& estimate,
                       const Eigen::Matrix3d& K,
                       double matching_time_ms,
                       double pose_est_time_ms) const
{
  // Clone image to draw in.
  cv::Mat ar_img = image.clone();

  // Extract reference to pose.
  const auto& pose = estimate.pose_W_C;

  if (estimate.isFound())
  {
    // Compute projection matrix P.
    const Eigen::Matrix<double, 3, 4> P = K * pose.inverse().matrix3x4();

    // Project 3D axis into image.
    Eigen::Vector2d o = (P*origin_).hnormalized();
    Eigen::Vector2d x = (P*X_).hnormalized();
    Eigen::Vector2d y = (P*Y_).hnormalized();
    Eigen::Vector2d z = (P*Z_).hnormalized();

    // Draw axis.
    cv::line(ar_img, cv::Point2d(o[0], o[1]), cv::Point2d(x[0], x[1]), color::red, 4);
    cv::line(ar_img, cv::Point2d(o[0], o[1]), cv::Point2d(y[0], y[1]), color::green, 4);
    cv::line(ar_img, cv::Point2d(o[0], o[1]), cv::Point2d(z[0], z[1]), color::blue, 4);

    // Print processing durations.
    std::stringstream corr_duration_txt;
    corr_duration_txt << std::fixed << std::setprecision(0);
    corr_duration_txt << "Matching: " << matching_time_ms << "ms";
    cv::putText(ar_img, corr_duration_txt.str(), {10, 20}, font::face, font::scale_small, color::red);

    std::stringstream pose_duration_txt;
    pose_duration_txt << std::fixed << std::setprecision(0);
    pose_duration_txt << "Pose est.: " << pose_est_time_ms << "ms";
    cv::putText(ar_img, pose_duration_txt.str(), {10, 40}, font::face, font::scale_small, color::red);

    // Print position.
    const auto& pos = pose.translation() * 100.0; // In cm.
    std::stringstream pos_txt;
    pos_txt << std::fixed << std::setprecision(1);
    pos_txt << "Pos (cm): (" << pos.x() << ", " << pos.y() << ", " << pos.z() << ")";
    cv::putText(ar_img, pos_txt.str(), {10, 60}, font::face, font::scale_small, color::green);

    // Print attitude.
    std::stringstream att_txt;
    Eigen::Vector3d att = attitudeFromR(pose.rotationMatrix());
    att_txt << std::fixed << std::setprecision(1);
    att_txt << "Att (deg): (" << att.x() << ", " << att.y() << ", " << att.z() << ")";
    cv::putText(ar_img, att_txt.str(), {10, 80}, font::face, font::scale_small, color::green);

    // Draw keypoints.
    const auto& inliers = estimate.image_inlier_points;
    for (const auto& inlier : inliers)
    {
      cv::drawMarker(ar_img, inlier, color::green, cv::MARKER_CROSS, 5);
    }

  }
  else
  {
    cv::putText(ar_img, "No tracking!", {20, 80}, font::face, font::scale_large, color::red, font::thickness_bold);
  }

  cv::imshow("AR example", ar_img);
}

Eigen::Vector3d ARExample::attitudeFromR(const Eigen::Matrix3d& R) const
{
  Eigen::Vector3d att;

  if (R(2, 0) < 1)
  {
    if (R(2, 0) > -1)
    {
      att.y() = std::asin(-R(2, 0));
      att.x() = std::atan2(R(2, 1) / std::cos(att.y()), R(2, 2) / std::cos(att.y()));
      att.z() = std::atan2(R(1, 0) / std::cos(att.y()), R(0, 0) / std::cos(att.y()));
    }
    else // R(2,0)==-1
    {
      att.x() = std::atan2(-R(1, 2), R(1, 1));
      att.y() = 0.5 * CV_PI;
      att.z() = 0;
    }
  }
  else // R(2,0) == 1
  {
    att.x() = std::atan2(-R(1, 2), R(1, 1));
    att.y() = -0.5 * CV_PI;
    att.z() = 0;
  }

  att.x() *= (180. / CV_PI);
  att.y() *= (180. / CV_PI);
  att.z() *= (180. / CV_PI);

  att.x() = std::fmod(360. + att.x(), 360.) - 180.;

  return att;
}

} // namespace visualise


#include <stan/math/fwd.hpp>
#include <gtest/gtest.h>
#include <boost/random/mersenne_twister.hpp>
#include <boost/math/distributions.hpp>
#include <vector>

TEST(ProbDistributionsMultinomial, fvar_double) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using stan::math::fvar;
  std::vector<int> ns;
  ns.push_back(1);
  ns.push_back(2);
  ns.push_back(3);
  Matrix<fvar<double>, Dynamic, 1> theta(3, 1);
  theta << 0.2, 0.3, 0.5;
  for (int i = 0; i < 3; i++)
    theta(i).d_ = 1.0;

  EXPECT_FLOAT_EQ(-2.002481, stan::math::multinomial_log(ns, theta).val_);
  EXPECT_FLOAT_EQ(17.666666, stan::math::multinomial_log(ns, theta).d_);
}

TEST(ProbDistributionsMultinomial, fvar_fvar_double) {
  using Eigen::Dynamic;
  using Eigen::Matrix;
  using stan::math::fvar;
  std::vector<int> ns;
  ns.push_back(1);
  ns.push_back(2);
  ns.push_back(3);
  Matrix<fvar<fvar<double> >, Dynamic, 1> theta(3, 1);
  theta << 0.2, 0.3, 0.5;
  for (int i = 0; i < 3; i++)
    theta(i).d_.val_ = 1.0;

  EXPECT_FLOAT_EQ(-2.002481, stan::math::multinomial_log(ns, theta).val_.val_);
  EXPECT_FLOAT_EQ(17.666666, stan::math::multinomial_log(ns, theta).d_.val_);
}


#ifndef BAOBZI_TEMPLATE_HPP
#define BAOBZI_TEMPLATE_HPP
#define _USE_MATH_DEFINES

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <iostream>
#include <limits>
#include <mutex>
#include <numeric>
#include <queue>
#include <vector>

#include <msgpack.hpp>
#define EIGEN_MATRIX_PLUGIN "baobzi/eigen_matrix_plugin.h"

#define EIGEN_MAX_ALIGN_BYTES 64
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/LU>
#include <unsupported/Eigen/CXX11/Tensor>

#include <baobzi/header.h>

/// Namespace for baobzi
namespace baobzi {
using index_t = uint32_t;    ///< Type specifying indexing into flattened tree
using coeff_data = double *; ///< Array to hold flattened coefficients

template <int DIM, int ORDER, int ISET>
class Function;

/// @brief Structure to represent geometric portion of Baobzi nodes
/// @tparam DIM number of dimensions of box
/// @tparam ISET Instruction set index (dummy variable to force alignment for different instruction sets)
template <int DIM, int ISET>
struct Box {
    using VecDimD = Eigen::Vector<double, DIM>; ///< DIM dimensional vector type

    VecDimD center;          ///< Center of box
    VecDimD inv_half_length; ///< 1.0 / half the dimension of the box

    Box<DIM, ISET>() = default; ///< default constructor for msgpack happiness
    /// @brief Constructor, just copies x, hl over
    Box<DIM, ISET>(const VecDimD &x, const VecDimD &hl)
        : center(x), inv_half_length(VecDimD::Ones().array() / hl.array()) {}

    /// @brief return vector of box half lengths along each dimension
    inline VecDimD half_length() const { return VecDimD::Ones().array() / inv_half_length.array(); }

    /// @brief MSGPACK serialization magic
    MSGPACK_DEFINE(center, inv_half_length);
};

/// @brief Return an estimate of the error for a given set of coefficients
/// @param[in] coeffs one or two dimensional Vector/Matrix of coefficients
/// @returns estimation of error given those coefficients
inline double standard_error(const Eigen::Ref<Eigen::MatrixXd> &coeffs) {
    double maxcoeff = 0.0;
    double scaling_factor = 1.0;
    if (coeffs.cols() == 1) {
        int n = coeffs.size();
        for (auto i = n - 2; i < n; ++i)
            maxcoeff = std::max(std::abs(coeffs(i, 0)), maxcoeff);
        scaling_factor = std::max(scaling_factor, std::abs(coeffs(0, 0)));
    } else {
        int n = coeffs.rows();
        for (auto i = 0; i < n; ++i)
            maxcoeff = std::max(std::abs(coeffs(i, n - i - 1)), maxcoeff);

        scaling_factor = std::max(scaling_factor, std::abs(coeffs(n - 1, 0)));
        scaling_factor = std::max(scaling_factor, std::abs(coeffs(0, n - 1)));
    }

    return maxcoeff / scaling_factor;
}

/// @brief Evaluate chebyshev polynomial given a box and a point inside that box
/// @tparam DIM dim of chebyshev polynomial to evaluate
/// @tparam ORDER order of chebyshev polynomial to evaluate
/// @tparam ISET Instruction set index (dummy variable to force alignment for different instruction sets)
/// @param[in] x position of point to evaluate (pre-normalized on interval from -1:1)
/// @param[in] coeffs_raw flat vector of coefficients
/// @returns value of interpolating function at x
template <int DIM, int ORDER, int ISET>
inline double cheb_eval(const Eigen::Vector<double, DIM> &x, const double *coeffs_raw);

template <int ORDER, int ISET>
inline double cheb_eval(const Eigen::Vector<double, 1> &x, const double *c) {
    // note (RB): uses clenshaw's method to avoid direct calculation of recurrence relation of
    // T_i, where res = \Sum_i T_i c_i
    const double x2 = 2 * x[0];

    double c0 = c[0];
    double c1 = c[1];
    for (int i = 2; i < ORDER; ++i) {
        double tmp = c1;
        c1 = c[i] - c0;
        c0 = tmp + c0 * x2;
    }

    return c1 + c0 * x[0];
}

template <int ORDER, int ISET>
inline double cheb_eval(const Eigen::Vector2d &x, const double *coeffs_raw) {
    // note (RB): There is code to do this with clenshaw's method (twice), but it doesn't seem
    // faster (isolated tests shows it's 3x faster, but that doesn't bear fruit in production
    // and this is, imho, clearer)
    Eigen::Matrix<double, 2, ORDER> Tns;
    Tns.col(0).setOnes();
    Tns.col(1) = x;
    for (int i = 2; i < ORDER; ++i)
        Tns.col(i) = 2 * x.array() * Tns.col(i - 1).array() - Tns.col(i - 2).array();

    Eigen::Map<const Eigen::Matrix<double, ORDER, ORDER>> coeffs(coeffs_raw);

    return Tns.row(0).transpose().dot(coeffs * Tns.row(1).transpose());
}

template <int ORDER, int ISET>
inline double cheb_eval(const Eigen::Vector3d &x, const double *coeffs_raw) {
    Eigen::Vector<double, ORDER> Tn[3];
    Tn[0][0] = Tn[1][0] = Tn[2][0] = 1.0;
    for (int i = 0; i < 3; ++i) {
        Tn[i][1] = x[i];
        for (int j = 2; j < ORDER; ++j)
            Tn[i][j] = 2 * x[i] * Tn[i][j - 1] - Tn[i][j - 2];
    }

    double res = 0.0;
    using map_t = Eigen::Map<const Eigen::Matrix<double, ORDER, ORDER>>;
    for (int i = 0; i < ORDER; ++i)
        res += Tn[0][i] * Tn[1].dot(map_t(coeffs_raw + i * ORDER * ORDER) * Tn[2]);
    return res;
}

/// @brief Node in baobzi::FunctionTree. If leaf, contains evaluation data, otherwise children
/// @tparam DIM dimension of function
/// @tparam ORDER order of evaluation polynomial
/// @tparam ISET instruction set index (dummy variable to force alignment for different instruction sets)
template <int DIM, int ORDER, int ISET>
class Node {
  public:
    using VecDimD = Eigen::Vector<double, DIM>;     ///< D dimensional vector type
    using VecOrderD = Eigen::Vector<double, ORDER>; ///< ORDER dimensional vector type
    using Func = Function<DIM, ORDER, ISET>;        ///< Type of boabzi function this belongs to

    Box<DIM, ISET> box_;                                          ///< Geometric position/size of this node
    uint64_t coeff_offset = std::numeric_limits<uint64_t>::max(); ///< Flattened chebyshev coeffs
    uint32_t first_child_idx = -1; ///< First child's index in a flattened list of all nodes

    Node<DIM, ORDER, ISET>() = default; ///< Default constructor for msgpack happiness

    /// @brief Construct node from box (without fitting)
    /// @param [in] box box this node represents
    Node<DIM, ORDER, ISET>(const Box<DIM, ISET> &box) : box_(box) {}

    /// @brief check if node is leaf
    /// @return true if leaf, false otherwise
    inline bool is_leaf() const { return coeff_offset != std::numeric_limits<uint64_t>::max(); }

    /// @brief Fit node to a given tolerance. If fit succeeds, set leaf and coeffs, otherwise ... don't
    ///
    /// Modifies: Node::leaf_, Node::coeffs_
    /// @param[in] input parameters for fit (function, tol, etc)
    /// @returns true if fit successful, false if not good enough
    std::vector<double> fit(const baobzi_input_t *input) {
        VecDimD half_length = box_.half_length();
        if constexpr (DIM == 1) {
            Eigen::Vector<double, ORDER> F;
            VecOrderD xvec = Func::get_cheb_nodes(box_.center[0] - half_length[0], box_.center[0] + half_length[0]);

            for (int i = 0; i < ORDER; ++i)
                F(i) = input->func(&xvec[i], input->data);

            Eigen::Vector<double, ORDER> coeffs = Func::VLU_.solve(F);

            if (standard_error(coeffs) > input->tol)
                return std::vector<double>();

            std::vector<double> coeffs_stl(coeffs.size());
            for (int i = 0; i < coeffs.size(); ++i)
                coeffs_stl[i] = coeffs(ORDER - i - 1);

            coeff_offset = 0;
            return coeffs_stl;
        }
        if constexpr (DIM == 2) {
            Eigen::Matrix<double, ORDER, ORDER> F;
            VecOrderD xvec = Func::get_cheb_nodes(box_.center[0] - half_length[0], box_.center[0] + half_length[0]);
            VecOrderD yvec = Func::get_cheb_nodes(box_.center[1] - half_length[1], box_.center[1] + half_length[1]);

            for (int i = 0; i < ORDER; ++i) {
                for (int j = 0; j < ORDER; ++j) {
                    double x[2] = {xvec[i], yvec[j]};
                    F(i, j) = input->func(x, input->data);
                }
            }

            Eigen::Matrix<double, ORDER, ORDER> coeffs = Func::VLU_.solve(F);
            coeffs = Func::VLU_.solve(coeffs.transpose()).transpose();

            if (standard_error(coeffs) > input->tol)
                return std::vector<double>();

            std::vector<double> coeffs_stl(coeffs.size());
            for (int i = 0; i < coeffs.size(); ++i)
                coeffs_stl[i] = coeffs(i);

            coeff_offset = 0;
            return coeffs_stl;
        }
        if constexpr (DIM == 3) {
            Eigen::Tensor<double, 3> F(ORDER, ORDER, ORDER);

            VecOrderD xvec = Func::get_cheb_nodes(box_.center[0] - half_length[0], box_.center[0] + half_length[0]);
            VecOrderD yvec = Func::get_cheb_nodes(box_.center[1] - half_length[1], box_.center[1] + half_length[1]);
            VecOrderD zvec = Func::get_cheb_nodes(box_.center[2] - half_length[2], box_.center[2] + half_length[2]);

            for (int i = 0; i < ORDER; ++i) {
                for (int j = 0; j < ORDER; ++j) {
                    for (int k = 0; k < ORDER; ++k) {
                        double x[3] = {xvec[i], yvec[j], zvec[k]};
                        F(i, j, k) = input->func(x, input->data);
                    }
                }
            }

            std::vector<double> coeffs(ORDER * ORDER * ORDER);
            Eigen::Tensor<double, 3> coeffs_tensor(ORDER, ORDER, ORDER);
            using matrix_t = Eigen::Matrix<double, ORDER, ORDER>;
            using map_t = Eigen::Map<matrix_t>;
            using tensor_t = Eigen::Tensor<double, 2>;
            for (int block = 0; block < ORDER; ++block) {
                tensor_t F_block_tensor = F.chip(block, 2);
                map_t F_block(F_block_tensor.data());

                matrix_t coeffs_tmp = Func::VLU_.solve(F_block);
                coeffs_tmp = Func::VLU_.solve(coeffs_tmp.transpose()).transpose();
                coeffs_tensor.chip(block, 2) = Eigen::TensorMap<tensor_t>(coeffs_tmp.data(), ORDER, ORDER);
            }
            for (int block = 0; block < ORDER; ++block) {
                Eigen::Tensor<double, 2> coeffs_tmp = coeffs_tensor.chip(block, 0);
                map_t coeffs_ysolve(coeffs_tmp.data());
                map_t(coeffs.data() + block * ORDER * ORDER) = Func::VLU_.solve(coeffs_ysolve.transpose()).transpose();
            }

            // Hack to use local coefficient array rather than global one
            coeff_offset = 0;
            for (int i = 0; i < ORDER; ++i) {
                for (int j = 0; j < ORDER; ++j) {
                    for (int k = 0; k < ORDER; ++k) {
                        VecDimD point =
                            (box_.center - half_length).array() +
                            2.0 * VecDimD{(double)i, (double)j, (double)k}.array() * half_length.array() / ORDER;

                        const double test_val = eval(point, coeffs.data());
                        const double actual_val = input->func(point.data(), input->data);
                        const double rel_error = std::abs((actual_val - test_val) / actual_val);

                        if (fabs(actual_val) > 1E-16 && rel_error > input->tol) {
                            coeff_offset = std::numeric_limits<uint64_t>::max();
                            return std::vector<double>();
                        }
                    }
                }
            }

            return coeffs;
        }
    }

    /// @brief eval node at point x
    /// @param[in] x point to evaluate at
    /// @param[in] coeffs flat/global coefficient array
    /// @returns function approximation at x
    inline double eval(const VecDimD &x, const double *coeffs) const {
        const VecDimD xinterp = (x - box_.center).array() * box_.inv_half_length.array();
        return cheb_eval<ORDER, ISET>(xinterp, coeffs + coeff_offset);
    }

    /// @brief Calculate memory usage of self (including unused space from vector allocation)
    /// @returns size in bytes of object instance
    inline std::size_t memory_usage() const { return sizeof(*this); }

    /// @brief MSGPACK serialization magic
    MSGPACK_DEFINE(box_, first_child_idx, coeff_offset);
};

/// @brief Represent a function in some domain as a tree of chebyshev nodes
/// @tparam DIM dimension of function
/// @tparam ORDER order of evaluation polynomial
/// @tparam ISET instruction set index (dummy variable to force alignment for different instruction sets)
template <int DIM, int ORDER, int ISET>
struct FunctionTree {
    static constexpr int NChild = 1 << DIM; ///< Number of children each node potentially has (2^D)
    static constexpr int Dim = DIM;         ///< Dimension of tree
    static constexpr int Order = ORDER;     ///< Order of tree

    using node_t = Node<DIM, ORDER, ISET>;      ///< DIM,ORDER node type
    using box_t = Box<DIM, ISET>;               ///< DIM box type
    using VecDimD = Eigen::Vector<double, DIM>; ///< D dimensional vector type

    std::vector<node_t> nodes_; ///< Flat list of all nodes in Tree (leaf or otherwise)
    int max_depth_;             ///< Maximum depth of tree

    /// @brief Construct tree
    /// @param[in] input parameters for fit (function, tol, etc)
    /// @param[in] coeffs flat/global coefficient vector
    /// @param[in] box box that this tree lives in
    FunctionTree<DIM, ORDER, ISET>(const baobzi_input_t *input, const Box<DIM, ISET> &box,
                                   std::vector<double> &coeffs) {
        std::queue<Box<DIM, ISET>> q;
        VecDimD half_width = box.half_length() * 0.5;
        q.push(box);

        index_t curr_child_idx = 1;
        max_depth_ = 0;
        while (!q.empty()) {
            int n_next = q.size();
            int node_index = nodes_.size();
            for (int i = 0; i < n_next; ++i) {
                box_t box = q.front();
                q.pop();

                nodes_.push_back(node_t(box));

                auto &node = nodes_[i + node_index];
                std::vector new_coeffs = node.fit(input);

                if (node.is_leaf()) {
                    node.coeff_offset = coeffs.size();
                    coeffs.insert(std::end(coeffs), std::begin(new_coeffs), std::end(new_coeffs));
                } else if (!node.is_leaf()) {
                    node.first_child_idx = curr_child_idx;
                    curr_child_idx += NChild;

                    VecDimD &center = node.box_.center;
                    for (index_t child = 0; child < NChild; ++child) {
                        VecDimD offset;

                        // Extract sign of each offset component from the bits of child
                        // Basically: permute all possible offsets
                        for (int j = 0; j < DIM; ++j) {
                            double signed_hw[2] = {-half_width[j], half_width[j]};
                            offset[j] = signed_hw[(child >> j) & 1];
                        }

                        q.push(Box<DIM, ISET>(center + offset, half_width));
                    }
                }
            }

            if (!q.empty())
                max_depth_++;

            half_width *= 0.5;
        }
    }

    FunctionTree<DIM, ORDER, ISET>() = default; ///< Default constructor for msgpack happiness

    /// @brief Find leaf node containing a point via standard pointer traversal
    /// @param[in] x point that the node will contain
    /// @return leaf node containing point x
    inline const node_t &find_node_traverse(const VecDimD &x) const {
        auto *node = &nodes_[0];
        auto *next_node = &nodes_[node->first_child_idx]; // attempt to force preload of potential next node
        while (!node->is_leaf()) {
            index_t child_idx = 0;
            for (int i = 0; i < DIM; ++i)
                child_idx = child_idx | ((x[i] > node->box_.center[i]) << i);

            node = next_node + child_idx;
            next_node = &nodes_[node->first_child_idx];
        }

        return *node;
    }

    /// @brief Get index of node at point x (relative to local nodes_ array)
    /// @param[in] x [DIM] point to lookup
    /// @returns index of node in nodes_ array containing x
    inline std::size_t get_node_index(const VecDimD &x) const {
        index_t curr_index = 0;
        while (!nodes_[curr_index].is_leaf()) {
            index_t child_idx = 0;
            for (int i = 0; i < DIM; ++i)
                child_idx = child_idx | ((x[i] > nodes_[curr_index].box_.center[i]) << i);

            curr_index = nodes_[curr_index].first_child_idx + child_idx;
        }

        return curr_index;
    }

    /// @brief Calculate total number of nodes in instance
    /// @return number of nodes in instance
    inline std::size_t size() const { return nodes_.size(); }

    /// @brief Calculate lowest depth of any node in instance (relative subtree node)
    /// @return lowest depth of all contained nodes
    inline int max_depth() const { return max_depth_; }

    /// @brief Calculate memory usage of self (including all contained nodes)
    /// @returns size in bytes of object instance
    inline std::size_t memory_usage() const {
        std::size_t memory_usage = sizeof(*this);
        for (const auto &node : nodes_)
            memory_usage += node.memory_usage();
        return memory_usage;
    }

    /// @brief eval function approximation at point
    /// @param[in] x point to evaluate function at
    /// @param[in] coeffs flat/global coefficient array
    /// @returns function approximation at point x
    inline double eval(const VecDimD &x, const double *coeffs) const { return find_node_traverse(x).eval(x, coeffs); }

    /// @brief msgpack serialization magic
    MSGPACK_DEFINE(nodes_);
};

/// @brief Represents a function in some domain as a grid of baobzi::FunctionTree objects
/// @tparam DIM dimension of function
/// @tparam ORDER order of evaluation polynomial
/// @tparam ISET instruction set index (dummy variable to force alignment for different instruction sets)
template <int DIM, int ORDER, int ISET = 0>
class Function {
  public:
    using VecDimD = Eigen::Vector<double, DIM>;            ///< DIM dimensional vector type
    using VecOrderD = Eigen::Vector<double, ORDER>;        ///< Order dimensional vector type
    using VanderMat = Eigen::Matrix<double, ORDER, ORDER>; ///< VanderMonde Matrix type
    using node_t = Node<DIM, ORDER, ISET>;                 ///< DIM,ORDER Node type (duh)
    using box_t = Box<DIM, ISET>;                          ///< DIM dimensional box type

    static constexpr int NChild = 1 << DIM; ///< Number of children each node potentially has (2^D)
    static constexpr int Dim = DIM;         ///< Input dimension of function
    static constexpr int Order = ORDER;     ///< Order of polynomial representation
    static constexpr int ISet = ISET;       ///< Instruction set (dummy param)

    static std::mutex statics_mutex;            ///< mutex for locking vandermonde/chebyshev initialization
    static VecOrderD cosarray_;                 ///< Cached array of cosine values at chebyshev nodes
    static Eigen::PartialPivLU<VanderMat> VLU_; ///< Cached LU decomposition of Vandermonde matrix

    box_t box_;          ///< box representing the domain of our function
    double tol_;         ///< Desired relative tolerance of our approximation
    VecDimD lower_left_; ///< Bottom 'corner' of our domain

    std::vector<FunctionTree<DIM, ORDER, ISET>> subtrees_; ///< Grid of FunctionTree objects that do the work
    Eigen::Vector<int, DIM> n_subtrees_;                   ///< Number of subtrees in each linear dimension of our space
    std::vector<int> subtree_node_offsets_; ///< n_subtrees array of offsets for where in the global array of node
                                            ///< pointers the global node pointer array starts
    std::vector<node_t *> node_pointers_;   ///< Vector of pointers to every node from every subtree
    VecDimD inv_bin_size_;                  ///< Inverse linear dimensions of the bins that our subtrees live

    std::vector<double> coeffs_; ///< Flat vector of all chebyshev coefficients from all leaf nodes

    bool split_multi_eval_ = true; ///< Split node-search and evaluation when evaluating multiple points

    /// Structure containing info about self creation :D
    struct {
        uint16_t base_depth = 0;   ///< depth of subtrees
        uint64_t n_evals_root = 0; ///< number of function evals before subtree calls
        uint32_t t_elapsed = 0;    ///< time in milliseconds to create object
    } stats_;

    /// @brief Calculate memory_usage of this object in bytes
    /// @returns Memory usage of baobzi object in bytes
    std::size_t memory_usage() const {
        std::size_t mem = sizeof(*this);
        mem += subtree_node_offsets_.capacity() * sizeof(subtree_node_offsets_[0]);
        mem += node_pointers_.capacity() * sizeof(node_pointers_[0]);
        mem += coeffs_.capacity() * sizeof(double);
        for (const auto &subtree : subtrees_)
            mem += subtree.memory_usage();
        return mem;
    }

    /// @brief Calculate and print various information about object instance to stdout
    void print_stats() const {
        std::size_t n_nodes = 0;
        std::size_t n_leaves = 0;
        std::size_t n_subtrees = subtrees_.size();
        int max_depth = 0;
        std::size_t mem = memory_usage();
        for (const auto &subtree : subtrees_) {
            n_nodes += subtree.size();
            max_depth = std::max(max_depth, subtree.max_depth());
            for (const auto &node : subtree.nodes_)
                n_leaves += node.is_leaf();
        }

        std::cout << "Baobzi tree represented by " << n_nodes << " nodes, of which " << n_leaves << " are leaves\n";
        std::cout << "Nodes are distributed across " << n_subtrees << " subtrees at an initial depth of "
                  << stats_.base_depth << " with a maximum subtree depth of " << max_depth << "\n";
        std::cout << "Total function evaluations required for fit: "
                  << n_nodes * (int)std::pow(ORDER, DIM) + stats_.n_evals_root << std::endl;
        std::cout << "Total time to create tree: " << stats_.t_elapsed << " milliseconds\n";
        std::cout << "Approximate memory usage of tree: " << (double)mem / (1024 * 1024) << " MiB" << std::endl;
    }

    /// @brief calculate vandermonde matrix
    /// @return Vandermonde matrix for chebyshev polynomials with order=ORDER
    static VanderMat calc_vandermonde() {
        VanderMat V;

        for (int j = 0; j < ORDER; ++j) {
            V(0, j) = 1;
            V(1, j) = cosarray_(j);
        }

        for (int i = 2; i < ORDER; ++i) {
            for (int j = 0; j < ORDER; ++j) {
                V(i, j) = double(2) * V(i - 1, j) * cosarray_(j) - V(i - 2, j);
            }
        }

        return V.transpose();
    }

    /// @brief calculate chebyshev nodes on bounds [lb, ub]
    /// @param[in] lb lower bound
    /// @param[in] ub upper bound
    /// @returns vector of chebyshev nodes scaled within [lb, ub]
    static inline VecOrderD get_cheb_nodes(double lb, double ub) {
        return 0.5 * ((lb + ub) + (ub - lb) * cosarray_.array());
    }

    /// @brief initialize static class variables
    ///
    /// Modifies baobzi::Function::cosarray_, baobzi::Function::VLU_
    static void init_statics() {
        static bool is_initialized = false;
        std::lock_guard<std::mutex> lock(statics_mutex);
        if (is_initialized)
            return;

        for (int i = 0; i < ORDER; ++i)
            cosarray_[ORDER - i - 1] = cos(M_PI * (i + 0.5) / ORDER);
        VLU_ = Eigen::PartialPivLU<VanderMat>(calc_vandermonde());
        is_initialized = true;
    }

    /// @brief Construct our Function object (fits recursively, can be slow)
    /// @param[in] input parameters for fit (function, tol, etc)
    /// @param[in] xp [dim] center of function domain
    /// @param[in] lp [dim] half length of function domain
    Function<DIM, ORDER, ISET>(const baobzi_input_t *input, const double *xp, const double *lp)
        : box_(VecDimD(xp), VecDimD(lp)), tol_(input->tol), split_multi_eval_(input->split_multi_eval) {
        auto t_start = std::chrono::steady_clock::now();
        init_statics();

        VecDimD l(lp);
        VecDimD x(xp);
        std::queue<box_t> q;
        std::queue<box_t> maybe_q;

        for (int i = 0; i < DIM; ++i)
            n_subtrees_[i] = l[i] / l.minCoeff();

        q.push(box_t(x, l));

        // Half-width of next children
        VecDimD half_width = l * 0.5;

        // Breadth first search. Step through each level of the tree and test fit all of the nodes
        // We exit when a level isn't completely filled with parent nodes (rather than leaves)
        // This way we can always avoid redundant traversals by jumping straight to a root node of a subtree
        while (!q.empty()) {
            int n_next = q.size();

            auto add_node_children_to_queue = [](std::queue<box_t> &theq, const VecDimD &center,
                                                 const VecDimD &half_width) {
                for (unsigned child = 0; child < NChild; ++child) {
                    VecDimD offset;

                    // Extract sign of each offset component from the bits of child
                    // Basically: permute all possible offsets
                    for (int j = 0; j < DIM; ++j) {
                        double signed_hw[2] = {-half_width[j], half_width[j]};
                        offset[j] = signed_hw[(child >> j) & 1];
                    }

                    theq.push(box_t(center + offset, half_width));
                }
            };

            std::vector<node_t> nodes;
            double leaf_fraction = 0.0;
            for (int i = 0; i < n_next; ++i) {
                box_t box = q.front();
                q.pop();

                nodes.emplace_back(node_t(box));
                auto &node = nodes.back();
                node.fit(input);

                if (!node.is_leaf()) {
                    add_node_children_to_queue(q, node.box_.center, half_width);
                } else {
                    leaf_fraction += 1.0;
                    add_node_children_to_queue(maybe_q, node.box_.center, half_width);
                }
            }
            stats_.n_evals_root += nodes.size() * std::pow(ORDER, DIM);

            leaf_fraction /= nodes.size();
            if (leaf_fraction < input->minimum_leaf_fraction) {
                while (!maybe_q.empty()) {
                    box_t box = maybe_q.front();
                    maybe_q.pop();
                    q.push(box);
                }
            }

            half_width *= 0.5;
            if ((1 << (DIM * (stats_.base_depth + 1))) == q.size()) {
                n_subtrees_ *= 2;
                stats_.base_depth++;
            } else
                break;
        }

        VecDimD bin_size;
        VecDimD half_length = box_.half_length();
        for (int j = 0; j < DIM; ++j) {
            bin_size[j] = 2.0 * half_length[j] / n_subtrees_[j];
            inv_bin_size_[j] = 0.5 * n_subtrees_[j] / half_length[j];
        }
        lower_left_ = box_.center - half_length;

        subtrees_.reserve(n_subtrees_.prod());
        for (int i_bin = 0; i_bin < n_subtrees_.prod(); ++i_bin) {
            Eigen::Vector<int, DIM> bins = get_bins(i_bin);

            VecDimD parent_center =
                (bins.template cast<double>().array() + 0.5) * bin_size.array() + lower_left_.array();

            Box<DIM, ISET> root_box = {parent_center, 0.5 * bin_size};
            subtrees_.push_back(FunctionTree<DIM, ORDER, ISET>(input, root_box, coeffs_));
        }

        auto t_end = std::chrono::steady_clock::now();
        auto t_elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(t_end - t_start);
        stats_.t_elapsed = t_elapsed.count();
        build_cache();
    }

    /// @brief Build any intermediate state necessary for computation
    void build_cache() {
        subtree_node_offsets_.resize(n_subtrees_.prod());
        subtree_node_offsets_[0] = 0;
        for (int i = 1; i < subtree_node_offsets_.size(); ++i)
            subtree_node_offsets_[i] = subtree_node_offsets_[i - 1] + subtrees_[i - 1].size();

        auto n_nodes_tot = std::accumulate(subtrees_.begin(), subtrees_.end(), (std::size_t)0,
                                           [](size_t prior, auto &subtree) { return prior + subtree.size(); });

        node_pointers_.resize(n_nodes_tot);

        int i = 0;
        for (auto &subtree : subtrees_)
            for (node_t &node : subtree.nodes_)
                node_pointers_[i++] = &node;
    }

    /// @brief default constructor for msgpack magic
    Function<DIM, ORDER, ISET>() { init_statics(); };

    /// @brief convert linear bin index to [dim] bin vector
    /// @param[in] i_bin linear index
    /// @returns [dim] bin vector
    inline Eigen::Vector<int, DIM> get_bins(const int i_bin) const {
        if constexpr (DIM == 1)
            return Eigen::Vector<int, DIM>{i_bin};
        else if constexpr (DIM == 2)
            return Eigen::Vector<int, DIM>{i_bin % n_subtrees_[0], i_bin / n_subtrees_[0]};
        else if constexpr (DIM == 3)
            return Eigen::Vector<int, DIM>{i_bin % n_subtrees_[0], (i_bin / n_subtrees_[0]) % n_subtrees_[1],
                                           i_bin / (n_subtrees_[0] * n_subtrees_[1])};
    }

    /// @brief find linear index of bin at a point
    /// @param[in] x [1] position to find bin
    /// @returns linear index of bin that x lives in
    inline int get_linear_bin(const Eigen::Vector<double, 1> &x) const {
        const double x_bin = x[0] - lower_left_[0];
        return x_bin * inv_bin_size_[0];
    }

    /// @brief find linear index of bin at a point
    /// @param[in] x [2] position to find bin
    /// @returns linear index of bin that x lives in
    inline int get_linear_bin(const Eigen::Vector2d &x) const {
        const VecDimD x_bin = x - lower_left_;
        const Eigen::Vector<int, DIM> bin = (x_bin.array() * inv_bin_size_.array()).template cast<int>();
        return bin[0] + n_subtrees_[0] * bin[1];
    }

    /// @brief find linear index of bin at a point
    /// @param[in] x [3] position to find bin
    /// @returns linear index of bin that x lives in
    inline int get_linear_bin(const Eigen::Vector3d &x) const {
        const VecDimD x_bin = x - lower_left_;
        const Eigen::Vector<int, DIM> bin = (x_bin.array() * inv_bin_size_.array()).template cast<int>();
        return bin[0] + n_subtrees_[0] * bin[1] + n_subtrees_[0] * n_subtrees_[1] * bin[2];
    }

    /// @brief get constant reference to leaf node that contains a point
    /// @param[in] x point of interest
    /// @returns constant reference to leaf node that contains x
    inline const node_t &find_node(const VecDimD &x) const {
        return subtrees_[get_linear_bin(x)].find_node_traverse(x);
    }

    /// @brief eval function approximation at point
    /// @param[in] x point to evaluate function at
    /// @returns function approximation at point x
    inline double eval(const VecDimD &x) const { return find_node(x).eval(x, coeffs_.data()); }

    /// @brief eval function approximation at point
    /// @param[in] xp [DIM] point to evaluate function at
    /// @returns function approximation at point xp
    inline double eval(const double *xp) const { return eval(VecDimD(xp)); }

    /// @brief get index of node (across all subnodes)
    /// @param[in] x [DIM] point to find the node of
    /// @returns index in global node array
    inline std::size_t get_global_node_index(const VecDimD &x) const {
        int i_sub = get_linear_bin(x);
        return subtree_node_offsets_[i_sub] + subtrees_[i_sub].get_node_index(x);
    }

    /// @brief eval function approximation at n_trg points
    /// @param[in] xp [DIM * n_trg] array of points to evaluate function at
    /// @param[out] res [n_trg] array of results
    /// @param[in] n_trg number of points to evaluate
    inline void eval(const double *xp, double *res, int n_trg) const {
        if (split_multi_eval_) {
            std::vector<std::pair<node_t *, VecDimD>> node_map(n_trg);
            for (int i = 0; i < n_trg; ++i) {
                VecDimD xi = VecDimD(xp + DIM * i);
                node_map[i] = std::make_pair(node_pointers_[get_global_node_index(xi)], xi);
            }

            for (int i_trg = 0; i_trg < n_trg; i_trg++)
                res[i_trg] = node_map[i_trg].first->eval(node_map[i_trg].second, coeffs_.data());
        } else
            for (int i_trg = 0; i_trg < n_trg; i_trg++)
                res[i_trg] = eval(VecDimD(xp + DIM * i_trg));
    }

    /// @brief eval function approximation at point
    /// @param[in] x [DIM] point to evaluate function at
    /// @returns function approximation at point x
    inline double operator()(const VecDimD &x) const { return eval(x); }

    /// @brief eval function approximation at point
    /// @param[in] x point to evaluate function at
    /// @returns function approximation at point x
    inline double operator()(const double *x) const { return eval(x); }

    /// @brief eval function approximation at n_trg points
    /// @param[in] xp [DIM * n_trg] array of points to evaluate function at
    /// @param[out] res [DIM * n_trg] array of results
    /// @param[in] n_trg number of points to evaluate
    inline void operator()(const double *xp, double *res, int n_trg) const { eval(xp, res, n_trg); }

    /// @brief save function approximation to file
    /// @param[in] filename path to save file at
    void save(const char *filename) const {
        std::ofstream ofs(filename, std::ofstream::binary | std::ofstream::out);
        baobzi_header_t params{Dim, Order, BAOBZI_HEADER_VERSION};
        msgpack::pack(ofs, params);
        msgpack::pack(ofs, *this);
    }

    /// @brief msgpack serialization magic
    MSGPACK_DEFINE_MAP(box_, subtrees_, n_subtrees_, tol_, lower_left_, inv_bin_size_, coeffs_, split_multi_eval_);
};

template <int DIM, int ORDER, int ISET>
std::mutex Function<DIM, ORDER, ISET>::statics_mutex;

template <int DIM, int ORDER, int ISET>
typename Function<DIM, ORDER, ISET>::VecOrderD Function<DIM, ORDER, ISET>::cosarray_;

template <int DIM, int ORDER, int ISET>
Eigen::PartialPivLU<typename Function<DIM, ORDER, ISET>::VanderMat> Function<DIM, ORDER, ISET>::VLU_;
} // namespace baobzi

#endif


#pragma once

#include <Eigen/Dense>
#include <iostream>
#include <vector>

#include <boost/random/mersenne_twister.hpp>
#include <boost/shared_ptr.hpp>

#include "sphere.hpp"
#include "dpmeans.hpp"
#include "dir.hpp"
#include "cat.hpp"

using namespace Eigen;
using std::cout;
using std::endl;

template<class T>
class DDPMeans : public DPMeans<T>
{
public:
  DDPMeans(const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx,
      T lambda, T Q, T tau, boost::mt19937* pRndGen);
  virtual ~DDPMeans();

//  void initialize(const Matrix<T,Dynamic,Dynamic>& x);
  virtual void updateLabels();
  virtual void updateCenters();

  virtual void nextTimeStep(const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx);
  virtual void updateState(); // after converging for a single time instant

  virtual uint32_t indOfClosestCluster(int32_t i);

  const static uint32_t UNASSIGNED = 4294967295;
protected:

  std::vector<T> ts_; // age of clusters - incremented each iteration
  std::vector<T> ws_; // weights of each cluster
  T Q_; // Q parameter
  T tau_; // tau parameter
  T Kprev_; // K before updateLabels()
  Matrix<T,Dynamic,Dynamic> psPrev_; // centroids from last set of datapoints

};

// -------------------------------- impl ----------------------------------
template<class T>
DDPMeans<T>::DDPMeans(const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx,
    T lambda, T Q, T tau, boost::mt19937* pRndGen)
  : DPMeans<T>(spx,0,lambda,pRndGen), Q_(Q), tau_(tau)
{
  // compute initial counts for weight initialization
//#pragma omp parallel for
//  for (uint32_t k=0; k<this->K_; ++k)
//    for(uint32_t i=0; i<this->N_; ++i)
//      if(this->z_(i) == k)
//        this->Ns_(k) ++;
//  for (uint32_t k=0; k<this->K_; ++k)
//  {
//    ws_.push_back(this->Ns_(k));
//    ts_.push_back(1);
//  }
  this->Kprev_ = 0; // so that centers are initialized directly from sample mean
  psPrev_ = this->ps_;
}

template<class T>
DDPMeans<T>::~DDPMeans()
{}

template<class T>
uint32_t DDPMeans<T>::indOfClosestCluster(int32_t i)
{
  int z_i = this->K_;
  T sim_closest = this->lambda_;
//  cout<<"K="<<this->K_<<" Ns:"<<this->Ns_.transpose()<<endl;
//  cout<<"cluster dists "<<i<<": "<<this->lambda_;
  for (uint32_t k=0; k<this->K_; ++k)
  {
    T sim_k = this->dist(this->ps_.col(k), this->spx_->col(i));
    if(this->Ns_(k) == 0) // cluster not instantiated yet in this timestep
    {
      //TODO use gamma
//      T gamma = 1.0/(1.0/ws_[z_i] + ts_[z_i]*tau_);
      sim_k = sim_k/(tau_*ts_[k]+1.) + Q_*ts_[k];
//      sim_k = sim_k/(tau_*ts_[k]+1.) + Q_*ts_[k];
    }
//    cout<<" "<<sim_k;
    if(this->closer(sim_k, sim_closest))
    {
      sim_closest = sim_k;
      z_i = k;
    }
  }
//  }cout<<endl;
  return z_i;
}

template<class T>
void DDPMeans<T>::updateLabels()
{
  // reset cluster counts -> all uninstantiated
//  for (uint32_t k=0; k<this->K_; ++k)
//    this->Ns_(k) = 0;
//#pragma omp parallel for
// TODO not sure how to parallelize
  for(uint32_t i=0; i<this->N_; ++i)
  {
    uint32_t z_i = indOfClosestCluster(i);
    if(z_i == this->K_)
    { // start a new cluster
      Matrix<T,Dynamic,Dynamic> psNew(this->D_,this->K_+1);
      psNew.leftCols(this->K_) = this->ps_;
      psNew.col(this->K_) = this->spx_->col(i);
      this->ps_ = psNew;
      this->K_ ++;
      this->Ns_.conservativeResize(this->K_);
      this->Ns_(z_i) = 1.;
    } else {
      if(this->Ns_[z_i] == 0)
      { // instantiated an old cluster
        T gamma = 1.0/(1.0/ws_[z_i] + ts_[z_i]*tau_);
        this->ps_.col(z_i)=(this->ps_.col(z_i)*gamma + this->spx_->col(i))/(gamma+1.);
      }
      this->Ns_(z_i) ++;
    }
    if(this->z_(i) != UNASSIGNED)
    {
      this->Ns_(this->z_(i)) --;
    }
    this->z_(i) = z_i;
  }
};

template<class T>
void DDPMeans<T>::updateCenters()
{
#pragma omp parallel for
  for(uint32_t k=0; k<this->K_; ++k)
  {
    Matrix<T,Dynamic,1> mean_k = this->computeCenter(k);
    if (this->Ns_(k) > 0)
    { // have data to update kth cluster
      if(k < this->Kprev_){
        T gamma = 1.0/(1.0/ws_[k] + ts_[k]*tau_);
        this->ps_.col(k) = (this->ps_.col(k)*gamma+mean_k*this->Ns_(k))/
          (gamma+this->Ns_(k));
      }else{
        this->ps_.col(k)=mean_k;
      }
    }
  }
};

template<class T>
void DDPMeans<T>::nextTimeStep(const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx)
{
  assert(this->D_ == spx->rows());
  this->spx_ = spx; // update the data
  this->N_ = spx->cols();
  this->z_.resize(this->N_);
  this->z_.fill(UNASSIGNED);
};

template<class T>
void DDPMeans<T>::updateState()
{
  for(uint32_t k=0; k<this->K_; ++k)
  {
    if (k<ws_.size() && this->Ns_(k) > 0)
    { // instantiated cluster from previous time;
      ws_[k] = 1./(1./ws_[k] + ts_[k]*tau_) + this->Ns_(k);
      ts_[k] = 0; // re-instantiated -> age is 0
    }else if(k >= ws_.size()){
      // new cluster
      ts_.push_back(0);
      ws_.push_back(this->Ns_(k));
    }
    ts_[k] ++; // increment all ages
    cout<<"cluster "<<k
      <<"\tN="<<this->Ns_(k)
      <<"\tage="<<ts_[k]
      <<"\tweight="<<ws_[k]<<endl;
    cout<<"  center: "<<this->ps_.col(k).transpose()<<endl;
  }
  psPrev_ = this->ps_;
  this->Kprev_ = this->K_;
};