module;

#include <array>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:preconditioning.equilibrium_coordinates;

export import :preconditioning.specification_border;

export namespace mean_field::preconditioning {
    namespace detail {
        template <typename CandidateList, typename Universe> struct EquilibriumCoordinateListIsSubset;

        template <typename... Candidates, typename Universe>
        struct EquilibriumCoordinateListIsSubset<utils::blocks::type_list<Candidates...>, Universe>
            : std::bool_constant<(utils::blocks::contains_type_v<Candidates, Universe> && ...)> { };
    } // namespace detail

    template <typename Component, typename Form>
    concept EquilibriumCoordinateComponentFor =
        PreconditionerComponent<Component> && utils::blocks::block_form_is_valid_v<Form> &&
        std::remove_cvref_t<Component>::CorrectionBlocks::size == Form::value_block_count &&
        std::remove_cvref_t<Component>::ResidualBlocks::size == Form::residual_block_count &&
        detail::EquilibriumCoordinateListIsSubset<
            typename std::remove_cvref_t<Component>::CorrectionBlocks,
            typename Form::value_blocks>::value &&
        detail::EquilibriumCoordinateListIsSubset<
            typename std::remove_cvref_t<Component>::ResidualBlocks,
            typename Form::residual_blocks>::value;

    struct EquilibriumCoordinateRange final {
        int equilibriumOffset{0};
        int preconditionerOffset{0};
        int size{0};

        constexpr bool operator==(const EquilibriumCoordinateRange &) const = default;
    };

    struct EquilibriumCoordinateMapStatistics final {
        std::uint64_t residualPacks{0};
        std::uint64_t residualUnpacks{0};
        std::uint64_t correctionPacks{0};
        std::uint64_t correctionUnpacks{0};
    };

    template <typename Form, typename Component>
        requires EquilibriumCoordinateComponentFor<Component, Form>
    class EquilibriumPreconditionerCoordinateMap final {
    private:
        using ComponentType                               = std::remove_cvref_t<Component>;
        using Layout                                      = utils::blocks::form_layout<Form>;

        static constexpr std::size_t correctionBlockCount = ComponentType::CorrectionBlocks::size;
        static constexpr std::size_t residualBlockCount   = ComponentType::ResidualBlocks::size;

    public:
        explicit EquilibriumPreconditionerCoordinateMap(const Layout &layout)
            : m_correctionRanges(MakeCorrectionRanges(
                  layout,
                  typename ComponentType::CorrectionBlocks{}
              )),
              m_residualRanges(MakeResidualRanges(
                  layout,
                  typename ComponentType::ResidualBlocks{}
              )),
              m_equilibriumStateSize(layout.value_offsets().Last()),
              m_equilibriumResidualSize(layout.residual_offsets().Last()),
              m_preconditionerCorrectionSize(TotalSize(m_correctionRanges)),
              m_preconditionerResidualSize(TotalSize(m_residualRanges)) {
            if (m_preconditionerCorrectionSize != m_equilibriumStateSize ||
                m_preconditionerResidualSize != m_equilibriumResidualSize) {
                throw std::logic_error(
                    "The typed preconditioner coordinate map does not span the complete equilibrium operator."
                );
            }
        }

        void PackResidual(
            const mfem::Vector &equilibriumResidual,
            mfem::Vector &preconditionerResidual
        ) const {
            VerifySizes(
                equilibriumResidual, m_equilibriumResidualSize, preconditionerResidual, m_preconditionerResidualSize,
                "residual pack"
            );
            EquilibriumToPreconditioner(equilibriumResidual, preconditionerResidual, m_residualRanges);
            ++m_statistics.residualPacks;
        }

        void UnpackResidual(
            const mfem::Vector &preconditionerResidual,
            mfem::Vector &equilibriumResidual
        ) const {
            VerifySizes(
                preconditionerResidual, m_preconditionerResidualSize, equilibriumResidual, m_equilibriumResidualSize,
                "residual unpack"
            );
            PreconditionerToEquilibrium(preconditionerResidual, equilibriumResidual, m_residualRanges);
            ++m_statistics.residualUnpacks;
        }

        void PackCorrection(
            const mfem::Vector &equilibriumCorrection,
            mfem::Vector &preconditionerCorrection
        ) const {
            VerifySizes(
                equilibriumCorrection, m_equilibriumStateSize, preconditionerCorrection, m_preconditionerCorrectionSize,
                "correction pack"
            );
            EquilibriumToPreconditioner(equilibriumCorrection, preconditionerCorrection, m_correctionRanges);
            ++m_statistics.correctionPacks;
        }

