module;

#include <algorithm>
#include <cmath>
#include <compare>
#include <cstdint>
#include <optional>
#include <utility>

#include <mfem.hpp>

export module mean_field:operators.prepared_central_density;

export import :field.mfem;
export import :model.compiled_fixed_central_density;

export namespace mean_field::operators {
    struct CentralDensityDependencyStamp final {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        constexpr auto operator<=>(const CentralDensityDependencyStamp &) const = default;
    };

    struct CentralDensityDependencies final {
        CentralDensityDependencyStamp enthalpy;

        constexpr auto operator<=>(const CentralDensityDependencies &) const = default;
    };

    struct PreparedCentralDensityReport final {
        bool refreshedCentralEnthalpy{false};
        bool refreshedBorder{false};
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return refreshedCentralEnthalpy || refreshedBorder || assembledResidual;
        }

        constexpr auto operator<=>(const PreparedCentralDensityReport &) const = default;
    };

    struct CentralDensityConstraintReport final {
        double targetDensity;
        double achievedDensity;
        double targetEnthalpy;
        double achievedEnthalpy;
        double enthalpyResidual;
        double scaledResidual;
    };

    struct CentralDensityJacobianInput final {
        const mfem::Vector &enthalpyVariation;
        double borderVariation;
    };

    struct CentralDensityJacobianOutput final {
        mfem::Vector &enthalpyAction;
        mfem::Vector &phaseAction;
    };

    struct CentralDensityJacobianTransposeInput final {
        const mfem::Vector &enthalpyResidualDual;
        double phaseResidualDual;
    };

    struct CentralDensityJacobianTransposeOutput final {
        mfem::Vector &enthalpyDual;
        mfem::Vector &borderDual;
    };

    /*
     * Bordered central-density phase condition
     *
     *   R_c(h) = h(0) - h(rho_c,target),
     *   R_h    <- R_h + lambda_c e_c.
     *
     * The point functional e_c selects the unique scalar H1 vertex at the
     * computational origin. Its coordinate transpose supplies the border
     * column, so this contribution is algebraically symmetric before any
     * independent scaling is applied by a solver.
     */
    class PreparedCentralDensityConstraint final {
    public:
        PreparedCentralDensityConstraint(
            field::FieldPointDofMap centerDof,
            const MPI_Comm communicator
        )
            : m_centerDof(std::move(centerDof)),
              m_communicator(communicator) {
        }

        PreparedCentralDensityReport Prepare(
            const models::CompiledFixedCentralDensity &constraint,
            const mfem::Vector &enthalpy,
            const double border,
            const CentralDensityDependencies &dependencies
        ) {
            MFEM_VERIFY(
                enthalpy.Size() == m_centerDof.field_size(),
                "The central-density phase received an enthalpy vector with the wrong size."
            );
            MFEM_VERIFY(std::isfinite(border), "The central-density phase received a non-finite border value.");

            const bool wasPrepared = m_isPrepared;
            PreparedCentralDensityReport report;

            if (!wasPrepared || dependencies.enthalpy != m_preparedDependencies.enthalpy) {
                double localCentralEnthalpy = 0.0;
                for (const int reducedDof : m_centerDof.reduced_dofs()) {
                    const double value = enthalpy(reducedDof);
                    MFEM_VERIFY(std::isfinite(value), "The central enthalpy is non-finite.");
                    localCentralEnthalpy += value;
                }
                m_centralEnthalpy               = GlobalSum(localCentralEnthalpy);
                report.refreshedCentralEnthalpy = true;
            }

            if (!wasPrepared || border != m_border) {
                m_border               = border;
                report.refreshedBorder = true;
            }

            const bool targetChanged =
                !m_constraint.has_value() || constraint.targetDensity() != m_constraint->targetDensity();
            if (targetChanged) {
                m_constraint = constraint;
            }

            if (report.refreshedCentralEnthalpy || report.refreshedBorder || targetChanged) {
                m_cachedPhaseResidual    = m_centralEnthalpy - m_constraint->targetEnthalpy().value();
                report.assembledResidual = true;
            }

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

        void AddResidual(
            mfem::Vector &enthalpyResidual,
            mfem::Vector &phaseResidual
        ) const {
            VerifyPrepared();
            VerifyOutputSizes(enthalpyResidual, phaseResidual);
            for (const int reducedDof : m_centerDof.reduced_dofs()) {
                enthalpyResidual(reducedDof) += m_border;
            }
            phaseResidual(0) = m_cachedPhaseResidual;
        }

        void ApplyJacobian(
            const CentralDensityJacobianInput &input,
            CentralDensityJacobianOutput output
        ) const {
            VerifyPrepared();
            MFEM_VERIFY(
                input.enthalpyVariation.Size() == m_centerDof.field_size(),
                "The central-density Jacobian received an enthalpy direction with the wrong size."
            );
            VerifyOutputSizes(output.enthalpyAction, output.phaseAction);

            double localPhaseAction = 0.0;
            for (const int reducedDof : m_centerDof.reduced_dofs()) {
                output.enthalpyAction(reducedDof) += input.borderVariation;
                localPhaseAction += input.enthalpyVariation(reducedDof);
            }
            output.phaseAction(0) = GlobalSum(localPhaseAction);
            ++m_jacobianApplicationCount;
        }

        void ApplyJacobianTranspose(
            const CentralDensityJacobianTransposeInput &input,
            CentralDensityJacobianTransposeOutput output
        ) const {
            VerifyPrepared();
            MFEM_VERIFY(
                input.enthalpyResidualDual.Size() == m_centerDof.field_size(),
                "The central-density transpose received an enthalpy residual dual with the wrong size."
            );
            MFEM_VERIFY(
                output.enthalpyDual.Size() == m_centerDof.field_size() && output.borderDual.Size() == 1,
                "The central-density transpose received output vectors with the wrong size."
            );

            double localBorderDual = 0.0;
            for (const int reducedDof : m_centerDof.reduced_dofs()) {
                output.enthalpyDual(reducedDof) += input.phaseResidualDual;
                localBorderDual += input.enthalpyResidualDual(reducedDof);
            }
            output.borderDual(0) += GlobalSum(localBorderDual);
            ++m_transposeApplicationCount;
        }

        [[nodiscard]] CentralDensityConstraintReport GetConstraintReport() const {
            VerifyPrepared();
            const double targetEnthalpy = m_constraint->targetEnthalpy().value();
            const double scale          = std::max(std::abs(targetEnthalpy), 1.0e-300);
            return {
                .targetDensity = m_constraint->targetDensity().value(),
                .achievedDensity =
                    m_constraint->densityFromEnthalpy(dimensions::SpecificEnthalpyValue{m_centralEnthalpy}).value(),
                .targetEnthalpy   = targetEnthalpy,
                .achievedEnthalpy = m_centralEnthalpy,
                .enthalpyResidual = m_cachedPhaseResidual,
                .scaledResidual   = m_cachedPhaseResidual / scale
            };
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_isPrepared;
        }

        [[nodiscard]] const field::FieldPointDofMap &GetCenterDof() const noexcept {
            return m_centerDof;
        }

        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept {
            return m_preparationCount;
        }

        [[nodiscard]] std::uint64_t GetJacobianApplicationCount() const noexcept {
            return m_jacobianApplicationCount;
        }

        [[nodiscard]] std::uint64_t GetTransposeApplicationCount() const noexcept {
            return m_transposeApplicationCount;
        }

    private:
        [[nodiscard]] double GlobalSum(const double localValue) const {
            double globalValue = 0.0;
            MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_communicator);
            return globalValue;
        }

        void VerifyOutputSizes(
            const mfem::Vector &enthalpyOutput,
            const mfem::Vector &phaseOutput
        ) const {
            MFEM_VERIFY(
                enthalpyOutput.Size() == m_centerDof.field_size() && phaseOutput.Size() == 1,
                "The central-density phase received output vectors with the wrong size."
            );
        }

        void VerifyPrepared() const {
            MFEM_VERIFY(m_isPrepared, "The central-density phase must be prepared before application.");
        }

        field::FieldPointDofMap m_centerDof;
        MPI_Comm m_communicator;
        std::optional<models::CompiledFixedCentralDensity> m_constraint;
        CentralDensityDependencies m_preparedDependencies;
        double m_centralEnthalpy{0.0};
        double m_border{0.0};
        double m_cachedPhaseResidual{0.0};
        std::uint64_t m_preparationCount{0};
        mutable std::uint64_t m_jacobianApplicationCount{0};
        mutable std::uint64_t m_transposeApplicationCount{0};
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators