/*********************************************************************
* Software License Agreement (BSD License)
*
*  Copyright (c) 2010, Rice University
*  All rights reserved.
*
*  Redistribution and use in source and binary forms, with or without
*  modification, are permitted provided that the following conditions
*  are met:
*
*   * Redistributions of source code must retain the above copyright
*     notice, this list of conditions and the following disclaimer.
*   * Redistributions in binary form must reproduce the above
*     copyright notice, this list of conditions and the following
*     disclaimer in the documentation and/or other materials provided
*     with the distribution.
*   * Neither the name of the Rice University nor the names of its
*     contributors may be used to endorse or promote products derived
*     from this software without specific prior written permission.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
*  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
*  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
*  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
*  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
*  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
*  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
*  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
*  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
*  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
*  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
*  POSSIBILITY OF SUCH DAMAGE.
*********************************************************************/

/* Author: Mark Moll */

#include <ompl/extensions/vortex/VortexSimpleSetup.h>
#include <ompl/base/GoalRegion.h>

#include <Vx/VxBox.h>
#include <Vx/VxPart.h>
#include <Vx/VxAssembly.h>
#include <Vx/VxState.h>

#include <ompl/config.h>
#include <iostream>

namespace ob = ompl::base;
namespace og = ompl::geometric;
namespace oc = ompl::control;


class RigidBodyEnvironment : public oc::VortexEnvironment
{
public:

    RigidBodyEnvironment(void) : oc::VortexEnvironment()
    {
        createWorld();
    }
    ~RigidBodyEnvironment()
    {
        visualizer_->stopViewer();
    }

    /**************************************************
     * Implementation of functions needed by planning *
     **************************************************/

    virtual unsigned int getControlDimension(void) const
    {
        return 3;
    }

    virtual void getControlBounds(std::vector<double> &lower, std::vector<double> &upper) const
    {
        static double maxForce = 1.;
        lower.resize(3);
        lower[0] = -maxForce;
        lower[1] = -maxForce;
        lower[2] = 9.81 - maxForce;

        upper.resize(3);
        upper[0] = maxForce;
        upper[1] = maxForce;
        upper[2] = 9.81 + maxForce;
    }

    virtual void applyControl(const double *control) const
    {
        Vx::VxReal3 force = { control[0], control[1], control[2] };
        static_cast<Vx::VxPart*>(assembly_->findEntityFromName("moving box"))->addForce(force);
    }

    /**************************************************/

    // OMPL does not require this function here; we implement it here
    // for convenience. This function is only Vortex code to create a
    // simulation environment. At the end of the function, there is a
    // call to setPlanningParameters(), which configures members of
    // the base class needed by planners.
    void createWorld(void);

    // Set parameters needed by the base class (such as the bodies
    // that make up to state of the system we are planning for)
    void setPlanningParameters(void);

    // convenience function to get the position of the box
    const Vx::VxVector3& getPosition(const ob::State* st) const;
};


// Define the goal we want to reach
class RigidBodyGoal : public ob::GoalRegion
{
public:
    RigidBodyGoal(const ob::SpaceInformationPtr &si) : ob::GoalRegion(si)
    {
        threshold_ = 1.;
    }

    virtual double distanceGoal(const ob::State *st) const
    {
        static const Vx::VxVector3 goalPos(30, 55, 35);
        const Vx::VxVector3& pos =
            si_->getStateSpace()->as<oc::VortexStateSpace>()->getEnvironment()->as<RigidBodyEnvironment>()->getPosition(st);
        return (goalPos - pos).norm();
    }
};


// Define how we project a state
class RigidBodyStateProjectionEvaluator : public ob::ProjectionEvaluator
{
public:

    RigidBodyStateProjectionEvaluator(const ob::StateSpace *space) : ob::ProjectionEvaluator(space)
    {
        cellSizes_.resize(3);
        cellSizes_[0] = 1;
        cellSizes_[1] = 1;
        cellSizes_[2] = 1;
    }

    virtual unsigned int getDimension(void) const
    {
        return 3;
    }

    virtual void project(const ob::State *state, ob::EuclideanProjection &projection) const
    {
        const Vx::VxVector3& pos =
            space_->as<oc::VortexStateSpace>()->getEnvironment()->as<RigidBodyEnvironment>()->getPosition(state);
        projection[0] = pos[0];
        projection[1] = pos[1];
        projection[2] = pos[2];
    }

};

// Define our own state space, to include a distance function we want and register a default projection
class RigidBodyStateSpace : public oc::VortexStateSpace
{
public:
    RigidBodyStateSpace(const oc::VortexEnvironmentPtr &env) : oc::VortexStateSpace(env)
    {
    }

