// Copyright (C) 2005, 2009 International Business Machines and others.
// All Rights Reserved.
// This code is published under the Common Public License.
//
// $Id: IpLowRankSSAugSystemSolver.cpp 1692 2010-06-21 19:10:07Z andreasw $
//
// Authors:  Andreas Waechter                IBM    2009-11-05
//             (based on IpLowRankAugSystemSolver.cpp rev 1571)

#include "IpLowRankSSAugSystemSolver.hpp"
#include "IpLowRankUpdateSymMatrix.hpp"
#include "IpDenseVector.hpp"

namespace Ipopt
{
#if COIN_IPOPT_VERBOSITY > 0
  static const Index dbg_verbosity = 0;
#endif

  LowRankSSAugSystemSolver::LowRankSSAugSystemSolver(
    AugSystemSolver& aug_system_solver,
    Index max_rank)
      :
      AugSystemSolver(),
      aug_system_solver_(&aug_system_solver),
      max_rank_(max_rank),
      w_tag_(0),
      w_factor_(0.),
      d_x_tag_(0),
      delta_x_(0.),
      d_s_tag_(0),
      delta_s_(0.),
      j_c_tag_(0),
      d_c_tag_(0),
      delta_c_(0.),
      j_d_tag_(0),
      d_d_tag_(0),
      delta_d_(0.)
  {
    DBG_START_METH("LowRankSSAugSystemSolver::LowRankSSAugSystemSolver()",dbg_verbosity);
    DBG_ASSERT(IsValid(aug_system_solver_));
  }

  LowRankSSAugSystemSolver::~LowRankSSAugSystemSolver()
  {
    DBG_START_METH("LowRankSSAugSystemSolver::~LowRankSSAugSystemSolver()",dbg_verbosity);
  }

  bool LowRankSSAugSystemSolver::InitializeImpl(const OptionsList& options,
      const std::string& prefix)
  {
    first_call_ = true;
    Wdiag_ = NULL;
    expanded_vu_ = NULL;
    J_c_ext_ = NULL;
    D_c_ext_ = NULL;
    y_c_ext_space_ = NULL;

    return aug_system_solver_->Initialize(Jnlst(), IpNLP(), IpData(), IpCq(),
                                          options, prefix);
  }

