/* SPDX-FileCopyrightText: Copyright (c) 2022 - 2025, 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
{
template <int dim>
struct ScratchDataRightHandSide
{
  ScratchDataRightHandSide(
      dealii::hp::FEValues<dim> const &displacement_hp_fe_values,
      dealii::hp::FEValues<dim> const &temperature_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()),
        temperature_hp_fe_values(
            temperature_hp_fe_values.get_mapping_collection(),
            temperature_hp_fe_values.get_fe_collection(),
            temperature_hp_fe_values.get_quadrature_collection(),
            temperature_hp_fe_values.get_update_flags())
  {
  }

  ScratchDataRightHandSide(ScratchDataRightHandSide<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()),
        temperature_hp_fe_values(
            scratch_data.temperature_hp_fe_values.get_mapping_collection(),
            scratch_data.temperature_hp_fe_values.get_fe_collection(),
            scratch_data.temperature_hp_fe_values.get_quadrature_collection(),
            scratch_data.temperature_hp_fe_values.get_update_flags())
  {
  }

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

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);
  }
};

} // 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,
        std::vector<double> const &reference_temperatures)
    : _communicator(communicator),
      _reference_temperatures(reference_temperatures),
      _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>
void MechanicalOperator<dim, n_materials, p_order, MaterialStates,
                        MemorySpaceType>::
    update_temperature(
        dealii::DoFHandler<dim> const &thermal_dof_handler,
        dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host> const
            &temperature,
        std::vector<bool> const &has_melted)
{
  _thermal_dof_handler = &thermal_dof_handler;
  _temperature = temperature;
  _has_melted = has_melted;
}

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
}

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);
  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);
  dealii::LA::distributed::Vector<double, dealii::MemorySpace::Host>
      assembled_rhs(locally_owned_dofs, locally_relevant_dofs, _communicator);
  std::vector<dealii::types::global_dof_index> local_dof_indices(dofs_per_cell);
  // If the list of reference temperatures is non-empty, we solve the
  // thermo-elastic problem.
  if (_reference_temperatures.size() > 0)
  {
    // _has_melted is using its own indices. The indices are computed using the
    // locally owned cells of the thermal DoFHandler that have an FE_Index
    // equal to zero. We need to translate these indices to be used with the
    // mechanical DoFHandler. This is simplified by the fact that both the
    // thermal and the mechanical simulation use the same Triangulation and have
    // the same cells locally owned.
    auto &triangulation = _dof_handler->get_triangulation();
    unsigned int const n_local_active_cells = triangulation.n_active_cells();
    std::vector<unsigned int> cell_indices(n_local_active_cells);
    unsigned int cell_index = 0;
    for (auto const &tria_cell :
         triangulation.active_cell_iterators() |
             dealii::IteratorFilters::LocallyOwnedCell())
    {
      dealii::TriaIterator<dealii::DoFCellAccessor<dim, dim, false>>
          temperature_cell(&triangulation, tria_cell->level(),
                           tria_cell->index(), _thermal_dof_handler);
      if (temperature_cell->active_fe_index() == 0)
      {
        dealii::TriaIterator<dealii::DoFCellAccessor<dim, dim, false>>
            displacement_cell(&triangulation, tria_cell->level(),
                              tria_cell->index(), _dof_handler);
        if (displacement_cell->active_fe_index() == 0)
        {
          cell_indices[displacement_cell->active_cell_index()] = cell_index;
        }
        ++cell_index;
      }
    }

    _temperature.update_ghost_values();

    std::vector<dealii::types::global_dof_index> temperature_local_dof_indices(
        _thermal_dof_handler->get_fe_collection().max_dofs_per_cell());
    double const initial_temperature = _reference_temperatures.back();

    auto cell_worker =
        [&](const typename dealii::DoFHandler<dim>::active_cell_iterator &cell,
            ScratchDataRightHandSide<dim> &scratch_data,
            CopyDataRightHandSide &copy_data)
    {
      // Get the temperature cell associated to the mechanical cell
      dealii::TriaIterator<dealii::DoFCellAccessor<dim, dim, false>>
          temperature_cell(&triangulation, cell->level(), cell->index(),
                           _thermal_dof_handler);

      // Get the appropriate reference temperature for the cell. If the cell
      // is not in the unmelted substrate, the reference temperature depends
      // on the material.
      double reference_temperature =
          _has_melted[cell_indices[cell->active_cell_index()]]
              ? _reference_temperatures[temperature_cell->material_id()]
              : initial_temperature;

      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();

      double const alpha = this->_material_properties.get_mechanical_property(
          cell, StateProperty::thermal_expansion_coef);
      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);

      dealii::hp::FEValues<dim> &temperature_hp_fe_values =
          scratch_data.temperature_hp_fe_values;
      temperature_hp_fe_values.reinit(temperature_cell);
      auto &temperature_fe_values =
          temperature_hp_fe_values.get_present_fe_values();
      temperature_cell->get_dof_indices(temperature_local_dof_indices);

      copy_data.reinit(cell);
      dealii::Vector<double> &cell_rhs = copy_data.cell_rhs;

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

      for (auto const q_point : fe_values.quadrature_point_indices())
      {
        double delta_T = -reference_temperature;
        for (unsigned int j = 0; j < temperature_fe_values.dofs_per_cell; ++j)
        {
          delta_T += temperature_fe_values.shape_value(j, q_point) *
                     _temperature(temperature_local_dof_indices[j]);
        }
        auto const B = (3. * lambda + 2 * mu) * alpha * delta_T;

        for (auto const i : fe_values.dof_indices())
        {
          cell_rhs(i) +=
              dealii::trace(fe_values[displacements].gradient(i, q_point)) * B *
              fe_values.JxW(q_point);
        }
      }
    };

    auto copier = [&](const CopyDataRightHandSide &cd)
    {
      _affine_constraints->distribute_local_to_global(
          cd.cell_rhs, cd.local_dof_indices, assembled_rhs);
    };

    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());

    dealii::hp::FEValues<dim> temperature_hp_fe_values(
        _thermal_dof_handler->get_fe_collection(), *_q_collection,
        dealii::update_values);

    ScratchDataRightHandSide<dim> scratch_data(displacement_hp_fe_values,
                                               temperature_hp_fe_values);

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

  // Add gravitational body force if present
  if (body_forces.size())
  {
    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)