#include "DelaunayTriangulation.h"
#include "DelaunayMemoryManager.h"
#include "math/mathutils.h"
#include <float.h>
#include <cassert>
#include <set>
#include <list>
#include <boost/bind.hpp>

namespace math
{
   //--------------------------------------------------------------------------

   DelaunayTriangulation::DelaunayTriangulation(double xmin, double ymin, double xmax, double ymax, EDelaunayLocationAlgorithms eAlgorithm)
      : _xmin(xmin), _ymin(ymin), _xmax(xmax), _ymax(ymax), _pStartTriangle(0), _eLocationAlgorithm(eAlgorithm)
   {
      assert(_xmin<_xmax);
      assert(_ymin<_ymax);

      _Init();
   }

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

   DelaunayTriangulation::~DelaunayTriangulation()
   {
      Clear();
      if (_pStartTriangle)
      {
         //std::cout << "<b>Removing Triangle</b>\n";
         _qLocationStructure->DeleteMemory(_pStartTriangle);

         // This assertion is only valid when one instance of DelaunayTriangulation exists!
         //assert(DelaunayMemoryManager::GetNumTriangles() == 0);
         //assert(DelaunayMemoryManager::GetNumVertices() == 0);

      }
   }

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

   void DelaunayTriangulation::_Init()
   {
      // calculate Triangle that surrounds area defined in _xmin, _ymin, _xmax, _ymax
      double Mx = _xmin + (_xmax-_xmin)/2.0;
      double My = _ymin + (_ymax-_ymin)/2.0;
      double r = sqrt((_xmax-Mx)*(_xmax-Mx)+(_ymax-My)*(_ymax-My));
      double Cy = _ymax + (_xmax-Mx)*(_xmax-Mx) / (_ymax-My);
      double Cx = Mx;
      double By = My -r;
      double Bx = Cx-(Cy-My+r)*(_xmax-Cx)/(_ymax-Cy);
      double Ay = My - r;
      double Ax = Cx-(Cy-My+r)*(_xmin-Cx)/(_ymax-Cy);

      DelaunayVertex* A = DelaunayMemoryManager::AllocVertex(Ax,Ay,0,-1);
      DelaunayVertex* B = DelaunayMemoryManager::AllocVertex(Bx,By,0,-1);
      DelaunayVertex* C = DelaunayMemoryManager::AllocVertex(Cx,Cy,0,-1);

      _pStartTriangle = DelaunayMemoryManager::AllocTriangle();

      _pStartTriangle->SetVertex(0, A);
      _pStartTriangle->SetVertex(1, B);
      _pStartTriangle->SetVertex(2, C);

      double xmax, ymax, xmin, ymin;
      xmax = ymax = -1e20;
      xmin = ymin = 1e20;

      xmin = math::Min<double>(Ax, xmin);
      xmin = math::Min<double>(Bx, xmin);
      xmin = math::Min<double>(Cx, xmin);
      ymin = math::Min<double>(Ay, ymin);
      ymin = math::Min<double>(By, ymin);
      ymin = math::Min<double>(Cy, ymin);
      xmax = math::Max<double>(Ax, xmax);
      xmax = math::Max<double>(Bx, xmax);
      xmax = math::Max<double>(Cx, xmax);
      ymax = math::Max<double>(Ay, ymax);
      ymax = math::Max<double>(By, ymax);
      ymax = math::Max<double>(Cy, ymax);

      _qLocationStructure = IDelaunayLocationStructure::CreateLocationStructure(xmin, ymin, xmax, ymax, _eLocationAlgorithm);
      if (_qLocationStructure)
      {
         _qLocationStructure->AddTriangle(_pStartTriangle);
      }

      if (!_pStartTriangle->IsCCW())
      {
         std::cout << "Error: Start-triangle is not ccw!\n";
         assert(false);
      }

      _bError = false;
   }

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

   void DelaunayTriangulation::_CollectTriangle(DelaunayTriangle* pTri)
   {
      _vecTriangles.push_back(pTri);
   }

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

   void DelaunayTriangulation::Clear()
   {
      if (_pStartTriangle)
      {
         _vecTriangles.clear();
         _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_CollectTriangle, this, _1));

         for (size_t i=0;i<_vecTriangles.size();i++)
         {
            _qLocationStructure->DeleteMemory(_vecTriangles[i]);
         }

