module;

#include <mfem.hpp>

module mean_field;

import :operators.prepared_displacement_residual;

namespace {
    using Dependencies = mean_field::operators::DisplacementResidualDependencies;

    [[nodiscard]] mean_field::operators::context::pressure_force::PressureForceDependencies
    make_pressure_dependencies(const Dependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .enthalpy     = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision},
            .displacement = {
                .identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision
            }
        };
    }

    [[nodiscard]] mean_field::operators::context::rotational_displacement_force::RotationalDisplacementForceDependencies
    make_rotational_dependencies(const Dependencies &dependencies) {
        return {
            .discretization =
                {.identity = dependencies.discretization.identity, .revision = dependencies.discretization.revision},
            .density = {.identity = dependencies.density.identity, .revision = dependencies.density.revision},
            .displacement =
                {.identity = dependencies.displacement.identity, .revision = dependencies.displacement.revision},
            .rotation = {.identity = dependencies.rotation.identity, .revision = dependencies.rotation.revision}
        };
    }

    void validate_shared_gravity_revisions(
        const mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &gravityContext,
        const Dependencies &dependencies
    ) {
        MFEM_VERIFY(
            gravityContext.IsPrepared(), "PreparedDisplacementResidualOperator requires the shared "
                                         "gravity linearization context to be prepared first."
        );

        const mean_field::operators::context::gravity_field::GravityFieldRevisions &gravityRevisions =
            gravityContext.GetRevisions();

        MFEM_VERIFY(
            gravityRevisions.discretization.value == dependencies.discretization.revision &&
                gravityRevisions.density.value == dependencies.density.revision &&
                gravityRevisions.displacement.value == dependencies.displacement.revision &&
                gravityRevisions.gravity_gradient.value == dependencies.gravityGradient.revision,
            "PreparedDisplacementResidualOperator received dependency "
            "revisions that do not match the shared gravity context."
        );
    }

    void validate_shared_identity_transition(
        const mean_field::operators::DisplacementResidualDependencyStamp &prepared,
        const mean_field::operators::DisplacementResidualDependencyStamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(prepared.identity == requested.identity || prepared.revision != requested.revision, message);
    }

    void add_compatible(
        mfem::Vector &destination,
        const mfem::Vector &source,
        const char *message
    ) {
        MFEM_VERIFY(destination.Size() == source.Size(), message);
        destination += source;
    }
} // namespace

