// ------------------------------------------------------------------------
//  Copyright (C)
//  ETH Zurich - Switzerland
//
//  Kevis-Kokitsi Maninis <kmaninis@vision.ee.ethz.ch>
//  Jordi Pont-Tuset <jponttuset@vision.ee.ethz.ch>
//  July 2016
// ------------------------------------------------------------------------
// This file is part of the COB package presented in:
//    K.K. Maninis, J. Pont-Tuset, P. Arbelaez and L. Van Gool
//    Convolutional Oriented Boundaries
//    European Conference on Computer Vision (ECCV), 2016
// Please consider citing the paper if you use this code.
// ------------------------------------------------------------------------
#ifndef UCM_FILL_GAPS_HPP
#define UCM_FILL_GAPS_HPP

#include <iostream>
#include "containers.hpp"
#include <Eigen/Dense>

void ucm_fill_gaps(Eigen::Map<Eigen::Array<double,Eigen::Dynamic,Eigen::Dynamic> >& ucm2)
{
    /* Sizes */
    std::size_t sx = ucm2.rows();
    std::size_t sy = ucm2.cols();

    /*-------------------------------------------------------------*/
    /* Create LUT of contour positions and sets of labels to merge */
    /*-------------------------------------------------------------*/
    map_type ucm_th_leaves_pairs;

    /* Scan horizontal contours*/
    for (std::size_t xx=0; xx<sx; xx+=2)
    {
        for (std::size_t yy=0; yy<sy; yy+=2)
        {
            double curr_max = 0;
            if(xx>0)
                curr_max = std::max(curr_max,ucm2(xx-1,yy));
            if(yy>0)
                curr_max = std::max(curr_max,ucm2(xx,yy-1));
            if(xx<sx-1)
                curr_max = std::max(curr_max,ucm2(xx+1,yy));
            if(yy<sy-1)
                curr_max = std::max(curr_max,ucm2(xx,yy+1));

            ucm2(xx,yy) = curr_max;
        }
    }
}

#endif


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_SWAR_FUNCTIONS_SIMD_SSE_SSE2_CUMSUM_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_SWAR_FUNCTIONS_SIMD_SSE_SSE2_CUMSUM_HPP_INCLUDED
#ifdef BOOST_SIMD_HAS_SSE2_SUPPORT
#include <boost/simd/toolbox/swar/functions/cumsum.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/dispatch/meta/strip.hpp>

/////////////////////////////////////////////////////////////////////////////
// Implementation when type A0 is type8_
/////////////////////////////////////////////////////////////////////////////
namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::cumsum_, boost::simd::tag::sse2_,
                         (A0),
                         ((simd_<type8_<A0>,boost::simd::tag::sse_>))
                        )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename dispatch::meta::as_integer<A0>::type sint;
      A0 a = a0;
      sint tmp = simd::bitwise_cast<sint>(a0);
      sint tmp1 = _mm_slli_si128(tmp,1);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,2);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,3);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,4);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,5);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,6);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,7);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,8);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,9);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,10);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,11);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,12);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,13);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,14);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp, 15);
      a += simd::bitwise_cast<A0>(tmp1);
      return a;
    }
  };

/////////////////////////////////////////////////////////////////////////////
// Implementation when type A0 is type64_
/////////////////////////////////////////////////////////////////////////////

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::cumsum_, boost::simd::tag::sse2_,
                         (A0),
                         ((simd_<type64_<A0>,boost::simd::tag::sse_>))
                        )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename dispatch::meta::as_integer<A0>::type sint;
      sint tmp = simd::bitwise_cast<sint>(a0);
      sint tmp1 = _mm_slli_si128(tmp,8);
      return a0+simd::bitwise_cast<A0>(tmp1);
    }
  };

/////////////////////////////////////////////////////////////////////////////
// Implementation when type A0 is type16_
/////////////////////////////////////////////////////////////////////////////

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::cumsum_, boost::simd::tag::sse2_,
                         (A0),
                         ((simd_<type16_<A0>,boost::simd::tag::sse_>))
                        )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename dispatch::meta::as_integer<A0>::type sint;
      A0 a = a0;
      sint tmp = simd::bitwise_cast<sint>(a0);
      sint tmp1 = _mm_slli_si128(tmp,2);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,4);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp, 6);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,8);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,10);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,12);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp,14);
      return a+simd::bitwise_cast<A0>(tmp1);
    }
  };

/////////////////////////////////////////////////////////////////////////////
// Implementation when type A0 is type32_
/////////////////////////////////////////////////////////////////////////////

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION(boost::simd::tag::cumsum_, boost::simd::tag::sse2_,
                         (A0),
                         ((simd_<type32_<A0>,boost::simd::tag::sse_>))
                        )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(1)
    {
      typedef typename dispatch::meta::as_integer<A0>::type sint;
      A0 a = a0;
      sint tmp = simd::bitwise_cast<sint>(a0);
      sint tmp1 = _mm_slli_si128(tmp, 4);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp, 8);
      a = a+simd::bitwise_cast<A0>(tmp1);
      tmp1 = _mm_slli_si128(tmp, 12);
      return a+simd::bitwise_cast<A0>(tmp1);
    }
  };
} } }
#endif
#endif


#include "animator.h"

#include <Eigen/Dense>
#include <QFile>
#include <QGLViewer/manipulatedCameraFrame.h>
#include <QGLViewer/vec.h>
#include <QKeyEvent>
#include <QtGlobal>

#include <algorithm>
#include <limits>
#include <stdexcept>
#include <utility>

#include <iostream>

const Animator::Vec3 Animator::MeshColor{Scalar(1), Scalar(1), Scalar(1)};

Animator::Animator(const QStringList &poseFiles, Animator::Scalar duration,
                   Animator::Scalar fps)
: QGLViewer()
, m_fps{fps}
, m_duration{duration}
, m_speed{0}
, m_backwards{false}
, m_framePoses()
, m_triangles()
, m_numVertices{0}
, m_currentFrame{0}
, m_vao(this)
, m_vertexVbo(QOpenGLBuffer::VertexBuffer)
, m_normalVbo(QOpenGLBuffer::VertexBuffer)
, m_colorVbo(QOpenGLBuffer::VertexBuffer)
, m_elemVbo(QOpenGLBuffer::IndexBuffer)
, m_program(this)
, m_posAttrLoc{-1}
, m_normalAttrLoc{-1}
, m_colorAttrLoc{-1}
, m_projMatAttrLoc{-1}
, m_modelViewMatAttrLoc{-1}
{
    Eigen::initParallel();
    if (poseFiles.size() < 2)
    {
        throw std::runtime_error("At least two poses are necessary");
    }
    if (fps <= Scalar(0))
    {
        throw std::runtime_error("The frame rate must be a positive value");
    }
    if (duration <= Scalar(0))
    {
        throw std::runtime_error("The duration must be a positive value");
    }

    // TODO do this in a different method while showing progress?
    qInfo() << "Reading pose files...";
    std::vector<std::shared_ptr<Mesh>> poses;
    poses.reserve(poseFiles.size());
    for (const auto &file : poseFiles) {
        poses.push_back(Animator::readMesh(file));
    }
    generateAnimation(poses);
}

Animator::~Animator()
{
    destroyVertexBuffers();
}

void Animator::init()
{
    setWindowTitle("Animator");

    initializeOpenGLFunctions();

    // Disable spinning
    camera()->frame()->setSpinningSensitivity(100.0);

    if (!hasOpenGLFeature(QOpenGLFunctions::Shaders) ||
        !hasOpenGLFeature(QOpenGLFunctions::Buffers))
    {
        throw std::runtime_error("Unsupported version of OpenGL");
    }

    setBackgroundColor(Qt::black);

    // Shaders
    loadShaders();
    createVertexBuffers();

    restoreStateFromFile();

    // Camera
    setupCamera();

    // Control
    setKeyDescription(Qt::Key_B, "Toggle backwards animation");
    setKeyDescription(Qt::Key_Plus, "Accelerate animation");
    setKeyDescription(Qt::Key_Minus, "Decelerate animation");

    // Set up animation
    allocateVertexBuffers();
    updateAnimationPeriod();
    startAnimation();
}

void Animator::loadShaders()
{

    m_program.addShaderFromSourceFile(QOpenGLShader::Vertex,
                                      ":/shaders/vertex.glsl");
    m_program.addShaderFromSourceFile(QOpenGLShader::Fragment,
                                      ":/shaders/fragment.glsl");
    m_program.link();

    // Attribute locations
    m_posAttrLoc = m_program.attributeLocation("position");
    m_normalAttrLoc = m_program.attributeLocation("normal");
    m_colorAttrLoc = m_program.attributeLocation("color");
    m_projMatAttrLoc = m_program.uniformLocation("projection");
    m_modelViewMatAttrLoc = m_program.uniformLocation("modelView");
}

void Animator::createVertexBuffers()
{
    m_vertexVbo.bind();
    m_vertexVbo.create();
    m_vertexVbo.setUsagePattern(QOpenGLBuffer::StreamDraw);
    m_normalVbo.create();
    m_normalVbo.bind();
    m_normalVbo.setUsagePattern(QOpenGLBuffer::StreamDraw);
    m_colorVbo.create();
    m_colorVbo.bind();
    m_colorVbo.setUsagePattern(QOpenGLBuffer::StreamDraw);
    m_elemVbo.create();
    m_elemVbo.bind();
    m_elemVbo.setUsagePattern(QOpenGLBuffer::StaticDraw);
}

void Animator::allocateVertexBuffers()
{
    auto vertexBufferSize = m_numVertices * 3 * sizeof(Scalar);

    // Initialize color
    std::vector<Vec3> colorData;
    colorData.reserve(m_numVertices);
    for (Idx iVertex = 0; iVertex < m_numVertices; iVertex++) {
        colorData.push_back(MeshColor);
    }
    m_colorVbo.bind();
    m_colorVbo.allocate(colorData.data(), vertexBufferSize);
    m_colorVbo.release();

    // Copy triangles data
    auto elementBufferSize = m_triangles.size() * sizeof(Idx);
    m_elemVbo.bind();
    m_elemVbo.allocate(m_triangles.data(), elementBufferSize);
    m_elemVbo.release();

    // Set first frame
    setFrame(0);

    bindVao();
}

void Animator::destroyVertexBuffers()
{
    m_vertexVbo.destroy();
    m_normalVbo.destroy();
    m_colorVbo.destroy();
    m_elemVbo.destroy();
}

void Animator::bindVao()
{
    // Bind VAO layout
    m_vao.bind();
    m_program.bind();

    m_program.enableAttributeArray(m_posAttrLoc);
    m_vertexVbo.bind();
    m_program.setAttributeBuffer(m_posAttrLoc, ScalarTypeEnum, 0, 3);
    m_vertexVbo.release();

    m_program.enableAttributeArray(m_normalAttrLoc);
    m_normalVbo.bind();
    m_program.setAttributeBuffer(m_normalAttrLoc, ScalarTypeEnum, 0, 3);
    m_normalVbo.release();

    m_program.enableAttributeArray(m_colorAttrLoc);
    m_colorVbo.bind();
    m_program.setAttributeBuffer(m_colorAttrLoc, ScalarTypeEnum, 0, 3);
    m_colorVbo.release();

    m_elemVbo.bind();

    m_program.release();
    m_vao.release();
}

