///////////////////////////////////////////////////////////////////////////////
// BSD 3-Clause License
//
// Copyright (C) 2019-2023, University of Edinburgh, LAAS-CNRS,
//                          Heriot-Watt University
// Copyright note valid unless otherwise stated in individual files.
// All rights reserved.
///////////////////////////////////////////////////////////////////////////////

#ifndef CROCODDYL_LEGGED_ROBOTS_FACTORY_HPP_
#define CROCODDYL_LEGGED_ROBOTS_FACTORY_HPP_

#include <example-robot-data/path.hpp>
#include <pinocchio/algorithm/model.hpp>
#include <pinocchio/parsers/srdf.hpp>
#include <pinocchio/parsers/urdf.hpp>

#include "crocoddyl/core/costs/cost-sum.hpp"
#include "crocoddyl/core/costs/residual.hpp"
#include "crocoddyl/core/integrator/euler.hpp"
#include "crocoddyl/core/mathbase.hpp"
#include "crocoddyl/core/residuals/control.hpp"
#include "crocoddyl/multibody/actions/contact-fwddyn.hpp"
#include "crocoddyl/multibody/actuations/floating-base.hpp"
#include "crocoddyl/multibody/contacts/contact-3d.hpp"
#include "crocoddyl/multibody/contacts/contact-6d.hpp"
#include "crocoddyl/multibody/contacts/multiple-contacts.hpp"
#include "crocoddyl/multibody/residuals/com-position.hpp"
#include "crocoddyl/multibody/residuals/frame-placement.hpp"
#include "crocoddyl/multibody/residuals/frame-translation.hpp"
#include "crocoddyl/multibody/residuals/state.hpp"
#include "crocoddyl/multibody/states/multibody.hpp"
#include "robot-ee-names.hpp"

