#include "loam_velodyne/BasicLaserOdometry.h"

#include "math_utils.h"
#include <pcl/filters/filter.h>
#include <Eigen/Eigenvalues>
#include <Eigen/QR>

namespace loam
{

using std::sin;
using std::cos;
using std::asin;
using std::atan2;
using std::sqrt;
using std::fabs;
using std::pow;

BasicLaserOdometry::BasicLaserOdometry(float scanPeriod, size_t maxIterations) :
   _scanPeriod(scanPeriod),
   _systemInited(false),
   _frameCount(0),
   _maxIterations(maxIterations),
   _deltaTAbort(0.1),
   _deltaRAbort(0.1),
   _cornerPointsSharp(new pcl::PointCloud<pcl::PointXYZI>()),
   _cornerPointsLessSharp(new pcl::PointCloud<pcl::PointXYZI>()),
   _surfPointsFlat(new pcl::PointCloud<pcl::PointXYZI>()),
   _surfPointsLessFlat(new pcl::PointCloud<pcl::PointXYZI>()),
   _laserCloud(new pcl::PointCloud<pcl::PointXYZI>()),
   _lastCornerCloud(new pcl::PointCloud<pcl::PointXYZI>()),
   _lastSurfaceCloud(new pcl::PointCloud<pcl::PointXYZI>()),
   _laserCloudOri(new pcl::PointCloud<pcl::PointXYZI>()),
   _coeffSel(new pcl::PointCloud<pcl::PointXYZI>())
{}

void BasicLaserOdometry::transformToStart(const pcl::PointXYZI& pi, pcl::PointXYZI& po)
{
   float s = (1.f / _scanPeriod) * (pi.intensity - int(pi.intensity));

   po.x = pi.x - s * _transform.pos.x();
   po.y = pi.y - s * _transform.pos.y();
   po.z = pi.z - s * _transform.pos.z();
   po.intensity = pi.intensity;

   Angle rx = -s * _transform.rot_x.rad();
   Angle ry = -s * _transform.rot_y.rad();
   Angle rz = -s * _transform.rot_z.rad();
   rotateZXY(po, rz, rx, ry);
}

size_t BasicLaserOdometry::transformToEnd(pcl::PointCloud<pcl::PointXYZI>::Ptr& cloud)
{
   size_t cloudSize = cloud->points.size();

   for (size_t i = 0; i < cloudSize; i++)
   {
      pcl::PointXYZI& point = cloud->points[i];

      float s = (1.f / _scanPeriod) * (point.intensity - int(point.intensity));

      point.x -= s * _transform.pos.x();
      point.y -= s * _transform.pos.y();
      point.z -= s * _transform.pos.z();
      point.intensity = int(point.intensity);

      Angle rx = -s * _transform.rot_x.rad();
      Angle ry = -s * _transform.rot_y.rad();
      Angle rz = -s * _transform.rot_z.rad();
      rotateZXY(point, rz, rx, ry);
      rotateYXZ(point, _transform.rot_y, _transform.rot_x, _transform.rot_z);

      point.x += _transform.pos.x() - _imuShiftFromStart.x();
      point.y += _transform.pos.y() - _imuShiftFromStart.y();
      point.z += _transform.pos.z() - _imuShiftFromStart.z();

      rotateZXY(point, _imuRollStart, _imuPitchStart, _imuYawStart);
      rotateYXZ(point, -_imuYawEnd, -_imuPitchEnd, -_imuRollEnd);
   }

   return cloudSize;
}

void BasicLaserOdometry::pluginIMURotation(const Angle& bcx, const Angle& bcy, const Angle& bcz,
                                           const Angle& blx, const Angle& bly, const Angle& blz,
                                           const Angle& alx, const Angle& aly, const Angle& alz,
                                           Angle &acx, Angle &acy, Angle &acz)
{
   float sbcx = bcx.sin();
   float cbcx = bcx.cos();
   float sbcy = bcy.sin();
   float cbcy = bcy.cos();
   float sbcz = bcz.sin();
   float cbcz = bcz.cos();

   float sblx = blx.sin();
   float cblx = blx.cos();
   float sbly = bly.sin();
   float cbly = bly.cos();
   float sblz = blz.sin();
   float cblz = blz.cos();

   float salx = alx.sin();
   float calx = alx.cos();
   float saly = aly.sin();
   float caly = aly.cos();
   float salz = alz.sin();
   float calz = alz.cos();

   float srx = -sbcx * (salx*sblx + calx * caly*cblx*cbly + calx * cblx*saly*sbly)
      - cbcx * cbcz*(calx*saly*(cbly*sblz - cblz * sblx*sbly)
                     - calx * caly*(sbly*sblz + cbly * cblz*sblx) + cblx * cblz*salx)
      - cbcx * sbcz*(calx*caly*(cblz*sbly - cbly * sblx*sblz)
                     - calx * saly*(cbly*cblz + sblx * sbly*sblz) + cblx * salx*sblz);
   acx = -asin(srx);

   float srycrx = (cbcy*sbcz - cbcz * sbcx*sbcy)*(calx*saly*(cbly*sblz - cblz * sblx*sbly)
                                                  - calx * caly*(sbly*sblz + cbly * cblz*sblx) + cblx * cblz*salx)
      - (cbcy*cbcz + sbcx * sbcy*sbcz)*(calx*caly*(cblz*sbly - cbly * sblx*sblz)
                                        - calx * saly*(cbly*cblz + sblx * sbly*sblz) + cblx * salx*sblz)
      + cbcx * sbcy*(salx*sblx + calx * caly*cblx*cbly + calx * cblx*saly*sbly);
   float crycrx = (cbcz*sbcy - cbcy * sbcx*sbcz)*(calx*caly*(cblz*sbly - cbly * sblx*sblz)
                                                  - calx * saly*(cbly*cblz + sblx * sbly*sblz) + cblx * salx*sblz)
      - (sbcy*sbcz + cbcy * cbcz*sbcx)*(calx*saly*(cbly*sblz - cblz * sblx*sbly)
                                        - calx * caly*(sbly*sblz + cbly * cblz*sblx) + cblx * cblz*salx)
      + cbcx * cbcy*(salx*sblx + calx * caly*cblx*cbly + calx * cblx*saly*sbly);
   acy = atan2(srycrx / acx.cos(), crycrx / acx.cos());

   float srzcrx = sbcx * (cblx*cbly*(calz*saly - caly * salx*salz) - cblx * sbly*(caly*calz + salx * saly*salz) + calx * salz*sblx)
      - cbcx * cbcz*((caly*calz + salx * saly*salz)*(cbly*sblz - cblz * sblx*sbly)
                     + (calz*saly - caly * salx*salz)*(sbly*sblz + cbly * cblz*sblx)
                     - calx * cblx*cblz*salz)
      + cbcx * sbcz*((caly*calz + salx * saly*salz)*(cbly*cblz + sblx * sbly*sblz)
                     + (calz*saly - caly * salx*salz)*(cblz*sbly - cbly * sblx*sblz)
                     + calx * cblx*salz*sblz);
   float crzcrx = sbcx * (cblx*sbly*(caly*salz - calz * salx*saly) - cblx * cbly*(saly*salz + caly * calz*salx) + calx * calz*sblx)
      + cbcx * cbcz*((saly*salz + caly * calz*salx)*(sbly*sblz + cbly * cblz*sblx)
                     + (caly*salz - calz * salx*saly)*(cbly*sblz - cblz * sblx*sbly)
                     + calx * calz*cblx*cblz)
      - cbcx * sbcz*((saly*salz + caly * calz*salx)*(cblz*sbly - cbly * sblx*sblz)
                     + (caly*salz - calz * salx*saly)*(cbly*cblz + sblx * sbly*sblz)
                     - calx * calz*cblx*sblz);
   acz = atan2(srzcrx / acx.cos(), crzcrx / acx.cos());
}

void BasicLaserOdometry::accumulateRotation(Angle cx, Angle cy, Angle cz,
                                            Angle lx, Angle ly, Angle lz,
                                            Angle &ox, Angle &oy, Angle &oz)
{
   float srx = lx.cos()*cx.cos()*ly.sin()*cz.sin()
      - cx.cos()*cz.cos()*lx.sin()
      - lx.cos()*ly.cos()*cx.sin();
   ox = -asin(srx);

   float srycrx = lx.sin()*(cy.cos()*cz.sin() - cz.cos()*cx.sin()*cy.sin())
      + lx.cos()*ly.sin()*(cy.cos()*cz.cos() + cx.sin()*cy.sin()*cz.sin())
      + lx.cos()*ly.cos()*cx.cos()*cy.sin();
   float crycrx = lx.cos()*ly.cos()*cx.cos()*cy.cos()
      - lx.cos()*ly.sin()*(cz.cos()*cy.sin() - cy.cos()*cx.sin()*cz.sin())
      - lx.sin()*(cy.sin()*cz.sin() + cy.cos()*cz.cos()*cx.sin());
   oy = atan2(srycrx / ox.cos(), crycrx / ox.cos());

   float srzcrx = cx.sin()*(lz.cos()*ly.sin() - ly.cos()*lx.sin()*lz.sin())
      + cx.cos()*cz.sin()*(ly.cos()*lz.cos() + lx.sin()*ly.sin()*lz.sin())
      + lx.cos()*cx.cos()*cz.cos()*lz.sin();
   float crzcrx = lx.cos()*lz.cos()*cx.cos()*cz.cos()
      - cx.cos()*cz.sin()*(ly.cos()*lz.sin() - lz.cos()*lx.sin()*ly.sin())
      - cx.sin()*(ly.sin()*lz.sin() + ly.cos()*lz.cos()*lx.sin());
   oz = atan2(srzcrx / ox.cos(), crzcrx / ox.cos());
}

void BasicLaserOdometry::updateIMU(pcl::PointCloud<pcl::PointXYZ> const& imuTrans)
{
   assert(4 == imuTrans.size());
   _imuPitchStart = imuTrans.points[0].x;
   _imuYawStart = imuTrans.points[0].y;
   _imuRollStart = imuTrans.points[0].z;

   _imuPitchEnd = imuTrans.points[1].x;
   _imuYawEnd = imuTrans.points[1].y;
   _imuRollEnd = imuTrans.points[1].z;

   _imuShiftFromStart = imuTrans.points[2];
   _imuVeloFromStart = imuTrans.points[3];
}

void BasicLaserOdometry::process()
{
   if (!_systemInited)
   {
      _cornerPointsLessSharp.swap(_lastCornerCloud);
      _surfPointsLessFlat.swap(_lastSurfaceCloud);

      _lastCornerKDTree.setInputCloud(_lastCornerCloud);
      _lastSurfaceKDTree.setInputCloud(_lastSurfaceCloud);

      _transformSum.rot_x += _imuPitchStart;
      _transformSum.rot_z += _imuRollStart;

      _systemInited = true;
      return;
   }

   pcl::PointXYZI coeff;
   bool isDegenerate = false;
   Eigen::Matrix<float, 6, 6> matP;

   _frameCount++;
   _transform.pos -= _imuVeloFromStart * _scanPeriod;

   size_t lastCornerCloudSize = _lastCornerCloud->points.size();
   size_t lastSurfaceCloudSize = _lastSurfaceCloud->points.size();

   if (lastCornerCloudSize > 10 && lastSurfaceCloudSize > 100)
   {
      std::vector<int> pointSearchInd(1);
      std::vector<float> pointSearchSqDis(1);
      std::vector<int> indices;

      pcl::removeNaNFromPointCloud(*_cornerPointsSharp, *_cornerPointsSharp, indices);
      size_t cornerPointsSharpNum = _cornerPointsSharp->points.size();
      size_t surfPointsFlatNum = _surfPointsFlat->points.size();

      _pointSearchCornerInd1.resize(cornerPointsSharpNum);
      _pointSearchCornerInd2.resize(cornerPointsSharpNum);
      _pointSearchSurfInd1.resize(surfPointsFlatNum);
      _pointSearchSurfInd2.resize(surfPointsFlatNum);
      _pointSearchSurfInd3.resize(surfPointsFlatNum);

      for (size_t iterCount = 0; iterCount < _maxIterations; iterCount++)
      {
         pcl::PointXYZI pointSel, pointProj, tripod1, tripod2, tripod3;
         _laserCloudOri->clear();
         _coeffSel->clear();

         for (int i = 0; i < cornerPointsSharpNum; i++)
         {
            transformToStart(_cornerPointsSharp->points[i], pointSel);

            if (iterCount % 5 == 0)
            {
               pcl::removeNaNFromPointCloud(*_lastCornerCloud, *_lastCornerCloud, indices);
               _lastCornerKDTree.nearestKSearch(pointSel, 1, pointSearchInd, pointSearchSqDis);

               int closestPointInd = -1, minPointInd2 = -1;
               if (pointSearchSqDis[0] < 25)
               {
                  closestPointInd = pointSearchInd[0];
                  int closestPointScan = int(_lastCornerCloud->points[closestPointInd].intensity);

                  float pointSqDis, minPointSqDis2 = 25;
                  for (int j = closestPointInd + 1; j < cornerPointsSharpNum; j++)
                  {
                     if (int(_lastCornerCloud->points[j].intensity) > closestPointScan + 2.5)
                     {
                        break;
                     }

                     pointSqDis = calcSquaredDiff(_lastCornerCloud->points[j], pointSel);

                     if (int(_lastCornerCloud->points[j].intensity) > closestPointScan)
                     {
                        if (pointSqDis < minPointSqDis2)
                        {
                           minPointSqDis2 = pointSqDis;
                           minPointInd2 = j;
                        }
                     }
                  }
                  for (int j = closestPointInd - 1; j >= 0; j--)
                  {
                     if (int(_lastCornerCloud->points[j].intensity) < closestPointScan - 2.5)
                     {
                        break;
                     }

                     pointSqDis = calcSquaredDiff(_lastCornerCloud->points[j], pointSel);

                     if (int(_lastCornerCloud->points[j].intensity) < closestPointScan)
                     {
                        if (pointSqDis < minPointSqDis2)
                        {
                           minPointSqDis2 = pointSqDis;
                           minPointInd2 = j;
                        }
                     }
                  }
               }

               _pointSearchCornerInd1[i] = closestPointInd;
               _pointSearchCornerInd2[i] = minPointInd2;
            }

            if (_pointSearchCornerInd2[i] >= 0)
            {
               tripod1 = _lastCornerCloud->points[_pointSearchCornerInd1[i]];
               tripod2 = _lastCornerCloud->points[_pointSearchCornerInd2[i]];

               float x0 = pointSel.x;
               float y0 = pointSel.y;
               float z0 = pointSel.z;
               float x1 = tripod1.x;
               float y1 = tripod1.y;
               float z1 = tripod1.z;
               float x2 = tripod2.x;
               float y2 = tripod2.y;
               float z2 = tripod2.z;

               float a012 = sqrt(((x0 - x1)*(y0 - y2) - (x0 - x2)*(y0 - y1))
                                 * ((x0 - x1)*(y0 - y2) - (x0 - x2)*(y0 - y1))
                                 + ((x0 - x1)*(z0 - z2) - (x0 - x2)*(z0 - z1))
                                 * ((x0 - x1)*(z0 - z2) - (x0 - x2)*(z0 - z1))
                                 + ((y0 - y1)*(z0 - z2) - (y0 - y2)*(z0 - z1))
                                 * ((y0 - y1)*(z0 - z2) - (y0 - y2)*(z0 - z1)));

               float l12 = sqrt((x1 - x2)*(x1 - x2) + (y1 - y2)*(y1 - y2) + (z1 - z2)*(z1 - z2));

               float la = ((y1 - y2)*((x0 - x1)*(y0 - y2) - (x0 - x2)*(y0 - y1))
                           + (z1 - z2)*((x0 - x1)*(z0 - z2) - (x0 - x2)*(z0 - z1))) / a012 / l12;

               float lb = -((x1 - x2)*((x0 - x1)*(y0 - y2) - (x0 - x2)*(y0 - y1))
                            - (z1 - z2)*((y0 - y1)*(z0 - z2) - (y0 - y2)*(z0 - z1))) / a012 / l12;

               float lc = -((x1 - x2)*((x0 - x1)*(z0 - z2) - (x0 - x2)*(z0 - z1))
                            + (y1 - y2)*((y0 - y1)*(z0 - z2) - (y0 - y2)*(z0 - z1))) / a012 / l12;

               float ld2 = a012 / l12; // Eq. (2)

               // TODO: Why writing to a variable that's never read?
               pointProj = pointSel;
               pointProj.x -= la * ld2;
               pointProj.y -= lb * ld2;
               pointProj.z -= lc * ld2;

               float s = 1;
               if (iterCount >= 5)
               {
                  s = 1 - 1.8f * fabs(ld2);
               }

               coeff.x = s * la;
               coeff.y = s * lb;
               coeff.z = s * lc;
               coeff.intensity = s * ld2;

               if (s > 0.1 && ld2 != 0)
               {
                  _laserCloudOri->push_back(_cornerPointsSharp->points[i]);
                  _coeffSel->push_back(coeff);
               }
            }
         }

         for (int i = 0; i < surfPointsFlatNum; i++)
         {
            transformToStart(_surfPointsFlat->points[i], pointSel);

            if (iterCount % 5 == 0)
            {
               _lastSurfaceKDTree.nearestKSearch(pointSel, 1, pointSearchInd, pointSearchSqDis);
               int closestPointInd = -1, minPointInd2 = -1, minPointInd3 = -1;
               if (pointSearchSqDis[0] < 25)
               {
                  closestPointInd = pointSearchInd[0];
                  int closestPointScan = int(_lastSurfaceCloud->points[closestPointInd].intensity);

                  float pointSqDis, minPointSqDis2 = 25, minPointSqDis3 = 25;
                  for (int j = closestPointInd + 1; j < surfPointsFlatNum; j++)
                  {
                     if (int(_lastSurfaceCloud->points[j].intensity) > closestPointScan + 2.5)
                     {
                        break;
                     }

                     pointSqDis = calcSquaredDiff(_lastSurfaceCloud->points[j], pointSel);

                     if (int(_lastSurfaceCloud->points[j].intensity) <= closestPointScan)
                     {
                        if (pointSqDis < minPointSqDis2)
                        {
                           minPointSqDis2 = pointSqDis;
                           minPointInd2 = j;
                        }
                     }
                     else
                     {
                        if (pointSqDis < minPointSqDis3)
                        {
                           minPointSqDis3 = pointSqDis;
                           minPointInd3 = j;
                        }
                     }
                  }
                  for (int j = closestPointInd - 1; j >= 0; j--)
                  {
                     if (int(_lastSurfaceCloud->points[j].intensity) < closestPointScan - 2.5)
                     {
                        break;
                     }

                     pointSqDis = calcSquaredDiff(_lastSurfaceCloud->points[j], pointSel);

                     if (int(_lastSurfaceCloud->points[j].intensity) >= closestPointScan)
                     {
                        if (pointSqDis < minPointSqDis2)
                        {
                           minPointSqDis2 = pointSqDis;
                           minPointInd2 = j;
                        }
                     }
                     else
                     {
                        if (pointSqDis < minPointSqDis3)
                        {
                           minPointSqDis3 = pointSqDis;
                           minPointInd3 = j;
                        }
                     }
                  }
               }

               _pointSearchSurfInd1[i] = closestPointInd;
               _pointSearchSurfInd2[i] = minPointInd2;
               _pointSearchSurfInd3[i] = minPointInd3;
            }

            if (_pointSearchSurfInd2[i] >= 0 && _pointSearchSurfInd3[i] >= 0)
            {
               tripod1 = _lastSurfaceCloud->points[_pointSearchSurfInd1[i]];
               tripod2 = _lastSurfaceCloud->points[_pointSearchSurfInd2[i]];
               tripod3 = _lastSurfaceCloud->points[_pointSearchSurfInd3[i]];

               float pa = (tripod2.y - tripod1.y) * (tripod3.z - tripod1.z)
                  - (tripod3.y - tripod1.y) * (tripod2.z - tripod1.z);
               float pb = (tripod2.z - tripod1.z) * (tripod3.x - tripod1.x)
                  - (tripod3.z - tripod1.z) * (tripod2.x - tripod1.x);
               float pc = (tripod2.x - tripod1.x) * (tripod3.y - tripod1.y)
                  - (tripod3.x - tripod1.x) * (tripod2.y - tripod1.y);
               float pd = -(pa * tripod1.x + pb * tripod1.y + pc * tripod1.z);

               float ps = sqrt(pa * pa + pb * pb + pc * pc);
               pa /= ps;
               pb /= ps;
               pc /= ps;
               pd /= ps;

               float pd2 = pa * pointSel.x + pb * pointSel.y + pc * pointSel.z + pd; //Eq. (3)??

               // TODO: Why writing to a variable that's never read? Maybe it should be used afterwards?
               pointProj = pointSel;
               pointProj.x -= pa * pd2;
               pointProj.y -= pb * pd2;
               pointProj.z -= pc * pd2;

               float s = 1;
               if (iterCount >= 5)
               {
                  s = 1 - 1.8f * fabs(pd2) / sqrt(calcPointDistance(pointSel));
               }

               coeff.x = s * pa;
               coeff.y = s * pb;
               coeff.z = s * pc;
               coeff.intensity = s * pd2;

               if (s > 0.1 && pd2 != 0)
               {
                  _laserCloudOri->push_back(_surfPointsFlat->points[i]);
                  _coeffSel->push_back(coeff);
               }
            }
         }

         int pointSelNum = _laserCloudOri->points.size();
         if (pointSelNum < 10)
         {
            continue;
         }

         Eigen::Matrix<float, Eigen::Dynamic, 6> matA(pointSelNum, 6);
         Eigen::Matrix<float, 6, Eigen::Dynamic> matAt(6, pointSelNum);
         Eigen::Matrix<float, 6, 6> matAtA;
         Eigen::VectorXf matB(pointSelNum);
         Eigen::Matrix<float, 6, 1> matAtB;
         Eigen::Matrix<float, 6, 1> matX;

         for (int i = 0; i < pointSelNum; i++)
         {
            const pcl::PointXYZI& pointOri = _laserCloudOri->points[i];
            coeff = _coeffSel->points[i];

            float s = 1;

            float srx = sin(s * _transform.rot_x.rad());
            float crx = cos(s * _transform.rot_x.rad());
            float sry = sin(s * _transform.rot_y.rad());
            float cry = cos(s * _transform.rot_y.rad());
            float srz = sin(s * _transform.rot_z.rad());
            float crz = cos(s * _transform.rot_z.rad());
            float tx = s * _transform.pos.x();
            float ty = s * _transform.pos.y();
            float tz = s * _transform.pos.z();

            float arx = (-s * crx*sry*srz*pointOri.x + s * crx*crz*sry*pointOri.y + s * srx*sry*pointOri.z
                         + s * tx*crx*sry*srz - s * ty*crx*crz*sry - s * tz*srx*sry) * coeff.x
               + (s*srx*srz*pointOri.x - s * crz*srx*pointOri.y + s * crx*pointOri.z
                  + s * ty*crz*srx - s * tz*crx - s * tx*srx*srz) * coeff.y
               + (s*crx*cry*srz*pointOri.x - s * crx*cry*crz*pointOri.y - s * cry*srx*pointOri.z
                  + s * tz*cry*srx + s * ty*crx*cry*crz - s * tx*crx*cry*srz) * coeff.z;

            float ary = ((-s * crz*sry - s * cry*srx*srz)*pointOri.x
                         + (s*cry*crz*srx - s * sry*srz)*pointOri.y - s * crx*cry*pointOri.z
                         + tx * (s*crz*sry + s * cry*srx*srz) + ty * (s*sry*srz - s * cry*crz*srx)
                         + s * tz*crx*cry) * coeff.x
               + ((s*cry*crz - s * srx*sry*srz)*pointOri.x
                  + (s*cry*srz + s * crz*srx*sry)*pointOri.y - s * crx*sry*pointOri.z
                  + s * tz*crx*sry - ty * (s*cry*srz + s * crz*srx*sry)
                  - tx * (s*cry*crz - s * srx*sry*srz)) * coeff.z;

            float arz = ((-s * cry*srz - s * crz*srx*sry)*pointOri.x + (s*cry*crz - s * srx*sry*srz)*pointOri.y
                         + tx * (s*cry*srz + s * crz*srx*sry) - ty * (s*cry*crz - s * srx*sry*srz)) * coeff.x
               + (-s * crx*crz*pointOri.x - s * crx*srz*pointOri.y
                  + s * ty*crx*srz + s * tx*crx*crz) * coeff.y
               + ((s*cry*crz*srx - s * sry*srz)*pointOri.x + (s*crz*sry + s * cry*srx*srz)*pointOri.y
                  + tx * (s*sry*srz - s * cry*crz*srx) - ty * (s*crz*sry + s * cry*srx*srz)) * coeff.z;

            float atx = -s * (cry*crz - srx * sry*srz) * coeff.x + s * crx*srz * coeff.y
               - s * (crz*sry + cry * srx*srz) * coeff.z;

            float aty = -s * (cry*srz + crz * srx*sry) * coeff.x - s * crx*crz * coeff.y
               - s * (sry*srz - cry * crz*srx) * coeff.z;

            float atz = s * crx*sry * coeff.x - s * srx * coeff.y - s * crx*cry * coeff.z;

            float d2 = coeff.intensity;

            matA(i, 0) = arx;
            matA(i, 1) = ary;
            matA(i, 2) = arz;
            matA(i, 3) = atx;
            matA(i, 4) = aty;
            matA(i, 5) = atz;
            matB(i, 0) = -0.05 * d2;
         }
         matAt = matA.transpose();
         matAtA = matAt * matA;
         matAtB = matAt * matB;

         matX = matAtA.colPivHouseholderQr().solve(matAtB);

         if (iterCount == 0)
         {
            Eigen::Matrix<float, 1, 6> matE;
            Eigen::Matrix<float, 6, 6> matV;
            Eigen::Matrix<float, 6, 6> matV2;

            Eigen::SelfAdjointEigenSolver< Eigen::Matrix<float, 6, 6> > esolver(matAtA);
            matE = esolver.eigenvalues().real();
            matV = esolver.eigenvectors().real();

            matV2 = matV;

            isDegenerate = false;
            float eignThre[6] = { 10, 10, 10, 10, 10, 10 };
            for (int i = 0; i < 6; i++)
            {
               if (matE(0, i) < eignThre[i])
               {
                  for (int j = 0; j < 6; j++)
                  {
                     matV2(i, j) = 0;
                  }
                  isDegenerate = true;
               }
               else
               {
                  break;
               }
            }
            matP = matV.inverse() * matV2;
         }

         if (isDegenerate)
         {
            Eigen::Matrix<float, 6, 1> matX2(matX);
            matX = matP * matX2;
         }

         _transform.rot_x = _transform.rot_x.rad() + matX(0, 0);
         _transform.rot_y = _transform.rot_y.rad() + matX(1, 0);
         _transform.rot_z = _transform.rot_z.rad() + matX(2, 0);
         _transform.pos.x() += matX(3, 0);
         _transform.pos.y() += matX(4, 0);
         _transform.pos.z() += matX(5, 0);

         if (!pcl_isfinite(_transform.rot_x.rad())) _transform.rot_x = Angle();
         if (!pcl_isfinite(_transform.rot_y.rad())) _transform.rot_y = Angle();
         if (!pcl_isfinite(_transform.rot_z.rad())) _transform.rot_z = Angle();

         if (!pcl_isfinite(_transform.pos.x())) _transform.pos.x() = 0.0;
         if (!pcl_isfinite(_transform.pos.y())) _transform.pos.y() = 0.0;
         if (!pcl_isfinite(_transform.pos.z())) _transform.pos.z() = 0.0;

         float deltaR = sqrt(pow(rad2deg(matX(0, 0)), 2) +
                             pow(rad2deg(matX(1, 0)), 2) +
                             pow(rad2deg(matX(2, 0)), 2));
         float deltaT = sqrt(pow(matX(3, 0) * 100, 2) +
                             pow(matX(4, 0) * 100, 2) +
                             pow(matX(5, 0) * 100, 2));

         if (deltaR < _deltaRAbort && deltaT < _deltaTAbort)
            break;
      }
   }

   Angle rx, ry, rz;
   accumulateRotation(_transformSum.rot_x,
                      _transformSum.rot_y,
                      _transformSum.rot_z,
                      -_transform.rot_x,
                      -_transform.rot_y.rad() * 1.05,
                      -_transform.rot_z,
                      rx, ry, rz);

   Vector3 v(_transform.pos.x() - _imuShiftFromStart.x(),
             _transform.pos.y() - _imuShiftFromStart.y(),
             _transform.pos.z() * 1.05 - _imuShiftFromStart.z());
   rotateZXY(v, rz, rx, ry);
   Vector3 trans = _transformSum.pos - v;

   pluginIMURotation(rx, ry, rz,
                     _imuPitchStart, _imuYawStart, _imuRollStart,
                     _imuPitchEnd, _imuYawEnd, _imuRollEnd,
                     rx, ry, rz);

   _transformSum.rot_x = rx;
   _transformSum.rot_y = ry;
   _transformSum.rot_z = rz;
   _transformSum.pos = trans;

   transformToEnd(_cornerPointsLessSharp);
   transformToEnd(_surfPointsLessFlat);

   _cornerPointsLessSharp.swap(_lastCornerCloud);
   _surfPointsLessFlat.swap(_lastSurfaceCloud);

   lastCornerCloudSize = _lastCornerCloud->points.size();
   lastSurfaceCloudSize = _lastSurfaceCloud->points.size();

   if (lastCornerCloudSize > 10 && lastSurfaceCloudSize > 100)
   {
      _lastCornerKDTree.setInputCloud(_lastCornerCloud);
      _lastSurfaceKDTree.setInputCloud(_lastSurfaceCloud);
   }

}

} // end namespace loam