void Animator::generateAnimation(
    const std::vector<std::shared_ptr<Animator::Mesh>> &poses)
{
    const auto &basePose = *poses[0];
    // Copy triangles structure
    m_triangles.reserve(basePose.n_faces() * 3);
    for (auto faceIt = basePose.faces_begin(); faceIt != basePose.faces_end();
         ++faceIt)
    {
        for (auto fvIt = basePose.cfv_iter(*faceIt); fvIt.is_valid(); ++fvIt) {
            m_triangles.push_back(fvIt->idx());
        }
    }
    m_numVertices = basePose.n_vertices();
    SurfaceMorph surfaceMorph = createSurfaceMorph(poses);
    generateAnimationFrames(basePose, surfaceMorph);

    qInfo() << "Animation generated";
}

Animator::SurfaceMorph Animator::createSurfaceMorph(
    const std::vector<std::shared_ptr<Animator::Mesh>> &poses)
{
    qInfo() << "Creating interpolator...";

    using namespace Eigen;

    if (poses.size() < 2)
    {
        throw std::runtime_error("At least two poses are necessary");
    }

    auto firstPose = poses[0];

    // Check every pose has the same number of vertices, edges and faces
    for (const auto &pose : poses) {
        if (pose->n_vertices() != firstPose->n_vertices() ||
            pose->n_halfedges() != firstPose->n_halfedges() ||
            pose->n_faces() != firstPose->n_faces())
        {
            throw std::runtime_error(
                "Every pose mesh must have the same structure.");
        }
    }

    // Create triangles matrix
    Matrix<Idx, 3, Dynamic> triangles(3, firstPose->n_faces());
    // Could parallelize, worth it?
    for (auto faceIt = firstPose->faces_begin();
         faceIt != firstPose->faces_end(); ++faceIt)
    {
        Idx iVertex = 0;
        for (auto fvIt = firstPose->fv_iter(*faceIt); fvIt.is_valid(); ++fvIt) {
            triangles(iVertex, faceIt->idx()) = static_cast<Idx>(fvIt->idx());
            iVertex++;
        }
    }

    // Create the surface morpher
    SurfaceMorph surfaceMorph(triangles);

    // Add poses
    for (const auto &pose : poses) {
        Matrix<Scalar, 3, Dynamic> poseCoords(3, pose->n_vertices());
        for (auto vertexIt = pose->vertices_begin();
             vertexIt != pose->vertices_end(); ++vertexIt)
        {
            const auto &p = pose->point(*vertexIt);
            poseCoords.col(vertexIt->idx()) << p[0], p[1], p[2];
        }
        surfaceMorph.addPose(poseCoords);
    }

    return surfaceMorph;
}

void Animator::generateAnimationFrames(
    const Animator::Mesh &basePose, const Animator::SurfaceMorph &surfaceMorph)
{
    using namespace Eigen;

    int numFrames = static_cast<int>(m_duration * m_fps);
    if (numFrames < 0)
    {
        throw std::runtime_error(
            "The duration and frame rate must be positive values");
    }

    // Interpolate poses
    m_framePoses.resize(numFrames);
    auto t = Matrix<Scalar, Dynamic, 1>::LinSpaced(
        numFrames, Scalar(0), Scalar(surfaceMorph.numPoses() - 1));
    int processedFrames = 0;
    int tenPercents = -1;
#pragma omp parallel
    {
        Mesh mutablePose{basePose};
#pragma omp for
        for (int iFrame = 0; iFrame < numFrames; iFrame++) {
            m_framePoses[iFrame] =
                generateFramePose(mutablePose, surfaceMorph, t[iFrame]);
#pragma omp critical
            {
                processedFrames++;
                int currentTen = (processedFrames * 10) / numFrames;
                if (currentTen > tenPercents)
                {
                    tenPercents = currentTen;
                    qInfo("Interpolating frames (%d%%)...", tenPercents * 10);
                }
            }
        }
    }
}

Animator::FramePose Animator::generateFramePose(
    Mesh &mutablePose, const SurfaceMorph &surfaceMorph, Scalar t)
{
    using namespace Eigen;

    // Compute interpolation
    Matrix<Scalar, 3, Dynamic> framePose = surfaceMorph.interpolatePoseAt(t);

    // Copy to mesh to compute normals
    for (auto vIt = mutablePose.vertices_begin();
         vIt != mutablePose.vertices_end(); ++vIt)
    {
        auto vIdx = vIt->idx();
        mutablePose.set_point(*vIt, Mesh::Point(framePose(0, vIdx),
                                                framePose(1, vIdx),
                                                framePose(2, vIdx)));
    }
    mutablePose.update_normals();

    // Save frame
    auto numVertices = mutablePose.n_vertices();
    std::vector<Scalar> coords;
    coords.reserve(numVertices * 3);
    std::vector<Scalar> normals;
    normals.reserve(numVertices * 3);
    for (auto vIt = mutablePose.vertices_begin();
         vIt != mutablePose.vertices_end(); ++vIt)
    {
        auto p = mutablePose.point(*vIt);
        auto n = mutablePose.normal(*vIt);
        coords.push_back(p[0]);
        coords.push_back(p[1]);
        coords.push_back(p[2]);
        normals.push_back(n[0]);
        normals.push_back(n[1]);
        normals.push_back(n[2]);
    }

    return {std::move(coords), std::move(normals)};
}

std::shared_ptr<Animator::Mesh> Animator::readMesh(const QString &meshFile)
{
    std::shared_ptr<Mesh> newMesh = std::make_shared<Mesh>();
    newMesh->request_vertex_normals();
    newMesh->request_face_normals();
    OpenMesh::IO::Options ropt;
    if (!OpenMesh::IO::read_mesh(*newMesh, meshFile.toStdString(), ropt))
    {
        throw std::runtime_error("Could not load mesh file.");
    }
    // If the file did not provide vertex normals, then calculate them
    if (!ropt.check(OpenMesh::IO::Options::VertexNormal))
    {
        newMesh->update_normals();
    }
    return newMesh;
}

void Animator::setFrame(unsigned int iFrame)
{
    auto vertexBufferSize = m_numVertices * 3 * sizeof(Scalar);
    const auto &framePose = m_framePoses[iFrame];

    m_vertexVbo.bind();
    m_vertexVbo.allocate(nullptr, vertexBufferSize);
    auto vertexVboMapped = m_vertexVbo.map(QOpenGLBuffer::WriteOnly);
    std::memcpy(vertexVboMapped, framePose.coords.data(), vertexBufferSize);
    m_vertexVbo.unmap();
    m_vertexVbo.release();

    m_normalVbo.bind();
    m_normalVbo.allocate(nullptr, vertexBufferSize);
    auto normalVboMapped = m_vertexVbo.map(QOpenGLBuffer::WriteOnly);
    std::memcpy(normalVboMapped, framePose.normals.data(), vertexBufferSize);
    m_vertexVbo.unmap();
    m_normalVbo.release();
}

void Animator::setupCamera()
{
    // Find poses bounds
    qglviewer::Vec pMin(std::numeric_limits<qreal>::max(),
                        std::numeric_limits<qreal>::max(),
                        std::numeric_limits<qreal>::max());
    qglviewer::Vec pMax(std::numeric_limits<qreal>::min(),
                        std::numeric_limits<qreal>::min(),
                        std::numeric_limits<qreal>::min());
    for (const auto &pose : m_framePoses) {
        for (Idx iCoord = 0; iCoord < pose.coords.size(); iCoord += 3) {
            pMin[0] = std::min(qreal(pose.coords[iCoord + 0]), pMin[0]);
            pMin[1] = std::min(qreal(pose.coords[iCoord + 1]), pMin[1]);
            pMin[2] = std::min(qreal(pose.coords[iCoord + 2]), pMin[2]);
            pMax[0] = std::max(qreal(pose.coords[iCoord + 0]), pMax[0]);
            pMax[1] = std::max(qreal(pose.coords[iCoord + 1]), pMax[1]);
            pMax[2] = std::max(qreal(pose.coords[iCoord + 2]), pMax[2]);
        }
    }

    // Set camera direction
    camera()->setUpVector(qglviewer::Vec(0, 1, 0));
    camera()->setViewDirection(qglviewer::Vec(0, 0, -1));

    // Set scene bounds
    camera()->setSceneBoundingBox(pMin, pMax);
    camera()->fitBoundingBox(pMin, pMax);
}

void Animator::updateAnimationPeriod()
{
    bool started{animationIsStarted()};

    auto speedFactor = static_cast<Scalar>(std::pow(1.2, m_speed));
    auto newPeriod =
        std::max(Scalar(0),
                 std::min(static_cast<Scalar>(std::numeric_limits<int>::max()),
                          std::round(Scalar(1000) / (m_fps * speedFactor))));
    if (started)
    {
        stopAnimation();
    }
    setAnimationPeriod(static_cast<int>(newPeriod));
    if (started)
    {
        startAnimation();
    }
}

void Animator::animate()
{
    if (m_framePoses.size() < 1)
    {
        return;
    }

    unsigned int animationFrame;
    if (m_backwards && m_framePoses.size() > 1)
    {
        m_currentFrame = m_currentFrame % (2 * m_framePoses.size() - 2);
        if (m_currentFrame < m_framePoses.size())
        {
            animationFrame = m_currentFrame;
        }
        else
        {
            animationFrame = 2 * m_framePoses.size() - m_currentFrame - 2;
        }
    }
    else
    {
        m_currentFrame = m_currentFrame % m_framePoses.size();
        animationFrame = m_currentFrame;
    }
    setFrame(animationFrame);
    bindVao();
    m_currentFrame++;
}

void Animator::preDraw()
{
    QGLViewer::preDraw();

    // Get cameras
    Scalar pm[4][4];
    camera()->getProjectionMatrix(reinterpret_cast<Scalar *>(pm));
    Scalar mvm[4][4];
    camera()->getModelViewMatrix(reinterpret_cast<Scalar *>(mvm));

    // Set cameras
    m_program.bind();
    m_program.setUniformValue(m_projMatAttrLoc, pm);
    m_program.setUniformValue(m_modelViewMatAttrLoc, mvm);
}

void Animator::preDrawStereo(bool leftBuffer)
{
    QGLViewer::preDrawStereo(leftBuffer);

    // Get cameras
    Scalar pm[4][4];
    glMatrixMode(GL_PROJECTION);
    glPushMatrix();
    camera()->loadProjectionMatrixStereo(leftBuffer);
    glGetFloatv(GL_PROJECTION_MATRIX, reinterpret_cast<Scalar *>(pm));
    glPopMatrix();

    Scalar mvm[4][4];
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    camera()->loadModelViewMatrixStereo(leftBuffer);
    glGetFloatv(GL_MODELVIEW_MATRIX, reinterpret_cast<Scalar *>(mvm));
    glPopMatrix();

    // Set cameras
    m_program.bind();
    m_program.setUniformValue(m_projMatAttrLoc, pm);
    m_program.setUniformValue(m_modelViewMatAttrLoc, mvm);
}