         _pStartTriangle = 0;
         _Init();
      }
   }

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

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

   namespace internal
   {
      inline double _CalcTriArea(const ElevationPoint& A, const ElevationPoint& B, const ElevationPoint& C)
      {
         return math::ccw(A.x, A.y, B.x, B.y, C.x, C.y);
         //return 0.5*((B.x - A.x)*(C.y - A.y) - (B.y - A.y)*(C.x - A.x));
      }
   }

   double DelaunayTriangulation::GetElevationAt(double x, double y, ElevationQuery& query_result)
   {
      if (x<=_xmax && x>=_xmin &&
         y<=_ymax && y>=_ymin)
      {
         query_result = EQ_INTERIOR;
         ePointTriangleRelation relation;
         DelaunayTriangle* pTri = _qLocationStructure->GetTriangleAt(x,y,relation);

         if (pTri && !pTri->IsSuperSimplex())
         {
            ElevationPoint P;
            P.x = x; P.y = y; P.elevation = 0.0;

            ElevationPoint A = pTri->GetVertex(0)->GetElevationPointCopy();
            ElevationPoint B = pTri->GetVertex(1)->GetElevationPointCopy();
            ElevationPoint C = pTri->GetVertex(2)->GetElevationPointCopy();

            double elva = A.elevation;
            double elvb = B.elevation;
            double elvc = C.elevation;

            switch(relation)
            {
            case PointTriangle_Edge0:  // Point on Edge 0
            case PointTriangle_Edge1:  // Point on Edge 1
            case PointTriangle_Edge2:  // Point on Edge 2
            case PointTriangle_Inside: // Point inside triangle
               {
                  A.elevation = 0.0;
                  B.elevation = 0.0;
                  C.elevation = 0.0;

                  double F_abc = internal::_CalcTriArea(A,B,C);
                  double F_pbc = internal::_CalcTriArea(P,B,C);
                  double F_apc = internal::_CalcTriArea(A,P,C);

                  double r = F_pbc/F_abc;
                  double s = F_apc/F_abc;
                  double t = 1.0-r-s;

                  A.elevation = elva;
                  B.elevation = elvb;
                  C.elevation = elvc;

                  double Pelv = r*A.elevation + s*B.elevation + t*C.elevation;
                  return Pelv;
               }
               break;
            case PointTriangle_Vertex0:  // Point lies on vertex 0
               return pTri->GetVertex(0)->elevation();
            case PointTriangle_Vertex1:  // Point lies on vertex 1
               return pTri->GetVertex(1)->elevation();
            case PointTriangle_Vertex2:  // Point lies on vertex 2
               return pTri->GetVertex(2)->elevation();
            default: // outside triangle / invalid triangle
               query_result = EQ_EXTERIOR;
               return 0.0;
            }
         }
         else
         {
            query_result = EQ_EXTERIOR;
            return 0.0;
         }
      }
      else
      {
         query_result = EQ_UNDEFINED; // outside valid triangulation area
         return 0.0;
      }

   }

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

   void DelaunayTriangulation::InsertPoint(const ElevationPoint& pt)
   {
      if (pt.x<=_xmax && pt.x>=_xmin &&
         pt.y<=_ymax && pt.y>=_ymin)
      {
         DelaunayVertex* pNewVertex = DelaunayMemoryManager::AllocVertex(pt);
         _pStartTriangle = _qLocationStructure->InsertVertex(pNewVertex, _pStartTriangle);
         _bError = false; // inserting a point invalidates errors!
      }
   }

   void DelaunayTriangulation::_InsertPointSetId(const ElevationPoint& pt, int id)
   {
      if (pt.x<=_xmax && pt.x>=_xmin &&
         pt.y<=_ymax && pt.y>=_ymin)
      {
         DelaunayVertex* pNewVertex = DelaunayMemoryManager::AllocVertex(pt);
         pNewVertex->SetId(id);
         _pStartTriangle = _qLocationStructure->InsertVertex(pNewVertex, _pStartTriangle);
         _bError = false; // inserting a point invalidates errors!
      }
   }

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

   void DelaunayTriangulation::_ResetVertexId(DelaunayTriangle* pTri)
   {
      assert(pTri);
      assert(pTri->GetVertex(0));
      assert(pTri->GetVertex(1));
      assert(pTri->GetVertex(2));

      pTri->GetVertex(0)->SetId(-1);
      pTri->GetVertex(1)->SetId(-1);
      pTri->GetVertex(2)->SetId(-1);
   }

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

   void DelaunayTriangulation::_CollectElevationPoints(DelaunayTriangle* pTri)
   {
      assert(pTri);
      assert(pTri->GetVertex(0));
      assert(pTri->GetVertex(1));
      assert(pTri->GetVertex(2));

      // ignore supersimplex triangles.
      if (!pTri->IsSuperSimplex())
      {
         for (int v=0;v<3;v++)
         {
            if (pTri->GetVertex(v)->GetId() == -1)
            {
               pTri->GetVertex(v)->SetId(_idcnt);
               _vPts.push_back(pTri->GetVertex(v)->GetElevationPoint());
               _idcnt++;
            }

            _vIndex.push_back(pTri->GetVertex(v)->GetId());
         }
      }
   }

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

   // before calling this make sure to clear _vCutEdge, _vIndex and _vPts
   void DelaunayTriangulation::_CollectTriangulationStructure(DelaunayTriangle* pTri)
   {
      assert(pTri);
      assert(pTri->GetVertex(0));
      assert(pTri->GetVertex(1));
      assert(pTri->GetVertex(2));

      // ignore supersimplex triangles.
      if (!pTri->IsSuperSimplex())
      {
         for (int v=0;v<3;v++)
         {
            if (pTri->GetVertex(v)->GetId() == -1)
            {
               pTri->GetVertex(v)->SetId(_idcnt);
               _vPts.push_back(pTri->GetVertex(v)->GetElevationPoint());
               _idcnt++;
            }

            _vIndex.push_back(pTri->GetVertex(v)->GetId());
         }

         if (pTri->GetTriangle(0) == 0)
         {
            _vCutEdge.push_back(std::pair<int, int>(pTri->GetVertex(0)->GetId(), pTri->GetVertex(1)->GetId()));
         }

         if (pTri->GetTriangle(1) == 0)
         {
            _vCutEdge.push_back(std::pair<int, int>(pTri->GetVertex(1)->GetId(), pTri->GetVertex(2)->GetId()));
         }

         if (pTri->GetTriangle(2) == 0)
         {
            _vCutEdge.push_back(std::pair<int, int>(pTri->GetVertex(2)->GetId(), pTri->GetVertex(0)->GetId()));
         }

      }

   }

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

   // Retrieve Triangulation as Point / Index List
   void DelaunayTriangulation::GetPointVec(std::vector<ElevationPoint>& vPoints)
   {

      // Reset All Vertices to 0
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_ResetVertexId, this, _1));

      _idcnt = 0; _vPts.clear(); _vIndex.clear(); _vCutEdge.clear();
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_CollectTriangulationStructure, this, _1));

      vPoints = _vPts;
      _vPts.clear();

   }

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

   void DelaunayTriangulation::GetTriangleIndices(std::vector<int>& vIndices)
   {
      vIndices = _vIndex;
   }

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

   void DelaunayTriangulation::GetCutEdges(std::vector<std::pair<int, int> >& vCutEdges)
   {
      vCutEdges = _vCutEdge;
   }
   //--------------------------------------------------------------------------

   std::vector<DelaunayTriangle*>& DelaunayTriangulation::GetAllTriangles()
   {
      _vecTriangles.clear();
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_CollectTriangle, this, _1));

      return _vecTriangles;
   }

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

   void DelaunayTriangulation::DeleteMemory(DelaunayTriangle* pTri)
   {
      _qLocationStructure->DeleteMemory(pTri);
   }

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

   void DelaunayTriangulation::_CutEdges(std::vector< std::pair<int,int> >& vCut)
   {
      // Traverse triangulation structure and delete corresponding triangles

      std::vector<DelaunayTriangle*>& vTri = DelaunayTriangulation::GetAllTriangles();

      std::set<DelaunayTriangle*> setDelete;

      for (size_t i=0;i<vTri.size();i++)
      {
         DelaunayTriangle* pTri = vTri[i];
         DelaunayTriangle* pTri0 = pTri->GetTriangle(0);
         DelaunayTriangle* pTri1 = pTri->GetTriangle(1);
         DelaunayTriangle* pTri2 = pTri->GetTriangle(2);

         int A,B,C;
         A = pTri->GetVertex(0)->GetId();
         B = pTri->GetVertex(1)->GetId();
         C = pTri->GetVertex(2)->GetId();

         // does this triangle contain a cut edge ?

         for (size_t c=0;c<vCut.size();c++)
         {
            int start = vCut[c].first;
            int end = vCut[c].second;

            if (A == start && B == end)
            {
               setDelete.insert(pTri0);
            }
            else if (B == start && C == end)
            {
               setDelete.insert(pTri1);
            }
            else if (C == start && A == end)
            {
               setDelete.insert(pTri2);
            }
         }
      }

      std::set<DelaunayTriangle*>::iterator it = setDelete.begin();

      while (it != setDelete.end())
      {
         DeleteMemory(*it);
         it++;
      }

   }

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

   void DelaunayTriangulation::SetEpsilon(double epsilon)
   {
      assert(_qLocationStructure);

      if (_qLocationStructure)
      {
         _qLocationStructure->SetEpsilon(epsilon);
      }
   }

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

   void DelaunayTriangulation::_ResetVertexErrors(DelaunayTriangle* pTri)
   {
      pTri->GetVertex(0)->GetElevationPoint().error = -1.0;
      pTri->GetVertex(1)->GetElevationPoint().error = -1.0;
      pTri->GetVertex(2)->GetElevationPoint().error = -1.0;
   }

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

   void DelaunayTriangulation::_GetCCWVertices(DelaunayTriangle* pTri, int vertex_index, std::vector<DelaunayVertex*>& outputVertices)
   {
      outputVertices.clear();
      std::list<DelaunayVertex*> lst;

      const DelaunayTriangle* pStartTriangle = pTri;

      int vtx = vertex_index;
      int triangle_index = (vtx+2)%3;
      lst.push_front(pTri->GetVertex((vtx+2)%3));

      DelaunayTriangle* A = pTri;
      DelaunayTriangle* C = pTri->GetTriangle(triangle_index);

      while(C && C != pStartTriangle && !C->IsSuperSimplex() && !A->IsSuperSimplex())
      {
         if (A->GetVertex(vtx) == C->GetVertex(0))
         {
            vtx = 0;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(1))
         {
            vtx = 1;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(2))
         {
            vtx = 2;
         }
         else
         {
            assert(false);
            break;
         }

         lst.push_front(C->GetVertex((vtx+2)%3));

         triangle_index = (vtx+2)%3;
         A = C;
         C = C->GetTriangle(triangle_index);
      }

      lst.push_back(pTri->GetVertex((vertex_index+1)%3));

      // ignore non ccw direction!
      // copy result to output
      std::list<DelaunayVertex*>::iterator it = lst.begin();
      while (it!=lst.end())
      {
         outputVertices.push_back(*it);
         it++;
      }

   }

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

   void DelaunayTriangulation::_CreateSurroundingPolygon(DelaunayTriangle* pTri, int vertex_index, std::vector<ElevationPoint>& outputPolygon)
   {
      outputPolygon.clear();
      std::list<ElevationPoint> lst;

      const DelaunayTriangle* pStartTriangle = pTri;

      int vtx = vertex_index;
      int triangle_index = (vtx+2)%3;
      lst.push_front(pTri->GetVertex((vtx+2)%3)->GetElevationPoint());

      DelaunayTriangle* A = pTri;
      DelaunayTriangle* B;
      DelaunayTriangle* C = pTri->GetTriangle(triangle_index);

      while(C && C != pStartTriangle && !C->IsSuperSimplex() && !A->IsSuperSimplex())
      {
         if (A->GetVertex(vtx) == C->GetVertex(0))
         {
            vtx = 0;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(1))
         {
            vtx = 1;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(2))
         {
            vtx = 2;
         }
         else
         {
            assert(false);
            break;
         }

         lst.push_front(C->GetVertex((vtx+2)%3)->GetElevationPoint());

         triangle_index = (vtx+2)%3;
         A = C;
         C = C->GetTriangle(triangle_index);
      }

      // other direction and only if pTri->GetTriangle(vtx) != pStartTriangle
      lst.push_back(pTri->GetVertex((vertex_index+1)%3)->GetElevationPoint());

      if (C != pStartTriangle)
      {
         vtx = vertex_index;
         triangle_index = vtx;

         A = pTri;
         B = pTri->GetTriangle(triangle_index);

         while(B && B != pStartTriangle && !B->IsSuperSimplex() && !A->IsSuperSimplex())
         {
            if (A->GetVertex(vtx) == B->GetVertex(0))
            {
               vtx = 0;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(1))
            {
               vtx = 1;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(2))
            {
               vtx = 2;
            }

            lst.push_back(B->GetVertex((vtx+1)%3)->GetElevationPoint());
            triangle_index = vtx;

            A = B;
            B = B->GetTriangle(triangle_index);
         }

      }

      // Convert list to vector (for returning values)
      if (lst.size()>2)
      {
         std::list<ElevationPoint>::iterator it = lst.begin();
         while (it!=lst.end())
         {
            outputPolygon.push_back(*it);
            it++;
         }
      }

   }

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

   void DelaunayTriangulation::_CalcVertexErrorsVtx(DelaunayTriangle* pTri, int vtx)
   {
      assert(vtx>=0 && vtx<=3);
      assert(pTri);

      ElevationPoint pt = pTri->GetVertex(vtx)->GetElevationPoint();

      // corner points always have max error
      if (pt.weight < -2
         || pTri->GetVertex((vtx+1)%3)->GetElevationPoint().weight < -2
         || pTri->GetVertex((vtx+2)%3)->GetElevationPoint().weight < -2)
      {
         pTri->GetVertex(vtx)->GetElevationPoint().error = DBL_MAX;
      }
      else
      {
         boost::shared_ptr<DelaunayTriangulation> qTriangulation;

         if (pTri->GetVertex(vtx)->GetElevationPoint().error <= -1.0)
         {
            std::vector<ElevationPoint> outputPolygon;
            _CreateSurroundingPolygon(pTri, vtx, outputPolygon);

            qTriangulation = boost::shared_ptr<DelaunayTriangulation>(new DelaunayTriangulation(_xmin, _ymin, _xmax, _ymax, _eLocationAlgorithm));
            qTriangulation->SetEpsilon(DBL_EPSILON); // can be removed later

            for (size_t i=0;i<outputPolygon.size();i++)
            {
               qTriangulation->InsertPoint(outputPolygon[i]);
            }

            math::ElevationQuery query;
            double elv = qTriangulation->GetElevationAt(pt.x, pt.y, query);

            if (query == EQ_INTERIOR)
            {
               double dError = fabs(elv - pt.elevation);
               pTri->GetVertex(vtx)->GetElevationPoint().error = dError;

               if (dError < _minError)
               {
                  _oVertexMinError.pTri = pTri;
                  _oVertexMinError.idx0 = vtx;
                  _minError = dError;
               }

            }
            else
            {
               pTri->GetVertex(vtx)->GetElevationPoint().error = DBL_MAX;
            }

         }

      }
   }

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

   void DelaunayTriangulation::_CalcVertexErrors(DelaunayTriangle* pTri)
   {
      assert(pTri);

      if (!pTri->IsSuperSimplex()) // ignore supersimplex triangles!
      {
         _CalcVertexErrorsVtx(pTri,0);
         _CalcVertexErrorsVtx(pTri,1);
         _CalcVertexErrorsVtx(pTri,2);
      }
   }

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

   void DelaunayTriangulation::CalculateVertexErrors()
   {
      _oVertexMinError.pTri = 0;
      _oVertexMinError.idx0 = -1;

      _minError = DBL_MAX;
      // Reset Errors:
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_ResetVertexErrors, this, _1));
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_CalcVertexErrors, this, _1));

      _bError = true;
      //std::cout << "MinError: " << _minError << "\n";
   }

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

   void DelaunayTriangulation::UpdateVertexErrors(std::vector<DelaunayVertex*>& vVertex)
   {

      ePointTriangleRelation e;
      DelaunayTriangle* pTri;
      int idx;

      for (size_t i=0;i<vVertex.size();i++)
      {
         idx = -1;
         pTri = _qLocationStructure->GetTriangleAt(vVertex[i]->x(),vVertex[i]->y(),e);
         if (pTri && !pTri->IsSuperSimplex())
         {
            if (pTri->GetVertex(0) == vVertex[i])
               idx = 0;
            else if (pTri->GetVertex(1) == vVertex[i])
               idx = 1;
            else if (pTri->GetVertex(2) == vVertex[i])
               idx = 2;

            if (idx != -1)
            {
               _CalcVertexErrorsVtx(pTri,idx);
            }
            else
            {
               assert(false); // vertex not found!! How can this be ??!??
            }
         }

      }

      _UpdateMinError();
   }

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

   void DelaunayTriangulation::_UpdateVertexErrors(DelaunayTriangle* pTri)
   {
      double dError;

      if (!pTri->IsSuperSimplex())
      {
         for (int i=0;i<3;i++)
         {
            dError = pTri->GetVertex(i)->GetElevationPoint().error;
            if (dError < _minError)
            {
               _oVertexMinError.pTri = pTri;
               _oVertexMinError.idx0 = i;
               _minError = dError;
            }
         }
      }
   }

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

   void DelaunayTriangulation::_UpdateMinError()
   {
      _minError = DBL_MAX;
      _qLocationStructure->Traverse(boost::bind(&DelaunayTriangulation::_UpdateVertexErrors, this, _1));
   }

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

   void DelaunayTriangulation::RemoveLeastErrorVertex()
   {
      if (_minError == DBL_MAX)
         return;

      if (_oVertexMinError.pTri)
      {
         // GetNeighbourVertices and save in list
         if (_oVertexMinError.pTri->GetVertex(_oVertexMinError.idx0)->weight() == -3)
         {
            std::cout << "**** FATAL ERROR ****: Removing corner vertex!!\n";
         }

         _RemoveVertex(_oVertexMinError.pTri, _oVertexMinError.idx0);
         _oVertexMinError.pTri = 0;
         _oVertexMinError.idx0 = -1;
      }

   }

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

   int DelaunayTriangulation::Simplify(double epsilon, int maxiterations)
   {
      std::vector<DelaunayVertex*> vVertex;
      int rmvsteps = 0;
      if (!_bError)
         CalculateVertexErrors();

       while (_minError <= epsilon && _minError != DBL_MAX && rmvsteps < maxiterations)
       {
          GetCCWVertices(_oVertexMinError.pTri, _oVertexMinError.idx0, vVertex);
          RemoveLeastErrorVertex(); rmvsteps++;
          UpdateVertexErrors(vVertex);
       }

       /*if (rmvsteps >= maxiterations)
       {
          std::cout << "Simplify-Condition: Max Iterations reached!\n";
       }

       if (_minError > epsilon)
       {
          std::cout << "Simplify-Condition: Epsilon condition reached!\n";
       }*/

       return rmvsteps;

   }

   //--------------------------------------------------------------------------
   void DelaunayTriangulation::Reduce(int nPoints)
   {
      std::vector<DelaunayVertex*> vVertex;
      int rmvsteps = 0;
      if (!_bError)
         CalculateVertexErrors();

      while (_minError != DBL_MAX && rmvsteps < nPoints)
      {
         GetCCWVertices(_oVertexMinError.pTri, _oVertexMinError.idx0, vVertex);
         RemoveLeastErrorVertex(); rmvsteps++;
         UpdateVertexErrors(vVertex);
      }
   }

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

   void DelaunayTriangulation::_GetVertexAt(double x, double y, DelaunayTriangle*& pTri, int& idx)
   {
      ePointTriangleRelation e;
      pTri = _qLocationStructure->GetTriangleAt(x,y,e);
      if (pTri)
      {
         const ElevationPoint& P0 = pTri->GetVertex(0)->GetElevationPoint();
         const ElevationPoint& P1 = pTri->GetVertex(1)->GetElevationPoint();
         const ElevationPoint& P2 = pTri->GetVertex(2)->GetElevationPoint();

         double dist0, dist1, dist2;

         dist0 = sqrt((P0.x-x)*(P0.x-x)+(P0.y-y)*(P0.y-y));
         dist1 = sqrt((P1.x-x)*(P1.x-x)+(P1.y-y)*(P1.y-y));
         dist2 = sqrt((P2.x-x)*(P2.x-x)+(P2.y-y)*(P2.y-y));

         if (dist0 <= dist1 && dist0 <= dist2)
         {
            idx = 0;
         }
         else if (dist1 <= dist0 && dist1 <= dist2)
         {
            idx = 1;
         }
         else //if (dist2 <= dist0 && dist2 <=dist1)
         {
            idx = 2;
         }
      }
      else
      {
         // no triangle found!!
         idx = -1;
      }
   }

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

   void DelaunayTriangulation::GetCCWTriangles(DelaunayTriangle* pTri, int vertex_index, std::vector<STriangleVertex>& outputVertices)
   {
      outputVertices.clear();
      std::list<STriangleVertex> lst;

      const DelaunayTriangle* pStartTriangle = pTri;

      STriangleVertex oElement;
      int vtx = vertex_index;
      int triangle_index = (vtx+2)%3;

      oElement.pTri = pTri;
      oElement.idx0 = (vertex_index+1)%3;
      lst.push_front(oElement);

      DelaunayTriangle* A = pTri;
      DelaunayTriangle* B;
      DelaunayTriangle* C = pTri->GetTriangle(triangle_index);

      while(C && C != pStartTriangle)
      {
         if (A->GetVertex(vtx) == C->GetVertex(0))
         {
            vtx = 0;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(1))
         {
            vtx = 1;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(2))
         {
            vtx = 2;
         }
         else
         {
            assert(false);
            break;
         }

         oElement.pTri = C;
         oElement.idx0 = (vtx+1)%3;
         lst.push_front(oElement);

         triangle_index = (vtx+2)%3;
         A = C;
         C = C->GetTriangle(triangle_index);
      }

      // move around other side:

      if (C != pStartTriangle)
      {
         vtx = vertex_index;
         triangle_index = vtx;

         A = pTri;
         B = pTri->GetTriangle(triangle_index);

         while(B && B != pStartTriangle)
         {
            if (A->GetVertex(vtx) == B->GetVertex(0))
            {
               vtx = 0;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(1))
            {
               vtx = 1;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(2))
            {
               vtx = 2;
            }

            oElement.pTri = B;
            oElement.idx0 = (vtx+1)%3;
            lst.push_back(oElement);
            triangle_index = vtx;

            A = B;
            B = B->GetTriangle(triangle_index);
         }

      }

      std::reverse(lst.begin(), lst.end());

      // copy result to output
      std::list<STriangleVertex>::iterator it = lst.begin();
      while (it!=lst.end())
      {
         outputVertices.push_back(*it);
         it++;
      }
   }

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

   void DelaunayTriangulation::GetCCWVertices(DelaunayTriangle* pTri, int vertex_index, std::vector<DelaunayVertex*>& outputVertices)
   {
      outputVertices.clear();
      std::list<DelaunayVertex*> lst;

      const DelaunayTriangle* pStartTriangle = pTri;

      int vtx = vertex_index;
      int triangle_index = (vtx+2)%3;

      lst.push_front(pTri->GetVertex((vertex_index+1)%3));

      DelaunayTriangle* A = pTri;
      DelaunayTriangle* B;
      DelaunayTriangle* C = pTri->GetTriangle(triangle_index);

      while(C && C != pStartTriangle)
      {
         if (A->GetVertex(vtx) == C->GetVertex(0))
         {
            vtx = 0;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(1))
         {
            vtx = 1;
         }
         else if (A->GetVertex(vtx) == C->GetVertex(2))
         {
            vtx = 2;
         }
         else
         {
            assert(false);
            break;
         }

         lst.push_front(C->GetVertex((vtx+1)%3));

         triangle_index = (vtx+2)%3;
         A = C;
         C = C->GetTriangle(triangle_index);
      }

      // move around other side:

      if (C != pStartTriangle)
      {
         vtx = vertex_index;
         triangle_index = vtx;

         A = pTri;
         B = pTri->GetTriangle(triangle_index);

         while(B && B != pStartTriangle)
         {
            if (A->GetVertex(vtx) == B->GetVertex(0))
            {
               vtx = 0;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(1))
            {
               vtx = 1;
            }
            else if (A->GetVertex(vtx) == B->GetVertex(2))
            {
               vtx = 2;
            }

            lst.push_back(B->GetVertex((vtx+1)%3));
            triangle_index = vtx;

            A = B;
            B = B->GetTriangle(triangle_index);
         }

      }

      std::reverse(lst.begin(), lst.end());

      // copy result to output
      std::list<DelaunayVertex*>::iterator it = lst.begin();
      while (it!=lst.end())
      {
         outputVertices.push_back(*it);
         it++;
      }
   }

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

   void DelaunayTriangulation::_RemoveVertex(DelaunayTriangle* pTri, int idx)
   {
      if (pTri && idx>=0 && idx<=3)
      {
         bool DebugOutput = false;
         DelaunayVertex* pVertex = pTri->GetVertex(idx); // this is the point P to be removed.

         std::vector<STriangleVertex> outputTriangles;
         GetCCWTriangles(pTri, idx, outputTriangles);

         if (outputTriangles.size() < 3 )
         {
            // this can't happen in a "normal" triangulation
            // unless you try to remove supersimplex corner.
            // however this can happen in a triangulation with holes. This
            // case is currently not supported!
            std::cout<< "*Error* Can't remove specified vertex! (Triangulation would be broken.)\n";
            return;
         }
         else
         {
            //std::cout << "-----------------------------\n";
            //std::cout << "Removing Point\n";
            do
            {
               size_t num_changes = 0;

               //std::cout << "outputTriangles: " << outputTriangles.size() << "\n";

               if (outputTriangles.size() > 3)
               {

                  std::vector<STriangleVertex>::iterator it = outputTriangles.begin();

                  while (it != outputTriangles.end() && outputTriangles.size()>3)
                  {

                     DelaunayTriangle* pCurrentTriangle = (*it).pTri;
                     int idx0 = (*it).idx0;
                     int idx1 = (idx0+1)%3;

                     DelaunayVertex* s0 = pCurrentTriangle->GetVertex(idx0);
                     DelaunayVertex* s1 = pCurrentTriangle->GetVertex(idx1);
                     DelaunayVertex* s2 = pCurrentTriangle->GetOppositeVertex(idx1);

                     // for debugging reasons only:
                     ElevationPoint s0p = s0->GetElevationPoint();
                     ElevationPoint s1p = s1->GetElevationPoint();
                     ElevationPoint s2p = s2->GetElevationPoint();
                     ElevationPoint sP = pVertex->GetElevationPoint();

                     double ccwpredicate = math::ccw(s0,s1,s2);

                     if (DebugOutput)
                        std::cout << "ccwpredicate0 = " << ccwpredicate << "\n";

                     if (ccwpredicate<=0)
                     {
                        // this triangle-combination is ignored!
                     }
                     else
                     {
                        ccwpredicate = math::ccw(s0,s2,pVertex);
                        if (DebugOutput)
                           std::cout << "ccwpredicate1 = " << ccwpredicate << "\n";
                        if (ccwpredicate<0) // P encloses Triangle ?
                        {

                           // this triangle-combination is ignored!
                        }
                        else
                        {
                           // Check if other points are inside cirumcircle!
                           bool bCircleTest = true;

                           std::vector<STriangleVertex>::iterator it2 = outputTriangles.begin();
                           while (it2 != outputTriangles.end())
                           {
                              if ((*it2).pTri != (*it).pTri)
                              {
                                 int tidx0 = (*it2).idx0;
                                 int tidx1 = (tidx0+1)%3;

                                 DelaunayVertex* p0 = (*it2).pTri->GetVertex(tidx0);
                                 DelaunayVertex* p1 = (*it2).pTri->GetVertex(tidx1);

                                 assert(p0 != pVertex);
                                 assert(p1 != pVertex);

                                 if (p0 != s0 && p0 != s1 && p0 != s2 &&
                                     p1 != s0 && p1 != s1 && p1 != s2)
                                 {
                                    double circ = math::InCircleValue(s0,s1,s2,p0);
                                    if (DebugOutput)
                                       std::cout << "circ0 = " << circ << "\n";
                                    if (circ >= DBL_EPSILON)
                                    {
                                       bCircleTest = false;
                                    }
                                    circ = math::InCircleValue(s0,s1,s2,p1);
                                    if (DebugOutput)
                                       std::cout << "circ1 = " << circ << "\n";
                                    if (circ >= DBL_EPSILON)
                                    {
                                       bCircleTest = false;
                                    }
                                 }
                              }
                              it2++;
                           }

                           if (bCircleTest)
                           {
                              // this triangle and (next) are swapped and removed from list.
                              // new (swapped) triangle containing P will be added to list.

                              DelaunayTriangle* C;
                              DelaunayTriangle* D;

                              if (pCurrentTriangle->FlipEdge(idx1, &C, &D))
                              {
                                 it = outputTriangles.erase(it);
                                 if (it==outputTriangles.end())
                                 {
                                    it = outputTriangles.begin();
                                 }
                                 it = outputTriangles.erase(it);

                                 STriangleVertex newElement;

                                 if (C->GetVertex(0) == pVertex)
                                 {
                                    newElement.pTri = C; newElement.idx0 = 1;
                                 }
                                 else if (C->GetVertex(1) == pVertex)
                                 {
                                    newElement.pTri = C; newElement.idx0 = 2;
                                 }
                                 else if (C->GetVertex(2) == pVertex)
                                 {
                                    newElement.pTri = C; newElement.idx0 = 0;
                                 }
                                 else if (D->GetVertex(0) == pVertex)
                                 {
                                    newElement.pTri = D; newElement.idx0 = 1;
                                 }
                                 else if (D->GetVertex(1) == pVertex)
                                 {
                                    newElement.pTri = D; newElement.idx0 = 2;
                                 }
                                 else if (D->GetVertex(2) == pVertex)
                                 {
                                    newElement.pTri = D; newElement.idx0 = 0;
                                 }
                                 else
                                 {
                                    assert(false);
                                 }

                                 it = outputTriangles.insert(it, newElement);
                                 num_changes++;
                              }
                              else
                              {
                                 // can't flip
                              }
                           }

                        }
                     }

                     it++;
                  }
               }

               if (outputTriangles.size()>3 && num_changes == 0)
               {
                  std::cout << "<b>*WARNING* Detected infinite loop!</b>\n";
                  pVertex->GetElevationPoint().error = -0.5;

                  if (DebugOutput)
                     break;
                  else
                     DebugOutput = true;
               }

               // 3 Remaining pairs: Remove 3 Triangles!!
               if (outputTriangles.size() == 3)
               {
                  STriangleVertex* st0 =  &outputTriangles[0];
                  STriangleVertex* st1 =  &outputTriangles[1];
                  STriangleVertex* st2 =  &outputTriangles[2];

                  // Create New Triangle
                  DelaunayTriangle* pNewTriangle = DelaunayMemoryManager::AllocTriangle();
                  pNewTriangle->SetVertex(0, st0->pTri->GetVertex(st0->idx0));
                  pNewTriangle->SetVertex(1, st1->pTri->GetVertex(st1->idx0));
                  pNewTriangle->SetVertex(2, st2->pTri->GetVertex(st2->idx0));

                  DelaunayTriangle* pNeighbour0 = st0->pTri->GetTriangle(st0->idx0);
                  int nr0 = st0->pTri->NeighbourReference(st0->idx0);

                  DelaunayTriangle* pNeighbour1 = st1->pTri->GetTriangle(st1->idx0);
                  int nr1 = st1->pTri->NeighbourReference(st1->idx0);

                  DelaunayTriangle* pNeighbour2 = st2->pTri->GetTriangle(st2->idx0);
                  int nr2 = st2->pTri->NeighbourReference(st2->idx0);

                  pNewTriangle->SetTriangle(0, pNeighbour0);
                  pNewTriangle->SetTriangle(1, pNeighbour1);
                  pNewTriangle->SetTriangle(2, pNeighbour2);

                  _qLocationStructure->DeleteMemory(st0->pTri);
                  _qLocationStructure->DeleteMemory(st1->pTri);
                  _qLocationStructure->DeleteMemory(st2->pTri);

                  if (pNeighbour0)
                     pNeighbour0->SetTriangle(nr0, pNewTriangle);

                  if (pNeighbour1)
                     pNeighbour1->SetTriangle(nr1, pNewTriangle);

                  if (pNeighbour2)
                     pNeighbour2->SetTriangle(nr2, pNewTriangle);

                  _qLocationStructure->AddTriangle(pNewTriangle);

                  //assert(pNewTriangle->IsCCW()); // the mosted hated assertion
               }

            }
            while (outputTriangles.size()>3);
         }

      }
   }

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

   void DelaunayTriangulation::RemoveVertex(double x, double y)
   {
      int idx;
      DelaunayTriangle* pTri;
      _GetVertexAt(x,y,pTri,idx);
      _RemoveVertex(pTri,idx);
   }

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

} // namespace