        void UnpackCorrection(
            const mfem::Vector &preconditionerCorrection,
            mfem::Vector &equilibriumCorrection
        ) const {
            VerifySizes(
                preconditionerCorrection, m_preconditionerCorrectionSize, equilibriumCorrection, m_equilibriumStateSize,
                "correction unpack"
            );
            PreconditionerToEquilibrium(preconditionerCorrection, equilibriumCorrection, m_correctionRanges);
            ++m_statistics.correctionUnpacks;
        }

        [[nodiscard]] int EquilibriumStateSize() const noexcept {
            return m_equilibriumStateSize;
        }

        [[nodiscard]] int EquilibriumResidualSize() const noexcept {
            return m_equilibriumResidualSize;
        }

        [[nodiscard]] int PreconditionerCorrectionSize() const noexcept {
            return m_preconditionerCorrectionSize;
        }

        [[nodiscard]] int PreconditionerResidualSize() const noexcept {
            return m_preconditionerResidualSize;
        }

        [[nodiscard]] const std::array<
            EquilibriumCoordinateRange,
            correctionBlockCount> &
        GetCorrectionRanges() const noexcept {
            return m_correctionRanges;
        }

        [[nodiscard]] const std::array<
            EquilibriumCoordinateRange,
            residualBlockCount> &
        GetResidualRanges() const noexcept {
            return m_residualRanges;
        }

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

    private:
        template <typename... Blocks>
        [[nodiscard]] static std::array<
            EquilibriumCoordinateRange,
            sizeof...(Blocks)>
        MakeCorrectionRanges(
            const Layout &layout,
            utils::blocks::type_list<Blocks...>
        ) {
            std::array<EquilibriumCoordinateRange, sizeof...(Blocks)> ranges{};
            int preconditionerOffset = 0;
            std::size_t range        = 0;
            (
                [&] {
                    constexpr int equilibriumBlock = utils::blocks::type_index_v<Blocks, typename Form::value_blocks>;
                    const int size                 = layout.size(utils::blocks::value_block<equilibriumBlock>{});
                    ranges[range++]                = {
                        .equilibriumOffset    = layout.offset(utils::blocks::value_block<equilibriumBlock>{}),
                        .preconditionerOffset = preconditionerOffset,
                        .size                 = size
                    };
                    preconditionerOffset += size;
                }(),
                ...);
            return ranges;
        }

        template <typename... Blocks>
        [[nodiscard]] static std::array<
            EquilibriumCoordinateRange,
            sizeof...(Blocks)>
        MakeResidualRanges(
            const Layout &layout,
            utils::blocks::type_list<Blocks...>
        ) {
            std::array<EquilibriumCoordinateRange, sizeof...(Blocks)> ranges{};
            int preconditionerOffset = 0;
            std::size_t range        = 0;
            (
                [&] {
                    constexpr int equilibriumBlock =
                        utils::blocks::type_index_v<Blocks, typename Form::residual_blocks>;
                    const int size  = layout.size(utils::blocks::residual_block<equilibriumBlock>{});
                    ranges[range++] = {
                        .equilibriumOffset    = layout.offset(utils::blocks::residual_block<equilibriumBlock>{}),
                        .preconditionerOffset = preconditionerOffset,
                        .size                 = size
                    };
                    preconditionerOffset += size;
                }(),
                ...);
            return ranges;
        }

        template <std::size_t Size>
        [[nodiscard]] static int TotalSize(
            const std::array<
                EquilibriumCoordinateRange,
                Size> &ranges
        ) noexcept {
            int size = 0;
            for (const auto &range : ranges) {
                size += range.size;
            }
            return size;
        }

        template <std::size_t Size>
        static void EquilibriumToPreconditioner(
            const mfem::Vector &equilibrium,
            mfem::Vector &preconditioner,
            const std::array<
                EquilibriumCoordinateRange,
                Size> &ranges
        ) {
            for (const auto &range : ranges) {
                for (int index = 0; index < range.size; ++index) {
                    preconditioner(range.preconditionerOffset + index) = equilibrium(range.equilibriumOffset + index);
                }
            }
        }

        template <std::size_t Size>
        static void PreconditionerToEquilibrium(
            const mfem::Vector &preconditioner,
            mfem::Vector &equilibrium,
            const std::array<
                EquilibriumCoordinateRange,
                Size> &ranges
        ) {
            for (const auto &range : ranges) {
                for (int index = 0; index < range.size; ++index) {
                    equilibrium(range.equilibriumOffset + index) = preconditioner(range.preconditionerOffset + index);
                }
            }
        }

