module;

#include <cmath>
#include <memory>
#include <stdexcept>

#include <mfem.hpp>
#include <mpi.h>

module mean_field;

import :deformation.surface_metric;
import :field.mfem;
import :utils.domain;

namespace mean_field::deformation {
    // Parameter coordinates need their own reference metric: a modal
    // coefficient is not the value at a boundary node. This assembles the
    // diagonal of J_surface^T (M_boundary tensor I) J_surface, without using
    // a physical prepared state or the nonphysical volume extension.
    mfem::Vector assembleSurfaceParameterMassDiagonal(
        const fem::FEM &finiteElements,
        const PreparedDomainDeformationRuntime &deformation
    ) {
        if (!finiteElements.okay() || !deformation.matchesCurrentDiscretization() ||
            !deformation.supportsSurfaceParameterMetric()) {
            throw std::invalid_argument("A surface parameter metric requires a current reference surface tangent.");
        }
        const auto dependencies = deformation.discretizationDependencies();
        if (dependencies.physicalMeshIdentity != finiteElements.mesh.get() ||
            dependencies.surfaceScalarSpaceIdentity != finiteElements.surfaceDeformationFes.get() ||
            dependencies.volumeDisplacementSpaceIdentity != finiteElements.displacementFes.get()) {
            throw std::invalid_argument("The surface parameter metric and deformation must use the same FEM.");
        }
        int ranks = 0;
        MPI_Comm_size(finiteElements.mesh->GetComm(), &ranks);
        if (ranks != 1) {
            throw std::invalid_argument("Surface parameter metric assembly currently requires one MPI rank.");
        }
        using Schema = utils::domain::CoreEnvelopeVacuumDomainSchema;
        const auto surfaceMap =
            field::make_stellar_surface_scalar_dof_map<Schema>(*finiteElements.surfaceDeformationFes);
        const int dimension = deformation.descriptor().surfaceDeformation.spatialDimension;
        if (deformation.surfaceDisplacementSize() != dimension * surfaceMap.local_size()) {
            throw std::invalid_argument("The surface tangent and scalar boundary map have incompatible traces.");
        }
        constexpr int attribute = Schema::template boundary_attribute<utils::domain::StellarSurface>();
        mfem::Array<int> marker(finiteElements.mesh->bdr_attributes.Max());
        if (attribute <= 0 || attribute > marker.Size()) {
            throw std::invalid_argument("The surface parameter metric requires the stellar boundary attribute.");
        }
        marker                = 0;
        marker[attribute - 1] = 1;
        mfem::ParBilinearForm mass(finiteElements.surfaceDeformationFes.get());
        mass.AddBoundaryIntegrator(new mfem::MassIntegrator(), marker);
        mass.Assemble();
        mass.Finalize();
        std::unique_ptr<mfem::HypreParMatrix> matrix(mass.ParallelAssemble());
        if (!matrix)
            throw std::runtime_error("The reference surface mass failed to assemble.");

        const int count = deformation.parameterCount();
        mfem::Vector parameters(count), direction(count), diagonal(count);
        mfem::Vector trace(deformation.surfaceDisplacementSize());
        mfem::Vector scalarTrace(surfaceMap.local_size());
        mfem::Vector ambient(surfaceMap.volume_true_dof_size()), action(ambient.Size());
        parameters = 0.0;
        for (int parameter = 0; parameter < count; ++parameter) {
            direction            = 0.0;
            direction(parameter) = 1.0;
            deformation.applySurfaceJacobian(parameters, direction, trace);
            double entry = 0.0;
            for (int component = 0; component < dimension; ++component) {
                for (int node = 0; node < scalarTrace.Size(); ++node) {
                    scalarTrace(node) = trace(dimension * node + component);
                }
                ambient = 0.0;
                surfaceMap.scatter(scalarTrace, ambient);
                matrix->Mult(ambient, action);
                entry += ambient * action;
            }
            if (!std::isfinite(entry) || entry <= 0.0) {
                throw std::runtime_error("A surface parameter has a nonpositive or nonfinite reference metric.");
            }
            diagonal(parameter) = entry;
        }
        return diagonal;
    }
} // namespace mean_field::deformation