#include <iostream>
#include <Eigen/Dense>
#include <random>
#include <vector>
#include <iomanip>
#include <fstream>
#include <bitset>
#include <string>
#include <map>

#ifndef QST_TOMOGRAPHY_HPP
#define QST_TOMOGRAPHY_HPP

namespace qst{

//Quantum State Tomography class
template<class NNState,class Observer,class Optimizer> class Tomography {

    NNState & NNstate_;                         // Neural network representation of the state
    Optimizer & opt_;                           // Optimizer
    Observer &obs_;                             // Observer

    int N_;                                     // Number of physical degrees of freedom
    int npar_;                                  // Number of variational parameters
    int nparLambda_;                            // Number of amplitude variational parameters
    int nparMu_;                                // Number of phase variational parameters
    int bs_;                                    // Batch size
    int cd_;                                    // Number of Gibbs stepos in contrastive divergence
    int epochs_;                                // Number of training iterations
    double lr_;                                 // Learning rate
    double l2_;                                 // L2 regularization constant

    Eigen::VectorXd grad_;                      // Gradients
    Eigen::VectorXcd rotated_grad_;             // Rotated gradients
    std::mt19937 rgen_;                         // Random number generator
    std::map<std::string,Eigen::MatrixXcd> U_;  // Structure containin the single unitary rotations
public:

    Tomography(Optimizer & opt,NNState & NNstate,Observer &obs,Parameters &par):opt_(opt),NNstate_(NNstate),obs_(obs),N_(NNstate.N()),bs_(par.bs_),cd_(par.cd_) {

        npar_=NNstate_.Npar();
        nparLambda_ = NNstate_.NparLambda();
        nparMu_ = NNstate_.NparMu();
        opt_.SetNpar(npar_);
        bs_ = par.bs_;
        cd_ = par.cd_;
        lr_ = par.lr_;
        l2_ = par.l2_;
        epochs_ = par.ep_;
        grad_.resize(npar_);
        rotated_grad_.resize(npar_);
    }

    //Compute gradient of KL divergence
    void ComputeGradient(const Eigen::MatrixXd & batchSamples,const std::vector<std::vector<std::string> >& batchBases){
        grad_.setZero();

        int bID = 0;
        //Positive Phase
        for(int k=0;k<bs_;k++){
            bID = 0;
            for(int j=0;j<N_;j++){ // Check if the basis is the reference one
                if (batchBases[k][j]!="Z"){
                    bID = 1;
                    break;
                }
            }
            if (bID==0){ // Positive phase - Lambda gradient in the reference basis
                grad_.head(nparLambda_) += NNstate_.LambdaGrad(batchSamples.row(k))/double(bs_);
            }
            else { // Positive phase - Lambda and Mu gradients for non-trivial bases
                NNstate_.rotatedGrad(batchBases[k],batchSamples.row(k),U_,rotated_grad_);
                //getRotatedGradient(batchBases[k],batchSamples.row(k),rotated_grad_);
                grad_.head(nparLambda_) += rotated_grad_.head(nparLambda_).real()/double(bs_);
                grad_.tail(nparMu_) -= rotated_grad_.tail(nparMu_).imag()/double(bs_);
            }
        }

        //Negative Phase
        NNstate_.Sample(cd_);
        for(int k=0;k<NNstate_.Nchains();k++){
            grad_.head(nparLambda_) -= NNstate_.LambdaGrad(NNstate_.VisibleStateRow(k))/double(NNstate_.Nchains());
        }
        opt_.getUpdates(grad_);
    }

    // Update rbm parameters
    void UpdateParameters(){
        auto pars=NNstate_.GetParameters();
        opt_.Update(pars);
        NNstate_.SetParameters(pars);
    }

    ////Run the tomography
    void Run(Eigen::MatrixXd & trainData,std::vector<std::vector<std::string> >& trainBases){
        //opt_.Reset();
        int index;
        int counter = 0;
        int trainSize = trainData.rows();
        int saveFrequency =  int(trainSize / bs_);
        Eigen::MatrixXd batch_samples;
        std::vector<std::vector<std::string> > batch_bases;
        std::uniform_int_distribution<int> distribution(0,trainSize-1);

        int epoch = 0;
        for(int i=0;i<epochs_;i++){
            // Randomize a batch and set the visible layer to a data point
            SetUpTrainingStep(trainData,batch_samples,trainBases,batch_bases,distribution);
            // Perform one step of optimization
            ComputeGradient(batch_samples,batch_bases);
            UpdateParameters();
            //Compute stuff and print
            if (counter == saveFrequency){
                epoch += 1;
                obs_.Scan(epoch);
                counter = 0;
            }
            counter++;
        }
    }

    //Set the value of the target wavefunction
    void setBasisRotations(std::map<std::string,Eigen::MatrixXcd> & U){
        U_ = U;
    }
    // Setup the training batch and visible layer initial configuration
    void SetUpTrainingStep(Eigen::MatrixXd & trainData,
                           Eigen::MatrixXd &batch_samples,
                           std::vector<std::vector<std::string> >& trainBases,
                           std::vector<std::vector<std::string> > &batch_bases,
                           std::uniform_int_distribution<int> & distribution){
            int index;
            // Initialize the visible layer to random data samples
            batch_samples.resize(NNstate_.Nchains(),N_);
            for(int k=0;k<NNstate_.Nchains();k++){
                index = distribution(rgen_);
                batch_samples.row(k) = trainData.row(index);
            }
            NNstate_.SetVisibleLayer(batch_samples);

            // Build the batch of data
            batch_samples.resize(bs_,N_);
            batch_bases.resize(bs_,std::vector<std::string>(N_));
            for(int k=0;k<bs_;k++){
                index = distribution(rgen_);
                batch_samples.row(k) = trainData.row(index);
                batch_bases[k] = trainBases[index];
            }
    }
};
}