void Animator::draw()
{
    QGLViewer::draw();

    if (m_framePoses.size() < 1)
    {
        return;
    }

    m_vao.bind();
    glDrawElements(GL_TRIANGLES, m_triangles.size(), IdxTypeEnum, nullptr);
    m_vao.release();
}

void Animator::postDraw()
{
    m_program.release();

    QGLViewer::postDraw();
}

void Animator::keyPressEvent(QKeyEvent *e)
{
    auto key = e->key();
    if (key == Qt::Key_B)
    {
        // Backwards animation
        m_backwards = !m_backwards;
        if (!m_backwards && m_currentFrame >= m_framePoses.size())
        {
            m_currentFrame = 0;
        }
    }
    if (key == Qt::Key_Plus)
    {
        m_speed++;
        updateAnimationPeriod();
    }
    else if (key == Qt::Key_Minus)
    {
        m_speed--;
        updateAnimationPeriod();
    }
    else
    {
        QGLViewer::keyPressEvent(e);
    }
}

QString Animator::helpString() const
{
    QFile helpFile(":/doc/help.html");
    if (!helpFile.open(QFile::ReadOnly | QFile::Text))
    {
        throw std::logic_error("Error opening internal resource");
    }
    QTextStream helpStream(&helpFile);
    QString helpText = helpStream.readAll();
    helpFile.close();
    return helpText;
}


// Copyright (c) 2016 Jack Grigg
// Copyright (c) 2016 The Zcash developers
// Copyright (c) 2019 BTC.COM
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#pragma once

#include "sodium.h"

#include <cstring>
#include <exception>
#include <functional>
#include <memory>
#include <set>
#include <vector>

#include <boost/static_assert.hpp>

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wsign-compare"
#pragma GCC diagnostic ignored "-Wunused-variable"

namespace equihash_zcash {

typedef crypto_generichash_blake2b_state eh_HashState;
typedef uint32_t eh_index;
typedef uint8_t eh_trunc;

void ExpandArray(const unsigned char* in, size_t in_len,
                 unsigned char* out, size_t out_len,
                 size_t bit_len, size_t byte_pad=0);
void CompressArray(const unsigned char* in, size_t in_len,
                   unsigned char* out, size_t out_len,
                   size_t bit_len, size_t byte_pad=0);

eh_index ArrayToEhIndex(const unsigned char* array);
eh_trunc TruncateIndex(const eh_index i, const unsigned int ilen);

std::vector<eh_index> GetIndicesFromMinimal(std::vector<unsigned char> minimal,
                                            size_t cBitLen);
std::vector<unsigned char> GetMinimalFromIndices(std::vector<eh_index> indices,
                                                 size_t cBitLen);

template<size_t WIDTH>
class StepRow
{
    template<size_t W>
    friend class StepRow;
    friend class CompareSR;

protected:
    unsigned char hash[WIDTH];

public:
    StepRow(const unsigned char* hashIn, size_t hInLen,
            size_t hLen, size_t cBitLen);
    ~StepRow() { }

    template<size_t W>
    StepRow(const StepRow<W>& a);

    bool IsZero(size_t len);
    std::string GetHex(size_t len) { return HexStr(hash, hash+len); }

    template<size_t W>
    friend bool HasCollision(StepRow<W>& a, StepRow<W>& b, int l);
};

class CompareSR
{
private:
    size_t len;

public:
    CompareSR(size_t l) : len {l} { }

    template<size_t W>
    inline bool operator()(const StepRow<W>& a, const StepRow<W>& b) { return memcmp(a.hash, b.hash, len) < 0; }
};

template<size_t WIDTH>
bool HasCollision(StepRow<WIDTH>& a, StepRow<WIDTH>& b, int l);

template<size_t WIDTH>
class FullStepRow : public StepRow<WIDTH>
{
    template<size_t W>
    friend class FullStepRow;

    using StepRow<WIDTH>::hash;

public:
    FullStepRow(const unsigned char* hashIn, size_t hInLen,
                size_t hLen, size_t cBitLen, eh_index i);
    ~FullStepRow() { }

    FullStepRow(const FullStepRow<WIDTH>& a) : StepRow<WIDTH> {a} { }
    template<size_t W>
    FullStepRow(const FullStepRow<W>& a, const FullStepRow<W>& b, size_t len, size_t lenIndices, int trim);
    FullStepRow& operator=(const FullStepRow<WIDTH>& a);

    inline bool IndicesBefore(const FullStepRow<WIDTH>& a, size_t len, size_t lenIndices) const { return memcmp(hash+len, a.hash+len, lenIndices) < 0; }
    std::vector<unsigned char> GetIndices(size_t len, size_t lenIndices,
                                          size_t cBitLen) const;

    template<size_t W>
    friend bool DistinctIndices(const FullStepRow<W>& a, const FullStepRow<W>& b,
                                size_t len, size_t lenIndices);
    template<size_t W>
    friend bool IsValidBranch(const FullStepRow<W>& a, const size_t len, const unsigned int ilen, const eh_trunc t);
};

template<size_t WIDTH>
class TruncatedStepRow : public StepRow<WIDTH>
{
    template<size_t W>
    friend class TruncatedStepRow;

    using StepRow<WIDTH>::hash;

public:
    TruncatedStepRow(const unsigned char* hashIn, size_t hInLen,
                     size_t hLen, size_t cBitLen,
                     eh_index i, unsigned int ilen);
    ~TruncatedStepRow() { }

    TruncatedStepRow(const TruncatedStepRow<WIDTH>& a) : StepRow<WIDTH> {a} { }
    template<size_t W>
    TruncatedStepRow(const TruncatedStepRow<W>& a, const TruncatedStepRow<W>& b, size_t len, size_t lenIndices, int trim);
    TruncatedStepRow& operator=(const TruncatedStepRow<WIDTH>& a);

    inline bool IndicesBefore(const TruncatedStepRow<WIDTH>& a, size_t len, size_t lenIndices) const { return memcmp(hash+len, a.hash+len, lenIndices) < 0; }
    std::shared_ptr<eh_trunc> GetTruncatedIndices(size_t len, size_t lenIndices) const;
};

enum EhSolverCancelCheck
{
    ListGeneration,
    ListSorting,
    ListColliding,
    RoundEnd,
    FinalSorting,
    FinalColliding,
    PartialGeneration,
    PartialSorting,
    PartialSubtreeEnd,
    PartialIndexEnd,
    PartialEnd
};

class EhSolverCancelledException : public std::exception
{
    virtual const char* what() const throw() {
        return "Equihash solver was cancelled";
    }
};

inline constexpr const size_t max(const size_t A, const size_t B) { return A > B ? A : B; }

inline constexpr size_t equihash_solution_size(unsigned int N, unsigned int K) {
    return (1 << K)*(N/(K+1)+1)/8;
}

template<unsigned int N, unsigned int K>
class Equihash
{
private:
    BOOST_STATIC_ASSERT(K < N);
    BOOST_STATIC_ASSERT(N % 8 == 0);
    BOOST_STATIC_ASSERT((N/(K+1)) + 1 < 8*sizeof(eh_index));

public:
    enum : size_t { IndicesPerHashOutput=512/N };
    enum : size_t { HashOutput=IndicesPerHashOutput*N/8 };
    enum : size_t { CollisionBitLength=N/(K+1) };
    enum : size_t { CollisionByteLength=(CollisionBitLength+7)/8 };
    enum : size_t { HashLength=(K+1)*CollisionByteLength };
    enum : size_t { FullWidth=2*CollisionByteLength+sizeof(eh_index)*(1 << (K-1)) };
    enum : size_t { FinalFullWidth=2*CollisionByteLength+sizeof(eh_index)*(1 << (K)) };
    enum : size_t { TruncatedWidth=max(HashLength+sizeof(eh_trunc), 2*CollisionByteLength+sizeof(eh_trunc)*(1 << (K-1))) };
    enum : size_t { FinalTruncatedWidth=max(HashLength+sizeof(eh_trunc), 2*CollisionByteLength+sizeof(eh_trunc)*(1 << (K))) };
    enum : size_t { SolutionWidth=(1 << K)*(CollisionBitLength+1)/8 };

    Equihash() { }

    int InitialiseState(eh_HashState& base_state);
#ifdef ENABLE_MINING
    bool BasicSolve(const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled);
    bool OptimisedSolve(const eh_HashState& base_state,
                        const std::function<bool(std::vector<unsigned char>)> validBlock,
                        const std::function<bool(EhSolverCancelCheck)> cancelled);
#endif
    bool IsValidSolution(const eh_HashState& base_state, std::vector<unsigned char> soln);
};

} // namespace

#include "equihash_zcash.tcc"
#include "equihash_zcash.inl"

namespace equihash_zcash {

static Equihash<96,3> Eh96_3;
static Equihash<200,9> Eh200_9;
static Equihash<96,5> Eh96_5;
static Equihash<48,5> Eh48_5;

#define EhZecInitialiseState(n, k, base_state)  \
    if (n == 96 && k == 3) {                 \
        equihash_zcash::Eh96_3.InitialiseState(base_state);  \
    } else if (n == 200 && k == 9) {         \
        equihash_zcash::Eh200_9.InitialiseState(base_state); \
    } else if (n == 96 && k == 5) {          \
        equihash_zcash::Eh96_5.InitialiseState(base_state);  \
    } else if (n == 48 && k == 5) {          \
        equihash_zcash::Eh48_5.InitialiseState(base_state);  \
    } else {                                 \
        throw std::invalid_argument("Unsupported Equihash parameters"); \
    }

#ifdef ENABLE_MINING
inline bool EhBasicSolve(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled)
{
    if (n == 96 && k == 3) {
        return Eh96_3.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 200 && k == 9) {
        return Eh200_9.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 96 && k == 5) {
        return Eh96_5.BasicSolve(base_state, validBlock, cancelled);
    } else if (n == 48 && k == 5) {
        return Eh48_5.BasicSolve(base_state, validBlock, cancelled);
    } else {
        throw std::invalid_argument("Unsupported Equihash parameters");
    }
}

inline bool EhBasicSolveUncancellable(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock)
{
    return EhBasicSolve(n, k, base_state, validBlock,
                        [](EhSolverCancelCheck pos) { return false; });
}

inline bool EhOptimisedSolve(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock,
                    const std::function<bool(EhSolverCancelCheck)> cancelled)
{
    if (n == 96 && k == 3) {
        return Eh96_3.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 200 && k == 9) {
        return Eh200_9.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 96 && k == 5) {
        return Eh96_5.OptimisedSolve(base_state, validBlock, cancelled);
    } else if (n == 48 && k == 5) {
        return Eh48_5.OptimisedSolve(base_state, validBlock, cancelled);
    } else {
        throw std::invalid_argument("Unsupported Equihash parameters");
    }
}

inline bool EhOptimisedSolveUncancellable(unsigned int n, unsigned int k, const eh_HashState& base_state,
                    const std::function<bool(std::vector<unsigned char>)> validBlock)
{
    return EhOptimisedSolve(n, k, base_state, validBlock,
                            [](EhSolverCancelCheck pos) { return false; });
}
#endif // ENABLE_MINING

#define EhZecIsValidSolution(n, k, base_state, soln, ret)   \
    if (n == 96 && k == 3) {                             \
        ret = equihash_zcash::Eh96_3.IsValidSolution(base_state, soln);  \
    } else if (n == 200 && k == 9) {                     \
        ret = equihash_zcash::Eh200_9.IsValidSolution(base_state, soln); \
    } else if (n == 96 && k == 5) {                      \
        ret = equihash_zcash::Eh96_5.IsValidSolution(base_state, soln);  \
    } else if (n == 48 && k == 5) {                      \
        ret = equihash_zcash::Eh48_5.IsValidSolution(base_state, soln);  \
    } else {                                             \
        throw std::invalid_argument("Unsupported Equihash parameters"); \
    }

} // namespace

