// SPDX-FileCopyrightText: 2022 The Ikarus Developers mueller@ibb.uni-stuttgart.de
// SPDX-License-Identifier: LGPL-3.0-or-later

//
#include <config.h>

#include "common.hh"
#include "testHelpers.hh"

#include <dune/common/test/testsuite.hh>
#include <dune/functions/functionspacebases/basistags.hh>
#include <dune/functions/functionspacebases/boundarydofs.hh>
#include <dune/functions/functionspacebases/lagrangebasis.hh>
#include <dune/functions/functionspacebases/powerbasis.hh>

#include "spdlog/spdlog.h"

#include <Eigen/Core>

#include <pyikarus/assembler/simpleAssemblers.hh>
#include <pyikarus/controlRoutines/loadControl.hh>
#include <pyikarus/finiteElements/mechanics/nonLinearElasticityFE.hh>
#include <pyikarus/linearAlgebra/dirichletValues.hh>
#include <pyikarus/linearAlgebra/nonLinearOperator.hh>
#include <pyikarus/solver/nonLinearSolver/newtonRaphson.hh>
#include <pyikarus/solver/nonLinearSolver/trustRegion.hh>
#include <pyikarus/utils/algorithms.hh>
#include <pyikarus/utils/drawing/griddrawer.hh>
#include <pyikarus/utils/init.hh>
#include <pyikarus/utils/observer/controlVTKWriter.hh>

using Dune::TestSuite;

template <typename Grid>
auto NonLinearElasticityLoadControlNRandTR() {
  TestSuite t("NonLinearElasticityLoadControlNRandTR" + Dune::className(Grid{}));
  auto grid     = createGrid<Grid>();
  auto gridView = grid->leafGridView();

  using namespace Ikarus;

  using namespace Dune::Functions::BasisFactory;
  auto basis = makeBasis(gridView, power<2>(lagrange<1>(), FlatInterleaved()));

  auto g = basis.localView();
  std::vector<Ikarus::NonLinearElasticityFE<decltype(basis)>> fes;
  auto volumeLoad = []([[maybe_unused]] auto& globalCoord, auto& lamb) {
    Eigen::Vector2d fext;
    fext.setZero();
    fext[1] = 2 * lamb;
    fext[0] = lamb;
    return fext;
  };
  for (auto& element : elements(gridView))
    fes.emplace_back(basis, element, 1000, 0.3, nullptr, nullptr, volumeLoad);
  auto basisP = std::make_shared<const decltype(basis)>(basis);
  Ikarus::DirichletValues dirichletValues(basisP);
  dirichletValues.fixBoundaryDOFs([&](auto& dirichletFlags, auto&& localIndex, auto&& localView, auto&& intersection) {
    if (std::abs(intersection.geometry().center()[1]) < 1e-8) dirichletFlags[localView.index(localIndex)] = true;
  });

  auto sparseAssembler = SparseFlatAssembler(fes, dirichletValues);

  Eigen::VectorXd d;
  d.setZero(basis.size());
  double lambda = 0.0;

  auto req = FErequirements().addAffordance(Ikarus::AffordanceCollections::elastoStatics);

  auto residualFunction = [&](auto&& disp_, auto&& lambdaLocal) -> auto& {
    req.insertGlobalSolution(Ikarus::FESolutions::displacement, disp_)
        .insertParameter(Ikarus::FEParameter::loadfactor, lambdaLocal);
    return sparseAssembler.getVector(req);
  };

  auto KFunction = [&](auto&& disp_, auto&& lambdaLocal) -> auto& {
    req.insertGlobalSolution(Ikarus::FESolutions::displacement, disp_)
        .insertParameter(Ikarus::FEParameter::loadfactor, lambdaLocal);
    return sparseAssembler.getMatrix(req);
  };

  auto energyFunction = [&](auto&& disp_, auto&& lambdaLocal) -> auto& {
    req.insertGlobalSolution(Ikarus::FESolutions::displacement, disp_)
        .insertParameter(Ikarus::FEParameter::loadfactor, lambdaLocal);
    return sparseAssembler.getScalar(req);
  };

  auto nonLinOp = Ikarus::NonLinearOperator(linearAlgebraFunctions(energyFunction, residualFunction, KFunction),
                                            parameter(d, lambda));

  const double gradTol = 1e-8;

  auto tr = Ikarus::makeTrustRegion(nonLinOp);
  tr->setup({.verbosity = 1,
             .maxiter   = 1000,
             .grad_tol  = gradTol,
             .corr_tol  = 1e-16,  // everything should converge to the gradient tolerance
             .useRand   = false,
             .rho_reg   = 1e8,
             .Delta0    = 1});

  auto vtkWriter = std::make_shared<ControlSubsamplingVertexVTKWriter<decltype(basis)>>(basis, d, 2);
  vtkWriter->setFileNamePrefix("Test2Dsolid");
  vtkWriter->setFieldInfo("Displacement", Dune::VTK::FieldInfo::Type::vector, 2);

  auto lc = Ikarus::LoadControl(tr, 1, {0, 2000});
  lc.subscribeAll(vtkWriter);
  const auto controlInfo = lc.run();
  nonLinOp.template update<0>();
  const auto maxDisp = std::ranges::max(d);
  const double energyExpected
      = (std::is_same_v<Grid, Grids::Yasp>)
            ? -1.4809559783564966e+03
            : ((std::is_same_v<Grid, Grids::Alu>) ? -1.4842107484533601e+03
                                                  : /* std::is_same_v<Grid, Grids::Iga> */ -8.1142552237939071e+02);
  const double maxDispExpected
      = (std::is_same_v<Grid, Grids::Yasp>)
            ? 0.786567027108437
            : ((std::is_same_v<Grid, Grids::Alu>) ? 0.78426066482258983
                                                  : /* std::is_same_v<Grid, Grids::Iga> */ 0.615624125459537153);
  std::cout << std::setprecision(20) << nonLinOp.value() << std::endl;
  t.check(Dune::FloatCmp::eq(energyExpected, nonLinOp.value()), "energyExpected == nonLinOp.value()")
      << "energyExpected: " << energyExpected << "\nnonLinOp.value(): " << nonLinOp.value();

  t.check(std::abs(maxDispExpected - maxDisp) < 1e-12, "maxDispExpected-maxDisp") << "maxDispExpected: \n"
                                                                                  << maxDispExpected << "\nmaxDisp: \n"
                                                                                  << maxDisp;

  nonLinOp.template update<1>();
  t.check(controlInfo.success, "Successful result");
  t.check(gradTol >= nonLinOp.derivative().norm(), "Gradient Tolerance should be larger than actual tolerance");
  return t;
}

int main(int argc, char** argv) {
  Ikarus::init(argc, argv);
  TestSuite t;

  t.subTest(NonLinearElasticityLoadControlNRandTR<Grids::Alu>());
  t.subTest(NonLinearElasticityLoadControlNRandTR<Grids::Yasp>());
  t.subTest(NonLinearElasticityLoadControlNRandTR<Grids::Iga>());

  return t.exit();
}
