#include "itkTractsToVectorImageFilter.h"

// VTK
#include <vtkPolyLine.h>
#include <vtkCellArray.h>
#include <vtkCellData.h>

// ITK
#include <itkTimeProbe.h>
#include <itkImageRegionIterator.h>

// misc
#define _USE_MATH_DEFINES
#include <math.h>
#include <boost/progress.hpp>

namespace itk{

static bool CompareVectorLengths(const vnl_vector_fixed< double, 3 >& v1, const vnl_vector_fixed< double, 3 >& v2)
{
    return (v1.magnitude()>v2.magnitude());
}

template< class PixelType >
TractsToVectorImageFilter< PixelType >::TractsToVectorImageFilter():
    m_AngularThreshold(0.7),
    m_Epsilon(0.999),
    m_MaskImage(NULL),
    m_NormalizeVectors(false),
    m_UseWorkingCopy(true),
    m_UseTrilinearInterpolation(false),
    m_MaxNumDirections(3),
    m_Thres(0.5),
    m_NumDirectionsImage(NULL)
{
    this->SetNumberOfRequiredOutputs(1);
}

template< class PixelType >
TractsToVectorImageFilter< PixelType >::~TractsToVectorImageFilter()
{
}

template< class PixelType >
vnl_vector_fixed<double, 3> TractsToVectorImageFilter< PixelType >::GetVnlVector(double point[])
{
    vnl_vector_fixed<double, 3> vnlVector;
    vnlVector[0] = point[0];
    vnlVector[1] = point[1];
    vnlVector[2] = point[2];
    return vnlVector;
}

template< class PixelType >
itk::Point<double, 3> TractsToVectorImageFilter< PixelType >::GetItkPoint(double point[])
{
    itk::Point<double, 3> itkPoint;
    itkPoint[0] = point[0];
    itkPoint[1] = point[1];
    itkPoint[2] = point[2];
    return itkPoint;
}

template< class PixelType >
void TractsToVectorImageFilter< PixelType >::GenerateData()
{
  mitk::BaseGeometry::Pointer geometry = m_FiberBundle->GetGeometry();

    // calculate new image parameters
    itk::Vector<double> spacing;
    itk::Point<double> origin;
    itk::Matrix<double, 3, 3> direction;
    ImageRegion<3> imageRegion;
    if (!m_MaskImage.IsNull())
    {
        spacing = m_MaskImage->GetSpacing();
        imageRegion = m_MaskImage->GetLargestPossibleRegion();
        origin = m_MaskImage->GetOrigin();
        direction = m_MaskImage->GetDirection();
    }
    else
    {
        spacing = geometry->GetSpacing();
        origin = geometry->GetOrigin();
        mitk::BaseGeometry::BoundsArrayType bounds = geometry->GetBounds();
        origin[0] += bounds.GetElement(0);
        origin[1] += bounds.GetElement(2);
        origin[2] += bounds.GetElement(4);

        for (int i=0; i<3; i++)
            for (int j=0; j<3; j++)
                direction[j][i] = geometry->GetMatrixColumn(i)[j];
        imageRegion.SetSize(0, geometry->GetExtent(0));
        imageRegion.SetSize(1, geometry->GetExtent(1));
        imageRegion.SetSize(2, geometry->GetExtent(2));

        m_MaskImage = ItkUcharImgType::New();
        m_MaskImage->SetSpacing( spacing );
        m_MaskImage->SetOrigin( origin );
        m_MaskImage->SetDirection( direction );
        m_MaskImage->SetRegions( imageRegion );
        m_MaskImage->Allocate();
        m_MaskImage->FillBuffer(1);
    }
    OutputImageType::RegionType::SizeType outImageSize = imageRegion.GetSize();
    m_OutImageSpacing = m_MaskImage->GetSpacing();
    m_ClusteredDirectionsContainer = ContainerType::New();

    // initialize crossings image
    m_CrossingsImage = ItkUcharImgType::New();
    m_CrossingsImage->SetSpacing( spacing );
    m_CrossingsImage->SetOrigin( origin );
    m_CrossingsImage->SetDirection( direction );
    m_CrossingsImage->SetRegions( imageRegion );
    m_CrossingsImage->Allocate();
    m_CrossingsImage->FillBuffer(0);

    // initialize num directions image
    m_NumDirectionsImage = ItkUcharImgType::New();
    m_NumDirectionsImage->SetSpacing( spacing );
    m_NumDirectionsImage->SetOrigin( origin );
    m_NumDirectionsImage->SetDirection( direction );
    m_NumDirectionsImage->SetRegions( imageRegion );
    m_NumDirectionsImage->Allocate();
    m_NumDirectionsImage->FillBuffer(0);

    // resample fiber bundle
    double minSpacing = 1;
    if(m_OutImageSpacing[0]<m_OutImageSpacing[1] && m_OutImageSpacing[0]<m_OutImageSpacing[2])
        minSpacing = m_OutImageSpacing[0];
    else if (m_OutImageSpacing[1] < m_OutImageSpacing[2])
        minSpacing = m_OutImageSpacing[1];
    else
        minSpacing = m_OutImageSpacing[2];

    if (m_UseWorkingCopy)
        m_FiberBundle = m_FiberBundle->GetDeepCopy();

    // resample fiber bundle for sufficient voxel coverage
    m_FiberBundle->ResampleFibers(minSpacing/3);

    // iterate over all fibers
    vtkSmartPointer<vtkPolyData> fiberPolyData = m_FiberBundle->GetFiberPolyData();
    vtkSmartPointer<vtkCellArray> vLines = fiberPolyData->GetLines();
    vLines->InitTraversal();
    int numFibers = m_FiberBundle->GetNumFibers();
    itk::TimeProbe clock;
    m_DirectionsContainer = ContainerType::New();

    if (m_UseTrilinearInterpolation)
        MITK_INFO << "Generating directions from tractogram (trilinear interpolation)";
    else
        MITK_INFO << "Generating directions from tractogram";

    boost::progress_display disp(numFibers);
    for( int i=0; i<numFibers; i++ )
    {
        ++disp;
        clock.Start();
        vtkIdType   numPoints(0);
        vtkIdType*  points(NULL);
        vLines->GetNextCell ( numPoints, points );
        if (numPoints<2)
            continue;

        itk::Index<3> index; index.Fill(0);
        itk::ContinuousIndex<double, 3> contIndex;
        vnl_vector_fixed<double, 3> dir, wDir;
        itk::Point<double, 3> worldPos;
        vnl_vector<double> v;
        for( int j=0; j<numPoints-1; j++)
        {
            double* temp = fiberPolyData->GetPoint(points[j]);
            worldPos = GetItkPoint(temp);
            v = GetVnlVector(temp);

            dir = GetVnlVector(fiberPolyData->GetPoint(points[j+1]))-v;
            dir.normalize();

            m_MaskImage->TransformPhysicalPointToIndex(worldPos, index);
            m_MaskImage->TransformPhysicalPointToContinuousIndex(worldPos, contIndex);

            if (m_MaskImage->GetPixel(index)==0)
                continue;

            if (!m_UseTrilinearInterpolation)
            {
                if (index[0] < 0 || (unsigned long)index[0] >= outImageSize[0])
                    continue;
                if (index[1] < 0 || (unsigned long)index[1] >= outImageSize[1])
                    continue;
                if (index[2] < 0 || (unsigned long)index[2] >= outImageSize[2])
                    continue;

                unsigned int idx = index[0] + outImageSize[0]*(index[1] + outImageSize[1]*index[2]);
                DirectionContainerType::Pointer dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, dir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), dir);
                }
                else
                {
                    dirCont->InsertElement(0, dir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }

                continue;
            }

            double frac_x = contIndex[0] - index[0];
            double frac_y = contIndex[1] - index[1];
            double frac_z = contIndex[2] - index[2];

            if (frac_x<0)
            {
                index[0] -= 1;
                frac_x += 1;
            }
            if (frac_y<0)
            {
                index[1] -= 1;
                frac_y += 1;
            }
            if (frac_z<0)
            {
                index[2] -= 1;
                frac_z += 1;
            }

            frac_x = 1-frac_x;
            frac_y = 1-frac_y;
            frac_z = 1-frac_z;

            // int coordinates inside image?
            if (index[0] < 0 || (unsigned long)index[0] >= outImageSize[0]-1)
                continue;
            if (index[1] < 0 || (unsigned long)index[1] >= outImageSize[1]-1)
                continue;
            if (index[2] < 0 || (unsigned long)index[2] >= outImageSize[2]-1)
                continue;

            DirectionContainerType::Pointer dirCont;
            int idx;
            wDir = dir;
            double weight = (  frac_x)*(  frac_y)*(  frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]   + outImageSize[0]*(index[1]  + outImageSize[1]*index[2]  );
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (  frac_x)*(1-frac_y)*(  frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]   + outImageSize[0]*(index[1]+1+ outImageSize[1]*index[2]  );
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (  frac_x)*(  frac_y)*(1-frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]   + outImageSize[0]*(index[1]  + outImageSize[1]*index[2]+outImageSize[1]);
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (  frac_x)*(1-frac_y)*(1-frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]   + outImageSize[0]*(index[1]+1+ outImageSize[1]*index[2]+outImageSize[1]);
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (1-frac_x)*(  frac_y)*(  frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]+1 + outImageSize[0]*(index[1]  + outImageSize[1]*index[2]  );
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (1-frac_x)*(  frac_y)*(1-frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]+1 + outImageSize[0]*(index[1]  + outImageSize[1]*index[2]+outImageSize[1]);
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (1-frac_x)*(1-frac_y)*(  frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]+1 + outImageSize[0]*(index[1]+1+ outImageSize[1]*index[2]  );
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }

