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

#ifndef __pinocchio_multibody_liegroup_cartesian_product_operation_hpp__
#define __pinocchio_multibody_liegroup_cartesian_product_operation_hpp__

#include <pinocchio/multibody/liegroup/liegroup-base.hpp>

namespace pinocchio
{
  template<int dim1, int dim2>
  struct eval_set_dim
  {
    enum
    {
      value = dim1 + dim2
    };
  };

  template<int dim>
  struct eval_set_dim<dim, Eigen::Dynamic>
  {
    enum
    {
      value = Eigen::Dynamic
    };
  };

  template<int dim>
  struct eval_set_dim<Eigen::Dynamic, dim>
  {
    enum
    {
      value = Eigen::Dynamic
    };
  };

  template<typename LieGroup1, typename LieGroup2>
  struct CartesianProductOperation;

  template<typename LieGroup1, typename LieGroup2>
  struct traits<CartesianProductOperation<LieGroup1, LieGroup2>>
  {
    typedef typename traits<LieGroup1>::Scalar Scalar;
    enum
    {
      Options = traits<LieGroup1>::Options,
      NQ = eval_set_dim < LieGroup1::NQ,
      LieGroup2::NQ > ::value,
      NV = eval_set_dim < LieGroup1::NV,
      LieGroup2::NV > ::value
    };
  };

  template<typename LieGroup1, typename LieGroup2>
  struct CartesianProductOperation
  : public LieGroupBase<CartesianProductOperation<LieGroup1, LieGroup2>>
  {
    PINOCCHIO_LIE_GROUP_TPL_PUBLIC_INTERFACE(CartesianProductOperation);

    CartesianProductOperation()
    : lg1()
    , lg2()
    {
    }
    // Get dimension of Lie Group vector representation
    //
    // For instance, for SO(3), the dimension of the vector representation is
    // 4 (quaternion) while the dimension of the tangent space is 3.
    Index nq() const
    {
      return lg1.nq() + lg2.nq();
    }

    // Get dimension of Lie Group tangent space
    Index nv() const
    {
      return lg1.nv() + lg2.nv();
    }

    ConfigVector_t neutral() const
    {
      ConfigVector_t n;
      n.resize(nq());
      Qo1(n) = lg1.neutral();
      Qo2(n) = lg2.neutral();
      return n;
    }

    std::string name() const
    {
      std::ostringstream oss;
      oss << lg1.name() << "*" << lg2.name();
      return oss.str();
    }

    template<class ConfigL_t, class ConfigR_t, class Tangent_t>
    void difference_impl(
      const Eigen::MatrixBase<ConfigL_t> & q0,
      const Eigen::MatrixBase<ConfigR_t> & q1,
      const Eigen::MatrixBase<Tangent_t> & d) const
    {
      lg1.difference(Q1(q0), Q1(q1), Vo1(d));
      lg2.difference(Q2(q0), Q2(q1), Vo2(d));
    }

    template<ArgumentPosition arg, class ConfigL_t, class ConfigR_t, class JacobianOut_t>
    void dDifference_impl(
      const Eigen::MatrixBase<ConfigL_t> & q0,
      const Eigen::MatrixBase<ConfigR_t> & q1,
      const Eigen::MatrixBase<JacobianOut_t> & J) const
    {
      J12(J).setZero();
      J21(J).setZero();

      lg1.template dDifference<arg>(Q1(q0), Q1(q1), J11(J));
      lg2.template dDifference<arg>(Q2(q0), Q2(q1), J22(J));
    }

    template<class ConfigIn_t, class Velocity_t, class ConfigOut_t>
    void integrate_impl(
      const Eigen::MatrixBase<ConfigIn_t> & q,
      const Eigen::MatrixBase<Velocity_t> & v,
      const Eigen::MatrixBase<ConfigOut_t> & qout) const
    {
      lg1.integrate(Q1(q), V1(v), Qo1(qout));
      lg2.integrate(Q2(q), V2(v), Qo2(qout));
    }

    template<class Config_t, class Jacobian_t>
    void integrateCoeffWiseJacobian_impl(
      const Eigen::MatrixBase<Config_t> & q, const Eigen::MatrixBase<Jacobian_t> & J) const
    {
      assert(J.rows() == nq() && J.cols() == nv() && "J is not of the right dimension");
      Jacobian_t & J_ = PINOCCHIO_EIGEN_CONST_CAST(Jacobian_t, J);
      J_.topRightCorner(lg1.nq(), lg2.nv()).setZero();
      J_.bottomLeftCorner(lg2.nq(), lg1.nv()).setZero();

      lg1.integrateCoeffWiseJacobian(Q1(q), J_.topLeftCorner(lg1.nq(), lg1.nv()));
      lg2.integrateCoeffWiseJacobian(Q2(q), J_.bottomRightCorner(lg2.nq(), lg2.nv()));
    }

    template<class Config_t, class Tangent_t, class JacobianOut_t>
    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) const
    {
      switch (op)
      {
      case SETTO:
        J12(J).setZero();
        J21(J).setZero();
        // fallthrough
      case ADDTO:
      case RMTO:
        lg1.dIntegrate_dq(Q1(q), V1(v), J11(J), op);
        lg2.dIntegrate_dq(Q2(q), V2(v), J22(J), op);
        break;
      default:
        assert(false && "Wrong Op requesed value");
        break;
      }
    }

