MeanField_SERiF/tests/sandbox/stellar_rayfem.hpp

#pragma once

#include <algorithm>
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>

#include "mfem/fem/gridfunc.hpp"
#include "mfem/mesh/submesh/submesh.hpp"
#include "mfem/mesh/submesh/transfermap.hpp"
#include "libRayfem/network/client.hpp"
#include "serif/backend/mfem/discretization/field_values.hpp"
#include "serif/backend/mfem/model/volume_displacement.hpp"
#include "serif/discretization/blocks/fields.hpp"

namespace serif::visualization::backend::mfem {
    /**
     * @brief Own visualization copies on the stellar interior's reference submesh.
     * @note update() requires the evaluator cache for this exact state. For StandardNewton,
     * use initial/resumed/accepted events, never finished or a rejected trial. The existing
     * nodal displacement is sent separately; neither the solver mesh nor state is changed.
     * Reconstructed density is sampled into a higher-order discontinuous display field. This
     * polynomial display approximation never enters the solve or physical diagnostics.
     */
    template <typename Discretization>
    class StellarFrame {
        using Density = discretization::blocks::density::mass;
        using Potential = discretization::blocks::gravity::potential;
        using Displacement = discretization::blocks::displacement::vector;
        using Shape = discretization::blocks::surface_deformation::coefficients;

        Discretization m_fe;
        ::mfem::GridFunction m_parent_density, m_parent_potential, m_parent_displacement;
        ::mfem::SubMesh m_mesh;
        std::unique_ptr<::mfem::L2_FECollection> m_display_density_collection;
        ::mfem::FiniteElementSpace m_density_space, m_potential_space, m_displacement_space;
        ::mfem::GridFunction m_density, m_potential, m_displacement;
        std::unique_ptr<::mfem::TransferMap> m_density_transfer;
        ::mfem::TransferMap m_potential_transfer, m_displacement_transfer;

        /** @brief Extract the material mesh while preserving parent element reference coordinates. */
        static ::mfem::SubMesh interior_mesh(const Discretization& fe) {
            const auto* parent = fe.template binding<Displacement>().space().GetMesh();
            if (fe.template binding<Density>().space().GetMesh() != parent ||
                fe.template binding<Potential>().space().GetMesh() != parent) {
                throw std::invalid_argument("Stellar visualization requires a common parent mesh");
            }
            ::mfem::Array<int> attributes;
            for (int attribute : fe.deformation.reference.stellar_attributes) attributes.Append(attribute);
            if (attributes.Size() == 0) throw std::invalid_argument("No stellar material attributes to visualize");
            return ::mfem::SubMesh::CreateFromDomain(*parent, attributes);
        }

        /** @brief Use additional display nodes only when an ordinary GridFunction cannot represent density. */
        static std::unique_ptr<::mfem::L2_FECollection> density_display_collection(const Discretization& fe) {
            const auto& binding = fe.template binding<Density>();
            if (binding.reconstruction() == nullptr) return {};
            return std::make_unique<::mfem::L2_FECollection>(binding.space().GetMaxElementOrder() + 3, 3, ::mfem::BasisType::GaussLobatto);
        }

        /** @brief Copy supported physical coefficients into an ordinary visualization field. */
        template <typename Term, typename Values>
        void scatter(const Values& values, ::mfem::GridFunction& field) {
            const auto& binding = m_fe.template binding<Term>();
            if (values.size() != static_cast<std::size_t>(binding.size())) {
                throw std::invalid_argument("State and visualization binding differ");
            }
            ::mfem::Vector reduced(binding.size());
            for (int i = 0; i < reduced.Size(); ++i) reduced[i] = values[i];
            field = binding.scatter(reduced);
        }