            wDir = dir;
            weight = (1-frac_x)*(1-frac_y)*(1-frac_z);
            if (weight>m_Thres)
            {
                wDir *= weight;
                idx = index[0]+1 + outImageSize[0]*(index[1]+1+ outImageSize[1]*index[2]+outImageSize[1]);
                dirCont = DirectionContainerType::New();
                if (m_DirectionsContainer->IndexExists(idx))
                {
                    dirCont = m_DirectionsContainer->GetElement(idx);
                    if (dirCont.IsNull())
                    {
                        dirCont = DirectionContainerType::New();
                        dirCont->InsertElement(0, wDir);
                        m_DirectionsContainer->InsertElement(idx, dirCont);
                    }
                    else
                        dirCont->InsertElement(dirCont->Size(), wDir);
                }
                else
                {
                    dirCont->InsertElement(0, wDir);
                    m_DirectionsContainer->InsertElement(idx, dirCont);
                }
            }
        }
        clock.Stop();
    }

    vtkSmartPointer<vtkCellArray> m_VtkCellArray = vtkSmartPointer<vtkCellArray>::New();
    vtkSmartPointer<vtkPoints>    m_VtkPoints = vtkSmartPointer<vtkPoints>::New();

    itk::ImageRegionIterator<ItkUcharImgType> dirIt(m_NumDirectionsImage, m_NumDirectionsImage->GetLargestPossibleRegion());
    itk::ImageRegionIterator<ItkUcharImgType> crossIt(m_CrossingsImage, m_CrossingsImage->GetLargestPossibleRegion());

    m_DirectionImageContainer = DirectionImageContainerType::New();
    unsigned int maxNumDirections = 0;

    MITK_INFO << "Clustering directions";
    boost::progress_display disp2(outImageSize[0]*outImageSize[1]*outImageSize[2]);
    for(crossIt.GoToBegin(); !crossIt.IsAtEnd(); ++crossIt)
    {
        ++disp2;
        OutputImageType::IndexType index = crossIt.GetIndex();
        int idx = index[0]+(index[1]+index[2]*outImageSize[1])*outImageSize[0];

        if (!m_DirectionsContainer->IndexExists(idx))
        {
            ++dirIt;
            continue;
        }
        DirectionContainerType::Pointer dirCont = m_DirectionsContainer->GetElement(idx);
        if (dirCont.IsNull() || index[0] < 0 || (unsigned long)index[0] >= outImageSize[0] || index[1] < 0 || (unsigned long)index[1] >= outImageSize[1] || index[2] < 0 || (unsigned long)index[2] >= outImageSize[2])
        {
            ++dirIt;
            continue;
        }

        std::vector< DirectionType > directions;

        for (unsigned int i=0; i<dirCont->Size(); i++)
            if (dirCont->ElementAt(i).magnitude()>0.0001)
                directions.push_back(dirCont->ElementAt(i));

        if (!directions.empty())
            directions = FastClustering(directions);

        std::sort( directions.begin(), directions.end(), CompareVectorLengths );

        if ( directions.size() > maxNumDirections )
        {
            for (unsigned int i=maxNumDirections; i<std::min<unsigned int>(directions.size(), m_MaxNumDirections); i++)
            {
                ItkDirectionImageType::Pointer directionImage = ItkDirectionImageType::New();
                directionImage->SetSpacing( spacing );
                directionImage->SetOrigin( origin );
                directionImage->SetDirection( direction );
                directionImage->SetRegions( imageRegion );
                directionImage->Allocate();
                Vector< float, 3 > nullVec; nullVec.Fill(0.0);
                directionImage->FillBuffer(nullVec);
                m_DirectionImageContainer->InsertElement(i, directionImage);
            }
            maxNumDirections = std::min<unsigned int>(directions.size(), m_MaxNumDirections);
        }

        unsigned int numDir = directions.size();
        if (numDir>m_MaxNumDirections)
            numDir = m_MaxNumDirections;

        for (unsigned int i=0; i<numDir; i++)
        {
            vtkSmartPointer<vtkPolyLine> container = vtkSmartPointer<vtkPolyLine>::New();
            itk::ContinuousIndex<double, 3> center;
            center[0] = index[0];
            center[1] = index[1];
            center[2] = index[2];
            itk::Point<double> worldCenter;
            m_MaskImage->TransformContinuousIndexToPhysicalPoint( center, worldCenter );
            DirectionType dir = directions.at(i);

            // set direction image pixel
            ItkDirectionImageType::Pointer directionImage = m_DirectionImageContainer->GetElement(i);
            Vector< float, 3 > pixel;
            pixel.SetElement(0, dir[0]);
            pixel.SetElement(1, dir[1]);
            pixel.SetElement(2, dir[2]);
            directionImage->SetPixel(index, pixel);

            // add direction to vector field (with spacing compensation)
            itk::Point<double> worldStart;
            worldStart[0] = worldCenter[0]-dir[0]/2*minSpacing;
            worldStart[1] = worldCenter[1]-dir[1]/2*minSpacing;
            worldStart[2] = worldCenter[2]-dir[2]/2*minSpacing;
            vtkIdType id = m_VtkPoints->InsertNextPoint(worldStart.GetDataPointer());
            container->GetPointIds()->InsertNextId(id);
            itk::Point<double> worldEnd;
            worldEnd[0] = worldCenter[0]+dir[0]/2*minSpacing;
            worldEnd[1] = worldCenter[1]+dir[1]/2*minSpacing;
            worldEnd[2] = worldCenter[2]+dir[2]/2*minSpacing;
            id = m_VtkPoints->InsertNextPoint(worldEnd.GetDataPointer());
            container->GetPointIds()->InsertNextId(id);
            m_VtkCellArray->InsertNextCell(container);
        }
        dirIt.Set(numDir);
        ++dirIt;
    }

    vtkSmartPointer<vtkPolyData> directionsPolyData = vtkSmartPointer<vtkPolyData>::New();
    directionsPolyData->SetPoints(m_VtkPoints);
    directionsPolyData->SetLines(m_VtkCellArray);
    m_OutputFiberBundle = mitk::FiberBundleX::New(directionsPolyData);
}

template< class PixelType >
std::vector< vnl_vector_fixed< double, 3 > > TractsToVectorImageFilter< PixelType >::FastClustering(std::vector< vnl_vector_fixed< double, 3 > >& inDirs)
{
    std::vector< vnl_vector_fixed< double, 3 > > outDirs;
    if (inDirs.empty())
        return outDirs;
    vnl_vector_fixed< double, 3 > oldMean, currentMean, workingMean;

    std::vector< vnl_vector_fixed< double, 3 > > normalizedDirs;
    std::vector< int > touched;
    for (unsigned int i=0; i<inDirs.size(); i++)
    {
        normalizedDirs.push_back(inDirs[i]);
        normalizedDirs.back().normalize();
    }

    // initialize
    double max = 0.0;
    touched.resize(inDirs.size(), 0);
    bool free = true;
    currentMean = inDirs[0];  // initialize first seed
    while (free)
    {
        oldMean.fill(0.0);

        // start mean-shift clustering
        double angle = 0.0;
        int counter = 0;
        while ((currentMean-oldMean).magnitude()>0.0001)
        {
            counter = 0;
            oldMean = currentMean;
            workingMean = oldMean;
            workingMean.normalize();
            currentMean.fill(0.0);
            for (unsigned int i=0; i<normalizedDirs.size(); i++)
            {
                angle = dot_product(workingMean, normalizedDirs[i]);
                if (angle>=m_AngularThreshold)
                {
                    currentMean += inDirs[i];
                    touched[i] = 1;
                    counter++;
                }
                else if (-angle>=m_AngularThreshold)
                {
                    currentMean -= inDirs[i];
                    touched[i] = 1;
                    counter++;
                }
            }
        }

        // found stable mean
        if (counter>0)
        {
            currentMean /= counter;
            double mag = currentMean.magnitude();

            if (mag>0)
            {
                if (mag>max)
                    max = mag;

                outDirs.push_back(currentMean);
            }
        }

        // find next unused seed
        free = false;
        for (unsigned int i=0; i<touched.size(); i++)
            if (touched[i]==0)
            {
                currentMean = inDirs[i];
                free = true;
            }
    }

    if (m_NormalizeVectors)
        for (unsigned int i=0; i<outDirs.size(); i++)
            outDirs[i].normalize();
    else if (max>0)
        for (unsigned int i=0; i<outDirs.size(); i++)
            outDirs[i] /= max;

    if (inDirs.size()==outDirs.size())
        return outDirs;
    else
        return FastClustering(outDirs);
}

template< class PixelType >
std::vector< vnl_vector_fixed< double, 3 > > TractsToVectorImageFilter< PixelType >::Clustering(std::vector< vnl_vector_fixed< double, 3 > >& inDirs)
{
    std::vector< vnl_vector_fixed< double, 3 > > outDirs;
    if (inDirs.empty())
        return outDirs;
    vnl_vector_fixed< double, 3 > oldMean, currentMean, workingMean;

    std::vector< vnl_vector_fixed< double, 3 > > normalizedDirs;
    std::vector< int > touched;
    for (std::size_t i=0; i<inDirs.size(); i++)
    {
        normalizedDirs.push_back(inDirs[i]);
        normalizedDirs.back().normalize();
    }

    // initialize
    double max = 0.0;
    touched.resize(inDirs.size(), 0);
    for (std::size_t j=0; j<inDirs.size(); j++)
    {
        currentMean = inDirs[j];
        oldMean.fill(0.0);

        // start mean-shift clustering
        double angle = 0.0;
        int counter = 0;
        while ((currentMean-oldMean).magnitude()>0.0001)
        {
            counter = 0;
            oldMean = currentMean;
            workingMean = oldMean;
            workingMean.normalize();
            currentMean.fill(0.0);
            for (std::size_t i=0; i<normalizedDirs.size(); i++)
            {
                angle = dot_product(workingMean, normalizedDirs[i]);
                if (angle>=m_AngularThreshold)
                {
                    currentMean += inDirs[i];
                    counter++;
                }
                else if (-angle>=m_AngularThreshold)
                {
                    currentMean -= inDirs[i];
                    counter++;
                }
            }
        }

        // found stable mean
        if (counter>0)
        {
            bool add = true;
            vnl_vector_fixed< double, 3 > normMean = currentMean;
            normMean.normalize();
            for (std::size_t i=0; i<outDirs.size(); i++)
            {
                vnl_vector_fixed< double, 3 > dir = outDirs[i];
                dir.normalize();
                if ((normMean-dir).magnitude()<=0.0001)
                {
                    add = false;
                    break;
                }
            }

            currentMean /= counter;
            if (add)
            {
                double mag = currentMean.magnitude();
                if (mag>0)
                {
                    if (mag>max)
                        max = mag;

                    outDirs.push_back(currentMean);
                }
            }
        }
    }

    if (m_NormalizeVectors)
        for (std::size_t i=0; i<outDirs.size(); i++)
            outDirs[i].normalize();
    else if (max>0)
        for (std::size_t i=0; i<outDirs.size(); i++)
            outDirs[i] /= max;

    if (inDirs.size()==outDirs.size())
        return outDirs;
    else
        return FastClustering(outDirs);
}

template< class PixelType >
TractsToVectorImageFilter< PixelType >::DirectionContainerType::Pointer TractsToVectorImageFilter< PixelType >::MeanShiftClustering(DirectionContainerType::Pointer dirCont)
{
    DirectionContainerType::Pointer container = DirectionContainerType::New();

    double max = 0;
    for (DirectionContainerType::ConstIterator it = dirCont->Begin(); it!=dirCont->End(); ++it)
    {
        vnl_vector_fixed<double, 3> mean = ClusterStep(dirCont, it.Value());

        if (mean.is_zero())
            continue;
        bool addMean = true;

        for (DirectionContainerType::ConstIterator it2 = container->Begin(); it2!=container->End(); ++it2)
        {
            vnl_vector_fixed<double, 3> dir = it2.Value();
            double angle = fabs(dot_product(mean, dir)/(mean.magnitude()*dir.magnitude()));
            if (angle>=m_Epsilon)
            {
                addMean = false;
                break;
            }
        }

        if (addMean)
        {
            if (m_NormalizeVectors)
                mean.normalize();
            else if (mean.magnitude()>max)
                max = mean.magnitude();
            container->InsertElement(container->Size(), mean);
        }
    }

    // max normalize voxel directions
    if (max>0 && !m_NormalizeVectors)
        for (std::size_t i=0; i<container->Size(); i++)
            container->ElementAt(i) /= max;

    if (container->Size()<dirCont->Size())
        return MeanShiftClustering(container);
    else
        return container;
}