/*

    Copyright (C) 2014 K. Carter Dodd (carter.dodd@gmail.com)

    This program is free software; you can redistribute it and/or
    modify it under the terms of the GNU General Public License as
    published by the Free Software Foundation; either version 2 of
    the License, or (at your option) any later version.

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

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

*/

#include "RadialSolver.hpp"
#include "Constants.hpp"

#include <cmath>
#include <algorithm>

#include <boost/thread/thread.hpp>
#include <boost/bind.hpp>

#include <iostream>

#include "dynarray.hpp"

#include <time.h>

using namespace std;

//----------------------------------------------------------------------------------------------------
// RadialSolver()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
RadialSolver<
    numr,
    numz,
    numv>::RadialSolver(
        double minr,
        double maxr,
        double maxz,
        double mass,
        double T0,
        const double (&Li)[numr],
        const std::string& M_file,
        const std::string& N_file,
        double up_down_symmetry,
        double left_right_symmetry)
    :
        mMinr(minr),
        mMaxr(maxr),
        mMaxz(maxz),
        mMass(mass),
        mT0(T0),
        mVth(sqrt(2*BOLTZMANN_CONST*mT0/mMass)),
        mBD_matrix_compiled(false),
        mUpDownSymmetry(up_down_symmetry),
        mLeftRightSymmetry(left_right_symmetry)
    {
        for(unsigned int nr = 0; nr < numr_bd; ++nr)
        {

            mSolN_BD_top[nr] = 0.5;
            mSolN_BD_top_tmp[nr] = 0.5;

            mSolN_BD_bottom[nr] = 0.5;
            mSolN_BD_bottom_tmp[nr] = 0.5;
        }

        for (unsigned int nz = 0; nz < numz_bd; ++nz)
        {
            mSolN_BD_outer[nz]=0.5;
            mSolN_BD_outer_tmp[nz]=0.5;

            mSolN_BD_inner[nz]=0.5;
            mSolN_BD_inner_tmp[nz]=0.5;
        }

        for(unsigned int nr = 0; nr < numr; ++nr)
        {

            mLi[nr] = Li[nr];

            for (unsigned int nz = 0; nz < numz_tot; ++nz)
            {

                    mSolutionN[nr][nz] = 1;
                    mSolutionVr[nr][nz] = 0;
                    mSolutionVz[nr][nz] = 0;

            }
        }

        for (unsigned int i=0; i < num_bd; ++i)
        {
            for (unsigned int j=0; j < num_bd; ++j)
            {
                mBD_matrix[i][j] = 0;
            }
        }

        mM.load_text_file(M_file);
        mN.load_text_file(N_file);

        mDtheta = PI_CONST/(total_numv-1);
        mDphi = PI_CONST/(total_numv-1);

        mDr = (mMaxr - mMinr)/(numr-1);

        mDz = mMaxz/(numz-1);

    }