  ESymSolverStatus LowRankSSAugSystemSolver::Solve(
    const SymMatrix* W,
    double W_factor,
    const Vector* D_x,
    double delta_x,
    const Vector* D_s,
    double delta_s,
    const Matrix* J_c,
    const Vector* D_c,
    double delta_c,
    const Matrix* J_d,
    const Vector* D_d,
    double delta_d,
    const Vector& rhs_x,
    const Vector& rhs_s,
    const Vector& rhs_c,
    const Vector& rhs_d,
    Vector& sol_x,
    Vector& sol_s,
    Vector& sol_c,
    Vector& sol_d,
    bool check_NegEVals,
    Index numberOfNegEVals)
  {
    DBG_START_METH("LowRankSSAugSystemSolver::Solve",dbg_verbosity);

    ESymSolverStatus retval;

    if (first_call_) {
      DBG_ASSERT(IsNull(Wdiag_));
      // Set up the diagonal matrix Wdiag_
      Index dimx = rhs_x.Dim();
      SmartPtr<DiagMatrixSpace> Wdiag_space = new DiagMatrixSpace(dimx);
      Wdiag_ = Wdiag_space->MakeNewDiagMatrix();
    }

    // This might be used with a linear solver that cannot detect the
    // inertia.  In that case, we should not asked for checking the
    // number of negative eigenvalues.
    if (!aug_system_solver_->ProvidesInertia()) {
      check_NegEVals = false;
    }

    if (first_call_ ||
        AugmentedSystemRequiresChange(W, W_factor, D_x, delta_x, D_s, delta_s,
                                      *J_c, D_c, delta_c, *J_d, D_d,
                                      delta_d) ) {
      retval = UpdateExtendedData(W, W_factor, D_x, delta_x, D_s, delta_s,
                                  *J_c, D_c, delta_c, *J_d, D_d, delta_d,
                                  rhs_x, rhs_s, rhs_c, rhs_d);
      if (retval != SYMSOLVER_SUCCESS) {
        return retval;
      }

      // Store the tags
      w_tag_ = W->GetTag();
      w_factor_ = W_factor;
      if (D_x) {
        d_x_tag_ = D_x->GetTag();
      }
      else {
        d_x_tag_ = 0;
      }
      delta_x_ = delta_x;
      if (D_s) {
        d_s_tag_ = D_s->GetTag();
      }
      else {
        d_s_tag_ = 0;
      }
      delta_s_ = delta_s;
      if (J_c) {
        j_c_tag_ = J_c->GetTag();
      }
      else {
        j_c_tag_ = 0;
      }
      if (D_c) {
        d_c_tag_ = D_c->GetTag();
      }
      else {
        d_c_tag_ = 0;
      }
      delta_c_ = delta_c;
      if (J_d) {
        j_d_tag_ = J_d->GetTag();
      }
      else {
        j_d_tag_ = 0;
      }
      if (D_d) {
        d_d_tag_ = D_d->GetTag();
      }
      else {
        d_d_tag_ = 0;
      }
      delta_d_ = delta_d;

      first_call_ = false;
    }

    // Extend the right hand side
    SmartPtr<CompoundVector> rhs_c_ext =
      y_c_ext_space_->MakeNewCompoundVector(true);
    rhs_c_ext->SetComp(0, rhs_c);
    rhs_c_ext->GetCompNonConst(1)->Set(0.);
    SmartPtr<CompoundVector> sol_c_ext =
      y_c_ext_space_->MakeNewCompoundVector(true);
    sol_c_ext->SetCompNonConst(0, sol_c);

    // Now solve the system for the given right hand side, using the
    // extended Jacobian_c and y_c data.
    numberOfNegEVals += negEvalsCorrection_;
    retval = aug_system_solver_->Solve(GetRawPtr(Wdiag_), W_factor,
                                       D_x, delta_x, D_s, delta_s,
                                       GetRawPtr(J_c_ext_),
                                       GetRawPtr(D_c_ext_), delta_c,
                                       J_d, D_d, delta_d,
                                       rhs_x, rhs_s, *rhs_c_ext, rhs_d,
                                       sol_x, sol_s, *sol_c_ext, sol_d,
                                       check_NegEVals, numberOfNegEVals);
    if (aug_system_solver_->ProvidesInertia()) {
      num_neg_evals_ =
        aug_system_solver_->NumberOfNegEVals() - negEvalsCorrection_;
    }
    if (retval != SYMSOLVER_SUCCESS) {
      Jnlst().Printf(J_DETAILED, J_SOLVE_PD_SYSTEM,
                     "LowRankSSAugSystemSolver: AugSystemSolver returned retval = %d for right hand side.\n", retval);
      return retval;
    }

    return retval;
  }

