#include "Common/Common.h"
#include "Demos/Visualization/MiniGL.h"
#include "Demos/Visualization/Selection.h"
#include "GL/glut.h"
#include "Simulation/TimeManager.h"
#include <Eigen/Dense>
#include "Simulation/SimulationModel.h"
#include "Simulation/TimeStepController.h"
#include <iostream>
#include "Demos/Visualization/Visualization.h"
#include "Simulation/DistanceFieldCollisionDetection.h"
#include "Utils/OBJLoader.h"
#include "Utils/Logger.h"
#include "Utils/Timing.h"
#include "Utils/FileSystem.h"
#include "Demos/Common/DemoBase.h"
#include "Demos/Common/TweakBarParameters.h"
#include "Simulation/Simulation.h"
#include "Utils/pymesh/MshLoader.h"

#include <map>

// Enable memory leak detection
#if defined(_DEBUG) && !defined(EIGEN_ALIGN)
#define new DEBUG_NEW
#endif

using namespace PBD;
using namespace Eigen;
using namespace std;
using namespace Utilities;

void initParameters();
void timeStep();
void buildModel();
void createMesh();
void render();
void reset();
void exportOBJ();
void TW_CALL setStiffness(const void *value, void *clientData);
void TW_CALL getStiffness(void *value, void *clientData);
void TW_CALL setPoissonRatio(const void *value, void *clientData);
void TW_CALL getPoissonRatio(void *value, void *clientData);
void TW_CALL setNormalizeStretch(const void *value, void *clientData);
void TW_CALL getNormalizeStretch(void *value, void *clientData);
void TW_CALL setNormalizeShear(const void *value, void *clientData);
void TW_CALL getNormalizeShear(void *value, void *clientData);
void TW_CALL setSimulationMethod(const void *value, void *clientData);
void TW_CALL getSimulationMethod(void *value, void *clientData);
void TW_CALL setContactTolerance(const void *value, void *clientData);
void TW_CALL getContactTolerance(void *value, void *clientData);
void TW_CALL setContactStiffnessRigidBody(const void *value, void *clientData);
void TW_CALL getContactStiffnessRigidBody(void *value, void *clientData);
void TW_CALL setContactStiffnessParticleRigidBody(const void *value, void *clientData);
void TW_CALL getContactStiffnessParticleRigidBody(void *value, void *clientData);

DemoBase *base;
DistanceFieldCollisionDetection cd;

const unsigned int width = 20;
const unsigned int depth = 3;
const unsigned int height = 3;
short simulationMethod = 2;

bool enableExportOBJ = false;
unsigned int exportFPS = 25;
Real nextFrameTime = 0.0;
unsigned int frameCounter = 1;

double GPLANE = 5.3;
Vector3r eyep(0.0, 30.0, 30.0);

void mcopy(unsigned int node1, unsigned int node2, ParticleData &pd)
{
    Real mass = pd.getMass(node1);
    Real invmass = pd.getInvMass(node1);
    Vector3r x0 = pd.getPosition0(node1);
    Vector3r x = pd.getPosition(node1);
    Vector3r v = pd.getVelocity(node1);
    Vector3r a = pd.getAcceleration(node1);
    Vector3r oldx = pd.getOldPosition(node1);
    Vector3r lastx = pd.getLastPosition(node1);

    pd.setMass(node2, mass);
    pd.setPosition0(node2, x0);
    pd.setPosition(node2, x);
    pd.setVelocity(node2, v);
    pd.setAcceleration(node2, a);
    pd.setOldPosition(node2, oldx);
    pd.setLastPosition(node2, lastx);
}

void loadbranch()
{
    PyMesh::MshLoader branch("/Users/polaris/Desktop/tetwildsample/branch02_.msh");
    //    PyMesh::MshLoader branch("/Users/polaris/Desktop/tetwildsample/branch04_.msh");

    for (auto &v : branch.vertices)
    {
        v *= 10.0;
        v += Vector3r(3.0, 5.0, 0.0);
    }

    SimulationModel *model = Simulation::getCurrent()->getModel();
    model->addTetModel((unsigned int) branch.vertices.size(), (unsigned int) branch.tets.size() / 4u, branch.vertices.data(), branch.tets.data());

    LOG_INFO << "Number of tets: " << branch.tets.size() / 4;
    LOG_INFO << "Number of vertices: " << branch.vertices.size();
}

void seperate()
{
    vector<unsigned int> tet1;
    vector<unsigned int> tet2;
    vector<unsigned int> fracture;

    SimulationModel *model = Simulation::getCurrent()->getModel();
    ParticleData &pd = model->getParticles();

    TetModel::ParticleMesh &pm = model->getTetModels()[0]->getParticleMesh();

    int totaltet = pm.numTets();
    for (int i = 0; i < totaltet; ++i)
    {
        std::vector<unsigned int> &tetinds = pm.getTets();
        unsigned int a = tetinds[i * 4 + 0];
        unsigned int b = tetinds[i * 4 + 1];
        unsigned int c = tetinds[i * 4 + 2];
        unsigned int d = tetinds[i * 4 + 3];
        Vector3r apos = pd.getPosition(a);
        Vector3r bpos = pd.getPosition(b);
        Vector3r cpos = pd.getPosition(c);
        Vector3r dpos = pd.getPosition(d);
        Vector3r center = (1.0 / 4.0) * (apos + bpos + cpos + dpos);

        double plane = GPLANE;
        double d1 = apos.x() - plane;
        double d2 = bpos.x() - plane;
        double d3 = cpos.x() - plane;
        double d4 = dpos.x() - plane;
        if ((d1 * d2 < 0) || (d1 * d3 < 0) || (d1 * d4 < 0) || (d2 * d3 < 0) || (d2 * d4 < 0) || (d3 * d4 < 0))
        {
            unsigned int nVert;
            unsigned int offset;
            unsigned int tail;

            tail = pd.size();
            vector<Vector3r> p1{center, pd.getPosition(b), pd.getPosition(c), pd.getPosition(d)};
            vector<unsigned int> f1{0, 1, 2, 3};
            model->addTetModel(4, (unsigned int) f1.size() / 4u, p1.data(), f1.data());
            mcopy(a, tail, pd);
            mcopy(b, tail + 1, pd);
            mcopy(c, tail + 2, pd);
            mcopy(d, tail + 3, pd);

            //            nVert = model->getTetModels().back()->getParticleMesh().numVertices();
            //            offset = model->getTetModels().back()->getIndexOffset();
            //            model->getTetModels().back()->setFrictionCoeff(static_cast<Real>(0.1));
            //            cd.addCollisionObjectWithoutGeometry(model->getTetModels().size() - 1, CollisionDetection::CollisionObject::TetModelCollisionObjectType,
            //                                                 &pd.getPosition(offset), nVert, true);

            tail = pd.size();
            vector<Vector3r> p2{pd.getPosition(a), center, pd.getPosition(c), pd.getPosition(d)};
            vector<unsigned int> f2{0, 1, 2, 3};
            model->addTetModel(4, (unsigned int) f2.size() / 4u, p2.data(), f2.data());
            mcopy(a, tail, pd);
            mcopy(b, tail + 1, pd);
            mcopy(c, tail + 2, pd);
            mcopy(d, tail + 3, pd);

            //            nVert = model->getTetModels().back()->getParticleMesh().numVertices();
            //            offset = model->getTetModels().back()->getIndexOffset();
            //            model->getTetModels().back()->setFrictionCoeff(static_cast<Real>(0.1));
            //            cd.addCollisionObjectWithoutGeometry(model->getTetModels().size() - 1, CollisionDetection::CollisionObject::TetModelCollisionObjectType,
            //                                                 &pd.getPosition(offset), nVert, true);

            tail = pd.size();
            vector<Vector3r> p3{pd.getPosition(a), pd.getPosition(b), center, pd.getPosition(d)};
            vector<unsigned int> f3{0, 1, 2, 3};
            model->addTetModel(4, (unsigned int) f3.size() / 4u, p3.data(), f3.data());
            mcopy(a, tail, pd);
            mcopy(b, tail + 1, pd);
            mcopy(c, tail + 2, pd);
            mcopy(d, tail + 3, pd);

            //            nVert = model->getTetModels().back()->getParticleMesh().numVertices();
            //            offset = model->getTetModels().back()->getIndexOffset();
            //            model->getTetModels().back()->setFrictionCoeff(static_cast<Real>(0.1));
            //            cd.addCollisionObjectWithoutGeometry(model->getTetModels().size() - 1, CollisionDetection::CollisionObject::TetModelCollisionObjectType,
            //                                                 &pd.getPosition(offset), nVert, true);

            tail = pd.size();
            vector<Vector3r> p4{pd.getPosition(a), pd.getPosition(b), pd.getPosition(c), center};
            vector<unsigned int> f4{0, 1, 2, 3};
            model->addTetModel(4, (unsigned int) f4.size() / 4u, p4.data(), f4.data());
            mcopy(a, tail, pd);
            mcopy(b, tail + 1, pd);
            mcopy(c, tail + 2, pd);
            mcopy(d, tail + 3, pd);

            //            nVert = model->getTetModels().back()->getParticleMesh().numVertices();
            //            offset = model->getTetModels().back()->getIndexOffset();
            //            model->getTetModels().back()->setFrictionCoeff(static_cast<Real>(0.1));
            //            cd.addCollisionObjectWithoutGeometry(model->getTetModels().size() - 1, CollisionDetection::CollisionObject::TetModelCollisionObjectType,
            //                                                 &pd.getPosition(offset), nVert, true);
        } else
        {
            if (d1 < 0)
            {
                tet1.push_back(a);
                tet1.push_back(b);
                tet1.push_back(c);
                tet1.push_back(d);
            } else
            {
                tet2.push_back(a);
                tet2.push_back(b);
                tet2.push_back(c);
                tet2.push_back(d);
            }
        }
    }

    model->addTetModelFromTetModelWithNewTets(Simulation::getCurrent()->getModel()->getTetModels()[0], tet2);
    //    model->addTetModelFromTetModelWithNewTets(Simulation::getCurrent()->getModel()->getTetModels()[0], fracture);

    pm.release();
    pm.changetetind(tet1);

    for (auto &tm : model->getTetModels())
    {
        tm->reflash();
    }

    model->clearConstraint();
    //    cd.cleanup();

    // init constraints
    for (unsigned int cm = 0; cm < model->getTetModels().size(); cm++)
    {
        const unsigned int nTets = model->getTetModels()[cm]->getParticleMesh().numTets();
        const unsigned int *tets = model->getTetModels()[cm]->getParticleMesh().getTets().data();
        const IndexedTetMesh::VertexTets *vTets = model->getTetModels()[cm]->getParticleMesh().getVertexTets().data();
        if (simulationMethod == 1)
        {
            const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
            const unsigned int nEdges = model->getTetModels()[cm]->getParticleMesh().numEdges();
            const IndexedTetMesh::Edge *edges = model->getTetModels()[cm]->getParticleMesh().getEdges().data();
            for (unsigned int i = 0; i < nEdges; i++)
            {
                const unsigned int v1 = edges[i].m_vert[0] + offset;
                const unsigned int v2 = edges[i].m_vert[1] + offset;

                model->addDistanceConstraint(v1, v2);
            }

            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v1 = tets[4 * i];
                const unsigned int v2 = tets[4 * i + 1];
                const unsigned int v3 = tets[4 * i + 2];
                const unsigned int v4 = tets[4 * i + 3];

                model->addVolumeConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 2)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
            for (unsigned int i = 0; i < nTets; i++)
            {
                // @author Xayah
                const unsigned int v1 = tets[4 * i] + offset;
                const unsigned int v2 = tets[4 * i + 1] + offset;
                const unsigned int v3 = tets[4 * i + 2] + offset;
                const unsigned int v4 = tets[4 * i + 3] + offset;

                model->addFEMTetConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 3)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v1 = tets[4 * i];
                const unsigned int v2 = tets[4 * i + 1];
                const unsigned int v3 = tets[4 * i + 2];
                const unsigned int v4 = tets[4 * i + 3];

                model->addStrainTetConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 4)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v[4] = {tets[4 * i], tets[4 * i + 1], tets[4 * i + 2], tets[4 * i + 3]};
                // Important: Divide position correction by the number of clusters
                // which contain the vertex.
                const unsigned int nc[4] = {vTets[v[0]].m_numTets, vTets[v[1]].m_numTets, vTets[v[2]].m_numTets, vTets[v[3]].m_numTets};
                model->addShapeMatchingConstraint(4, v, nc);
            }
        }
        model->getTetModels()[cm]->updateMeshNormals(pd);
    }
}