//----------------------------------------------------------------------------------------------------
// solutionN()
//----------------------------------------------------------------------------------------------------
template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::solutionN(unsigned int nr, unsigned int nz)
    {
        return mSolutionN[nr][nz];
    }

//----------------------------------------------------------------------------------------------------
// solutionVr()
//----------------------------------------------------------------------------------------------------
template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::solutionVr(unsigned int nr, unsigned int nz)
    {
        return mSolutionVr[nr][nz];
    }

//----------------------------------------------------------------------------------------------------
// solutionVz()
//----------------------------------------------------------------------------------------------------
template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::solutionVz(unsigned int nr, unsigned int nz)
    {
        return mSolutionVz[nr][nz];
    }

//----------------------------------------------------------------------------------------------------
// boundaryN()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::boundaryN(
        unsigned int n,
        unsigned int bd)
    {
        // outer
        if (bd == bdOUTER)
        {
            return mSolN_BD_outer[n];
        }

        // inner
        if (bd == bdINNER)
        {
            return mSolN_BD_inner[n];
        }

        // top
        if (bd == bdTOP)
        {
            return mSolN_BD_top[n];
        }

        // bottom
        if (bd == bdBOTTOM)
        {
            return mSolN_BD_bottom[n];
        }

        return -1;
    }

//----------------------------------------------------------------------------------------------------
// ionLevel()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::avg_ion_rate()
    {

        double sum=0;
        double weights=0;
        double r_weight = 0;

        for (unsigned int nr = 1; nr < numr-1; ++nr)
        {
            r_weight = nr*mDr+mMinr;

            for(unsigned int nz=1; nz < numz_tot-1; ++nz)
            {
                sum += mLi[nr]*mSolutionN[nr][nz]*r_weight;
                weights += mSolutionN[nr][nz]*r_weight;
            }
        }

        double avg = sum/weights;

        return avg;

    }

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::avg_n()
    {

        double sum=0;
        double weights=0;
        double r_weight = 0;

        for (unsigned int nr = 1; nr < numr-1; ++nr)
        {
            r_weight = nr*mDr+mMinr;

            for(unsigned int nz=1; nz < numz_tot-1; ++nz)
            {
                sum += mSolutionN[nr][nz]*r_weight;
                weights += r_weight;
            }
        }

        double avg = sum/weights;

        return avg;

    }