    virtual double distance(const ob::State *s1, const ob::State *s2) const
    {
        return (env_->as<RigidBodyEnvironment>()->getPosition(s1) 
              - env_->as<RigidBodyEnvironment>()->getPosition(s2)).norm();
    }

    virtual void registerProjections(void)
    {
        registerDefaultProjection(ob::ProjectionEvaluatorPtr(new RigidBodyStateProjectionEvaluator(this)));
    }
};


int main(int, char **)
{
    // create the Vortex environment
    oc::VortexEnvironmentPtr env(new RigidBodyEnvironment());

    // create the state space and the control space for planning
    ob::StateSpacePtr stateSpacePtr = ob::StateSpacePtr(new RigidBodyStateSpace(env));

    // this will take care of setting a proper collision checker and the
    // starting state for the planner as the initial Vortex state
    oc::VortexSimpleSetup ss(stateSpacePtr);

    // set the goal we would like to reach
    ss.setGoal(ob::GoalPtr(new RigidBodyGoal(ss.getSpaceInformation())));

    ob::RealVectorBounds bounds(3);
    bounds.setLow(-1);
    bounds.setHigh(70);
    stateSpacePtr->as<RigidBodyStateSpace>()->setWorkspaceBounds(bounds);


    ss.setup();
    ss.print();
    
    // ss.simulate(10);
    // return 0;

    if (ss.solve(30))
    {
        bool running;
        oc::PathControl p(ss.getSolutionPath());
        //p.interpolate(); // uncomment if you want to plot the path

        // print out box positions along the path
        for (unsigned int i = 0 ; i < p.states.size() ; ++i)
        {
            const Vx::VxVector3 pos = env->as<RigidBodyEnvironment>()->getPosition(p.states[i]);
            std::cout << pos[0] << ' ' << pos[1] << ' ' << pos[2] << ' ';
            if (i==0)
                // null controls applied for zero seconds to get to start state
                std::cout << "0 0 0 0";
            else
            {
                // print controls and control duration needed to get from state i-1 to state i
                const double* c = p.controls[i-1]->as<oc::RealVectorControlSpace::ControlType>()->values;
                std::cout << c[0] << ' ' << c[1] << ' ' << c[2] <<' ' << p.controlDurations[i-1];
            }
            std::cout << std::endl;
        }
        // loop over the box's motion
        do
        {
            running = ss.playSolutionPath();
        } while (running);
    }
    return 0;
}







/***********************************************
 * Member function implementations             *
 ***********************************************/

void RigidBodyEnvironment::createWorld(void)
{
    // BEGIN SETTING UP A VORTEX ENVIRONMENT
    // ***********************************
    Vx::VxUniverse* universe = new Vx::VxUniverse();
    universe->setGravity(0., 0., -9.81);
    world_->addUniverse(universe);

    visualizer_->addDataFilePathList(VORTEX_RESOURCE_DIR);
    visualizer_->createDefaultWindow(0, NULL, "Vortex Rigid Body Planning Demo");
    visualizer_->setUniverse(world_->getUniverse(0));
    visualizer_->setFrameRateLock(true);
    visualizer_->createHelp();
    visualizer_->toggleHelp();

    Vx::VxAssembly* assembly = new Vx::VxAssembly();
    Vx::VxPart* box = new Vx::VxPart(1);
    Vx::VxReal3 boxSize = {.2, .2, .1};
    float boxColor[4] = {1., .1, .1, 1.};

    box->setName("moving box");
    box->addGeometry(new Vx::VxBox(boxSize), 0);
    assembly->addEntity(box);
    assembly_ = assembly; // assembly_ contains all movable parts
    universe->addAssembly(assembly_);
    Vx::VxNode boxNode = visualizer_->createBox(boxSize, boxColor);
    box->setNode(boxNode);
    box->setPosition(0., 0., 0.);
    box->setOrientationQuaternion(1., 0., 0., 0.);

    visualizer_->setCameraLookAtAndPosition(Vx::VxVector3(30,55,20), Vx::VxVector3(-5,-5,-5));
    // *********************************
    // END SETTING UP AN Vortex ENVIRONMENT

    setPlanningParameters();
}

void RigidBodyEnvironment::setPlanningParameters(void)
{
    // Fill in parameters for OMPL:
    stepSize_ = 0.05;
    minControlSteps_ = 5;
    maxControlSteps_ = 50;
}

const Vx::VxVector3& RigidBodyEnvironment::getPosition(const ob::State* st) const
{
    const Vx::VxState* vxstate = st->as<oc::VortexStateSpace::StateType>()->vxstate_;
    vxstate->restoreState(assembly_);
    return static_cast<Vx::VxAssembly*>(assembly_)->getEntity(0)->getPosition();
}