// main
int main(int argc, char **argv)
{
    REPORT_MEMORY_LEAKS

    base = new DemoBase();
    base->init(argc, argv, "Bar collision demo");

    SimulationModel *model = new SimulationModel();
    model->init();
    Simulation::getCurrent()->setModel(model);

    buildModel();

    initParameters();

    Simulation::getCurrent()->setSimulationMethodChangedCallback([&]()
                                                                 {
                                                                     reset();
                                                                     initParameters();
                                                                     base->getSceneLoader()->readParameterObject(Simulation::getCurrent()->getTimeStep());
                                                                 });

    // OpenGL
    MiniGL::setClientIdleFunc(50, timeStep);
    MiniGL::setKeyFunc(0, 'r', reset);
    MiniGL::setKeyFunc(1, 's', seperate);
    MiniGL::setClientSceneFunc(render);
    MiniGL::setViewport(40.0f, 0.1f, 500.0f, eyep, Vector3r(5.0, 0.0, 0.0));

    TwType enumType2 = TwDefineEnum("SimulationMethodType", NULL, 0);
    TwAddVarCB(MiniGL::getTweakBar(), "SimulationMethod", enumType2, setSimulationMethod, getSimulationMethod, &simulationMethod,
               " label='Simulation method' enum='0 {None}, 1 {Volume constraints}, 2 {FEM based PBD}, 3 {Strain based dynamics (no inversion handling)}, 4 {Shape matching (no inversion handling)}' group=Simulation");
    TwAddVarCB(MiniGL::getTweakBar(), "Stiffness", TW_TYPE_REAL, setStiffness, getStiffness, model,
               " label='Stiffness'  min=0.0 step=0.1 precision=4 group='Simulation' ");
    TwAddVarCB(MiniGL::getTweakBar(), "PoissonRatio", TW_TYPE_REAL, setPoissonRatio, getPoissonRatio, model,
               " label='Poisson ratio XY'  min=0.0 step=0.1 precision=4 group='Simulation' ");
    TwAddVarCB(MiniGL::getTweakBar(), "NormalizeStretch", TW_TYPE_BOOL32, setNormalizeStretch, getNormalizeStretch, model,
               " label='Normalize stretch' group='Strain based dynamics' ");
    TwAddVarCB(MiniGL::getTweakBar(), "NormalizeShear", TW_TYPE_BOOL32, setNormalizeShear, getNormalizeShear, model,
               " label='Normalize shear' group='Strain based dynamics' ");
    TwAddVarCB(MiniGL::getTweakBar(), "ContactTolerance", TW_TYPE_REAL, setContactTolerance, getContactTolerance, &cd,
               " label='Contact tolerance'  min=0.0 step=0.001 precision=3 group=Simulation ");
    TwAddVarCB(MiniGL::getTweakBar(), "ContactStiffnessRigidBody", TW_TYPE_REAL, setContactStiffnessRigidBody, getContactStiffnessRigidBody, model,
               " label='Contact stiffness RB'  min=0.0 step=0.1 precision=2 group=Simulation ");
    TwAddVarCB(MiniGL::getTweakBar(), "ContactStiffnessParticleRigidBody", TW_TYPE_REAL, setContactStiffnessParticleRigidBody, getContactStiffnessParticleRigidBody,
               model, " label='Contact stiffness Particle-RB'  min=0.0 step=0.1 precision=2 group=Simulation ");

    glutMainLoop();

    base->cleanup();

    Utilities::Timing::printAverageTimes();
    Utilities::Timing::printTimeSums();

    delete Simulation::getCurrent();
    delete base;
    delete model;

    return 0;
}

void initParameters()
{
    TwRemoveAllVars(MiniGL::getTweakBar());
    TweakBarParameters::cleanup();

    MiniGL::initTweakBarParameters();

    TwAddVarRW(MiniGL::getTweakBar(), "ExportOBJ", TW_TYPE_BOOL32, &enableExportOBJ, " label='Export OBJ'");
    TwAddVarRW(MiniGL::getTweakBar(), "ExportFPS", TW_TYPE_UINT32, &exportFPS, " label='Export FPS'");

    TweakBarParameters::createParameterGUI();
    TweakBarParameters::createParameterObjectGUI(base);
    TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent());
    TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getModel());
    TweakBarParameters::createParameterObjectGUI(Simulation::getCurrent()->getTimeStep());
}

void reset()
{
    Utilities::Timing::printAverageTimes();
    Utilities::Timing::reset();

    Simulation::getCurrent()->reset();
    base->getSelectedParticles().clear();

    Simulation::getCurrent()->getModel()->cleanup();
    Simulation::getCurrent()->getTimeStep()->getCollisionDetection()->cleanup();

//    nextFrameTime = 0.0;
//    frameCounter = 1;

    buildModel();
}

void timeStep()
{
    const Real pauseAt = base->getValue<Real>(DemoBase::PAUSE_AT);
    if ((pauseAt > 0.0) && (pauseAt < TimeManager::getCurrent()->getTime()))
        base->setValue(DemoBase::PAUSE, true);

    if (base->getValue<bool>(DemoBase::PAUSE))
        return;

    // Simulation code
    SimulationModel *model = Simulation::getCurrent()->getModel();
    const unsigned int numSteps = base->getValue<unsigned int>(DemoBase::NUM_STEPS_PER_RENDER);
    for (unsigned int i = 0; i < numSteps; i++)
    {
        START_TIMING("SimStep");
        Simulation::getCurrent()->getTimeStep()->step(*model);
        STOP_TIMING_AVG;

//        exportOBJ();
    }

    for (unsigned int i = 0; i < model->getTetModels().size(); i++)
        model->getTetModels()[i]->updateMeshNormals(model->getParticles());
}

void loadObj(const std::string &filename, VertexData &vd, IndexedFaceMesh &mesh, const Vector3r &scale)
{
    std::vector<OBJLoader::Vec3f> x;
    std::vector<OBJLoader::Vec3f> normals;
    std::vector<OBJLoader::Vec2f> texCoords;
    std::vector<MeshFaceIndices> faces;
    OBJLoader::Vec3f s = {(float) scale[0], (float) scale[1], (float) scale[2]};
    OBJLoader::loadObj(filename, &x, &faces, &normals, &texCoords, s);

    mesh.release();
    const unsigned int nPoints = (unsigned int) x.size();
    const unsigned int nFaces = (unsigned int) faces.size();
    const unsigned int nTexCoords = (unsigned int) texCoords.size();
    mesh.initMesh(nPoints, nFaces * 2, nFaces);
    vd.reserve(nPoints);
    for (unsigned int i = 0; i < nPoints; i++)
    {
        vd.addVertex(Vector3r(x[i][0], x[i][1], x[i][2]));
    }
    for (unsigned int i = 0; i < nTexCoords; i++)
    {
        mesh.addUV(texCoords[i][0], texCoords[i][1]);
    }
    for (unsigned int i = 0; i < nFaces; i++)
    {
        // Reduce the indices by one
        int posIndices[3];
        int texIndices[3];
        for (int j = 0; j < 3; j++)
        {
            posIndices[j] = faces[i].posIndices[j] - 1;
            if (nTexCoords > 0)
            {
                texIndices[j] = faces[i].texIndices[j] - 1;
                mesh.addUVIndex(texIndices[j]);
            }
        }

        mesh.addFace(&posIndices[0]);
    }
    mesh.buildNeighbors();

    mesh.updateNormals(vd, 0);
    mesh.updateVertexNormals(vd);

    LOG_INFO << "Number of triangles: " << nFaces;
    LOG_INFO << "Number of vertices: " << nPoints;
}

void buildModel()
{
    TimeManager::getCurrent()->setTimeStepSize(static_cast<Real>(0.005));
    SimulationModel *model = Simulation::getCurrent()->getModel();

    createMesh();

    // create static rigid body
    string fileName = FileSystem::normalizePath(base->getDataPath() + "/models/cube.obj");
    IndexedFaceMesh mesh;
    VertexData vd;
    loadObj(fileName, vd, mesh, Vector3r::Ones());

    string fileNameTorus = FileSystem::normalizePath(base->getDataPath() + "/models/torus.obj");
    IndexedFaceMesh meshTorus;
    VertexData vdTorus;
    loadObj(fileNameTorus, vdTorus, meshTorus, Vector3r::Ones());

    SimulationModel::RigidBodyVector &rb = model->getRigidBodies();
    rb.resize(1);

    // floor
    rb[0] = new RigidBody();
    rb[0]->initBody(1.0, Vector3r(0.0, -5.5, 0.0), Quaternionr(1.0, 0.0, 0.0, 0.0), vd, mesh, Vector3r(100.0, 1.0, 100.0));
    rb[0]->setMass(0.0);

//    // torus
//    rb[1] = new RigidBody();
//    rb[1]->initBody(1.0, Vector3r(5.0, -1.5, 0.0), Quaternionr(1.0, 0.0, 0.0, 0.0), vdTorus, meshTorus, Vector3r(3.0, 3.0, 3.0));
//    rb[1]->setMass(0.0);
//    rb[1]->setFrictionCoeff(static_cast<Real>(0.1));

    Simulation::getCurrent()->getTimeStep()->setCollisionDetection(*model, &cd);
    cd.setTolerance(static_cast<Real>(0.05));

    const std::vector<Vector3r> *vertices1 = rb[0]->getGeometry().getVertexDataLocal().getVertices();
    const unsigned int nVert1 = static_cast<unsigned int>(vertices1->size());
    cd.addCollisionBox(0, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices1)[0], nVert1, Vector3r(100.0, 1.0, 100.0));

//    const std::vector<Vector3r> *vertices2 = rb[1]->getGeometry().getVertexDataLocal().getVertices();
//    const unsigned int nVert2 = static_cast<unsigned int>(vertices2->size());
//    cd.addCollisionTorus(1, CollisionDetection::CollisionObject::RigidBodyCollisionObjectType, &(*vertices2)[0], nVert2, Vector2r(3.0, 1.5));

    SimulationModel::TetModelVector &tm = model->getTetModels();
    ParticleData &pd = model->getParticles();
    for (unsigned int i = 0; i < tm.size(); i++)
    {
        const unsigned int nVert = tm[i]->getParticleMesh().numVertices();
        unsigned int offset = tm[i]->getIndexOffset();
        tm[i]->setFrictionCoeff(static_cast<Real>(0.1));
        cd.addCollisionObjectWithoutGeometry(i, CollisionDetection::CollisionObject::TetModelCollisionObjectType, &pd.getPosition(offset), nVert, true);
    }
    //    seperate();
}