namespace mean_field::operators {
    PreparedDisplacementResidualOperator::PreparedDisplacementResidualOperator(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const eos::Polytrope &barotrope,
        const context::gravity_field::GravityFieldLinearizationContext &gravityContext
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_gravityContext(gravityContext),
          m_pressureOperator(
              f,
              domainMapper,
              barotrope
          ),
          m_gravityOperator(
              f,
              domainMapper,
              gravityContext
          ),
          m_rotationalOperator(
              f,
              domainMapper
          ) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedDisplacementResidualOperator requires a mesh.");

        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.displacementFes != nullptr && m_fem.gravityFluxFes != nullptr &&
                m_fem.enthalpyFes != nullptr,
            "PreparedDisplacementResidualOperator requires density, "
            "displacement, gravity-gradient, and enthalpy finite-element "
            "spaces."
        );

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "PreparedDisplacementResidualOperator received a mapper with "
            "the wrong dimension."
        );
    }

    PreparedDisplacementResidualReport PreparedDisplacementResidualOperator::Prepare(
        const DisplacementResidualStateView &state,
        const DisplacementResidualDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        validate_shared_gravity_revisions(m_gravityContext, dependencies);

        if (m_isPrepared) {
            /*
             * GravityFieldLinearizationContext currently tracks revisions
             * but not semantic identities. Require an identity replacement
             * to be accompanied by a visible revision change so it cannot
             * silently reuse the old shared density, geometry, or flux.
             */
            validate_shared_identity_transition(
                m_preparedDependencies.discretization, dependencies.discretization,
                "A new displacement-residual discretization identity must "
                "also change the shared gravity revision."
            );

            validate_shared_identity_transition(
                m_preparedDependencies.density, dependencies.density,
                "A new displacement-residual density identity must also "
                "change the shared gravity revision."
            );

            validate_shared_identity_transition(
                m_preparedDependencies.displacement, dependencies.displacement,
                "A new displacement-residual displacement identity must "
                "also change the shared gravity revision."
            );

            validate_shared_identity_transition(
                m_preparedDependencies.gravityGradient, dependencies.gravityGradient,
                "A new displacement-residual gravity-gradient identity "
                "must also change the shared gravity revision."
            );
        }

        const mfem::Vector &density      = m_gravityContext.GetDensity();
        const mfem::Vector &displacement = m_gravityContext.GetGeometryContext().GetDisplacement();

        m_isPrepared                     = false;

        PreparedDisplacementResidualReport report;

        report.pressure = m_pressureOperator.Prepare(
            {.enthalpy = state.enthalpy, .displacement = displacement}, make_pressure_dependencies(dependencies)
        );

        report.gravity  = m_gravityOperator.Prepare();

        report.rotation = m_rotationalOperator.Prepare(
            {.density = density, .displacement = displacement}, make_rotational_dependencies(dependencies), rotation
        );

        if (report.DidAnyChildWork() || m_cachedResidual.Size() != m_fem.displacementFes->GetTrueVSize()) {
            AssembleResidual();
            report.assembledResidual = true;
        }

        MFEM_VERIFY(
            m_cachedResidual.Size() == m_fem.displacementFes->GetTrueVSize(),
            "PreparedDisplacementResidualOperator produced a cached "
            "residual with the wrong size."
        );

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;
        return report;
    }

    void PreparedDisplacementResidualOperator::AssembleResidual() {
        mfem::Vector pressureResidual;
        mfem::Vector gravityResidual;
        mfem::Vector rotationalResidual;

        m_pressureOperator.BuildResidual(pressureResidual);
        m_gravityOperator.BuildResidual(gravityResidual);
        m_rotationalOperator.BuildResidual(rotationalResidual);

        m_cachedResidual = pressureResidual;

        add_compatible(
            m_cachedResidual, gravityResidual,
            "Cannot combine pressure and gravity displacement residuals "
            "with different sizes."
        );

        add_compatible(
            m_cachedResidual, rotationalResidual,
            "Cannot combine mechanical displacement residuals with "
            "different sizes."
        );

        ++m_residualPreparationCount;
    }

    void PreparedDisplacementResidualOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();
        residual = m_cachedResidual;
        ++m_residualApplicationCount;
    }

    void PreparedDisplacementResidualOperator::ApplyDensityJacobianAction(
        const mfem::Vector &densityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        mfem::Vector rotationalAction;

        m_gravityOperator.ApplyDensityJacobianAction(densityVariation, action);

        m_rotationalOperator.ApplyDensityJacobianAction(densityVariation, rotationalAction);

        add_compatible(
            action, rotationalAction,
            "Cannot combine gravity and rotation density-column actions "
            "with different sizes."
        );

        ++m_actionStatistics.densityApplications;
    }

    void PreparedDisplacementResidualOperator::ApplyDisplacementJacobianAction(
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        mfem::Vector gravityAction;
        mfem::Vector rotationalAction;

        m_pressureOperator.ApplyDisplacementJacobianAction(displacementVariation, action);

        m_gravityOperator.ApplyDisplacementJacobianAction(displacementVariation, gravityAction);

        m_rotationalOperator.ApplyDisplacementJacobianAction(displacementVariation, rotationalAction);

        add_compatible(
            action, gravityAction,
            "Cannot combine pressure and gravity displacement-column "
            "actions with different sizes."
        );

        add_compatible(
            action, rotationalAction,
            "Cannot combine mechanical displacement-column actions with "
            "different sizes."
        );

        ++m_actionStatistics.displacementApplications;
    }

    void PreparedDisplacementResidualOperator::ApplyGravityGradientJacobianAction(
        const mfem::Vector &gravityGradientVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_gravityOperator.ApplyGravityGradientJacobianAction(gravityGradientVariation, action);

        ++m_actionStatistics.gravityGradientApplications;
    }

    void PreparedDisplacementResidualOperator::ApplyEnthalpyJacobianAction(
        const mfem::Vector &enthalpyVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_pressureOperator.ApplyEnthalpyJacobianAction(enthalpyVariation, action);

        ++m_actionStatistics.enthalpyApplications;
    }

    void PreparedDisplacementResidualOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        const mfem::Vector &gravityGradientVariation,
        const mfem::Vector &enthalpyVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        mfem::Vector gravityAction;
        mfem::Vector rotationalAction;

        m_pressureOperator.ApplyCompleteJacobianAction(enthalpyVariation, displacementVariation, action);

        m_gravityOperator.ApplyCompleteJacobianAction(
            densityVariation, displacementVariation, gravityGradientVariation, gravityAction
        );

        m_rotationalOperator.ApplyCompleteJacobianAction(densityVariation, displacementVariation, rotationalAction);

        add_compatible(
            action, gravityAction,
            "Cannot combine pressure and gravity complete Jacobian "
            "actions with different sizes."
        );

        add_compatible(
            action, rotationalAction,
            "Cannot combine mechanical complete Jacobian actions with "
            "different sizes."
        );

        ++m_actionStatistics.densityApplications;
        ++m_actionStatistics.displacementApplications;
        ++m_actionStatistics.gravityGradientApplications;
        ++m_actionStatistics.enthalpyApplications;
        ++m_actionStatistics.completeApplications;
    }

    bool PreparedDisplacementResidualOperator::IsPrepared() const noexcept {
        if (!m_isPrepared || !m_pressureOperator.IsPrepared() || !m_gravityOperator.IsPrepared() ||
            !m_rotationalOperator.IsPrepared() || !m_gravityContext.IsPrepared()) {
            return false;
        }

        const context::gravity_field::GravityFieldRevisions &gravityRevisions = m_gravityContext.GetRevisions();

        return gravityRevisions.discretization.value == m_preparedDependencies.discretization.revision &&
               gravityRevisions.density.value == m_preparedDependencies.density.revision &&
               gravityRevisions.displacement.value == m_preparedDependencies.displacement.revision &&
               gravityRevisions.gravity_gradient.value == m_preparedDependencies.gravityGradient.revision;
    }

    std::uint64_t PreparedDisplacementResidualOperator::GetResidualPreparationCount() const noexcept {
        return m_residualPreparationCount;
    }

    std::uint64_t PreparedDisplacementResidualOperator::GetResidualApplicationCount() const noexcept {
        return m_residualApplicationCount;
    }

    const PreparedDisplacementResidualActionStatistics &
    PreparedDisplacementResidualOperator::GetActionStatistics() const noexcept {
        return m_actionStatistics;
    }

    const PreparedPressureForceOperator &PreparedDisplacementResidualOperator::GetPressureOperator() const noexcept {
        return m_pressureOperator;
    }

    const PreparedGravityDisplacementForceOperator &
    PreparedDisplacementResidualOperator::GetGravityOperator() const noexcept {
        return m_gravityOperator;
    }

    const PreparedRotationalDisplacementForceOperator &
    PreparedDisplacementResidualOperator::GetRotationalOperator() const noexcept {
        return m_rotationalOperator;
    }

    const fem::FEM &PreparedDisplacementResidualOperator::GetFEM() const noexcept {
        return m_fem;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    PreparedDisplacementResidualOperator::GetGravityContext() const noexcept {
        return m_gravityContext;
    }

    void PreparedDisplacementResidualOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedDisplacementResidualOperator must be prepared for "
                          "the current shared gravity-context revisions before residual "
                          "or Jacobian application."
        );
    }

    PreparedDisplacementResidualJacobianOperator::PreparedDisplacementResidualJacobianOperator(
        const DisplacementResidualLayout &layout,
        const PreparedDisplacementResidualOperator &preparedOperator
    )
        : mfem::Operator(
              layout.residual_offsets().Last(),
              layout.value_offsets().Last()
          ),
          m_layout(layout),
          m_preparedOperator(preparedOperator) {
        const fem::FEM &f = m_preparedOperator.GetFEM();

        MFEM_VERIFY(
            f.densityFes != nullptr && f.displacementFes != nullptr && f.gravityFluxFes != nullptr &&
                f.gravityPotentialFes != nullptr && f.enthalpyFes != nullptr,
            "Prepared displacement-residual MFEM adapter requires every "
            "finite-element space in the barotropic equilibrium layout."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto gravityPotentialValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);

        constexpr auto enthalpyValue =
            utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);

        constexpr auto barotropicConstantValue =
            utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        constexpr auto gravityGradientResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto gravityPotentialResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);

        constexpr auto densityResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto enthalpyResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);

        constexpr auto massResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);

        MFEM_VERIFY(
            m_layout.size(densityValue) == f.densityFes->GetTrueVSize() &&
                m_layout.size(displacementValue) == f.displacementFes->GetTrueVSize() &&
                m_layout.size(gravityGradientValue) == f.gravityFluxFes->GetTrueVSize() &&
                m_layout.size(gravityPotentialValue) == f.gravityPotentialFes->GetTrueVSize() &&
                m_layout.size(barotropicConstantValue) == 1,
            "Prepared displacement-residual MFEM adapter received "
            "incompatible barotropic value-block sizes."
        );

        MFEM_VERIFY(
            m_layout.size(enthalpyValue) == m_preparedOperator.GetPressureOperator().GetEnthalpySize(),
            "Prepared displacement-residual MFEM adapter received an "
            "incompatible enthalpy value block."
        );

        MFEM_VERIFY(
            m_layout.size(gravityGradientResidual) == f.gravityFluxFes->GetTrueVSize() &&
                m_layout.size(gravityPotentialResidual) == f.gravityPotentialFes->GetTrueVSize() &&
                m_layout.size(densityResidual) == f.densityFes->GetTrueVSize() &&
                m_layout.size(displacementResidual) == f.displacementFes->GetTrueVSize() &&
                m_layout.size(enthalpyResidual) == f.enthalpyFes->GetTrueVSize() && m_layout.size(massResidual) == 1,
            "Prepared displacement-residual MFEM adapter received "
            "incompatible barotropic residual-block sizes."
        );

        MFEM_VERIFY(
            Height() == m_layout.residual_offsets().Last() && Width() == m_layout.value_offsets().Last(),
            "Prepared displacement-residual MFEM adapter has inconsistent "
            "operator dimensions."
        );
    }

    void PreparedDisplacementResidualJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(), "Prepared displacement-residual MFEM adapter requires a "
                                             "prepared row operator."
        );

        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared displacement-residual MFEM adapter received a "
                                         "direction with the wrong size."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto enthalpyValue =
            utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        const mfem::Vector densityVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(densityValue), m_layout.size(densityValue)
        );

        const mfem::Vector displacementVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(displacementValue),
            m_layout.size(displacementValue)
        );

        const mfem::Vector gravityGradientVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(gravityGradientValue),
            m_layout.size(gravityGradientValue)
        );

        const mfem::Vector enthalpyVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(enthalpyValue),
            m_layout.size(enthalpyValue)
        );

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(
            densityVariation, displacementVariation, gravityGradientVariation, enthalpyVariation, displacementAction
        );

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared displacement-residual MFEM adapter produced an "
            "action with the wrong size."
        );

        action.SetSize(Height());
        action                   = 0.0;

        const int residualOffset = m_layout.offset(displacementResidual);

        for (int entry = 0; entry < displacementAction.Size(); ++entry) {
            action(residualOffset + entry) = displacementAction(entry);
        }
    }

    const DisplacementResidualLayout &PreparedDisplacementResidualJacobianOperator::GetLayout() const noexcept {
        return m_layout;
    }
} // namespace mean_field::operators