/*
 * Copyright (c) 2011-2019, The DART development contributors
 * All rights reserved.
 *
 * The list of contributors can be found at:
 *   https://github.com/dartsim/dart/blob/master/LICENSE
 *
 * This file is provided under the following "BSD-style" License:
 *   Redistribution and use in source and binary forms, with or
 *   without modification, are permitted provided that the following
 *   conditions are met:
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * 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.
 *   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.
 */

#include "dart/constraint/PGSLCPSolver.hpp"

#ifndef NDEBUG
#  include <iomanip>
#  include <iostream>
#endif

#include "dart/external/odelcpsolver/lcp.h"

#include "dart/common/Console.hpp"
#include "dart/constraint/ConstrainedGroup.hpp"
#include "dart/constraint/ConstraintBase.hpp"
#include "dart/lcpsolver/Lemke.hpp"

namespace dart {
namespace constraint {

//==============================================================================
PGSLCPSolver::PGSLCPSolver(double _timestep) : LCPSolver(_timestep)
{
}

//==============================================================================
PGSLCPSolver::~PGSLCPSolver()
{
}

//==============================================================================
void PGSLCPSolver::solve(ConstrainedGroup* _group)
{
  // If there is no constraint, then just return true.
  std::size_t numConstraints = _group->getNumConstraints();
  if (numConstraints == 0)
    return;

  // Build LCP terms by aggregating them from constraints
  std::size_t n = _group->getTotalDimension();
  int nSkip = dPAD(n);
  double* A = new double[n * nSkip];
  double* x = new double[n];
  double* b = new double[n];
  double* w = new double[n];
  double* lo = new double[n];
  double* hi = new double[n];
  int* findex = new int[n];

  // Set w to 0 and findex to -1
#ifndef NDEBUG
  std::memset(A, 0.0, n * nSkip * sizeof(double));
#endif
  std::memset(w, 0.0, n * sizeof(double));
  std::memset(findex, -1, n * sizeof(int));

  // Compute offset indices
  std::size_t* offset = new std::size_t[n];
  offset[0] = 0;
  //  std::cout << "offset[" << 0 << "]: " << offset[0] << std::endl;
  for (std::size_t i = 1; i < numConstraints; ++i)
  {
    const ConstraintBasePtr& constraint = _group->getConstraint(i - 1);
    assert(constraint->getDimension() > 0);
    offset[i] = offset[i - 1] + constraint->getDimension();
    //    std::cout << "offset[" << i << "]: " << offset[i] << std::endl;
  }

  // For each constraint
  ConstraintInfo constInfo;
  constInfo.invTimeStep = 1.0 / mTimeStep;
  for (std::size_t i = 0; i < numConstraints; ++i)
  {
    const ConstraintBasePtr& constraint = _group->getConstraint(i);

    constInfo.x = x + offset[i];
    constInfo.lo = lo + offset[i];
    constInfo.hi = hi + offset[i];
    constInfo.b = b + offset[i];
    constInfo.findex = findex + offset[i];
    constInfo.w = w + offset[i];

    // Fill vectors: lo, hi, b, w
    constraint->getInformation(&constInfo);

    // Fill a matrix by impulse tests: A
    constraint->excite();
    for (std::size_t j = 0; j < constraint->getDimension(); ++j)
    {
      // Adjust findex for global index
      if (findex[offset[i] + j] >= 0)
        findex[offset[i] + j] += offset[i];

      // Apply impulse for mipulse test
      constraint->applyUnitImpulse(j);

      // Fill upper triangle blocks of A matrix
      int index = nSkip * (offset[i] + j) + offset[i];
      constraint->getVelocityChange(A + index, true);
      for (std::size_t k = i + 1; k < numConstraints; ++k)
      {
        index = nSkip * (offset[i] + j) + offset[k];
        _group->getConstraint(k)->getVelocityChange(A + index, false);
      }

      // Filling symmetric part of A matrix
      for (std::size_t k = 0; k < i; ++k)
      {
        for (std::size_t l = 0; l < _group->getConstraint(k)->getDimension();
             ++l)
        {
          int index1 = nSkip * (offset[i] + j) + offset[k] + l;
          int index2 = nSkip * (offset[k] + l) + offset[i] + j;

          A[index1] = A[index2];
        }
      }
    }

    assert(isSymmetric(
        n, A, offset[i], offset[i] + constraint->getDimension() - 1));

    constraint->unexcite();
  }

  assert(isSymmetric(n, A));

  // Print LCP formulation
  //  dtdbg << "Before solve:" << std::endl;
  //  print(n, A, x, lo, hi, b, w, findex);
  //  std::cout << std::endl;

  // Solve LCP using ODE's Dantzig algorithm
  //  dSolveLCP(n, A, x, b, w, 0, lo, hi, findex);
  PGSOption option;
  option.setDefault();
  solvePGS(n, nSkip, 0, A, x, b, lo, hi, findex, &option);

  // Print LCP formulation
  //  dtdbg << "After solve:" << std::endl;
  //  print(n, A, x, lo, hi, b, w, findex);
  //  std::cout << std::endl;

  // Apply constraint impulses
  for (std::size_t i = 0; i < numConstraints; ++i)
  {
    const ConstraintBasePtr& constraint = _group->getConstraint(i);
    constraint->applyImpulse(x + offset[i]);
    constraint->excite();
  }

  delete[] offset;

  delete[] A;
  delete[] x;
  delete[] b;
  delete[] w;
  delete[] lo;
  delete[] hi;
  delete[] findex;
}

//==============================================================================
#ifndef NDEBUG
bool PGSLCPSolver::isSymmetric(std::size_t _n, double* _A)
{
  std::size_t nSkip = dPAD(_n);
  for (std::size_t i = 0; i < _n; ++i)
  {
    for (std::size_t j = 0; j < _n; ++j)
    {
      if (std::abs(_A[nSkip * i + j] - _A[nSkip * j + i]) > 1e-6)
      {
        std::cout << "A: " << std::endl;
        for (std::size_t k = 0; k < _n; ++k)
        {
          for (std::size_t l = 0; l < nSkip; ++l)
          {
            std::cout << std::setprecision(4) << _A[k * nSkip + l] << " ";
          }
          std::cout << std::endl;
        }

        std::cout << "A(" << i << ", " << j << "): " << _A[nSkip * i + j]
                  << std::endl;
        std::cout << "A(" << j << ", " << i << "): " << _A[nSkip * j + i]
                  << std::endl;
        return false;
      }
    }
  }

  return true;
}

//==============================================================================
bool PGSLCPSolver::isSymmetric(
    std::size_t _n, double* _A, std::size_t _begin, std::size_t _end)
{
  std::size_t nSkip = dPAD(_n);
  for (std::size_t i = _begin; i <= _end; ++i)
  {
    for (std::size_t j = _begin; j <= _end; ++j)
    {
      if (std::abs(_A[nSkip * i + j] - _A[nSkip * j + i]) > 1e-6)
      {
        std::cout << "A: " << std::endl;
        for (std::size_t k = 0; k < _n; ++k)
        {
          for (std::size_t l = 0; l < nSkip; ++l)
          {
            std::cout << std::setprecision(4) << _A[k * nSkip + l] << " ";
          }
          std::cout << std::endl;
        }

        std::cout << "A(" << i << ", " << j << "): " << _A[nSkip * i + j]
                  << std::endl;
        std::cout << "A(" << j << ", " << i << "): " << _A[nSkip * j + i]
                  << std::endl;
        return false;
      }
    }
  }

  return true;
}

//==============================================================================
void PGSLCPSolver::print(
    std::size_t _n,
    double* _A,
    double* _x,
    double* /*lo*/,
    double* /*hi*/,
    double* b,
    double* w,
    int* findex)
{
  std::size_t nSkip = dPAD(_n);
  std::cout << "A: " << std::endl;
  for (std::size_t i = 0; i < _n; ++i)
  {
    for (std::size_t j = 0; j < nSkip; ++j)
    {
      std::cout << std::setprecision(4) << _A[i * nSkip + j] << " ";
    }
    std::cout << std::endl;
  }

  std::cout << "b: ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << std::setprecision(4) << b[i] << " ";
  }
  std::cout << std::endl;