  ESymSolverStatus LowRankSSAugSystemSolver::UpdateExtendedData(
    const SymMatrix* W,
    double W_factor,
    const Vector* D_x,
    double delta_x,
    const Vector* D_s,
    double delta_s,
    const Matrix& J_c,
    const Vector* D_c,
    double delta_c,
    const Matrix& J_d,
    const Vector* D_d,
    double delta_d,
    const Vector& proto_rhs_x,
    const Vector& proto_rhs_s,
    const Vector& proto_rhs_c,
    const Vector& proto_rhs_d)
  {
    DBG_START_METH("LowRankSSAugSystemSolver::UpdateExtendedData",
                   dbg_verbosity);

    DBG_ASSERT(W_factor == 0.0 || W_factor == 1.0);
    ESymSolverStatus retval = SYMSOLVER_SUCCESS;

    // Get the low update information out of W
    const LowRankUpdateSymMatrix* LR_W =
      static_cast<const LowRankUpdateSymMatrix*> (W);
    DBG_ASSERT(dynamic_cast<const LowRankUpdateSymMatrix*>(W));
    DBG_PRINT_MATRIX(2, "LR_W", *LR_W);

    // If we don't have it yet, create the ExpandedMultiVectorMatrix
    SmartPtr<const Matrix> P_LM = LR_W->P_LowRank();
    SmartPtr<const VectorSpace> LR_VecSpace = LR_W->LowRankVectorSpace();
    if (IsNull(expanded_vu_)) {
      SmartPtr<const ExpansionMatrix> exp_matrix;
      if (IsValid(P_LM)) {
        exp_matrix = static_cast<const ExpansionMatrix*>(GetRawPtr(P_LM));
        DBG_ASSERT(dynamic_cast<const ExpansionMatrix*>(GetRawPtr(P_LM)));
      }
      SmartPtr<ExpandedMultiVectorMatrixSpace> expanded_vu_space =
        new ExpandedMultiVectorMatrixSpace(max_rank_, *LR_VecSpace, exp_matrix);
      expanded_vu_ = expanded_vu_space->MakeNewExpandedMultiVectorMatrix();

      // Create extended y_c quantities to include the V and U matrices
      DBG_ASSERT(IsNull(J_c_ext_));
      SmartPtr<CompoundMatrixSpace> J_c_ext_space =
        new CompoundMatrixSpace(2, 1, proto_rhs_c.Dim()+max_rank_,
                                proto_rhs_x.Dim());
      J_c_ext_space->SetBlockRows(0, proto_rhs_c.Dim());
      J_c_ext_space->SetBlockRows(1, max_rank_);
      J_c_ext_space->SetBlockCols(0, proto_rhs_x.Dim());
      J_c_ext_space->SetCompSpace(0, 0, *J_c.OwnerSpace());
      J_c_ext_space->SetCompSpace(1, 0, *expanded_vu_space);

      J_c_ext_ = J_c_ext_space->MakeNewCompoundMatrix();

      DBG_ASSERT(IsNull(D_c_ext_));
      DBG_ASSERT(IsNull(y_c_ext_space_));
      y_c_ext_space_ = new CompoundVectorSpace(2, proto_rhs_c.Dim()+max_rank_);
      y_c_ext_space_->SetCompSpace(0, *proto_rhs_c.OwnerSpace());
      SmartPtr<DenseVectorSpace> D_c_rank_space =
        new DenseVectorSpace(max_rank_);
      y_c_ext_space_->SetCompSpace(1, *D_c_rank_space);
      D_c_ext_ = y_c_ext_space_->MakeNewCompoundVector(true);
    }

    SmartPtr<const Vector> B0;
    SmartPtr<const MultiVectorMatrix> V;
    SmartPtr<const MultiVectorMatrix> U;
    if (W_factor == 1.0) {
      V = LR_W->GetV();
      U = LR_W->GetU();
      B0 = LR_W->GetDiag();
    }

    if (IsNull(B0)) {
      SmartPtr<Vector> zero_B0 = (IsValid(P_LM)) ? LR_VecSpace->MakeNew() : proto_rhs_x.MakeNew();
      zero_B0->Set(0.0);
      B0 = GetRawPtr(zero_B0);
    }

    // set up the Hessian for the underlying augmented system solver
    // without the low-rank update
    if (IsValid(P_LM) && LR_W->ReducedDiag()) {
      DBG_ASSERT(IsValid(B0));
      SmartPtr<Vector> fullx = proto_rhs_x.MakeNew();
      P_LM->MultVector(1., *B0, 0., *fullx);
      Wdiag_->SetDiag(*fullx);
    }
    else {
      Wdiag_->SetDiag(*B0);
      DBG_PRINT_VECTOR(2, "B0", *B0);
    }

    SmartPtr<Vector> D_c_rank_vec = D_c_ext_->GetCompNonConst(1);
    SmartPtr<DenseVector> D_c_rank =
      static_cast<DenseVector*>(GetRawPtr(D_c_rank_vec));
    DBG_ASSERT(dynamic_cast<DenseVector*>(GetRawPtr(D_c_rank_vec)));
    Number* D_c_rank_vals = D_c_rank->Values();
    Index irank = 0;
    if (IsValid(V)) {
      Index nV = V->NCols();
      negEvalsCorrection_ = nV;
      ASSERT_EXCEPTION(irank + nV, INTERNAL_ABORT, "max_rank too small for V");
      for (Index i=0; i<nV; i++) {
        SmartPtr<const Vector> vec = V->GetVector(i);
        expanded_vu_->SetVector(irank, vec);
        D_c_rank_vals[irank] = -1.;
        irank++;
      }
    }
    else {
      negEvalsCorrection_ = 0;
    }
    if (IsValid(U)) {
      Index nU = U->NCols();
      ASSERT_EXCEPTION(irank + nU, INTERNAL_ABORT, "max_rank too small for V");
      for (Index i=0; i<nU; i++) {
        SmartPtr<const Vector> vec = U->GetVector(i);
        expanded_vu_->SetVector(irank, vec);
        D_c_rank_vals[irank] = 1.;
        irank++;
      }
    }
    for (; irank<max_rank_; irank++) {
      expanded_vu_->SetVector(irank, NULL);
      D_c_rank_vals[irank] = 1.;
    }
    if (D_c) {
      D_c_ext_->SetComp(0, *D_c);
    }
    else {
      SmartPtr<Vector> zero_c = proto_rhs_c.MakeNew();
      zero_c->Set(0.);
      D_c_ext_->SetComp(0, *zero_c);
    }
    J_c_ext_->SetComp(0, 0, J_c);
    J_c_ext_->SetComp(1, 0, *expanded_vu_);

    return retval;
  }

