// Copyright 2023, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "neml2/solvers/NonlinearSystem.h"
#include "neml2/misc/math.h"

namespace neml2
{
OptionSet
NonlinearSystem::expected_options()
{
  OptionSet options;
  options.set<bool>("automatic_scaling") = false;
  options.set<Real>("automatic_scaling_tol") = 0.01;
  options.set<unsigned int>("automatic_scaling_miter") = 20;
  return options;
}

NonlinearSystem::NonlinearSystem(const OptionSet & options)
  : _autoscale(options.get<bool>("automatic_scaling")),
    _autoscale_tol(options.get<Real>("automatic_scaling_tol")),
    _autoscale_miter(options.get<unsigned int>("automatic_scaling_miter")),
    _scaling_matrices_initialized(false)
{
}

void
NonlinearSystem::init_scaling(const bool verbose)
{
  if (!_autoscale)
    return;

  if (_scaling_matrices_initialized)
    return;

  using namespace torch::indexing;

  // First compute the Jacobian
  assemble(false, true);

  auto Jp = _Jacobian.clone();
  _row_scaling = BatchTensor::ones_like(_solution);
  _col_scaling = BatchTensor::ones_like(_solution);

  if (verbose)
    std::cout << "Before automatic scaling cond(J) = " << std::scientific
              << torch::max(torch::linalg_cond(Jp)).item<Real>() << std::endl;

  for (unsigned int itr = 0; itr < _autoscale_miter; itr++)
  {
    // check for convergence
    auto rR = torch::max(torch::abs(1.0 - 1.0 / torch::sqrt(std::get<0>(Jp.max(-1))))).item<Real>();
    auto rC = torch::max(torch::abs(1.0 - 1.0 / torch::sqrt(std::get<0>(Jp.max(-2))))).item<Real>();
    if (verbose)
      std::cout << "ITERATION " << itr << ", ROW ILLNESS = " << std::scientific << rR
                << ", COL ILLNESS = " << std::scientific << rC << std::endl;
    if (rR < _autoscale_tol && rC < _autoscale_tol)
      break;

    // scale rows and columns
    for (TorchSize i = 0; i < _ndof; i++)
    {
      auto ar = 1.0 / torch::sqrt(torch::max(Jp.base_index({i})));
      auto ac = 1.0 / torch::sqrt(torch::max(Jp.base_index({Slice(), i})));
      _row_scaling.base_index({i}) *= ar;
      _col_scaling.base_index({i}) *= ac;
      Jp.base_index({i}) *= ar;
      Jp.base_index({Slice(), i}) *= ac;
    }
  }

  _scaling_matrices_initialized = true;

  if (verbose)
    std::cout << " After automatic scaling cond(J) = " << std::scientific
              << torch::max(torch::linalg_cond(Jp)).item<Real>() << std::endl;
}

BatchTensor
NonlinearSystem::scale_residual(const BatchTensor & r) const
{
  neml_assert_dbg(
      _autoscale == _scaling_matrices_initialized,
      _autoscale ? "Automatic scaling is requested but scaling matrices have not been initialized."
                 : "Automatic scaling is not requested but scaling matrices were initialized.");
  return _row_scaling * r;
}

BatchTensor
NonlinearSystem::scale_Jacobian(const BatchTensor & J) const
{
  neml_assert_dbg(
      _autoscale == _scaling_matrices_initialized,
      _autoscale ? "Automatic scaling is requested but scaling matrices have not been initialized."
                 : "Automatic scaling is not requested but scaling matrices were initialized.");
  return math::bmm(math::bmm(math::base_diag_embed(_row_scaling), J),
                   math::base_diag_embed(_col_scaling));
}

BatchTensor
NonlinearSystem::scale_direction(const BatchTensor & p) const
{
  neml_assert_dbg(
      _autoscale == _scaling_matrices_initialized,
      _autoscale ? "Automatic scaling is requested but scaling matrices have not been initialized."
                 : "Automatic scaling is not requested but scaling matrices were initialized.");
  return _autoscale ? _col_scaling * p : p;
}

void
NonlinearSystem::set_solution(const BatchTensor & x)
{
  _solution.variable_data().copy_(x);
}

BatchTensor
NonlinearSystem::residual(const BatchTensor & x)
{
  set_solution(x);
  residual();
  return residual_view();
}

void
NonlinearSystem::residual()
{
  assemble(true, false);

  if (_autoscale)
    _scaled_residual = scale_residual(_residual);
}

BatchTensor
NonlinearSystem::Jacobian(const BatchTensor & x)
{
  set_solution(x);
  Jacobian();
  return Jacobian_view();
}

void
NonlinearSystem::Jacobian()
{
  assemble(false, true);

  if (_autoscale)
    _scaled_Jacobian = scale_Jacobian(_Jacobian);
}

std::tuple<BatchTensor, BatchTensor>
NonlinearSystem::residual_and_Jacobian(const BatchTensor & x)
{
  set_solution(x);
  residual_and_Jacobian();
  return {residual_view(), Jacobian_view()};
}

void
NonlinearSystem::residual_and_Jacobian()
{
  assemble(true, true);

  if (_autoscale)
  {
    _scaled_residual = scale_residual(_residual);
    _scaled_Jacobian = scale_Jacobian(_Jacobian);
  }
}

BatchTensor
NonlinearSystem::residual_norm() const
{
  return math::linalg::vector_norm(residual_view());
}
} // namespace neml2