void render()
{
    base->render();
}

void createMesh()
{
    Vector3r points[width * height * depth];
    for (unsigned int i = 0; i < width; i++)
    {
        for (unsigned int j = 0; j < height; j++)
        {
            for (unsigned int k = 0; k < depth; k++)
            {
                points[i * height * depth + j * depth + k] = 0.3 * Vector3r((Real) i, (Real) j + static_cast<Real>(5.0), (Real) k);
            }
        }
    }

    vector<unsigned int> indices;
    for (unsigned int i = 0; i < width - 1; i++)
    {
        for (unsigned int j = 0; j < height - 1; j++)
        {
            for (unsigned int k = 0; k < depth - 1; k++)
            {
                // For each block, the 8 corners are numerated as:
                //     4*-----*7
                //     /|    /|
                //    / |   / |
                //  5*-----*6 |
                //   | 0*--|--*3
                //   | /   | /
                //   |/    |/
                //  1*-----*2
                unsigned int p0 = i * height * depth + j * depth + k;
                unsigned int p1 = p0 + 1;
                unsigned int p3 = (i + 1) * height * depth + j * depth + k;
                unsigned int p2 = p3 + 1;
                unsigned int p7 = (i + 1) * height * depth + (j + 1) * depth + k;
                unsigned int p6 = p7 + 1;
                unsigned int p4 = i * height * depth + (j + 1) * depth + k;
                unsigned int p5 = p4 + 1;

                // Ensure that neighboring tetras are sharing faces
                if ((i + j + k) % 2 == 1)
                {
                    indices.push_back(p2);
                    indices.push_back(p1);
                    indices.push_back(p6);
                    indices.push_back(p3);
                    indices.push_back(p6);
                    indices.push_back(p3);
                    indices.push_back(p4);
                    indices.push_back(p7);
                    indices.push_back(p4);
                    indices.push_back(p1);
                    indices.push_back(p6);
                    indices.push_back(p5);
                    indices.push_back(p3);
                    indices.push_back(p1);
                    indices.push_back(p4);
                    indices.push_back(p0);
                    indices.push_back(p6);
                    indices.push_back(p1);
                    indices.push_back(p4);
                    indices.push_back(p3);
                } else
                {
                    indices.push_back(p0);
                    indices.push_back(p2);
                    indices.push_back(p5);
                    indices.push_back(p1);
                    indices.push_back(p7);
                    indices.push_back(p2);
                    indices.push_back(p0);
                    indices.push_back(p3);
                    indices.push_back(p5);
                    indices.push_back(p2);
                    indices.push_back(p7);
                    indices.push_back(p6);
                    indices.push_back(p7);
                    indices.push_back(p0);
                    indices.push_back(p5);
                    indices.push_back(p4);
                    indices.push_back(p0);
                    indices.push_back(p2);
                    indices.push_back(p7);
                    indices.push_back(p5);
                }
            }
        }
    }
    SimulationModel *model = Simulation::getCurrent()->getModel();
    //    model->addTetModel(width * height * depth, (unsigned int) indices.size() / 4u, points, indices.data());
    //
    //    vector<Vector3r> p1{Vector3r(5, 5, 5), Vector3r(5, 5, 4), Vector3r(5, 4, 5), Vector3r(4, 5, 5)};
    //    vector<unsigned int> f1{0, 1, 2, 3};
    //    model->addTetModel(4, (unsigned int) f1.size() / 4u, p1.data(), f1.data());
    //
    //    vector<Vector3r> p2{Vector3r(7, 7, 7), Vector3r(7, 7, 6), Vector3r(7, 6, 7), Vector3r(6, 7, 7)};
    //    vector<unsigned int> f2{0, 1, 2, 3};
    //    model->addTetModel(4, (unsigned int) f2.size() / 4u, p2.data(), f2.data());

    loadbranch();

    ParticleData &pd = model->getParticles();
    for (unsigned int i = 0; i < pd.getNumberOfParticles(); i++)
    {
        pd.setMass(i, 1.0);
        if (pd.getPosition0(i).x() < GPLANE - 0.3)
            pd.setMass(i, 0.0);
    }

    // init constraints
    for (unsigned int cm = 0; cm < model->getTetModels().size(); cm++)
    {
        const unsigned int nTets = model->getTetModels()[cm]->getParticleMesh().numTets();
        const unsigned int *tets = model->getTetModels()[cm]->getParticleMesh().getTets().data();
        const IndexedTetMesh::VertexTets *vTets = model->getTetModels()[cm]->getParticleMesh().getVertexTets().data();
        if (simulationMethod == 1)
        {
            const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
            const unsigned int nEdges = model->getTetModels()[cm]->getParticleMesh().numEdges();
            const IndexedTetMesh::Edge *edges = model->getTetModels()[cm]->getParticleMesh().getEdges().data();
            for (unsigned int i = 0; i < nEdges; i++)
            {
                const unsigned int v1 = edges[i].m_vert[0] + offset;
                const unsigned int v2 = edges[i].m_vert[1] + offset;

                model->addDistanceConstraint(v1, v2);
            }

            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v1 = tets[4 * i];
                const unsigned int v2 = tets[4 * i + 1];
                const unsigned int v3 = tets[4 * i + 2];
                const unsigned int v4 = tets[4 * i + 3];

                model->addVolumeConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 2)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            const unsigned int offset = model->getTetModels()[cm]->getIndexOffset();
            for (unsigned int i = 0; i < nTets; i++)
            {
                // @author Xayah
                const unsigned int v1 = tets[4 * i] + offset;
                const unsigned int v2 = tets[4 * i + 1] + offset;
                const unsigned int v3 = tets[4 * i + 2] + offset;
                const unsigned int v4 = tets[4 * i + 3] + offset;

                model->addFEMTetConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 3)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v1 = tets[4 * i];
                const unsigned int v2 = tets[4 * i + 1];
                const unsigned int v3 = tets[4 * i + 2];
                const unsigned int v4 = tets[4 * i + 3];

                model->addStrainTetConstraint(v1, v2, v3, v4);
            }
        } else if (simulationMethod == 4)
        {
            TetModel::ParticleMesh &mesh = model->getTetModels()[cm]->getParticleMesh();
            for (unsigned int i = 0; i < nTets; i++)
            {
                const unsigned int v[4] = {tets[4 * i], tets[4 * i + 1], tets[4 * i + 2], tets[4 * i + 3]};
                // Important: Divide position correction by the number of clusters
                // which contain the vertex.
                const unsigned int nc[4] = {vTets[v[0]].m_numTets, vTets[v[1]].m_numTets, vTets[v[2]].m_numTets, vTets[v[3]].m_numTets};
                model->addShapeMatchingConstraint(4, v, nc);
            }
        }
        model->getTetModels()[cm]->updateMeshNormals(pd);
    }

    LOG_INFO << "Number of tets: " << indices.size() / 4;
    LOG_INFO << "Number of vertices: " << width * height * depth;

}

void exportMeshOBJ(const std::string &exportFileName, const unsigned int nVert, const Vector3r *pos, const unsigned int nTri, const unsigned int *faces)
{
    // Open the file
    std::ofstream outfile(exportFileName);
    if (!outfile)
    {
        LOG_WARN << "Cannot open a file to save OBJ mesh.";
        return;
    }

    // Header
    outfile << "# Created by the PositionBasedDynamics library\n";
    outfile << "g default\n";

    // Vertices
    {
        for (auto j = 0u; j < nVert; j++)
        {
            const Vector3r &x = pos[j];
            outfile << "v " << x[0] << " " << x[1] << " " << x[2] << "\n";
        }
    }

    // faces
    {
        for (auto j = 0; j < nTri; j++)
        {
            outfile << "f " << faces[3 * j + 0] + 1 << " " << faces[3 * j + 1] + 1 << " " << faces[3 * j + 2] + 1 << "\n";
        }
    }
    outfile.close();
}

void exportOBJ()
{
    if (!enableExportOBJ)
        return;

    cout << "RENDER A FRAME" << endl;

    if (TimeManager::getCurrent()->getTime() < nextFrameTime)
        return;

    nextFrameTime += 1.0 / (Real) exportFPS;

    //////////////////////////////////////////////////////////////////////////
    // rigid bodies
    //////////////////////////////////////////////////////////////////////////

    //    std::string exportPath = base->getOutputPath() + "/export";
    //    FileSystem::makeDirs(exportPath);

    std::string exportPath = "/Users/polaris/Desktop/objs/branchfracture";
    FileSystem::makeDirs(exportPath);

    SimulationModel *model = Simulation::getCurrent()->getModel();
    const ParticleData &pd = model->getParticles();
    for (unsigned int i = 0; i < model->getTriangleModels().size(); i++)
    {
        const IndexedFaceMesh &mesh = model->getTriangleModels()[i]->getParticleMesh();
        const unsigned int offset = model->getTriangleModels()[i]->getIndexOffset();
        const Vector3r *x = model->getParticles().getVertices()->data();

        std::string fileName = "triangle_model";
        fileName = fileName + std::to_string(i) + "_" + std::to_string(frameCounter) + ".obj";
        std::string exportFileName = FileSystem::normalizePath(exportPath + "/" + fileName);

        exportMeshOBJ(exportFileName, mesh.numVertices(), &x[offset], mesh.numFaces(), mesh.getFaces().data());
    }

    for (unsigned int i = 0; i < model->getTetModels().size(); i++)
    {
        // @author Xayah
        const IndexedFaceMesh &mesh = model->getTetModels()[i]->getSurfaceMesh();
        const unsigned int offset = model->getTetModels()[i]->getIndexOffset();
        const Vector3r *x = model->getTetModels()[i]->getVisVertices().getVertices()->data();

        std::vector<unsigned int> faces = model->getTetModels()[i]->getSurfaceMesh().getFaces();
        map<unsigned int, unsigned int> mymap;

        for (int j = 0; j < faces.size(); ++j)
        {
            ++mymap[faces[j]];
        }

        vector<Vector3r> vertices;
        std::vector<unsigned int> indices;

        int num = 1;
        for (auto it = mymap.begin(); it != mymap.end(); ++it)
        {
            vertices.push_back(pd.getPosition(offset + (it->first)));
            it->second = num++;
        }

        for (int j = 0; j < faces.size(); ++j)
        {
            indices.push_back(mymap[faces[j]]);
        }

        std::string fileName = "tet_model";
        fileName = fileName + std::to_string(i) + "_" + std::to_string(frameCounter) + ".obj";
        std::string exportFileName = FileSystem::normalizePath(exportPath + "/" + fileName);

        exportMeshOBJ(exportFileName, vertices.size(), vertices.data(), indices.size() / 3u, indices.data());
    }

    for (unsigned int i = 0; i < model->getRigidBodies().size(); i++)
    {
        const IndexedFaceMesh &mesh = model->getRigidBodies()[i]->getGeometry().getMesh();
        const Vector3r *x = model->getRigidBodies()[i]->getGeometry().getVertexData().getVertices()->data();

        std::string fileName = "rigid_body";
        fileName = fileName + std::to_string(i) + "_" + std::to_string(frameCounter) + ".obj";
        std::string exportFileName = FileSystem::normalizePath(exportPath + "/" + fileName);

        exportMeshOBJ(exportFileName, mesh.numVertices(), x, mesh.numFaces(), mesh.getFaces().data());
    }

    frameCounter++;
}

