module;

#include <concepts>
#include <cstdint>
#include <type_traits>
#include <utility>

export module mean_field:preconditioning.backend;

export namespace mean_field::preconditioning {
    enum class OperatorCategory : std::uint8_t {
        identity,
        mass_like,
        elliptic_like,
        surface_like,
        mixed,
        dense_border
    };

    enum class OperatorValueStructure : std::uint8_t { scalar, vector, block };

    enum class OperatorSymmetry : std::uint8_t { symmetric, nonsymmetric };

    enum class OperatorDefiniteness : std::uint8_t {
        positive_definite,
        positive_semidefinite,
        indefinite,
        unspecified
    };

    enum class OperatorRepresentation : std::uint8_t { none, diagonal, matrix_free, assembled_sparse, assembled_dense };
    enum class OperatorDistribution : std::uint8_t { not_applicable, local, distributed_true_dof };
    enum class OperatorFESpace : std::uint8_t { not_applicable, h1, h_curl, h_div, l2, product };
    enum class OperatorNullspace : std::uint8_t { none, constant_mode, supplied_basis };
    enum class ApplicationContract : std::uint8_t { stationary_linear, flexible };
    enum class SymmetryRequirement : std::uint8_t { none, symmetric };
    enum class NullspaceRequirement : std::uint8_t { none, constant_mode_supported, supplied_basis_required };
    enum class SurrogateRequirement : std::uint8_t { none, diagonal, assembled_dense, assembled_sparse };

    enum class PreparationDependency : std::uint8_t {
        discretization    = 1U << 0U,
        geometry          = 1U << 1U,
        equation_of_state = 1U << 2U,
        linearization     = 1U << 3U
    };

    template <PreparationDependency... Dependencies> struct PreparationDependencies final {
        static constexpr std::uint8_t mask = (std::uint8_t{0} | ... | static_cast<std::uint8_t>(Dependencies));

        [[nodiscard]] static consteval bool contains(const PreparationDependency dependency) noexcept {
            return (mask & static_cast<std::uint8_t>(dependency)) != 0U;
        }
    };

    template <typename Candidate> struct IsPreparationDependencies : std::false_type { };

    template <PreparationDependency... Dependencies>
    struct IsPreparationDependencies<PreparationDependencies<Dependencies...>> : std::true_type { };

    template <typename Candidate>
    concept PreparationDependenciesType = IsPreparationDependencies<std::remove_cvref_t<Candidate>>::value;

    using NoPreparationDependencies     = PreparationDependencies<>;

    template <
        OperatorCategory Category,
        OperatorValueStructure ValueStructure,
        OperatorSymmetry Symmetry,
        OperatorDefiniteness Definiteness,
        OperatorRepresentation Representation,
        OperatorDistribution Distribution,
        OperatorFESpace FESpace     = OperatorFESpace::not_applicable,
        OperatorNullspace Nullspace = OperatorNullspace::none>
    struct OperatorCharacteristics final {
        static constexpr OperatorCategory category             = Category;
        static constexpr OperatorValueStructure valueStructure = ValueStructure;
        static constexpr OperatorSymmetry symmetry             = Symmetry;
        static constexpr OperatorDefiniteness definiteness     = Definiteness;
        static constexpr OperatorRepresentation representation = Representation;
        static constexpr OperatorDistribution distribution     = Distribution;
        static constexpr OperatorFESpace finiteElementSpace    = FESpace;
        static constexpr OperatorNullspace nullspace           = Nullspace;
    };

    template <typename Candidate> struct IsOperatorCharacteristics : std::false_type { };

    template <
        OperatorCategory Category,
        OperatorValueStructure ValueStructure,
        OperatorSymmetry Symmetry,
        OperatorDefiniteness Definiteness,
        OperatorRepresentation Representation,
        OperatorDistribution Distribution,
        OperatorFESpace FESpace,
        OperatorNullspace Nullspace>
    struct IsOperatorCharacteristics<OperatorCharacteristics<
        Category,
        ValueStructure,
        Symmetry,
        Definiteness,
        Representation,
        Distribution,
        FESpace,
        Nullspace>> : std::true_type { };

    template <typename Candidate>
    concept OperatorCharacteristicsType   = IsOperatorCharacteristics<std::remove_cvref_t<Candidate>>::value;

    using IdentityOperatorCharacteristics = OperatorCharacteristics<
        OperatorCategory::identity,
        OperatorValueStructure::block,
        OperatorSymmetry::symmetric,
        OperatorDefiniteness::positive_definite,
        OperatorRepresentation::none,
        OperatorDistribution::not_applicable>;

