#pragma once

#include <span>
#include <type_traits>
#include <vector>

#include "drake/common/drake_assert.h"
#include "drake/common/drake_copyable.h"
#include "drake/common/eigen_types.h"

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

/* Contiguous storage for a pool of fixed-size Eigen objects.

This is essentially a wrapper around std::vector<EigenType>, where EigenType is
something like Matrix3d, Vector4d, etc.

All of the public types and methods are an implementation of the EigenPool API.
Refer to that class for documentation. */
template <typename EigenType>
class EigenPoolFixedSizeStorage {
 public:
  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(EigenPoolFixedSizeStorage);

  static_assert(EigenType::SizeAtCompileTime != Eigen::Dynamic);

  using Scalar = typename EigenType::Scalar;
  using ConstMatrixView = const EigenType&;
  using MatrixView = EigenType&;

  EigenPoolFixedSizeStorage();
  ~EigenPoolFixedSizeStorage();

  int size() const { return ssize(data_); }
  ConstMatrixView operator[](int i) const { return data_[i]; }
  MatrixView operator[](int i) { return data_[i]; }
  void Resize(int num_matrices, int rows, int cols);
  void Clear();
  void SetZero();
  void Add(int rows, int cols);

 private:
  // Contiguous storage for all Eigen matrices.
  std::vector<EigenType> data_;
};

/* Contiguous storage for a pool of dynamic-size Eigen objects.

Each matrix in the pool can have a different size, but they must have the same
type (e.g. MatrixXd, VectorXd, VectorX<AutoDiffXd>, etc).

All of the public types and methods are an implementation of the EigenPool API.
Refer to that class for documentation. */
template <typename EigenType>
class EigenPoolDynamicSizeStorage {
 public:
  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(EigenPoolDynamicSizeStorage);

  static_assert(EigenType::SizeAtCompileTime == Eigen::Dynamic);

  using Scalar = typename EigenType::Scalar;
  using ConstMatrixView = Eigen::Map<const EigenType>;
  using MatrixView = Eigen::Map<EigenType>;

  EigenPoolDynamicSizeStorage();
  ~EigenPoolDynamicSizeStorage();

  int size() const { return blocks_.size(); }
  ConstMatrixView operator[](int i) const {
    return ConstMatrixView(&data_[blocks_[i].index], blocks_[i].rows,
                           blocks_[i].cols);
  }
  MatrixView operator[](int i) {
    return MatrixView(&data_[blocks_[i].index], blocks_[i].rows,
                      blocks_[i].cols);
  }
  void Resize(int num_matrices, int rows, int cols);
  /* Slight API change vs the outer class: if either RowsAtCompileTime or
  ColsAtCompileTime is not Eigen::Dynamic, then the respective `rows` or `cols`
  argument must be empty (vs a matching `int` in the outer class). */
  void Resize(int num_matrices, std::span<const int> rows,
              std::span<const int> cols);
  void SetZero();
  void Clear();
  void Add(int rows, int cols);

 private:
  struct MatrixData {
    int index{0};  // Offset into data_.
    int rows{0};   // Number of rows.
    int cols{0};   // Number of columns.
  };

  // Contiguous storage for all Eigen scalars.
  std::vector<Scalar> data_;

  // Properly sized maps to each matrix.
  std::vector<MatrixData> blocks_;
};

/* A replacement for std::vector<Eigen::Matrix<...>> with a contiguous memory
layout.

TODO(#23770) The pool puts all of the scalars in *this* pool next to each other,
but we can do even better by putting the meta-data std::vector alongside the
scalars and/or sharing the allocation across multiple pools.

@tparam EigenType The type of Eigen::Matrix being stored; valid choices are:
- Matrix3<T>
- Matrix3X<T>
- Matrix6<T>
- Matrix6X<T>
- MatrixX<T>
- Vector3<T>
- Vector6<T>
- VectorX<T>
where T is either `double` or `AutoDiffXd`. */
template <typename EigenType>
class EigenPool {
 public:
  DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(EigenPool);

  static_assert(is_eigen_type<EigenType>::value);

  /* The underlying scalar type of the Eigen::Matrix in this pool. */
  using Scalar = typename EigenType::Scalar;

  /* The private implementation class. */
  using Storage =
      std::conditional_t<EigenType::SizeAtCompileTime == Eigen::Dynamic,
                         EigenPoolDynamicSizeStorage<EigenType>,
                         EigenPoolFixedSizeStorage<EigenType>>;

  /* The return type of const operator[]. For fixed-size storage, this is a C++
  reference to the EigenType. For dynamic-size stoage, this is an Eigen::Map. */
  using ConstMatrixView = typename Storage::ConstMatrixView;