//----------------------------------------------------------------------------------------------------
// solve_bd()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::solve_bd(
        double precision,
        unsigned int max_iterations,
        bool multithread)
    {

        if (!mBD_matrix_compiled)
        {
            cerr << "First time solving boundary: compiling boundary matrix." << endl;

            // must first compile boundary matrix

            // set boundary values to unity to compile matrix so that values can be scaled to solution
            for(unsigned int nr = 0; nr < numr_bd; ++nr)
            {
                mSolN_BD_top[nr] = 1;
                mSolN_BD_bottom[nr] = 1;

            }

            for (unsigned int nz = 0; nz < numz_bd; ++nz)
            {
                mSolN_BD_outer[nz]=1;
                mSolN_BD_inner[nz]=1;
            }

            if (multithread)
            {
                // solve in parallel

                boost::thread_group calcThreads;

                for (unsigned int nz = 0; nz < numz_bd; ++nz)
                {
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::compile_bd, this, nz, bdOUTER));
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::compile_bd, this, nz, bdINNER));
                }

                for (unsigned int nr = 0; nr < numr_bd; ++nr)
                {
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::compile_bd, this, nr, bdTOP));
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::compile_bd, this, nr, bdBOTTOM));
                }

                calcThreads.join_all();

            }else{

                // solve serially
                for (unsigned int nz = 0; nz < numz_bd; ++nz)
                {
                    compile_bd(nz, bdOUTER);
                    compile_bd(nz, bdINNER);
                }

                for (unsigned int nr = 0; nr < numr_bd; ++nr)
                {
                    compile_bd(nr, bdTOP);
                    compile_bd(nr, bdBOTTOM);
                }
            }

            mBD_matrix_compiled = true;
            cerr << "Boundary matrix compiled." << endl;
        }

        unsigned int iterations = 0;

        double maxdev;

        // keep track of how long the solution takes to compute
        time_t t_begin = time(0);

        do{
            //cerr << "Solving boundary: ";

            for(unsigned int nr = 0; nr < numr_bd; ++nr)
            {
                mSolN_BD_top_tmp[nr] = mSolN_BD_top[nr];
                mSolN_BD_bottom_tmp[nr] = mSolN_BD_bottom[nr];
            }

            for (unsigned int nz = 0; nz < numz_bd; ++nz)
            {
                mSolN_BD_outer_tmp[nz]=mSolN_BD_outer[nz];
                mSolN_BD_inner_tmp[nz]=mSolN_BD_inner[nz];
            }

            // solve boundary

            if (multithread)
            {
                // solve in parallel

                boost::thread_group calcThreads;

                for (unsigned int nz = 0; nz < numz_bd; ++nz)
                {
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::solve_bd, this, nz, bdOUTER));
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::solve_bd, this, nz, bdINNER));
                }

                for (unsigned int nr = 0; nr < numr_bd; ++nr)
                {
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::solve_bd, this, nr, bdTOP));
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::solve_bd, this, nr, bdBOTTOM));
                }

                calcThreads.join_all();

            }else{

                // solve serially
                for (unsigned int nz = 0; nz < numz_bd; ++nz)
                {
                    solve_bd(nz, bdOUTER);
                    solve_bd(nz, bdINNER);
                }

                for (unsigned int nr = 0; nr < numr_bd; ++nr)
                {
                    solve_bd(nr, bdTOP);
                    solve_bd(nr, bdBOTTOM);
                }
            }

            maxdev=0;

            for(unsigned int nr = 0; nr < numr_bd; ++nr)
            {
                double dev = 2*abs(mSolN_BD_top_tmp[nr] - mSolN_BD_top[nr])/(mSolN_BD_top_tmp[nr] + mSolN_BD_top[nr]);

                maxdev = max(maxdev, dev);

                dev = 2*abs(mSolN_BD_bottom_tmp[nr] - mSolN_BD_bottom[nr])/(mSolN_BD_bottom_tmp[nr] + mSolN_BD_bottom[nr]);

                maxdev = max(maxdev, dev);
            }

            for (unsigned int nz = 0; nz < numz_bd; ++nz)
            {
                double dev = 2*abs(mSolN_BD_outer_tmp[nz] - mSolN_BD_outer[nz])/(mSolN_BD_outer_tmp[nz] + mSolN_BD_outer[nz]);

                maxdev = max(maxdev, dev);

                dev = 2*abs(mSolN_BD_inner_tmp[nz] - mSolN_BD_inner[nz])/(mSolN_BD_inner_tmp[nz] + mSolN_BD_inner[nz]);

                maxdev = max(maxdev, dev);

            }

            iterations++;

            //time_t t_cur = time(0);

            //cerr << maxdev << " deviation, iteration " << iterations << ", time = " << t_cur - t_begin << "s." << endl;

        }while(maxdev > precision && iterations < max_iterations);

        time_t t_end = time(0);

        if (iterations == max_iterations)
        {
            cerr << "Max iterations reached in attempt to solve boundary." << endl;
        }else{
            cerr << "Solution for boundary found to " << maxdev << " deviation in " << iterations << " steps and " << t_end - t_begin << "s." << endl;
        }

    }

//----------------------------------------------------------------------------------------------------
// solve_bd()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::solve_bd(
        unsigned int n,
        unsigned int bd)
    {

        unsigned int this_bd_index = bd_index(n, bd);

        double out_flux = 0;

        for(unsigned int nr = 0; nr < numr_bd; ++nr)
        {
            out_flux += mSolN_BD_top_tmp[nr]*mBD_matrix[this_bd_index][bd_index(nr, bdTOP)];
            out_flux += mSolN_BD_bottom_tmp[nr]*mBD_matrix[this_bd_index][bd_index(nr, bdBOTTOM)];
        }

        for (unsigned int nz = 0; nz < numz_bd; ++nz)
        {
            out_flux += mSolN_BD_outer_tmp[nz]*mBD_matrix[this_bd_index][bd_index(nz, bdOUTER)];
            out_flux += mSolN_BD_inner_tmp[nz]*mBD_matrix[this_bd_index][bd_index(nz, bdINNER)];
        }

        if (bd == bdTOP)
        {
            double r = n*mDr+mMinr;

            double source_symmetric = pow(neTe_profile_function(r, 1.0, 0.4), 0.5)*neTe_profile_function(r, 1.1, 0.5)*r*2/(mMinr + mMaxr);
            double source_asymmetric = max(-sin(TWOPI_CONST*(r-mMinr)/(mMaxr-mMinr)), 0.0);

            mSolN_BD_top[n] = out_flux*sqrt(PI_CONST)/mVth + source_symmetric /* +  mUpDownSymmetry*source_symmetric + (1.0-mUpDownSymmetry)*mLeftRightSymmetry*source_asymmetric*/;

        }else if (bd == bdBOTTOM){

            double r = n*mDr+mMinr;

            double source_symmetric = pow(neTe_profile_function(r, 1.0, 0.4), 0.5)*neTe_profile_function(r, 1.1, 0.5)*r*2/(mMinr + mMaxr);
            double source_asymmetric = max(sin(TWOPI_CONST*(r-mMinr)/(mMaxr-mMinr)), 0.0);

            mSolN_BD_bottom[n] = out_flux*sqrt(PI_CONST)/mVth/* +  mUpDownSymmetry*source_symmetric + (1.0-mUpDownSymmetry)*(1.0-mLeftRightSymmetry)*source_asymmetric*/;

        }else if (bd == bdINNER){

            mSolN_BD_inner[n] = out_flux*sqrt(PI_CONST)/mVth;

        }else{

            mSolN_BD_outer[n] = out_flux*sqrt(PI_CONST)/mVth;

        }

    }

//----------------------------------------------------------------------------------------------------
// compile_bd()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::compile_bd(
        unsigned int n,
        unsigned int bd)
    {

        double r, z;
        unsigned int min_theta, max_theta, min_phi, max_phi;
        double total_theta;
        bool need_vr_vz; // vr == true, vz == false
        double dot_product_sign;

        // outer
        if (bd == bdOUTER)
        {
            r = mMaxr;
            z = n*mDz - mMaxz + 0.5*mDz;

            need_vr_vz = true;

            min_theta = 0;
            max_theta = numv;
            total_theta = PI_CONST;

            dot_product_sign = 1;

            min_phi = 0;
            max_phi = total_numv;

        }

        // inner
        if (bd == bdINNER)
        {

            r = mMinr;
            z = n*mDz - mMaxz + 0.5*mDz;;

            need_vr_vz = true;

            min_theta = numv-1;
            max_theta = total_numv;

            total_theta = PI_CONST;

            dot_product_sign = -1;

            min_phi = 0;
            max_phi = total_numv;

        }

        // top
        if (bd == bdTOP)
        {
            r = n*mDr + mMinr + 0.5*mDr;
            z = mMaxz;

            need_vr_vz = false;

            min_phi = 0;
            max_phi = numv;

            dot_product_sign = 1;

            min_theta = 0;
            max_theta = total_numv;

            total_theta = TWOPI_CONST;

        }

        // bottom
        if (bd == bdBOTTOM)
        {
            r = n*mDr + mMinr + 0.5*mDr;
            z = -mMaxz;

            need_vr_vz = false;

            min_phi = numv-1;
            max_phi = total_numv;

            dot_product_sign = -1;

            min_theta = 0;
            max_theta = total_numv;

            total_theta = TWOPI_CONST;

        }

        double I1 = 0;
        double I2 = 0;
        double I = 0;
        unsigned int bdi;

        double int_fact;
        double int_edge_fact;

        double V_factor = 2.0*mVth/pow(PI_CONST, 3.0/2.0);

        double theta, phi;

        unsigned int this_bd_index = bd_index(n, bd);

        for(unsigned int phi_i = min_phi; phi_i < max_phi; ++phi_i)
        {
            if (phi_i == 0)
            {
                int_fact = total_theta*(1-sin(mDphi)/mDphi);

                integrate(r, z, 0, 0, I1, I2, bdi);

                if (need_vr_vz)
                {
                    I = 0;
                }else{
                    I = dot_product_sign*int_fact*V_factor*I2;
                }

                mBD_matrix[this_bd_index][bdi] += I;

            }else if (phi_i == total_numv-1){

                int_fact = total_theta*(1-sin(mDphi)/mDphi);

                integrate(r, z, 0, PI_CONST, I1, I2, bdi);

                if (need_vr_vz)
                {
                    I = 0;
                }else{
                    I = -dot_product_sign*int_fact*V_factor*I2;
                }

                mBD_matrix[this_bd_index][bdi] += I;

            }else{

                phi = phi_i*mDphi;

                if (phi_i == min_phi)
                {
                    // if it got here then min_phi != 0, so integration factor is different

                    int_fact = mDtheta*(cos(phi) - (sin(phi+mDphi)-sin(phi))/mDphi);

                }else if (phi_i == max_phi){
                    // same deal except for other limit

                    int_fact = mDtheta*((sin(phi)-sin(phi-mDphi))/mDphi - cos(phi));

                }else{

                    // not on edge of phi integration
                    int_fact = 2*mDtheta*sin(phi)*(1-cos(mDphi))/mDphi;
                }

                for (unsigned int theta_i = min_theta; theta_i < max_theta; ++theta_i)
                {
                    if (theta_i == 0)
                    {
                        integrate(r, z, 0, phi, I1, I2, bdi);

                        if (need_vr_vz)
                        {
                            I = dot_product_sign*sin(phi)*int_fact*V_factor*I2;
                        }else{
                            I = dot_product_sign*cos(phi)*int_fact*V_factor*I2;
                        }

                        mBD_matrix[this_bd_index][bdi] += I;

                    }else if (theta_i == total_numv-1){

                        integrate(r, z, PI_CONST, phi, I1, I2, bdi);

                        if (need_vr_vz)
                        {
                            I = -dot_product_sign*sin(phi)*int_fact*V_factor*I2;
                        }else{
                            I = dot_product_sign*cos(phi)*int_fact*V_factor*I2;
                        }

                        mBD_matrix[this_bd_index][bdi] += I;

                    }else{

                        if (theta_i == min_theta || min_theta == max_theta)
                        {
                            // if limits are not 0 and pi, then need to alter integration factor
                            int_edge_fact = 1/2;
                        }else{
                            int_edge_fact = 1;
                        }

                        theta = theta_i*mDtheta;

                        integrate(r, z, theta, phi, I1, I2, bdi);

                        if (need_vr_vz)
                        {
                            I = 2*dot_product_sign*int_edge_fact*sin(phi)*cos(theta)*int_fact*V_factor*I2;
                        }else{
                            I = 2*dot_product_sign*int_edge_fact*cos(phi)*int_fact*V_factor*I2;
                        }

                        mBD_matrix[this_bd_index][bdi] += I;

                    }

                }
            }

        }

    }