#pragma GCC diagnostic pop


/* Copyright 2015 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Loss Op

#include <stdio.h>
#include <cfloat>
#include <math.h>
#include <vector>
#include <ctime>
#include <cstdlib>
#include <Eigen/Core>
#include <Eigen/Geometry>

#include "third_party/eigen3/Eigen/Core"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_shape.h"

#define POSE_CHANNELS 4

using namespace tensorflow;
typedef Eigen::ThreadPoolDevice CPUDevice;

REGISTER_OP("Averagedistance")
    .Attr("T: {float, double}")
    .Attr("margin: float")
    .Input("bottom_prediction: T")
    .Input("bottom_target: T")
    .Input("bottom_weight: T")
    .Input("bottom_point: T")
    .Input("bottom_symmetry: T")
    .Output("loss: T")
    .Output("bottom_diff: T");

REGISTER_OP("AveragedistanceGrad")
    .Attr("T: {float, double}")
    .Attr("margin: float")
    .Input("bottom_diff: T")
    .Input("grad: T")
    .Output("output: T");

template <typename Device, typename T>
class AveragedistanceOp : public OpKernel {
 public:
  explicit AveragedistanceOp(OpKernelConstruction* context) : OpKernel(context)
  {
    // Get the margin
    OP_REQUIRES_OK(context,
                   context->GetAttr("margin", &margin_));
    // Check that margin is positive
    OP_REQUIRES(context, margin_ >= 0,
                errors::InvalidArgument("Need margin >= 0, got ", margin_));
  }

  // bottom_prediction: (batch_size, 4 * num_classes)
  // bottom_point: (num_classes, num_points, 3)
  void Compute(OpKernelContext* context) override
  {
    // Grab the input tensor
    const Tensor& bottom_prediction = context->input(0);
    const T* prediction = bottom_prediction.flat<T>().data();

    const Tensor& bottom_target = context->input(1);
    const T* target = bottom_target.flat<T>().data();

    const Tensor& bottom_weight = context->input(2);
    const T* weight = bottom_weight.flat<T>().data();

    const Tensor& bottom_point = context->input(3);
    const T* point = bottom_point.flat<T>().data();

    const Tensor& bottom_symmetry = context->input(4);
    const T* symmetry = bottom_symmetry.flat<T>().data();

    // data should have 4 dimensions.
    OP_REQUIRES(context, bottom_prediction.dims() == 2,
                errors::InvalidArgument("prediction must be 2-dimensional"));

    OP_REQUIRES(context, bottom_target.dims() == 2,
                errors::InvalidArgument("target must be 2-dimensional"));

    OP_REQUIRES(context, bottom_weight.dims() == 2,
                errors::InvalidArgument("weight must be 2-dimensional"));

    OP_REQUIRES(context, bottom_point.dims() == 3,
                errors::InvalidArgument("point must be 3-dimensional"));

    OP_REQUIRES(context, bottom_symmetry.dims() == 1,
                errors::InvalidArgument("symmetry must be 1-dimensional"));

    // batch size
    int batch_size = bottom_prediction.dim_size(0);
    int num_classes = bottom_point.dim_size(0);
    int num_points = bottom_point.dim_size(1);

    // Create output loss tensor
    int dim = 1;
    TensorShape output_shape;
    TensorShapeUtils::MakeShape(&dim, 1, &output_shape);

    Tensor* top_data_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &top_data_tensor));
    auto top_data = top_data_tensor->template flat<T>();

    // bottom diff
    TensorShape output_shape_diff = bottom_prediction.shape();
    Tensor* bottom_diff_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(1, output_shape_diff, &bottom_diff_tensor));
    T* bottom_diff = bottom_diff_tensor->template flat<T>().data();
    memset(bottom_diff, 0, batch_size * POSE_CHANNELS * num_classes *sizeof(T));

    T loss = 0;
    // for each object
    for (int n = 0; n < batch_size; n++)
    {
      // find the class label and pose of this object
      int index_cls = -1;
      Eigen::Quaternionf pose_gt;
      Eigen::Quaternionf pose_u;
      for (int i = 0; i < POSE_CHANNELS * num_classes; i += POSE_CHANNELS)
      {
        int index = n * POSE_CHANNELS * num_classes + i;
        if (weight[index] > 0)
        {
          index_cls = i / POSE_CHANNELS;

          pose_gt.w() = target[index + 0];
          pose_gt.x() = target[index + 1];
          pose_gt.y() = target[index + 2];
          pose_gt.z() = target[index + 3];

          pose_u.w() = prediction[index + 0];
          pose_u.x() = prediction[index + 1];
          pose_u.y() = prediction[index + 2];
          pose_u.z() = prediction[index + 3];
          break;
        }
      }
      if (index_cls == -1)
        continue;

      // rotation matrix
      Eigen::Matrix3f Rgt = pose_gt.toRotationMatrix();
      Eigen::Matrix3f Ru = pose_u.toRotationMatrix();

      // derivatives of Ru to quaternion
      Eigen::Matrix3f Ru_w;
      Ru_w << 2 * pose_u.w(), -2 * pose_u.z(), 2 * pose_u.y(),
              2 * pose_u.z(), 2 * pose_u.w(), -2 * pose_u.x(),
              -2 * pose_u.y(), 2 * pose_u.x(), 2 * pose_u.w();

      Eigen::Matrix3f Ru_x;
      Ru_x << 2 * pose_u.x(), 2 * pose_u.y(), 2 * pose_u.z(),
              2 * pose_u.y(), -2 * pose_u.x(), -2 * pose_u.w(),
              2 * pose_u.z(), 2 * pose_u.w(), -2 * pose_u.x();

      Eigen::Matrix3f Ru_y;
      Ru_y << -2 * pose_u.y(), 2 * pose_u.x(), 2 * pose_u.w(),
              2 * pose_u.x(), 2 * pose_u.y(), 2 * pose_u.z(),
              -2 * pose_u.w(), 2 * pose_u.z(), -2 * pose_u.y();

      Eigen::Matrix3f Ru_z;
      Ru_z << -2 * pose_u.z(), -2 * pose_u.w(), 2 * pose_u.x(),
              2 * pose_u.w(), -2 * pose_u.z(), 2 * pose_u.y(),
              2 * pose_u.x(), 2 * pose_u.y(), 2 * pose_u.z();

      // for each point
      for (int i = 0; i < num_points; i++)
      {
        int index = index_cls * num_points * 3 + i * 3;
        Eigen::Vector3f x3d(point[index], point[index + 1], point[index + 2]);
        Eigen::Vector3f diff = Ru * x3d - Rgt * x3d;
        T distance = diff.dot(diff);
        if (distance < margin_)
          continue;
        loss += (distance - margin_) / 2.0;

        // compute the gradient from this point
        Eigen::Matrix3f f0;
        f0 << x3d[0], x3d[1], x3d[2],
              0, 0, 0,
              0, 0, 0;

        Eigen::Matrix3f f1;
        f1 << 0, 0, 0,
              x3d[0], x3d[1], x3d[2],
              0, 0, 0;

        Eigen::Matrix3f f2;
        f2 << 0, 0, 0,
              0, 0, 0,
              x3d[0], x3d[1], x3d[2];

        Eigen::Matrix3f f = diff[0] * f0 + diff[1] * f1 + diff[2] * f2;

        int index_diff = n * POSE_CHANNELS * num_classes + POSE_CHANNELS * index_cls;
        bottom_diff[index_diff + 0] += f.cwiseProduct(Ru_w).sum() / (batch_size * num_points);
        bottom_diff[index_diff + 1] += f.cwiseProduct(Ru_x).sum() / (batch_size * num_points);
        bottom_diff[index_diff + 2] += f.cwiseProduct(Ru_y).sum() / (batch_size * num_points);
        bottom_diff[index_diff + 3] += f.cwiseProduct(Ru_z).sum() / (batch_size * num_points);
      }
    }
    loss /= batch_size * num_points;
    top_data(0) = loss;
  }
 private:
  float margin_;
};

REGISTER_KERNEL_BUILDER(Name("Averagedistance").Device(DEVICE_CPU).TypeConstraint<float>("T"), AveragedistanceOp<CPUDevice, float>);
REGISTER_KERNEL_BUILDER(Name("Averagedistance").Device(DEVICE_CPU).TypeConstraint<double>("T"), AveragedistanceOp<CPUDevice, double>);

// GPU implementation for forward pass
bool AveragedistanceForwardLaucher(OpKernelContext* context,
    const float* bottom_prediction, const float* bottom_target, const float* bottom_weight, const float* bottom_point,
    const float* bottom_symmetry, const int batch_size, const int num_classes, const int num_points, const float margin,
    float* top_data, float* bottom_diff, const Eigen::GpuDevice& d);

static void AveragedistanceKernel(
    OpKernelContext* context, const Tensor* bottom_prediction, const Tensor* bottom_target, const Tensor* bottom_weight,
    const Tensor* bottom_point, const Tensor* bottom_symmetry, const int batch_size, const int num_classes, const int num_points, const float margin,
    const TensorShape& tensor_output_shape, const TensorShape& tensor_output_shape_diff)
{
  Tensor* top_data = nullptr;
  Tensor* bottom_diff = nullptr;
  OP_REQUIRES_OK(context, context->allocate_output(0, tensor_output_shape, &top_data));
  OP_REQUIRES_OK(context, context->allocate_output(1, tensor_output_shape_diff, &bottom_diff));

  if (!context->status().ok()) {
    return;
  }

   AveragedistanceForwardLaucher(context,
    bottom_prediction->flat<float>().data(), bottom_target->flat<float>().data(), bottom_weight->flat<float>().data(),
    bottom_point->flat<float>().data(), bottom_symmetry->flat<float>().data(), batch_size, num_classes, num_points, margin,
    top_data->flat<float>().data(), bottom_diff->flat<float>().data(), context->eigen_device<Eigen::GpuDevice>());
}

template <class T>
class AveragedistanceOp<Eigen::GpuDevice, T> : public OpKernel {
 public:
  typedef Eigen::GpuDevice Device;

  explicit AveragedistanceOp(OpKernelConstruction* context) : OpKernel(context)
  {
    // Get the margin
    OP_REQUIRES_OK(context,
                   context->GetAttr("margin", &margin_));
    // Check that margin is positive
    OP_REQUIRES(context, margin_ >= 0,
                errors::InvalidArgument("Need margin >= 0, got ", margin_));
  }

  void Compute(OpKernelContext* context) override
  {
    // Grab the input tensor
    const Tensor& bottom_prediction = context->input(0);
    const Tensor& bottom_target = context->input(1);
    const Tensor& bottom_weight = context->input(2);
    const Tensor& bottom_point = context->input(3);
    const Tensor& bottom_symmetry = context->input(4);

    // data should have 4 dimensions.
    OP_REQUIRES(context, bottom_prediction.dims() == 2,
                errors::InvalidArgument("prediction must be 2-dimensional"));

    OP_REQUIRES(context, bottom_target.dims() == 2,
                errors::InvalidArgument("target must be 2-dimensional"));

    OP_REQUIRES(context, bottom_weight.dims() == 2,
                errors::InvalidArgument("weight must be 2-dimensional"));

    OP_REQUIRES(context, bottom_point.dims() == 3,
                errors::InvalidArgument("point must be 3-dimensional"));

    OP_REQUIRES(context, bottom_symmetry.dims() == 1,
                errors::InvalidArgument("symmetry must be 1-dimensional"));

    // batch size
    int batch_size = bottom_prediction.dim_size(0);
    int num_classes = bottom_point.dim_size(0);
    int num_points = bottom_point.dim_size(1);

//    int num_dimensions = bottom_prediction.shape().dims();
//    std::cout << "DIMS: " << num_dimensions << std::endl;
//    for(int ii_dim=0; ii_dim<num_dimensions; ii_dim++) {
//        std::cout << bottom_prediction.shape().dim_size(ii_dim)) << std::endl;
//    }

    // Create output tensors
    // loss
    int dim = 1;
    TensorShape output_shape;
// SLoss
//    TensorShapeUtils::MakeShape(&dim, 1, &output_shape);
// SLoss modified
    TensorShapeUtils::MakeShape(&batch_size, 1, &output_shape);

    // bottom diff
    TensorShape output_shape_diff = bottom_prediction.shape();

    AveragedistanceKernel(context, &bottom_prediction, &bottom_target, &bottom_weight, &bottom_point, &bottom_symmetry, batch_size, num_classes,
      num_points, margin_, output_shape, output_shape_diff);
  }
 private:
  float margin_;
};

REGISTER_KERNEL_BUILDER(Name("Averagedistance").Device(DEVICE_GPU).TypeConstraint<float>("T"), AveragedistanceOp<Eigen::GpuDevice, float>);

// compute gradient
template <class Device, class T>
class AveragedistanceGradOp : public OpKernel {
 public:
  explicit AveragedistanceGradOp(OpKernelConstruction* context) : OpKernel(context)
  {
    // Get the margin
    OP_REQUIRES_OK(context,
                   context->GetAttr("margin", &margin_));
    // Check that margin is positive
    OP_REQUIRES(context, margin_ >= 0,
                errors::InvalidArgument("Need margin >= 0, got ", margin_));
  }

  void Compute(OpKernelContext* context) override
  {
    const Tensor& bottom_diff = context->input(0);
    auto bottom_diff_flat = bottom_diff.flat<T>();

    const Tensor& out_backprop = context->input(1);
    T loss = out_backprop.flat<T>()(0);

    // data should have 4 dimensions.
    OP_REQUIRES(context, bottom_diff.dims() == 2,
                errors::InvalidArgument("bottom diff must be 2-dimensional"));

    // batch size
    int batch_size = bottom_diff.dim_size(0);
    // number of channels
    int num_channels = bottom_diff.dim_size(1);

    // construct the output shape
    TensorShape output_shape = bottom_diff.shape();
    Tensor* output = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
    auto top_data = output->template flat<T>();

    for (int i = 0; i < batch_size * num_channels; i++)
      top_data(i) = loss * bottom_diff_flat(i);
  }
 private:
  float margin_;
};

REGISTER_KERNEL_BUILDER(Name("AveragedistanceGrad").Device(DEVICE_CPU).TypeConstraint<float>("T"), AveragedistanceGradOp<CPUDevice, float>);

bool AveragedistanceBackwardLaucher(const float* top_diff, const float* bottom_diff, const int batch_size,
    const int channels, float* output, const Eigen::GpuDevice& d);

static void AveragedistanceGradKernel(
    OpKernelContext* context, const Tensor* bottom_diff, const Tensor* out_backprop,
    const int batch_size, const int channels,
    const TensorShape& tensor_output_shape)
{
  Tensor* output = nullptr;
  OP_REQUIRES_OK(context, context->allocate_output(0, tensor_output_shape, &output));

  if (!context->status().ok()) {
    return;
  }

  AveragedistanceBackwardLaucher(
    out_backprop->flat<float>().data(), bottom_diff->flat<float>().data(),
    batch_size, channels, output->flat<float>().data(), context->eigen_device<Eigen::GpuDevice>());
}

template <class T>
class AveragedistanceGradOp<Eigen::GpuDevice, T> : public OpKernel {
 public:
  typedef Eigen::GpuDevice Device;

  explicit AveragedistanceGradOp(OpKernelConstruction* context) : OpKernel(context)
  {
    // Get the margin
    OP_REQUIRES_OK(context,
                   context->GetAttr("margin", &margin_));
    // Check that margin is positive
    OP_REQUIRES(context, margin_ >= 0,
                errors::InvalidArgument("Need margin >= 0, got ", margin_));
  }

  void Compute(OpKernelContext* context) override
  {
    const Tensor& bottom_diff = context->input(0);
    const Tensor& out_backprop = context->input(1);

    // data should have 2 dimensions.
    OP_REQUIRES(context, bottom_diff.dims() == 2,
                errors::InvalidArgument("bottom diff must be 2-dimensional"));

    // batch size
    int batch_size = bottom_diff.dim_size(0);
    // number of channels
    int num_channels = bottom_diff.dim_size(1);

    // construct the output shape
    TensorShape output_shape = bottom_diff.shape();

    // run the kernel
    AveragedistanceGradKernel(
      context, &bottom_diff, &out_backprop, batch_size, num_channels, output_shape);
  }
 private:
  float margin_;
};

REGISTER_KERNEL_BUILDER(Name("AveragedistanceGrad").Device(DEVICE_GPU).TypeConstraint<float>("T"), AveragedistanceGradOp<Eigen::GpuDevice, float>);


#include "Blocks/pch.h"
#include "Blocks/World/World.h"

#include <shared_mutex>
#include <boost/log/trivial.hpp>
#include <FastNoise/FastNoise.h>

#include "BlocksEngine/Core/Actor.h"
#include "BlocksEngine/Core/Dispatch/DispatchQueue.h"
#include "BlocksEngine/Core/Dispatch/DispatchWorkGroup.h"
#include "BlocksEngine/Main/Game.h"

using namespace Blocks;
using namespace BlocksEngine;

World::World(std::weak_ptr<Transform> playerTransform,
             const uint8_t chunkLoadDistance)
    : chunkViewDistance_{chunkLoadDistance},
      playerTransform_{std::move(playerTransform)}
{
}

void World::Start()
{
    SetEventTypes(EventType::Update);
    loadingScreen_ = GetActor()->AddComponent<LoadingScreen>();
    GenerateWorld();

    // Comment this to show a loading screen whilst world is loading.
    //loadingScreen_->LevelLoaded();
}

void World::Update()
{
    const auto transform = playerTransform_.lock();
    if (!transform) return;
    const Chunk::ChunkCoords chunkCoords = ChunkCoordFromPosition(transform->GetPosition());

    if (chunkCoords != lastChunkCoords_)
    {
        UpdateChunks();
        lastChunkCoords_ = chunkCoords;
    }
}

void World::SetPlayerTransform(std::shared_ptr<Transform> playerTransform) noexcept
{
    playerTransform_ = std::move(playerTransform);
}

Vector2<int> World::ChunkCoordFromPosition(
    const Vector3<float>& position) const noexcept
{
    int x = static_cast<int>(std::floor(position.x / Chunk::Width));
    int z = static_cast<int>(std::floor(position.z / Chunk::Depth));
    return {x, z};
}

const Block& World::GetBlock(const Vector3<int> position) const noexcept
{
    const Chunk::ChunkCoords coords = ChunkCoordFromPosition(position);
    const auto chunk = chunks_.find(coords);
    if (chunk == chunks_.end())
    {
        return Block::Air;
    }

    return chunk->second->GetWorldBlock(position);
}

void World::GenerateWorld() noexcept
{
    BOOST_LOG_TRIVIAL(debug) << "Starting World Generator";
    const auto workGroup{std::make_shared<DispatchWorkGroup>()};

    for (int i = 0; i < chunkViewDistance_; i++)
    {
        for (int j = 0; j < chunkViewDistance_; j++)
        {
            const auto chunk = CreateChunk({i - chunkViewDistance_ / 2, j - chunkViewDistance_ / 2});
            activeChunkCoords_.insert(chunk->GetCoords());
            workGroup->AddWorkItem(CreateGenerationRequestForChunk(chunk), DispatchQueue::Background());
        }
    }

    workGroup->AddCallback(GetGame()->MainDispatchQueue(), std::make_shared<DispatchWorkItem>([this]
    {
        OnWorldGenerated();
    }));

    workGroup->Execute();
}

void World::OnWorldGenerated()
{
    BOOST_LOG_TRIVIAL(info) << "World generated";

    std::vector<std::shared_ptr<Chunk>> chunks(chunks_.size());
    std::ranges::transform(chunks_, chunks.begin(), [](auto pair) { return pair.second; });
    const auto meshRequestGroup = CreateMeshRequestGroup(std::move(chunks));

    meshRequestGroup->AddCallback(GetGame()->MainDispatchQueue(), std::make_shared<DispatchWorkItem>([this]
    {
        OnWorldLoaded();
    }));
    meshRequestGroup->Execute();
}

Chunk::ChunkData World::GenerateChunk(const std::shared_ptr<Chunk> chunk) const
{
    auto fnPerlin = FastNoise::New<FastNoise::Perlin>();

    const Chunk::ChunkCoords coords = chunk->GetCoords();
    std::vector<float> noiseOutput(Chunk::Width * Chunk::Depth);
    fnPerlin->GenUniformGrid2D(noiseOutput.data(), coords.x * Chunk::Width, coords.y * Chunk::Depth, Chunk::Width,
                               Chunk::Depth, 0.04f, 48295);
    //fnPerlin->GenUniformGrid3D(noiseOutput.data(), 0, 0, 0, 16, 16, 16, 0.2f, 1337);

    constexpr int center = 25;
    constexpr int delta = 20;

    auto blocks = std::vector<uint8_t>(Chunk::Size);
    for (int i = 0; i < Chunk::Width; i++)
    {
        for (int j = 0; j < Chunk::Height; j++)
        {
            for (int k = 0; k < Chunk::Depth; k++)
            {
                const int targetHeight = center + static_cast<int>(
                    std::round(noiseOutput[k * Chunk::Depth + i] * delta));
                if (targetHeight > 17)
                {
                    blocks[Chunk::GetFlatIndex(i, j, k)] = j < targetHeight ? (j == targetHeight - 1 ? 2 : 1) : 0;
                }
                else
                {
                    blocks[Chunk::GetFlatIndex(i, j, k)] = j < targetHeight ? 3 : 0;
                }
            }
        }
    }
    return blocks;
}

std::shared_ptr<Chunk> World::CreateChunk(Chunk::ChunkCoords coords)
{
    auto name = L"Chunk: " + std::to_wstring(coords.x) + L", " + std::to_wstring(coords.y);
    const std::shared_ptr<Actor> actor = GetGame()->AddActor(std::move(name));
    const std::shared_ptr<Chunk> chunk = actor->AddComponent<Chunk>(*this, coords);
    chunks_[coords] = chunk;

    return chunk;
}

// TODO: Notify will get called before the blocks are assigned which is wrong
std::shared_ptr<DispatchWorkItem> World::CreateGenerationRequestForChunk(
    std::shared_ptr<Chunk> chunk) const
{
    return std::make_shared<DispatchWorkItem>([this, chunk]
    {
        auto blocks = GenerateChunk(chunk);

        BOOST_LOG_TRIVIAL(debug) << "Generating Chunk: " << chunk;

        const auto workItem = std::make_shared<DispatchWorkItem>([chunk, blocks = std::move(blocks)]()
        mutable
            {
                chunk->SetBlocks(std::move(blocks));
                BOOST_LOG_TRIVIAL(debug) << "Blocks assigned for chunk: " << chunk;
            });
        GetGame()->MainDispatchQueue()->Async(workItem);
    });
}

std::shared_ptr<DispatchWorkGroup> World::CreateMeshRequestGroup(
    const std::vector<std::shared_ptr<Chunk>> chunks) const
{
    auto workGroup = std::make_shared<DispatchWorkGroup>();

    for (auto& chunk : chunks)
    {
        workGroup->AddWorkItem(chunk->RegenerateMesh(), DispatchQueue::Background());
    }

    return workGroup;
}

void World::OnWorldLoaded() const
{
    BOOST_LOG_TRIVIAL(info) << "World successfully loaded";
    loadingScreen_->LevelLoaded();
}

void World::UpdateChunks()
{
    const auto workGroup = std::make_shared<DispatchWorkGroup>();
    const Chunk::ChunkCoords playerCoords = ChunkCoordFromPosition(playerTransform_.lock()->GetPosition());

    const std::unordered_set<Chunk::ChunkCoords, ChunkHash> previousActiveChunks = std::move(activeChunkCoords_);
    std::vector<std::shared_ptr<Chunk>> newChunks{};

    for (int i = 0; i < chunkViewDistance_; i++)
    {
        for (int j = 0; j < chunkViewDistance_; j++)
        {
            const Chunk::ChunkCoords chunkCoords = Vector2(
                i - chunkViewDistance_ / 2, j - chunkViewDistance_ / 2) + playerCoords;

            activeChunkCoords_.insert(chunkCoords);

            // Check if chunk already exists. If it doesn't create a new one.
            const auto chunk = chunks_.find(chunkCoords);
            if (chunk == chunks_.end())
            {
                const auto c = CreateChunk(chunkCoords);
                newChunks.push_back(c);
                workGroup->AddWorkItem(CreateGenerationRequestForChunk(c), DispatchQueue::Background());
            }

            // If it does enable the chunk and add it to active chunks
            chunks_[chunkCoords]->Enable();
        }
    }

    for (auto& coords : previousActiveChunks)
    {
        if (!activeChunkCoords_.contains(coords))
        {
            chunks_[coords]->Disable();
        }
    }

    workGroup->AddCallback(GetGame()->MainDispatchQueue(),
                           std::make_shared<DispatchWorkItem>([this, newChunks = move(newChunks)]()
                           mutable
                               {
                                   const auto meshRequestGroup = CreateMeshRequestGroup(std::move(newChunks));
                                   meshRequestGroup->Execute();
                               }));

    workGroup->Execute();
}


//
//  Getplastic.hpp
//  hybrid_fem_bie
//
//  Created by Max on 3/20/18.
//
//

#ifndef Getplastic_hpp
#define Getplastic_hpp

#include <stdio.h>
#include <Eigen/Eigen>

using namespace Eigen;

void Getplastic(double E, double nu , VectorXd &strain, VectorXd &stress , VectorXd &stress_n, VectorXd &strain_n,double sxx_initial,
                double syy_initial, double sxy_initial, double cohes, double blkfric, MatrixXd &de_p);
#endif /* Getplastic_hpp */