    template<class Config_t, class Tangent_t, class JacobianOut_t>
    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) const
    {
      switch (op)
      {
      case SETTO:
        J12(J).setZero();
        J21(J).setZero();
        // fallthrough
      case ADDTO:
      case RMTO:
        lg1.dIntegrate_dv(Q1(q), V1(v), J11(J), op);
        lg2.dIntegrate_dv(Q2(q), V2(v), J22(J), op);
        break;
      default:
        assert(false && "Wrong Op requesed value");
        break;
      }
    }

    template<class Config_t, class Tangent_t, class JacobianIn_t, class JacobianOut_t>
    void dIntegrateTransport_dq_impl(
      const Eigen::MatrixBase<Config_t> & q,
      const Eigen::MatrixBase<Tangent_t> & v,
      const Eigen::MatrixBase<JacobianIn_t> & J_in,
      const Eigen::MatrixBase<JacobianOut_t> & J_out) const
    {
      JacobianOut_t & Jout = PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J_out);
      JacobianOut_t & Jin = PINOCCHIO_EIGEN_CONST_CAST(JacobianIn_t, J_in);
      lg1.dIntegrateTransport_dq(
        Q1(q), V1(v), Jin.template topRows<LieGroup1::NV>(),
        Jout.template topRows<LieGroup1::NV>());
      lg2.dIntegrateTransport_dq(
        Q2(q), V2(v), Jin.template bottomRows<LieGroup2::NV>(),
        Jout.template bottomRows<LieGroup2::NV>());
    }

    template<class Config_t, class Tangent_t, class JacobianIn_t, class JacobianOut_t>
    void dIntegrateTransport_dv_impl(
      const Eigen::MatrixBase<Config_t> & q,
      const Eigen::MatrixBase<Tangent_t> & v,
      const Eigen::MatrixBase<JacobianIn_t> & J_in,
      const Eigen::MatrixBase<JacobianOut_t> & J_out) const
    {
      JacobianOut_t & Jout = PINOCCHIO_EIGEN_CONST_CAST(JacobianOut_t, J_out);
      JacobianOut_t & Jin = PINOCCHIO_EIGEN_CONST_CAST(JacobianIn_t, J_in);
      lg1.dIntegrateTransport_dv(
        Q1(q), V1(v), Jin.template topRows<LieGroup1::NV>(),
        Jout.template topRows<LieGroup1::NV>());
      lg2.dIntegrateTransport_dv(
        Q2(q), V2(v), Jin.template bottomRows<LieGroup2::NV>(),
        Jout.template bottomRows<LieGroup2::NV>());
    }

    template<class Config_t, class Tangent_t, class Jacobian_t>
    void dIntegrateTransport_dq_impl(
      const Eigen::MatrixBase<Config_t> & q,
      const Eigen::MatrixBase<Tangent_t> & v,
      const Eigen::MatrixBase<Jacobian_t> & Jin) const
    {
      Jacobian_t & J = PINOCCHIO_EIGEN_CONST_CAST(Jacobian_t, Jin);
      lg1.dIntegrateTransport_dq(Q1(q), V1(v), J.template topRows<LieGroup1::NV>());
      lg2.dIntegrateTransport_dq(Q2(q), V2(v), J.template bottomRows<LieGroup2::NV>());
    }

    template<class Config_t, class Tangent_t, class Jacobian_t>
    void dIntegrateTransport_dv_impl(
      const Eigen::MatrixBase<Config_t> & q,
      const Eigen::MatrixBase<Tangent_t> & v,
      const Eigen::MatrixBase<Jacobian_t> & Jin) const
    {
      Jacobian_t & J = PINOCCHIO_EIGEN_CONST_CAST(Jacobian_t, Jin);
      lg1.dIntegrateTransport_dv(Q1(q), V1(v), J.template topRows<LieGroup1::NV>());
      lg2.dIntegrateTransport_dv(Q2(q), V2(v), J.template bottomRows<LieGroup2::NV>());
    }

    template<class ConfigL_t, class ConfigR_t>
    Scalar squaredDistance_impl(
      const Eigen::MatrixBase<ConfigL_t> & q0, const Eigen::MatrixBase<ConfigR_t> & q1) const
    {
      return lg1.squaredDistance(Q1(q0), Q1(q1)) + lg2.squaredDistance(Q2(q0), Q2(q1));
    }

    template<class Config_t>
    void normalize_impl(const Eigen::MatrixBase<Config_t> & qout) const
    {
      lg1.normalize(Qo1(qout));
      lg2.normalize(Qo2(qout));
    }

    template<class Config_t>
    bool isNormalized_impl(const Eigen::MatrixBase<Config_t> & qin, const Scalar & prec) const
    {
      return lg1.isNormalized(Qo1(qin), prec) && lg2.isNormalized(Qo2(qin), prec);
    }

    template<class Config_t>
    void random_impl(const Eigen::MatrixBase<Config_t> & qout) const
    {
      lg1.random(Qo1(qout));
      lg2.random(Qo2(qout));
    }

    template<class ConfigL_t, class ConfigR_t, class ConfigOut_t>
    void randomConfiguration_impl(
      const Eigen::MatrixBase<ConfigL_t> & lower,
      const Eigen::MatrixBase<ConfigR_t> & upper,
      const Eigen::MatrixBase<ConfigOut_t> & qout) const
    {
      lg1.randomConfiguration(Q1(lower), Q1(upper), Qo1(qout));
      lg2.randomConfiguration(Q2(lower), Q2(upper), Qo2(qout));
    }

    template<class ConfigL_t, class ConfigR_t>
    bool isSameConfiguration_impl(
      const Eigen::MatrixBase<ConfigL_t> & q0,
      const Eigen::MatrixBase<ConfigR_t> & q1,
      const Scalar & prec) const
    {
      return lg1.isSameConfiguration(Q1(q0), Q1(q1), prec)
             && lg2.isSameConfiguration(Q2(q0), Q2(q1), prec);
    }

    bool isEqual_impl(const CartesianProductOperation & other) const
    {
      return lg1 == other.lg1 && lg2 == other.lg2;
    }

    LieGroup1 lg1;
    LieGroup2 lg2;

  private:
    // VectorSpaceOperationTpl<-1> within CartesianProductOperation will not work
    // if Eigen version is lower than 3.2.1