//----------------------------------------------------------------------------------------------------
// solve()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::solve(bool multithread)
    {
        if (multithread)
        {
            // solve in parallel

            boost::thread_group calcThreads;

            for (unsigned int nr = 1; nr < numr-1; ++nr)
            {
                for (unsigned int nz = 1; nz < numz_tot-1; ++nz)
                {
                    calcThreads.create_thread(boost::bind(&RadialSolver<numr,numz,numv>::solve, this, nr, nz));
                }

            }

            calcThreads.join_all();

        }else{
            // solve serially

            for (unsigned int nr = 1; nr < numr-1; ++nr)
            {
                for (unsigned int nz = 1; nz < numz_tot-1; ++nz)
                {

                    solve(nr, nz);
                }

            }
        }
    }

//----------------------------------------------------------------------------------------------------
// solve()
//----------------------------------------------------------------------------------------------------

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::solve(
        unsigned int nr,
        unsigned int nz)
    {

        double r = nr*mDr + mMinr;
        double z = nz*mDz - mMaxz;

        // reset to zero

        mSolutionN[nr][nz] = 0;
        mSolutionVr[nr][nz] = 0;
        mSolutionVz[nr][nz] = 0;

        double I1 = 0;
        double I2 = 0;
        unsigned int bdi;

        double int_fact;
        double N_factor = 2.0/pow(PI_CONST, 3.0/2.0);
        double V_factor = 2.0*mVth/pow(PI_CONST, 3.0/2.0);
        double theta, phi;

        for(unsigned int phi_i = 0; phi_i < total_numv; ++phi_i)
        {
            if (phi_i == 0)
            {
                int_fact = TWOPI_CONST*(1-sin(mDphi)/mDphi);

                integrate(r, z, 0, 0, I1, I2, bdi);

                mSolutionN[nr][nz] += int_fact*N_factor*I1;
                mSolutionVz[nr][nz] += int_fact*V_factor*I2;

            }else if (phi_i == total_numv-1){

                int_fact = TWOPI_CONST*(1-sin(mDphi)/mDphi);

                integrate(r, z, 0, PI_CONST, I1, I2, bdi);

                mSolutionN[nr][nz] += int_fact*N_factor*I1;
                mSolutionVz[nr][nz] += -int_fact*V_factor*I2;

            }else{

                phi = phi_i*mDphi;

                int_fact = 2*mDtheta*sin(phi)*(1-cos(mDphi))/mDphi;

                for (unsigned int theta_i = 0; theta_i < total_numv; ++theta_i)
                {
                    if (theta_i == 0)
                    {
                        integrate(r, z, 0, phi, I1, I2, bdi);

                        mSolutionN[nr][nz] += int_fact*N_factor*I1;
                        mSolutionVr[nr][nz] += sin(phi)*int_fact*V_factor*I2;
                        mSolutionVz[nr][nz] += cos(phi)*int_fact*V_factor*I2;

                    }else if (theta_i == total_numv-1){

                        integrate(r, z, PI_CONST, phi, I1, I2, bdi);

                        mSolutionN[nr][nz] += int_fact*N_factor*I1;
                        mSolutionVr[nr][nz] += -sin(phi)*int_fact*V_factor*I2;
                        mSolutionVz[nr][nz] += cos(phi)*int_fact*V_factor*I2;

                    }else{

                        theta = theta_i*mDtheta;

                        integrate(r, z, theta, phi, I1, I2, bdi);

                        mSolutionN[nr][nz] += 2*int_fact*N_factor*I1;
                        mSolutionVr[nr][nz] += 2*sin(phi)*cos(theta)*int_fact*V_factor*I2;
                        mSolutionVz[nr][nz] += 2*cos(phi)*int_fact*V_factor*I2;
                    }

                }
            }

        }

    }

//----------------------------------------------------------------------------------------------------
// integrate()
//----------------------------------------------------------------------------------------------------
template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
void RadialSolver<
    numr,
    numz,
    numv>::integrate(
        double r,
        double z,
        double theta,
        double phi,
        double& nM,
        double& nN,
        unsigned int& bdi)
    {
        // setup variables defining integral path

        double sintheta = sin(theta);
        double costheta = cos(theta);

        double sinphi = sin(phi);
        double cosphi = cos(phi);

        double theta_transition = asin(mMinr/r);

        double L_perp = -1;

        if (theta <= theta_transition)
        {
            //inner wall
            L_perp = r*costheta - sqrt(mMinr*mMinr-r*r*sintheta*sintheta);

        }else{
            // outer wall
            L_perp = r*costheta + sqrt(mMaxr*mMaxr-r*r*sintheta*sintheta);
        }

        double h1 = mMaxz+z;
        double h2 = mMaxz-z;

        double phi1 = atan(L_perp/h1);
        double phi2 = atan(h2/L_perp) + HALFPI_CONST;

        double L = -1;

        if (z == mMaxz && phi == HALFPI_CONST)
        {
            L = L_perp;

        }else if ((r == mMinr || r == mMaxr) && phi == 0){

            L = h1;

        }else if ((r == mMinr || r == mMaxr) && phi == PI_CONST){

            L = h2;

        }else{

            if (phi < phi1)
            {
                // floor
                L = h1/cosphi;

            }else if (phi < phi2){
                // wall
                L = L_perp/sinphi;

            }else{
                // ceiling
                L = -h2/cosphi;
            }
        }

        // get value at wall using linear interpolation
        double n0_wall = 0;

        double r_intersect = min(max(sqrt(r*r + L*L*sinphi*sinphi - 2*r*L*sinphi*costheta), mMinr), mMaxr);
        double delta_r = r - r_intersect;
        double delta_r2 = delta_r*delta_r;
        double delta_z = -L*cosphi;
        double delta_z2 = delta_z*delta_z;
        double z_intersect = min(max(z + delta_z, -mMaxz), mMaxz);

        double frac1, frac2;
        unsigned int nri, nzi;
        unsigned int nri1, nri2;

        if (phi < phi1 || phi > phi2)
        {
            // floor or ceiling

            nri = floor((r_intersect - mMinr)/mDr);

            if (phi < phi1)
            {
                // bottom

                n0_wall = mSolN_BD_bottom[nri];

                bdi = bd_index(nri, bdBOTTOM);

            }else{

                // top

                n0_wall = mSolN_BD_top[nri];

                bdi = bd_index(nri, bdTOP);
            }

        }else{
            // wall

            nzi = floor((z_intersect+mMaxz)/mDz);

            if (theta <= theta_transition)
            {
                //inner wall

                n0_wall = mSolN_BD_inner[nzi];

                bdi = bd_index(nzi, bdINNER);

            }else{

                // outer wall

                n0_wall = mSolN_BD_outer[nzi];

                bdi = bd_index(nzi, bdOUTER);

            }

        }

        if (n0_wall < 0)
        {
            cerr << "n0_wall not defined" << endl;
        }

        unsigned int N = 200;
        double dl = L/((double) N-1);
        double l = 0;
        double ri;
        double x;

        if (L == 0.0)
        {
            cerr << "L = 0" << endl;
        }

        dynarray<double> ion_rate_function(N, 0.0, L);

        for (unsigned int i=0; i < N; ++i)
        {
            // calculate element indecies
            l = i*dl;
            x = l - L;

            ri = sqrt(r*r + x*x*sinphi*sinphi + 2*r*x*sinphi*costheta);

            if (ri > mMaxr)
            {
                //cerr << "r > " << mMaxr << ": " << ri << endl;

                ri = mMaxr;
            }

            if (ri < mMinr)
            {
                //cerr << "r < " << mMinr << ": " << ri << endl;

                ri = mMinr;
            }

            nri1 = floor((ri - mMinr)/mDr);
            nri2 = nri1+1;
            frac2 = (ri - (mDr*nri1 + mMinr))/mDr;
            frac1 = 1-frac2;

            ion_rate_function(i) = frac1*mLi[nri1] + frac2*mLi[nri2];

            //ion_rate_function(i) = 0;

        }

        double vk = ion_rate_function.integral_total();

        double alpha = vk/mVth;

        if (alpha > mM.xmax())
        {
            nM = 0.0;
            nN = 0.0;
        }else{
            nM = n0_wall*mM.interp(alpha);
            nN = n0_wall*mN.interp(alpha);
        }

    }

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
unsigned int RadialSolver<
    numr,
    numz,
    numv>::
    bd_index(
        unsigned int n,
        unsigned int bd)
    {
        if (bd == bdOUTER)
        {
            return n;
        }

        if (bd == bdINNER)
        {
            return n + numz_bd;
        }

        if (bd == bdTOP)
        {
            return n + 2*numz_bd;
        }

        if (bd == bdBOTTOM)
        {
            return n + 2*numz_bd + numr_bd;
        }

        throw -1;
    }

template<
    unsigned int numr,
    unsigned int numz,
    unsigned int numv>
double RadialSolver<
    numr,
    numz,
    numv>::
    neTe_profile_function(
        double r,
        double rp,
        double gamma)
    {
        double xp = (rp - mMinr)/(mMaxr - mMinr);
        double x = (r - mMinr)/(mMaxr - mMinr);

        double a = (0.5 - xp)/(xp*xp - xp);
        double y = a*x*x + (1-a)*x;

        return (gamma/2)*(1-cos(PI_CONST*y)) + (1-gamma/2)*pow(sin(PI_CONST*y), 3);
    }


#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <math.h>

#include "dlib/data_io.h"
#include "dlib/string.h"
#include <dlib/image_transforms.h>
#include <dlib/dir_nav.h>
#include <dlib/dnn.h>
#include <dlib/data_io.h>
#include <dlib/image_processing/frontal_face_detector.h>

namespace py = pybind11;
using namespace dlib;

const char* VERSION = "1.0.1";

// Age Predictor
// ----------------------------------------------------------------------------------------

// This block of statements defines a Resnet-10 architecture for the age predictor.
// We will use 81 classes (0-80 years old) to predict the age of a face.
const unsigned long number_of_age_classes = 81;

// The resnet basic block.
template<
	int num_filters,
	template<typename> class BN,  // some kind of batch normalization or affine layer
	int stride,
	typename SUBNET
>
using basicblock = BN<con<num_filters, 3, 3, 1, 1, relu<BN<con<num_filters, 3, 3, stride, stride, SUBNET>>>>>;

// A residual making use of the skip layer mechanism.
template<
	template<int, template<typename> class, int, typename> class BLOCK,  // a basic block defined before
	int num_filters,
	template<typename> class BN,  // some kind of batch normalization or affine layer
	typename SUBNET
> // adds the block to the result of tag1 (the subnet)
using residual = add_prev1<BLOCK<num_filters, BN, 1, tag1<SUBNET>>>;

// A residual that does subsampling (we need to subsample the output of the subnet, too).
template<
	template<int, template<typename> class, int, typename> class BLOCK,  // a basic block defined before
	int num_filters,
	template<typename> class BN,
	typename SUBNET
>
using residual_down = add_prev2<avg_pool<2, 2, 2, 2, skip1<tag2<BLOCK<num_filters, BN, 2, tag1<SUBNET>>>>>>;

// Residual block with optional downsampling and batch normalization.
template<
	template<template<int, template<typename> class, int, typename> class, int, template<typename>class, typename> class RESIDUAL,
	template<int, template<typename> class, int, typename> class BLOCK,
	int num_filters,
	template<typename> class BN,
	typename SUBNET
>
using residual_block = relu<RESIDUAL<BLOCK, num_filters, BN, SUBNET>>;

template<int num_filters, typename SUBNET>
using aresbasicblock_down = residual_block<residual_down, basicblock, num_filters, affine, SUBNET>;

// Some useful definitions to design the affine versions for inference.
template<typename SUBNET> using aresbasicblock256 = residual_block<residual, basicblock, 256, affine, SUBNET>;
template<typename SUBNET> using aresbasicblock128 = residual_block<residual, basicblock, 128, affine, SUBNET>;
template<typename SUBNET> using aresbasicblock64  = residual_block<residual, basicblock, 64, affine, SUBNET>;

