module;

#include <array>
#include <cmath>
#include <cstdint>
#include <memory>
#include <utility>

#include <mfem.hpp>

module mean_field;

import :operators.prepared_central_density_stellar_equilibrium;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using PhysicalForm = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;
    using BorderedForm = mean_field::operators::CentralDensityStellarEquilibriumForm;

    [[nodiscard]] std::array<
        int,
        BorderedForm::value_block_count>
    make_value_sizes(const mean_field::operators::StellarEquilibriumLayout &physicalLayout) {
        std::array<int, BorderedForm::value_block_count> sizes{};
        for (int block = 0; block < PhysicalForm::value_block_count; ++block) {
            sizes[block] = physicalLayout.value_offsets()[block + 1] - physicalLayout.value_offsets()[block];
        }
        sizes[PhysicalForm::value_block_count] = 1;
        return sizes;
    }

    [[nodiscard]] std::array<
        int,
        BorderedForm::residual_block_count>
    make_residual_sizes(const mean_field::operators::StellarEquilibriumLayout &physicalLayout) {
        std::array<int, BorderedForm::residual_block_count> sizes{};
        for (int block = 0; block < PhysicalForm::residual_block_count; ++block) {
            sizes[block] = physicalLayout.residual_offsets()[block + 1] - physicalLayout.residual_offsets()[block];
        }
        sizes[PhysicalForm::residual_block_count] = 1;
        return sizes;
    }

    [[nodiscard]] mean_field::operators::CentralDensityDependencies
    make_phase_dependencies(const mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        return {.enthalpy = {.identity = dependencies.enthalpy.identity, .revision = dependencies.enthalpy.revision}};
    }

    void validate_finite_scalar(
        const double value,
        const char *message
    ) {
        MFEM_VERIFY(std::isfinite(value), message);
    }
} // namespace

namespace mean_field::operators {
    field::FieldPointDofMap PreparedCentralDensityStellarEquilibriumOperator::MakeCenterDofMap(const fem::FEM &f) {
        MFEM_VERIFY(
            f.mesh != nullptr && f.enthalpyFes != nullptr,
            "The central-density phase requires the mesh and enthalpy finite-element space."
        );
        const field::FieldDofMap enthalpyMap = field::make_field_dof_map<field::Enthalpy, DomainSchema>(*f.enthalpyFes);
        mfem::Vector origin(f.mesh->SpaceDimension());
        origin = 0.0;
        return field::make_field_point_dof_map<field::Enthalpy>(*f.enthalpyFes, enthalpyMap, origin, 1.0e-12);
    }

    PreparedCentralDensityStellarEquilibriumOperator::PreparedCentralDensityStellarEquilibriumOperator(
        fem::FEM &f,
        std::unique_ptr<PreparedStellarEquilibriumOperator> physicalOperator,
        models::CompiledFixedCentralDensity centralDensity,
        field::FieldPointDofMap centerDof
    )
        : mfem::Operator(
              physicalOperator->Height() + 1,
              physicalOperator->Width() + 1
          ),
          m_physicalOperator(std::move(physicalOperator)),
          m_centralDensity(std::move(centralDensity)),
          m_phaseConstraint(
              std::move(centerDof),
              f.mesh->GetComm()
          ),
          m_rootManifest(
              make_value_sizes(m_physicalOperator->GetLayout()),
              make_residual_sizes(m_physicalOperator->GetLayout()),
              m_physicalOperator->GetTargetMass(),
              m_physicalOperator->GetSurfaceConstraintOperator().GetPhysicalCondition().targetPressure,
              m_physicalOperator->GetSurfaceConstraintOperator().GetSurfaceRows().size(),
              CentralDensityManifestInput{
                  .targetDensity  = m_centralDensity.targetDensity().value(),
                  .targetEnthalpy = m_centralDensity.targetEnthalpy().value(),
                  .centerDofCount = 1
              }
          ) {
        MFEM_VERIFY(
            Width() == m_rootManifest.layout().value_offsets().Last() &&
                Height() == m_rootManifest.layout().residual_offsets().Last(),
            "The central-density bordered root has inconsistent dimensions."
        );
    }

