/* SPDX-FileCopyrightText: Copyright (c) 2022 - 2026, the adamantine authors.
 * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 */

#include <MechanicalOperator.hh>
#include <instantiation.hh>
#include <utils.hh>

#include <deal.II/base/index_set.h>
#include <deal.II/base/symmetric_tensor.h>
#include <deal.II/base/tensor.h>
#include <deal.II/differentiation/ad.h>
#include <deal.II/dofs/dof_tools.h>
#include <deal.II/fe/fe_update_flags.h>
#include <deal.II/grid/filtered_iterator.h>
#include <deal.II/grid/tria_iterator.h>
#include <deal.II/hp/fe_values.h>
#include <deal.II/hp/q_collection.h>
#include <deal.II/lac/affine_constraints.h>
#if DEAL_II_VERSION_GTE(9, 7, 0) && defined(DEAL_II_TRILINOS_WITH_TPETRA)
#include <deal.II/lac/affine_constraints.templates.h>
#endif
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/sparsity_tools.h>
#include <deal.II/meshworker/mesh_loop.h>
#include <deal.II/physics/elasticity/standard_tensors.h>

#ifdef ADAMANTINE_WITH_CALIPER
#include <caliper/cali.h>
#endif

namespace adamantine
{
namespace
{
struct CopyDataRightHandSide
{
  CopyDataRightHandSide(unsigned int dofs_per_cell)
  {
    cell_rhs.reinit(dofs_per_cell);
    local_dof_indices.resize(dofs_per_cell);
  }
  dealii::Vector<double> cell_rhs;
  std::vector<dealii::types::global_dof_index> local_dof_indices;

  template <class Iterator>
  void reinit(const Iterator &cell)
  {
    cell_rhs = 0.;
    cell->get_dof_indices(local_dof_indices);
  }
};

template <int dim>
struct ScratchDataMatrix
{
  ScratchDataMatrix(dealii::hp::FEValues<dim> const &displacement_hp_fe_values)
      : displacement_hp_fe_values(
            displacement_hp_fe_values.get_mapping_collection(),
            displacement_hp_fe_values.get_fe_collection(),
            displacement_hp_fe_values.get_quadrature_collection(),
            displacement_hp_fe_values.get_update_flags())
  {
  }

  ScratchDataMatrix(ScratchDataMatrix<dim> const &scratch_data)
      : displacement_hp_fe_values(
            scratch_data.displacement_hp_fe_values.get_mapping_collection(),
            scratch_data.displacement_hp_fe_values.get_fe_collection(),
            scratch_data.displacement_hp_fe_values.get_quadrature_collection(),
            scratch_data.displacement_hp_fe_values.get_update_flags())
  {
  }

  dealii::hp::FEValues<dim> displacement_hp_fe_values;
};

struct CopyDataMatrix
{
  CopyDataMatrix(unsigned int dofs_per_cell)
  {
    cell_matrix.reinit(dofs_per_cell, dofs_per_cell);
    local_dof_indices.resize(dofs_per_cell);
  }
  dealii::FullMatrix<double> cell_matrix;
  std::vector<dealii::types::global_dof_index> local_dof_indices;

  template <class Iterator>
  void reinit(const Iterator &cell)
  {
    cell_matrix = 0.;
    cell->get_dof_indices(local_dof_indices);
  }
};

struct CopyDataInternalForce
{
  CopyDataInternalForce(unsigned int dofs_per_cell)
  {
    cell_force.reinit(dofs_per_cell);
    local_dof_indices.resize(dofs_per_cell);
  }

  dealii::Vector<double> cell_force;
  std::vector<dealii::types::global_dof_index> local_dof_indices;

