#pragma once

#include <concepts>
#include <utility>
#include <vector>

#include "drake/common/drake_assert.h"
#include "drake/common/drake_copyable.h"
#include "drake/common/eigen_types.h"
#include "drake/multibody/contact_solvers/block_sparse_lower_triangular_or_symmetric_matrix.h"
#include "drake/multibody/contact_solvers/icf/eigen_pool.h"
#include "drake/multibody/contact_solvers/icf/holonomic_constraints_data_pool.h"
#include "drake/multibody/contact_solvers/icf/icf_data.h"
#include "drake/multibody/contact_solvers/icf/reduced_mapping.h"

namespace drake {
namespace multibody {
namespace contact_solvers {
namespace icf {
namespace internal {

// Forward declaration to break circular dependencies.
template <typename T>
class IcfModel;

/* Shifts a spatial impulse applied to a rigid body from a point B to a point A,
both fixed to the same body. Given the spatial impulse Γ_B applied at B and the
position p_AB of B relative to A, returns the equivalent spatial impulse Γ_A at
A:
  Γ_A = Φ(p_AB)⋅Γ_B,
where Φ(p) is the rigid shift operator:
   Φ(p) =
   | I₃  pₓ |
   | 0   I₃ |

All quantities must be expressed in the same frame.

Γ is represented in code as "Gamma". */
template <typename T>
Vector6<T> ShiftSpatialImpulse(const Vector6<T>& Gamma_B,
                               const Vector3<T>& p_AB) {
  Vector6<T> Gamma_A;
  Gamma_A.template head<3>() =
      Gamma_B.template head<3>() + p_AB.cross(Gamma_B.template tail<3>());
  Gamma_A.template tail<3>() = Gamma_B.template tail<3>();
  return Gamma_A;
}

/* CRTP base class implementing the shared machinery for a pool of holonomic
constraints between pairs of bodies, A and B. Weld, ball, and distance
constraints are all holonomic; each has a convex cost:
    ℓ(vc) = ½(v̂ − vc)ᵀ⋅R⁻¹⋅(v̂ − vc),
where vc ∈ ℝᴺ is the constraint velocity, v̂ a bias velocity from the
constraint function g₀, and R a (near-rigid or compliant) regularization. The
impulse is γ = R⁻¹⋅(v̂ − vc), applied to body B and -γ to body A.

This base owns all the constraint-agnostic machinery: the regularization R, the
bias v̂, the cost, gradient, block Hessian, the sparsity pattern, and model
reduction.

Each concrete pool supplies a set of CRTP "hooks" providing the
structure-specific kinematics. They are declared and documented as deleted
member functions in the "CRTP hooks" section of this class (each concrete pool
shadows them with a real implementation).

@tparam_nonsymbolic_scalar
@tparam N The number of constraint equations.
@tparam Derived The concrete constraint pool. */
template <typename T, int N, typename Derived>
class HolonomicConstraintsPool {
 public:
  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(HolonomicConstraintsPool);

  ~HolonomicConstraintsPool();

  using ConstraintVector = Eigen::Matrix<T, N, 1>;
  using DataPool = HolonomicConstraintsDataPool<T, N>;

  /* CRTP hooks. Each concrete pool provides the per-constraint hooks below,
  which this base invokes on the derived pool (see derived()). They are declared
  here as deleted member functions so their shared contract is documented in a
  single place; each concrete pool shadows them with a real implementation. */

  /* Computes the constraint velocity vc ∈ ℝᴺ of the k-th constraint from the
  body spatial velocities.
  @param k The index of the constraint within the pool.
  @param V_WB The spatial velocity of body B, expressed in the world frame.
  @param V_WA The spatial velocity of body A, expressed in the world frame, or
              null when body A is anchored.
  @returns The constraint velocity vc ∈ ℝᴺ. */
  ConstraintVector CalcConstraintVelocity(
      int k, const Vector6<T>& V_WB, const Vector6<T>* V_WA) const = delete;

  /* Computes the spatial impulses at the body origins produced by the
  generalized impulse `gamma` of the k-th constraint.
  @param k The index of the constraint within the pool.
  @param gamma The generalized impulse of the k-th constraint.
  @param[out] Gamma_Bo The spatial impulse on body B at its origin Bo.
  @param[out] Gamma_Ao The spatial impulse on body A at its origin Ao, or null
                       when body A is anchored. */
  void CalcSpatialImpulses(int k, const ConstraintVector& gamma,
                           Vector6<T>* Gamma_Bo,
                           Vector6<T>* Gamma_Ao) const = delete;

  /* Computes the body-space Hessian blocks G_Xp = Φ(p_X)ᵀ⋅G⋅Φ(p_X) of the k-th
  constraint, with the constraint-space Hessian G = R⁻¹.
  @param k The index of the constraint within the pool.
  @param R_inv The inverse regularization R⁻¹ of the k-th constraint.
  @param[out] G_Bp The Hessian block for body B.
  @param[out] G_Ap The Hessian block for body A.
  @param[out] G_cross The Hessian block for the A/B cross term.
  @note G_Ap and G_cross are either both non-null (body A dynamic) or both null
  (body A anchored). */
  void CalcHessianBlocks(int k, const T& R_inv, Matrix6<T>* G_Bp,
                         Matrix6<T>* G_Ap, Matrix6<T>* G_cross) const = delete;

  /* Returns the derived pool's typed data pool stored within `data`.
  @param data The IcfData holding this pool's per-constraint data. */
  const DataPool& GetDataPool(const IcfData<T>& data) const = delete;

  /* (Re)sizes the derived pool's own geometry storage to hold the given number
  of constraints.
  @param num_constraints The number of constraints to store. */
  void ResizeGeometry(int num_constraints) = delete;