#endif


/*******************************************************************************
 *
 * Implementation of bignums based on GMP, the Gnu Multiple Precision Arithmetic
 * Library (http://gmplib.org).
 *
 * Author: Arnaud J. Venet (arnaud.j.venet@nasa.gov)
 *
 * Notices:
 *
 * Copyright (c) 2011-2014 United States Government as represented by the
 * Administrator of the National Aeronautics and Space Administration.
 * All Rights Reserved.
 *
 * Disclaimers:
 *
 * No Warranty: THE SUBJECT SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTY OF
 * ANY KIND, EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED
 * TO, ANY WARRANTY THAT THE SUBJECT SOFTWARE WILL CONFORM TO SPECIFICATIONS,
 * ANY IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
 * OR FREEDOM FROM INFRINGEMENT, ANY WARRANTY THAT THE SUBJECT SOFTWARE WILL BE
 * ERROR FREE, OR ANY WARRANTY THAT DOCUMENTATION, IF PROVIDED, WILL CONFORM TO
 * THE SUBJECT SOFTWARE. THIS AGREEMENT DOES NOT, IN ANY MANNER, CONSTITUTE AN
 * ENDORSEMENT BY GOVERNMENT AGENCY OR ANY PRIOR RECIPIENT OF ANY RESULTS,
 * RESULTING DESIGNS, HARDWARE, SOFTWARE PRODUCTS OR ANY OTHER APPLICATIONS
 * RESULTING FROM USE OF THE SUBJECT SOFTWARE.  FURTHER, GOVERNMENT AGENCY
 * DISCLAIMS ALL WARRANTIES AND LIABILITIES REGARDING THIRD-PARTY SOFTWARE,
 * IF PRESENT IN THE ORIGINAL SOFTWARE, AND DISTRIBUTES IT "AS IS."
 *
 * Waiver and Indemnity:  RECIPIENT AGREES TO WAIVE ANY AND ALL CLAIMS AGAINST
 * THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL
 * AS ANY PRIOR RECIPIENT.  IF RECIPIENT'S USE OF THE SUBJECT SOFTWARE RESULTS
 * IN ANY LIABILITIES, DEMANDS, DAMAGES, EXPENSES OR LOSSES ARISING FROM SUCH
 * USE, INCLUDING ANY DAMAGES FROM PRODUCTS BASED ON, OR RESULTING FROM,
 * RECIPIENT'S USE OF THE SUBJECT SOFTWARE, RECIPIENT SHALL INDEMNIFY AND HOLD
 * HARMLESS THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS,
 * AS WELL AS ANY PRIOR RECIPIENT, TO THE EXTENT PERMITTED BY LAW.
 * RECIPIENT'S SOLE REMEDY FOR ANY SUCH MATTER SHALL BE THE IMMEDIATE,
 * UNILATERAL TERMINATION OF THIS AGREEMENT.
 *
 ******************************************************************************/

#pragma once

#include <climits>
#include <gmpxx.h>
#include <boost/functional/hash.hpp>
#include <crab/common/types.hpp>

namespace ikos {

class z_number {
  friend class q_number;

private:
  mpz_class _n;

public:

  z_number(mpz_class n) : _n(n) {}

  static z_number from_ulong(unsigned long n) {
    mpz_class b(n);
    return z_number(b);
  }

  static z_number from_slong(signed long n) {
    mpz_class b(n);
    return z_number(b);
  }

  // overloaded typecast operators
  explicit operator long() const {
    if (_n.fits_slong_p ()) {
      return _n.get_si ();
    }
    else {
      CRAB_ERROR("mpz_class ", _n.get_str(), " does not fit into a signed long integer");
    }
  }

  explicit operator int() const {
    if (_n.fits_sint_p ()) {
      // get_si returns a signed long so we cast it to int
      return (int) _n.get_si ();
    }
    else {
      CRAB_ERROR("mpz_class ", _n.get_str(), " does not fit into a signed integer");
    }
  }

  explicit operator mpz_class() const {
    return _n;
  }

public:

  z_number() : _n(0) {}

  z_number(std::string s) {
    try {
      this->_n = s;
    } catch (std::invalid_argument& e) {
      CRAB_ERROR ("z_number: invalid string in constructor", s);
    }
  }

  z_number(signed long long int n) : _n((signed long int) n) {
    if (n > LONG_MAX) {
      CRAB_ERROR(n, " cannot fit into a signed long int: use another mpz_class constructor");
    }
  }

  std::string get_str () const {
    return _n.get_str();
  }

  bool fits_sint() const {
    return _n.fits_sint_p();
  }

  bool fits_slong() const {
    return _n.fits_slong_p();
  }

  z_number operator+(z_number x) const {
    mpz_class r = this->_n + x._n;
    return z_number(r);
  }

  z_number operator*(z_number x) const {
    mpz_class r = this->_n * x._n;
    return z_number(r);
  }

  z_number operator-(z_number x) const {
    mpz_class r = this->_n - x._n;
    return z_number(r);
  }

  z_number operator-() const {
    mpz_class r = -this->_n;
    return z_number(r);
  }

  z_number operator/(z_number x) const {
    if (x._n == 0) {
      CRAB_ERROR("z_number: division by zero [1]");
    } else {
      mpz_class r = this->_n / x._n;
      return z_number(r);
    }
  }

  z_number operator%(z_number x) const {
    if (x._n == 0) {
      CRAB_ERROR("z_number: division by zero [2]");
    } else {
      mpz_class r = this->_n % x._n;
      return z_number(r);
    }
  }

  z_number& operator+=(z_number x) {
    this->_n += x._n;
    return *this;
  }

  z_number& operator*=(z_number x) {
    this->_n *= x._n;
    return *this;
  }

  z_number& operator-=(z_number x) {
    this->_n -= x._n;
    return *this;
  }

  z_number& operator/=(z_number x) {
    if (x._n == 0) {
      CRAB_ERROR("z_number: division by zero [3]");
    } else {
      this->_n /= x._n;
      return *this;
    }
  }

  z_number& operator%=(z_number x) {
    if (x._n == 0) {
      CRAB_ERROR("z_number: division by zero [4]");
    } else {
      this->_n %= x._n;
      return *this;
    }
  }

  z_number& operator--() {
    --(this->_n);
    return *this;
  }

  z_number& operator++() {
    ++(this->_n);
    return *this;
  }

  z_number operator++(int) {
    z_number r(*this);
    ++(*this);
    return r;
  }

  z_number operator--(int) {
    z_number r(*this);
    --(*this);
    return r;
  }

  bool operator==(z_number x) const { return this->_n == x._n; }

  bool operator!=(z_number x) const { return this->_n != x._n; }

  bool operator<(z_number x) const { return this->_n < x._n; }

  bool operator<=(z_number x) const { return this->_n <= x._n; }

  bool operator>(z_number x) const { return this->_n > x._n; }

  bool operator>=(z_number x) const { return this->_n >= x._n; }

  z_number operator&(z_number x) const { return z_number(this->_n & x._n); }

  z_number operator|(z_number x) const { return z_number(this->_n | x._n); }

  z_number operator^(z_number x) const { return z_number(this->_n ^ x._n); }

  z_number operator<<(z_number x) const {
    mpz_t tmp;
    mpz_init(tmp);
    mpz_mul_2exp(tmp, this->_n.get_mpz_t(), mpz_get_ui(x._n.get_mpz_t()));
    mpz_class result(tmp);
    return z_number(result);
  }

  z_number operator>>(z_number x) const {
    mpz_class tmp(this->_n);
    return z_number(tmp.operator>>=(mpz_get_ui(x._n.get_mpz_t())));
  }

  z_number fill_ones() const {
    assert(this->_n >= 0);
    if (this->_n == 0) {
      return z_number(0);
    }

    mpz_class result;
    for (result = 1; result < this->_n; result = 2 * result + 1)
      ;
    return z_number(result);
  }

  void write(crab::crab_os& o) { o << this->_n.get_str(); }

}; // class z_number

inline crab::crab_os& operator<<(crab::crab_os& o, z_number z) {
  z.write(o);
  return o;
}

inline std::size_t hash_value(const z_number& n) {
  boost::hash<std::string> hasher;
  return hasher(n.get_str());
}

class q_number {
private:
  mpq_class _n;

public:

  q_number() : _n(0) {}

  q_number(mpq_class n) : _n(n) {}

  q_number(std::string s) {
    try {
      this->_n = s;
      this->_n.canonicalize();
    } catch (std::invalid_argument& e) {
      CRAB_ERROR("q_number: invalid string in constructor ",s);
    }
  }

  q_number(double n): _n(n) { this->_n.canonicalize(); }

  q_number(z_number n) : _n(n._n) { this->_n.canonicalize(); }

  q_number(z_number n, z_number d) : _n(n._n, d._n) { this->_n.canonicalize(); }

  explicit operator mpq_class() const {
    return _n;
  }

  std::string get_str () const {
    return _n.get_str();
  }

  q_number operator+(q_number x) const {
    mpq_class r = this->_n + x._n;
    return q_number(r);
  }

  q_number operator*(q_number x) const {
    mpq_class r = this->_n * x._n;
    return q_number(r);
  }