#include "MasterCommunication.h"
#include "DCServoCommunicator.h"

#include "DummyTrajectoryGenerator.h"
#include "PathAndMoveBuilder.h"

#include <numeric>
#include <cmath>
#define pi (M_PI)

#include <chrono>
#include <thread>
#include <algorithm>
#include <functional>
#include <mutex>

template <class T>
T interpolate(const T& a, const T& b, float t)
{
    return a * (1 - t) + b * t;
}

template <class T>
class SamplingHandler
{
public:
    SamplingHandler(std::function<T()> inInc, double inputDt) :
            inputIncrementer(std::move(inInc)),
            inputDt(inputDt)
    {
        n = inputIncrementer();
        np1 = inputIncrementer();
    }

    void increment(float dt)
    {
        interpolT += dt;

        while (interpolT > inputDt)
        {
            interpolT -= inputDt;

            n = np1;
            np1 = inputIncrementer();
        }
    }

    T getSample()
    {
        return interpolate(n, np1, interpolT);
    }

private:
    std::function<T()> inputIncrementer;
    double inputDt;
    double interpolT{0};
    T n{};
    T np1{};
};

class Robot
{
public:
    static const size_t dof = 6;

    std::unique_ptr<SimulateCommunication> communicationSim{std::make_unique<SimulateCommunication>()};