template< class PixelType >
vnl_vector_fixed<double, 3> TractsToVectorImageFilter< PixelType >::ClusterStep(DirectionContainerType::Pointer dirCont, vnl_vector_fixed<double, 3> currentMean)
{
    vnl_vector_fixed<double, 3> newMean; newMean.fill(0);

    for (DirectionContainerType::ConstIterator it = dirCont->Begin(); it!=dirCont->End(); ++it)
    {
        vnl_vector_fixed<double, 3> dir = it.Value();
        double angle = dot_product(currentMean, dir)/(currentMean.magnitude()*dir.magnitude());
        if (angle>=m_AngularThreshold)
            newMean += dir;
        else if (-angle>=m_AngularThreshold)
            newMean -= dir;
    }

    if (fabs(dot_product(currentMean, newMean)/(currentMean.magnitude()*newMean.magnitude()))>=m_Epsilon || newMean.is_zero())
        return newMean;
    else
        return ClusterStep(dirCont, newMean);
}
}


#include "precompiled.h"
#include "finite_difference_weights.h"

#ifndef AUTOMATIC_PRECOMPILATION
#include <utility>
#include <tbb/tick_count.h>
#include <tbb/partitioner.h>
#include <boost/throw_exception.hpp>
#endif

#include "Eigen.h"
#include "utility.h"
#include "Math.h"

#ifndef AUTOMATIC_PRECOMPILATION
#include <pretty_printer.h>
#endif

#include "finite_difference_stencils.h"
#include "logging.h"
#include "finite_difference_weights_algo.h"
#include "grid_striding_info_inline.h"
#include "global_config.h"
#include "exceptions.h"
#include "PDE.h"

namespace fipster { namespace finite_difference_weights {

	// #################################################
	// ##############   PHASE 2:			############
	// ##############   Node Factories      ############
	// #################################################

	fdweights::sender_ptr fdweights::setup_node( const arg_t& arg )
	{
		auto PDE=arg.first.PDE;

		auto pde = dynamic_pointer_cast<const BSiso>(PDE);
		if(pde) return create_node(body_t<const BSiso>(arg,pde));

		// check for different PDEs
		// auto pde = dynamic_cast<const DifferenPDE* >(PDE);
		// if(pde) return create_node(body_t<const ...>(arg,pde));

		FIPSTER_THROW_EXCEPTION(runtime_error("PDE type "+PDE->type+" not implemented!"));
		return 0;
	}

	// #################################################
	// ##############   PHASE 3:			############
	// ##############   Node Bodies         ############
	// #################################################

	using namespace Eigen;

	template<class PDE_core>
	body_t<PDE_core>::body_t( const fd_weights_arg_t& args
														, shared_ptr<const PDE_core> pde_ptr )
		:	pde(pde_ptr),args(args),D(args.second->D)
	{}

	//############################################################################
	//Constructs the multi_indices corresponding to partial derivatives
	int references::generatePDindices(uint order)
	{
		uint i__,i__1,i__2,i__3,k,j;
		bool more;
		/* Calculate number of partial derivatives to include */
		int numpds = 0;
		i__1 = order;
		for (i__ = 1; i__ <= i__1; ++i__) {
			i__2 = i__ + D - 1;
			i__3 = D - 1;
			numpds += (int) binom(i__2, i__3);
		}

		multi_indices.resize(numpds,D);
		position_t ind(D);
		more=0;
		j = 0;
		/* Calculate multi indices for partial derivatives */
		i__1 = numpds;
		int h,t;
		for (i__ = 0; i__ < i__1; ++i__) {
			if (! more) {
				++j;
				if (j > order)
					BOOST_THROW_EXCEPTION(invalid_argument(
					"Number of Partial derivatives is to high. Specified order exceeded"));
			}
			comp_next(j, D, &ind[0], &more, &h, &t);
			for (k = 0; k < D; ++k) {
				multi_indices(i__,k) = ind[k];
			}
		}

		return numpds;
	}

	//############################################################################
	// Compute line-wise constant factorials in denominator for weight optimization systems
	void references::generateFactorialDenoms(double HigherOrderWeight){

		factorialdenoms.resize(multi_indices.rows());

		uint j;

		for (int i = 0; i < factorialdenoms.size(); ++i) {
			uint sum=0;
			factorialdenoms[i]=1.0;
			for (j=0; j < D; ++j) {
				factorialdenoms[i] /= factorial(multi_indices(i,j));
				if(multi_indices(i,j)!=0) sum++;
			}
			if (sum>2)
				factorialdenoms[i] *= HigherOrderWeight;
		}
	}

	template<class F,class G, class H, class I >
	void test_A_o(const F& A_o, const G& A_io, const H& gsi, const I& stencil_points){
		auto& grid=*gsi.grid;
		auto end_b = grid.template b_end<J_SGS>();
		set<boundary_index<J_SGS-1>> obinds;
		for(auto bit = boundary_iterator<J_SGS>(grid); bit != end_b; ++bit)
		{
			obinds.clear();

			//loop over neighbors
			const auto s_end=end(gsi.sorted_neighs);
			for(auto s=begin(gsi.sorted_neighs);s!=s_end;++s){
				//get neighbor position
				auto pos=bit.pos()+stencil_points.at(s->second);
				//check, if points lies on outer boundary:
				boundary_index<J_SGS-1> obind;
				if(obind.make_ind_from_pos<true>(pos,grid)){
					FIPSTER_ASSERT(A_o.coeff(bit.b_ind,obind)
										==A_io.at(s->second).at(bit.g_ind));
					if(A_io.at(s->second).at(bit.g_ind)!=0)
						obinds.insert(obind);
				}
			}

			//compare number of nonzeros in this A_o row and the number of unique accesses
			FIPSTER_ASSERT(obinds.size()==
											(uint)A_o->innerVector(bit.b_ind.ind).nonZeros());
		}
	}

	//############################################################################
	/** Generate the finite difference weights for the selected stencil,
		to get a finite difference approximation of the differential operator.

		It's exactly done as described in  the paper
		"Ito, K., & Toivanen, J. (2009). Lagrange Multiplier Approach with
		Optimized Finite Difference Stencils for Pricing American Options
		under Stochastic Volatility. SIAM Journal on Scientific Computing,
		31(4), 2646. doi:10.1137/07070574X"

		With the following additions:
		1. The sign choices specified in PDE.h dicate that the weights
			constitute a NEGATIVE M-MATRIX.
		2. This is achieved by using non-negative least squares for the off-diagonal weights,
			as opposed to the proposed solution of a quadratic programming problem.
		3. Extension to higher dimensional PDEs
		4. Eq. 40 contains (dx_i)^delta_i. I use 0^0 = 1 (Reason: If there is no
			derivative present in one direction (delta_i=0) then it is
			ignored (i.e. factor 1). In all other cases a stepsize of
			zero in one direction (dx_i=0) will yield a 0 in the matrix d.)

		 \todo verbesserung: wie vorgeschlagen, die partiellen Ableitungen
		 die man mit einem gröberen Gitter erhällt für die optimierung
		 benutzen...

	\return 1. A 2D array, where each row corresponds to one point in the multi-dim
			grid and the i-th column represent the weights of the i-th stencil point,
			if sorted lexicically by its direction vector. The last weight
			corresponds to the unfinished (see note) center weight.
			2. The so called "grid_striding_info"

	\note To get the final center weight all non-center weights have
			to be subtracted.

	*/
	template<class PDE_core>
	typename body_t<PDE_core>::result_sptr_t body_t<PDE_core>::operator()( const input_t& )
	{
		const grid_ptr& grid = args.second;

		//gcc: was editable_result(new result_t) (i.e. possibly without initizalitaion?)
		fd_weights_result_ptr_t result = make_shared<result_t>();

		//create object to hold stuff
		references r(D);

		//#### Create Stencil Points (excl. center point), sorted, such that:
		// [ points(i) , points(points.size()-i-1) ] is
		// the i-th direction

		try{

			r.stencil_points.construct(D,args.first.stencil_type);

		}catch(boost::exception& e){
			e<<boost::errinfo_file_name(("Stencil Type "+toS(args.first.stencil_type)).c_str());
			throw;
		}

		uint proper_neighs=r.stencil_points.proper_neighs();

		//#### Allocate result memory (+1 for center weights)
		auto end_ind = result->A_io.resize(*grid,r.stencil_points.nDirs());

		//#### Construct the multi_indices corresponding to partial derivatives
		uint numpds = r.generatePDindices(args.first.partial_derivative_order);

		thread_logger()<<"WEIGHTS| for "+toS(proper_neighs)+" neighbors and "+toS(numpds)+" of PDE terms"<<endl;

		if(0){cout<<"X"<<endl;
		cout<<"WEIGHTS| for "+toS(proper_neighs)+" neighbors and "+toS(numpds)+" of PDE terms"<<endl;
		cout<<args.first<<endl;
		}
		//#### Check, if weight (unconstrained) determination is not overdetermined:
		if (numpds == proper_neighs)
			thread_logger()<<"WEIGHTS| (unconstrained weights) are NOT overdetermined"<<endl;
		else if (numpds < proper_neighs)
			thread_logger()<<"WEIGHTS| (unconstrained weights) are UNDERDETERMINED and not unique!"<<endl;

		//####  Compute line wise constant factorials in denominator
		r.generateFactorialDenoms(args.first.higher_order_weight);

		//###### Optimize weights in parallel ( minimze norm: | D*w - b |)
		{	using namespace tbb;
			optimize_body<fd_weights_result_t::A_io_t,body_t<PDE_core>> body
				(r,result->A_io,pde,*grid);
			thread_logger()<<"WEIGHTS| nSites: "<<end_ind.ind<<endl;
			tick_count t0 = tick_count::now();
			parallel_reduce(
				blocked_range<grid_index<sg>>(grid_index<sg>(),end_ind,100),
				body/*,simple_partitioner()*/);
			tick_count t1 = tick_count::now();
			thread_logger()<<"WEIGHTS| DONE, maxNorm:"<<scientific <<setprecision(3)<<
				body.max_norm <<" in " <<(t1-t0).seconds()<<"s"<<endl;
		}

		//###### create information for decoupled tridiagonal systems in each splitting direction
		//###### AND A_o
		tick_count t0 = tick_count::now();

		auto b_end = grid->b_end<sg>();

		auto& gsi = *new grid_striding_info_collection<sg>(r.stencil_points.nDirs(),r.stencil_points,grid);
		thread_logger()<< "&gsi: " << &gsi << endl;
		result->grid_striding_infos.reset(&gsi);

		auto& A_io=result->A_io;
		auto& A_o=*result->A_o;
		A_o.resize(*grid);

		//test for uniqueness of the systems:
		//map<pair<int,int>,int> m;

		A_o->reserve(VectorXi::Constant(grid->Nbound[J_SGS],proper_neighs));

		//iterate over boundary and check for system beginnings
		for(connected_set_setup_iterator<sg> it(*grid);it!=b_end;++it){

			//cout<<"at "<<it.inner_start_g_ind().ind<<endl;

			auto i = r.stencil_points.dir_it();
			for(;i.valid();++i){//iterate through all directions
				auto& dir = r.stencil_points[i];
				if(it.apply_dir(dir)){//if a system in this direction exists

					//cout<<"dir: "<<dir.transpose()<<endl;

					/*
					//test for uniqueness of the system:
					//The naive test, that identifies a system
					//via its order start and end point on the inner boundary
					//fails in D>2, as a corner (of the inner boundary) can
					//belong to more than one system (in D>2)!
					//int a=it.inner_start_b_ind().ind;int b=it.inner_end_b_ind.ind;
					//use this instead:
					int a=it.outer_start_b_ind().ind;int b=it.outer_end_b_ind().ind;
					if(a==2 && b==2){
						cout<<"pos: "<<it.it.pos.transpose()<<endl;
						cout<<"dir: "<<dir.transpose()<<endl;
					}
					m[make_pair(min(a,b),max(a,b))]++;
					if(m[make_pair(min(a,b),max(a,b))]>1){
						cout<<m[make_pair(min(a,b),max(a,b))]<<" Systems found between "<<min(a,b)<<" and "<<max(a,b)<<endl;
					}
					*/

					//add subsystem
					gsi.at(i).pushback_connected_set(it);

					bool debug = false;///\todo: set false

					//add A_o corresponding to start of system
					A_o.insert(it.inner_start_b_ind(),it.outer_start_b_ind,
								(debug?1:A_io.at_opposite(i)[it.inner_start_g_ind()]));

					//cout<<it.inner_start_g_ind().ind<<": "<<A_io.at_opposite(i)[it.inner_start_g_ind()]<<endl;
					//cout<<it.inner_end_g_ind.ind<<": "<<A_io.at(i)[it.inner_end_g_ind]<<endl;

					//add A_o corresponding to end of system
					A_o.insert(it.inner_end_b_ind,it.outer_end_b_ind,
								(debug?1:A_io.at(i)[it.inner_end_g_ind]));

					///\todo: check for natural boundaries (i.e. all A_o=0 on bound)
				}
			}
		}

		A_o->makeCompressed();

		if(global_config::get().test_A_o){
			test_A_o(A_o,A_io,gsi,r.stencil_points);
		}

		//A_o.print();

		tick_count t1 = tick_count::now();
		thread_logger()<<"WEIGHTS| DONE systems information in "<<scientific <<setprecision(3)
			<<(t1-t0).seconds()<<"s"<<endl;

		return result;
	}

}}