    PreparedCentralDensityStellarEquilibriumReport PreparedCentralDensityStellarEquilibriumOperator::Prepare(
        const mfem::Vector &state,
        const StellarEquilibriumDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        MFEM_VERIFY(state.Size() == Width(), "The central-density bordered root received a state with the wrong size.");
        const auto stateView        = m_rootManifest.stateView(state);
        const mfem::Vector enthalpy = stateView.block(utils::blocks::enthalpy_field.specific_term);
        const mfem::Vector border   = stateView.block(utils::blocks::fixed_central_density_phase.central_value_term);
        validate_finite_scalar(border(0), "The central-density bordered root received a non-finite border value.");

        mfem::Vector physicalState(const_cast<mfem::real_t *>(state.GetData()), m_physicalOperator->Width());

        m_isPrepared = false;
        PreparedCentralDensityStellarEquilibriumReport report;
        report.physical = m_physicalOperator->Prepare(physicalState, dependencies, rotation);
        report.phase =
            m_phaseConstraint.Prepare(m_centralDensity, enthalpy, border(0), make_phase_dependencies(dependencies));

        if (report.physical.assembledResidual || report.phase.DidAnyWork() || m_cachedResidual.Size() != Height()) {
            AssembleResidual();
            report.assembledResidual = true;
        }

        m_isPrepared = true;
        return report;
    }

    void PreparedCentralDensityStellarEquilibriumOperator::AssembleResidual() {
        mfem::Vector physicalResidual;
        m_physicalOperator->BuildResidual(physicalResidual);

        m_cachedResidual.SetSize(Height());
        m_cachedResidual = 0.0;
        mfem::Vector physicalDestination(m_cachedResidual.GetData(), physicalResidual.Size());
        physicalDestination           = physicalResidual;

        const auto residualView       = m_rootManifest.residualView(m_cachedResidual);
        mfem::Vector enthalpyResidual = residualView.block(utils::blocks::enthalpy_field.specific_term);
        mfem::Vector phaseResidual = residualView.block(utils::blocks::fixed_central_density_phase.central_value_term);
        m_phaseConstraint.AddResidual(enthalpyResidual, phaseResidual);
    }

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

    void PreparedCentralDensityStellarEquilibriumOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        VerifyPrepared();
        MFEM_VERIFY(
            direction.Size() == Width(), "The central-density bordered root received a direction with the wrong size."
        );
        const auto directionView             = m_rootManifest.directionView(direction);
        const mfem::Vector enthalpyDirection = directionView.block(utils::blocks::enthalpy_field.specific_term);
        const mfem::Vector borderDirection =
            directionView.block(utils::blocks::fixed_central_density_phase.central_value_term);
        validate_finite_scalar(
            borderDirection(0), "The central-density bordered root received a non-finite border direction."
        );

        mfem::Vector physicalDirection(const_cast<mfem::real_t *>(direction.GetData()), m_physicalOperator->Width());
        mfem::Vector physicalAction;
        m_physicalOperator->Mult(physicalDirection, physicalAction);

        action.SetSize(Height());
        action = 0.0;
        mfem::Vector physicalDestination(action.GetData(), physicalAction.Size());
        physicalDestination         = physicalAction;

        const auto actionView       = m_rootManifest.residualView(action);
        mfem::Vector enthalpyAction = actionView.block(utils::blocks::enthalpy_field.specific_term);
        mfem::Vector phaseAction    = actionView.block(utils::blocks::fixed_central_density_phase.central_value_term);
        m_phaseConstraint.ApplyJacobian(
            {.enthalpyVariation = enthalpyDirection, .borderVariation = borderDirection(0)},
            {.enthalpyAction = enthalpyAction, .phaseAction = phaseAction}
        );
    }

    bool PreparedCentralDensityStellarEquilibriumOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_physicalOperator->IsPrepared() && m_phaseConstraint.IsPrepared();
    }

    const CentralDensityStellarEquilibriumLayout &
    PreparedCentralDensityStellarEquilibriumOperator::GetLayout() const noexcept {
        return m_rootManifest.layout();
    }

    const CentralDensityStellarEquilibriumRootManifest &
    PreparedCentralDensityStellarEquilibriumOperator::GetRootManifest() const noexcept {
        return m_rootManifest;
    }

    const PreparedStellarEquilibriumOperator &
    PreparedCentralDensityStellarEquilibriumOperator::GetPhysicalOperator() const noexcept {
        return *m_physicalOperator;
    }

    const PreparedCentralDensityConstraint &
    PreparedCentralDensityStellarEquilibriumOperator::GetCentralDensityConstraint() const noexcept {
        return m_phaseConstraint;
    }

    RootConstraintReport PreparedCentralDensityStellarEquilibriumOperator::GetFixedMassReport() const {
        VerifyPrepared();
        return m_physicalOperator->GetFixedMassReport();
    }

    CentralDensityConstraintReport PreparedCentralDensityStellarEquilibriumOperator::GetCentralDensityReport() const {
        VerifyPrepared();
        return m_phaseConstraint.GetConstraintReport();
    }

    void PreparedCentralDensityStellarEquilibriumOperator::VerifyPrepared() const {
        MFEM_VERIFY(IsPrepared(), "The central-density bordered root must be prepared before application.");
    }
} // namespace mean_field::operators