  q_number operator-(q_number x) const {
    mpq_class r = this->_n - x._n;
    return q_number(r);
  }

  q_number operator-() const {
    mpq_class r = -this->_n;
    return q_number(r);
  }

  q_number operator/(q_number x) const {
    if (x._n == 0) {
      CRAB_ERROR("q_number: division by zero [1]");
    } else {
      mpq_class r = this->_n / x._n;
      return q_number(r);
    }
  }

  q_number& operator+=(q_number x) {
    this->_n += x._n;
    return *this;
  }

  q_number& operator*=(q_number x) {
    this->_n *= x._n;
    return *this;
  }

  q_number& operator-=(q_number x) {
    this->_n -= x._n;
    return *this;
  }

  q_number& operator/=(q_number x) {
    if (x._n == 0) {
      CRAB_ERROR("q_number: division by zero [2]");
    } else {
      this->_n /= x._n;
      return *this;
    }
  }

  q_number& operator--() {
    --(this->_n);
    return *this;
  }

  q_number& operator++() {
    ++(this->_n);
    return *this;
  }

  q_number operator--(int) {
    q_number r(*this);
    --(*this);
    return r;
  }

  q_number operator++(int) {
    q_number r(*this);
    ++(*this);
    return r;
  }

  bool operator==(q_number x) const { return this->_n == x._n; }

  bool operator!=(q_number x) const { return this->_n != x._n; }

  bool operator<(q_number x) const { return this->_n < x._n; }

  bool operator<=(q_number x) const { return this->_n <= x._n; }

  bool operator>(q_number x) const { return this->_n > x._n; }

  bool operator>=(q_number x) const { return this->_n >= x._n; }

  z_number numerator() const { return z_number(this->_n.get_num()); }

  z_number denominator() const { return z_number(this->_n.get_den()); }

  z_number round_to_upper() const {
    z_number num = numerator();
    z_number den = denominator();
    z_number q = num / den;
    z_number r = num % den;
    if (r == 0 || *this < 0) {
      return q;
    } else {
      return q + 1;
    }
  }

  z_number round_to_lower() const {
    z_number num = numerator();
    z_number den = denominator();
    z_number q = num / den;
    z_number r = num % den;
    if (r == 0 || *this > 0) {
      return q;
    } else {
      return q - 1;
    }
  }

  void write(crab::crab_os& o) { o << this->_n.get_str(); }

}; // class q_number

inline crab::crab_os& operator<<(crab::crab_os& o, q_number q) {
  q.write(o);
  return o;
}

inline std::size_t hash_value(const q_number& n) {
  boost::hash<std::string> hasher;
  return hasher(n.get_str());
}
}


#ifndef STAN_MATH_PRIM_ARR_FUNCTOR_INTEGRATE_ODE_RK45_HPP
#define STAN_MATH_PRIM_ARR_FUNCTOR_INTEGRATE_ODE_RK45_HPP

#include <stan/math/prim/arr/err/check_nonzero_size.hpp>
#include <stan/math/prim/arr/err/check_ordered.hpp>
#include <stan/math/prim/arr/functor/coupled_ode_system.hpp>
#include <stan/math/prim/arr/functor/coupled_ode_observer.hpp>
#include <stan/math/prim/scal/fun/value_of.hpp>
#include <stan/math/prim/scal/err/check_less.hpp>
#include <stan/math/prim/scal/err/check_finite.hpp>
#include <stan/math/prim/scal/err/invalid_argument.hpp>
#include <stan/math/prim/scal/meta/return_type.hpp>
#include <boost/version.hpp>
#if BOOST_VERSION == 106400
#include <boost/serialization/array_wrapper.hpp>
#endif
#include <boost/numeric/odeint.hpp>
#include <ostream>
#include <vector>

namespace stan {
namespace math {

/**
 * Return the solutions for the specified system of ordinary
 * differential equations given the specified initial state,
 * initial times, times of desired solution, and parameters and
 * data, writing error and warning messages to the specified
 * stream.
 *
 * Warning: If the system of equations is stiff, roughly
 * defined by having varying time scales across dimensions, then
 * this solver is likely to be slow.
 *
 * This function is templated to allow the initial times to be
 * either data or autodiff variables and the parameters to be data
 * or autodiff variables.  The autodiff-based implementation for
 * reverse-mode are defined in namespace stan::math
 * and may be invoked via argument-dependent lookup by including
 * their headers.
 *
 * This function uses the Dormand-Prince
 * method as implemented in Boost's
 * boost::numeric::odeint::runge_kutta_dopri5 integrator.
 *
 * @tparam F type of ODE system function.
 * @tparam T1 type of scalars for initial values.
 * @tparam T2 type of scalars for parameters.
 * @param[in] f functor for the base ordinary differential equation.
 * @param[in] y0 initial state.
 * @param[in] t0 initial time.
 * @param[in] ts times of the desired solutions, in strictly
 * increasing order, all greater than the initial time.
 * @param[in] theta parameter vector for the ODE.
 * @param[in] x continuous data vector for the ODE.
 * @param[in] x_int integer data vector for the ODE.
 * @param[out] msgs the print stream for warning messages.
 * @param[in] relative_tolerance relative tolerance parameter
 *   for Boost's ode solver. Defaults to 1e-6.
 * @param[in] absolute_tolerance absolute tolerance parameter
 *   for Boost's ode solver. Defaults to 1e-6.
 * @param[in] max_num_steps maximum number of steps to take within
 *   the Boost ode solver.
 * @return a vector of states, each state being a vector of the
 * same size as the state variable, corresponding to a time in ts.
 */
template <typename F, typename T1, typename T2>
std::vector<std::vector<typename stan::return_type<T1, T2>::type> >
integrate_ode_rk45(const F& f, const std::vector<T1>& y0, double t0,
                   const std::vector<double>& ts, const std::vector<T2>& theta,
                   const std::vector<double>& x, const std::vector<int>& x_int,
                   std::ostream* msgs = nullptr,
                   double relative_tolerance = 1e-6,
                   double absolute_tolerance = 1e-6, int max_num_steps = 1E6) {
  using boost::numeric::odeint::integrate_times;
  using boost::numeric::odeint::make_dense_output;
  using boost::numeric::odeint::max_step_checker;
  using boost::numeric::odeint::runge_kutta_dopri5;

  check_finite("integrate_ode_rk45", "initial state", y0);
  check_finite("integrate_ode_rk45", "initial time", t0);
  check_finite("integrate_ode_rk45", "times", ts);
  check_finite("integrate_ode_rk45", "parameter vector", theta);
  check_finite("integrate_ode_rk45", "continuous data", x);

  check_nonzero_size("integrate_ode_rk45", "times", ts);
  check_nonzero_size("integrate_ode_rk45", "initial state", y0);
  check_ordered("integrate_ode_rk45", "times", ts);
  check_less("integrate_ode_rk45", "initial time", t0, ts[0]);

  if (relative_tolerance <= 0)
    invalid_argument("integrate_ode_rk45", "relative_tolerance,",
                     relative_tolerance, "", ", must be greater than 0");
  if (absolute_tolerance <= 0)
    invalid_argument("integrate_ode_rk45", "absolute_tolerance,",
                     absolute_tolerance, "", ", must be greater than 0");
  if (max_num_steps <= 0)
    invalid_argument("integrate_ode_rk45", "max_num_steps,", max_num_steps, "",
                     ", must be greater than 0");

  // creates basic or coupled system by template specializations
  coupled_ode_system<F, T1, T2> coupled_system(f, y0, theta, x, x_int, msgs);

  // first time in the vector must be time of initial state
  std::vector<double> ts_vec(ts.size() + 1);
  ts_vec[0] = t0;
  for (size_t n = 0; n < ts.size(); n++)
    ts_vec[n + 1] = ts[n];

  std::vector<std::vector<double> > y_coupled(ts_vec.size());
  coupled_ode_observer observer(y_coupled);

  // the coupled system creates the coupled initial state
  std::vector<double> initial_coupled_state = coupled_system.initial_state();

  const double step_size = 0.1;
  integrate_times(
      make_dense_output(absolute_tolerance, relative_tolerance,
                        runge_kutta_dopri5<std::vector<double>, double,
                                           std::vector<double>, double>()),
      boost::ref(coupled_system), initial_coupled_state, boost::begin(ts_vec),
      boost::end(ts_vec), step_size, observer, max_step_checker(max_num_steps));

  // remove the first state corresponding to the initial value
  y_coupled.erase(y_coupled.begin());

  // the coupled system also encapsulates the decoupling operation
  return coupled_system.decouple_states(y_coupled);
}

}  // namespace math

}  // namespace stan

#endif


#include "parser.h"

#include <boost/spirit/home/x3.hpp>

namespace x3 = boost::spirit::x3;

namespace new_math {

x3::rule<class expression, uint64_t> const expression = "expression";
x3::rule<class term, uint64_t> const term = "term";

auto const set = [](auto &context) { x3::_val(context) = x3::_attr(context); };
auto const add = [](auto &context) { x3::_val(context) += x3::_attr(context); };
auto const mul = [](auto &context) { x3::_val(context) *= x3::_attr(context); };

auto const expression_def = term[set] >>
                            *(('+' >> term[add]) | ('*' >> term[mul]));
auto const term_def = x3::uint_ | ('(' >> expression >> ')');

BOOST_SPIRIT_DEFINE(expression, term);

} // namespace new_math

namespace advanced_math {

auto const set = [](auto &context) { x3::_val(context) = x3::_attr(context); };
auto const add = [](auto &context) { x3::_val(context) += x3::_attr(context); };
auto const mul = [](auto &context) { x3::_val(context) *= x3::_attr(context); };

x3::rule<class expression, uint64_t> const expression = "expression";
x3::rule<class term, uint64_t> const term = "term";
x3::rule<class factor, uint64_t> const factor = "factor";

auto const expression_def = factor[set] >> *('*' >> factor[mul]);
auto const factor_def = term[set] >> *('+' >> term[add]);
auto const term_def = x3::uint_ | ('(' >> expression >> ')');

BOOST_SPIRIT_DEFINE(expression, term, factor);

} // namespace advanced_math

uint64_t parse_expression(std::string input, Math math) {
  uint64_t result = 0;

  auto first = input.begin();
  auto last = input.end();

  auto passed = math == Math::New
                    ? x3::phrase_parse(first, last, new_math::expression,
                                       x3::space, result)
                    : x3::phrase_parse(first, last, advanced_math::expression,
                                       x3::space, result);

  if (not passed || first != last) {
    throw std::invalid_argument{"parse error"};
  }

  return result;
}


#include "tkdCmdParser.h"

#include "vnl/vnl_vector.h"
#include "vnl/vnl_matrix.h"
#include "vnl/vnl_matrix_fixed.h"
#include <vnl/algo/vnl_cholesky.h>
#include <vnl/vnl_transpose.h>

#include <cmath>
#include <ctime>

#include <list>

#include <itkImage.h>
#include <itkImageFileReader.h>
#include <itkImageFileWriter.h>
#include <itkImageRegionIteratorWithIndex.h>
#include <itkImageLinearIteratorWithIndex.h>
#include <itkImageRegionIterator.h>
#include <itkDiffusionTensor3D.h>

#include <boost/tokenizer.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random.hpp>
#include <boost/math/distributions/gamma.hpp>

/**
 * dki namespace.
 */
namespace dki
{

	struct parameters
	{
		std::string inputFileName;
		std::string maskFileName;
		std::string bvecsFileName;
		std::string bvalsFileName;
		std::string outputFileName;
		std::string sep;
		int iterations;
		bool regularize;
	};

	typedef double PixelType;
	typedef unsigned char MaskPixelType;

	typedef itk::Image< PixelType, 4 > ImageType;
	typedef itk::Image< MaskPixelType, 3 > MaskImageType;

	typedef itk::Image< PixelType, 3 > OutputImageType;

	typedef itk::DiffusionTensor3D< PixelType > DTITensorType;
	typedef vnl_vector_fixed< PixelType, 15 > DKITensorType;

	typedef itk::ImageFileReader< ImageType > ReaderType;
	typedef itk::ImageFileReader< MaskImageType > MaskReaderType;

	typedef vnl_vector< PixelType > VectorType;
	typedef vnl_matrix< PixelType > MatrixType;

	typedef itk::ImageLinearConstIteratorWithIndex< ImageType > ConstIterator4DType;
	typedef itk::ImageRegionConstIteratorWithIndex< MaskImageType > ConstIterator3DType;
	typedef itk::ImageRegionIteratorWithIndex< OutputImageType > Iterator3DType;

	typedef itk::ImageFileWriter< OutputImageType > WriterType;

	typedef boost::tokenizer< boost::char_separator< char > > TokType;
	typedef boost::mt19937 random_number_type;

	/**
	 * Return 2-rand DTI tensor.
	 *
	 * ITK tensor layout =>
	 * 					| 0  1  2  |
	 *			        | X  3  4  |
	 *       			| X  X  5  |
	 */
	static DTITensorType GetDTITensor( const VectorType& theta )
	{
		DTITensorType D( 0.0 );

		D( 0, 0 ) = theta( 1 );
		D( 0, 1 ) = theta( 4 );
		D( 0, 2 ) = theta( 6 );
		D( 1, 1 ) = theta( 2 );
		D( 1, 2 ) = theta( 5 );
		D( 2, 2 ) = theta( 3 );

		return D;
	}