/******************************************************************************
 *   Copyright (C) 2006-2021 by the resistivity.net development team          *
 *   Carsten Rücker carsten@resistivity.net                                   *
 *                                                                            *
 *   Licensed under the Apache License, Version 2.0 (the "License");          *
 *   you may not use this file except in compliance with the License.         *
 *   You may obtain a copy of the License at                                  *
 *                                                                            *
 *       http://www.apache.org/licenses/LICENSE-2.0                           *
 *                                                                            *
 *   Unless required by applicable law or agreed to in writing, software      *
 *   distributed under the License is distributed on an "AS IS" BASIS,        *
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. *
 *   See the License for the specific language governing permissions and      *
 *   limitations under the License.                                           *
 *                                                                            *
 ******************************************************************************/

#include "bertJacobian.h"

#include <calculateMultiThread.h>
#include <elementmatrix.h>
#include <memwatch.h>
#include <meshentities.h>
#include <shape.h>
#include <stopwatch.h>

#include <shape.h>

namespace GIMLI{

#if USE_BOOST_THREAD
    #include <boost/thread.hpp>
    boost::mutex eraseMutex__;
#else
    #include <thread>
    #include <mutex>
    std::mutex eraseMutex__;
#endif

template < class ValueType > class CreateSensitivityColMT : public GIMLI::BaseCalcMT{
public:
  CreateSensitivityColMT(Matrix < ValueType >          & S,
                         const std::vector < Cell * >  & para,
                         const DataContainerERT        & data,
                         const Matrix < ValueType >    & pots,
                         const std::map< long, uint >  & currPatternIdx,
                         const RVector                 & weights,
                         const RVector                 & k,
                         bool calc1,
                         bool verbose)
    : BaseCalcMT(verbose), S_(&S), para_(&para), //cellMapIndex_ (&cellMapIndex),
        data_(&data), pots_(&pots), currPatternIdx_(&currPatternIdx),
        weights_(&weights), k_(&k), calc1_(calc1){
            nData_ = data.size();
            nElecs_ = data.sensorCount();
    }

    virtual ~CreateSensitivityColMT(){}

    virtual void calc(){

        if (calc1_){
            calc1();
        }else {
            calc2();
        }
    }

    virtual void calc2(){
        ElementMatrix < double > S_i;
        ElementMatrix < double > S1_i;

        Cell * cell = NULL;
        int modelIdx = 0;

        const Vector < ValueType > *va;
        const Vector < ValueType > *vb;
        const Vector < ValueType > *vm;
        const Vector < ValueType > *vn;
        Vector < ValueType > dummy((*pots_)[0].size(), ValueType(0));

        const RVector *da = &(*data_)("a");
        const RVector *db = &(*data_)("b");
        const RVector *dm = &(*data_)("m");
        const RVector *dn = &(*data_)("n");

        for (Index cellID = start_; cellID < end_; cellID ++) {

            cell    = (*para_)[cellID];
            modelIdx = cell->marker();

            if (modelIdx < 0) continue;

            S1_i.ux2uy2uz2(*cell);

            for (Index kIdx = 0; kIdx < weights_->size(); kIdx ++){
                S_i.u2(*cell);
                S_i *= (*k_)[kIdx] * (*k_)[kIdx];
                S_i += S1_i;
                int a = 0, b = 0, m = 0, n = 0;
                for (Index dataIdx = 0; dataIdx < nData_; dataIdx ++ ){

                    a = (int)(*da)[dataIdx];
                    b = (int)(*db)[dataIdx];
                    m = (int)(*dm)[dataIdx];
                    n = (int)(*dn)[dataIdx];

                    if (a > -1) va = &(*pots_)[a + nElecs_ * kIdx]; else va = &dummy;
                    if (b > -1) vb = &(*pots_)[b + nElecs_ * kIdx]; else vb = &dummy;
                    if (m > -1) vm = &(*pots_)[m + nElecs_ * kIdx]; else vm = &dummy;
                    if (n > -1) vn = &(*pots_)[n + nElecs_ * kIdx]; else vn = &dummy;

                    (*S_)[dataIdx][modelIdx] +=
                       S_i.mult((*va), (*vb), (*vm), (*vn)) * (*weights_)[kIdx];
                }
            }
        }
    }

    virtual void calc1(){
        // log(Debug, "Thread #" + str(_threadNumber) + ": on CPU " + str(schedGetCPU()) +
        //            " slice " + str(start_) + ":" + str(end_));
        bool haveCurrentPatterns = false;

        if (currPatternIdx_->size() * weights_->size() == pots_->rows()) {
        //** we have current and measurements pattern instead of pol pol potentials;
            haveCurrentPatterns = true;
        }

        ElementMatrix < double > S_i;
        Cell * cell = NULL;
        int modelIdx = 0;
        // Index si, sj;

        const Vector < ValueType > *va;
        const Vector < ValueType > *vb;
        const Vector < ValueType > *vm;
        const Vector < ValueType > *vn;

        const RVector *da = &(*data_)("a");
        const RVector *db = &(*data_)("b");
        const RVector *dm = &(*data_)("m");
        const RVector *dn = &(*data_)("n");

        Vector < ValueType > dummy((*pots_)[0].size(), ValueType(0));

        for (Index cellID = start_; cellID < end_; cellID ++) {

            cell    = (*para_)[cellID];
            modelIdx = cell->marker();

            if (modelIdx < 0) continue;

            if (verbose_) {
                // 	cout << "\r";
                // 	for (int i = 0; i < tNr; i ++) cout << "\t\t\t";
                // 	cout <<	cellID << "/" << para_->size() - 1;
            }

            S_i.ux2uy2uz2(*cell);
            Index cellNodeCount = cell->nodeCount();

//             ValueType tmpPotA = ValueType(0);
//             ValueType tmpPotM = ValueType(0);
            ValueType sum = ValueType(0);

            int a = 0, b = 0, m = 0, n = 0;

            double weightsFactor = 1.0;
            //** if weights_->size() > 1, assuming 2.5D so we need to double the weights
            //** we integrate from 0 to \infty but need we need from -\infty to \infty
            if (weights_->size() > 1) weightsFactor = 2.0;

            for (Index dataIdx = 0; dataIdx < nData_; dataIdx ++ ){

                if (haveCurrentPatterns){
                    a = currPatternIdx_->find(data_->electrodeToCurrentPattern(a, b))->second;
                    m = currPatternIdx_->find(data_->electrodeToCurrentPattern(m, n))->second;
                    b = -1;
                    n = -1;
                } else {
                    a = (int)(*da)[dataIdx];
                    b = (int)(*db)[dataIdx];
                    m = (int)(*dm)[dataIdx];
                    n = (int)(*dn)[dataIdx];
                }

                for (Index kIdx = 0; kIdx < weights_->size(); kIdx ++){
                    sum = ValueType(0);

                    if (a > -1) va = &(*pots_)[a + nElecs_ * kIdx]; else va = &dummy;
                    if (b > -1) vb = &(*pots_)[b + nElecs_ * kIdx]; else vb = &dummy;
                    if (m > -1) vm = &(*pots_)[m + nElecs_ * kIdx]; else vm = &dummy;
                    if (n > -1) vn = &(*pots_)[n + nElecs_ * kIdx]; else vn = &dummy;

                    (*S_)[dataIdx][modelIdx] += S_i.mult((*va), (*vb), (*vm), (*vn)) * (weightsFactor * (*weights_)[kIdx]);
                    continue;
//                     std::cout << cell->id() << std::endl;
                    for (Index i = 0; i < cellNodeCount; i ++){
//                         si = S_i.idx(i);
//                         std::cout << cell->node(i).id() << std::endl;
                        for (Index j = 0; j < cellNodeCount; j ++){
//                             sj = S_i.idx(j);
                            // very most time criticle section here
//                             tmpPotA  = (*va)[si] - (*vb)[si];
//                             tmpPotM  = (*vm)[sj] - (*vn)[sj];
//
//                             sum += S_i.getVal(i, j) * tmpPotA * tmpPotM;

//                             sum += S_i.getVal(i, j) *
//                                     ((*va)[si] - (*vb)[si]) *
//                                     ((*vm)[sj] - (*vn)[sj]);
                            sum += S_i.getVal(i, j) *
                                    ((*va)[S_i.idx(i)] - (*vb)[S_i.idx(i)]) *
                                    ((*vm)[S_i.idx(j)] - (*vn)[S_i.idx(j)]);

                      /*  std::cout << i<<" "<<j<<" "<< S_i.getVal(i, j)<<" "<< (*va)[S_i.idx(i)]<<" "<<
                        (*vb)[S_i.idx(i)]<<" "<< (*vm)[S_i.idx(j)]<<" "<< (*vn)[S_i.idx(j)] << std::endl;
                      */
                        }
                    }
                    if (isInfNaN(sum)){
                        std::cerr << WHERE_AM_I << std::endl;
                    }
// 	  if ((*cellMapIndex_)[cellID] - 2 < 0 ||
// 	       (*cellMapIndex_)[cellID] - 2 > nModel-1){
// 	    std::cerr << WHERE_AM_I << " index out of bounds: [0 -- " << nModel-1 << "]" << cellMapIndex[cellID] - 2 << std::endl;
// 	  }

//	  (*S_)[dataIdx][(*cellMapIndex_)[cellID] - 2] += sum * (data_->k(dataIdx) * 2.0 * (*weights_)[kIdx]);

                    /*! Dangerous!!! mt writing into matrix*/
	       //(*S_)[dataIdx][(*cellMapIndex_)[cellID]] += sum * (2.0 * (*weights_)[kIdx]);
                    {
//                   #ifdef HAVE_LIBBOOST_THREAD
//                   boost::mutex::scoped_lock lock(eraseMutex__); // slows down alot
//                   #endif

                        (*S_)[dataIdx][modelIdx] += sum * (weightsFactor * (*weights_)[kIdx]);
//                         std::cout << "b: " << dataIdx<<" "<<modelIdx<<" "<<(*S_)[dataIdx][modelIdx]<< " "
//                         << sum * (weightsFactor * (*weights_)[kIdx]) << std::endl;
                    }
//                     exit(1);
                } // for each k
//                      exit(1);
            } // for each data
        } // for each cellID
    }

protected:
    Matrix < ValueType >            * S_;
    const std::vector < Cell * >    * para_;
    const DataContainerERT          * data_;
    const Matrix < ValueType >      * pots_;
    const std::map< long, uint >    * currPatternIdx_;
    const RVector                   * weights_;
    const RVector                   * k_;
    uint                            nData_;
    uint                            nElecs_;
    bool                            calc1_;

};

bool lessCellMarker(const Cell * c1, const Cell * c2) { return c1->marker() < c2->marker(); }

template < class ValueType >
void createSensitivityCol_(Matrix < ValueType > & S,
                          const Mesh & mesh,
                          const DataContainerERT & data,
                          const Matrix < ValueType > & pots,
                          const RVector & weights,
                          const RVector & k,
                          std::vector < std::pair < Index, Index > > & matrixClusterIds,
                          uint nThreads, bool verbose){

MEMINFO

    Index nData  = data.size();
    Index nModel = max(mesh.cellMarkers()) + 1;
    Index maxRows = weights.size() * data.sensorCount();

    if (pots.rows() >= maxRows){
//         if (pots[0].size() < mesh.nodeCount()){
//             std::stringstream str; str << WHERE_AM_I << " potential matrix colsize to small. "
//                                        << pots[0].size()  << "< " << mesh.nodeCount() << std::endl;
//         }
    } else {
        std::stringstream str1; str1 << WHERE_AM_I << " potential matrix rowsize to small."
                                   << pots.rows() << " < " << maxRows << std::endl;
        throwLengthError(str1.str());
    }

    Stopwatch swatch(true);
    std::map< long, uint > currPatternIdx;
    //std::cout << "CreateSensitivityColMT " << nThreads << std::endl;

    std::vector< Cell * > cells(mesh.findCellByMarker(0, -1));
    std::sort(cells.begin(), cells.end(), lessCellMarker);

    double maxMemSize = max(0.0, getEnvironment("SENSMATMAXMEM", 0.0, verbose));
    double maxSizeNeeded = mByte((double)nData * nModel * sizeof(double));

    if (maxMemSize > 0 && verbose){
        std::cout << "Size of S: " << maxSizeNeeded << " MB" << std::endl;
    }

    //** avoid MT problems
    for (std::vector< Cell * >::iterator it = cells.begin();
         it != cells.end(); it ++){
        (*it)->pShape()->invJacobian();
    }

//     ShapeFunctionCache::instance().shapeFunctions(cells[0]->shape());
//     ShapeFunctionCache::instance().deriveShapeFunctions(cells[0]->shape(), 0);
//     ShapeFunctionCache::instance().deriveShapeFunctions(cells[0]->shape(), 1);
//     ShapeFunctionCache::instance().deriveShapeFunctions(cells[0]->shape(), 2);
//     ShapeFunctionCache::instance().shapeFunctions(*cells[0]);
//     ShapeFunctionCache::instance().deriveShapeFunctions(*cells[0], 0);
//     ShapeFunctionCache::instance().deriveShapeFunctions(*cells[0], 1);
//     ShapeFunctionCache::instance().deriveShapeFunctions(*cells[0], 2);
//     //cells[0]->createShapefunctionts();

    if (maxMemSize > 0 && maxMemSize < maxSizeNeeded){

        uint modelCluster = std::floor((double)nModel / (maxSizeNeeded / maxMemSize));

        if (modelCluster < 1) {
            throwError(WHERE_AM_I + " sorry, size of single sensitivity-row exceeds memory limitations.");
        }

        std::cout << "Size of S cluster: " << mByte((double)nData * modelCluster * sizeof(ValueType)) << " MB" << std::endl;
        std::cout << "Using model cluster " << nModel << " x " << modelCluster << std::endl;

        S.resize(nData, modelCluster);

        matrixClusterIds.clear();
        matrixClusterIds.push_back(std::pair < Index, Index >(nData, nModel));

        bool calc1 = getEnvironment("SENSMAT1", false, true);
        for (uint i = 0; i < nModel; i += modelCluster ){
MEMINFO
            Index start = i;
            Index end   = min(start + modelCluster, nModel);
            std::cout << " " << start << " " << end<< std::endl;

            S.resize(nData, end - start);

            std::vector< Cell * > cellsCluster(mesh.findCellByMarker(start, end));
            std::sort(cellsCluster.begin(), cellsCluster.end(), lessCellMarker);

MEMINFO
            // subtract marker start index
            for (std::vector< Cell * >::iterator it = cellsCluster.begin(); it != cellsCluster.end(); it ++){
                (*it)->setMarker((*it)->marker() - start);
            }

            S *= ValueType(0);
MEMINFO

            distributeCalc(CreateSensitivityColMT< ValueType >(S, cellsCluster,
                                                               data, pots,
                                                               currPatternIdx,
                                                               weights, k,
                                                               calc1,
                                                               verbose),
                           cellsCluster.size(), nThreads, verbose);

MEMINFO

            //** fight against the Lorenz butterfly
            //** 1e-8 is to coarse, need adaptive tolerance
            //S.round(1e-8);
MEMINFO

            S.save("sensPart_" + str(start) + "-" + str(end));

            matrixClusterIds.push_back(std::pair < Index, Index >(start, end));

            // add marker start index
            for (std::vector< Cell * >::iterator it = cellsCluster.begin(); it != cellsCluster.end(); it ++){
                (*it)->setMarker((*it)->marker() + start);
            }
MEMINFO
        }

        S.clear();
    } else {

        // __MS(nData << " "<< S.rows() << " "<<  nModel << " "<<  S.cols())
        if (S.rows() != nData || S.cols() != nModel) S.resize(nData, nModel);
        S *= ValueType(0);
MEMINFO

        if (verbose){
            std::cout << "S(" << numberOfCPU() << "/" << nThreads; //**check!!!
            #if USE_BOOST_THREAD
            std::cout << "-boost::mt";
            #else
            std::cout << "-std::mt";
            #endif
            std::cout << "): " << swatch.duration() << ":";
//swatch.stop(verbose);
        }
        bool calc1 = getEnvironment("SENSMAT1", false, true);
        distributeCalc(CreateSensitivityColMT< ValueType >(S, cells, data,
                                                           pots, currPatternIdx,
                                                           weights, k, calc1, verbose),
                        cells.size(), nThreads, verbose);
         if (verbose){
             swatch.stop(verbose);
         }
MEMINFO
        //** fight against the Lorenz butterfly
        //** 1e-8 is to coarse, need adaptive tolerance
        //S.round(1e-8);
    }
}

void createSensitivityCol(RMatrix & S,
                          const Mesh & mesh,
                          const DataContainerERT & data,
                          const RMatrix & pots,
                          const RVector & weights,
                          const RVector & k,
                          std::vector < std::pair < Index, Index > > & matrixClusterIds,
                          uint nThreads, bool verbose){
    createSensitivityCol_(S, mesh, data, pots, weights, k, matrixClusterIds, nThreads, verbose);
}

void createSensitivityCol(CMatrix & S,
                          const Mesh & mesh,
                          const DataContainerERT & data,
                          const CMatrix & pots,
                          const RVector & weights,
                          const RVector & k,
                          std::vector < std::pair < Index, Index > > & matrixClusterIds,
                          uint nThreads, bool verbose){
    createSensitivityCol_(S, mesh, data, pots, weights, k, matrixClusterIds, nThreads, verbose);
}

void sensitivityDCFEMSingle(const std::vector < Cell * > & para, const RVector & p1, const RVector & p2,
		       RVector & sens, bool verbose){
    uint nCells = para.size();
    if (sens.size() != nCells) sens.resize(nCells);

    ElementMatrix < double > S_i;
    double sum = 0.0, a_jk = 0.0, tmppot = 0.0;
    //  cout << nCells << std::endl;

    for (uint i = 0; i < nCells; i ++){
    //cout << "Nr. " << i << std::endl;
        S_i.ux2uy2uz2(*para[i]);

        sum = 0.0;
        for (int j = 0, jmax = para[i]->nodeCount(); j < jmax; j ++){
            for (int k = 0, kmax = para[i]->nodeCount(); k < kmax; k ++){
	       a_jk = S_i.getVal(j, k);
	       tmppot = p1[S_i.idx(j)] * p2[S_i.idx(k)];
	       sum += a_jk * tmppot;
	//	cout << "\tS_mn: " << a_jk << "\tp1*p2: " << tmppot << "\tp*S_mn: " << a_jk * tmppot << "\tsum: " << sum << std::endl;
            }
        }
        sens[i] = sum;
    }
}

RVector prepExportSensitivityData(const Mesh & mesh, const RVector & data, double logdrop){
    Index nModel = unique(sort(mesh.cellMarkers())).size();

    ASSERT_EQUAL(nModel, data.size())

    //data have always the right length since it comes from S directly
    RVector modelSizes(nModel, 0.0);
    for (Index i = 0; i < mesh.cellCount(); i ++ ){
        modelSizes[mesh.cell(i).marker()] += mesh.cell(i).size();
    }

    return logTransDropTol(data/modelSizes, logdrop, true)(mesh.cellMarkers());

    // //RVector tmp(data/mesh.cellSizes());
    // if ((uint)data.size() != (uint)mesh.cellCount()){

    //     throwLengthError(WHERE_AM_I + " Datasize missmatch: " + str(mesh.cellCount())+
    //                         " " + str(data.size()));
    // } else {
    //     //for (uint i = 0; i < tmp.size(); i ++) tmp[i] = tmp[i] / mesh.cell(i).shape().domainSize();
    // }
    // RVector tmp(data/mesh.cellSizes());

    // RVector s(sign(tmp));

    // double tmpMax = max(abs(tmp));
    // tmp /= tmpMax;

    // for (uint i = 0; i < tmp.size(); i ++) {
    //     tmp[i] = std::fabs(tmp[i] / logdrop);
    //     if (tmp[i] < 1.0) tmp[i] = 1.0;
    // }

    // tmp = log10(tmp);
    // tmp /= max(tmp) * s;
    // return tmp;
}

void exportSensitivityVTK(const std::string & fileName,
                          const Mesh & mesh, const RVector & data,
                          double logdrop){
    std::map< std::string, RVector > res;
    res.insert(std::make_pair("Sensitivity" ,
                              prepExportSensitivityData(mesh, data, logdrop)));
    mesh.exportVTK(fileName, res);
}

// void exportSensMatrixDC(const std::string & filename, const Mesh & mesh, const RMatrix & S) {
//     exportSensMatrixDC(filename, mesh, S
// }

void exportSensMatrixDC(const std::string & filename, const Mesh & mesh,
                        const RMatrix & S, const IVector & idx,
                        double logdrop) {
    std::map< std::string, RVector > res;

    for (std::map < std::string, RVector >::const_iterator
            it = mesh.dataMap().begin();
            it != mesh.dataMap().end(); it ++){
        res.insert(std::make_pair(it->first, it->second));
    }

    std::string add;
//     RVector tmp(mesh.cellCount());

    for (size_t i = 0; i < S.rows(); i ++) {
        if (i < 100000) add = "0";
        if (i < 10000) add = "00";
        if (i < 1000) add = "000";
        if (i < 100) add = "0000";
        if (i < 10) add = "00000";

//         for (uint j = 0; j < tmp.size(); j ++) tmp[j] = S[i][mesh.cell(j).marker()];

//#res.insert(std::make_pair("sens-" + add + str(i), log10(RVector(abs(tmp)))));

            res.insert(std::make_pair("sens-" + add + str(i),
                        prepExportSensitivityData(mesh, S[i], logdrop)));

   }
   mesh.exportVTK(filename, res);
}

RVector coverageDC(const RMatrix & sensMatrix) {
    RVector cov;
    if (sensMatrix.rows() > 0) {
        cov.resize(sensMatrix.cols(), 0.0);

        for (size_t i = 0; i < sensMatrix.rows(); i ++) {
            cov += abs(sensMatrix[i]);
        }
    } else {
        std::cout << "Sensmatrix invalid" << std::endl;
    }
    return cov;
}

RVector coverageDCtrans(const MatrixBase & S, const RVector & dd, const RVector & mm) {
    RVector cov;

    if (S.rows() > 0) {
        cov.resize(S.cols(), 0.0);
    } else {
        std::cout << "Sensmatrix invalid" << std::endl;
    }

    if (S.rtti() == GIMLI_MATRIX_RTTI){
        const RMatrix *Sl = dynamic_cast < const RMatrix * >(&S);

        for (size_t i = 0; i < S.rows(); i ++) {
            cov += abs((*Sl)[i] * dd[i]);
        }
    } else if (S.rtti() == GIMLI_SPARSE_MAP_MATRIX_RTTI){

        const RSparseMapMatrix * Sl = dynamic_cast< const RSparseMapMatrix * >(&S);

        for (RSparseMapMatrix::const_iterator it = Sl->begin(); it != Sl->end(); it ++){
            Index row = (*it).first.first;
            Index col = (*it).first.second;
            cov[col] += (*it).second * dd[row];
        }
    } else {
        CERR_TO_IMPL
    }

    return cov / abs(mm);
}

RVector createCoverage(const MatrixBase & S, const Mesh & mesh){
    return createCoverage(S, mesh, RVector(S.rows(), 1.0), RVector(S.cols(), 1.0));
}

RVector createCoverage(const MatrixBase & S, const Mesh & mesh,
                       const RVector & response, const RVector & model){

    RVector covModel(coverageDCtrans(S, 1.0 / response, 1.0 / model));
	RVector covMesh(covModel(mesh.cellMarkers()));
    if (mesh.cellCount() == model.size()) {
        covMesh /= mesh.cellSizes();
    } else {
        RVector modelCellSizes(covMesh.size(), 0.0);
        for (Index i = 0; i < mesh.cellCount(); i ++){
            Cell *c = &mesh.cell(i);
            modelCellSizes[c->marker()] += c->shape().domainSize();
        }
        if (min(modelCellSizes) > TOLERANCE){
            covMesh /= modelCellSizes;
        } else {
            log(Error, "Coverage fails:" + str(mesh.cellCount()) + " " + str(model.size()));
        }
    }

    return covMesh;
}
} // namespace GIMLI