  bool LowRankSSAugSystemSolver::AugmentedSystemRequiresChange(
    const SymMatrix* W,
    double W_factor,
    const Vector* D_x,
    double delta_x,
    const Vector* D_s,
    double delta_s,
    const Matrix& J_c,
    const Vector* D_c,
    double delta_c,
    const Matrix& J_d,
    const Vector* D_d,
    double delta_d)
  {
    DBG_START_METH("LowRankSSAugSystemSolver::AugmentedSystemRequiresChange",
                   dbg_verbosity);

#if COIN_IPOPT_VERBOSITY > 0

    bool Wtest = (W && W->GetTag() != w_tag_);
    bool iWtest = (!W && w_tag_ != 0);
    bool wfactor_test = (W_factor != w_factor_);
    bool D_xtest = (D_x && D_x->GetTag() != d_x_tag_);
    bool iD_xtest = (!D_x && d_x_tag_ != 0);
    bool delta_xtest = (delta_x != delta_x_);
    bool D_stest = (D_s && D_s->GetTag() != d_s_tag_);
    bool iD_stest = (!D_s && d_s_tag_ != 0);
    bool delta_stest = (delta_s != delta_s_);
    bool J_ctest = (J_c.GetTag() != j_c_tag_);
    bool D_ctest = (D_c && D_c->GetTag() != d_c_tag_);
    bool iD_ctest = (!D_c && d_c_tag_ != 0);
    bool delta_ctest = (delta_c != delta_c_);
    bool J_dtest = (J_d.GetTag() != j_d_tag_);
    bool D_dtest = (D_d && D_d->GetTag() != d_d_tag_);
    bool iD_dtest = (!D_d && d_d_tag_ != 0);
    bool delta_dtest = (delta_d != delta_d_);
#endif

    DBG_PRINT((2,"Wtest = %d\n", Wtest));
    DBG_PRINT((2,"iWtest = %d\n", iWtest));
    DBG_PRINT((2,"wfactor_test = %d\n", wfactor_test));
    DBG_PRINT((2,"D_xtest = %d\n", D_xtest));
    DBG_PRINT((2,"iD_xtest = %d\n", iD_xtest));
    DBG_PRINT((2,"delta_xtest = %d\n", delta_xtest));
    DBG_PRINT((2,"D_stest = %d\n", D_stest));
    DBG_PRINT((2,"iD_stest = %d\n", iD_stest));
    DBG_PRINT((2,"delta_stest = %d\n", delta_stest));
    DBG_PRINT((2,"J_ctest = %d\n", J_ctest));
    DBG_PRINT((2,"D_ctest = %d\n", D_ctest));
    DBG_PRINT((2,"iD_ctest = %d\n", iD_ctest));
    DBG_PRINT((2,"delta_ctest = %d\n", delta_ctest));
    DBG_PRINT((2,"J_dtest = %d\n", J_dtest));
    DBG_PRINT((2,"D_dtest = %d\n", D_dtest));
    DBG_PRINT((2,"iD_dtest = %d\n", iD_dtest));
    DBG_PRINT((2,"delta_dtest = %d\n", delta_dtest));

    if ( (W && W->GetTag() != w_tag_)
         || (!W && w_tag_ != 0)
         || (W_factor != w_factor_)
         || (D_x && D_x->GetTag() != d_x_tag_)
         || (!D_x && d_x_tag_ != 0)
         || (delta_x != delta_x_)
         || (D_s && D_s->GetTag() != d_s_tag_)
         || (!D_s && d_s_tag_ != 0)
         || (delta_s != delta_s_)
         || (J_c.GetTag() != j_c_tag_)
         || (D_c && D_c->GetTag() != d_c_tag_)
         || (!D_c && d_c_tag_ != 0)
         || (delta_c != delta_c_)
         || (J_d.GetTag() != j_d_tag_)
         || (D_d && D_d->GetTag() != d_d_tag_)
         || (!D_d && d_d_tag_ != 0)
         || (delta_d != delta_d_) ) {
      return true;
    }

    return false;
  }

  Index LowRankSSAugSystemSolver::NumberOfNegEVals() const
  {
    DBG_ASSERT(!first_call_);
    return num_neg_evals_;
  }

  bool LowRankSSAugSystemSolver::ProvidesInertia() const
  {
    return aug_system_solver_->ProvidesInertia();
  }

  bool LowRankSSAugSystemSolver::IncreaseQuality()
  {
    return aug_system_solver_->IncreaseQuality();
  }

} // namespace Ipopt