  template <class Iterator>
  void reinit(const Iterator &cell)
  {
    cell_force = 0.;
    cell->get_dof_indices(local_dof_indices);
  }
};

} // namespace

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
MechanicalOperator<dim, n_materials, p_order, MaterialStates, MemorySpaceType>::
    MechanicalOperator(
        MPI_Comm const &communicator,
        MaterialProperty<dim, n_materials, p_order, MaterialStates,
                         MemorySpaceType> &material_properties)
    : _communicator(communicator), _material_properties(material_properties)
{
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::
    reinit(dealii::DoFHandler<dim> const &dof_handler,
           dealii::AffineConstraints<double> const &affine_constraints,
           dealii::hp::QCollection<dim> const &q_collection,
           std::vector<std::shared_ptr<BodyForce<dim>>> const &body_forces)
{
  _dof_handler = &dof_handler;
  _affine_constraints = &affine_constraints;
  _q_collection = &q_collection;
  assemble_matrix();
  assemble_rhs(body_forces);
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host>
MechanicalOperator<dim, n_materials, p_order, MaterialStates, MemorySpaceType>::
    assemble_internal_force(
        std::vector<std::vector<dealii::SymmetricTensor<2, dim>>> const &stress)
        const
{
  auto const locally_owned_dofs = _dof_handler->locally_owned_dofs();
  auto const locally_relevant_dofs =
      dealii::DoFTools::extract_locally_relevant_dofs(*_dof_handler);
  unsigned int const dofs_per_cell =
      _dof_handler->get_fe_collection().max_dofs_per_cell();
  dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host>
      internal_force(locally_owned_dofs, locally_relevant_dofs, _communicator);

  dealii::hp::FEValues<dim> displacement_hp_fe_values(
      _dof_handler->get_fe_collection(), *_q_collection,
      dealii::update_gradients | dealii::update_JxW_values);

  auto cell_worker =
      [&](const typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
          ScratchDataMatrix<dim> &scratch_data,
          CopyDataInternalForce &copy_data)
  {
    scratch_data.displacement_hp_fe_values.reinit(cell);
    auto const &fe_values =
        scratch_data.displacement_hp_fe_values.get_present_fe_values();
    copy_data.reinit(cell);

    dealii::FEValuesExtractors::Vector const displacements(0);
    auto const &cell_stress = stress[cell->active_cell_index()];
    for (auto const i : fe_values.dof_indices())
      for (auto const q_point : fe_values.quadrature_point_indices())
        copy_data.cell_force(i) +=
            fe_values[displacements].symmetric_gradient(i, q_point) *
            cell_stress[q_point] * fe_values.JxW(q_point);
  };

  auto copier = [&](CopyDataInternalForce const &copy_data)
  {
    _affine_constraints->distribute_local_to_global(
        copy_data.cell_force, copy_data.local_dof_indices, internal_force);
  };

  dealii::IteratorFilters::ActiveFEIndexEqualTo filter(
      0, /* locally owned */ true);
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      begin(filter, _dof_handler->begin_active());
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      end(filter, _dof_handler->end());

  ScratchDataMatrix<dim> scratch_data(displacement_hp_fe_values);
  dealii::MeshWorker::mesh_loop(begin, end, cell_worker, copier, scratch_data,
                                CopyDataInternalForce(dofs_per_cell),
                                dealii::MeshWorker::assemble_own_cells);
  internal_force.compress(dealii::VectorOperation::add);

  return internal_force;
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host>
MechanicalOperator<dim, n_materials, p_order, MaterialStates, MemorySpaceType>::
    assemble_residual(
        std::vector<std::vector<dealii::SymmetricTensor<2, dim>>> const &stress)
        const
{
  auto residual = assemble_internal_force(stress);
  residual -= _system_rhs;
  return residual;
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::
    assemble_jacobian(
        std::vector<std::vector<dealii::SymmetricTensor<4, dim>>> const
            &constitutive_tangent)
{
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_BEGIN("assemble mechanical Jacobian");
#endif

  _system_matrix = 0.;

  auto const dofs_per_cell =
      _dof_handler->get_fe_collection().max_dofs_per_cell();
  dealii::hp::FEValues<dim> displacement_hp_fe_values(
      _dof_handler->get_fe_collection(), *_q_collection,
      dealii::update_gradients | dealii::update_JxW_values);

  auto cell_worker =
      [&](const typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
          ScratchDataMatrix<dim> &scratch_data, CopyDataMatrix &copy_data)
  {
    scratch_data.displacement_hp_fe_values.reinit(cell);
    auto const &fe_values =
        scratch_data.displacement_hp_fe_values.get_present_fe_values();
    auto const &cell_tangent = constitutive_tangent[cell->active_cell_index()];
    copy_data.reinit(cell);

    dealii::FEValuesExtractors::Vector const displacements(0);
    for (auto const i : fe_values.dof_indices())
      for (auto const j : fe_values.dof_indices())
        for (auto const q_point : fe_values.quadrature_point_indices())
        {
          auto const strain_i =
              fe_values[displacements].symmetric_gradient(i, q_point);
          auto const strain_j =
              fe_values[displacements].symmetric_gradient(j, q_point);
          copy_data.cell_matrix(i, j) += strain_i *
                                         (cell_tangent[q_point] * strain_j) *
                                         fe_values.JxW(q_point);
        }
  };

  auto copier = [&](const CopyDataMatrix &copy_data)
  {
    _affine_constraints->distribute_local_to_global(
        copy_data.cell_matrix, copy_data.local_dof_indices, _system_matrix);
  };

  dealii::IteratorFilters::ActiveFEIndexEqualTo filter(
      0, /* locally owned */ true);
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      begin(filter, _dof_handler->begin_active());
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      end(filter, _dof_handler->end());

  ScratchDataMatrix<dim> scratch_data(displacement_hp_fe_values);
  dealii::MeshWorker::mesh_loop(begin, end, cell_worker, copier, scratch_data,
                                CopyDataMatrix(dofs_per_cell),
                                dealii::MeshWorker::assemble_own_cells);

  _system_matrix.compress(dealii::VectorOperation::add);
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_END("assemble mechanical Jacobian");
#endif

#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_BEGIN("initialize mechanical Jacobian preconditioner");
#endif
  _preconditioner.clear();
  _preconditioner.initialize(_system_matrix);
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_END("initialize mechanical Jacobian preconditioner");
#endif
  _system_matrix_is_elastic = false;
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::ensure_elastic_matrix()
{
  if (!_system_matrix_is_elastic)
    assemble_matrix();
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::assemble_matrix()
{
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_BEGIN("assemble mechanical matrix");
#endif

  // Create the sparsity pattern. Since we use a Trilinos matrix we don't need
  // the sparsity pattern to outlive the sparse matrix.
  auto locally_owned_dofs = _dof_handler->locally_owned_dofs();
  auto locally_relevant_dofs =
      dealii::DoFTools::extract_locally_relevant_dofs(*_dof_handler);
  dealii::DynamicSparsityPattern dsp(locally_relevant_dofs);
  dealii::DoFTools::make_sparsity_pattern(*_dof_handler, dsp,
                                          *_affine_constraints, false);
  dealii::SparsityTools::distribute_sparsity_pattern(
      dsp, locally_owned_dofs, _communicator, locally_relevant_dofs);

  _system_matrix.reinit(locally_owned_dofs, dsp, _communicator);

  dealii::hp::FEValues<dim> displacement_hp_fe_values(
      _dof_handler->get_fe_collection(), *_q_collection,
      dealii::update_values | dealii::update_gradients |
          dealii::update_JxW_values);

  unsigned int const dofs_per_cell =
      _dof_handler->get_fe_collection().max_dofs_per_cell();

  // Loop over the locally owned cells that are not FE_Nothing and assemble the
  // sparse matrix and the right-hand-side
  auto cell_worker =
      [&](const typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
          ScratchDataMatrix<dim> &scratch_data, CopyDataMatrix &copy_data)
  {
    dealii::hp::FEValues<dim> &displacement_hp_fe_values =
        scratch_data.displacement_hp_fe_values;
    displacement_hp_fe_values.reinit(cell);
    auto const &fe_values = displacement_hp_fe_values.get_present_fe_values();
    auto const &fe = fe_values.get_fe();

    // Assemble the local martrix
    copy_data.reinit(cell);
    dealii::FullMatrix<double> &cell_matrix = copy_data.cell_matrix;
    double const lambda = this->_material_properties.get_mechanical_property(
        cell, StateProperty::lame_first_parameter);
    double const mu = this->_material_properties.get_mechanical_property(
        cell, StateProperty::lame_second_parameter);
    for (auto const i : fe_values.dof_indices())
    {
      auto const component_i = fe.system_to_component_index(i).first;
      for (auto const j : fe_values.dof_indices())
      {
        auto const component_j = fe.system_to_component_index(j).first;
        for (auto const q_point : fe_values.quadrature_point_indices())
        {
          cell_matrix(i, j) +=
              ((fe_values.shape_grad(i, q_point)[component_i] *
                fe_values.shape_grad(j, q_point)[component_j] * lambda) +
               (fe_values.shape_grad(i, q_point)[component_j] *
                fe_values.shape_grad(j, q_point)[component_i] * mu) +
               ((component_i == component_j)
                    ? mu * fe_values.shape_grad(i, q_point) *
                          fe_values.shape_grad(j, q_point)
                    : 0.)) *
              fe_values.JxW(q_point);
        }
      }
    }
  };

  auto copier = [&](const CopyDataMatrix &cd)
  {
    _affine_constraints->distribute_local_to_global(
        cd.cell_matrix, cd.local_dof_indices, _system_matrix);
  };

  dealii::IteratorFilters::ActiveFEIndexEqualTo filter(
      0,
      /* locally owned */ true);
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      begin(filter, _dof_handler->begin_active());
  dealii::FilteredIterator<
      typename dealii::DoFHandler<dim>::active_cell_iterator>
      end(filter, _dof_handler->end());

  ScratchDataMatrix<dim> scratch_data(displacement_hp_fe_values);

  dealii::MeshWorker::mesh_loop(begin, end, cell_worker, copier, scratch_data,
                                CopyDataMatrix(dofs_per_cell),
                                dealii::MeshWorker::assemble_own_cells);

  _system_matrix.compress(dealii::VectorOperation::add);
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_END("assemble mechanical matrix");
#endif

#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_BEGIN("initialize mechanical matrix preconditioner");
#endif

// Constant modes reduce cg iterations by 40%.
// FIXME: With Tpetra, we run into a level 0 scratch memory issue on Frontier.
#if DEAL_II_VERSION_GTE(9, 7, 0) && defined(DEAL_II_TRILINOS_WITH_TPETRA)
  _preconditioner.clear();
  _preconditioner.initialize(_system_matrix);

#else
  // Provide the constant displacement modes to AMG for the elasticity problem.
  dealii::FEValuesExtractors::Vector const displacement_components(0);
  auto const component_mask =
      _dof_handler->get_fe_collection().component_mask(displacement_components);

  std::vector<std::vector<bool>> constant_modes;
  dealii::DoFTools::extract_constant_modes(*_dof_handler, component_mask,
                                           constant_modes);

  typename TrilinosPreconditionerType::AdditionalData amg_data;
  amg_data.constant_modes = constant_modes;
  _preconditioner.clear();
  _preconditioner.initialize(_system_matrix, amg_data);
#endif

#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_END("initialize mechanical matrix preconditioner");
#endif
  _system_matrix_is_elastic = true;
}

template <int dim, int n_materials, int p_order, typename MaterialStates,
          typename MemorySpaceType>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::
    assemble_rhs(
        std::vector<std::shared_ptr<BodyForce<dim>>> const &body_forces)
{
#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_BEGIN("assemble mechanical rhs");
#endif

  auto locally_owned_dofs = _dof_handler->locally_owned_dofs();
  auto locally_relevant_dofs =
      dealii::DoFTools::extract_locally_relevant_dofs(*_dof_handler);
  dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host>
      assembled_rhs(locally_owned_dofs, locally_relevant_dofs, _communicator);

  // Add gravitational body force if present
  if (body_forces.size())
  {
    unsigned int const dofs_per_cell =
        _dof_handler->get_fe_collection().max_dofs_per_cell();
    dealii::hp::FEValues<dim> displacement_hp_fe_values(
        _dof_handler->get_fe_collection(), *_q_collection,
        dealii::update_values | dealii::update_gradients |
            dealii::update_JxW_values);
    std::vector<dealii::types::global_dof_index> local_dof_indices(
        dofs_per_cell);
    dealii::Vector<double> cell_rhs(dofs_per_cell);

    for (auto const &cell : _dof_handler->active_cell_iterators() |
                                dealii::IteratorFilters::ActiveFEIndexEqualTo(
                                    0, /* locally owned */ true))
    {
      cell_rhs = 0.;

      displacement_hp_fe_values.reinit(cell);
      auto const &fe_values = displacement_hp_fe_values.get_present_fe_values();

      dealii::Tensor<1, dim, double> body_force;
      for (auto &force : body_forces)
      {
        body_force += force->eval(cell);
      }

      dealii::FEValuesExtractors::Vector const displacements(0);

      for (auto const i : fe_values.dof_indices())
      {
        for (auto const q_point : fe_values.quadrature_point_indices())
          cell_rhs(i) += body_force *
                         fe_values[displacements].value(i, q_point) *
                         fe_values.JxW(q_point);
      }

      cell->get_dof_indices(local_dof_indices);
      _affine_constraints->distribute_local_to_global(
          cell_rhs, local_dof_indices, assembled_rhs);
    }
    assembled_rhs.compress(dealii::VectorOperation::add);
  }

  // When solving the system, we don't want ghost entries
  _system_rhs.reinit(_dof_handler->locally_owned_dofs(), _communicator);
  _system_rhs = assembled_rhs;

#ifdef ADAMANTINE_WITH_CALIPER
  CALI_MARK_END("assemble mechanical rhs");
#endif
}
} // namespace adamantine

INSTANTIATE_DIM_NMAT_PORDER_MATERIALSTATES_HOST(MechanicalOperator)
INSTANTIATE_DIM_NMAT_PORDER_MATERIALSTATES_DEVICE(MechanicalOperator)