#include <memory>
#include <chrono>
#include <tuple>
#include <unistd.h>
#include <gflags/gflags.h>
#include <string.h>
#include <assert.h>
#include <Eigen/StdVector>

#include "drake/solvers/snopt_solver.h"
#include "drake/systems/analysis/simulator.h"
#include "drake/systems/framework/diagram.h"
#include "drake/systems/framework/diagram_builder.h"
#include "drake/multibody/parsing/parser.h"
#include <drake/multibody/inverse_kinematics/inverse_kinematics.h>
#include "drake/geometry/geometry_visualization.h"
#include "drake/solvers/solve.h"

#include "common/find_resource.h"
#include "systems/trajectory_optimization/dircon/dircon.h"
#include "multibody/kinematic/world_point_evaluator.h"
#include "solvers/nonlinear_constraint.h"
#include "multibody/kinematic/distance_evaluator.h"
#include "multibody/multibody_utils.h"
#include "multibody/visualization_utils.h"
#include "multibody/kinematic/kinematic_constraints.h"

#include "examples/Spirit/animate_spirit.h"
#include "examples/Spirit/spirit_utils.h"

DEFINE_double(duration, 1, "The stand duration");
DEFINE_double(front2BackToeDistance, 0.35, "Nominal distance between the back and front toes.");
DEFINE_double(side2SideToeDistance, 0.2, "Nominal distance between the back and front toes.");
DEFINE_double(bodyHeight, 0.104, "The spirit body start height (defined in URDF)");
DEFINE_double(lowerHeight,0.15, "The sitting height of the bottom of the robot");
DEFINE_double(upperHeight, 0.35, "The standing height.");
DEFINE_double(inputCost, 3, "The standing height.");
DEFINE_double(velocityCost, 10, "The standing height.");
DEFINE_double(eps, 1e-2, "The wiggle room.");
DEFINE_double(optTol, 1e-4,"Optimization Tolerance");
DEFINE_double(feasTol, 1e-4,"Feasibility Tolerance");
DEFINE_bool(autodiff, false, "Double or autodiff version");
DEFINE_bool(runInitTraj, false, "Animate initial conditions?");
// Parameters which enable dircon-improving features
DEFINE_bool(scale_constraint, true, "Scale the nonlinear constraint values");
DEFINE_bool(scale_variable, false, "Scale the decision variable");

