//
// Copyright (c) 2016-2020 CNRS INRIA
//

#pragma once

// IWYU pragma: private, include "pinocchio/multibody/liegroup.hpp"

#ifdef PINOCCHIO_LSP
  #undef PINOCCHIO_LSP
  #include "pinocchio/multibody/liegroup.hpp"
#endif // PINOCCHIO_LSP

namespace pinocchio
{
  template<int Dim, typename Scalar, int Options = context::Options>
  struct VectorSpaceOperationTpl;

  template<int Dim, typename _Scalar, int _Options>
  struct traits<VectorSpaceOperationTpl<Dim, _Scalar, _Options>>
  {
    typedef _Scalar Scalar;
    static constexpr int Options = _Options;
    static constexpr int NQ = Dim;
    static constexpr int NV = Dim;
  };

  template<int Dim, typename _Scalar, int _Options>
  struct VectorSpaceOperationTpl
  : public LieGroupBase<VectorSpaceOperationTpl<Dim, _Scalar, _Options>>
  {
    PINOCCHIO_LIE_GROUP_TPL_PUBLIC_INTERFACE(VectorSpaceOperationTpl);

    /// Constructor
    /// \param size size of the vector space: should be the equal to template
    ///        argument for static sized vector-spaces.
    VectorSpaceOperationTpl(int size = boost::static_signed_max<0, Dim>::value)
    : size_(size)
    {
      assert(size_.value() >= 0);
    }

    /// Constructor
    /// \param other other VectorSpaceOperationTpl from which to retrieve size
    VectorSpaceOperationTpl(const VectorSpaceOperationTpl & other)
    : Base()
    , size_(other.size_.value())
    {
      assert(size_.value() >= 0);
    }

    VectorSpaceOperationTpl & operator=(const VectorSpaceOperationTpl & other)
    {
      size_.setValue(other.size_.value());
      assert(size_.value() >= 0);
      return *this;
    }

    Index nq() const
    {
      return size_.value();
    }
    Index nv() const
    {
      return size_.value();
    }

    ConfigVector_t neutral() const
    {
      return ConfigVector_t::Zero(size_.value());
    }

    std::string name() const
    {
      std::ostringstream oss;
      oss << "R^" << nq();
      return oss.str();
    }

    template<class ConfigL_t, class ConfigR_t, class Tangent_t>
    static void difference_impl(
      const Eigen::MatrixBase<ConfigL_t> & q0,
      const Eigen::MatrixBase<ConfigR_t> & q1,
      const Eigen::MatrixBase<Tangent_t> & d)
    {
      PINOCCHIO_EIGEN_CONST_CAST(Tangent_t, d) = q1 - q0;
    }

    template<ArgumentPosition arg, class ConfigL_t, class ConfigR_t, class JacobianOut_t>
    void dDifference_impl(
      const Eigen::MatrixBase<ConfigL_t> &,
      const Eigen::MatrixBase<ConfigR_t> &,
      const Eigen::MatrixBase<JacobianOut_t> & J) const
    {
      if (arg == ARG0)
        PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J) =
          -JacobianMatrix_t::Identity(size_.value(), size_.value());
      else if (arg == ARG1)
        PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J).setIdentity();
    }

    template<
      ArgumentPosition arg,
      class ConfigL_t,
      class ConfigR_t,
      class JacobianIn_t,
      class JacobianOut_t>
    static void dDifference_product_impl(
      const ConfigL_t &,
      const ConfigR_t &,
      const JacobianIn_t & Jin,
      JacobianOut_t & Jout,
      bool,
      const AssignmentOperatorType op)
    {
      switch (op)
      {
      case SETTO:
        if (arg == ARG0)
          Jout = -Jin;
        else
          Jout = Jin;
        return;
      case ADDTO:
        if (arg == ARG0)
          Jout -= Jin;
        else
          Jout += Jin;
        return;
      case RMTO:
        if (arg == ARG0)
          Jout += Jin;
        else
          Jout -= Jin;
        return;
      }
    }

    template<class ConfigIn_t, class Velocity_t, class ConfigOut_t>
    static void integrate_impl(
      const Eigen::MatrixBase<ConfigIn_t> & q,
      const Eigen::MatrixBase<Velocity_t> & v,
      const Eigen::MatrixBase<ConfigOut_t> & qout)
    {
      PINOCCHIO_EIGEN_CONST_CAST(ConfigOut_t, qout) = q + v;
    }

    template<class Config_t, class Jacobian_t>
    static void integrateCoeffWiseJacobian_impl(
      const Eigen::MatrixBase<Config_t> &, const Eigen::MatrixBase<Jacobian_t> & J)
    {
      PINOCCHIO_EIGEN_CONST_CAST(Jacobian_t, J).setIdentity();
    }

    template<class Config_t, class Tangent_t, class JacobianOut_t>
    static void dIntegrate_dq_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<JacobianOut_t> & J,
      const AssignmentOperatorType op = SETTO)
    {
      Eigen::MatrixBase<JacobianOut_t> & Jout = PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J);
      switch (op)
      {
      case SETTO:
        Jout.setIdentity();
        break;
      case ADDTO:
        Jout.diagonal().array() += Scalar(1);
        break;
      case RMTO:
        Jout.diagonal().array() -= Scalar(1);
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t, class Tangent_t, class JacobianOut_t>
    static void dIntegrate_dv_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<JacobianOut_t> & J,
      const AssignmentOperatorType op = SETTO)
    {
      Eigen::MatrixBase<JacobianOut_t> & Jout = PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J);
      switch (op)
      {
      case SETTO:
        Jout.setIdentity();
        break;
      case ADDTO:
        Jout.diagonal().array() += Scalar(1);
        break;
      case RMTO:
        Jout.diagonal().array() -= Scalar(1);
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t, class Tangent_t, class JacobianIn_t, class JacobianOut_t>
    static void dIntegrate_product_impl(
      const Config_t &,
      const Tangent_t &,
      const JacobianIn_t & Jin,
      JacobianOut_t & Jout,
      bool,
      const ArgumentPosition,
      const AssignmentOperatorType op)
    {
      switch (op)
      {
      case SETTO:
        Jout = Jin;
        break;
      case ADDTO:
        Jout += Jin;
        break;
      case RMTO:
        Jout -= Jin;
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t, class Tangent_t, class JacobianIn_t, class JacobianOut_t>
    static void dIntegrateTransport_dq_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<JacobianIn_t> & Jin,
      const Eigen::MatrixBase<JacobianOut_t> & Jout)
    {
      PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, Jout) = Jin;
    }

    template<class Config_t, class Tangent_t, class JacobianIn_t, class JacobianOut_t>
    static void dIntegrateTransport_dv_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<JacobianIn_t> & Jin,
      const Eigen::MatrixBase<JacobianOut_t> & Jout)
    {
      PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, Jout) = Jin;
    }

    template<class Config_t, class Tangent_t, class Jacobian_t>
    static void dIntegrateTransport_dq_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<Jacobian_t> & /*J*/)
    {
    }

    template<class Config_t, class Tangent_t, class Jacobian_t>
    static void dIntegrateTransport_dv_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<Tangent_t> & /*v*/,
      const Eigen::MatrixBase<Jacobian_t> & /*J*/)
    {
    }

    // template <class ConfigL_t, class ConfigR_t>
    // static context::Scalar squaredDistance_impl(const Eigen::MatrixBase<ConfigL_t> & q0,
    // const Eigen::MatrixBase<ConfigR_t> & q1)

    template<class Config_t, class TangentMap_t>
    static void tangentMap_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      Eigen::MatrixBase<TangentMap_t> & TM,
      const AssignmentOperatorType op)
    {
      switch (op)
      {
      case SETTO:
        TM.setIdentity();
        break;
      case ADDTO:
        TM.diagonal().array() += Scalar(1);
        break;
      case RMTO:
        TM.diagonal().array() -= Scalar(1);
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t, class MatrixIn_t, class MatrixOut_t>
    static void tangentMapProduct_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<MatrixIn_t> & Min,
      Eigen::MatrixBase<MatrixOut_t> & Mout,
      const AssignmentOperatorType op)
    {
      switch (op)
      {
      case SETTO:
        Mout = Min;
        break;
      case ADDTO:
        Mout += Min;
        break;
      case RMTO:
        Mout -= Min;
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t, class MatrixIn_t, class MatrixOut_t>
    static void tangentMapTransposeProduct_impl(
      const Eigen::MatrixBase<Config_t> & /*q*/,
      const Eigen::MatrixBase<MatrixIn_t> & Min,
      Eigen::MatrixBase<MatrixOut_t> & Mout,
      const AssignmentOperatorType op)
    {
      switch (op)
      {
      case SETTO:
        Mout = Min;
        break;
      case ADDTO:
        Mout += Min;
        break;
      case RMTO:
        Mout -= Min;
        break;
      default:
        PINOCCHIO_UNREACHABLE();
      }
    }

    template<class Config_t>
    static void normalize_impl(const Eigen::MatrixBase<Config_t> & /*qout*/)
    {
    }

    template<class Config_t>
    static bool
    isNormalized_impl(const Eigen::MatrixBase<Config_t> & /*qout*/, const Scalar & /*prec*/)
    {
      return true;
    }

    template<class Config_t>
    static void random_impl(const Eigen::MatrixBase<Config_t> & qout)
    {
      PINOCCHIO_EIGEN_CONST_CAST(Config_t, qout).setRandom();
    }

    template<class ConfigL_t, class ConfigR_t, class ConfigOut_t>
    void randomConfiguration_impl(
      const Eigen::MatrixBase<ConfigL_t> & lower_pos_limit,
      const Eigen::MatrixBase<ConfigR_t> & upper_pos_limit,
      const Eigen::MatrixBase<ConfigOut_t> & qout) const
    {
      ConfigOut_t & res = PINOCCHIO_EIGEN_CONST_CAST(ConfigOut_t, qout).derived();
      for (int i = 0; i < nq(); ++i)
      {
        if (
          check_expression_if_real<Scalar, false>(
            lower_pos_limit[i] == -std::numeric_limits<typename ConfigL_t::Scalar>::infinity()
            || upper_pos_limit[i] == std::numeric_limits<typename ConfigR_t::Scalar>::infinity()))
        {
          std::ostringstream error;
          error << "non bounded limit. Cannot uniformly sample joint at rank " << i;
          throw std::range_error(error.str());
        }
        res[i] =
          lower_pos_limit[i] + ((upper_pos_limit[i] - lower_pos_limit[i]) * rand()) / RAND_MAX;
      }
    }

    bool isEqual_impl(const VectorSpaceOperationTpl & other) const
    {
      return size_.value() == other.size_.value();
    }

  private:
    Eigen::internal::variable_if_dynamic<Index, Dim> size_;
  }; // struct VectorSpaceOperationTpl

} // namespace pinocchio