#if EIGEN_VERSION_AT_LEAST(3, 2, 1)
  #define REMOVE_IF_EIGEN_TOO_LOW(x) x
#else
  #define REMOVE_IF_EIGEN_TOO_LOW(x)
#endif

    template<typename Config>
    typename Config ::template ConstFixedSegmentReturnType<LieGroup1::NQ>::Type
    Q1(const Eigen::MatrixBase<Config> & q) const
    {
      return q.derived().template head<LieGroup1::NQ>(REMOVE_IF_EIGEN_TOO_LOW(lg1.nq()));
    }
    template<typename Config>
    typename Config ::template ConstFixedSegmentReturnType<LieGroup2::NQ>::Type
    Q2(const Eigen::MatrixBase<Config> & q) const
    {
      return q.derived().template tail<LieGroup2::NQ>(REMOVE_IF_EIGEN_TOO_LOW(lg2.nq()));
    }
    template<typename Tangent>
    typename Tangent::template ConstFixedSegmentReturnType<LieGroup1::NV>::Type
    V1(const Eigen::MatrixBase<Tangent> & v) const
    {
      return v.derived().template head<LieGroup1::NV>(REMOVE_IF_EIGEN_TOO_LOW(lg1.nv()));
    }
    template<typename Tangent>
    typename Tangent::template ConstFixedSegmentReturnType<LieGroup2::NV>::Type
    V2(const Eigen::MatrixBase<Tangent> & v) const
    {
      return v.derived().template tail<LieGroup2::NV>(REMOVE_IF_EIGEN_TOO_LOW(lg2.nv()));
    }

    template<typename Config>
    typename Config ::template FixedSegmentReturnType<LieGroup1::NQ>::Type
    Qo1(const Eigen::MatrixBase<Config> & q) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Config, q).template head<LieGroup1::NQ>(
        REMOVE_IF_EIGEN_TOO_LOW(lg1.nq()));
    }
    template<typename Config>
    typename Config ::template FixedSegmentReturnType<LieGroup2::NQ>::Type
    Qo2(const Eigen::MatrixBase<Config> & q) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Config, q).template tail<LieGroup2::NQ>(
        REMOVE_IF_EIGEN_TOO_LOW(lg2.nq()));
    }
    template<typename Tangent>
    typename Tangent::template FixedSegmentReturnType<LieGroup1::NV>::Type
    Vo1(const Eigen::MatrixBase<Tangent> & v) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Tangent, v)
        .template head<LieGroup1::NV>(REMOVE_IF_EIGEN_TOO_LOW(lg1.nv()));
    }
    template<typename Tangent>
    typename Tangent::template FixedSegmentReturnType<LieGroup2::NV>::Type
    Vo2(const Eigen::MatrixBase<Tangent> & v) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Tangent, v)
        .template tail<LieGroup2::NV>(REMOVE_IF_EIGEN_TOO_LOW(lg2.nv()));
    }

    template<typename Jac>
    Eigen::Block<Jac, LieGroup1::NV, LieGroup1::NV> J11(const Eigen::MatrixBase<Jac> & J) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Jac, J)
        .template topLeftCorner<LieGroup1::NV, LieGroup1::NV>(lg1.nv(), lg1.nv());
    }
    template<typename Jac>
    Eigen::Block<Jac, LieGroup1::NV, LieGroup2::NV> J12(const Eigen::MatrixBase<Jac> & J) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Jac, J)
        .template topRightCorner<LieGroup1::NV, LieGroup2::NV>(lg1.nv(), lg2.nv());
    }
    template<typename Jac>
    Eigen::Block<Jac, LieGroup2::NV, LieGroup1::NV> J21(const Eigen::MatrixBase<Jac> & J) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Jac, J)
        .template bottomLeftCorner<LieGroup2::NV, LieGroup1::NV>(lg2.nv(), lg1.nv());
    }
    template<typename Jac>
    Eigen::Block<Jac, LieGroup2::NV, LieGroup2::NV> J22(const Eigen::MatrixBase<Jac> & J) const
    {
      return PINOCCHIO_EIGEN_CONST_CAST(Jac, J)
        .template bottomRightCorner<LieGroup2::NV, LieGroup2::NV>(lg2.nv(), lg2.nv());
    }
#undef REMOVE_IF_EIGEN_TOO_LOW

  }; // struct CartesianProductOperation

} // namespace pinocchio

#endif // ifndef __pinocchio_multibody_liegroup_cartesian_product_operation_hpp__