using drake::AutoDiffXd;
using drake::multibody::MultibodyPlant;
using drake::geometry::SceneGraph;
using drake::multibody::Parser;
using drake::trajectories::PiecewisePolynomial;

using Eigen::Vector3d;
using Eigen::VectorXd;
using Eigen::Matrix3d;
using Eigen::MatrixXd;

namespace dairlib {
namespace {

using systems::trajectory_optimization::DirconModeSequence;
using systems::trajectory_optimization::DirconMode;
using systems::trajectory_optimization::Dircon;
using systems::trajectory_optimization::KinematicConstraintType;

using std::vector;
using std::cout;
using std::endl;

template <typename T>
void addConstraints(const MultibodyPlant<T>& plant, Dircon<T>& trajopt){
  // Get position and velocity dictionaries
  int num_knotpoints =10; ///DEBUG
  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  auto x0 = trajopt.initial_state();
  auto xmid = trajopt.state_vars(0, (num_knotpoints - 1) / 2);
  // auto xf = trajopt.final_state();

  // Initial body positions
  trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_y")));
  trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, x0(positions_map.at("base_z")));

  // Mid body positions
  trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_y")));
  trajopt.AddBoundingBoxConstraint(FLAGS_upperHeight-FLAGS_eps, FLAGS_upperHeight+FLAGS_eps, xmid(positions_map.at("base_z")));

  return;
}