// Common input for standard resnets.
template<typename INPUT>
using aresnet_input = max_pool<3, 3, 2, 2, relu<affine<con<64, 7, 7, 2, 2, INPUT>>>>;

// Resnet-10 architecture for estimating.
template<typename SUBNET>
using aresnet10_level1 = aresbasicblock256<aresbasicblock_down<256, SUBNET>>;
template<typename SUBNET>
using aresnet10_level2 = aresbasicblock128<aresbasicblock_down<128, SUBNET>>;
template<typename SUBNET>
using aresnet10_level3 = aresbasicblock64<SUBNET>;
// The resnet 10 backbone.
template<typename INPUT>
using aresnet10_backbone = avg_pool_everything<
	aresnet10_level1<
	aresnet10_level2<
	aresnet10_level3<
	aresnet_input<INPUT>>>>>;

using apredictor_t = loss_multiclass_log<fc<number_of_age_classes, aresnet10_backbone<input_rgb_image>>>;

// Helper function to estimage the age
uint8_t get_estimated_age(matrix<float, 1, number_of_age_classes>& p, float& confidence)
{
	float estimated_age = (0.25f * p(0));
	confidence = p(0);

	for (uint16_t i = 1; i < number_of_age_classes; i++) {
		estimated_age += (static_cast<float>(i) * p(i));
		if (p(i) > confidence) confidence = p(i);
	}

	return std::lround(estimated_age);
}

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

// This block of statements defines the network as proposed for the CNN Model I.
// We however removed the "dropout" regularization on the activations of convolutional layers.
template <int N, template <typename> class BN, int stride, typename SUBNET>
using block = BN<con<N, 3, 3, stride, stride, relu<BN<con<N, 3, 3, stride, stride, SUBNET>>>>>;

template <int N, typename SUBNET> using res_ = relu<block<N, bn_con, 1, SUBNET>>;
template <int N, typename SUBNET> using ares_ = relu<block<N, affine, 1, SUBNET>>;

template <typename SUBNET> using alevel1 = avg_pool<2, 2, 2, 2, ares_<64, SUBNET>>;
template <typename SUBNET> using alevel2 = avg_pool<2, 2, 2, 2, ares_<32, SUBNET>>;

using agender_type = loss_multiclass_log<fc<2, multiply<relu<fc<16, multiply<alevel1<alevel2< input_rgb_image_sized<32>>>>>>>>>;

enum label_ : uint16_t
{
	female_label,
	male_label,
};

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

// convert a numpy array into a dlib::matrix<dlib::rgb_pixel>
const dlib::matrix<dlib::rgb_pixel> from_numpy(const py::array_t<unsigned char>& ndarray) {
    const int num_channels = 3;
    auto raw = ndarray.unchecked<num_channels>();
    std::vector<unsigned char> rgb(3);
    dlib::matrix<dlib::rgb_pixel> image(raw.shape(0), raw.shape(1));
    for (ssize_t i = 0; i < raw.shape(0); i++)
    {
        for (ssize_t j = 0; j < raw.shape(1); j++)
        {
            image(i, j) = dlib::rgb_pixel(
                static_cast<unsigned char>(raw(i, j, 0)),
                static_cast<unsigned char>(raw(i, j, 1)),
                static_cast<unsigned char>(raw(i, j, 2)));
        }
    }
    return image;
}

class AgeAndGender {

	public:
		virtual ~AgeAndGender() { }
		virtual void load_shape_predictor(std::string filename);
		virtual void load_dnn_gender_classifier(std::string filename);
		virtual void load_dnn_age_predictor(std::string filename);
		virtual py::list predict(
			const py::array_t<unsigned char>& photo_numpy_array,
			py::list face_bounding_boxes
		);
		virtual std::vector<dlib::rectangle> from_py_list_with_tuple_to_vector_with_rectangles(py::list face_bounding_boxes);

	private:
      shape_predictor sp;
      agender_type gender_predictor_net;
      apredictor_t age_predictor_net;
      frontal_face_detector detector;
};

void AgeAndGender::load_shape_predictor(std::string filename) {
	detector = get_frontal_face_detector();
	deserialize(filename) >> sp;
}

void AgeAndGender::load_dnn_gender_classifier(std::string filename) {
	deserialize(filename) >> gender_predictor_net;
}

void AgeAndGender::load_dnn_age_predictor(std::string filename) {
	deserialize(filename) >> age_predictor_net;
}

std::vector<dlib::rectangle> AgeAndGender::from_py_list_with_tuple_to_vector_with_rectangles(py::list face_bounding_boxes) {
	std::vector<dlib::rectangle> result;
	for (auto item : face_bounding_boxes) {
		std::vector<int> face_box;
		for(py::handle num : item) {
			face_box.push_back(num.cast<int>());
		}

		result.push_back(
			rectangle(face_box[3], face_box[0], face_box[1], face_box[2])
		);
	}

	return result;
}

py::list AgeAndGender::predict(
	const py::array_t<unsigned char>& photo_numpy_array,
	py::list face_bounding_boxes = py::list()
) {
	// Load the source image.
	matrix<rgb_pixel> in = from_numpy(photo_numpy_array);
	//load_image(in, photo);

	// As proposed in the paper, use Softmax for the last layer.
	softmax<agender_type::subnet_type> gender_snet;
	gender_snet.subnet() = gender_predictor_net.subnet();

	softmax<apredictor_t::subnet_type> age_snet;
	age_snet.subnet() = age_predictor_net.subnet();

	// Evaluate the gender
	int32_t cur_face = 0;
	float confidence = 0.0f;
	enum label_ gender;
	py::list result;

	std::vector<dlib::rectangle> all_faces;
	if(face_bounding_boxes.size()) {
		all_faces = AgeAndGender::from_py_list_with_tuple_to_vector_with_rectangles(
			face_bounding_boxes
		);
	} else {
		all_faces = detector(in);
	}

	for (auto face : all_faces) {
		auto shape = sp(in, face);

		if (shape.num_parts()) {
			matrix<rgb_pixel> face_chip;

			extract_image_chip(in, get_face_chip_details(shape, 32), face_chip);
			matrix<float, 1, 2> p = mat(gender_snet(face_chip));

			if (p(0) < p(1)) {
				gender = male_label;
				confidence = p(1);
			} else {
				gender = female_label;
				confidence = p(0);
			}

			py::list face_rect;
			face_rect.append(face.left());
			face_rect.append(face.top());
			face_rect.append(face.right());
			face_rect.append(face.bottom());

			// Age predictor
			extract_image_chip(in, get_face_chip_details(shape, 64), face_chip);
			matrix<float, 1, number_of_age_classes> p_age = mat(age_snet(face_chip));

			result.append(py::dict());
			result[cur_face]["gender"] = py::dict();
			result[cur_face]["gender"]["value"] = ((gender == female_label) ? "female" : "male");
			result[cur_face]["gender"]["confidence"] = (int)floorf(confidence * 100.0f);

			result[cur_face]["age"] = py::dict();
			result[cur_face]["age"]["value"] = get_estimated_age(p_age, confidence);
			result[cur_face]["age"]["confidence"] = (int)floorf(confidence * 100.0f);
			result[cur_face]["face"] = face_rect;
			cur_face++;
		}
	}

	return result;
}

PYBIND11_MODULE(age_and_gender, m) {
	m.doc() = "Predict Age and Gender using Python";
	m.attr("__version__") = VERSION;
    py::class_<AgeAndGender>(m, "AgeAndGender")
        .def("load_shape_predictor", &AgeAndGender::load_shape_predictor)
        .def("load_dnn_gender_classifier", &AgeAndGender::load_dnn_gender_classifier)
        .def("load_dnn_age_predictor", &AgeAndGender::load_dnn_age_predictor)
        .def(
			"predict",
			&AgeAndGender::predict,
			py::arg("photo_numpy_array"),
			py::arg("face_bounding_boxes") = py::list()
		)
        .def(py::init<>());
}


/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#include <itkImageFileWriter.h>
#include <itkResampleImageFilter.h>
#include <itkFiniteDiffOdfMaximaExtractionFilter.h>

#include <mitkBaseDataIOFactory.h>
#include <mitkDiffusionImage.h>
#include <mitkQBallImage.h>
#include <mitkImageCast.h>
#include <mitkImageToItk.h>
#include <mitkTensorImage.h>

#include <mitkCoreObjectFactory.h>
#include <mitkFiberBundleXWriter.h>
#include <itkShCoefficientImageImporter.h>
#include <itkFlipImageFilter.h>
#include <boost/lexical_cast.hpp>
#include <boost/algorithm/string.hpp>
#include <mitkTestingMacros.h>
#include <mitkIOUtil.h>

using namespace std;

int mitkPeakExtractionTest(int argc, char* argv[])
{
    MITK_TEST_BEGIN("mitkStreamlineTrackingTest");

    MITK_TEST_CONDITION_REQUIRED(argc>3,"check for input data")

    string shCoeffFileName = argv[1];
    string maskFileName = argv[2];
    string referenceFileName = argv[3];

    MITK_INFO << "SH-coefficient file: " << shCoeffFileName;
    MITK_INFO << "Mask file: " << maskFileName;
    MITK_INFO << "Reference fiber file: " << referenceFileName;

    try
    {
        mitk::CoreObjectFactory::GetInstance();

        mitk::Image::Pointer image = mitk::IOUtil::LoadImage(shCoeffFileName);
        mitk::Image::Pointer mitkMaskImage = mitk::IOUtil::LoadImage(maskFileName);

        typedef itk::Image<unsigned char, 3>  ItkUcharImgType;
        typedef itk::FiniteDiffOdfMaximaExtractionFilter< float, 4, 20242 > MaximaExtractionFilterType;
        MaximaExtractionFilterType::Pointer filter = MaximaExtractionFilterType::New();

        MITK_INFO << "Casting mask image ...";
        ItkUcharImgType::Pointer itkMask = ItkUcharImgType::New();
        mitk::CastToItkImage<ItkUcharImgType>(mitkMaskImage, itkMask);
        filter->SetMaskImage(itkMask);

        MITK_INFO << "Casting SH image ...";
        typedef mitk::ImageToItk< MaximaExtractionFilterType::CoefficientImageType > CasterType;
        CasterType::Pointer caster = CasterType::New();
        caster->SetInput(image);
        caster->Update();
        filter->SetInput(caster->GetOutput());
        filter->SetMaxNumPeaks(2);
        filter->SetPeakThreshold(0.4);
        filter->SetAbsolutePeakThreshold(0.01);
        filter->SetAngularThreshold(25);
        filter->SetNormalizationMethod(MaximaExtractionFilterType::MAX_VEC_NORM);
        filter->SetNumberOfThreads(1);
        MITK_INFO << "Starting extraction ...";
        filter->Update();
        mitk::FiberBundleX::Pointer fib1 = filter->GetOutputFiberBundle();

        MITK_INFO << "Loading reference ...";
        const std::string s1="", s2="";
        std::vector<mitk::BaseData::Pointer> infile = mitk::BaseDataIO::LoadBaseDataFromFile( referenceFileName, s1, s2, false );
        mitk::FiberBundleX::Pointer fib2 = dynamic_cast<mitk::FiberBundleX*>(infile.at(0).GetPointer());

        // TODO: reduce epsilon. strange issues with differing values between windows and linux.
        MITK_TEST_CONDITION_REQUIRED(fib1->Equals(fib2), "Check if tractograms are equal.");
    }
    catch (itk::ExceptionObject e)
    {
        MITK_INFO << e;
        return EXIT_FAILURE;
    }
    catch (std::exception e)
    {
        MITK_INFO << e.what();
        return EXIT_FAILURE;
    }
    catch (...)
    {
        MITK_INFO << "ERROR!?!";
        return EXIT_FAILURE;
    }

    MITK_TEST_END();
}


#include <iostream>
#include <vector>