void TW_CALL setStiffness(const void *value, void *clientData)
{
    const Real val = *(const Real *) (value);
    ((SimulationModel *) clientData)->setValue<Real>(SimulationModel::SOLID_STIFFNESS, val);
}

void TW_CALL getStiffness(void *value, void *clientData)
{
    *(Real *) (value) = ((SimulationModel *) clientData)->getValue<Real>(SimulationModel::SOLID_STIFFNESS);
}

void TW_CALL setPoissonRatio(const void *value, void *clientData)
{
    const Real val = *(const Real *) (value);
    ((SimulationModel *) clientData)->setValue<Real>(SimulationModel::SOLID_POISSON_RATIO, val);
}

void TW_CALL getPoissonRatio(void *value, void *clientData)
{
    *(Real *) (value) = ((SimulationModel *) clientData)->getValue<Real>(SimulationModel::SOLID_POISSON_RATIO);
}

void TW_CALL setNormalizeStretch(const void *value, void *clientData)
{
    const bool val = *(const bool *) (value);
    ((SimulationModel *) clientData)->setValue<Real>(SimulationModel::SOLID_NORMALIZE_STRETCH, val);
}

void TW_CALL getNormalizeStretch(void *value, void *clientData)
{
    *(bool *) (value) = ((SimulationModel *) clientData)->getValue<Real>(SimulationModel::SOLID_NORMALIZE_STRETCH);
}

void TW_CALL setNormalizeShear(const void *value, void *clientData)
{
    const bool val = *(const bool *) (value);
    ((SimulationModel *) clientData)->setValue<Real>(SimulationModel::SOLID_NORMALIZE_SHEAR, val);
}

void TW_CALL getNormalizeShear(void *value, void *clientData)
{
    *(bool *) (value) = ((SimulationModel *) clientData)->getValue<Real>(SimulationModel::SOLID_NORMALIZE_SHEAR);
}

void TW_CALL setSimulationMethod(const void *value, void *clientData)
{
    const short val = *(const short *) (value);
    *((short *) clientData) = val;
    reset();
}

void TW_CALL getSimulationMethod(void *value, void *clientData)
{
    *(short *) (value) = *((short *) clientData);
}

void TW_CALL setContactStiffnessRigidBody(const void *value, void *clientData)
{
    const Real val = *(const Real *) (value);
    ((SimulationModel *) clientData)->setContactStiffnessRigidBody(val);
}

void TW_CALL getContactStiffnessRigidBody(void *value, void *clientData)
{
    *(Real *) (value) = ((SimulationModel *) clientData)->getContactStiffnessRigidBody();
}

void TW_CALL setContactStiffnessParticleRigidBody(const void *value, void *clientData)
{
    const Real val = *(const Real *) (value);
    ((SimulationModel *) clientData)->setContactStiffnessParticleRigidBody(val);
}

void TW_CALL getContactStiffnessParticleRigidBody(void *value, void *clientData)
{
    *(Real *) (value) = ((SimulationModel *) clientData)->getContactStiffnessParticleRigidBody();
}

void TW_CALL setContactTolerance(const void *value, void *clientData)
{
    const Real val = *(const Real *) (value);
    ((DistanceFieldCollisionDetection *) clientData)->setTolerance(val);
}

void TW_CALL getContactTolerance(void *value, void *clientData)
{
    *(Real *) (value) = ((DistanceFieldCollisionDetection *) clientData)->getTolerance();
}


#ifndef SAMPLER_HPP
#define SAMPLER_HPP
#include <Eigen/Geometry>

namespace InputGen{

template <typename _Scalar,
          template <class> class T,
          class _Primitive>
template <class SampleContainer, class PrimitiveContainer>
void
PrimitiveSampler<_Scalar, T, _Primitive>::generateSamples(
        SampleContainer&    scontainer,
        const PrimitiveContainer& pcontainer){
    typedef typename PrimitiveContainer::value_type::vec vec;
    typedef typename SampleContainer::value_type Sample;

    typename PrimitiveContainer::const_iterator it;

    // Iterate over all the primitives and generate samples using uniform spacing
    for(it = pcontainer.cbegin(); it != pcontainer.cend(); it++){

        unsigned int nbSampleX = int(std::abs((*it).dim()(0)) / spacing);
        unsigned int nbSampleY = int(std::abs((*it).dim()(1)) / spacing);

        const uint& primitiveUID = (*it).uid();

        vec midPoint   = (*it).getMidPoint();
        vec tangentVec = (*it).getTangentVector();
        vec orthoVec   = (*it).normal().cross(tangentVec);

        if (nbSampleX == 0 && nbSampleY == 0) // generate a single sample at the primitive midpoint
            Base::addSample( scontainer,
                             Sample(midPoint, (*it).normal(), primitiveUID),
                             *it);
        else{

            if(nbSampleY == 0){
                Base::addSample( scontainer,
                                 Sample(midPoint, (*it).normal(), primitiveUID),
                                 *it);
            }else{
                for (unsigned int j = 1; j< nbSampleY/2+1; j++){
                    vec offsetY = Scalar(j)*spacing*orthoVec;
                    Base::addSample( scontainer,
                                     Sample(midPoint + offsetY, (*it).normal(), primitiveUID),
                                     *it);
                }
            }
            for (unsigned int i = 1; i< nbSampleX/2+1; i++){
                vec offsetX = Scalar(i)*spacing*tangentVec;
                Base::addSample( scontainer,
                                 Sample(midPoint + offsetX, (*it).normal(), primitiveUID),
                                 *it);
                Base::addSample( scontainer,
                                 Sample(midPoint - offsetX, (*it).normal(), primitiveUID),
                                 *it);

                for (unsigned int j = 1; j< nbSampleY/2+1; j++){
                    vec offsetY = Scalar(j)*spacing*orthoVec;
                    Base::addSample( scontainer,
                                     Sample(midPoint + offsetX + offsetY, (*it).normal(), primitiveUID),
                                     *it);
                    Base::addSample( scontainer,
                                     Sample(midPoint + offsetX - offsetY, (*it).normal(), primitiveUID),
                                     *it);
                    Base::addSample( scontainer,
                                     Sample(midPoint - offsetX + offsetY, (*it).normal(), primitiveUID),
                                     *it);
                    Base::addSample( scontainer,
                                     Sample(midPoint - offsetX - offsetY, (*it).normal(), primitiveUID),
                                     *it);
                }
                // \FIXME There is here a case not covered, when nbSampleX == 0 and nbSampleY != 0
            }
        }
    }

}

template<class Derived1, class Derived2>
inline
typename Derived1::Scalar cross2D(const Derived1 & v1, const Derived2 & v2)
{
    return (v1(0)*v2(1)) - (v1(1)*v2(0));
}

template <typename _Scalar,
          template <class> class T,
          class _Primitive>
template <class SampleContainer, class PrimitiveContainer>
void
PunctualSampler<_Scalar, T, _Primitive>::generateSamples(
              SampleContainer&    scontainer,
        const PrimitiveContainer& pcontainer){
    typedef typename PrimitiveContainer::value_type::vec vec;
    typedef typename SampleContainer::value_type Sample;

    typename PrimitiveContainer::const_iterator it;

    _Scalar angleStep = _Scalar(2.) * _Scalar(M_PI) / _Scalar(nbSamples);

    for (unsigned int i = 0; i!= this->nbSamples; i++){
        _Scalar angle = _Scalar(i) * angleStep;
        //direction of the ray from the sample source (this->pos)
        vec d (std::cos(angle), std::sin(angle), _Scalar(0.));

        _Scalar alphaRef = std::numeric_limits<Scalar>::max();
        uint primitiveUID = 0;

        typename PrimitiveContainer::const_iterator itFound = pcontainer.end();

        // Iterate over all the primitives, compute the intersection with the ray, create a sample on the closest primitive
        for(it = pcontainer.begin(); it != pcontainer.end(); it++){
            // a and b are the end points of the current segment
            const vec a = (*it).coord();
            const vec b = (*it).getEndPoint();

            // we use this technique to compute the intersection
            // http://mathforum.org/library/drmath/view/62814.html

            const _Scalar left  = cross2D(d, a-b);

            if(left != _Scalar(0.)){
                const _Scalar right = cross2D((b-this->pos), a-b);

                // alpha is the distance between the source and the intersection point with the ray
                _Scalar alpha = right / left;

                if (alpha > _Scalar(0.f)){

                    const vec inter = (alpha*d) + this->pos;

                    // now we need to check if we are inside the segment
                    // inter is the position of the intersection, so must be in [0:1] in parametric domain
                    const _Scalar beta  =  (inter - b).norm() / (a-b).norm();
                    if ((inter - b).dot(a-b) >=_Scalar(0.) && beta <= _Scalar(1.)){

                        // now we check for intersection or simply add the segment
                        if (this->occlusion){
                            if (alpha < alphaRef){
                                alphaRef     = alpha;
                                itFound      = it;
                                primitiveUID = (*it).uid();
                            }
                        }else
                            Base::addSample( scontainer,
                                             Sample(inter, alpha*alpha* -d, primitiveUID),
                                             *it);
                    }
                }
            }
        }

        if (this->occlusion && itFound != pcontainer.end()){
            Base::addSample( scontainer,
                             Sample((alphaRef*d) + this->pos, alphaRef*alphaRef*-d, primitiveUID),
                             *itFound);
        }

            // alpha is the position of the intersection in the ray parametric space, so the distance to the source

            //        unsigned int nbSampleX = int(std::abs((*it).dim()(0)) / spacing);
            //        unsigned int nbSampleY = int(std::abs((*it).dim()(1)) / spacing);

            //        const uint& primitiveUID = (*it).uid();

            //        vec midPoint   = (*it).getMidPoint();
            //        vec tangentVec = (*it).getTangentVector();

            //        if (nbSampleX == 0 && nbSampleY == 0) // generate a single sample at the primitive midpoint
            //            scontainer.push_back(Sample(midPoint, (*it).normal(), primitiveUID));
            //        else{

            //            if (nbSampleX != 0 && nbSampleY == 0){
            //                scontainer.push_back(Sample(midPoint, (*it).normal(), primitiveUID));
            //                for (unsigned int i = 1; i< nbSampleX/2+1; i++){
            //                    vec offset = Scalar(i)*spacing*tangentVec;
            //                    scontainer.push_back(Sample(midPoint + offset, (*it).normal(), primitiveUID));
            //                    scontainer.push_back(Sample(midPoint - offset, (*it).normal(), primitiveUID));
            //                }
            //                //if ((*it).dim())
            //            }else if (nbSampleX == 0 && nbSampleY != 0){
            //                std::cerr << "[Sampler] Unsupported configuration" << std::endl;
            //            }else{
            //                std::cerr << "[Sampler] Unsupported configuration" << std::endl;
            //            }
            //        }

    }

}

}

#endif // SAMPLER_HPP


#include <boost/python/def.hpp>
#include <boost/python/args.hpp>
#include <scitbx/slatec/lib_cpp.h>

#if defined(__linux__) && defined(__GNUC__)
#define SCITBX_MATH_BOOST_PYTHON_CMATH_LGAMMA
#endif

namespace scitbx { namespace math {

namespace {

#if defined(SCITBX_MATH_BOOST_PYTHON_CMATH_LGAMMA)
  inline double
  cmath_lgamma(double const& x) { return lgamma(x); }
#endif

} // namespace <anonymous>

namespace boost_python {