  std::cout << "w: ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << w[i] << " ";
  }
  std::cout << std::endl;

  std::cout << "x: ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << _x[i] << " ";
  }
  std::cout << std::endl;

  //  std::cout << "lb: ";
  //  for (int i = 0; i < dim; ++i)
  //  {
  //    std::cout << lb[i] << " ";
  //  }
  //  std::cout << std::endl;

  //  std::cout << "ub: ";
  //  for (int i = 0; i < dim; ++i)
  //  {
  //    std::cout << ub[i] << " ";
  //  }
  //  std::cout << std::endl;

  std::cout << "frictionIndex: ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << findex[i] << " ";
  }
  std::cout << std::endl;

  double* Ax = new double[_n];

  for (std::size_t i = 0; i < _n; ++i)
  {
    Ax[i] = 0.0;
  }

  for (std::size_t i = 0; i < _n; ++i)
  {
    for (std::size_t j = 0; j < _n; ++j)
    {
      Ax[i] += _A[i * nSkip + j] * _x[j];
    }
  }

  std::cout << "Ax   : ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << Ax[i] << " ";
  }
  std::cout << std::endl;

  std::cout << "b + w: ";
  for (std::size_t i = 0; i < _n; ++i)
  {
    std::cout << b[i] + w[i] << " ";
  }
  std::cout << std::endl;

  delete[] Ax;
}
#endif