#include <algorithm>
#include <unordered_map>
#include <set>
#include <utility>
#include <stdexcept>
#include <cmath>
#include <chrono>
#include <boost/filesystem.hpp>

#include "dynet/lstm.h"
#include "dynet/training.h"
#include "dynet/dynet.h"
#include "dynet/dict.h"

#include <boost/functional/hash.hpp>
#include <boost/program_options.hpp>
#include <boost/serialization/unordered_set.hpp>

#include "serialization.h"
#include "nn/tagger.h"
#include "utils.h"

#include "dependency.h"
#include "conll.h"
#include "activation_function.h"

int main(int argc, char **argv)
{
    std::string train_path;
    std::string dev_path;
    std::string model_path;
    unsigned n_iteration;

    bool probabilistic;

    bool eval_on_dev;

    bool dropout;
    double dropout_p;

    std::string ignore_dynet_mem;
    std::string ignore_dynet_wd;

    unsigned n_stack;
    unsigned n_layer;
    int lstm_dim;
    int word_dim;

    ActivationFunctionOption activation_function;

    namespace po = boost::program_options;
    po::options_description desc("Options");
    desc.add_options()
        ("train", po::value<std::string>(&train_path)->required(), "")
        ("model", po::value<std::string>(&model_path)->required(), "")
        ("iteration", po::value<unsigned>(&n_iteration)->default_value(20), "")
        ("eval-on-dev", po::value<bool>(&eval_on_dev)->default_value(false), "")
        ("dev-path", po::value<std::string>(&dev_path)->default_value(""), "")
        // dropout
        ("dropout", po::value<bool>(&dropout)->default_value(false), "")
        ("dropout-p", po::value<double>(&dropout_p)->default_value(0.5), "")
        // dynet
        ("dynet-mem", po::value<std::string>(&ignore_dynet_mem), "")
        ("dynet-weight-decay", po::value<std::string>(&ignore_dynet_wd), "")
        // nn options
        ("activation-function", po::value<ActivationFunctionOption>(&activation_function), "")
        ("lstm-dim", po::value<int>(&lstm_dim)->default_value(125))
        ("word-dim", po::value<int>(&word_dim)->default_value(100))
        ("n-stack", po::value<unsigned>(&n_stack)->default_value(1u), "")
        ("n-layer", po::value<unsigned>(&n_layer)->default_value(1u), "")
        ("probabilistic", po::value<bool>(&probabilistic)->default_value(false), "")
    ;

    po::positional_options_description pod;
    pod.add("train", 1);
    pod.add("model", 1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).options(desc).positional(pod).run(), vm);
    po::notify(vm);

    dynet::initialize(argc, argv);

    ConllSettings conll_settings;
    read_object(model_path + ".conll_settings.param", conll_settings);

    std::cerr << "Reading train data..." << std::endl << std::flush;
    Conll conll_train(conll_settings);
    std::vector<IntSentence> train_data;
    conll_train.read(train_path);
    conll_train.as_int_sentence([&](const IntSentence& s) { train_data.push_back(s); });

    std::vector<IntSentence> dev_data;
    if (eval_on_dev)
    {
        Conll conll_dev(conll_settings);
        conll_dev.read(dev_path);
        conll_dev.as_int_sentence([&](const IntSentence& s) { dev_data.push_back(s); });
    }

    // TODO: other trainer
    dynet::Model model;
    dynet::AdamTrainer trainer(model);
    trainer.sparse_updates_enabled = false;

    RNNSettings node_rnn_settings;
    node_rnn_settings.padding = true;
    node_rnn_settings.dim = lstm_dim;
    node_rnn_settings.n_stack = n_stack;
    node_rnn_settings.n_layer = n_layer;

    NeuralTaggerSettings nn_settings;
    nn_settings.n_word = conll_settings.word_dict.size();
    nn_settings.n_pos = conll_settings.pos_dict.size();
    nn_settings.word_unknown = conll_settings.word_dict.convert("*UNKNOWN*");
    nn_settings.word_dim = word_dim;
    nn_settings.activation_function = activation_function.value;

    NeuralTagger<dynet::LSTMBuilder> rnn(model, nn_settings, node_rnn_settings);

    for (auto const& sentence : train_data)
    {
        // Count word frequence (used for dropout)
        for (auto const& token : sentence)
            rnn.m_word_count[token.word] += 1;
    }

    double loss = 0.0;

    save_object(model_path + ".tagger.nn_settings", nn_settings);
    save_object(model_path + ".tagger.rnn_settings", node_rnn_settings);

    std::vector<unsigned> dev_prec;

    for (unsigned iteration = 0 ; iteration <= n_iteration ; ++iteration)
    {
        std::cerr << "Iteration: " << iteration << std::endl;
        std::cerr << std::flush;

        unsigned n_correct = 0;
        unsigned n_total = 0;

        std::random_shuffle(std::begin(train_data), std::end(train_data));

        for (const IntSentence& sentence : train_data)
        {
            dynet::ComputationGraph cg;

            // loss output
            std::vector<Expression> errs;

            rnn.compute(
                cg,
                sentence,
                [&] (unsigned index, dynet::expr::Expression& expr) -> void
                {
                    if (probabilistic)
                    {
                        errs.push_back(pickneglogsoftmax(expr, sentence[index+1].pos));

                        auto vec = dynet::as_vector(expr.value());
                        int predicted = std::distance(std::begin(vec), std::max_element(std::begin(vec), std::end(vec)));
                        if (predicted == sentence[index+1].pos)
                            n_correct += 1;
                    }
                    else
                    {
                        dynet::expr::Expression expr2 = tanh(expr);
                        auto vec = dynet::as_vector(expr2.value());

                        // loss augmented inference
                        for (unsigned i = 0u ; i < vec.size() ; ++i)
                        {
                            if ((int) i != sentence[index+1].pos)
                                vec[i] += 1.0;
                        }

                        int predicted = std::distance(std::begin(vec), std::max_element(std::begin(vec), std::end(vec)));
                        if (predicted == sentence[index+1].pos)
                        {
                            n_correct += 1;
                        }
                        else
                        {
                            errs.push_back(dynet::expr::pick(expr2, predicted) - pick(expr2, sentence[index+1].pos));
                        }
                    }
                },
                true, // word dropout
                dropout,
                dropout_p
            );

            n_total += sentence.size();

            // backprop
            if (errs.size() > 0)
            {
                Expression sum_errs = dynet::expr::sum(errs);
                loss += as_scalar(cg.get_value(sum_errs.i));
                cg.backward(sum_errs.i);
                trainer.update(1.0);
            }
        }

        trainer.update_epoch();
        trainer.status();

        std::cerr << "E = " << (loss / (double) n_total) << " ppl=" << exp(loss / (double) n_total) << " (acc=" << n_correct / (double) n_total << ")" << std::endl;
        std::cerr << std::flush;

        save_object(model_path + ".tagger.param." + std::to_string(iteration), model);

        // Evaluate on dev set
        if (eval_on_dev)
        {
            unsigned n_good_pos = 0u;
            unsigned n_total_pos = 0u;

            std::cerr << "Evaluating on dev..." << std::endl << std::flush;

            for (auto const& sentence : dev_data)
            {
                dynet::ComputationGraph cg;

                rnn.compute(
                    cg,
                    sentence,
                    [&] (unsigned index, dynet::expr::Expression& expr) -> void
                    {
                        unused_parameter(expr);

                        if (probabilistic)
                        {
                            auto vec = dynet::as_vector(dynet::expr::softmax(expr).value());
                            int predicted = std::distance(std::begin(vec), std::max_element(std::begin(vec), std::end(vec)));
                            if (predicted == sentence[index+1].pos)
                                ++ n_good_pos;
                        }
                        else
                        {
                            auto vec = dynet::as_vector(dynet::expr::tanh(expr).value());
                            int predicted = std::distance(std::begin(vec), std::max_element(std::begin(vec), std::end(vec)));
                            if (predicted == sentence[index+1].pos)
                                ++ n_good_pos;
                        }
                    }
                );

                n_total_pos += sentence.size();
            }

            std::cerr << " -> POS precision: " << n_good_pos / (double) n_total_pos << std::endl;
            dev_prec.push_back(n_good_pos);
        }
        std::cerr << "------" << std::endl;
    }

    if (eval_on_dev)
    {
        unsigned max = std::distance(std::begin(dev_prec), std::max_element(std::begin(dev_prec), std::end(dev_prec)));
        boost::filesystem::copy_file(
                model_path + ".tagger.param." + std::to_string(max),
                model_path + ".tagger.param",
                boost::filesystem::copy_option::overwrite_if_exists
        );

    }
    else
    {
        // juste resave the last model
        save_object(model_path + ".tagger.param", model);
    }
}


// Copyright (c) 2016-2020 The ZCash developers
// Copyright (c) 2020 The PIVX developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "test/test_iqcash.h"

#include <array>
#include <stdexcept>

#include "sapling/note.hpp"
#include "sapling/noteencryption.hpp"
#include "sapling/prf.h"
#include "sapling/address.hpp"
#include "sapling/util.h"
#include "crypto/sha256.h"

#include <boost/test/unit_test.hpp>
#include <librustzcash.h>
#include <sodium.h>

BOOST_FIXTURE_TEST_SUITE(noteencryption_tests, BasicTestingSetup)

class TestNoteDecryption : public ZCNoteDecryption {
public:
    TestNoteDecryption(uint256 sk_enc) : ZCNoteDecryption(sk_enc) {}

    void change_pk_enc(uint256 to) {
        pk_enc = to;
    }
};

BOOST_AUTO_TEST_CASE(note_plain_text)
{
    using namespace libzcash;
    auto xsk = SaplingSpendingKey(uint256()).expanded_spending_key();
    auto fvk = xsk.full_viewing_key();
    auto ivk = fvk.in_viewing_key();
    SaplingPaymentAddress addr = *ivk.address({0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});

    std::array<unsigned char, ZC_MEMO_SIZE> memo;
    for (size_t i = 0; i < ZC_MEMO_SIZE; i++) {
        // Fill the message with dummy data
        memo[i] = (unsigned char) i;
    }

    SaplingNote note(addr, 39393);
    auto cmu_opt = note.cm();
    if (!cmu_opt) {
        BOOST_ERROR("SaplingNote cm failed");
    }
    uint256 cmu = cmu_opt.get();
    SaplingNotePlaintext pt(note, memo);

    auto res = pt.encrypt(addr.pk_d);
    if (!res) {
        BOOST_ERROR("SaplingNotePlaintext encrypt failed");
    }

    auto enc = res.get();

    auto ct = enc.first;
    auto encryptor = enc.second;
    auto epk = encryptor.get_epk();

    // Try to decrypt with incorrect commitment
    BOOST_CHECK(!SaplingNotePlaintext::decrypt(
        ct,
        ivk,
        epk,
        uint256()
    ));

    // Try to decrypt with correct commitment
    auto foo = SaplingNotePlaintext::decrypt(
        ct,
        ivk,
        epk,
        cmu
    );

    if (!foo) {
        BOOST_ERROR("SaplingNotePlaintext decrypt failed");
    }

    auto bar = foo.get();

    BOOST_CHECK(bar.value() == pt.value());
    BOOST_CHECK(bar.memo() == pt.memo());
    BOOST_CHECK(bar.d == pt.d);
    BOOST_CHECK(bar.rcm == pt.rcm);

    auto foobar = bar.note(ivk);

    if (!foobar) {
        // Improve this message
        BOOST_ERROR("Invalid note");
    }

    auto new_note = foobar.get();

    BOOST_CHECK(note.value() == new_note.value());
    BOOST_CHECK(note.d == new_note.d);
    BOOST_CHECK(note.pk_d == new_note.pk_d);
    BOOST_CHECK(note.r == new_note.r);
    BOOST_CHECK(note.cm() == new_note.cm());

    SaplingOutgoingPlaintext out_pt;
    out_pt.pk_d = note.pk_d;
    out_pt.esk = encryptor.get_esk();

    auto ovk = random_uint256();
    auto cv = random_uint256();
    auto cm = random_uint256();

    auto out_ct = out_pt.encrypt(
        ovk,
        cv,
        cm,
        encryptor
    );

    auto decrypted_out_ct = out_pt.decrypt(
        out_ct,
        ovk,
        cv,
        cm,
        encryptor.get_epk()
    );

    if (!decrypted_out_ct) {
        BOOST_ERROR("SaplingOutgoingPlaintext decrypt failed");
    }

    auto decrypted_out_ct_unwrapped = decrypted_out_ct.get();

    BOOST_CHECK(decrypted_out_ct_unwrapped.pk_d == out_pt.pk_d);
    BOOST_CHECK(decrypted_out_ct_unwrapped.esk == out_pt.esk);

    // Test sender won't accept invalid commitments
    BOOST_CHECK(!SaplingNotePlaintext::decrypt(
        ct,
        epk,
        decrypted_out_ct_unwrapped.esk,
        decrypted_out_ct_unwrapped.pk_d,
        uint256()
    ));

    // Test sender can decrypt the note ciphertext.
    foo = SaplingNotePlaintext::decrypt(
        ct,
        epk,
        decrypted_out_ct_unwrapped.esk,
        decrypted_out_ct_unwrapped.pk_d,
        cmu
    );

    if (!foo) {
        BOOST_ERROR("Sender decrypt note ciphertext failed.");
    }

    bar = foo.get();

    BOOST_CHECK(bar.value() == pt.value());
    BOOST_CHECK(bar.memo() == pt.memo());
    BOOST_CHECK(bar.d == pt.d);
    BOOST_CHECK(bar.rcm == pt.rcm);
}