  void wrap_slatec()
  {
    using namespace boost::python;
    slatec_clear_error();
    def("slatec_dgamma", slatec::dgamma, (arg("x")));
    def("slatec_dlngam", slatec::dlngam, (arg("x")));
    def("slatec_dlnrel", slatec::dlnrel, (arg("x")));
    def("slatec_dbinom", slatec::dbinom, (arg("n"), arg("m")));
#if defined(SCITBX_MATH_BOOST_PYTHON_CMATH_LGAMMA)
    def("cmath_lgamma", cmath_lgamma, (arg("x")));
#endif
  }

}}} // namespace scitbx::math::boost_python


#ifndef STAN_MATH_PRIM_SCAL_FUN_ASINH_HPP
#define STAN_MATH_PRIM_SCAL_FUN_ASINH_HPP

#include <stan/math/prim/scal/fun/constants.hpp>
#include <stan/math/prim/scal/fun/is_nan.hpp>
#include <stan/math/prim/scal/meta/likely.hpp>
#include <stan/math/prim/scal/fun/boost_policy.hpp>
#include <boost/math/special_functions/asinh.hpp>

namespace stan {
  namespace math {

    /**
     * Return the inverse hyperbolic sine of the specified value.
     * Returns infinity for infinity argument and -infinity for
     * -infinity argument.
     * Returns nan for nan argument.
     *
     * @param[in] x Argument.
     * @return Inverse hyperbolic sine of the argument.
     */
    inline double asinh(double x) {
      if (unlikely(is_nan(x)))
        return x;
      else
        return boost::math::asinh(x, boost_policy_t());
    }

    /**
     * Integer version of asinh.
     *
     * @param[in] x Argument.
     * @return Inverse hyperbolic sine of the argument.
     */
    inline double asinh(int x) {
      return asinh(static_cast<double>(x));
    }

  }
}
#endif


/*
 * protonReconstruct.cpp
 *
 *  Created on: Dec 22, 2011
 *      Author: u051747
 */
#include <iostream>
#include "parameterparser.h"
#include "cuNDArray.h"
#include "hoCuNDArray.h"

#include "cuNDArray.h"
#include "cuNDArray_math.h"
#include "cuNDArray_fileio.h"

#include "protonSubsetOperator.h"
#include "hoNDArray_fileio.h"
#include "check_CUDA.h"

#include "hoCuGPBBSolver.h"
#include "hoCuPartialDerivativeOperator.h"
#include "hoCuNDArray_blas.h"

#include "osSARTSolver.h"
#include "protonDROPSolver.h"
#include "osSPSSolver.h"
#include "hoOSGPBBSolver.h"
#include "hoOSCGSolver.h"
#include "hoOSCGPBBSolver.h"
#include "hoCuBILBSolver.h"
//#include "hoABIBBSolver.h"
#include "hoCuNCGSolver.h"
#include "cuNCGSolver.h"
#include "hdf5_utils.h"
#include "BILBSolver.h"
#include "cuSolverUtils.h"
#include "osLALMSolverD.h"
#include "osLALMSolver.h"
#include "osMOMSolverD.h"
#include "osMOMSolverD3.h"

#include "encodingOperatorContainer.h"
#include "hoCuOperator.h"
#include "hoImageOperator.h"
#include "identityOperator.h"
#include <boost/program_options.hpp>
#include "vector_td_io.h"

#include "hoCuTvOperator.h"
#include "hoCuTvPicsOperator.h"
#include "projectionSpaceOperator.h"

using namespace std;
using namespace Gadgetron;
typedef float _real;

typedef solver<hoCuNDArray<_real>, hoCuNDArray<_real> > baseSolver;
namespace po = boost::program_options;
int main( int argc, char** argv)
{

  //
  // Parse command line
  //

  _real background =  0.00106;
  std::string dataName;
	std::string outputFile;
	vector_td<int,3> dimensions;
	vector_td<float,3> physical_dims;
	vector_td<float,3> origin;
	float beta,gamma;
	float huber;
	int iterations;
	int device;
	int subsets;
	bool use_hull,use_weights, use_non_negativity;
	float tv_weight;
	po::options_description desc("Allowed options");
	desc.add_options()
			("help", "produce help message")
			("data,D", po::value<std::string>(&dataName)->default_value("data.hdf5"), "HDF5 file containing projections and splines")
			("dimensions,d", po::value<vector_td<int,3> >(&dimensions)->default_value(vector_td<int,3>(512,512,1)), "Pixel dimensions of the image")
			("size,S", po::value<vector_td<float,3> >(&physical_dims)->default_value(vector_td<float,3>(20,20,5)), "Dimensions of the image")
			("center,c", po::value<vector_td<float,3> >(&origin)->default_value(vector_td<float,3>(0,0,0)), "Center of the reconstruction")
			("iterations,i", po::value<int>(&iterations)->default_value(10), "Dimensions of the image")
			("output,f", po::value<std::string>(&outputFile)->default_value("image.hdf5"), "Output filename")
			("prior,P", po::value<std::string>(),"Prior image filename")
			("prior-weight,k",po::value<float>(),"Weight of the prior image")
			("device",po::value<int>(&device)->default_value(0),"Number of the device to use (0 indexed)")
			("TV,T",po::value<float>(&tv_weight)->default_value(0),"TV Weight ")
			("subsets,n", po::value<int>(&subsets)->default_value(10), "Number of subsets to use")
			("beta",po::value<float>(&beta)->default_value(1),"Step size for SART")
			("gamma",po::value<float>(&gamma)->default_value(1e-1),"Relaxation Gamma")
			("huber",po::value<float>(&huber)->default_value(0),"Huber value")
			("use_hull",po::value<bool>(&use_hull)->default_value(true),"Estimate convex hull of object")
			("use_weights",po::value<bool>(&use_weights)->default_value(false),"Use weights if available. ")
			("use_non_negativity",po::value<bool>(&use_non_negativity)->default_value(true),"Use non-negativity constraint. ")
			("dump","Will dump all iterations")

	;

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

	if (vm.count("help")) {
		cout << desc << "\n";
		return 1;
	}
	std::cout << "Command line options:" << std::endl;
	for (po::variables_map::iterator it = vm.begin(); it != vm.end(); ++it){
		boost::any a = it->second.value();
		std::cout << it->first << ": ";
		if (a.type() == typeid(std::string)) std::cout << it->second.as<std::string>();
		else if (a.type() == typeid(int)) std::cout << it->second.as<int>();
		else if (a.type() == typeid(float)) std::cout << it->second.as<float>();
		else if (a.type() == typeid(bool)) std::cout << it->second.as<bool>();
		else if (a.type() == typeid(vector_td<float,3>)) std::cout << it->second.as<vector_td<float,3> >();
		else if (a.type() == typeid(vector_td<int,3>)) std::cout << it->second.as<vector_td<int,3> >();
		else std::cout << "Unknown type" << std::endl;
		std::cout << std::endl;
	}
	cudaSetDevice(device);
	//cudaDeviceReset();

/*
  hoCuGPBBSolver< _real> solver;

  solver.set_max_iterations( iterations);

  solver.set_output_mode( hoCuGPBBSolver< _real>::OUTPUT_VERBOSE );*/
	//osSARTSolver<cuNDArray<_real> > solver;
	//osSPSSolver<cuNDArray<_real> > solver;
	//protonDROPSolver<cuNDArray > solver;
	//BILBSolver<cuNDArray<float> > solver;
	//osSPSSolver<cuNDArray<_real> > solver;
	//cuNCGSolver<_real> solver;
	//hoOSCGPBBSolver<hoCuNDArray<_real> > solver;
	//hoABIBBSolver<hoCuNDArray<_real> > solver;
	//hoOSCGSolver<hoCuNDArray<_real> > solver;
	//hoCuBILBSolver<hoCuNDArray<_real> > solver;
	osMOMSolverD<cuNDArray<float>> solver;
	//osLALMSolver<cuNDArray<float>> solver;

	//solver.set_m(24);
	//solver.set_beta(1.9f);
	//solver.set_beta(beta);
	//solver.set_gamma(gamma);

  solver.set_non_negativity_constraint(use_non_negativity);
  solver.set_max_iterations(iterations);
  solver.set_huber(huber);
  solver.set_reg_steps(5);
  solver.set_dump(vm.count("dump") > 0);

  solver.set_tau(1e-4);
  //solver.set_beta(0.1);
  //solver.set_tau(gamma);
  //solver.set_regularization_iterations(1);
  //solver.set_damping(beta);

  //solver.set_tc_tolerance(1e-10f);
  std::vector<size_t> rhs_dims(&dimensions[0],&dimensions[3]); //Quick and dirty vector_td to vector

  boost::shared_ptr<protonDataset<cuNDArray> >  data(new protonDataset<cuNDArray>(dataName,use_weights));

  data = protonDataset<cuNDArray>::shuffle_dataset(data,subsets);

  data->preprocess(rhs_dims,physical_dims,use_hull);

  std::cout << "Size: " << data->get_projections()->get_number_of_elements() << std::endl;

  boost::shared_ptr< protonSubsetOperator<cuNDArray> > E (new protonSubsetOperator<cuNDArray>(data->get_subsets(), physical_dims) );

  E->set_domain_dimensions(&rhs_dims);
  E->set_codomain_dimensions(data->get_projections()->get_dimensions().get());

  if (tv_weight > 0){

  	auto Dx = boost::make_shared<cuPartialDerivativeOperator<float,3>>(0);
  	Dx->set_weight(tv_weight);
  	Dx->set_domain_dimensions(&rhs_dims);
  	Dx->set_codomain_dimensions(&rhs_dims);

  	auto Dy = boost::make_shared<cuPartialDerivativeOperator<float,3>>(1);
  	Dy->set_weight(tv_weight);
  	Dy->set_domain_dimensions(&rhs_dims);
  	Dy->set_codomain_dimensions(&rhs_dims);

  	if (rhs_dims[2] > 1){

  		auto Dz = boost::make_shared<cuPartialDerivativeOperator<float,2>>(2);
  		Dz->set_weight(tv_weight);
  		Dz->set_domain_dimensions(&rhs_dims);
  		Dz->set_codomain_dimensions(&rhs_dims);
  		solver.add_regularization_group({Dx,Dy,Dz});
  	} else {
  		solver.add_regularization_group({Dx,Dy});
  	}
  }

/*
  auto precon = boost::make_shared<cuNDArray<float>>(rhs_dims);
  fill(precon.get(),1.0f);

  solver.set_preconditioning_image(precon);
  solver.set_tau(0.001);
  solver.set_damping(1e-4);
  */
/*
  boost::shared_ptr<hoCuNDArray<_real > > prior;
  if (vm.count("prior")){
 	  std::cout << "Prior image regularization in use" << std::endl;
		prior = boost::static_pointer_cast<hoCuNDArray<_real > >(read_nd_array<_real >(vm["prior"].as<std::string>().c_str()));

		prior->reshape(&rhs_dims);
		_real offset = _real(0.01);
		//cuNDA_add(offset,prior.get());

		if (vm.count("prior-weight")){

		//boost::shared_ptr<hoImageOperator<_real> > I (new hoImageOperator<_real > ());
		//I->compute(prior.get());
			boost::shared_ptr<identityOperator<hoCuNDArray<_real> > > Itmp ( new identityOperator<hoCuNDArray<_real> >);

			boost::shared_ptr<projectionSpaceOperator<hoCuNDArray<_real> > > I (new projectionSpaceOperator<hoCuNDArray<_real> >(Itmp));

		I->set_weight(vm["prior-weight"].as<float>());

		I->set_codomain_dimensions(&rhs_dims);
		I->set_domain_dimensions(&rhs_dims);
		I->set_projections(prior);
		//hoCuNDArray<_real> tmp = *prior;

		//I->mult_M(prior.get(),&tmp);
		solver.add_regularization_operator(I);

		}
		solver.set_x0(prior);
  }
  */
/*
  if (vm.count("TV")){
	  std::cout << "Total variation regularization in use" << std::endl;
	  boost::shared_ptr<cuTvOperator<float,3> > tv(new cuTvOperator<float,3>);
	  //tv->set_weight(vm["TV"].as<float>());
	  //solver.add_nonlinear_operator(tv);
	  solver.set_reg_op(tv);
  }
*/

  solver.set_encoding_operator(E);
  solver.set_output_mode(baseSolver::OUTPUT_VERBOSE);
	//hoCuNDA_clear(projections.get());
	//CHECK_FOR_CUDA_ERROR();

/*
  {
  	cuNDArray<float> precon(rhs_dims);
  	fill(&precon,1.0f);
  	linearOperator<cuNDArray<float>>* E2 = E.get();
  	E2->mult_MH_M(&precon,&precon);
  	clamp_min(&precon,1e-6);
  	reciprocal_inplace(&precon);

  	auto old = boost::make_shared<cuNDArray<float>>(rhs_dims);
  	fill(old.get(),1.0f);

  	auto eigen_vec = boost::make_shared<cuNDArray<float>>(rhs_dims);
  	for (int i = 0; i < 20; i ++){
  		E2->mult_MH_M(old.get(),eigen_vec.get());
  		//*eigen_vec *= precon;

  		std::cout << "Val: " << dot(old.get(),eigen_vec.get()) << std::endl;
  		*eigen_vec /= nrm2(eigen_vec.get());
  		std::swap(eigen_vec,old);

  	}

  }
*/
  //auto precon = boost::make_shared<cuNDArray<float>>(rhs_dims);
  //fill(precon.get(),1.0f);
  //solver.set_preconditioning_image(precon);

	//float res = dot(projections.get(),projections.get());

  if (data->get_weights()) *data->get_projections() *= *data->get_weights();

  //fill(data->get_projections().get(),1.0f);
  boost::shared_ptr<cuNDArray<float> > result;
  {
  	  GPUTimer timer("Reconstruction time");
	result = solver.solve(data->get_projections().get());
  }

	boost::shared_ptr< hoNDArray<float> > host_result = result->to_host();
	//write_nd_array<_real>(result.get(), (char*)parms.get_parameter('f')->get_string_value());
	std::stringstream ss;
	for (int i = 0; i < argc; i++){
		ss << argv[i] << " ";
	}
	saveNDArray2HDF5<3>(host_result.get(),outputFile,physical_dims,origin,ss.str(), solver.get_max_iterations());
}