namespace crocoddyl {
namespace benchmark {

template <typename Scalar>
void build_contact_action_models(
    RobotEENames robotNames,
    boost::shared_ptr<crocoddyl::ActionModelAbstractTpl<Scalar> >& runningModel,
    boost::shared_ptr<crocoddyl::ActionModelAbstractTpl<Scalar> >&
        terminalModel) {
  typedef
      typename crocoddyl::DifferentialActionModelContactFwdDynamicsTpl<Scalar>
          DifferentialActionModelContactFwdDynamics;
  typedef typename crocoddyl::IntegratedActionModelEulerTpl<Scalar>
      IntegratedActionModelEuler;
  typedef typename crocoddyl::ActuationModelFloatingBaseTpl<Scalar>
      ActuationModelFloatingBase;
  typedef typename crocoddyl::CostModelSumTpl<Scalar> CostModelSum;
  typedef typename crocoddyl::ContactModelMultipleTpl<Scalar>
      ContactModelMultiple;
  typedef typename crocoddyl::CostModelAbstractTpl<Scalar> CostModelAbstract;
  typedef typename crocoddyl::ContactModelAbstractTpl<Scalar>
      ContactModelAbstract;
  typedef typename crocoddyl::CostModelResidualTpl<Scalar> CostModelResidual;
  typedef typename crocoddyl::ResidualModelFramePlacementTpl<Scalar>
      ResidualModelFramePlacement;
  typedef typename crocoddyl::ResidualModelCoMPositionTpl<Scalar>
      ResidualModelCoMPosition;
  typedef typename crocoddyl::ResidualModelStateTpl<Scalar> ResidualModelState;
  typedef typename crocoddyl::ResidualModelControlTpl<Scalar>
      ResidualModelControl;
  typedef typename crocoddyl::ContactModel6DTpl<Scalar> ContactModel6D;
  typedef typename crocoddyl::ContactModel3DTpl<Scalar> ContactModel3D;
  typedef typename crocoddyl::MathBaseTpl<Scalar>::Vector2s Vector2s;
  typedef typename crocoddyl::MathBaseTpl<Scalar>::Vector3s Vector3s;
  typedef typename crocoddyl::MathBaseTpl<Scalar>::VectorXs VectorXs;
  typedef typename crocoddyl::MathBaseTpl<Scalar>::Matrix3s Matrix3s;

  pinocchio::ModelTpl<double> modeld;
  pinocchio::urdf::buildModel(robotNames.urdf_path,
                              pinocchio::JointModelFreeFlyer(), modeld);
  modeld.lowerPositionLimit.head<7>().array() = -1;
  modeld.upperPositionLimit.head<7>().array() = 1.;
  pinocchio::srdf::loadReferenceConfigurations(modeld, robotNames.srdf_path,
                                               false);

  pinocchio::ModelTpl<Scalar> model(modeld.cast<Scalar>());
  boost::shared_ptr<crocoddyl::StateMultibodyTpl<Scalar> > state =
      boost::make_shared<crocoddyl::StateMultibodyTpl<Scalar> >(
          boost::make_shared<pinocchio::ModelTpl<Scalar> >(model));

  VectorXs default_state(model.nq + model.nv);
  default_state << model.referenceConfigurations[robotNames.reference_conf],
      VectorXs::Zero(model.nv);

  boost::shared_ptr<ActuationModelFloatingBase> actuation =
      boost::make_shared<ActuationModelFloatingBase>(state);

  boost::shared_ptr<CostModelAbstract> comCost =
      boost::make_shared<CostModelResidual>(
          state, boost::make_shared<ResidualModelCoMPosition>(
                     state, Vector3s::Zero(), actuation->get_nu()));
  boost::shared_ptr<CostModelAbstract> goalTrackingCost =
      boost::make_shared<CostModelResidual>(
          state, boost::make_shared<ResidualModelFramePlacement>(
                     state, model.getFrameId(robotNames.ee_name),
                     pinocchio::SE3Tpl<Scalar>(
                         Matrix3s::Identity(),
                         Vector3s(Scalar(.0), Scalar(.0), Scalar(.4))),
                     actuation->get_nu()));
  boost::shared_ptr<CostModelAbstract> xRegCost =
      boost::make_shared<CostModelResidual>(
          state, boost::make_shared<ResidualModelState>(state, default_state,
                                                        actuation->get_nu()));
  boost::shared_ptr<CostModelAbstract> uRegCost =
      boost::make_shared<CostModelResidual>(
          state,
          boost::make_shared<ResidualModelControl>(state, actuation->get_nu()));

  // Create a cost model per the running and terminal action model.
  boost::shared_ptr<CostModelSum> runningCostModel =
      boost::make_shared<CostModelSum>(state, actuation->get_nu());
  boost::shared_ptr<CostModelSum> terminalCostModel =
      boost::make_shared<CostModelSum>(state, actuation->get_nu());

  // Then let's added the running and terminal cost functions
  runningCostModel->addCost("gripperPose", goalTrackingCost, Scalar(1));
  //   runningCostModel->addCost("comPos", comCost, Scalar(1e-7));
  runningCostModel->addCost("xReg", xRegCost, Scalar(1e-4));
  runningCostModel->addCost("uReg", uRegCost, Scalar(1e-4));
  terminalCostModel->addCost("gripperPose", goalTrackingCost, Scalar(1));

  boost::shared_ptr<ContactModelMultiple> contact_models =
      boost::make_shared<ContactModelMultiple>(state, actuation->get_nu());

  for (std::size_t i = 0; i < robotNames.contact_names.size(); ++i) {
    switch (robotNames.contact_types[i]) {
      case Contact3D: {
        boost::shared_ptr<ContactModelAbstract> support_contact =
            boost::make_shared<ContactModel3D>(
                state, model.getFrameId(robotNames.contact_names[i]),
                Eigen::Vector3d::Zero(), pinocchio::LOCAL_WORLD_ALIGNED,
                actuation->get_nu(), Vector2s(Scalar(0.), Scalar(50.)));
        contact_models->addContact(
            model.frames[model.getFrameId(robotNames.contact_names[i])].name,
            support_contact);
        break;
      }
      case Contact6D: {
        boost::shared_ptr<ContactModelAbstract> support_contact =
            boost::make_shared<ContactModel6D>(
                state, model.getFrameId(robotNames.contact_names[i]),
                pinocchio::SE3Tpl<Scalar>::Identity(),
                pinocchio::LOCAL_WORLD_ALIGNED, actuation->get_nu(),
                Vector2s(Scalar(0.), Scalar(50.)));
        contact_models->addContact(
            model.frames[model.getFrameId(robotNames.contact_names[i])].name,
            support_contact);
        break;
      }
      default: {
        break;
      }
    }
  }

  // Next, we need to create an action model for running and terminal nodes
  boost::shared_ptr<DifferentialActionModelContactFwdDynamics> runningDAM =
      boost::make_shared<DifferentialActionModelContactFwdDynamics>(
          state, actuation, contact_models, runningCostModel);
  boost::shared_ptr<DifferentialActionModelContactFwdDynamics> terminalDAM =
      boost::make_shared<DifferentialActionModelContactFwdDynamics>(
          state, actuation, contact_models, terminalCostModel);

  runningModel =
      boost::make_shared<IntegratedActionModelEuler>(runningDAM, Scalar(5e-3));
  terminalModel =
      boost::make_shared<IntegratedActionModelEuler>(terminalDAM, Scalar(5e-3));
}

}  // namespace benchmark
}  // namespace crocoddyl

#endif  // CROCODDYL_BIPED_FACTORY_HPP_