    Robot(Communication* communication, const std::array<bool, 7> simulate = {false}, double cycleTime = 0.018) :
            dcServoArray{{{1, simulate[0] ? communicationSim.get() : communication},
                            {2, simulate[1] ? communicationSim.get() : communication},
                            {3, simulate[2] ? communicationSim.get() : communication},
                            {4, simulate[3] ? communicationSim.get() : communication},
                            {5, simulate[4] ? communicationSim.get() : communication},
                            {6, simulate[5] ? communicationSim.get() : communication}}},
            gripperServo{7, simulate[6] ? communicationSim.get() : communication},
            cycleTime(cycleTime)
    {
        dcServoArray[0].setOffsetAndScaling(2 * pi / 4096.0, 0.950301, 0);
        dcServoArray[1].setOffsetAndScaling(2 * pi / 4096.0, -2.042107923, pi / 2);
        dcServoArray[2].setOffsetAndScaling(2 * pi / 4096.0, 1.0339, pi / 2);
        dcServoArray[3].setOffsetAndScaling(-2 * pi / 4096.0, -1.47531 + 0.242229, 0.0);
        dcServoArray[4].setOffsetAndScaling(2 * pi / 4096.0, 1.23394 + 0.393326, 0.0);
        dcServoArray[5].setOffsetAndScaling(-2 * pi / 4096.0, -1.45191 + 0.191361, 0.0);

        dcServoArray[0].setControlSpeed(10);
        dcServoArray[0].setBacklashControlSpeed(10, 3.0, 0.00);
        dcServoArray[0].setFrictionCompensation(0);
        dcServoArray[1].setControlSpeed(20);
        dcServoArray[1].setBacklashControlSpeed(10, 3.0, 0.00);
        dcServoArray[1].setFrictionCompensation(0);
        dcServoArray[2].setControlSpeed(10);
        dcServoArray[2].setBacklashControlSpeed(10, 3.0, 0.00);
        dcServoArray[2].setFrictionCompensation(0);
        dcServoArray[3].setControlSpeed(20);
        dcServoArray[3].setBacklashControlSpeed(10, 3.0, 0.0);
        dcServoArray[3].setFrictionCompensation(0);
        dcServoArray[4].setControlSpeed(20);
        dcServoArray[4].setBacklashControlSpeed(10, 3.0, 0.0);
        dcServoArray[4].setFrictionCompensation(0);
        dcServoArray[5].setControlSpeed(20);
        dcServoArray[5].setBacklashControlSpeed(10, 3.0, 0.0);
        dcServoArray[5].setFrictionCompensation(0);

        gripperServo.setOffsetAndScaling(pi / 1900.0, 0.0, 0.0);

        while (std::any_of(std::begin(dcServoArray), std::end(dcServoArray), [](auto& d)
                {
                    return !d.isInitComplete();
                }) || !gripperServo.isInitComplete())
        {
            std::for_each(std::begin(dcServoArray), std::end(dcServoArray), [](auto& d)
                {
                    d.run();
                });

            gripperServo.run();
        }

        std::transform(std::begin(dcServoArray), std::end(dcServoArray), std::begin(currentPosition), [](auto& d)
            {
                return d.getPosition();
            });

        for (size_t i = 0; i != 2; ++i)
        {
            gripperServo.setReference(pi / 2.0, 0.0, 0.0);
            gripperServo.run();
            gripperServo.getPosition();
        }

        t = std::thread{&Robot::run, this};
    }

    virtual ~Robot()
    {
        shutdown();
    }

    void run()
    {
        using namespace std::chrono;
        high_resolution_clock::time_point sleepUntilTimePoint = high_resolution_clock::now();
        high_resolution_clock::duration clockDurationCycleTime(
                duration_cast<high_resolution_clock::duration>(duration<double>(cycleTime)));

        while(!shuttingDown)
        {
            cycleSleepTime = std::chrono::duration<double>(sleepUntilTimePoint - high_resolution_clock::now()).count();
            std::this_thread::sleep_until(sleepUntilTimePoint);
            sleepUntilTimePoint += clockDurationCycleTime;

            std::function<void(double, Robot*)> tempSendHandlerFunction;
            std::function<void(double, Robot*)> tempReadHandlerFunction;
            {
                const std::lock_guard<std::mutex> lock(handlerFunctionMutex);
                tempSendHandlerFunction = sendCommandHandlerFunction;
                tempReadHandlerFunction = readResultHandlerFunction;
            }

            tempSendHandlerFunction(cycleTime, this);

            std::for_each(std::begin(dcServoArray), std::end(dcServoArray), [](auto& d)
                {
                    d.run();
                });

            gripperServo.run();

            std::transform(std::begin(dcServoArray), std::end(dcServoArray), std::begin(currentPosition), [](auto& d)
                {
                    return d.getPosition();
                });

            gripperServo.getPosition();

            tempReadHandlerFunction(cycleTime, this);
        }
    }