	/**
	 * Return 4-rand DKI tensor (for now a fixed vector with 15 elements).
	 *
	 * theta_dki = [W_1111, W_2222, W_3333,
	 * 				W_1112, W_2223, W_2221,
	 * 				W_2223, W_3331, W_3332,
	 * 				W_1122, W_1133, W_2233,
	 * 				W_1123, W_1223, W_1233]
	 */
	DKITensorType GetDKITensor( const VectorType& theta )
	{
		DKITensorType tensor;
		for( unsigned int i = 0; i < 15; i++ )
			tensor( i ) = theta( 7 + i );

		return tensor;
	}

	/**
	 * Predicate to remove non positive semi-definite DTI/DKI tensors.
	 */
	class non_positive_definite
	{
		public:

		/**
		 * Check for non positive semi-definite tensor.
		 */
		bool operator() ( const VectorType& beta, bool dki_tensor = false )
		{
			DTITensorType DT = GetDTITensor( beta );

			DTITensorType::EigenValuesArrayType eigenValues;
			DT.ComputeEigenValues( eigenValues );

			bool non_negative = false;

			for( unsigned int i = 0; i < 3; i++ )
				if( eigenValues[ i ] < 0 )
					non_negative = true;

			if( dki_tensor )
			{
				DKITensorType DK = GetDKITensor( beta );
				// TODO calculate if DK tensor is postive semi-definite ...
			}

			return non_negative;
		}
	};

	/**
	 * Date, 06-01-2011.
	 *
	 * Implement Gibbs sampler for linear DTI fits.
	 *
	 * Ref: "Introduction to Applied Bayesian Statistics and Estimation for Social Scientists, Scott M. Lynch, 2007. Springer"
	 */
	class DKIFit
    {
        public:

		ImageType::Pointer Input;
		MaskImageType::Pointer Mask;

		MatrixType X_dti; // dti design matrix
		MatrixType X_dki; // dki design matrix
		MatrixType Z; // covariance matrix
		VectorType B; // bvals
		MatrixType R; // bvecs

		OutputImageType::Pointer S0_dti;
		OutputImageType::Pointer S0_dki;

		OutputImageType::Pointer FA_dti;
		OutputImageType::Pointer FA_dki;

		OutputImageType::Pointer Trace_dti;
		OutputImageType::Pointer Trace_dki;

		OutputImageType::Pointer AKC;
		OutputImageType::Pointer Kmax;
		OutputImageType::Pointer Kmin;

		/**
		 * Start DKI fit.
		 */
		void Run( const parameters& args )
        {
			SetImage( args.inputFileName );
			SetMask( args.maskFileName, Input.GetPointer() );
			SetBVals( args.bvalsFileName, args.sep );
			SetBVecs( args.bvecsFileName, args.sep );
			SetDTIDesignMatrix();
			SetDKIDesignMatrix();
			AllocateOutput( Input.GetPointer() );
			Fit( args.iterations, args.regularize );
			Write( args.outputFileName );
        }

        protected:

		/**
		 * Return cholesky decomposition for symmetric matrix.
		 *
		 * Sqrt( A ) using cholesky decomposition is twice as fast as svd/qr.
		 *
		 * The cholesky decomposition decomposes symmetric A = LL'
		 * where L is lower triangular.
		 */
		MatrixType Cholesky( const MatrixType& A )
		{
			/* Test Cholesky ...
			MatrixType A( 3, 3 );
			A( 0, 0 ) = 2; A( 0, 1 ) = 0; A( 0, 2 ) = 0.1;
			A( 1, 0 ) = 0; A( 1, 1 ) = 2; A( 1, 2 ) = 0;
			A( 2, 0 ) = 0.1; A( 2, 1 ) = 0; A( 2, 2 ) = 4;
			MatrixType L = Cholesky( A );
			std::cout << "Cholesky: " << L * L.transpose() << std::endl;
			*/
			vnl_cholesky::Operation op = vnl_cholesky::quiet;

			vnl_cholesky chol( A, op );

			return chol.upper_triangle(); // A = U'U
		}

		/**
		 * Normal distribution sampler.
		 */
		VectorType SampleNormal( random_number_type& ran, unsigned int total, double mean, double sigma )
		{
			using namespace boost;

			// select Gaussian probability distribution
			normal_distribution< double > norm_dist( mean, sigma );

			// bind random number generator to distribution, forming a function
			variate_generator< random_number_type&, normal_distribution< double > > sampler( ran, norm_dist );

			VectorType samples( total );

			for( unsigned int i = 0; i < total; i++ )
				samples( i ) = sampler(); // sample from the distribution

			return samples;
		}

		/**
		 * Inverse gamma distribution sampler.
		 *
		 * Implemented as 1 / gamma, with boost library 1.45 switch to inverse_gamma_distribution()
		 *
		 * shape =~ alpha
		 * scale =~ beta
		 */
		VectorType SampleIGamma( random_number_type& ran, unsigned int total, double alpha, double beta )
		{
			using namespace boost;

		    gamma_distribution< double > gamma_dist( alpha );

		    variate_generator< random_number_type&, gamma_distribution< double > > sampler( ran, gamma_dist );

		    VectorType samples( total );

		    for( unsigned int i = 0; i < total; i++ )
		    	samples( i ) = 1 / ( beta * sampler() ); // sample gamma distribution with scale = 1, and invert

		    return samples;
		}

		/**
		 * Design matrix with ln(S0) as constant .
		 *
		 * See: "C.G. Koay et al / Journal of Magnetic Resonance 128 (2006) pag. 116"
		 *
		 * DTI tensor thus: ln(S0), Dxx, Dyy, Dzz, Dxy, Dyz, Dxz.
		 */
		void SetDTIDesignMatrix()
		{
			MatrixType A( R.rows(), 7, 0 ); // 5 x 7

			for ( unsigned int i = 0; i < R.rows(); i++ )
			{
				double b = B( i );
				double Gx = R( i, 0 );
				double Gy = R( i, 1 );
				double Gz = R( i, 2 );

				A( i, 0 ) = 1; // ln(S0 )

				A( i, 1 ) = -b * Gx * Gx;

				A( i, 2 ) = -b * Gy * Gy;

				A( i, 3 ) = -b * Gz * Gz;

				A( i, 4 ) = -2 * b * Gx * Gy;

				A( i, 5 ) = -2 * b * Gy * Gz;

				A( i, 6 ) = -2 * b * Gx * Gz;
			}

			X_dti = A;
		}

		/**
		 * Bayesian Linear least squares fit.
         *
		 * // TODO add covariance matrix W => to get linear 'weighted' regression.
		 * beta = (X' * X)^-1 * X' * ln( y )
		 */
		VectorType FitVoxel( const MatrixType& X, const VectorType& y, unsigned int iterations,
				unsigned int k, unsigned int n, bool regularize )
		{
			// (X' * X)^-1

			vnl_svd< PixelType > svd( X.transpose() * X );
			MatrixType XTXI = svd.inverse();

			// ln( y )

			VectorType Y( y.size() );
			for( unsigned int i = 0; i < y.size(); i++ )
				Y( i ) = log( y( i ) );

			// (X' * X)^-1 * X' * ln( y )

			VectorType beta = XTXI * X.transpose() * Y;

			// sample sigma from it's inverse gamma marginal distribution

			// Create a Mersenne twister random number generator
			random_number_type eng( static_cast< unsigned int > ( std::time( 0 ) ) );

			PixelType shape = 0.5 * ( n - k );
			PixelType scale = 0.5 * dot_product( ( Y - X * beta ), ( Y - X * beta ) );
			VectorType s2 = SampleIGamma( eng, iterations, shape, scale );

			// -----------------------------------
			// sample distribution for beta (MVN):
			// -----------------------------------
			// mean => inv(X'X) * (X'Y)				[fitted beta]
			// variance => sigma^2 * inv(X'X)		[sqrt using cholesky decomp.]

			std::list< VectorType > bs;
			for( unsigned int i = 0; i < iterations; i++ )
				bs.push_back( beta + SampleNormal( eng, k, 0, 1 ) * Cholesky( s2( i ) * XTXI ) );

			// remove non positive semi definite tensors
			if (regularize )
				bs.remove_if( non_positive_definite() );

			// if nothing left, return zero tensor ...

			if( bs.empty() )
			{
				//std::cout << "Voxel found with sampled tensors not positive semi-definite!" << std::endl;
				return VectorType( beta.size(), 0 ); // return empty tensor
			}

			// get mean tensor

			VectorType mean_beta( beta.size(), 0 );

			for( std::list< VectorType >::iterator it = bs.begin(); it != bs.end(); ++it )
				mean_beta += *it;

			return mean_beta /= bs.size();
		}

		/**
		 * Read bvals from file.
		 */
		void SetBVals( const std::string& bvalsFileName, const std::string& sep )
		{
			// init
			VectorType b( GetNumberOfRows( bvalsFileName ), 0 );

			// open data ...
			std::ifstream in( bvalsFileName.c_str() );
			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read from: " << bvalsFileName << std::endl;
				exit( EXIT_FAILURE );
			}

			std::string line;
			unsigned int rowIndex = 0;

			while ( getline( in, line ) )
			{
				TokType tok( line, boost::char_separator< char >( sep.c_str() ) );
				try
				{
					for ( TokType::iterator id = tok.begin(); id != tok.end(); ++id )
					{
						b( rowIndex ) = boost::lexical_cast< PixelType >( *id );
					}
					rowIndex++;
				}
				catch ( boost::bad_lexical_cast& e )
				{
					std::cout << "*** WARNING ***: could not parse " << bvalsFileName << std::endl;
					e.what();
					exit( EXIT_FAILURE );
				}
			}

			in.close();

			B = b;
		}

		void SetBVecs( const std::string& bvecsFileName, const std::string& sep )
		{
			// init
			MatrixType r( GetNumberOfRows( bvecsFileName ), 3, 0 );

			// open data ...
			std::ifstream in( bvecsFileName.c_str() );
			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read from: " << bvecsFileName << std::endl;
				exit( EXIT_FAILURE );
			}

			std::string line;
			unsigned int rowIndex = 0;

			while ( getline( in, line ) )
			{
				TokType tok( line, boost::char_separator< char >( sep.c_str() ) );
				try
				{
					unsigned int colIndex = 0;
					for ( TokType::iterator id = tok.begin(); id != tok.end(); ++id )
					{
						r( rowIndex, colIndex ) = boost::lexical_cast< PixelType >( *id );
						colIndex++;
					}
					rowIndex++;
				}
				catch ( boost::bad_lexical_cast& e )
				{
					std::cout << "*** WARNING ***: could not parse " << bvecsFileName << std::endl;
					std::cout << e.what() << std::endl;
					exit( EXIT_FAILURE );
				}
			}

			in.close();