#include "bytecode_generators.h"

#include "instructions.h"
#include "syntax_descriptors.h"
#include "unit_assert.h"
#include "type_primitive.h"
#include "primitive_cast.h"
#include <boost/lexical_cast.hpp>
#include <boost/algorithm/string.hpp>

using namespace compiler;
using namespace compiler::syntax;
using namespace compiler::instructions;

void discard_if(AST ast, bytecode_appender &appender)
{
	if (ast->discard_result && !ast->type.base_type->is_void())
		appender.append(instruction_ptr(new instruction_pop()));
}

template<typename T>
bool build_assign_common(const string &op, bytecode_appender &appender)
{
	if (op == "*=")
		appender.append(instruction_ptr(new instruction_mul<T>()));
	else if (op == "+=")
		appender.append(instruction_ptr(new instruction_add<T>()));
	else if (op == "-=")
		appender.append(instruction_ptr(new instruction_sub<T>()));
	else if (op == "/=")
		appender.append(instruction_ptr(new instruction_div<T>()));
	else
		return false;
	return true;
}

template<typename T>
bool build_assign_integral(const string &op, bytecode_appender &appender)
{
	if (build_assign_common<T>(op, appender))
		return true;
	else if (op == "&=")
		appender.append(instruction_ptr(new instruction_and<T>()));
	else if (op == "|=")
		appender.append(instruction_ptr(new instruction_or<T>()));
	else if (op == "^=")
		appender.append(instruction_ptr(new instruction_xor<T>()));
	else if (op == "<<=")
		appender.append(instruction_ptr(new instruction_shl<T>()));
	else if (op == ">>=")
		appender.append(instruction_ptr(new instruction_shr<T>()));
	else if (op == "%=")
		appender.append(instruction_ptr(new instruction_rem<T>()));
	else
		return false;
	return true;
}

bool build_assign_bool(const string &op, bytecode_appender &appender)
{
	if (op == "^=")
		appender.append(instruction_ptr(new instruction_xor<bool>()));
	else if (op == "|=")
		appender.append(instruction_ptr(new instruction_or<bool>()));
	else if (op == "&=")
		appender.append(instruction_ptr(new instruction_and<bool>()));
	else
		return false;
	return true;
}

void compiler::bytecode_assign::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_assign>(ast->desc);

	assert_eq(ast->children.size(), 2);

	appender.build_expression(ast->children[0]);
	if (desc.op != "=")
		appender.append(instruction_ptr(new instruction_duplicate()));
	appender.build_expression(ast->children[1]);

	auto type_name = ast->children[0]->type.base_type->id;
	assert_eq(ast->children[0]->type.base_type->id, ast->children[1]->type.base_type->id);

	if (desc.op != "=")
	{
		bool flag = false;

		if (type_name == type_int->id)
			flag = build_assign_integral<int>(desc.op, appender);
		else if (type_name == type_short->id)
			flag = build_assign_integral<short>(desc.op, appender);
		else if (type_name == type_long->id)
			flag = build_assign_integral<long>(desc.op, appender);
		else if (type_name == type_int64->id)
			flag = build_assign_integral<long long>(desc.op, appender);
		else if (type_name == type_float->id)
			flag = build_assign_common<float>(desc.op, appender);
		else if (type_name == type_double->id)
			flag = build_assign_common<double>(desc.op, appender);
		else if (type_name == type_long_double->id)
			flag = build_assign_common<long double>(desc.op, appender);
		else if (type_name == type_bool->id)
			flag = build_assign_bool(desc.op, appender);

		if (!flag)
			throw runtime_error(string_format("Unexpected assign operator %s", desc.op.c_str()));
	}
	appender.append(instruction_ptr(new instruction_store()));
}

template<typename T>
bool build_binary_common(const string &op, bytecode_appender &appender)
{
	if (op == "*")
		appender.append(instruction_ptr(new instruction_mul<T>()));
	else if (op == "+")
		appender.append(instruction_ptr(new instruction_add<T>()));
	else if (op == "-")
		appender.append(instruction_ptr(new instruction_sub<T>()));
	else if (op == "/")
		appender.append(instruction_ptr(new instruction_div<T>()));
	else if (op == "==")
		appender.append(instruction_ptr(new instruction_eq<T>()));
	else if (op == "!=")
		appender.append(instruction_ptr(new instruction_neq<T>()));
	else if (op == ">")
		appender.append(instruction_ptr(new instruction_greater<T>()));
	else if (op == "<")
		appender.append(instruction_ptr(new instruction_less<T>()));
	else if (op == "<=")
		appender.append(instruction_ptr(new instruction_leq<T>()));
	else if (op == ">=")
		appender.append(instruction_ptr(new instruction_geq<T>()));
	else
		return false;
	return true;
}

template<typename T>
bool build_binary_integral(const string &op, bytecode_appender &appender)
{
	if (build_binary_common<T>(op, appender))
		return true;
	else if (op == "<<")
		appender.append(instruction_ptr(new instruction_shl<T>()));
	else if (op == ">>")
		appender.append(instruction_ptr(new instruction_shr<T>()));
	else if (op == "&&")
		appender.append(instruction_ptr(new instruction_logical_and<T>()));
	else if (op == "||")
		appender.append(instruction_ptr(new instruction_logical_or<T>()));
	else if (op == "^")
		appender.append(instruction_ptr(new instruction_xor<T>()));
	else if (op == "|")
		appender.append(instruction_ptr(new instruction_or<T>()));
	else if (op == "&")
		appender.append(instruction_ptr(new instruction_and<T>()));
	else if (op == "%")
		appender.append(instruction_ptr(new instruction_rem<T>()));
	else
		return false;
	return true;
}

bool build_binary_bool(const string &op, bytecode_appender &appender)
{
	if (op == "&&")
		appender.append(instruction_ptr(new instruction_logical_and<bool>()));
	else if (op == "||")
		appender.append(instruction_ptr(new instruction_logical_or<bool>()));
	else if (op == "^")
		appender.append(instruction_ptr(new instruction_xor<bool>()));
	else if (op == "|")
		appender.append(instruction_ptr(new instruction_or<bool>()));
	else if (op == "&")
		appender.append(instruction_ptr(new instruction_and<bool>()));
	else
		return false;
	return true;
}

void compiler::bytecode_binary_operator::build(bytecode_appender &appender, AST ast)
{
	auto desc = any_cast<descriptor_binary_operator>(ast->desc);

	assert_eq(ast->children.size(), 2);

	if (desc.op == ",")
	{
		appender.build_expression(ast->children[0]);
		discard_if(ast->children[0], appender);
		appender.build_expression(ast->children[1]);
	}
	else
	{
		appender.build_expression(ast->children[0]);
		appender.build_expression(ast->children[1]);

		auto type_name = ast->children[0]->type.base_type->id;
		if (ast->children[0]->type.base_type->is_pointer())
		{
			auto pointer = dynamic_pointer_cast<type_pointer>(ast->children[0]->type.base_type);
			assert_cond(ast->children[1]->type.base_type->is_primitive());
			auto type = dynamic_pointer_cast<type_primitive>(ast->children[1]->type.base_type);
			assert_cond(type->is_integral());
			appender.append(instruction_ptr(new instruction_move(pointer->get_base_type()->size(), desc.op == "+")));
		}
		else if (ast->children[0]->type.base_type->is_array())
		{
			auto pointer = dynamic_pointer_cast<type_array>(ast->children[0]->type.base_type);
			assert_cond(ast->children[1]->type.base_type->is_primitive());
			auto type = dynamic_pointer_cast<type_primitive>(ast->children[1]->type.base_type);
			assert_cond(type->is_integral());
			appender.append(instruction_ptr(new instruction_move(pointer->get_base_type()->size(), desc.op == "+")));
		}
		else
		{
			assert_eq(ast->children[0]->type.base_type->id, ast->children[1]->type.base_type->id);

			bool flag;
			if (type_name == type_int->id)
				flag = build_binary_integral<int>(desc.op, appender);
			else if (type_name == type_short->id)
				flag = build_binary_integral<short>(desc.op, appender);
			else if (type_name == type_long->id)
				flag = build_binary_integral<long>(desc.op, appender);
			else if (type_name == type_int64->id)
				flag = build_binary_integral<long long>(desc.op, appender);
			else if (type_name == type_float->id)
				flag = build_binary_common<float>(desc.op, appender);
			else if (type_name == type_double->id)
				flag = build_binary_common<double>(desc.op, appender);
			else if (type_name == type_long_double->id)
				flag = build_binary_common<long double>(desc.op, appender);
			else if (type_name == type_bool->id)
				flag = build_binary_bool(desc.op, appender);
			else
				assert_cond(false); // Must be processed in semantic analyzer.

			if (!flag)
				appender.get_program()->compilation_error(ast->t, "Type %s does not accept operator %s", type_name.c_str(), desc.op.c_str());
		}
	}
}