  /// See runSpiritSquat()
    // std::unique_ptr<MultibodyPlant<T>> plant_ptr,
    // MultibodyPlant<double>* plant_double_ptr,
    // std::unique_ptr<SceneGraph<double>> scene_graph_ptr,
    // double duration,
    // PiecewisePolynomial<double> init_x_traj,
    // PiecewisePolynomial<double> init_u_traj,
    // vector<PiecewisePolynomial<double>> init_l_traj,
    // vector<PiecewisePolynomial<double>> init_lc_traj,
    // vector<PiecewisePolynomial<double>> init_vc_traj
  ///
  /// Given an initial guess for state, control, and forces optimizes a standing behaviour for the spirit robot
template <typename T>
void runSpiritSquat(
    std::unique_ptr<MultibodyPlant<T>> plant_ptr,
    MultibodyPlant<double>* plant_double_ptr,
    std::unique_ptr<SceneGraph<double>> scene_graph_ptr,
    double duration,
    PiecewisePolynomial<double> init_x_traj,
    PiecewisePolynomial<double> init_u_traj,
    vector<PiecewisePolynomial<double>> init_l_traj,
    vector<PiecewisePolynomial<double>> init_lc_traj,
    vector<PiecewisePolynomial<double>> init_vc_traj
    ) {

  drake::systems::DiagramBuilder<double> builder;
  MultibodyPlant<T>& plant = *plant_ptr;
  SceneGraph<double>& scene_graph =
      *builder.AddSystem(std::move(scene_graph_ptr));

  // Get position and velocity dictionaries
  auto positions_map = multibody::makeNameToPositionsMap(plant);
  auto velocities_map = multibody::makeNameToVelocitiesMap(plant);
  /// For Spirit front left leg->0, back left leg->1, front right leg->2, back right leg->3
  /// Get the frame of each toe and attach a world point to the toe tip (frame is at toe ball center).

  int num_legs = 4;
  double toeRadius = 0.02; // Radius of toe ball
  Vector3d toeOffset(toeRadius,0,0); // vector to "contact point"
  double mu = 1; // Coeff of friction

  int num_knotpoints_per_mode = 10;

  auto sequence = DirconModeSequence<T>(plant);

  dairlib::ModeSequenceHelper msh;

  msh.addMode( // FIRST MODE 1
    (Eigen::Matrix<bool,1,4>() << 1,1,1,1 ).finished(), // contact bools
    num_knotpoints_per_mode,  // number of knot points in the collocation
    Eigen::Vector3d::UnitZ(), // normal
    Eigen::Vector3d::Zero(),  // world offset
    mu //friction
    );

  // auto [modeVector, toeEvals, toeEvalSets] = createSpiritModeSequence(plant, msh.modes , msh.knots , msh.normals , msh.offsets, msh.mus, msh.minTs, msh.maxTs);
  auto [modeVector, toeEvals, toeEvalSets] = createSpiritModeSequence(plant, msh);

  for (auto& mode : modeVector){
    for (int i = 0; i < num_legs; i++ ){
      mode->MakeConstraintRelative(i,0);
      mode->MakeConstraintRelative(i,1);
    }
    mode->SetDynamicsScale(
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17}, 1.0 / 150.0);
    mode->SetKinVelocityScale(
      {0, 1, 2, 3}, {0, 1, 2}, 1.0 / 500.0 * 500 * 1 / 1);
    sequence.AddMode(mode.get());
  }

  ///Setup trajectory optimization
  auto trajopt = Dircon<T>(sequence);

  // Set up Trajectory Optimization options
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Print file", "../snopt.out");
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Major iterations limit", 20000);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Iterations limit", 100000);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(),
                           "Major optimality tolerance",
                           FLAGS_optTol);  // target optimality
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Major feasibility tolerance", FLAGS_feasTol);
  trajopt.SetSolverOption(drake::solvers::SnoptSolver::id(), "Verify level",
                           0);  // 0

    // Add duration constraint, currently constrained not bounded
  trajopt.AddDurationBounds(0, duration*2);
  // Initialize the trajectory control state and forces
  for (int j = 0; j < sequence.num_modes(); j++) {
    trajopt.drake::systems::trajectory_optimization::MultipleShooting::
        SetInitialTrajectory(init_u_traj, init_x_traj);
    trajopt.SetInitialForceTrajectory(j, init_l_traj[j], init_lc_traj[j],
                                      init_vc_traj[j]);
  }

  /// Setup all the optimization constraints
  int num_knotpoints = trajopt.N(); // number of knot points total in the collocation
  int n_v = plant.num_velocities();
  // int n_q = plant.num_positions();
  auto u = trajopt.input();
  auto x = trajopt.state();
  auto x0 = trajopt.initial_state();
  auto xmid = trajopt.state_vars(0, (num_knotpoints - 1) / 2);
  auto xf = trajopt.final_state();
  addConstraints(plant,trajopt);

  // // Initial body positions
  // trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  // trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, x0(positions_map.at("base_y")));
  // trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, x0(positions_map.at("base_z")));

  // // Mid body positions
  // trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
  // trajopt.AddBoundingBoxConstraint(-0, 0, xmid(positions_map.at("base_y")));
  // trajopt.AddBoundingBoxConstraint(FLAGS_upperHeight-FLAGS_eps, FLAGS_upperHeight+FLAGS_eps, xmid(positions_map.at("base_z")));

  // Final Position Constraints
  bool isPeriodic = 1;
  if (isPeriodic){
    trajopt.AddLinearConstraint( x0(positions_map.at("base_x"))==xf(positions_map.at("base_x")) );
    trajopt.AddLinearConstraint( x0(positions_map.at("base_y"))==xf(positions_map.at("base_y")) );
    trajopt.AddLinearConstraint( x0(positions_map.at("base_z"))==xf(positions_map.at("base_z")) );
  }else{
    trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, xf(positions_map.at("base_x"))); // Give the initial condition room to choose the x_init position (helps with positive knee constraint)
    trajopt.AddBoundingBoxConstraint(-FLAGS_eps, FLAGS_eps, xf(positions_map.at("base_y")));
    trajopt.AddBoundingBoxConstraint(FLAGS_lowerHeight-FLAGS_eps, FLAGS_lowerHeight+FLAGS_eps, xf(positions_map.at("base_z")));
  }

  /// Decide whether or not to constrain the body orientation at all knotpoints or only at the beginning and end
  bool isConstrainOrientation = false;
  if (isConstrainOrientation){
    // Body pose constraints (keep the body flat) at all the knotpoints including the ends
    for (int i = 0; i < num_knotpoints; i++) {
      auto xi = trajopt.state(i);
      trajopt.AddBoundingBoxConstraint(1, 1, xi(positions_map.at("base_qw")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qx")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qy")));
      trajopt.AddBoundingBoxConstraint(0, 0, xi(positions_map.at("base_qz")));
    }
  }else{
    // Body pose constraints (keep the body flat) at initial state
    trajopt.AddBoundingBoxConstraint(1, 1, x0(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, x0(positions_map.at("base_qz")));

    // Body pose constraints (keep the body flat) at mid state
    trajopt.AddBoundingBoxConstraint(1, 1, xmid(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, xmid(positions_map.at("base_qz")));

    // Body pose constraints (keep the body flat) at final state
    trajopt.AddBoundingBoxConstraint(1, 1, xf(positions_map.at("base_qw")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qx")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qy")));
    trajopt.AddBoundingBoxConstraint(0, 0, xf(positions_map.at("base_qz")));
  }

  /// Start/End velocity constraints of the behavior
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   x0.tail(n_v));
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   xmid.tail(n_v));
  trajopt.AddBoundingBoxConstraint(VectorXd::Zero(n_v), VectorXd::Zero(n_v),
                                   xf.tail(n_v));

  for (int i = 0; i < num_knotpoints; i++){
    auto xi = trajopt.state(i);
    // legs lined up (front and back hips equal)
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_8") ) == xi( positions_map.at("joint_9") ) );
    trajopt.AddLinearConstraint( xi( positions_map.at("joint_10") ) == xi( positions_map.at("joint_11") )  );
  }
  setSpiritSymmetry(plant, trajopt, "sagittal");
  setSpiritJointLimits(plant, trajopt);
  setSpiritActuationLimits(plant,trajopt);

 double upperLegLength = 0.206; // length of the upper leg link