  /* The return type of operator[]. For fixed-size storage, this is a C++
  reference to the EigenType. For dynamic-size stoage, this is an Eigen::Map. */
  using MatrixView = typename Storage::MatrixView;

  /* Constructs an empty pool. */
  EigenPool();

  ~EigenPool();

  /* Returns the number of matrices in this pool. */
  int size() const { return ssize(storage_); }

  /* Returns a const reference (or Eigen::Map) to the i-th matrix.
  @pre 0 <= i < size() */
  const ConstMatrixView operator[](int i) const {
    DRAKE_ASSERT(0 <= i && i < size());
    return storage_[i];
  }

  /* Returns a mutable reference (or Eigen::Map) to the i-th matrix.
  @pre 0 <= i < size() */
  MatrixView operator[](int i) {
    DRAKE_ASSERT(0 <= i && i < size());
    return storage_[i];
  }

  /* @name Re-sizing

  Resize methods allocate memory only if the current capacity is insufficient,
  and do not reduce the capacity. Any newly allocated coefficients are not
  guaranteed to be initialized.

  Note that even though Resize() is not a cheap / inline operation, these outer
  class wrappers are still defined inline so that the DRAKE_DEMAND checks can
  be elided by the compiler. The Storage methods that these delegate to are
  properly non-inline, so having these forwarding functions defined inline is
  not a big problem.

  @{ */

  /* Resizes a pool using a homogeneous size for all matrices. If either or both
  of RowsAtCompileTime or ColsAtCompileTime are not Eigen::Dynamic, then the
  respective `rows` or `cols` argument must match the compile-time size. */
  void Resize(int num_matrices, int rows, int cols) {
    DRAKE_DEMAND(EigenType::RowsAtCompileTime == Eigen::Dynamic ||
                 rows == EigenType::RowsAtCompileTime);
    DRAKE_DEMAND(EigenType::ColsAtCompileTime == Eigen::Dynamic ||
                 cols == EigenType::ColsAtCompileTime);
    storage_.Resize(num_matrices, rows, cols);
  }

  /* Resizes a pool using heterogenous row sizes for the matrices but a fixed
  size for `cols` that must match ColsAtCompileTime. The default `cols = 1` is
  useful when `EigenType = Eigen::VectorX<T>`, i.e., column vectors. */
  void Resize(int num_matrices, std::span<const int> rows, int cols = 1)
    requires(EigenType::RowsAtCompileTime == Eigen::Dynamic &&
             EigenType::ColsAtCompileTime >= 0)
  {
    DRAKE_DEMAND(cols == EigenType::ColsAtCompileTime);
    storage_.Resize(num_matrices, rows, {});
  }

  /* Resizes a pool using heterogenous column sizes for the matrices but a fixed
  size for `rows` that must match RowsAtCompileTime. */
  void Resize(int num_matrices, int rows, std::span<const int> cols)
    requires(EigenType::RowsAtCompileTime >= 0 &&
             EigenType::ColsAtCompileTime == Eigen::Dynamic)
  {
    DRAKE_DEMAND(rows == EigenType::RowsAtCompileTime);
    storage_.Resize(num_matrices, {}, cols);
  }

  /* Resizes a pool using heterogenous sizes for the matrices. Both
  RowsAtCompileTime and ColsAtCompileTime must be Eigen::Dynamic. */
  void Resize(int num_matrices, std::span<const int> rows,
              std::span<const int> cols)
    requires(EigenType::RowsAtCompileTime == Eigen::Dynamic &&
             EigenType::ColsAtCompileTime == Eigen::Dynamic)
  {
    storage_.Resize(num_matrices, rows, cols);
  }

  /* @} */

  /* Sets this pool's size() to zero. Capacity is not changed, and thus memory
  is not freed. */
  void Clear() { storage_.Clear(); }

  /* Zeroes out all matrices in this pool. */
  void SetZero() { storage_.SetZero(); }

  /* Appends a new matrix of the specified size to the end of the pool. If
  either or both of RowsAtCompileTime or ColsAtCompileTime are not
  Eigen::Dynamic, then the respective `rows` or `cols` argument must match the
  compile-time size. Coefficients are not initialized. */
  void Add(int rows, int cols) {
    DRAKE_DEMAND(EigenType::RowsAtCompileTime == Eigen::Dynamic ||
                 rows == EigenType::RowsAtCompileTime);
    DRAKE_DEMAND(EigenType::ColsAtCompileTime == Eigen::Dynamic ||
                 cols == EigenType::ColsAtCompileTime);
    storage_.Add(rows, cols);
  }

 private:
  /* The underlying data, one of EigenPool{Fixed,Dynamic}SizeStorage. */
  Storage storage_;
};

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