module;

#include <array>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <span>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include <mfem.hpp>

export module mean_field:normalization.operators;

export import :normalization.physical_riesz;

export namespace mean_field::normalization {
    class DiagonalNormalization final {
    public:
        DiagonalNormalization(
            mfem::Vector stateToNormalized,
            mfem::Vector residualToNormalized
        )
            : m_stateToNormalized(std::move(stateToNormalized)),
              m_residualToNormalized(std::move(residualToNormalized)) {
            ValidateFactors(m_stateToNormalized, "state");
            ValidateFactors(m_residualToNormalized, "residual");
        }

        [[nodiscard]] static DiagonalNormalization Identity(
            const int stateSize,
            const int residualSize
        ) {
            if (stateSize < 0 || residualSize < 0) {
                throw std::invalid_argument("Normalization dimensions cannot be negative.");
            }
            mfem::Vector state(stateSize);
            mfem::Vector residual(residualSize);
            state = 1.0;
            residual = 1.0;
            return {std::move(state), std::move(residual)};
        }

        [[nodiscard]] int StateSize() const noexcept {
            return m_stateToNormalized.Size();
        }

        [[nodiscard]] int ResidualSize() const noexcept {
            return m_residualToNormalized.Size();
        }

        [[nodiscard]] const mfem::Vector &StateFactors() const noexcept {
            return m_stateToNormalized;
        }

        [[nodiscard]] const mfem::Vector &ResidualFactors() const noexcept {
            return m_residualToNormalized;
        }

        void NormalizeState(
            const mfem::Vector &physical,
            mfem::Vector &normalized
        ) const {
            Apply(m_stateToNormalized, physical, normalized, false, "state");
        }

        void DenormalizeState(
            const mfem::Vector &normalized,
            mfem::Vector &physical
        ) const {
            Apply(m_stateToNormalized, normalized, physical, true, "state");
        }

        void NormalizeResidual(
            const mfem::Vector &physical,
            mfem::Vector &normalized
        ) const {
            Apply(m_residualToNormalized, physical, normalized, false, "residual");
        }

        void DenormalizeResidual(
            const mfem::Vector &normalized,
            mfem::Vector &physical
        ) const {
            Apply(m_residualToNormalized, normalized, physical, true, "residual");
        }

        [[nodiscard]] double LocalStateNormSquared(const mfem::Vector &physical) const {
            return LocalNormSquared(m_stateToNormalized, physical, "state");
        }

        [[nodiscard]] double LocalResidualNormSquared(const mfem::Vector &physical) const {
            return LocalNormSquared(m_residualToNormalized, physical, "residual");
        }

    private:
        static void ValidateFactors(
            const mfem::Vector &factors,
            const char *role
        ) {
            for (int index = 0; index < factors.Size(); ++index) {
                if (!std::isfinite(factors(index)) || factors(index) <= 0.0) {
                    throw std::invalid_argument(
                        std::string("The ") + role + " normalization factors must be finite and positive."
                    );
                }
            }
        }

        static void Apply(
            const mfem::Vector &factors,
            const mfem::Vector &input,
            mfem::Vector &output,
            const bool inverse,
            const char *role
        ) {
            if (input.Size() != factors.Size()) {
                throw std::invalid_argument(std::string("The ") + role + " vector has the wrong size.");
            }
            const bool exactAlias = input.GetData() == output.GetData() && input.Size() == output.Size();
            if (!exactAlias) {
                output.SetSize(input.Size());
            }
            for (int index = 0; index < input.Size(); ++index) {
                const double value = input(index);
                output(index) = inverse ? value / factors(index) : factors(index) * value;
            }
        }

        [[nodiscard]] static double LocalNormSquared(
            const mfem::Vector &factors,
            const mfem::Vector &physical,
            const char *role
        ) {
            if (physical.Size() != factors.Size()) {
                throw std::invalid_argument(std::string("The ") + role + " vector has the wrong size.");
            }
            double normSquared = 0.0;
            for (int index = 0; index < physical.Size(); ++index) {
                const double normalized = factors(index) * physical(index);
                normSquared += normalized * normalized;
            }
            return normSquared;
        }

