text: /*
   * Copyright (c) 2019, The Regents of the University of California (Regents).
   * All rights reserved.
   *
   * Redistribution and use in source and binary forms, with or without
   * modification, are permitted provided that the following conditions are
   * met:
   *
   *    1. Redistributions of source code must retain the above copyright
   *       notice, this list of conditions and the following disclaimer.
   *
   *    2. Redistributions in binary form must reproduce the above
   *       copyright notice, this list of conditions and the following
   *       disclaimer in the documentation and/or other materials provided
   *       with the distribution.
   *
   *    3. Neither the name of the copyright holder nor the names of its
   *       contributors may be used to endorse or promote products derived
   *       from this software without specific prior written permission.
   *
   * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS AS IS
   * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
   * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
   * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
   * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
   * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
   * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
   * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
   * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
   * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
   * POSSIBILITY OF SUCH DAMAGE.
   *
   * Please contact the author(s) of this library if you have any questions.
   * Authors: David Fridovich-Keil   ( dfk@eecs.berkeley.edu )
   */
  
  ///////////////////////////////////////////////////////////////////////////////
  //
  // Multi-player dynamical system comprised of several single player subsystems.
  //
  ///////////////////////////////////////////////////////////////////////////////
  
  #include <ilqgames/dynamics/concatenated_flat_system.h>
  #include <ilqgames/utils/linear_dynamics_approximation.h>
  #include <ilqgames/utils/types.h>
  
  #include <glog/logging.h>
  
  namespace ilqgames {
  
  ConcatenatedFlatSystem::ConcatenatedFlatSystem(
      const FlatSubsystemList& subsystems)
      : MultiPlayerFlatSystem(std::accumulate(
            subsystems.begin(), subsystems.end(), 0,
            [](Dimension total,
               const std::shared_ptr<SinglePlayerFlatSystem>& subsystem) {
              CHECK_NOTNULL(subsystem.get());
              return total + subsystem->XDim();
            })),
        subsystems_(subsystems) {
    // Populate subsystem start dimensions.
    subsystem_start_dims_.push_back(0);
    for (const auto& subsystem : subsystems_) {
      subsystem_start_dims_.push_back(subsystem_start_dims_.back() +
                                      subsystem->XDim());
    }
  }
  
  VectorXf ConcatenatedFlatSystem::Evaluate(
      const VectorXf& x, const std::vector<VectorXf>& us) const {
    CHECK_EQ(us.size(), NumPlayers());
  
    // Populate 'xdot' one subsystem at a time.
    VectorXf xdot(xdim_);
    Dimension dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      xdot.segment(dims_so_far, subsystem->XDim()) =
          subsystem->Evaluate(x.segment(dims_so_far, subsystem->XDim()), us[ii]);
      dims_so_far += subsystem->XDim();
    }
  
    return xdot;
  }
  
  void ConcatenatedFlatSystem::ComputeLinearizedSystem() const {
    // Populate a block-diagonal A, as well as Bs.
    LinearDynamicsApproximation linearization(*this);
  
    Dimension dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      const Dimension udim = subsystem->UDim();
      subsystem->LinearizedSystem(
          linearization.A.block(dims_so_far, dims_so_far, xdim, xdim),
          linearization.Bs[ii].block(dims_so_far, 0, xdim, udim));
  
      dims_so_far += xdim;
    }
  
    discrete_linear_system_.reset(new LinearDynamicsApproximation(linearization));
  
    // Reconstruct the continuous system.
    linearization.A -= MatrixXf::Identity(xdim_, xdim_);
    linearization.A /= time::kTimeStep;
    for (size_t ii = 0; ii < NumPlayers(); ii++)
      linearization.Bs[ii] /= time::kTimeStep;
    continuous_linear_system_.reset(
        new LinearDynamicsApproximation(linearization));
  }
  
  MatrixXf ConcatenatedFlatSystem::InverseDecouplingMatrix(
      const VectorXf& x) const {
    // Populate a block-diagonal inverse decoupling matrix M.
    MatrixXf M = MatrixXf::Zero(TotalUDim(), TotalUDim());
  
    Dimension x_dims_so_far = 0;
    Dimension u_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      const Dimension udim = subsystem->UDim();
      M.block(u_dims_so_far, u_dims_so_far, udim, udim) =
          subsystem->InverseDecouplingMatrix(x.segment(x_dims_so_far, xdim));
  
      x_dims_so_far += xdim;
      u_dims_so_far += udim;
    }
  
    return M;
  }
  
  VectorXf ConcatenatedFlatSystem::AffineTerm(const VectorXf& x) const {
    // Populate the affine term m for each subsystem.
    VectorXf m(TotalUDim());
  
    Dimension x_dims_so_far = 0;
    Dimension u_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      const Dimension udim = subsystem->UDim();
      m.segment(u_dims_so_far, udim) =
          subsystem->AffineTerm(x.segment(x_dims_so_far, xdim));
  
      x_dims_so_far += xdim;
      u_dims_so_far += udim;
    }
  
    return m;
  }
  
  std::vector<VectorXf> ConcatenatedFlatSystem::LinearizingControls(
      const VectorXf& x, const std::vector<VectorXf>& vs) const {
    std::vector<VectorXf> us(NumPlayers());
  
    Dimension x_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      us[ii] =
          subsystem->LinearizingControl(x.segment(x_dims_so_far, xdim), vs[ii]);
  
      x_dims_so_far += xdim;
    }
  
    return us;
  }
  
  VectorXf ConcatenatedFlatSystem::LinearizingControl(const VectorXf& x,
                                                      const VectorXf& v,
                                                      PlayerIndex player) const {
    return subsystems_[player]->LinearizingControl(
        x.segment(subsystem_start_dims_[player], SubsystemXDim(player)), v);
  }
  
  VectorXf ConcatenatedFlatSystem::ToLinearSystemState(const VectorXf& x) const {
    VectorXf xi(xdim_);
  
    Dimension x_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      xi.segment(x_dims_so_far, xdim) =
          subsystem->ToLinearSystemState(x.segment(x_dims_so_far, xdim));
  
      x_dims_so_far += xdim;
    }
  
    return xi;
  }
  
  VectorXf ConcatenatedFlatSystem::ToLinearSystemState(
      const VectorXf& x, PlayerIndex subsystem_idx) const {
    CHECK_LT(subsystem_idx, subsystems_.size());
    return subsystems_[subsystem_idx]->ToLinearSystemState(x.segment(
        subsystem_start_dims_[subsystem_idx], SubsystemXDim(subsystem_idx)));
  }
  
  VectorXf ConcatenatedFlatSystem::FromLinearSystemState(
      const VectorXf& xi) const {
    VectorXf x(xdim_);
  
    Dimension x_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      x.segment(x_dims_so_far, xdim) =
          subsystem->FromLinearSystemState(xi.segment(x_dims_so_far, xdim));
  
      x_dims_so_far += xdim;
    }
  
    return x;
  }
  
  VectorXf ConcatenatedFlatSystem::FromLinearSystemState(
      const VectorXf& xi, PlayerIndex subsystem_idx) const {
    CHECK_LT(subsystem_idx, subsystems_.size());
    return subsystems_[subsystem_idx]->FromLinearSystemState(xi.segment(
        subsystem_start_dims_[subsystem_idx], SubsystemXDim(subsystem_idx)));
  }
  
  VectorXf ConcatenatedFlatSystem::SubsystemStates(
      const VectorXf& x, PlayerIndex subsystem_idx) const {
    CHECK_LT(subsystem_idx, subsystems_.size());
    return x.segment(subsystem_start_dims_[subsystem_idx],
                     SubsystemXDim(subsystem_idx));
  }
  
  bool ConcatenatedFlatSystem::IsLinearSystemStateSingular(
      const VectorXf& xi) const {
    Dimension x_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem = subsystems_[ii];
      const Dimension xdim = subsystem->XDim();
      if (subsystem->IsLinearSystemStateSingular(xi.segment(x_dims_so_far, xdim)))
        return true;
      x_dims_so_far += xdim;
    }
  
    return false;
  }
  
  void ConcatenatedFlatSystem::ChangeCostCoordinates(
      const VectorXf& xi, std::vector<QuadraticCostApproximation>* q) const {
    CHECK_NOTNULL(q);
    CHECK_EQ(q->size(), NumPlayers());
  
    // For each player we record dx_i/dxi and d2x_i/dxi2.
    std::vector<std::vector<VectorXf>> first_partials(NumPlayers());
    std::vector<std::vector<MatrixXf>> second_partials(NumPlayers());
    Dimension xi_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem_ii = subsystems_[ii];
      const Dimension xdim = subsystem_ii->XDim();
  
      first_partials[ii] = std::vector<VectorXf>(xdim, VectorXf::Zero(xdim));
      second_partials[ii] =
          std::vector<MatrixXf>(xdim, MatrixXf::Zero(xdim, xdim));
  
      DCHECK_LT(xi_dims_so_far + xdim, xi.size());
      subsystem_ii->Partial(xi.segment(xi_dims_so_far, xdim), &first_partials[ii],
                            &second_partials[ii]);
      xi_dims_so_far += xdim;
    }
  
    // For loop for hessian.
    // Vector that is going to contain all the cost hessians for each player.
    std::vector<MatrixXf> hess_xs(NumPlayers(), MatrixXf::Zero(xdim_, xdim_));
    Dimension rows_so_far = 0;
  
    // Iterating over primary player indexes.
    for (PlayerIndex pp = 0; pp < NumPlayers(); pp++) {
      // Iterating over that player's number of dimensions.
      for (Dimension ii = 0; ii < SubsystemXDim(pp); ii++) {
        const Dimension ii_rows = ii + rows_so_far;
  
        // Iterating over secondary player indexes.
        Dimension cols_so_far = rows_so_far;
        for (PlayerIndex qq = pp; qq < NumPlayers(); qq++) {
          // Iterating over that player's number of dimensions.
          for (Dimension jj = 0; jj < SubsystemXDim(qq); jj++) {
            const Dimension jj_cols = jj + cols_so_far;
  
            // Iterating over each player's cost.
            for (PlayerIndex rr = 0; rr < NumPlayers(); rr++) {
              const MatrixXf& Q = (*q)[rr].state.hess;
              const VectorXf& l = (*q)[rr].state.grad;
              MatrixXf& xhess = hess_xs[rr];
              for (Dimension kk = 0; kk < SubsystemXDim(pp); kk++) {
                float total = 0.0;
  
                if (pp == qq) {
                  total += l(kk + rows_so_far) * second_partials[pp][kk](ii, jj);
                }
  
                float temp = 0.0;
                for (Dimension ll = 0; ll < SubsystemXDim(qq); ll++) {
                  temp += Q(kk + rows_so_far, ll + cols_so_far) *
                          first_partials[qq][ll](jj);
                }
  
                total += temp * first_partials[pp][kk](ii);
  
                DCHECK_LT(ii_rows, xdim_);
                DCHECK_LT(jj_cols, xdim_);
                xhess(ii_rows, jj_cols) += total;
                xhess(jj_cols, ii_rows) += total;
              }
            }
          }
  
          // Increment columns so far to track next subsystem.
          cols_so_far += SubsystemXDim(qq);
        }
      }
  
      // Increment rows so far to track next subsystem.
      rows_so_far += SubsystemXDim(pp);
    }
  
    // Update Qs to match these hessians.
    for (PlayerIndex pp = 0; pp < NumPlayers(); pp++) {
      //    std::cout << "pp: " << pp << "\n" << hess_xs[pp] << std::endl <<
      //    "-----------------\n";
      // if (pp == 1) {
      //   std::cout << "pp: " << pp << "\n"
      //             << hess_xs[pp] << std::endl
      //             << "-----------------\n";
      // }
      (*q)[pp].state.hess.swap(hess_xs[pp]);
    }
  
    // For loop for gradient.
    // Vector that is going to contain all the cost gradients for each player.
    std::vector<VectorXf> grad_xs(NumPlayers(), VectorXf::Zero(xdim_));
    rows_so_far = 0;
  
    // Iterating over primary player indexes.
    for (PlayerIndex pp = 0; pp < NumPlayers(); pp++) {
      const VectorXf& l = (*q)[pp].state.grad;
  
      // Iterating over that player's number of dimensions.
      for (Dimension ii = 0; ii < SubsystemXDim(pp); ii++) {
        // Iterating over each player's cost.
        for (Dimension jj = 0; jj < SubsystemXDim(pp); jj++) {
          DCHECK_LT(jj + rows_so_far, l.size());
          grad_xs[pp](ii + rows_so_far) +=
              l(jj + rows_so_far) * first_partials[pp][jj](ii);
        }
      }
  
      // Increment rows so far to track next subsystem.
      rows_so_far += SubsystemXDim(pp);
    }
  
    // Update ls to match these gradients.
    for (PlayerIndex pp = 0; pp < NumPlayers(); pp++) {
      // if (pp == 1) {
      //   std::cout << "pp: " << pp << "computed" <<  "\n"
      //             << grad_xs[pp] << std::endl
      //             << "-----------------\n";
      //   std::cout << "pp: " << pp << "original \n"
      //             << (*q)[pp].l << std::endl
      //             << "-----------------\n";
      // }
      // CHECK_LT(std::abs(grad_xs[pp](0) - (*q)[pp].l(0)), 1e-4);
      // CHECK_LT(std::abs(grad_xs[pp](1) - (*q)[pp].l(1)), 1e-4);
  
      (*q)[pp].state.grad.swap(grad_xs[pp]);
    }
  
    // Now modify the cost hessians, i.e. 'Rs'.
    // NOTE: this depends only on the decoupling matrix.
    ChangeControlCostCoordinates(xi, q);
  }
  
  void ConcatenatedFlatSystem::ChangeControlCostCoordinates(
      const VectorXf& xi, std::vector<QuadraticCostApproximation>* q) const {
    CHECK_NOTNULL(q);
  
    // Get nonlinear system state and all decoupling matrices.
    const VectorXf x = FromLinearSystemState(xi);
    std::vector<MatrixXf> M_invs(NumPlayers());
  
    Dimension x_dims_so_far = 0;
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      const auto& subsystem_ii = subsystems_[ii];
      const Dimension xdim = subsystem_ii->XDim();
      M_invs[ii] =
          subsystem_ii->InverseDecouplingMatrix(x.segment(x_dims_so_far, xdim));
      x_dims_so_far += xdim;
    }
  
    // Convert Rs.
    for (size_t ii = 0; ii < NumPlayers(); ii++) {
      for (auto& element : (*q)[ii].control) {
        element.second.hess = M_invs[element.first].transpose() *
                              element.second.hess * M_invs[element.first];
      }
    }
  }
  
  float ConcatenatedFlatSystem::DistanceBetween(const VectorXf& x0,
                                                const VectorXf& x1) const {
    Dimension dims_so_far = 0;
    float total = 0.0;
  
    // Accumulate total across all subsystems.
    for (const auto& subsystem : subsystems_) {
      const Dimension xdim = subsystem->XDim();
      total += subsystem->DistanceBetween(x0.segment(dims_so_far, xdim),
                                          x1.segment(dims_so_far, xdim));
  
      dims_so_far += xdim;
    }
  
    return total;
  }
  
  std::vector<Dimension> ConcatenatedFlatSystem::PositionDimensions() const {
    std::vector<Dimension> dims;
  
    for (const auto& s : subsystems_) {
      const std::vector<Dimension> sub_dims = s->PositionDimensions();
      dims.insert(dims.end(), sub_dims.begin(), sub_dims.end());
    }
  
    return dims;
  }
  
  }  // namespace ilqgames
  