			R = r;
		}

		/**
		 * Build X_dki from bvecs and bvals input files.
		 *
		 * DTI tensor elements:
		 * --------------------
		 * [D_11, D_12, D_13, D_22, D_23, D_33]
		 *
		 * DKI tensor elements:
		 * --------------------
		 * [W_1111, W_2222, W_3333,
		 * 	W_1112, W_2223, W_2221, W_2223, W_3331, W_3332,
		 *  W_1122, W_1133, W_2233,
		 * 	W_1123, W_1223, W_1233]
		 */
		void SetDKIDesignMatrix()
		{
			// check ...
			if( B.size() != R.rows() )
			{
				std::cerr << "*** ERROR ***: bvals and bvecs size does not match!" << std::endl;
				exit( EXIT_FAILURE );
			}

			// init ( e.g. [52 x 22] ) ...
			X_dki = MatrixType( R.rows(), 22, 0 );

			for ( unsigned int i = 0; i < R.rows(); i++ )
			{
				double b = B( i );
				double Gx = R( i, 0 );
				double Gy = R( i, 1 );
				double Gz = R( i, 2 );

				// ln( S0 )

				X_dki( i, 0 ) = 1;

				// DTI tensor

				X_dki( i, 1 ) = -b * Gx * Gx;
				X_dki( i, 2 ) = -b * Gy * Gy;
				X_dki( i, 3 ) = -b * Gz * Gz;
				X_dki( i, 4 ) = -2 * b * Gx * Gy;
				X_dki( i, 5 ) = -2 * b * Gy * Gz;
				X_dki( i, 6 ) = -2 * b * Gx * Gz;

				// DKI tensor

				PixelType c = ( b * b ) / 6;

				X_dki( i, 7 ) = c * Gx * Gx * Gx * Gx; // W_1111
				X_dki( i, 8 ) = c * Gy * Gy * Gy * Gy; // W_2222
				X_dki( i, 9 ) = c * Gz * Gz * Gz * Gz; // W_3333

				X_dki( i, 10 ) = 4 * c * Gx * Gx * Gx * Gy; // W_1112
				X_dki( i, 11 ) = 4 * c * Gx * Gx * Gx * Gy; // W_2223
				X_dki( i, 12 ) = 4 * c * Gx * Gy * Gy * Gy; // W_2221
				X_dki( i, 13 ) = 4 * c * Gy * Gy * Gy * Gz; // W_2223
				X_dki( i, 14 ) = 4 * c * Gx * Gz * Gz * Gz; // W_3331
				X_dki( i, 15 ) = 4 * c * Gy * Gz * Gz * Gz; // W_3332

				X_dki( i, 16 ) = 6 * c * Gx * Gx * Gy * Gy; // W_1122
				X_dki( i, 17 ) = 6 * c * Gx * Gx * Gz * Gz; // W_1133
				X_dki( i, 18 ) = 6 * c * Gy * Gy * Gz * Gz; // W_2233

				X_dki( i, 19 ) = 12 * c * Gx * Gx * Gy * Gz; // W_1123
				X_dki( i, 20 ) = 12 * c * Gx * Gy * Gy * Gz; // W_1223
				X_dki( i, 21 ) = 12 * c * Gx * Gy * Gz * Gz; // W_1233
			}
		}

		/**
		 * Return S0.
		 */
		PixelType GetS0( const VectorType& theta )
		{
			return std::exp( theta( 0 ) );
		}

		/**
		 * Return largest AKC value.
		 *
		 * See:
		 * "Principal invariant and inherent parameters of diffusion kurtosis tensors"
		 * "L. Qi et al / J. Math. Anal. Appl. 349 (2009) 165 - 180."
		 *
		 * return [mean(k), min(k), max(k)]
		 */
		VectorType GetKurtosisValues( const PixelType MD, const VectorType& D_app, const DKITensorType& W )
		{
			// get number of b-zeros ...
			unsigned int b0s = 0;
			for( unsigned int i = 0; i < B.size(); i++ )
				if( B( i ) == 0 )
					b0s++;

			VectorType K( R.rows() - b0s, 0 );

			unsigned int dwi_index = 0;

			for ( unsigned int i = 0; i < R.rows(); i++ )
			{
				if( B( i ) != 0 ) // only dwi ...
				{
					double Gx = R( i, 0 );
					double Gy = R( i, 1 );
					double Gz = R( i, 2 );

					PixelType tmp0 = MD * MD / ( D_app( i ) * D_app( i ) ); // D_app contains b0s ...

					PixelType tmp1 = W( 0 ) * Gx * Gx * Gx * Gx + W( 1 ) * Gy * Gy * Gy * Gy + W( 2 ) * Gz * Gz * Gz * Gz;

					PixelType tmp2 = 4 * ( W( 3 ) * Gx * Gx * Gx * Gy + W( 4 ) * Gx * Gx * Gx * Gy + W( 5 ) * Gx * Gy * Gy * Gy + W( 6 ) * Gy * Gy * Gy * Gz + W( 7 ) * Gx * Gz * Gz * Gz + W( 8 ) * Gy * Gz * Gz * Gz );

					PixelType tmp3 = 6 * ( W( 9 ) * Gx * Gx * Gy * Gy + W( 10 ) * Gx * Gx * Gz * Gz + W( 11 ) * Gy * Gy * Gz * Gz );

					PixelType tmp4 = 12 * ( W( 12 ) * Gx * Gx * Gy * Gz + W( 13 ) * Gx * Gy * Gy * Gz + W( 14 ) * Gx * Gy * Gz * Gz );

					K( dwi_index ) = tmp0 * ( tmp1 + tmp2 + tmp3 + tmp4 );

					dwi_index++;
				}
			}

			VectorType output( 3, 0 ); // average K, min K, max K

			output( 0 ) = K.mean();
			output( 1 ) = K.min_value();
			output( 2 ) = K.max_value();

			return output;
		}

		/**
		 * Return number of training data points.
		 */
		unsigned int GetNumberOfRows( const std::string& inputFile )
		{
			std::ifstream in( inputFile.c_str() );

			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read from: " << inputFile << std::endl;
				exit( EXIT_FAILURE );
			}
			std::string line;
			unsigned int result = 0;

			while ( getline( in, line ) )
				result++;

			in.close();
			return result;
		}

		/**
		 * MCMC DTI fitting.
		 */
		void Fit( unsigned int iterations, bool regularize )
		{
			ConstIterator4DType it( Input, Input->GetLargestPossibleRegion() );
			ConstIterator3DType mit( Mask, Mask->GetLargestPossibleRegion() );

			Iterator3DType itS0_dti( S0_dti, S0_dti->GetLargestPossibleRegion() );
			Iterator3DType itS0_dki( S0_dki, S0_dki->GetLargestPossibleRegion() );

			Iterator3DType itFA_dti( FA_dti, FA_dti->GetLargestPossibleRegion() );
			Iterator3DType itFA_dki( FA_dki, FA_dki->GetLargestPossibleRegion() );

			Iterator3DType itTrace_dti( Trace_dti, Trace_dti->GetLargestPossibleRegion() );
			Iterator3DType itTrace_dki( Trace_dki, Trace_dki->GetLargestPossibleRegion() );

			Iterator3DType itAKC( AKC, AKC->GetLargestPossibleRegion() );
			Iterator3DType itKmin( Kmin, Kmin->GetLargestPossibleRegion() );
			Iterator3DType itKmax( Kmax, Kmax->GetLargestPossibleRegion() );

			it.SetDirection( 3 );
			it.GoToBegin();
			mit.GoToBegin();

			itS0_dti.GoToBegin();
			itS0_dki.GoToBegin();

			itFA_dti.GoToBegin();
			itFA_dki.GoToBegin();

			itTrace_dti.GoToBegin();
			itTrace_dki.GoToBegin();

			itAKC.GoToBegin();
			itKmin.GoToBegin();
			itKmax.GoToBegin();

			unsigned int totalDWI = ( Input->GetLargestPossibleRegion().GetSize() )[3];

			unsigned int sliceIndex = 0;

			// for each series of DWI ...
			while ( !it.IsAtEnd(), !mit.IsAtEnd() )
			{
				if ( mit.Get() != 0 )
				{
					// slice index
					if( sliceIndex != mit.GetIndex()[ 2 ] )
					{
						std::cout << "Processing slice: " << sliceIndex << std::endl;
						sliceIndex = mit.GetIndex()[ 2 ];
					}

					VectorType y( totalDWI );

					while ( !it.IsAtEndOfLine() )
					{
						y( ( it.GetIndex() )[3] ) = it.Get();
						++it;
					}

					// DTI => 7 elements (ln(S0) = 1, theta_D = 6)
					// DTI fit only ...

					unsigned int k = 7;
					unsigned int n = y.size();
					VectorType theta_dti = FitVoxel( X_dti, y, iterations, k, n, regularize );
					itS0_dti.Set( GetS0( theta_dti ) );
					DTITensorType tensor_simple = GetDTITensor( theta_dti );
					itFA_dti.Set( tensor_simple.GetFractionalAnisotropy() );
					itTrace_dti.Set( tensor_simple.GetTrace() );

					// DKI => 22 elements (ln(S0) = 1, theta_D = 6, theta_K = 15)
					// DTI and DKI fit in one...

					k += 15;

					VectorType theta_dki = FitVoxel( X_dki, y, iterations, k, n, regularize );
					itS0_dki.Set( GetS0( theta_dki ) );

					DTITensorType tensor_dti = GetDTITensor( theta_dki ); // DTI tensor from DKI fit ...
					itFA_dki.Set( tensor_dti.GetFractionalAnisotropy() );
					itTrace_dki.Set( tensor_dti.GetTrace() );

					DKITensorType tensor_dki = GetDKITensor( theta_dki );
					VectorType tmp = GetKurtosisValues( tensor_dti.GetTrace(), FitDapp( y ), tensor_dki );

					itAKC.Set( tmp( 0 ) );
					itKmin.Set( tmp( 1 ) );
					itKmax.Set( tmp( 2 ) );
				}

				it.NextLine();
				++mit;

				++itS0_dti;
				++itS0_dki;

				++itFA_dti;
				++itFA_dki;

				++itTrace_dti;
				++itTrace_dki;

				++itAKC;
				++itKmin;
				++itKmax;
			}
		}

		/**
		 * Return apparant diffusion per direction.
		 *
		 * ln( Sb / S0 ) / -b = D_app
		 */
		VectorType FitDapp( const VectorType& y )
		{
			// average S0 ...

			PixelType S0 = 0;
			PixelType totalS0;

			for( unsigned int i = 0; i < B.size(); i++ )
			{
				if( B( i ) == 0 )
				{
					S0 += y( i );
					totalS0++;
				}
			}
			if( totalS0 == 0 )
			{
				std::cerr << "*** ERROR ***: no b-zeros!" << std::endl;
				exit( EXIT_FAILURE );
			}
			S0 /= totalS0;

			// D_app ...

			VectorType D_app( R.rows(), 0 );

			for( unsigned int i = 0; i < B.size(); i++ )
			{
				if( B( i ) != 0 )
					D_app( i ) = log( y( i )  / S0 ) / - B( i );
			}

			return D_app;
		}

		/**
		 * Return covariate matrix Z.
		 */
		MatrixType BuildCovarianceMatrix( const VectorType& y )
		{
			MatrixType Z( y.size(), y.size(), 0 );

			// fill diagonal
			for( unsigned int i = 0; i < Z.rows(); i++ )
				Z( i, i ) = y( i ) * y( i );

			return Z;
		}

		/**
		 * Allocate all output images to 0.
		 */
		void AllocateOutput( ImageType::ConstPointer input )
		{
			MaskImageType::Pointer output = MaskImageType::New();

			OutputImageType::RegionType region3D;
			OutputImageType::IndexType index3D;
			OutputImageType::SizeType size3D;
			OutputImageType::SpacingType spacing3D;
			OutputImageType::PointType origin3D;

			ImageType::RegionType region = input->GetLargestPossibleRegion();
			ImageType::SizeType size = region.GetSize();
			ImageType::IndexType index = region.GetIndex();
			ImageType::SpacingType spacing = input->GetSpacing();
			ImageType::PointType origin = input->GetOrigin();

			size3D[0] = size[0];
			size3D[1] = size[1];
			size3D[2] = size[2];
			index3D[0] = index[0];
			index3D[1] = index[1];
			index3D[2] = index[2];
			origin3D[0] = origin[0];
			origin3D[1] = origin[1];
			origin3D[2] = origin[2];
			spacing3D[0] = spacing[0];
			spacing3D[1] = spacing[1];
			spacing3D[2] = spacing[2];

			region3D.SetSize( size3D );
			region3D.SetIndex( index3D );

			// create
			S0_dti = OutputImageType::New();
			S0_dki = OutputImageType::New();

			FA_dti = OutputImageType::New();
			FA_dki = OutputImageType::New();

			Trace_dti = OutputImageType::New();
			Trace_dki = OutputImageType::New();

			AKC = OutputImageType::New();
			Kmax = OutputImageType::New();
			Kmin = OutputImageType::New();

			// set
			S0_dti->SetRegions( region3D );
			S0_dti->SetSpacing( spacing3D );
			S0_dti->SetOrigin( origin3D );
			S0_dti->Allocate();
			S0_dti->FillBuffer( 0 );

			S0_dki->SetRegions( region3D );
			S0_dki->SetSpacing( spacing3D );
			S0_dki->SetOrigin( origin3D );
			S0_dki->Allocate();
			S0_dki->FillBuffer( 0 );

			FA_dti->SetRegions( region3D );
			FA_dti->SetSpacing( spacing3D );
			FA_dti->SetOrigin( origin3D );
			FA_dti->Allocate();
			FA_dti->FillBuffer( 0 );

			FA_dki->SetRegions( region3D );
			FA_dki->SetSpacing( spacing3D );
			FA_dki->SetOrigin( origin3D );
			FA_dki->Allocate();
			FA_dki->FillBuffer( 0 );

			Trace_dti->SetRegions( region3D );
			Trace_dti->SetSpacing( spacing3D );
			Trace_dti->SetOrigin( origin3D );
			Trace_dti->Allocate();
			Trace_dti->FillBuffer( 0 );

			Trace_dki->SetRegions( region3D );
			Trace_dki->SetSpacing( spacing3D );
			Trace_dki->SetOrigin( origin3D );
			Trace_dki->Allocate();
			Trace_dki->FillBuffer( 0 );

			AKC->SetRegions( region3D );
			AKC->SetSpacing( spacing3D );
			AKC->SetOrigin( origin3D );
			AKC->Allocate();
			AKC->FillBuffer( 0 );

			Kmax->SetRegions( region3D );
			Kmax->SetSpacing( spacing3D );
			Kmax->SetOrigin( origin3D );
			Kmax->Allocate();
			Kmax->FillBuffer( 0 );

			Kmin->SetRegions( region3D );
			Kmin->SetSpacing( spacing3D );
			Kmin->SetOrigin( origin3D );
			Kmin->Allocate();
			Kmin->FillBuffer( 0 );
		}

		/**
		 * Set input image.
		 */
		void SetImage( const std::string& inputFileName )
		{
			if ( !inputFileName.empty() )
			{
				ReaderType::Pointer reader = ReaderType::New();
				reader->SetFileName( inputFileName );
				reader->Update();
				Input = reader->GetOutput();
			}
			else
			{
				std::cerr << "Could not read input!" << std::endl;
				exit( EXIT_FAILURE );
			}
		}

		/**
		 * Return mask if file given, else create empty mask from input file.
		 */
		void SetMask( const std::string& maskFileName, ImageType::ConstPointer input )
		{
			if ( !maskFileName.empty() )
			{
				MaskReaderType::Pointer reader = MaskReaderType::New();
				reader->SetFileName( maskFileName );
				reader->Update();
				Mask = reader->GetOutput();
			}
			else
			{
				MaskImageType::Pointer output = MaskImageType::New();
				MaskImageType::RegionType region3D;
				MaskImageType::IndexType index3D;
				MaskImageType::SizeType size3D;
				MaskImageType::SpacingType spacing3D;
				MaskImageType::PointType origin3D;

				ImageType::RegionType region = input->GetLargestPossibleRegion();
				ImageType::SizeType size = region.GetSize();
				ImageType::IndexType index = region.GetIndex();
				ImageType::SpacingType spacing = input->GetSpacing();
				ImageType::PointType origin = input->GetOrigin();

				size3D[0] = size[0];
				size3D[1] = size[1];
				size3D[2] = size[2];
				index3D[0] = index[0];
				index3D[1] = index[1];
				index3D[2] = index[2];
				origin3D[0] = origin[0];
				origin3D[1] = origin[1];
				origin3D[2] = origin[2];
				spacing3D[0] = spacing[0];
				spacing3D[1] = spacing[1];
				spacing3D[2] = spacing[2];

				region3D.SetSize( size3D );
				region3D.SetIndex( index3D );

				// set
				output->SetRegions( region3D );
				output->SetSpacing( spacing3D );
				output->SetOrigin( origin3D );
				output->Allocate();
				output->FillBuffer( 1 );

				Mask = output;
			}
		}

		/**
		 * Write output images.
		 */
		void Write( const std::string& outputFileName )
		{
			// Writer
			WriterType::Pointer writer = WriterType::New();

			// S0
			writer->SetFileName( ( outputFileName + std::string( "_S0_dti.nii.gz" ) ).c_str() );
			writer->SetInput( S0_dti );
			writer->Update();

			writer->SetFileName( ( outputFileName + std::string( "_S0_dki.nii.gz" ) ).c_str() );
			writer->SetInput( S0_dki );
			writer->Update();

			// FA
			writer->SetFileName( ( outputFileName + std::string( "_FA_dti.nii.gz" ) ).c_str() );
			writer->SetInput( FA_dti );
			writer->Update();

			writer->SetFileName( ( outputFileName + std::string( "_FA_dki.nii.gz" ) ).c_str() );
			writer->SetInput( FA_dki );
			writer->Update();

			// Trace
			writer->SetFileName( ( outputFileName + std::string( "_Trace_dti.nii.gz" ) ).c_str() );
			writer->SetInput( Trace_dti );
			writer->Update();

			writer->SetFileName( ( outputFileName + std::string( "_Trace_dki.nii.gz" ) ).c_str() );
			writer->SetInput( Trace_dki );
			writer->Update();

			// AKC
			writer->SetFileName( ( outputFileName + std::string( "_AKC.nii.gz" ) ).c_str() );
			writer->SetInput( AKC );
			writer->Update();

			// Kmin
			writer->SetFileName( ( outputFileName + std::string( "_Kmin.nii.gz" ) ).c_str() );
			writer->SetInput( Kmin );
			writer->Update();

			// Kmax
			writer->SetFileName( ( outputFileName + std::string( "_Kmax.nii.gz" ) ).c_str() );
			writer->SetInput( Kmax );
			writer->Update();
		}

    }; // end class DKIFit

} // end namespace dki