    namespace backend {
        struct FixedCycles final {
            int cycles{1};
        };

        struct SolveToTolerance final {
            double relativeTolerance{1.0e-8};
            int maximumCycles{100};
        };

        template <typename Candidate> struct ApplicationModeTraits {
            static constexpr bool registered = false;
        };

        template <> struct ApplicationModeTraits<FixedCycles> {
            static constexpr bool registered                         = true;
            static constexpr ApplicationContract applicationContract = ApplicationContract::stationary_linear;
        };

        template <> struct ApplicationModeTraits<SolveToTolerance> {
            static constexpr bool registered                         = true;
            static constexpr ApplicationContract applicationContract = ApplicationContract::flexible;
        };

        template <typename Candidate>
        concept ApplicationMode = ApplicationModeTraits<std::remove_cvref_t<Candidate>>::registered;

        struct Identity final { };
        struct Diagonal final { };
        struct DenseDirect final { };

        struct MatrixFreeChebyshev final {
            int order{2};
            int powerIterations{10};
            double powerTolerance{1.0e-8};
            int powerSeed{12345};
        };

        template <ApplicationMode Mode = FixedCycles> struct HypreBoomerAMG final {
            using ApplicationModeType = Mode;

            Mode application{};

            constexpr HypreBoomerAMG() = default;

            constexpr explicit HypreBoomerAMG(Mode applicationMode) : application(std::move(applicationMode)) {
            }
        };

        template <ApplicationMode Mode> HypreBoomerAMG(Mode) -> HypreBoomerAMG<Mode>;

        template <typename Candidate> struct Traits {
            static constexpr bool registered = false;

            using PreparationDependencies    = NoPreparationDependencies;
        };

        template <> struct Traits<Identity> {
            static constexpr bool registered                           = true;
            static constexpr ApplicationContract applicationContract   = ApplicationContract::stationary_linear;
            static constexpr bool supportsSerialExecution              = true;
            static constexpr bool supportsDistributedExecution         = true;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::none;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::none;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::none;
            static constexpr bool requiresAssembledSparseSurrogate     = false;

            using PreparationDependencies                              = NoPreparationDependencies;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports = Characteristics::category == OperatorCategory::identity &&
                                             Characteristics::representation == OperatorRepresentation::none;
        };

        template <> struct Traits<Diagonal> {
            static constexpr bool registered                           = true;
            static constexpr ApplicationContract applicationContract   = ApplicationContract::stationary_linear;
            static constexpr bool supportsSerialExecution              = true;
            static constexpr bool supportsDistributedExecution         = true;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::symmetric;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::none;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::diagonal;
            static constexpr bool requiresAssembledSparseSurrogate     = false;

            using PreparationDependencies =
                preconditioning::PreparationDependencies<PreparationDependency::linearization>;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports =
                (Characteristics::category == OperatorCategory::mass_like ||
                 Characteristics::category == OperatorCategory::elliptic_like ||
                 Characteristics::category == OperatorCategory::surface_like) &&
                Characteristics::symmetry == OperatorSymmetry::symmetric &&
                (Characteristics::representation == OperatorRepresentation::diagonal ||
                 (Characteristics::category == OperatorCategory::mass_like &&
                  Characteristics::representation == OperatorRepresentation::matrix_free)) &&
                Characteristics::nullspace == OperatorNullspace::none &&
                (Characteristics::distribution == OperatorDistribution::local ||
                 Characteristics::distribution == OperatorDistribution::distributed_true_dof);
        };

        template <> struct Traits<MatrixFreeChebyshev> {
            static constexpr bool registered                           = true;
            static constexpr ApplicationContract applicationContract   = ApplicationContract::stationary_linear;
            static constexpr bool supportsSerialExecution              = true;
            static constexpr bool supportsDistributedExecution         = true;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::symmetric;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::none;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::diagonal;
            static constexpr bool requiresAssembledSparseSurrogate     = false;

            using PreparationDependencies                              = preconditioning::PreparationDependencies<
                PreparationDependency::discretization,
                PreparationDependency::geometry,
                PreparationDependency::linearization>;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports =
                Characteristics::category == OperatorCategory::mass_like &&
                Characteristics::symmetry == OperatorSymmetry::symmetric &&
                Characteristics::definiteness == OperatorDefiniteness::positive_definite &&
                Characteristics::representation == OperatorRepresentation::matrix_free &&
                Characteristics::finiteElementSpace == OperatorFESpace::h_div &&
                Characteristics::nullspace == OperatorNullspace::none &&
                (Characteristics::distribution == OperatorDistribution::local ||
                 Characteristics::distribution == OperatorDistribution::distributed_true_dof);
        };