bool build_unary_pointer_pointer(AST ast, const string &op, bytecode_appender &appender)
{
	if (op == "&")
		appender.append(instruction_ptr(new instruction_reference(ast->children[0]->type)));
	else if (op == "*")
		appender.append(instruction_ptr(new instruction_dereference_pointer()));
	else
		return false;
	return true;
}

bool build_unary_pointer_array(AST ast, const string &op, bytecode_appender &appender)
{
	if (op == "&")
		appender.append(instruction_ptr(new instruction_reference(ast->children[0]->type)));
	else if (op == "*")
		appender.append(instruction_ptr(new instruction_dereference_array()));
	else
		return false;
	return true;
}

template<typename T>
bool build_unary_pointer(AST ast, const string &op, bytecode_appender &appender)
{
	if (op == "&")
		appender.append(instruction_ptr(new instruction_reference(ast->children[0]->type)));
	else if (op == "*")
		appender.append(instruction_ptr(new instruction_dereference<T>()));
	else
		return false;
	return true;
}

template<typename T>
bool build_unary_common(AST ast, const string &op, bytecode_appender &appender)
{
	if (op == "&")
		appender.append(instruction_ptr(new instruction_reference(ast->children[0]->type)));
	else if (op == "-")
		appender.append(instruction_ptr(new instruction_neg<T>()));
	else
		return false;
	return true;
}

template<typename T>
bool build_unary_integral(AST ast, const string &op, bytecode_appender &appender)
{
	if (build_unary_common<T>(ast, op, appender))
		return true;
	else  if (op == "!")
		appender.append(instruction_ptr(new instruction_not<T>()));
	else if (op == "~")
		appender.append(instruction_ptr(new instruction_invert<T>()));
	else
		return false;
	return true;
}

bool build_unary_bool(AST ast, const string &op, bytecode_appender &appender)
{
	if (op == "&")
		appender.append(instruction_ptr(new instruction_reference(ast->children[0]->type)));
	else if (op == "!")
		appender.append(instruction_ptr(new instruction_not<bool>()));
	else
		return false;
	return true;
}

void compiler::bytecode_unary_operator::build(bytecode_appender &appender, AST ast)
{
	auto desc = any_cast<descriptor_unary_operator>(ast->desc);

	assert_eq(ast->children.size(), 1);

	appender.build_expression(ast->children[0]);

	auto type_name = ast->children[0]->type.base_type->id;

	bool flag = false;
	if (type_name == type_int->id)
		flag = build_unary_integral<int>(ast, desc.op, appender);
	else if (type_name == type_short->id)
		flag = build_unary_integral<short>(ast, desc.op, appender);
	else if (type_name == type_long->id)
		flag = build_unary_integral<long>(ast, desc.op, appender);
	else if (type_name == type_int64->id)
		flag = build_unary_integral<long long>(ast, desc.op, appender);
	else if (type_name == type_bool->id)
		flag = build_unary_bool(ast, desc.op, appender);
	else if (type_name == type_float->id)
		flag = build_unary_common<float>(ast, desc.op, appender);
	else if (type_name == type_double->id)
		flag = build_unary_common<double>(ast, desc.op, appender);
	else if (type_name == type_long_double->id)
		flag = build_unary_common<long double>(ast, desc.op, appender);
	else if (ast->children[0]->type.base_type->is_pointer())
	{
		type_pointer *p = dynamic_cast<type_pointer*>(ast->children[0]->type.base_type.get());

		string type_name = p->get_base_type()->id;
		if (type_name == type_int->id)
			flag = build_unary_pointer<int>(ast, desc.op, appender);
		else if (type_name == type_short->id)
			flag = build_unary_pointer<short>(ast, desc.op, appender);
		else if (type_name == type_long->id)
			flag = build_unary_pointer<long>(ast, desc.op, appender);
		else if (type_name == type_char->id)
			flag = build_unary_pointer<char>(ast, desc.op, appender);
		else if (type_name == type_int64->id)
			flag = build_unary_pointer<long long>(ast, desc.op, appender);
		else if (type_name == type_bool->id)
			flag = build_unary_pointer<long long>(ast, desc.op, appender);
		else if (type_name == type_float->id)
			flag = build_unary_pointer<float>(ast, desc.op, appender);
		else if (type_name == type_double->id)
			flag = build_unary_pointer<double>(ast, desc.op, appender);
		else if (type_name == type_long_double->id)
			flag = build_unary_pointer<long double>(ast, desc.op, appender);
		else if (p->get_base_type()->is_pointer())
			flag = build_unary_pointer_pointer(ast, desc.op, appender);
		else if (p->get_base_type()->is_array())
			flag = build_unary_pointer_array(ast, desc.op, appender);
	}
	else
		assert_cond(false); // Must be processed in semantic analyzer.

	if (!flag)
		throw runtime_error(string_format("Unexpected unary operator %s", desc.op.c_str()));
}

void compiler::bytecode_func_call::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_func_call>(ast->desc);

	for (auto &child : ast->children)
	{
		appender.build_expression(child);
	}

	appender.append(instruction_ptr(new instruction_invoke(desc.name, ast->children.size())));
}

void compiler::bytecode_constant::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_constant>(ast->desc);

	if (boost::starts_with(desc.value, "\"")) // a string
	{
		auto chars = parse_constant_string(desc.value.substr(1, desc.value.size() - 2));
		appender.get_constant_string_pool().set(appender.get_constant_string_pool_ptr(), chars.data(), chars.size());
		appender.append(instruction_ptr(new instruction_constant<char*>(appender.get_constant_string_pool().get<char>(appender.get_constant_string_pool_ptr()))));
		appender.increase_constant_string_pool_ptr(chars.size());
	}
	else if (boost::starts_with(desc.value, "'")) // a char
	{
		auto chars = parse_constant_string(desc.value.substr(1, desc.value.size() - 2));
		if (chars.size() != 2)
			appender.get_program()->compilation_error(ast->t, "Char length exceeds 1");
		appender.append(instruction_ptr(new instruction_constant<char>(chars[0])));
	}
	else if (desc.value == "true")
	{
		appender.append(instruction_ptr(new instruction_constant<int>(1)));
	}
	else if (desc.value == "false")
	{
		appender.append(instruction_ptr(new instruction_constant<int>(0)));
	}
	else
	{
		try
		{
			boost::to_lower(desc.value);
			if (desc.value.find("e") != string::npos)
			{
				vector<string> strs;
				boost::split(strs, desc.value, boost::is_any_of("e"));
				assert_eq(strs.size(), 2);
				int e = boost::lexical_cast<int>(strs[1]);
				double m = boost::lexical_cast<double>(strs[0]);
				appender.append(instruction_ptr(new instruction_constant<double>(m * pow(10, e))));
			}
			else if (desc.value.find(".") == string::npos)
			{
				if (boost::starts_with(desc.value, "0x"))
				{
					if (boost::ends_with(desc.value, "ll"))
						appender.append(instruction_ptr(new instruction_constant<long long>(boost::lexical_cast<lexical_hex<long long>>(desc.value.substr(0, desc.value.size() - 2)))));
					else if (boost::ends_with(desc.value, "l"))
						appender.append(instruction_ptr(new instruction_constant<long>(boost::lexical_cast<lexical_hex<long>>(desc.value.substr(0, desc.value.size() - 1)))));
					else
						appender.append(instruction_ptr(new instruction_constant<int>(boost::lexical_cast<lexical_hex<int>>(desc.value.substr(0, desc.value.size())))));
				}
				else if (boost::starts_with(desc.value, "0"))
				{
					if (boost::ends_with(desc.value, "ll"))
						appender.append(instruction_ptr(new instruction_constant<long long>(boost::lexical_cast<lexical_oct<long long>>(desc.value.substr(0, desc.value.size() - 2)))));
					else if (boost::ends_with(desc.value, "l"))
						appender.append(instruction_ptr(new instruction_constant<long>(boost::lexical_cast<lexical_oct<long>>(desc.value.substr(0, desc.value.size() - 1)))));
					else
						appender.append(instruction_ptr(new instruction_constant<int>(boost::lexical_cast<lexical_oct<int>>(desc.value.substr(0, desc.value.size())))));
				}
				else
				{
					if (boost::ends_with(desc.value, "ll"))
						appender.append(instruction_ptr(new instruction_constant<long long>(boost::lexical_cast<long long>(desc.value.substr(0, desc.value.size() - 2)))));
					else if (boost::ends_with(desc.value, "l"))
						appender.append(instruction_ptr(new instruction_constant<long>(boost::lexical_cast<long>(desc.value.substr(0, desc.value.size() - 1)))));
					else
						appender.append(instruction_ptr(new instruction_constant<int>(boost::lexical_cast<int>(desc.value.substr(0, desc.value.size())))));
				}
			}
			else
			{
				if (boost::ends_with(desc.value, "l"))
					appender.append(instruction_ptr(new instruction_constant<long double>(boost::lexical_cast<long double>(desc.value.substr(0, desc.value.size() - 1)))));
				else if (boost::ends_with(desc.value, "f"))
					appender.append(instruction_ptr(new instruction_constant<float>(boost::lexical_cast<float>(desc.value.substr(0, desc.value.size() - 1)))));
				else
					appender.append(instruction_ptr(new instruction_constant<double>(boost::lexical_cast<double>(desc.value.substr(0, desc.value.size())))));
			}
		}
		catch (boost::bad_lexical_cast&)
		{
			appender.get_program()->compilation_error(ast->t, "Not a valid constant");
		}
	}
}

void compiler::bytecode_var::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_var>(ast->desc);

	for (auto it = appender.get_variables().rbegin(); it != appender.get_variables().rend(); ++it)
		if ((*it).variables.count(desc.name))
		{
			auto re = (*it).variables[desc.name];
			appender.append(get_instruction_load(re.second, re.first));
			return;
		}

	if (appender.get_global_variables()->variables.count(desc.name))
	{
		auto re = appender.get_global_variables()->variables[desc.name];
		appender.append(get_instruction_load(re.second, re.first));
		return;
	}

	throw runtime_error(string_format("Variable %s not found.", desc.name.c_str()));
}

void compiler::bytecode_if::build(bytecode_appender & appender, AST ast)
{
	assert_cond(2 <= ast->children.size() && ast->children.size() <= 3);

	appender.build_expression(ast->children[0]);
	int if_exp = appender.append(instruction_ptr()), else_exp;
	appender.build_statements(ast->children[1]);
	if (ast->children.size() == 3)
		else_exp = appender.append(instruction_ptr());
	appender.modify(if_exp, instruction_ptr(new instruction_if(appender.get_current_instruction_number())));

	if (ast->children.size() == 3)
	{
		appender.build_statements(ast->children[2]);
		appender.modify(else_exp, instruction_ptr(new instruction_jump(appender.get_current_instruction_number())));
	}
}

void compiler::bytecode_while::build(bytecode_appender & appender, AST ast)
{
	assert_eq(ast->children.size(), 2);

	int start = appender.get_current_instruction_number();
	appender.build_expression(ast->children[0]);

	int if_exp = appender.append(instruction_ptr());
	appender.build_statements(ast->children[1]);
	appender.append(instruction_ptr(new instruction_jump(start)));

	appender.modify(if_exp, instruction_ptr(new instruction_if(appender.get_current_instruction_number())));
}

void compiler::bytecode_return::build(bytecode_appender & appender, AST ast)
{
	assert_cond(ast->children.size() <= 1);

	if (ast->children.size() == 1)
		appender.build_expression(ast->children[0]);

	appender.append(instruction_ptr(new instruction_return()));
}

