#pragma once

#include <vector>

#include "drake/common/default_scalars.h"
#include "drake/common/drake_copyable.h"
#include "drake/common/eigen_types.h"
#include "drake/multibody/math/spatial_algebra.h"

namespace drake {
namespace multibody {

namespace internal {
template <typename T>
class MultibodyTree;
template <typename T>
class MultibodyTreeSystem;
}  // namespace internal

/// Holds a set of forces applied to a MultibodyTree system. Forces
/// can include generalized forces as well as mobilized body (mobod) spatial
/// forces. MultibodyPlant::CalcGeneralizedForces() can be used to compute
/// the _total_ generalized force, combining generalized_forces() with the
/// generalized force equivalent of the body_forces().
///
/// @tparam_default_scalar
template <typename T>
class MultibodyForces {
 public:
  DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(MultibodyForces);

  /// Constructs a force object to store a set of forces to be applied to
  /// the multibody model for `plant`. Forces are initialized to zero, meaning
  /// no forces are applied. `plant` must have been already finalized with
  /// MultibodyPlant::Finalize() or this constructor will abort.
  explicit MultibodyForces(const internal::MultibodyTreeSystem<T>& plant);

  /// (Advanced) Tree overload.
  explicit MultibodyForces(const internal::MultibodyTree<T>& model);

  /// (Advanced) Constructs a force object with a given number of mobilized
  /// bodies and corresponding mobilities (number of generalized velocities).
  /// This constructor may be useful for constructing the MultibodyForces
  /// structure before a MultibodyPlant has been constructed. However, this
  /// must be used cautiously since the number of mobilized bodies and their
  /// mobilities are not known until after Finalize().
  MultibodyForces(int nb, int nv);

  ~MultibodyForces();

  /// Sets `this` to store zero forces (no applied forces).
  MultibodyForces<T>& SetZero();

  /// Returns the number of bodies for which `this` force object applies.
  /// Determined at construction from the given model MultibodyTree object.
  int num_bodies() const { return static_cast<int>(F_B_W_.size()); }

  /// Returns the number of generalized velocities for the model to which these
  /// forces apply. The number of generalized forces in a multibody model always
  /// equals the number of generalized velocities.
  /// Determined at construction from the given model MultibodyTree object.
  int num_velocities() const { return static_cast<int>(tau_.size()); }

  /// (Advanced) Returns a constant reference to the vector of generalized
  /// forces stored by `this` forces object.
  const VectorX<T>& generalized_forces() const { return tau_; }

  /// (Advanced) Mutable version of generalized_forces().
  VectorX<T>& mutable_generalized_forces() { return tau_; }

  /// (Advanced) Returns a constant reference to the vector of spatial body
  /// forces `F_BBo_W` on each body B in the model, at the body's frame
  /// origin `Bo`, expressed in the world frame W.
  /// @note Entries are ordered by MobodIndex.
  const std::vector<SpatialForce<T>>& body_forces() const { return F_B_W_; }

  /// (Advanced) Mutable version of body_forces().
  std::vector<SpatialForce<T>>& mutable_body_forces() { return F_B_W_; }

  /// Adds into `this` the force contribution stored in `addend`.
  void AddInForces(const MultibodyForces<T>& addend);

  /// Utility that checks that the forces stored by `this` object have the
  /// proper sizes to represent the set of forces for the given `plant`.
  /// @returns true if `this` forces object has the proper sizes for the given
  /// `plant`.
  bool CheckHasRightSizeForModel(
      const internal::MultibodyTreeSystem<T>& plant) const;

  /// (Advanced) Tree overload.
  bool CheckHasRightSizeForModel(const internal::MultibodyTree<T>& model) const;

 private:
  // Vector holding, for each mobilized body (mobod) in the MultibodyTree, the
  // externally applied force F_Bi_W on the i-th mobod Bi, expressed in the
  // world frame W, and assumed to act at Bi's origin Bio. Indexed by
  // MobodIndex.
  std::vector<SpatialForce<T>> F_B_W_;

  // Vector of generalized forces applied on each mobilizer in the
  // MultibodyTree. Ordered identically to generalized velocities v.
  VectorX<T> tau_;
};

}  // namespace multibody
}  // namespace drake

DRAKE_DECLARE_CLASS_TEMPLATE_INSTANTIATIONS_ON_DEFAULT_SCALARS(
    class ::drake::multibody::MultibodyForces);