        mfem::Vector m_stateToNormalized;
        mfem::Vector m_residualToNormalized;
    };

    /* Detection-safe public operation for an ordinary third-party runtime
     * policy. The exact policy is recovered from the problem type and must own
     * every method in its compiled plan. Its implementation remains beside
     * the policy and is found by ADL, so adding a normalization family does
     * not edit a library registry or switch. */
    template <typename Problem>
    concept RuntimePreparedNormalizationOperation =
        requires(const std::remove_cvref_t<Problem> &problem) {
            typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType;
            typename std::remove_cvref_t<Problem>::FormType;
            requires RuntimePreparedNormalizationFor<
                typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType,
                typename std::remove_cvref_t<Problem>::FormType>;
            {
                problem.GetNormalizationPrescription()
            } -> std::same_as<const typename std::remove_cvref_t<Problem>::NormalizationPrescriptionType &>;
            {
                prepareStellarNormalization(
                    problem.GetNormalizationPrescription(),
                    problem)
            } -> std::same_as<DiagonalNormalization>;
        };

    template <typename Form>
        requires utils::blocks::block_form_is_valid_v<Form>
    class DiagonalNormalizationBuilder final {
    public:
        explicit DiagonalNormalizationBuilder(const utils::blocks::form_layout<Form> &layout)
            : m_layout(&layout),
              m_stateFactors(layout.value_offsets().Last()),
              m_residualFactors(layout.residual_offsets().Last()) {
        }

        explicit DiagonalNormalizationBuilder(
            utils::blocks::form_layout<Form> &&
        ) = delete;

        explicit DiagonalNormalizationBuilder(
            const utils::blocks::form_layout<Form> &&
        ) = delete;

        template <typename Block>
            requires utils::blocks::contains_type_v<Block, typename Form::value_blocks>
        void SetValueBlock(
            const double physicalScale,
            const mfem::Vector &primalGramDiagonal
        ) {
            constexpr int block = utils::blocks::type_index_v<Block, typename Form::value_blocks>;
            RequireUnassigned(m_valueAssigned[block], "value");
            AssignBlock(
                m_stateFactors,
                m_layout->value_offsets()[block],
                m_layout->value_offsets()[block + 1] - m_layout->value_offsets()[block],
                physicalScale,
                primalGramDiagonal,
                false
            );
            m_valueAssigned[block] = true;
        }

        template <typename Block>
            requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
        void SetResidualBlock(
            const double physicalScale,
            const mfem::Vector &primalGramDiagonal
        ) {
            constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
            RequireUnassigned(m_residualAssigned[block], "residual");
            AssignBlock(
                m_residualFactors,
                m_layout->residual_offsets()[block],
                m_layout->residual_offsets()[block + 1] - m_layout->residual_offsets()[block],
                physicalScale,
                primalGramDiagonal,
                true
            );
            m_residualAssigned[block] = true;
        }

        template <typename Block>
            requires utils::blocks::contains_type_v<Block, typename Form::value_blocks>
        void SetValueGlobal(const double physicalScale) {
            constexpr int block = utils::blocks::type_index_v<Block, typename Form::value_blocks>;
            SetConstantMetricValueBlock<Block>(physicalScale, BlockSize(m_layout->value_offsets(), block));
        }

        template <typename Block>
            requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
        void SetResidualGlobal(const double physicalScale) {
            constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
            mfem::Vector metric(BlockSize(m_layout->residual_offsets(), block));
            metric = 1.0;
            SetResidualBlock<Block>(physicalScale, metric);
        }