// Quick check for allowed heights for Spirit
  try{
    if (FLAGS_lowerHeight<FLAGS_bodyHeight/2){
        throw "Body to close to floor" ;
    }
    if (FLAGS_upperHeight>=upperLegLength*2){
        throw "Body can't stand to that height without leaving the ground" ;
    }
  }catch(char const* exc){
      std::cerr<< exc <<std::endl;
  }

  ///Setup the traditional cost function
  const double R = FLAGS_inputCost;  // Cost on input effort
  const MatrixXd Q = FLAGS_velocityCost  * MatrixXd::Identity(n_v, n_v); // Cost on velocity
  trajopt.AddRunningCost( x.tail(n_v).transpose() * Q * x.tail(n_v) );
  trajopt.AddRunningCost( u.transpose()*R*u );

  ///Add regularization costs
  trajopt.AddRunningCost( x0(positions_map.at("base_x")) * 10 * x0(positions_map.at("base_x")) ); // x position cost
  // trajopt.AddRunningCost( (x.segment(12,2)-x.segment(14,2)).transpose() * 1 * (x.segment(12,2)-x.segment(14,2)) );
  // trajopt.AddRunningCost( (x.segment(16,2)-x.segment(18,2)).transpose() * 1 * (x.segment(16,2)-x.segment(18,2)) );

  /// Setup the visualization during the optimization
  int num_ghosts = 3;// Number of ghosts in visualization. NOTE: there are limitations on number of ghosts based on modes and knotpoints
  std::vector<unsigned int> visualizer_poses; // Ghosts for visualizing during optimization
  visualizer_poses.push_back(num_ghosts);

  trajopt.CreateVisualizationCallback(
      dairlib::FindResourceOrThrow("examples/Spirit/spirit_drake.urdf"),
      visualizer_poses, 0.2); // setup which URDF, how many poses, and alpha transparency

  /// Run the optimization using your initial guess
  auto start = std::chrono::high_resolution_clock::now();
  const auto result = Solve(trajopt, trajopt.initial_guess());
  auto finish = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> elapsed = finish - start;
  std::cout << "Solve time:" << elapsed.count() <<std::endl;
  std::cout << "Cost:" << result.get_optimal_cost() <<std::endl;
  std::cout << (result.is_success() ? "Optimization Success" : "Optimization Fail") << std::endl;

  /// Run animation of the final trajectory
  const drake::trajectories::PiecewisePolynomial<double> pp_xtraj =
      trajopt.ReconstructStateTrajectory(result);
  multibody::connectTrajectoryVisualizer(plant_double_ptr,
      &builder, &scene_graph, pp_xtraj);
  auto diagram = builder.Build();
  while (true) {

    drake::systems::Simulator<double> simulator(*diagram);
    simulator.set_target_realtime_rate(0.25);
    simulator.Initialize();
    simulator.AdvanceTo(pp_xtraj.end_time());
    sleep(2);
  }
}
}  // namespace
}  // namespace dairlib

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  std::srand(time(0));  // Initialize random number generator.

  auto plant = std::make_unique<MultibodyPlant<double>>(0.0);
  auto plant_vis = std::make_unique<MultibodyPlant<double>>(0.0);
  auto scene_graph = std::make_unique<SceneGraph<double>>();
  Parser parser(plant.get());
  Parser parser_vis(plant_vis.get(), scene_graph.get());
  std::string full_name =
      dairlib::FindResourceOrThrow("examples/Spirit/spirit_drake.urdf");

  parser.AddModelFromFile(full_name);
  parser_vis.AddModelFromFile(full_name);

  plant->mutable_gravity_field().set_gravity_vector(-9.81 *
      Eigen::Vector3d::UnitZ());

  plant->Finalize();
  plant_vis->Finalize();
  Eigen::VectorXd x0 = Eigen::VectorXd::Zero(plant->num_positions() +
                       plant->num_velocities());

  int nu = plant->num_actuators();
  int nx = plant->num_positions() + plant->num_velocities();
  int nq = plant->num_positions();
  int nv = plant->num_velocities();
  int N = 20; // number of timesteps

  std::vector<MatrixXd> init_x;
  std::vector<MatrixXd> init_u;
  std::vector<PiecewisePolynomial<double>> init_l_traj;
  std::vector<PiecewisePolynomial<double>> init_lc_traj;
  std::vector<PiecewisePolynomial<double>> init_vc_traj;

  // Initialize state trajectory
  std::vector<double> init_time;

  VectorXd xInit(nx);
  VectorXd xMid(nx);
  VectorXd xState(nx);
  xInit = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());
  xMid = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());
  xState = Eigen::VectorXd::Zero(plant->num_positions() + plant->num_velocities());

  auto positions_map = dairlib::multibody::makeNameToPositionsMap(*plant);
  auto velocities_map = dairlib::multibody::makeNameToVelocitiesMap(*plant);
  auto actuators_map = dairlib::multibody::makeNameToActuatorsMap(*plant);
  int num_joints = 12;

  // Print joint dictionary
  std::cout<<"**********************Joints***********************"<<std::endl;
  for (auto const& element : positions_map)
    std::cout << element.first << " = " << element.second << std::endl;
  for (auto const& element : velocities_map)
    std::cout << element.first << " = " << element.second << std::endl;
  for (auto const& element : actuators_map)
    std::cout << element.first << " = " << element.second << std::endl;
  std::cout<<"***************************************************"<<std::endl;

  dairlib::nominalSpiritStand( *plant, xInit,  FLAGS_lowerHeight); //Update xInit
  dairlib::nominalSpiritStand( *plant, xMid,  FLAGS_upperHeight); //Update xMid

  VectorXd deltaX(nx);
  VectorXd averageV(nx);
  deltaX = xMid-xInit;
  averageV = 2* deltaX / FLAGS_duration;
  xInit.tail(nv-3) = (averageV.head(nq)).tail(nq-4); //Ignoring Orientation make velocity the average

  double time = 0;
  double dt = FLAGS_duration/(N-1);

  // Initial pose
  xState = xInit;

  for (int i = 0; i < N; i++) {
    time=i*dt; // calculate iteration's time
    init_time.push_back(time);

    // Switch the direction of the stand to go back to the initial state (not actually properly periodic initial)
    if ( i > (N-1)/2 ){
        xState.tail(nv) = -xInit.tail(nv);
    }
    // Integrate the positions based on constant velocity  for joints and xyz
    for (int j = 0; j < num_joints; j++){
          xState(positions_map.at("joint_" + std::to_string(j))) =
                        xState(positions_map.at("joint_" + std::to_string(j))) + xState(nq + velocities_map.at("joint_" + std::to_string(j)+"dot" )) * dt;
    }
    xState(positions_map.at("base_x")) =
                        xState(positions_map.at("base_x")) + xState(nq + velocities_map.at("base_vx")) * dt;

    xState(positions_map.at("base_y")) =
                        xState(positions_map.at("base_y")) + xState(nq + velocities_map.at("base_vy")) * dt;

    xState(positions_map.at("base_z")) =
                        xState(positions_map.at("base_z")) + xState(nq + velocities_map.at("base_vz")) * dt;
    // Save timestep state into matrix
    init_x.push_back(xState);
    init_u.push_back(Eigen::VectorXd::Zero(nu));
  }
  // Make matrix into trajectory
  auto init_x_traj = PiecewisePolynomial<double>::ZeroOrderHold(init_time, init_x);
  auto init_u_traj = PiecewisePolynomial<double>::ZeroOrderHold(init_time, init_u);

  // Four contacts so forces are 12 dimensional
  Eigen::VectorXd init_l_vec(12);
  // Initial guess
  init_l_vec << 0, 0, 3*9.81, 0, 0, 3*9.81, 0, 0, 3*9.81, 0, 0, 3*9.81; //gravity and mass distributed

  //Initialize force trajectories
  int num_modes = 1; // MAKE DYNAMIC DEBUG
  for (int j = 0; j < num_modes; j++) {
    std::vector<MatrixXd> init_l_j;
    std::vector<MatrixXd> init_lc_j;
    std::vector<MatrixXd> init_vc_j;
    std::vector<double> init_time_j;
    for (int i = 0; i < N; i++) {
      init_time_j.push_back(i*FLAGS_duration/(N-1));
      init_l_j.push_back(init_l_vec);
      init_lc_j.push_back(init_l_vec);
      init_vc_j.push_back(VectorXd::Zero(12));
    }

    auto init_l_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_l_j);
    auto init_lc_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_lc_j);
    auto init_vc_traj_j = PiecewisePolynomial<double>::ZeroOrderHold(init_time_j,init_vc_j);

    init_l_traj.push_back(init_l_traj_j);
    init_lc_traj.push_back(init_lc_traj_j);
    init_vc_traj.push_back(init_vc_traj_j);
  }

  // if (FLAGS_autodiff) {
  //   std::unique_ptr<MultibodyPlant<drake::AutoDiffXd>> plant_autodiff =
  //       drake::systems::System<double>::ToAutoDiffXd(*plant);
  //   dairlib::runSpiritSquat<drake::AutoDiffXd>(
  //     std::move(plant_autodiff), plant_vis.get(), std::move(scene_graph),
  //     FLAGS_duration, init_x_traj, init_u_traj, init_l_traj,
  //     init_lc_traj, init_vc_traj);
  // } else
  if (FLAGS_runInitTraj){
    dairlib::runAnimate<double>(
      std::move(plant), plant_vis.get(), std::move(scene_graph), init_x_traj);
  }else {
    dairlib::runSpiritSquat<double>(
      std::move(plant), plant_vis.get(), std::move(scene_graph),
      FLAGS_duration, init_x_traj, init_u_traj, init_l_traj,
      init_lc_traj, init_vc_traj);
  }

}


#include "AST.h"
#include "InterpreterContext.h"
#include "StringPool.h"
#include "VariablesScope.h"
#include <iostream>
#include <boost/range/adaptor/reversed.hpp>
#include <boost/range/algorithm.hpp>
#include <cmath>

namespace
{

template <class TFunction>
CValue ExecuteSafely(TFunction && fn)
{
    try
    {
        return fn();
    }
    catch (std::exception const&)
    {
        return CValue::FromError(std::current_exception());
    }
}

class CSinFunction : public IFunctionAst
{
public:
    CValue Call(CInterpreterContext &context, const std::vector<CValue> &arguments) const override
    {
        (void)context;
        return ExecuteSafely([&] {
            double radians = arguments.at(0).AsDouble();
            return CValue::FromDouble(std::sin(radians));
        });
    }

    unsigned GetNameId() const override
    {
        return 0;
    }
};

class CRandFunction : public IFunctionAst
{
public:
    CValue Call(CInterpreterContext &context, const std::vector<CValue> &arguments) const override
    {
        (void)context;
        return ExecuteSafely([&] {
            double minimum = arguments.at(0).AsDouble();
            double maximum = arguments.at(1).AsDouble();
            if (minimum > maximum)
            {
                return CValue::FromErrorMessage("invalid arguments for rand - range maximum is lesser than minimum.");
            }
            double rand0to1 = double(std::rand()) / std::numeric_limits<unsigned>::max();

            return CValue::FromDouble((maximum - minimum) * rand0to1 + minimum);
        });
    }

    unsigned GetNameId() const override
    {
        return 0;
    }
};

}

CInterpreterContext::CInterpreterContext(std::ostream &output, std::ostream &errors, CStringPool &pool)
    : m_pool(pool)
    , m_output(output)
    , m_errors(errors)
{
    AddBuiltin("sin", std::unique_ptr<IFunctionAst>(new CSinFunction));
    AddBuiltin("rand", std::unique_ptr<IFunctionAst>(new CRandFunction));
}

CInterpreterContext::~CInterpreterContext()
{
}

void CInterpreterContext::DefineVariable(unsigned nameId, const CValue &value)
{
    if (ValidateValue(value))
    {
        m_scopes.back()->AssignVariable(nameId, value);
    }
}

void CInterpreterContext::AssignVariable(unsigned nameId, const CValue &value)
{
    if (ValidateValue(value))
    {
        if (CVariablesScope *pScope = FindScopeWithVariable(nameId))
        {
            pScope->AssignVariable(nameId, value);
        }
        else
        {
            DefineVariable(nameId, value);
        }
    }
}

CValue CInterpreterContext::GetVariableValue(unsigned nameId) const
{
    if (CVariablesScope *pScope = FindScopeWithVariable(nameId))
    {
        return *pScope->GetVariableValue(nameId);
    }
    return CValue::FromErrorMessage("unknown variable " + m_pool.GetString(nameId));
}

void CInterpreterContext::PushScope(std::unique_ptr<CVariablesScope> &&scope)
{
    m_scopes.emplace_back(std::move(scope));
}

std::unique_ptr<CVariablesScope> CInterpreterContext::PopScope()
{
    std::unique_ptr<CVariablesScope> ret(m_scopes.back().release());
    m_scopes.pop_back();
    return ret;
}

size_t CInterpreterContext::GetScopesCount() const
{
    return m_scopes.size();
}

IFunctionAst *CInterpreterContext::GetFunction(unsigned nameId) const
{
    try
    {
        return m_functions.at(nameId);
    }
    catch (std::exception const&)
    {
        return nullptr;
    }
}

void CInterpreterContext::AddFunction(unsigned nameId, IFunctionAst *function)
{
    if (function)
    {
        m_functions[nameId] = function;
    }
}

std::string CInterpreterContext::GetStringLiteral(unsigned stringId) const
{
    return m_pool.GetString(stringId);
}

void CInterpreterContext::PrintResults(const std::vector<CValue> &values)
{
    m_output << "  ";
    for (const auto &value : values)
    {
        if (ValidateValue(value))
        {
            m_output << value.ToString();
        }
    }
    m_output << std::endl;
}

void CInterpreterContext::PrintError(const std::string &message)
{
    m_errors << "  Error: " << message << std::endl;
}

bool CInterpreterContext::ValidateValue(const CValue &value)
{
    try
    {
        value.RethrowIfException();
        return true;
    }
    catch (std::exception const& ex)
    {
        PrintError(ex.what());
        return false;
    }
}

void CInterpreterContext::SetReturnValue(boost::optional<CValue> const& valueOpt)
{
    m_returnValueOpt = valueOpt;
}

boost::optional<CValue> CInterpreterContext::GetReturnValue() const
{
    return m_returnValueOpt;
}

CVariablesScope *CInterpreterContext::FindScopeWithVariable(unsigned nameId) const
{
    auto range = boost::adaptors::reverse(m_scopes);
    auto it = boost::find_if(range, [=](const auto &pScope) {
        return pScope->HasVariable(nameId);
    });
    if (it != range.end())
    {
        return it->get();
    }
    return nullptr;
}

void CInterpreterContext::AddBuiltin(const std::string &name, std::unique_ptr<IFunctionAst> &&function)
{
    m_builtins.emplace_back(std::move(function));
    unsigned nameRand = m_pool.Insert(name);
    m_functions[nameRand] = m_builtins.back().get();
}