        template <> struct Traits<DenseDirect> {
            static constexpr bool registered                           = true;
            static constexpr ApplicationContract applicationContract   = ApplicationContract::stationary_linear;
            static constexpr bool supportsSerialExecution              = true;
            static constexpr bool supportsDistributedExecution         = false;
            static constexpr SymmetryRequirement symmetryRequirement   = SymmetryRequirement::none;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::none;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::assembled_dense;
            static constexpr bool requiresAssembledSparseSurrogate     = false;

            using PreparationDependencies =
                preconditioning::PreparationDependencies<PreparationDependency::linearization>;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports =
                Characteristics::category == OperatorCategory::dense_border &&
                Characteristics::representation == OperatorRepresentation::assembled_dense &&
                Characteristics::distribution == OperatorDistribution::local &&
                Characteristics::nullspace == OperatorNullspace::none;
        };

        template <ApplicationMode Mode> struct Traits<HypreBoomerAMG<Mode>> {
            static constexpr bool registered                         = true;
            static constexpr ApplicationContract applicationContract = ApplicationModeTraits<Mode>::applicationContract;
            static constexpr bool supportsSerialExecution            = false;
            static constexpr bool supportsDistributedExecution       = true;
            static constexpr SymmetryRequirement symmetryRequirement = SymmetryRequirement::symmetric;
            static constexpr NullspaceRequirement nullspaceRequirement = NullspaceRequirement::constant_mode_supported;
            static constexpr SurrogateRequirement surrogateRequirement = SurrogateRequirement::assembled_sparse;
            static constexpr bool requiresAssembledSparseSurrogate     = true;

            using PreparationDependencies                              = preconditioning::PreparationDependencies<
                PreparationDependency::discretization,
                PreparationDependency::geometry,
                PreparationDependency::equation_of_state,
                PreparationDependency::linearization>;

            template <OperatorCharacteristicsType Characteristics>
            static constexpr bool supports =
                Characteristics::category == OperatorCategory::elliptic_like &&
                Characteristics::valueStructure == OperatorValueStructure::scalar &&
                Characteristics::symmetry == OperatorSymmetry::symmetric &&
                (Characteristics::definiteness == OperatorDefiniteness::positive_definite ||
                 Characteristics::definiteness == OperatorDefiniteness::positive_semidefinite) &&
                Characteristics::representation == OperatorRepresentation::assembled_sparse &&
                Characteristics::distribution == OperatorDistribution::distributed_true_dof &&
                (Characteristics::finiteElementSpace == OperatorFESpace::h1 ||
                 Characteristics::finiteElementSpace == OperatorFESpace::l2) &&
                (Characteristics::nullspace == OperatorNullspace::none ||
                 Characteristics::nullspace == OperatorNullspace::constant_mode);
        };

        template <typename Candidate>
        concept Registered = Traits<std::remove_cvref_t<Candidate>>::registered;

        namespace detail {
            template <
                typename Backend,
                typename Characteristics,
                bool = Registered<Backend> && OperatorCharacteristicsType<Characteristics>>
            struct IsCompatible : std::false_type { };

            template <typename Backend, typename Characteristics>
            struct IsCompatible<Backend, Characteristics, true>
                : std::bool_constant<
                      Traits<std::remove_cvref_t<Backend>>::template supports<std::remove_cvref_t<Characteristics>>> {
            };
        } // namespace detail

        template <typename Backend, typename Characteristics>
        inline constexpr bool isCompatible = detail::IsCompatible<Backend, Characteristics>::value;

        template <typename Backend, typename Characteristics>
        concept Compatible = isCompatible<Backend, Characteristics>;

        template <Registered Backend>
        inline constexpr ApplicationContract applicationContract =
            Traits<std::remove_cvref_t<Backend>>::applicationContract;

        template <typename Backend>
        concept ArnoldiAdmissible = Registered<Backend> && applicationContract<std::remove_cvref_t<Backend>> ==
                                                               ApplicationContract::stationary_linear;

        template <Registered Backend>
        inline constexpr bool requiresAssembledSparseSurrogate =
            Traits<std::remove_cvref_t<Backend>>::requiresAssembledSparseSurrogate;
    } // namespace backend
} // namespace mean_field::preconditioning