        template <typename Block>
            requires utils::blocks::contains_type_v<Block, typename Form::residual_blocks>
        void SetHybridResidualBlock(
            const double physicalScale,
            const mfem::Vector &bulkPrimalGramDiagonal,
            const std::span<const int> pointRows,
            const double pointMetric = 1.0
        ) {
            constexpr int block = utils::blocks::type_index_v<Block, typename Form::residual_blocks>;
            const int size = BlockSize(m_layout->residual_offsets(), block);
            if (bulkPrimalGramDiagonal.Size() != size) {
                throw std::invalid_argument("The hybrid residual Gram diagonal has the wrong size.");
            }
            ValidateMetric(pointMetric);

            std::vector<bool> isPointRow(static_cast<std::size_t>(size), false);
            for (const int row : pointRows) {
                if (row < 0 || row >= size) {
                    throw std::out_of_range("A hybrid point row lies outside its residual block.");
                }
                if (isPointRow[static_cast<std::size_t>(row)]) {
                    throw std::invalid_argument("A hybrid point row was supplied more than once.");
                }
                isPointRow[static_cast<std::size_t>(row)] = true;
            }

            mfem::Vector metric(size);
            for (int row = 0; row < size; ++row) {
                metric(row) = isPointRow[static_cast<std::size_t>(row)]
                                  ? pointMetric
                                  : bulkPrimalGramDiagonal(row);
            }
            SetResidualBlock<Block>(physicalScale, metric);
        }

        [[nodiscard]] DiagonalNormalization Build() && {
            for (const bool assigned : m_valueAssigned) {
                if (!assigned) {
                    throw std::logic_error("The normalization is missing a value block.");
                }
            }
            for (const bool assigned : m_residualAssigned) {
                if (!assigned) {
                    throw std::logic_error("The normalization is missing a residual block.");
                }
            }
            return {std::move(m_stateFactors), std::move(m_residualFactors)};
        }

    private:
        template <typename Block>
        void SetConstantMetricValueBlock(
            const double physicalScale,
            const int size
        ) {
            mfem::Vector metric(size);
            metric = 1.0;
            SetValueBlock<Block>(physicalScale, metric);
        }

        [[nodiscard]] static int BlockSize(
            const mfem::Array<int> &offsets,
            const int block
        ) noexcept {
            return offsets[block + 1] - offsets[block];
        }

        static void RequireUnassigned(
            const bool assigned,
            const char *role
        ) {
            if (assigned) {
                throw std::logic_error(std::string("The ") + role + " block normalization was assigned twice.");
            }
        }

        static void ValidateMetric(const double metric) {
            if (!std::isfinite(metric) || metric <= 0.0) {
                throw std::invalid_argument("Every Riesz Gram diagonal entry must be finite and positive.");
            }
        }

        static void AssignBlock(
            mfem::Vector &factors,
            const int offset,
            const int size,
            const double physicalScale,
            const mfem::Vector &primalGramDiagonal,
            const bool dual
        ) {
            if (!std::isfinite(physicalScale) || physicalScale <= 0.0) {
                throw std::invalid_argument("A physical normalization scale must be finite and positive.");
            }
            if (primalGramDiagonal.Size() != size) {
                throw std::invalid_argument("A Riesz Gram diagonal has the wrong block size.");
            }
            for (int index = 0; index < size; ++index) {
                const double metric = primalGramDiagonal(index);
                ValidateMetric(metric);
                const double rieszFactor = std::sqrt(metric);
                const double factor = dual
                                          ? 1.0 / (physicalScale * rieszFactor)
                                          : rieszFactor / physicalScale;
                if (!std::isfinite(factor) || factor <= 0.0) {
                    throw std::overflow_error("A normalization factor is not finite and positive.");
                }
                factors(offset + index) = factor;
            }
        }

        const utils::blocks::form_layout<Form> *m_layout;
        mfem::Vector m_stateFactors;
        mfem::Vector m_residualFactors;
        std::array<bool, Form::value_block_count> m_valueAssigned{};
        std::array<bool, Form::residual_block_count> m_residualAssigned{};
    };