        static void VerifySizes(
            const mfem::Vector &source,
            const int expectedSourceSize,
            const mfem::Vector &destination,
            const int expectedDestinationSize,
            const char *operation
        ) {
            if (source.Size() != expectedSourceSize || destination.Size() != expectedDestinationSize) {
                throw std::invalid_argument(
                    std::string("The equilibrium preconditioner ") + operation + " received an incompatible vector."
                );
            }
        }

        std::array<EquilibriumCoordinateRange, correctionBlockCount> m_correctionRanges;
        std::array<EquilibriumCoordinateRange, residualBlockCount> m_residualRanges;
        int m_equilibriumStateSize;
        int m_equilibriumResidualSize;
        int m_preconditionerCorrectionSize;
        int m_preconditionerResidualSize;
        mutable EquilibriumCoordinateMapStatistics m_statistics;
    };

    struct PreparedStellarPreconditionerStatistics final {
        std::uint64_t applications{0};
        std::uint64_t residualCoordinateMappings{0};
        std::uint64_t correctionCoordinateMappings{0};
    };

    template <equilibrium::DiscretizedStellarEquilibriumProblem Problem, SpecificationBorderBlockType Block>
        requires EquilibriumCoordinateComponentFor<Block, typename std::remove_cvref_t<Problem>::FormType>
    class PreparedStellarPreconditioner final : public mfem::Solver {
    private:
        using ProblemType = std::remove_cvref_t<Problem>;
        using BlockType   = std::remove_cvref_t<Block>;

    public:
        using Form                  = typename ProblemType::FormType;
        using BackendType           = typename BlockType::BackendType;
        using GroupedPreconditioner = PreparedSpecificationBorderBlock<ProblemType, BlockType>;
        using CoordinateMap         = EquilibriumPreconditionerCoordinateMap<Form, BlockType>;

        PreparedStellarPreconditioner(
            const ProblemType &problem,
            BlockType block
        )
            : mfem::Solver(problem.StateSize()),
              m_grouped(
                  problem,
                  std::move(block)
              ),
              m_coordinates(problem.GetManifest().layout()),
              m_groupedResidual(m_coordinates.PreconditionerResidualSize()),
              m_groupedCorrection(m_coordinates.PreconditionerCorrectionSize()) {
            if (problem.StateSize() != problem.EquationSize() || m_grouped.Height() != problem.StateSize() ||
                m_grouped.Width() != problem.EquationSize()) {
                throw std::logic_error(
                    "The prepared stellar preconditioner is incompatible with the complete equilibrium operator."
                );
            }
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Height() || operation.Width() != Width()) {
                throw std::invalid_argument(
                    "The prepared stellar preconditioner received an incompatible equilibrium operator."
                );
            }
            m_grouped.SetOperator(operation);
        }

        void Mult(
            const mfem::Vector &equilibriumResidual,
            mfem::Vector &equilibriumCorrection
        ) const override {
            if (equilibriumResidual.Size() != Width() || equilibriumCorrection.Size() != Height()) {
                throw std::invalid_argument(
                    "The prepared stellar preconditioner requires compatible, preallocated equilibrium vectors."
                );
            }
            m_coordinates.PackResidual(equilibriumResidual, m_groupedResidual);
            ++m_statistics.residualCoordinateMappings;
            m_grouped.Mult(m_groupedResidual, m_groupedCorrection);
            m_coordinates.UnpackCorrection(m_groupedCorrection, equilibriumCorrection);
            ++m_statistics.correctionCoordinateMappings;
            ++m_statistics.applications;
        }

        [[nodiscard]] SpecificationBorderBlockPreparationReport Refresh() {
            return m_grouped.Refresh();
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_grouped.IsCurrent();
        }

        [[nodiscard]] const BlockType &GetBlock() const noexcept {
            return m_grouped.GetBlock();
        }

        [[nodiscard]] const GroupedPreconditioner &GetGroupedPreconditioner() const noexcept {
            return m_grouped;
        }

        [[nodiscard]] const CoordinateMap &GetCoordinateMap() const noexcept {
            return m_coordinates;
        }

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

    private:
        GroupedPreconditioner m_grouped;
        CoordinateMap m_coordinates;
        mutable mfem::Vector m_groupedResidual;
        mutable mfem::Vector m_groupedCorrection;
        mutable PreparedStellarPreconditionerStatistics m_statistics;
    };

    template <
        equilibrium::DiscretizedStellarEquilibriumProblem Problem,
        SpecificationBorderBlockType Block>
        requires EquilibriumCoordinateComponentFor<
            Block,
            typename std::remove_cvref_t<Problem>::FormType>
    [[nodiscard]] auto prepare(
        const Problem &problem,
        Block block
    ) {
        return PreparedStellarPreconditioner<Problem, Block>{problem, std::move(block)};
    }
} // namespace mean_field::preconditioning