void compiler::bytecode_primitive_cast::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_primitive_cast>(ast->desc);
	assert_eq(ast->children.size(), 1);
	appender.build_expression(ast->children[0]);
	auto ins = cast_primitive(desc.from_type.base_type, desc.to_type.base_type);
	if (ins != nullptr)
		appender.append(ins);
}

void build_inc(bytecode_appender & appender, const string &type_name, int delta)
{
	if (type_name == type_char->id)
		appender.append(instruction_ptr(new instruction_inc<char>(delta)));
	else if (type_name == type_int->id)
		appender.append(instruction_ptr(new instruction_inc<int>(delta)));
	else if (type_name == type_short->id)
		appender.append(instruction_ptr(new instruction_inc<short>(delta)));
	else if (type_name == type_long->id)
		appender.append(instruction_ptr(new instruction_inc<long>(delta)));
	else if (type_name == type_int64->id)
		appender.append(instruction_ptr(new instruction_inc<long long>(delta)));
	else
		assert_cond(0);
}

void compiler::bytecode_inc::build(bytecode_appender & appender, AST ast)
{
	auto desc = any_cast<descriptor_inc>(ast->desc);
	assert_eq(ast->children.size(), 1);

	appender.build_expression(ast->children[0]);
	if (desc.increase_first) // ++x
	{
		appender.append(instruction_ptr(new instruction_duplicate()));
		auto type_name = ast->children[0]->type.base_type->id;
		build_inc(appender, type_name, desc.delta);
		appender.append(instruction_ptr(new instruction_store()));
	}
	else // x++
	{
		appender.append(instruction_ptr(new instruction_duplicate()));
		appender.append(instruction_ptr(new instruction_duplicate()));
		auto type_name = ast->children[0]->type.base_type->id;
		build_inc(appender, type_name, desc.delta);
		appender.append(instruction_ptr(new instruction_store()));
		appender.append(instruction_ptr(new instruction_pop()));
	}
}

void compiler::bytecode_ternary::build(bytecode_appender & appender, AST ast)
{
	assert_eq(ast->children.size(), 3);

	appender.build_expression(ast->children[0]);
	int if_exp = appender.append(instruction_ptr());
	appender.build_expression(ast->children[1]);
	int else_exp = appender.append(instruction_ptr());
	appender.modify(if_exp, instruction_ptr(new instruction_if(appender.get_current_instruction_number())));
	appender.build_expression(ast->children[2]);
	appender.modify(else_exp, instruction_ptr(new instruction_jump(appender.get_current_instruction_number())));
}


#include "detail/random.hpp"
#include "model.hpp"
#include "observation.hpp"
#include "simulate.hpp"
#include "track.hpp"

#include <boost/random.hpp>
#include <cmath>
#include <ctime>
#include <deque>
#include <Eigen/Dense>

using namespace Eigen;

namespace biggles { namespace simulate {

boost::shared_ptr<track> generate_track(
    const model::parameters& parameters,
    const point_2d& initial_position,
    const point_2d& initial_velocity,
    time_stamp first_time_stamp,
    float position_sigma,
    float velocity_sigma)
{
    std::deque<state_t> states;
    return generate_track(states, parameters, initial_position, initial_velocity, first_time_stamp,
        position_sigma, velocity_sigma);
}

boost::shared_ptr<track> generate_track(
    std::deque<state_t> &states,
    const model::parameters& parameters,
    const point_2d& initial_position,
    const point_2d& initial_velocity,
    time_stamp first_time_stamp,
    float position_sigma,
    float velocity_sigma)
{
    boost::uniform_real<float> uni_real_dist(0,1);
    boost::variate_generator<boost::mt19937&, boost::uniform_real<float> >
        uni_real(detail::random_generator, uni_real_dist);

    boost::normal_distribution<float> norm_dist(0,1);
    boost::variate_generator<boost::mt19937&, boost::normal_distribution<float> >
        norm(detail::random_generator, norm_dist);

    // the observation vector type: [x, y]
    typedef Vector2f obs_t;

    // the state evolution matrix: X_{k+1} = A X_k + <noise>
    Matrix4f A;
    //float d(uni_real() < 0.5 ? 1 : 0); // the dynamic drag
    const float d(1.0f);
    A << 1, d, 0, 0,
         0, 1, 0, 0,
         0, 0, 1, d,
         0, 0, 0, 1;

    // the observation matrix: Y_k = B X_k + <noise>
    Matrix<float, 2, 4> B;
    B << 1, 0, 0, 0,
         0, 0, 1, 0;

    // initial state
    state_t X;
    X << x(initial_position), x(initial_velocity), y(initial_position), y(initial_velocity);

    // collection we'll store generated observations in
    std::deque<observation> observations;

    // extract the parameters we're interested in
    Matrix2f R(model::observation_error_covariance(parameters));
    float p_s(model::frame_to_frame_survival_probability(parameters));
    float p_d(model::generate_observation_probability(parameters));

    float sqrt_R_00(sqrt(R(0,0))), sqrt_R_11(sqrt(R(1,1)));

    // make sure we generate at least one observation. Keep time stamp in global scope because we use it later as the
    // duration.
    time_stamp time_stamp(0);
    for(time_stamp = first_time_stamp; observations.empty() || (uni_real() <= p_s); ++time_stamp)
    {
        states.push_back(X);
        // generate observation
        obs_t Y = B * X;

        // noise observation - FIXME: does not use cross-covariance
        Y[0] += sqrt_R_00 * norm();
        Y[1] += sqrt_R_11 * norm();

        // record observation if an observation is made
        if(uni_real() <= p_d)
            observations.push_back(new_obs(Y[0], Y[1], time_stamp));

        // evolve state
        X = A * X;

        // noise state
        X[0] += position_sigma * norm();
        X[1] += velocity_sigma * norm();
        X[2] += position_sigma * norm();
        X[3] += velocity_sigma * norm();
    }

    return boost::shared_ptr<track>(
        new track(first_time_stamp, time_stamp,
                  observations.begin(),
                  observations.end(), d));
}

void generate_partition(partition_ptr_t& output_part_ptr,
                        const model::parameters& params,
                        time_stamp start_time_stamp,
                        time_stamp end_time_stamp,
                        float image_width,
                        float image_height,
                        float position_sigma,
                        float velocity_sigma,
                        float initial_velocity_sigma)
{
    // the various random variate generators.
    boost::uniform_real<float> uni_real_dist(0,1);
    boost::variate_generator<boost::mt19937&, boost::uniform_real<float> >
        uni_real(detail::random_generator, uni_real_dist);

    boost::poisson_distribution<int, float> birth_count_dist(model::mean_new_tracks_per_frame(params));
    boost::variate_generator<boost::mt19937&, boost::poisson_distribution<int, float> >
        birth_count(detail::random_generator, birth_count_dist);

    boost::poisson_distribution<int, float> clutter_count_dist(model::mean_false_observations_per_frame(params));
    boost::variate_generator<boost::mt19937&, boost::poisson_distribution<int, float> >
        clutter_count(detail::random_generator, clutter_count_dist);

    boost::normal_distribution<float> norm_dist(0,1);
    boost::variate_generator<boost::mt19937&, boost::normal_distribution<float> >
        norm(detail::random_generator, norm_dist);

    boost::shared_ptr<track_collection> tracks(new track_collection());
    clutter_ptr clutter(new clutter_t());

    // the simulation work proper
    size_t expected_clutter_size(0);
    for(time_stamp time_stamp = start_time_stamp; time_stamp < end_time_stamp; ++time_stamp)
    {
        size_t n_new_tracks(birth_count());
        for(size_t track_idx = 0; track_idx < n_new_tracks; ++track_idx)
        {
            boost::shared_ptr<track> new_track = generate_track(
                params,
                point_2d(image_width * uni_real(), image_height * uni_real()),
                point_2d(initial_velocity_sigma * norm(), initial_velocity_sigma * norm()),
                time_stamp, position_sigma, velocity_sigma);

            if(new_track->size() >= 2)
                tracks->insert(new_track);
        }

        size_t n_clutter_obs(clutter_count());
        for(size_t clutter_idx = 0; clutter_idx < n_clutter_obs; ++clutter_idx)
        {
            clutter->insert( new_obs(image_width * uni_real(), image_height * uni_real(), time_stamp));
        }
        BOOST_ASSERT(clutter->count_at_time_stamp(time_stamp) == n_clutter_obs);

        expected_clutter_size += n_clutter_obs;
        BOOST_ASSERT(clutter->size() == expected_clutter_size);
    }

    // set output
    output_part_ptr = partition_ptr_t(new partition(tracks, clutter, partition::expansion_2d(0.f, image_width, 0.f, image_height)));
}

void crop_partition(partition_ptr_t& output_part_ptr,
                    const partition_ptr_t& input_part_ptr,
                    float first_x, float last_x,
                    float first_y, float last_y,
                    time_stamp first_time_stamp,
                    time_stamp last_time_stamp)
{
    clutter_ptr new_clutter(new clutter_t());
    boost::shared_ptr<track_collection> new_tracks(new track_collection());

    // crop clutter observations
    for (clutter_t::const_iterator it = input_part_ptr->clutter().begin();
                    it not_eq input_part_ptr->clutter().end(); ++it) {
        BOOST_FOREACH(const observation& o, it->second) {
            if((x(o) < first_x) || (x(o) >= last_x))
                continue;
            if((y(o) < first_y) || (y(o) >= last_y))
                continue;
            if((t(o) < first_time_stamp) || (t(o) >= last_time_stamp))
                continue;
            new_clutter->insert(o);
        }
    }

    // create new cropped tracks and insert them into the track collection
    BOOST_FOREACH(const boost::shared_ptr<const track>& p_track, input_part_ptr->tracks())
    {
        // create a new set of cropped observations
        std::deque<observation> cropped_observations;
        BOOST_FOREACH(const observation& o, *p_track)
        {
            if((x(o) < first_x) || (x(o) >= last_x))
                continue;
            if((y(o) < first_y) || (y(o) >= last_y))
                continue;
            if((t(o) < first_time_stamp) || (t(o) >= last_time_stamp))
                continue;
            cropped_observations.push_back(o);
        }

        // if no observations remain, remove this track
        if(0 == cropped_observations.size())
            continue;

        // create a new track based on old track.
        boost::shared_ptr<track> new_track(
                new track(std::max(p_track->first_time_stamp(), first_time_stamp),
                          std::min(p_track->last_time_stamp(), last_time_stamp),
                          cropped_observations.begin(),
                          cropped_observations.end(), p_track->dynamic_drag()));

        if(new_track->size() >= 2)
            new_tracks->insert(new_track);
    }

    // assign new tracks and clutter to output partition
    output_part_ptr = partition_ptr_t(
        new partition(new_tracks, new_clutter, partition::expansion_2d(first_x, last_x, first_y, last_y)));
}

void demote_all_tracks_from_partition(partition_ptr_t& output_part_ptr,
                                      const partition_ptr_t& input_part_ptr)
{
    // copy the clutter ready to demote observations
    clutter_ptr clutter(new clutter_t(input_part_ptr->clutter()));

    // a new, empty set of tracks
    boost::shared_ptr<track_collection> tracks(new track_collection());

    // for each input track...
    BOOST_FOREACH(const boost::shared_ptr<const track>& p_track, input_part_ptr->tracks())
    {
        // copy it's observations to the clutter
        clutter->insert(p_track->begin(), p_track->end());
    }

    // return new partition
    output_part_ptr = partition_ptr_t(new partition(tracks, clutter, input_part_ptr->expansion()));
}

} }