/**
 * DKI fitting routine.
 */
int main( int argc, char ** argv )
{
    tkd::CmdParser p( argv[0], "Fit voxel-wise DTI using Gibbs sampling and postitive semi-definite tensor regularization." );

    dki::parameters args;
    args.sep = " ";
    args.iterations = 1000;
    args.regularize = false;

    p.AddArgument( args.inputFileName, "input" )
        ->AddAlias( "i" )
        ->SetDescription( "Input 4D image" )
        ->SetRequired( true );

    p.AddArgument( args.bvecsFileName, "bvecs" )
        ->AddAlias( "r" )
        ->SetDescription( "Diffusion gradient vectors (FSL format)" )
        ->SetRequired( true );

    p.AddArgument( args.bvalsFileName, "bvals" )
        ->AddAlias( "b" )
        ->SetDescription( "Diffusion b-value vector (FSL format)" )
        ->SetRequired( true );

    p.AddArgument( args.outputFileName, "output" )
        ->AddAlias( "o" )
        ->SetDescription( "Output filename base" )
        ->SetRequired( true );

    p.AddArgument( args.maskFileName, "mask" )
        ->AddAlias( "m" )
        ->SetDescription( "Mask 3D image" );

    p.AddArgument( args.sep, "separation" )
        ->AddAlias( "s" )
        ->SetDescription( "Separation string in bvecs file (default: ' ')" );

    p.AddArgument( args.iterations, "iterations" )
        ->AddAlias( "it" )
        ->SetDescription( "Number of MCMC sampling iterations per voxel (default: 1000)" );

    p.AddArgument( args.regularize, "regularize" )
        ->AddAlias( "reg" )
        ->SetDescription( "Regularize posteriors by removal of non postivive semi-definite tensors (default: false)" );

    if ( !p.Parse( argc, argv ) )
    {
        p.PrintUsage( std::cout );
        return EXIT_FAILURE;
    }

    dki::DKIFit fit;
    fit.Run( args );

    return EXIT_SUCCESS;
}


#pragma once
#ifndef OPENGM_MINSTCUTBOOST_HXX
#define OPENGM_MINSTCUTBOOST_HXX

#ifndef BOOST_DISABLE_ASSERTS
#define BOOST_DISABLE_ASSERTS
#define USED_BOOST_DISABLE_ASSERTS
#endif

#include <queue>
#include <cassert>

#include <boost/graph/graph_traits.hpp>
#include <boost/graph/one_bit_color_map.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/typeof/typeof.hpp>

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/edmonds_karp_max_flow.hpp>
#include <boost/graph/push_relabel_max_flow.hpp>
#include <boost/graph/boykov_kolmogorov_max_flow.hpp>

namespace opengm {

   enum BoostMaxFlowAlgorithm {
      PUSH_RELABEL, EDMONDS_KARP, KOLMOGOROV
   };

   /// \brief Boost solvers for the min st-cut framework GraphCut
   template<class NType, class VType, BoostMaxFlowAlgorithm mfalg>
   class MinSTCutBoost {
   public:
      // Type-Definitions
      typedef NType node_type;
      typedef VType ValueType;
      typedef boost::vecS OutEdgeList;
      typedef boost::vecS VertexList;
      typedef boost::adjacency_list_traits<OutEdgeList, VertexList, boost::directedS> graph_traits;
      typedef graph_traits::edge_descriptor edge_descriptor;
      typedef graph_traits::vertex_descriptor vertex_descriptor;

      /// \cond HIDDEN_SYMBOLS
      struct Edge {
         Edge() : capacity(ValueType()), residual(ValueType()), reverse(edge_descriptor())
            {}
         ValueType capacity;
         ValueType residual;
         edge_descriptor reverse;
      };
      /// \endcond

      typedef boost::adjacency_list<OutEdgeList, VertexList, boost::directedS, size_t, Edge> graph_type;
      typedef typename boost::graph_traits<graph_type>::edge_iterator edge_iterator;
      typedef typename boost::graph_traits<graph_type>::out_edge_iterator out_edge_iterator;

      // Methods
      MinSTCutBoost();
      MinSTCutBoost(size_t numberOfNodes, size_t numberOfEdges);
      void addEdge(node_type, node_type, ValueType);
      void calculateCut(std::vector<bool>&);

   private:
      // Members
      graph_type graph_;
      size_t numberOfNodes_;
      size_t numberOfEdges_;
      static const NType S = 0;
      static const NType T = 1;
   };

   //*********************
   //** Implementation  **
   //*********************

   template<class NType, class VType, BoostMaxFlowAlgorithm mfalg>
   MinSTCutBoost<NType, VType, mfalg>::MinSTCutBoost() {
      numberOfNodes_ = 2;
      numberOfEdges_ = 0;
   }

   template<class NType, class VType, BoostMaxFlowAlgorithm mfalg>
   MinSTCutBoost<NType, VType, mfalg>::MinSTCutBoost(size_t numberOfNodes, size_t numberOfEdges) {
      numberOfNodes_ = numberOfNodes;
      numberOfEdges_ = numberOfEdges;
      graph_ = graph_type(numberOfNodes_);
      //std::cout << "#nodes : " << numberOfNodes_ << std::endl;
   }

   template<class NType, class VType, BoostMaxFlowAlgorithm mfalg>
   void MinSTCutBoost<NType, VType, mfalg>::addEdge(node_type n1, node_type n2, ValueType cost) {
      assert(n1 < numberOfNodes_);
      assert(n2 < numberOfNodes_);
      assert(cost >= 0);
      std::pair<edge_descriptor, bool> e = add_edge(n1, n2, graph_);
      std::pair<edge_descriptor, bool> er = add_edge(n2, n1, graph_);
      graph_[e.first].capacity += cost;
      graph_[e.first].reverse = er.first;
      graph_[er.first].reverse = e.first;
      //std::cout << n1 << "->" << n2 << " : " << cost << std::endl;
   }

   template<class NType, class VType, BoostMaxFlowAlgorithm mfalg>
   void MinSTCutBoost<NType, VType, mfalg>::calculateCut(std::vector<bool>& segmentation) {
      if (mfalg == KOLMOGOROV) {//Kolmogorov
         std::vector<boost::default_color_type> color(num_vertices(graph_));
         std::vector<edge_descriptor> pred(num_vertices(graph_));
         std::vector<vertex_descriptor> dist(num_vertices(graph_));
         boykov_kolmogorov_max_flow(graph_,
            get(&Edge::capacity, graph_),
            get(&Edge::residual, graph_),
            get(&Edge::reverse, graph_),
            &pred[0],
            &color[0],
            &dist[0],
            get(boost::vertex_index, graph_),
            S, T
            );
         // find (s,t)-cut set
         segmentation.resize(num_vertices(graph_));
         for (size_t j = 2; j < num_vertices(graph_); ++j) {
            if (color[j] == boost::black_color || color[j] == boost::gray_color) {
               segmentation[j] = false;
            } else if (color[j] == boost::white_color) {
               segmentation[j] = true;
            }
         }
      }
      else if (mfalg == PUSH_RELABEL) {// PushRelable

         push_relabel_max_flow(graph_, S, T,
            get(&Edge::capacity, graph_),
            get(&Edge::residual, graph_),
            get(&Edge::reverse, graph_),
            get(boost::vertex_index_t(), graph_)
            );
         // find (s,t)-cut set
         segmentation.resize(num_vertices(graph_), true);
         segmentation[S] = false; // source
         segmentation[T] = false; // sink
         typedef typename boost::property_map<graph_type, boost::vertex_index_t>::type VertexIndexMap;
         VertexIndexMap vertexIndexMap = get(boost::vertex_index, graph_);
         std::queue<vertex_descriptor> q;
         q.push(*(vertices(graph_).first)); // source
         while (!q.empty()) {
            out_edge_iterator current, end;
            boost::tie(current, end) = out_edges(q.front(), graph_);
            q.pop();
            while (current != end) {
               if (graph_[*current].residual > 0) {
                  vertex_descriptor v = target(*current, graph_);
                  if (vertexIndexMap[v] > 1 && segmentation[vertexIndexMap[v]] == true) {
                     segmentation[vertexIndexMap[v]] = false;
                     q.push(v);
                  }
               }
               ++current;
            }
         }
      }
      else if (mfalg == EDMONDS_KARP) {//EdmondsKarp
         std::vector<boost::default_color_type> color(num_vertices(graph_));
         std::vector<edge_descriptor> pred(num_vertices(graph_));
         edmonds_karp_max_flow(graph_, S, T,
            get(&Edge::capacity, graph_),
            get(&Edge::residual, graph_),
            get(&Edge::reverse, graph_),
            &color[0], &pred[0]
            );
         // find (s,t)-cut set
         segmentation.resize(num_vertices(graph_));
         for (size_t j = 2; j < num_vertices(graph_); ++j) {
            if (color[j] == boost::black_color) {
               segmentation[j] = false;
            } else if (color[j] == boost::white_color) {
               segmentation[j] = true;
            } else {
               throw std::runtime_error("At least one vertex is labeled neither black nor white.");
            }
         }
      }
      else {//UNKNOWN MaxFlowalgorithm
         throw std::runtime_error("Unknown MaxFlow-algorithm in MinSTCutBoost.hxx");
      }
      return;
   }

} // namespace opengm

#ifdef USED_BOOST_DISABLE_ASSERTS
#undef BOOST_DISABLE_ASSERTS
#undef USED_BOOST_DISABLE_ASSERTS
#endif

#endif // #ifndef OPENGM_MINSTCUTBOOST_HXX