BOOST_AUTO_TEST_SUITE_END()


/*    Copyright (c) 2010-2019, Delft University of Technology
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 *    References
 *      Musegaas, P. (2012). Optimization of Space Trajectories Including Multiple Gravity Assists
 *          and Deep Space Maneuvers. MSc Thesis, Delft University of Technology, Delft,
 *          The Netherlands.
 *
 */

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MAIN

#include <limits>

#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/unit_test.hpp>

#include <Eigen/Core>

#include <tudat/astro/basic_astro/physicalConstants.h>
#include <tudat/basics/testMacros.h>

#include "tudat/astro/ephemerides/constantEphemeris.h"
#include "tudat/astro/mission_segments/transferNode.h"
#include "tudat/astro/mission_segments/transferLeg.h"

namespace tudat
{
namespace unit_tests
{

//! Test implementation of departure leg within MGA trajectory model
BOOST_AUTO_TEST_SUITE( test_departure_leg_mga_1dsm_position )

//! Test delta-V computation
BOOST_AUTO_TEST_CASE( testVelocities )
{
    // Set tolerance. Due to iterative nature in the process, (much) higher accuracy cannot be
    // achieved.
    const double tolerance = 1.0e-6;

    // Expected test result based on the first leg of the ideal Messenger trajectory as modelled by
    // GTOP software distributed and downloadable from the ESA website, or within the PaGMO
    // Astrotoolbox.
    const double expectedDsmDeltaV = 910.801673341206;
    const Eigen::Vector3d expectedEndVelocity ( 17969.3166254715, -23543.6915939142,
                                             6.38384671663485 );

    // Specify the required parameters.
    // Set the planetary positions and velocities.

    const Eigen::Vector6d planet1State =
            ( Eigen::Vector6d( ) <<
              -148689402143.081, 7454242895.97607, 0.0,
              -1976.48781307596, -29863.4035321021, 0.0 ).finished( );
    std::shared_ptr< ephemerides::Ephemeris > constantEphemeris1 =
            std::make_shared< ephemerides::ConstantEphemeris >( planet1State );

    const Eigen::Vector6d planet2State =
            ( Eigen::Vector6d( ) <<
              -128359548637.032, -78282803797.7343, 0.,
              TUDAT_NAN, TUDAT_NAN, TUDAT_NAN ).finished( );
    std::shared_ptr< ephemerides::Ephemeris > constantEphemeris2 =
            std::make_shared< ephemerides::ConstantEphemeris >( planet2State );

    // Set the time of flight, which has to be converted from JD (in GTOP) to seconds (in Tudat).
    const double timeOfFlight = 399.999999715 * physical_constants::JULIAN_DAY;

    // Set the Sun gravitational parameter.
    const double sunGravitationalParameter = 1.32712428e20;

    // Set DSM parameters.
    const double dsmTimeOfFlightFraction = 0.234594654679;
    const double inPlaneAngle = 1.69629206466940;
    const double outOfPlaneAngle = 0.00019299969467;
    const double dimensionlessRadiusDsm = 0.915891737859598;

    Eigen::VectorXd legFreeParameters =
            ( Eigen::VectorXd( 6 ) << 0.0, timeOfFlight,
              dsmTimeOfFlightFraction, dimensionlessRadiusDsm, inPlaneAngle, outOfPlaneAngle ).finished( );

    // Set test case.
    mission_segments::DsmPositionBasedTransferLeg transferLeg(
           constantEphemeris1, constantEphemeris2,
            sunGravitationalParameter );
    transferLeg.updateLegParameters( legFreeParameters );

    BOOST_CHECK_CLOSE_FRACTION( expectedDsmDeltaV, transferLeg.getLegDeltaV( ), tolerance );
    TUDAT_CHECK_MATRIX_CLOSE_FRACTION( expectedEndVelocity, transferLeg.getArrivalVelocity( ), tolerance );

    double dsmTime = transferLeg.getDsmTime( );
    double timeTolerance = 1.0E-9;
    Eigen::Vector3d dsmLocation = transferLeg.getDsmLocation( );

    // Test state output from departure to DSM
    {
        // Get data on 10 equispace points on trajectory
        std::map< double, Eigen::Vector6d > statesAlongTrajectory;
        std::vector< double > outputTimes = utilities::linspace( 0.0, transferLeg.getDsmTime( ) - timeTolerance, 10 );
        transferLeg.getStatesAlongTrajectory( statesAlongTrajectory, outputTimes );

        // Check initial and final time on output list
        BOOST_CHECK_SMALL( statesAlongTrajectory.begin( )->first, 1.0E-14 );
        BOOST_CHECK_CLOSE_FRACTION( statesAlongTrajectory.rbegin( )->first, dsmTime - timeTolerance, 1.0E-14 );

        // Check if Keplerian state (slow elements) is the same for each output point
        Eigen::Vector6d previousKeplerianState = Eigen::Vector6d::Constant( TUDAT_NAN );
        for( auto it : statesAlongTrajectory )
        {
            Eigen::Vector6d currentCartesianState = it.second;
            Eigen::Vector6d currentKeplerianState = tudat::orbital_element_conversions::convertCartesianToKeplerianElements(
                        currentCartesianState, sunGravitationalParameter );
            if( previousKeplerianState == previousKeplerianState )
            {
                TUDAT_CHECK_MATRIX_CLOSE_FRACTION(
                            ( previousKeplerianState.segment( 0, 5 ) ),
                            ( currentKeplerianState.segment( 0, 5 ) ),
                            1.0E-14 );

            }
            previousKeplerianState = currentKeplerianState;
        }

        // Check if output meets boundary conditions
        for( int i = 0; i < 3; i++ )
        {
            BOOST_CHECK_SMALL( std::fabs( statesAlongTrajectory.begin( )->second( i ) - planet1State( i ) ), 1.0E-2 );
            BOOST_CHECK_SMALL( std::fabs( statesAlongTrajectory.rbegin( )->second( i ) - dsmLocation( i ) ), 5.0E-2 );
        }
    }

    // Test state output from DSM to arrival
    {
        // Get data on 10 equispace points on trajectory
        std::map< double, Eigen::Vector6d > statesAlongTrajectory;
        std::vector< double > outputTimes = utilities::linspace( transferLeg.getDsmTime( ) + timeTolerance, timeOfFlight, 10 );
        transferLeg.getStatesAlongTrajectory( statesAlongTrajectory, outputTimes );

        // Check initial and final time on output list
        BOOST_CHECK_CLOSE_FRACTION( statesAlongTrajectory.begin( )->first, dsmTime + timeTolerance, 1.0E-14 );
        BOOST_CHECK_CLOSE_FRACTION( statesAlongTrajectory.rbegin( )->first, timeOfFlight, 1.0E-14 );

        // Check if Keplerian state (slow elements) is the same for each output point
        Eigen::Vector6d previousKeplerianState = Eigen::Vector6d::Constant( TUDAT_NAN );
        for( auto it : statesAlongTrajectory )
        {
            Eigen::Vector6d currentCartesianState = it.second;
            Eigen::Vector6d currentKeplerianState = tudat::orbital_element_conversions::convertCartesianToKeplerianElements(
                        currentCartesianState, sunGravitationalParameter );
            if( previousKeplerianState == previousKeplerianState )
            {
                TUDAT_CHECK_MATRIX_CLOSE_FRACTION(
                            ( previousKeplerianState.segment( 0, 5 ) ),
                            ( currentKeplerianState.segment( 0, 5 ) ),
                            1.0E-14 );

            }
            previousKeplerianState = currentKeplerianState;
        }

        // Check if output meets boundary conditions
        for( int i = 0; i < 3; i++ )
        {
            BOOST_CHECK_SMALL( std::fabs( statesAlongTrajectory.begin( )->second( i ) - dsmLocation( i ) ), 5.0E-2 );
            BOOST_CHECK_SMALL( std::fabs( statesAlongTrajectory.rbegin( )->second( i ) - planet2State( i ) ), 1.0E-2 );
        }

    }
}

BOOST_AUTO_TEST_SUITE_END( )

} // namespace unit_tests
} // namespace tudat


// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// 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, see <http://www.gnu.org/licenses/>.

#ifndef UTIL_HH
#define UTIL_HH

#include "config.h"

#include <cmath>
#include <iostream>

#include "cache.hh"

#include <boost/range/counting_range.hpp>

namespace graph_tool
{
using namespace boost;

template <class T1, class T2>
inline double lbinom(T1 N, T2 k)
{
    if (N == 0 || k == 0 || k >= N)
        return 0;
    assert(N > 0);
    assert(k > 0);
    return ((std::lgamma(N + 1) - std::lgamma(k + 1)) - std::lgamma(N - k + 1));
}

template <bool Init=true, class T1, class T2>
inline double lbinom_fast(T1 N, T2 k)
{
    if (N == 0 || k == 0 || k > N)
        return 0;
    return ((lgamma_fast<Init>(N + 1) - lgamma_fast<Init>(k + 1)) - lgamma_fast<Init>(N - k + 1));
}

template <class T1, class T2>
inline double lbinom_careful(T1 N, T2 k)
{
    if (N == 0 || k == 0 || k >= N)
        return 0;
    double lgN = std::lgamma(N + 1);
    double lgk = std::lgamma(k + 1);
    if (lgN - lgk > 1e8)
    {
        // We have N >> k. Use Stirling's approximation: ln N! ~ N ln N - N
        // and reorder
        return - N * log1p(-k / N) - k * log1p(-k / N) - k - lgk + k * log(N);
    }
    else
    {
        return lgN - std::lgamma(N - k + 1) - lgk;
    }
}

template <class T>
inline auto lbeta(T x, T y)
{
    return (std::lgamma(x) + std::lgamma(y)) - std::lgamma(x + y);
}

template <class Vec, class PosMap, class Val>
void remove_element(Vec& vec, PosMap& pos, Val val)
{
    auto& back = vec.back();
    auto& j = pos[back];
    auto i = pos[val];
    vec[i] = back;
    j = i;
    vec.pop_back();
}

template <class Vec, class PosMap, class Val>
void add_element(Vec& vec, PosMap& pos, Val val)
{
    pos[val] = vec.size();
    vec.push_back(val);
}

template <class Vec, class PosMap, class Val>
bool has_element(Vec& vec, PosMap& pos, Val val)
{
    size_t i = pos[val];
    return (i < vec.size() && vec[i] == val);
}

} // namespace graph_tool

#endif // UTIL_HH