bool solvePGS(
    int n,
    int nskip,
    int /*nub*/,
    double* A,
    double* x,
    double* b,
    double* lo,
    double* hi,
    int* findex,
    PGSOption* option)
{
  // LDLT solver will work !!!
  // if (nub == n)
  //{
  //	return LDLTSolver(n,nskip,A,x,b)
  //}

  int i, j, iter, idx, n_new;
  bool sentinel;
  double old_x, new_x, hi_tmp, lo_tmp, dummy, ea;
  double* A_ptr;
  double one_minus_sor_w = 1.0 - (option->sor_w);

  //--- ORDERING & SCALING & INITIAL LOOP & Test
  int* order = new int[n];

  n_new = 0;
  sentinel = true;
  for (i = 0; i < n; i++)
  {
    // ORDERING
    if (A[nskip * i + i] < option->eps_div)
    {
      x[i] = 0.0;
      continue;
    }
    order[n_new++] = i;

    // INITIAL LOOP
    A_ptr = A + nskip * i;
    new_x = b[i];
    old_x = x[i];

    for (j = 0; j < i; j++)
      new_x -= A_ptr[j] * x[j];
    for (j = i + 1; j < n; j++)
      new_x -= A_ptr[j] * x[j];

    new_x = new_x / A[nskip * i + i];

    if (findex[i] >= 0) // friction index
    {
      hi_tmp = hi[i] * x[findex[i]];
      lo_tmp = -hi_tmp;

      if (new_x > hi_tmp)
        x[i] = hi_tmp;
      else if (new_x < lo_tmp)
        x[i] = lo_tmp;
      else
        x[i] = new_x;
    }
    else // no friction index
    {
      if (new_x > hi[i])
        x[i] = hi[i];
      else if (new_x < lo[i])
        x[i] = lo[i];
      else
        x[i] = new_x;
    }

    // TEST
    if (sentinel)
    {
      ea = std::abs(x[i] - old_x);
      if (ea > option->eps_res)
        sentinel = false;
    }
  }
  if (sentinel)
  {
    delete[] order;
    return true;
  }

  // SCALING
  for (i = 0; i < n_new; i++)
  {
    idx = order[i];

    dummy = 1.0 / A[nskip * idx + idx]; // diagonal element
    b[idx] *= dummy;
    for (j = 0; j < n; j++)
      A[nskip * idx + j] *= dummy;
  }

  //--- ITERATION LOOP
  for (iter = 1; iter < option->itermax; iter++)
  {
    //--- RANDOMLY_REORDER_CONSTRAINTS
#if LCP_PGS_RANDOMLY_REORDER_CONSTRAINTS
    if ((iter & 7) == 0)
    {
      int tmp, swapi;
      for (i = 1; i < n_new; i++)
      {
        tmp = order[i];
        swapi = dRandInt(i + 1);
        order[i] = order[swapi];
        order[swapi] = tmp;
      }
    }
#endif

    sentinel = true;

    //-- ONE LOOP
    for (i = 0; i < n_new; i++)
    {
      idx = order[i];

      A_ptr = A + nskip * idx;
      new_x = b[idx];
      old_x = x[idx];

      for (j = 0; j < idx; j++)
        new_x -= A_ptr[j] * x[j];
      for (j = idx + 1; j < n; j++)
        new_x -= A_ptr[j] * x[j];

      new_x = (option->sor_w * new_x) + (one_minus_sor_w * old_x);

      if (findex[idx] >= 0) // friction index
      {
        hi_tmp = hi[idx] * x[findex[idx]];
        lo_tmp = -hi_tmp;

        if (new_x > hi_tmp)
          x[idx] = hi_tmp;
        else if (new_x < lo_tmp)
          x[idx] = lo_tmp;
        else
          x[idx] = new_x;
      }
      else // no friction index
      {
        if (new_x > hi[idx])
          x[idx] = hi[idx];
        else if (new_x < lo[idx])
          x[idx] = lo[idx];
        else
          x[idx] = new_x;
      }

      if (sentinel && std::abs(x[idx]) > option->eps_div)
      {
        ea = std::abs((x[idx] - old_x) / x[idx]);
        if (ea > option->eps_ea)
          sentinel = false;
      }
    }

    if (sentinel)
      break;
  }
  delete[] order;
  return sentinel;
}

#define LCP_PGS_RANDOMLY_REORDER_CONSTRAINTS 1
#define LCP_PGS_OPTION_DEFAULT_ITERMAX 30
#define LCP_PGS_OPTION_DEFAULT_SOR_W 0.9
#define LCP_PGS_OPTION_DEFAULT_EPS_EA 1E-3
#define LCP_PGS_OPTION_DEFAULT_EPS_RESIDUAL 1E-6
#define LCP_PGS_OPTION_DEFAULT_EPS_DIVIDE 1E-9

void PGSOption::setDefault()
{
  itermax = LCP_PGS_OPTION_DEFAULT_ITERMAX;
  sor_w = LCP_PGS_OPTION_DEFAULT_SOR_W;
  eps_ea = LCP_PGS_OPTION_DEFAULT_EPS_EA;
  eps_res = LCP_PGS_OPTION_DEFAULT_EPS_RESIDUAL;
  eps_div = LCP_PGS_OPTION_DEFAULT_EPS_DIVIDE;
}

} // namespace constraint
} // namespace dart