  /* Copies the k-th constraint's geometry into `reduced`, flipping the roles of
  bodies A and B when `flip` is true.
  @param reduced The reduced pool receiving the constraint's geometry.
  @param k The index of the constraint within this pool.
  @param flip Whether to flip the roles of bodies A and B. */
  void ReduceGeometryInto(Derived* reduced, int k, bool flip) const = delete;

  /* Whether the constraint function g₀ negates when a reduced-model constraint
  flips A/B (see ReduceInto). This is true when g₀ is a "relative" quantity that
  reverses under swapping the two constraint points P and Q — e.g. the weld's
  (a_PQ, p_PoQo) or the ball's p_PQ. A concrete pool whose g₀ is invariant under
  the swap (e.g. the distance's scalar d − ℓ) hides this with a `false` value.
  The constraint velocity vc is invariant under the flip either way (the derived
  flips its geometry to keep it so), so g₀ and vc must transform consistently
  for the bias v̂ = −g₀/(δt + τ) to stay correct.

  Every concrete pool must explicitly shadow this function with the appropriate
  value (there is Reduce test coverage that exercises the flip). */
  static constexpr bool FlipNegatesG0() = delete;

  /* @see IsAbstractConstraintsPool. */
  const IcfModel<T>& model() const { return *model_; }
  int num_constraints() const { return ssize(body_pairs_); }
  void AccumulateGradient(const IcfData<T>& data, VectorX<T>* gradient) const;
  void AccumulateHessian(
      const IcfData<T>& data,
      contact_solvers::internal::BlockSparseSymmetricMatrix<MatrixX<T>>*
          hessian) const;
  void ReduceInto(const ReducedMapping& mapping, Derived* reduced_pool) const;

  /* Resizes the pool (generic data + the derived's geometry) to store the given
  number of constraints. Constraints hold invalid data until Set(). */
  void Resize(int num_constraints);

  /* Computes the sparsity pattern for the pool. Clique i is connected to
  clique j > i iff sparsity[i] contains j. */
  void CalcSparsityPattern(std::vector<std::vector<int>>* sparsity) const;

  /* Precomputes the iteration-invariant regularization R, bias v̂, and Hessian
  blocks. Must be called after all Set() calls and whenever the time step
  changes (R and v̂ depend on δt). */
  void PrecomputeHessianBlocks();

  /* Computes problem data (impulses γ and cost) from the body spatial
  velocities V_WB. */
  void CalcData(const EigenPool<Vector6<T>>& V_WB, DataPool* data) const;

  /* Computes the first and second derivatives of the cost ℓ̃(α) = ℓ(v + α⋅w).
  @param data Constraint data computed at v + α⋅w.
  @param U_WB Body spatial velocities evaluated at v + α⋅w. */
  void CalcCostAlongLine(const DataPool& data,
                         const EigenPool<Vector6<T>>& U_WB, T* dcost,
                         T* d2cost) const;

  /* Testing-only access to the generic per-constraint data. */
  const std::vector<std::pair<int, int>>& body_pairs() const {
    return body_pairs_;
  }
  const EigenPool<ConstraintVector>& g0() const { return g0_; }
  const std::vector<T>& stiffness() const { return stiffness_; }
  const std::vector<T>& damping() const { return damping_; }
  const std::vector<T>& R() const { return R_; }
  const EigenPool<ConstraintVector>& v_hat() const { return v_hat_; }
  int hessian_blocks_size() const { return ssize(hessian_blocks_); }

 protected:
  explicit HolonomicConstraintsPool(const IcfModel<T>* parent_model);

  /* Resizes only the generic base-owned arrays (not the derived's geometry). */
  void ResizeCommon(int num_constraints);

  /* Stores the generic per-constraint data. Called by the derived's Set() after
  it stores its geometry.
  @pre index in range, bodyB not anchored.
  @pre stiffness > 0
  @pre damping >= 0 */
  void SetCommon(int index, int bodyA, int bodyB, const ConstraintVector& g0,
                 const T& stiffness, const T& damping);

  // Precomputed, iteration-invariant Hessian blocks. Rank-N per constraint.
  struct HessianBlock {
    int c_B{-1};         // Clique index for body B (always valid).
    int c_A{-1};         // Clique index for body A (-1 if anchored).
    MatrixX<T> H_BB;     // Diagonal block for body B's clique.
    MatrixX<T> H_AA;     // Diagonal block for body A's clique (if dynamic).
    MatrixX<T> H_cross;  // Off-diagonal (or same-clique cross) block.
    int cross_row{-1};
    int cross_col{-1};
    bool A_is_dynamic{false};
  };

  const IcfModel<T>* const model_;  // The parent model.

  // Body pairs (bodyA, bodyB); bodyB is always dynamic (not anchored).
  std::vector<std::pair<int, int>> body_pairs_;

  // Constraint function at the start of the step, g₀ ∈ ℝᴺ, per constraint.
  EigenPool<ConstraintVector> g0_;

  // User compliance per constraint (stiffness +∞ / damping 0 ⇒ near-rigid).
  std::vector<T> stiffness_;
  std::vector<T> damping_;

  // Scalar regularization R and bias v̂ per constraint, computed in
  // PrecomputeHessianBlocks() (depend on the time step, so recomputed on
  // UpdateTimeStep).
  std::vector<T> R_;
  EigenPool<ConstraintVector> v_hat_;

  std::vector<HessianBlock> hessian_blocks_;

 private:
  Derived& mutable_derived() { return *static_cast<Derived*>(this); }
  const Derived& derived() const { return *static_cast<const Derived*>(this); }
};

}  // namespace internal
}  // namespace icf
}  // namespace contact_solvers
}  // namespace multibody
}  // namespace drake