    class ScaledJacobianOperator final : public mfem::Operator {
    public:
        ScaledJacobianOperator(
            const mfem::Operator &physicalJacobian,
            const DiagonalNormalization &normalization
        )
            : mfem::Operator(normalization.ResidualSize(), normalization.StateSize()),
              m_physicalJacobian(&physicalJacobian),
              m_normalization(&normalization),
              m_physicalDirection(normalization.StateSize()),
              m_physicalAction(normalization.ResidualSize()) {
            if (physicalJacobian.Width() != normalization.StateSize() ||
                physicalJacobian.Height() != normalization.ResidualSize()) {
                throw std::invalid_argument("The physical Jacobian and normalization dimensions do not agree.");
            }
        }

        ScaledJacobianOperator(
            mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledJacobianOperator(
            const mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledJacobianOperator(
            const mfem::Operator &,
            DiagonalNormalization &&
        ) = delete;

        ScaledJacobianOperator(
            const mfem::Operator &,
            const DiagonalNormalization &&
        ) = delete;

        void Mult(
            const mfem::Vector &normalizedDirection,
            mfem::Vector &normalizedAction
        ) const override {
            m_normalization->DenormalizeState(normalizedDirection, m_physicalDirection);
            m_physicalJacobian->Mult(m_physicalDirection, m_physicalAction);
            m_normalization->NormalizeResidual(m_physicalAction, normalizedAction);
        }

    private:
        const mfem::Operator *m_physicalJacobian;
        const DiagonalNormalization *m_normalization;
        mutable mfem::Vector m_physicalDirection;
        mutable mfem::Vector m_physicalAction;
    };

    class ScaledInverseOperator final : public mfem::Operator {
    public:
        ScaledInverseOperator(
            const mfem::Operator &physicalInverse,
            const DiagonalNormalization &normalization
        )
            : mfem::Operator(normalization.StateSize(), normalization.ResidualSize()),
              m_physicalInverse(&physicalInverse),
              m_normalization(&normalization),
              m_physicalResidual(normalization.ResidualSize()),
              m_physicalCorrection(normalization.StateSize()) {
            if (physicalInverse.Width() != normalization.ResidualSize() ||
                physicalInverse.Height() != normalization.StateSize()) {
                throw std::invalid_argument("The physical inverse and normalization dimensions do not agree.");
            }
        }

        ScaledInverseOperator(
            mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledInverseOperator(
            const mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledInverseOperator(
            const mfem::Operator &,
            DiagonalNormalization &&
        ) = delete;

        ScaledInverseOperator(
            const mfem::Operator &,
            const DiagonalNormalization &&
        ) = delete;

        void Mult(
            const mfem::Vector &normalizedResidual,
            mfem::Vector &normalizedCorrection
        ) const override {
            m_normalization->DenormalizeResidual(normalizedResidual, m_physicalResidual);
            m_physicalInverse->Mult(m_physicalResidual, m_physicalCorrection);
            m_normalization->NormalizeState(m_physicalCorrection, normalizedCorrection);
        }

    private:
        const mfem::Operator *m_physicalInverse;
        const DiagonalNormalization *m_normalization;
        mutable mfem::Vector m_physicalResidual;
        mutable mfem::Vector m_physicalCorrection;
    };

    struct ScaledPreconditionerStatistics final {
        std::uint64_t operatorBindings{0};
        std::uint64_t applications{0};
    };

    /*
     * Solver-compatible realization of R^{-1} M^{-1} L^{-1}. The wrapped
     * inverse always sees the dimensional Jacobian, even when an MFEM Krylov
     * solver binds this object to the normalized Jacobian L J R.
     */
    class ScaledPreconditioner final : public mfem::Solver {
    public:
        ScaledPreconditioner(
            mfem::Solver &physicalInverse,
            const mfem::Operator &physicalJacobian,
            const mfem::Operator &normalizedJacobian,
            const DiagonalNormalization &normalization
        )
            : mfem::Solver(
                  normalization.StateSize(),
                  normalization.ResidualSize(),
                  physicalInverse.iterative_mode
              ),
              m_physicalInverse(&physicalInverse),
              m_physicalJacobian(&physicalJacobian),
              m_expectedNormalizedJacobian(&normalizedJacobian),
              m_normalization(&normalization),
              m_physicalResidual(normalization.ResidualSize()),
              m_physicalCorrection(normalization.StateSize()) {
            if (physicalInverse.Width() != normalization.ResidualSize() ||
                physicalInverse.Height() != normalization.StateSize() ||
                physicalJacobian.Width() != normalization.StateSize() ||
                physicalJacobian.Height() != normalization.ResidualSize()) {
                throw std::invalid_argument(
                    "The physical preconditioner, Jacobian, and normalization dimensions do not agree."
                );
            }
            SetOperator(normalizedJacobian);
        }

        ScaledPreconditioner(
            mfem::Solver &,
            mfem::Operator &&,
            const mfem::Operator &,
            const DiagonalNormalization &
        ) = delete;

        ScaledPreconditioner(
            mfem::Solver &,
            const mfem::Operator &&,
            const mfem::Operator &,
            const DiagonalNormalization &
        ) = delete;

        ScaledPreconditioner(
            mfem::Solver &,
            const mfem::Operator &,
            mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledPreconditioner(
            mfem::Solver &,
            const mfem::Operator &,
            const mfem::Operator &&,
            const DiagonalNormalization &
        ) = delete;

        ScaledPreconditioner(
            mfem::Solver &,
            const mfem::Operator &,
            const mfem::Operator &,
            DiagonalNormalization &&
        ) = delete;

        ScaledPreconditioner(
            mfem::Solver &,
            const mfem::Operator &,
            const mfem::Operator &,
            const DiagonalNormalization &&
        ) = delete;

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

        void SetOperator(const mfem::Operator &normalizedJacobian) override {
            if (normalizedJacobian.Width() != Width() || normalizedJacobian.Height() != Height()) {
                throw std::invalid_argument(
                    "The scaled preconditioner received an incompatible normalized Jacobian."
                );
            }
            if (&normalizedJacobian != m_expectedNormalizedJacobian) {
                throw std::invalid_argument(
                    "The scaled preconditioner cannot be rebound to a different normalized Jacobian."
                );
            }
            m_physicalInverse->SetOperator(*m_physicalJacobian);
            m_normalizedJacobian = &normalizedJacobian;
            ++m_statistics.operatorBindings;
        }

        void Mult(
            const mfem::Vector &normalizedResidual,
            mfem::Vector &normalizedCorrection
        ) const override {
            if (m_normalizedJacobian == nullptr) {
                throw std::logic_error("The scaled preconditioner has not been bound to a normalized Jacobian.");
            }
            if (normalizedResidual.Size() != Width() || normalizedCorrection.Size() != Height()) {
                throw std::invalid_argument(
                    "The scaled preconditioner requires compatible, preallocated normalized vectors."
                );
            }
            m_normalization->DenormalizeResidual(normalizedResidual, m_physicalResidual);
            m_physicalInverse->Mult(m_physicalResidual, m_physicalCorrection);
            m_normalization->NormalizeState(m_physicalCorrection, normalizedCorrection);
            ++m_statistics.applications;
        }

        [[nodiscard]] const mfem::Solver &GetPhysicalInverse() const noexcept {
            return *m_physicalInverse;
        }

        [[nodiscard]] const mfem::Operator &GetPhysicalJacobian() const noexcept {
            return *m_physicalJacobian;
        }

        [[nodiscard]] const mfem::Operator &GetNormalizedJacobian() const {
            if (m_normalizedJacobian == nullptr) {
                throw std::logic_error("The scaled preconditioner has not been bound to a normalized Jacobian.");
            }
            return *m_normalizedJacobian;
        }

        [[nodiscard]] const ScaledPreconditionerStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        mfem::Solver *m_physicalInverse;
        const mfem::Operator *m_physicalJacobian;
        const mfem::Operator *m_expectedNormalizedJacobian;
        const mfem::Operator *m_normalizedJacobian{nullptr};
        const DiagonalNormalization *m_normalization;
        mutable mfem::Vector m_physicalResidual;
        mutable mfem::Vector m_physicalCorrection;
        mutable ScaledPreconditionerStatistics m_statistics;
    };
} // namespace mean_field::normalization