    Eigen::Matrix<double, dof, 1> getPosition()
    {
        return currentPosition;
    }

    void setHandlerFunctions(const std::function<void(double, Robot*)>& newSendCommandHandlerFunction,
            const std::function<void(double, Robot*)>& newReadResultHandlerFunction)
    {
        const std::lock_guard<std::mutex> lock(handlerFunctionMutex);

        sendCommandHandlerFunction = newSendCommandHandlerFunction;
        readResultHandlerFunction = newReadResultHandlerFunction;
    }

    void removeHandlerFunctions()
    {
        setHandlerFunctions([](double cycleTime, Robot* robot){}, [](double cycleTime, Robot* robot){});
    }

    void shutdown()
    {
        if (!shuttingDown)
        {
            shuttingDown = true;
            t.join();
        }
    }

    double getCycleSleepTime()
    {
        return cycleSleepTime;
    }

    std::array<DCServoCommunicator, dof> dcServoArray;
    DCServoCommunicator gripperServo;

private:
    double cycleTime;
    double cycleSleepTime{0};
    Eigen::Matrix<double, dof, 1> currentPosition{Eigen::Matrix<double, dof, 1>::Zero()};

    bool shuttingDown{false};

    std::thread t{};

    std::mutex handlerFunctionMutex{};
    std::function<void(double, Robot*)> sendCommandHandlerFunction{[](double cycleTime, Robot* robot){}};
    std::function<void(double, Robot*)> readResultHandlerFunction{[](double cycleTime, Robot* robot){}};
};

void playPath(Robot& robot,
        PathAndMoveBuilder& pathBuilder,
        const double playbackSpeed = 1.0,
        const std::array<bool, 6> activeMove = {true, true, true, true, true, true}, std::ostream& outStream = std::cout)
{
    if (playbackSpeed > 1.0)
    {
        throw -1;
    }

    double dt{0.001};

    RobotParameters::DynamicRobotDynamics dynamics{dt};

    DummyTrajectoryGenerator trajGen{&dynamics, 0.012 * 2};

    auto startPos = robot.getPosition();
    trajGen.setStart(startPos, 0.1);
    pathBuilder.renderTo(trajGen, startPos);

    trajGen.calculateTrajectory();

    bool doneRunning = false;
    bool reachedEndOfTrajectory = false;

    auto iter = std::cbegin(trajGen);
    auto endIter = std::cend(trajGen);
    auto outK = *iter;
    ++iter;
    auto outKp1 = *iter;
    auto pwm = outK.u;
    double playbackSpeedT = 0;
    SamplingHandler<TrajectoryItem<6, double> > sampler([&]()
            {
                auto outJ = interpolate(outK, outKp1, playbackSpeedT);
                playbackSpeedT += playbackSpeed;
                if (playbackSpeedT >= 1.0)
                {
                    playbackSpeedT -= 1.0;
                    outK = outKp1;
                    if (iter == endIter)
                    {
                        reachedEndOfTrajectory = true;
                    }
                    ++iter;
                    outKp1 = *iter;
                }
                auto outJp1 = interpolate(outK, outKp1, playbackSpeedT);
                outJ.v *= playbackSpeed;
                outJp1.v *= playbackSpeed;
                dynamics.recalculateFreedForward(outJ, outJp1);

                outJ.u[3] = 0;
                outJ.u[4] = 0;
                outJ.u[5] = 0;

                pwm = dynamics.recalcPwm(outJ.u, outJ.v);
                return outJ;
            }, dt);

    auto sendCommandHandlerFunction = [&sampler, &activeMove](double dt, Robot* robot)
        {
            auto& servos = robot->dcServoArray;

            auto trajItem = sampler.getSample();
            sampler.increment(dt);

            for (size_t i = 0; i != servos.size(); ++i)
            {
                if (activeMove[i])
                {
                    servos[i].setReference(trajItem.p[i], trajItem.v[i], trajItem.u[i]);
                }
            }
        };

    double t = 0;
    EigenVectord6 lastTempC;
    auto readResultHandlerFunction = [&t, &lastTempC, &doneRunning, &reachedEndOfTrajectory, &pwm, &outStream](double dt, Robot* robot)
        {
            auto& servos = robot->dcServoArray;

            bool communicationError = std::any_of(std::begin(servos), std::end(servos), [](auto& d)
                    {
                        return !d.isCommunicationOk();
                    });

            EigenVectord6 posJ;
            std::transform(std::cbegin(servos), std::cend(servos), std::begin(posJ),
                    [](const auto& c){return c.getPosition();});

            EigenVectord6 velJ;
            std::transform(std::cbegin(servos), std::cend(servos), std::begin(velJ),
                    [](const auto& c){return c.getBacklashCompensation();});

            EigenVectord6 erroJ;
            std::transform(std::cbegin(servos), std::cend(servos), std::begin(erroJ),
                    [](const auto& c){return c.getControlError(true);});

            EigenVectord6 controlSignal;
            std::transform(std::cbegin(servos), std::cend(servos), std::begin(controlSignal),
                    [](const auto& c){return c.getControlError(false);});

            JointSpaceCoordinate tempJ{posJ};
            CartesianCoordinate tempC{tempJ};

            if (t == 0)
            {
                lastTempC = tempC.c;
            }
            EigenVectord6 tempCVel = (tempC.c - lastTempC) / dt;
            lastTempC = tempC.c;

            const EigenVectord6& viewPos = tempC.c;
            const EigenVectord6& viewVel = tempCVel;

            std::string printVecName;
            auto printVecFunc = [&printVecName, &outStream](int i, const auto& v)
                    {
                        outStream << " " << printVecName << i << ":" << v;
                        return ++i;
                    };

            outStream << "t:" << t;
            printVecName = "p";
            std::accumulate(std::cbegin(viewPos), std::cend(viewPos), 0, printVecFunc);
            printVecName = "pj";
            std::accumulate(std::cbegin(posJ), std::cend(posJ), 0, printVecFunc);
            printVecName = "v";
            std::accumulate(std::cbegin(viewVel), std::cend(viewVel), 0, printVecFunc);
            printVecName = "vj";
            std::accumulate(std::cbegin(velJ), std::cend(velJ), 0, printVecFunc);
            printVecName = "e";
            std::accumulate(std::cbegin(erroJ), std::cend(erroJ), 0, printVecFunc);
            printVecName = "u";
            std::accumulate(std::cbegin(controlSignal), std::cend(controlSignal), 0, printVecFunc);
            outStream << "\n";

            t += dt;

            if (reachedEndOfTrajectory || communicationError)
            {
                robot->removeHandlerFunctions();
                doneRunning = true;
            }
        };

    robot.setHandlerFunctions(sendCommandHandlerFunction, readResultHandlerFunction);

    while(!doneRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

void moveGripper(Robot& robot, double pos, double time = 1.0)
{
    pos = asin(2.0 * pos - 1.0);

    bool doneRunning = false;

    double t = 0;
    double startPos = robot.gripperServo.getPosition();
    double amp = pos - startPos;
    auto sendCommandHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            t = std::min(t + dt / time, 1.0);

            double pRef = startPos + amp * (1.0 - cos(pi * t)) / 2.0;
            robotPointer->gripperServo.setReference(pRef, 0.0, 0.0);
        };

    auto readResultHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            if (t >= 1.0)
            {
                robotPointer->removeHandlerFunctions();
                doneRunning = true;
            }
        };

    robot.setHandlerFunctions(sendCommandHandlerFunction, readResultHandlerFunction);

    while(!doneRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

void recordeOpticalEncoderData(Robot& robot, size_t i, float pwm, float time, std::ostream& outStream = std::cout)
{
    bool doneRunning = false;

    double t = 0;

    auto sendCommandHandlerFunction = [&i, &pwm, &t](double dt, Robot* robot)
        {
            auto& servos = robot->dcServoArray;

            if (pwm != 0.0)
            {
                servos[i].setOpenLoopControlSignal(std::min(pwm, static_cast<float>(t * pwm)), true);
            }
        };

    auto readResultHandlerFunction = [&t, &doneRunning, &i, time, &outStream](double dt, Robot* robot)
        {
            t += dt;
            auto& servos = robot->dcServoArray;
            auto opticalEncoderData = servos[i].getOpticalEncoderChannelData();
            outStream << opticalEncoderData.a << ", "
                      << opticalEncoderData.b << ", "
                      << opticalEncoderData.minCostIndex << ", "
                      << opticalEncoderData.minCost << ", "
                      << servos[i].getVelocity() / servos[i].getScaling() << "\n";

            if (true)
            {
                std::cout << static_cast<int>(100 * t / time) << "%\r";
            }

            if (t >= time)
            {
                robot->removeHandlerFunctions();
                doneRunning = true;
            }
        };

    robot.setHandlerFunctions(sendCommandHandlerFunction, readResultHandlerFunction);

    while(!doneRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

void recordeMomentOfInertia(Robot& robot, size_t i, float amp, float freq, std::ostream& outStream = std::cout)
{
    std::for_each(std::begin(robot.dcServoArray), std::end(robot.dcServoArray), [](auto& d)
        {
            d.disableBacklashControl();
        });

    bool doneRunning = false;

    double t = 0;
    float pos = 0;
    float vel = 0;
    float acc = 0;

    auto startPos = robot.getPosition();

    auto sendCommandHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            t += dt;

            const double freqScaling = freq * (2.0 * pi);
            pos = startPos[i] + amp * (1 - cos(t * freqScaling));
            vel = amp * sin(t * freqScaling) * freqScaling;
            acc = amp * cos(t * freqScaling) * freqScaling * freqScaling;

            auto& servos = robotPointer->dcServoArray;

            for (size_t j = 0; j != servos.size(); ++j)
            {
                if (i == j)
                {
                    servos[j].setReference(pos, vel, 0);
                }
                else
                {
                    servos[j].setReference(startPos[j], 0, 0);
                }
            }
        };

    const double runTime = std::ceil(15 * freq) / freq;

    auto readResultHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            auto& servo = robotPointer->dcServoArray[i];

            outStream << "t: " << t << ", "
                        << "p: " << servo.getPosition() << ", "
                        << "v: " << servo.getVelocity() << ", "
                        << "u: " << servo.getControlSignal() << ", "
                        << "acc: " << acc << "\n";

            if (t >= runTime)
            {
                robotPointer->removeHandlerFunctions();
                doneRunning = true;
            }
        };

    robot.setHandlerFunctions(sendCommandHandlerFunction, readResultHandlerFunction);

    while(!doneRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

void recordeSystemIdentData(Robot& robot, size_t i, double pwmAmp, double frictionPwm,
        std::ostream& outStream = std::cout)
{
    std::for_each(std::begin(robot.dcServoArray), std::end(robot.dcServoArray), [](auto& d)
        {
            d.disableBacklashControl();
        });

    bool doneRunning = false;

    auto pwmTestVec = std::vector<double>();

    int nr = 10;
    for (int i = 4; i != nr + 1; ++i)
    {
        pwmTestVec.push_back(i * pwmAmp / nr + frictionPwm);
        pwmTestVec.push_back(frictionPwm);
        pwmTestVec.push_back(-(i * pwmAmp / nr + frictionPwm));
        pwmTestVec.push_back(-frictionPwm);
    }

    double t = 0;
    double pwm = 0;
    auto sendCommandHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            auto& servo = robotPointer->dcServoArray[i];

            if (t < pwmTestVec.size())
            {
                pwm = pwmTestVec.at(static_cast<size_t>(t));
            }
            servo.setOpenLoopControlSignal(pwm, true);
        };

    double runTime = pwmTestVec.size();
    auto readResultHandlerFunction = [&](double dt, Robot* robotPointer)
        {
            t += dt;
            auto& servo = robotPointer->dcServoArray[i];

            outStream << "t: " << t << ", "
                        << "p: " << servo.getPosition(false) / servo.getScaling() << ", "
                        << "c: " << servo.getCurrent() << ", "
                        << "pwm: " << pwm << "\n";

            if (t >= runTime)
            {
                robotPointer->removeHandlerFunctions();
                doneRunning = true;
            }
        };

    robot.setHandlerFunctions(sendCommandHandlerFunction, readResultHandlerFunction);

    while(!doneRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

PathAndMoveBuilder createPath()
{
    PathAndMoveBuilder pathBuilder;

    VelocityLimiter jointVelocityLimiter(0.1, EigenVectord6{1, 1, 1, 1, 1, 1}, 3.0 / 0.5);
    VelocityLimiter cartesianVelocityLimiter(0.025, EigenVectord6{1, 1, 1, 0, 0, 0}, 0.2 / 0.02);
    cartesianVelocityLimiter.add(0.1, EigenVectord6{1, 1, 1, 0, 0, 0});
    cartesianVelocityLimiter.add(0.01, EigenVectord6{0, 0, 0, 1, 1, 1}, 1.0 / 0.1);
    cartesianVelocityLimiter.add(0.4, EigenVectord6{0, 0, 0, 1, 1, 1});
    std::shared_ptr<JointSpaceDeviationLimiter> noDeviationLimiterJoint =
            std::make_shared<JointSpaceDeviationLimiter>(std::numeric_limits<double>::max());
    std::shared_ptr<JointSpaceDeviationLimiter> deviationLimiterJoint =
            std::make_shared<JointSpaceDeviationLimiter>(0.0001);
    std::shared_ptr<CartesianSpaceDeviationLimiter> deviationLimiterCartesian =
            std::make_shared<CartesianSpaceDeviationLimiter>(0.0001);
    deviationLimiterCartesian->add(0.0001, EigenVectord6{0.0, 0.0, 0.0, 1.0, 1.0, 1.0});

    EigenVectord6 xV{0.1, 0.0, 0.0, 0, 0, 0};
    EigenVectord6 yV{0.0, 0.08, 0.0, 0, 0, 0};
    EigenVectord6 zV{0.0, 0.0, 0.08, 0, 0, 0};
    EigenVectord6 raV{0, 0, 0, 0.9, 0.0, 0.0};
    EigenVectord6 rbV{0, 0, 0, 0.0, 0.8, 0.0};

    JointSpaceCoordinate jointSpaceHome{{0.0, pi / 2, pi / 2, 0, 0, 0}};
    CartesianCoordinate tempCartCoord(jointSpaceHome);

    EigenVectord6 crossCenter{0.0, 0.1, -0.1, 0, 0, 0};
    crossCenter += tempCartCoord.c;

    CartesianCoordinate tempCartCoord2(crossCenter);
    JointSpaceCoordinate crossCenterJointSpace{tempCartCoord2};

    pathBuilder.append(JointSpaceLinearPath::create(crossCenterJointSpace,
        jointVelocityLimiter, jointVelocityLimiter, deviationLimiterJoint));

    if (true)
    {
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + xV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + xV + 0.05 * zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - xV + 0.05 * zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - xV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));

        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + yV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + yV + 0.05 * zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - yV + 0.05 * zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - yV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));

        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + zV + 0.05 * yV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - zV + 0.05 * yV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - zV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));

        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + raV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - raV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));

        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter + rbV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter - rbV},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));
        pathBuilder.append(CartesianSpaceLinearPath::create(
                CartesianCoordinate{crossCenter},
                cartesianVelocityLimiter, cartesianVelocityLimiter, deviationLimiterCartesian));

        pathBuilder.append(JointSpaceLinearPath::create(jointSpaceHome,
            jointVelocityLimiter, jointVelocityLimiter, deviationLimiterJoint));
    }

    pathBuilder.append(JointSpaceLinearPath::create({-0.5, 0.9, 0.9, 1.4, 0, 0},
            jointVelocityLimiter, VelocityLimiter{0.1}, deviationLimiterJoint));

    pathBuilder.append(JointSpaceLinearPath::create({0.5, pi / 2 + 0, pi / 2 + 0, 0, 1.4, 0},
            jointVelocityLimiter, VelocityLimiter{0.1}, deviationLimiterJoint));

    pathBuilder.append(JointSpaceLinearPath::create({-0.5, 0.9, 0.9, 0, 0, 1.4},
            jointVelocityLimiter, VelocityLimiter{0.1}, deviationLimiterJoint));

    pathBuilder.append(JointSpaceLinearPath::create({0, pi / 2 + 0, pi / 2 + 0, 0, 0, 0},
            jointVelocityLimiter, VelocityLimiter{0.1}, deviationLimiterJoint));

    return pathBuilder;
}