        /** @brief Sample reconstructed density into a display-only polynomial on each material element. */
        template <typename Values>
        void sample_density(const Values& values) {
            using namespace utils::errors;

            const auto& binding = m_fe.template binding<Density>();
            serif_ensure<SERiFErrorCode::INVALID_DOF_COUNT>(values.size() == static_cast<std::size_t>(binding.size()),
                "Density state and visualization binding have different coefficient counts.");
            ::mfem::Vector reduced(binding.size());
            std::ranges::copy(values, reduced.GetData());
            discretization::backend::mfem::FieldValues<Density> density(binding, reduced);
            const auto& parent_elements = m_mesh.GetParentElementIDMap();
            ::mfem::Array<int> dofs;
            ::mfem::Vector local;
            for (int element = 0; element < m_mesh.GetNE(); ++element) {
                const auto& nodes = m_density_space.GetFE(element)->GetNodes();
                density.values(parent_elements[element], nodes, local);
                m_density_space.GetElementDofs(element, dofs);
                m_density.SetSubVector(dofs, local);
            }
        }

    public:
        /** @brief Retain geometry and field copies for successive accepted solver states. */
        explicit StellarFrame(Discretization fe)
            : m_fe(std::move(fe)),
              m_parent_density(&m_fe.template binding<Density>().space()),
              m_parent_potential(&m_fe.template binding<Potential>().space()),
              m_parent_displacement(&m_fe.template binding<Displacement>().space()),
              m_mesh(interior_mesh(m_fe)),
              m_display_density_collection(density_display_collection(m_fe)),
              m_density_space(&m_mesh, m_display_density_collection ? m_display_density_collection.get() : m_parent_density.FESpace()->FEColl(),
                  m_parent_density.FESpace()->GetVDim(), m_parent_density.FESpace()->GetOrdering()),
              m_potential_space(&m_mesh, m_parent_potential.FESpace()->FEColl(),
                  m_parent_potential.FESpace()->GetVDim(), m_parent_potential.FESpace()->GetOrdering()),
              m_displacement_space(&m_mesh, m_parent_displacement.FESpace()->FEColl(),
                  m_parent_displacement.FESpace()->GetVDim(), m_parent_displacement.FESpace()->GetOrdering()),
              m_density(&m_density_space), m_potential(&m_potential_space), m_displacement(&m_displacement_space),
              m_potential_transfer(m_parent_potential, m_potential),
              m_displacement_transfer(m_parent_displacement, m_displacement) {
            if (!m_display_density_collection) m_density_transfer = std::make_unique<::mfem::TransferMap>(m_parent_density, m_density);
        }

        StellarFrame(const StellarFrame&) = delete;
        StellarFrame& operator=(const StellarFrame&) = delete;
        StellarFrame(StellarFrame&&) = delete;
        StellarFrame& operator=(StellarFrame&&) = delete;

        /** @brief Publish one accepted state through its matching cached displacement. */
        template <typename State, typename Cache>
        void update(const State& state, const Cache& cache) {
            namespace deformation = geometry::deformation;
            const auto& linearization = cache.template get<deformation::keys::Linearization>();
            if (!std::ranges::equal(linearization.parameters, state.template get<Shape>())) {
                throw std::invalid_argument("Visualization state and cached deformation differ");
            }
            if (m_density_transfer) {
                scatter<Density>(state.template get<Density>(), m_parent_density);
                m_density_transfer->Transfer(m_parent_density, m_density);
            } else {
                sample_density(state.template get<Density>());
            }
            scatter<Potential>(state.template get<Potential>(), m_parent_potential);
            m_parent_displacement = m_fe.template binding<Displacement>().scatter(
                cache.template get<deformation::backend::mfem::volume_displacement_keys::Coefficients>());
            m_potential_transfer.Transfer(m_parent_potential, m_potential);
            m_displacement_transfer.Transfer(m_parent_displacement, m_displacement);
        }

        /** @brief Return the unchanged computational mesh sent to the viewer. */
        [[nodiscard]] ::mfem::Mesh& mesh() { return m_mesh; }
        /** @brief Return the separate physical displacement applied by the viewer. */
        [[nodiscard]] ::mfem::GridFunction& displacement() { return m_displacement; }
        /** @brief Label a reconstructed density display explicitly as sampled data. */
        [[nodiscard]] std::vector<RayFEM::NamedField> fields() {
            return {{m_density_transfer ? "density" : "density (sampled)", &m_density}, {"potential", &m_potential}, {"displacement", &m_displacement}};
        }
    };
}