#include <fstream>
#include <boost/program_options.hpp>
namespace po = boost::program_options;

int main(int argc, char* argv[])
{
    // Declare the supported options.
    po::options_description options("Allowed options");
    options.add_options()
        ("servoNr", po::value<int>(), "servo nr")
        ("amp", po::value<double>(), "amplitude for recMomentOfInertia")
        ("frq", po::value<double>(), "frequency for recMomentOfInertia")
        ("pwmAmp", po::value<int>(), "pwm amplitude for recOpticalEncoder and recSystemIdentData")
        ("fricPwmAmp", po::value<int>(), "pwm amplitude to overcome friction")
        ("recOpticalEncoder", "record optical encoder data of given servo")
        ("recSystemIdentData", "record system ident data of given servo")
        ("recMomentOfInertia", "record moment of inertia data of given servo")
        ("playPath", "play the path defined in createPath()")
        ("output", po::value<std::string>(), "data output file")
        ("simulate", "simulate servos");

    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, options), vm);
    po::notify(vm);

    int numberOfArgs = std::accumulate(std::cbegin(vm), std::cend(vm), 0,
            [](int i, const auto& v){return ++i;});
    if (numberOfArgs == 0)
    {
        std::cout << options;
        return 0;
    }

    int servoNr = -1;
    if (vm.count("servoNr"))
    {
        servoNr = vm["servoNr"].as<int>();
    }

    double amp = -1.0;
    if (vm.count("amp"))
    {
        amp = vm["amp"].as<double>();
    }

    double frq = -1.0;
    if (vm.count("frq"))
    {
        frq = vm["frq"].as<double>();
    }

    int pwmAmp = -1;
    if (vm.count("pwmAmp"))
    {
        pwmAmp = vm["pwmAmp"].as<int>();
    }

    int fricPwmAmp = -1;
    if (vm.count("fricPwmAmp"))
    {
        fricPwmAmp = vm["fricPwmAmp"].as<int>();
    }

    std::unique_ptr<std::ofstream> outFileStream{nullptr};
    std::ostream* outStream = &std::cout;
    if (vm.count("output"))
    {
        std::string fileName = vm["output"].as<std::string>();
        outFileStream = std::make_unique<std::ofstream>(fileName);
        outStream = outFileStream.get();
    }

    std::unique_ptr<Communication> communication{nullptr};
    if (vm.count("simulate"))
    {
        communication = std::make_unique<SimulateCommunication>();
    }
    else
    {
        try
        {
            communication = std::make_unique<SerialCommunication>("/dev/ttyACM0");
        }
        catch (std::exception& e)
        {
            std::cout << "could not connect to robot serial port\n";
            std::cout << e.what();

            return 0;
        }
    }

    double comCycleTime = 0.018;
    std::array<bool, 7> servoSimulated{{false, false, false, false, false, false, true}};
    if (servoNr != -1)
    {
        if (vm.count("recSystemIdentData") == false)
        {
            comCycleTime = 0.004;
        }
        servoSimulated.fill(true);
        servoSimulated[servoNr - 1] = false;
    }
    Robot robot(communication.get(), servoSimulated, comCycleTime);

    if (vm.count("recOpticalEncoder"))
    {
        if (servoNr != -1)
        {
            if (pwmAmp == -1)
            {
                std::cout << "pwmAmp = ?\n";
                std::cin >> pwmAmp;
            }
            recordeOpticalEncoderData(robot, std::max(servoNr - 1, 0), pwmAmp, 200, *outStream);
        }
        else
        {
            std::cout << "no servo selected\n";
        }
    }
    else if (vm.count("recSystemIdentData"))
    {
        if (servoNr != -1)
        {
            if (pwmAmp == -1)
            {
                std::cout << "pwmAmp = ?\n";
                std::cin >> pwmAmp;
            }
            if (fricPwmAmp == -1)
            {
                std::cout << "fricPwmAmp = ?\n";
                std::cin >> fricPwmAmp;
            }
            recordeSystemIdentData(robot, std::max(servoNr - 1, 0), pwmAmp, fricPwmAmp, *outStream);
        }
        else
        {
            std::cout << "no servo selected\n";
        }
    }
    else if (vm.count("recMomentOfInertia"))
    {
        if (servoNr != -1)
        {
            if (amp == -1.0)
            {
                std::cout << "amp = ?\n";
                std::cin >> amp;
            }
            if (frq < 0.0)
            {
                std::cout << "frq = ?\n";
                std::cin >> frq;
            }
            recordeMomentOfInertia(robot, std::max(servoNr - 1, 0), amp, frq, *outStream);
        }
        else
        {
            std::cout << "no servo selected\n";
        }
    }
    else if (vm.count("playPath"))
    {
        try
        {
            PathAndMoveBuilder pathBuilder{createPath()};
            playPath(robot, pathBuilder, 1.0, {true, true, true, true, true, true}, *outStream);
        }
        catch (std::exception& e)
        {
            std::cout << e.what();
        }
    }
    else
    {
        std::cout << options;
    }

    robot.shutdown();

    return 0;
}