module;

#include <concepts>
#include <type_traits>

export module mean_field:normalization.plan;

export import :utils.blocks;

export namespace mean_field::normalization {
    struct NormalizationPrescriptionTag { };

    template <typename Candidate>
    concept NormalizationPrescription =
        std::derived_from<
            std::remove_cvref_t<Candidate>,
            NormalizationPrescriptionTag>;

    enum class CoordinateKind { value, residual };

    enum class RieszTopology {
        identity,
        scalar_volume_l2,
        vector_volume_l2,
        scalar_boundary_l2,
        hybrid_scalar_volume_point_rows,
        global_scalar
    };

    enum class PhysicalScaleKind {
        dimensionless,
        density,
        length,
        acceleration,
        inverse_time_squared,
        specific_energy,
        pressure,
        mass,
        force,
        angular_velocity,
        angular_momentum
    };

    struct IdentityCoordinate final { };

    /*
     * Honest compile-time method for a coordinate whose positive diagonal
     * factor is supplied at runtime by one exact normalization prescription.
     * Unlike IdentityCoordinate, this category makes no claim about the
     * numerical value of that factor.  The owner type prevents one policy from
     * silently presenting another policy's runtime map as its own plan.
     */
    template <NormalizationPrescription Prescription>
    struct RuntimePreparedCoordinate final {
        using PrescriptionType = std::remove_cvref_t<Prescription>;
    };

    template <RieszTopology Topology, PhysicalScaleKind Scale> struct PhysicalRieszCoordinate final {
        static constexpr RieszTopology topology       = Topology;
        static constexpr PhysicalScaleKind scale      = Scale;
    };

    struct UnsupportedPhysicalRieszCoordinate final { };

    template <typename Block> struct PhysicalRieszBlockTraits {
        using Method                         = UnsupportedPhysicalRieszCoordinate;
        static constexpr bool registered     = false;
    };

    template <CoordinateKind Kind, typename BlockList, typename MethodType>
    struct CoordinateComponent final {
        using Blocks                         = BlockList;
        using Method                         = MethodType;
        static constexpr CoordinateKind kind = Kind;

        using ValueBlocks = std::conditional_t<
            Kind == CoordinateKind::value,
            BlockList,
            utils::blocks::type_list<>>;
        using ResidualBlocks = std::conditional_t<
            Kind == CoordinateKind::residual,
            BlockList,
            utils::blocks::type_list<>>;
    };

    namespace detail {
        template <typename Candidate> struct IsTypeList : std::false_type { };

        template <typename... Types>
        struct IsTypeList<utils::blocks::type_list<Types...>> : std::true_type { };

        template <typename List, typename Base> struct IsUniqueDerivedBlockList : std::false_type { };

        template <typename Base, typename... Blocks>
        struct IsUniqueDerivedBlockList<utils::blocks::type_list<Blocks...>, Base>
            : std::bool_constant<
                  (std::derived_from<Blocks, Base> && ...) &&
                  utils::blocks::types_are_unique_v<utils::blocks::type_list<Blocks...>>> { };

        template <typename Method> struct IsCoordinateMethod : std::false_type { };

        template <> struct IsCoordinateMethod<IdentityCoordinate> : std::true_type { };

        template <NormalizationPrescription Prescription>
        struct IsCoordinateMethod<RuntimePreparedCoordinate<Prescription>>
            : std::true_type { };

        template <RieszTopology Topology, PhysicalScaleKind Scale>
        struct IsCoordinateMethod<PhysicalRieszCoordinate<Topology, Scale>> : std::true_type { };

        template <typename Method, typename Block> struct MethodSupportsBlock : std::false_type { };

        template <typename Block>
        struct MethodSupportsBlock<IdentityCoordinate, Block>
            : std::bool_constant<std::derived_from<Block, utils::blocks::block>> { };

        template <NormalizationPrescription Prescription, typename Block>
        struct MethodSupportsBlock<RuntimePreparedCoordinate<Prescription>, Block>
            : std::bool_constant<std::derived_from<Block, utils::blocks::block>> { };

        template <RieszTopology Topology, PhysicalScaleKind Scale, typename Block>
        struct MethodSupportsBlock<PhysicalRieszCoordinate<Topology, Scale>, Block>
            : std::bool_constant<
                  PhysicalRieszBlockTraits<Block>::registered &&
                  std::same_as<
                      typename PhysicalRieszBlockTraits<Block>::Method,
                      PhysicalRieszCoordinate<Topology, Scale>>> { };

        template <typename Method, typename List> struct MethodSupportsEveryBlock : std::false_type { };

        template <typename Method, typename... Blocks>
        struct MethodSupportsEveryBlock<Method, utils::blocks::type_list<Blocks...>>
            : std::bool_constant<(MethodSupportsBlock<Method, Blocks>::value && ...)> { };

        template <typename Candidate, typename = void> struct ComponentTraits {
            static constexpr bool valid = false;
        };

        template <typename Candidate>
        struct ComponentTraits<
            Candidate,
            std::void_t<
                typename Candidate::Blocks,
                typename Candidate::Method,
                typename Candidate::ValueBlocks,
                typename Candidate::ResidualBlocks,
                decltype(Candidate::kind)>> {
            using Blocks         = typename Candidate::Blocks;
            using Method         = typename Candidate::Method;
            using ValueBlocks    = typename Candidate::ValueBlocks;
            using ResidualBlocks = typename Candidate::ResidualBlocks;

            static constexpr bool hasValidKind =
                std::same_as<std::remove_cv_t<decltype(Candidate::kind)>, CoordinateKind>;

            static constexpr bool hasValidBlockList = [] {
                if constexpr (!hasValidKind || !IsTypeList<Blocks>::value) {
                    return false;
                } else if constexpr (Candidate::kind == CoordinateKind::value) {
                    return IsUniqueDerivedBlockList<Blocks, utils::blocks::value_block_base>::value;
                } else if constexpr (Candidate::kind == CoordinateKind::residual) {
                    return IsUniqueDerivedBlockList<Blocks, utils::blocks::residual_block_base>::value;
                } else {
                    return false;
                }
            }();

            static constexpr bool hasCoherentCoordinateLists = [] {
                if constexpr (!hasValidKind || !IsTypeList<ValueBlocks>::value ||
                              !IsTypeList<ResidualBlocks>::value) {
                    return false;
                } else if constexpr (Candidate::kind == CoordinateKind::value) {
                    return std::same_as<ValueBlocks, Blocks> &&
                           std::same_as<ResidualBlocks, utils::blocks::type_list<>>;
                } else if constexpr (Candidate::kind == CoordinateKind::residual) {
                    return std::same_as<ValueBlocks, utils::blocks::type_list<>> &&
                           std::same_as<ResidualBlocks, Blocks>;
                } else {
                    return false;
                }
            }();

            static constexpr bool valid = hasValidKind && IsTypeList<Blocks>::value &&
                                          IsCoordinateMethod<Method>::value && hasValidBlockList &&
                                          hasCoherentCoordinateLists &&
                                          MethodSupportsEveryBlock<Method, Blocks>::value;
        };

        template <typename... Lists> struct Concatenate;

        template <> struct Concatenate<> {
            using Type = utils::blocks::type_list<>;
        };

        template <typename... Types> struct Concatenate<utils::blocks::type_list<Types...>> {
            using Type = utils::blocks::type_list<Types...>;
        };

        template <typename... Left, typename... Right, typename... Remaining>
        struct Concatenate<utils::blocks::type_list<Left...>, utils::blocks::type_list<Right...>, Remaining...> {
            using Type = typename Concatenate<utils::blocks::type_list<Left..., Right...>, Remaining...>::Type;
        };

        template <typename... Lists> using ConcatenateT = typename Concatenate<Lists...>::Type;

        template <typename List, typename Type> struct Append;

        template <typename... Types, typename Appended>
        struct Append<utils::blocks::type_list<Types...>, Appended> {
            using Type = utils::blocks::type_list<Types..., Appended>;
        };

        template <typename List, typename Type> using AppendT = typename Append<List, Type>::Type;

        template <typename List, typename Type>
        using AppendUniqueT = std::conditional_t<
            utils::blocks::contains_type_v<Type, List>,
            List,
            AppendT<List, Type>>;

        template <typename Source, typename Excluded> struct ListDifference;

        template <typename Excluded>
        struct ListDifference<utils::blocks::type_list<>, Excluded> {
            using Type = utils::blocks::type_list<>;
        };

        template <typename Head, typename... Tail, typename Excluded>
        struct ListDifference<utils::blocks::type_list<Head, Tail...>, Excluded> {
        private:
            using Remaining = typename ListDifference<utils::blocks::type_list<Tail...>, Excluded>::Type;

        public:
            using Type = std::conditional_t<
                utils::blocks::contains_type_v<Head, Excluded>,
                Remaining,
                ConcatenateT<utils::blocks::type_list<Head>, Remaining>>;
        };

        template <typename Source, typename Excluded>
        using ListDifferenceT = typename ListDifference<Source, Excluded>::Type;

        template <typename Remaining, typename Original, typename Repeated> struct CollectRepeatedTypes;

        template <typename Original, typename Repeated>
        struct CollectRepeatedTypes<utils::blocks::type_list<>, Original, Repeated> {
            using Type = Repeated;
        };

        template <typename Head, typename... Tail, typename Original, typename Repeated>
        struct CollectRepeatedTypes<utils::blocks::type_list<Head, Tail...>, Original, Repeated> {
        private:
            using Next = std::conditional_t<
                (utils::blocks::type_count_v<Head, Original> > 1),
                AppendUniqueT<Repeated, Head>,
                Repeated>;

        public:
            using Type = typename CollectRepeatedTypes<utils::blocks::type_list<Tail...>, Original, Next>::Type;
        };

        template <typename List>
        using RepeatedTypesT = typename CollectRepeatedTypes<
            List,
            List,
            utils::blocks::type_list<>>::Type;

        template <typename Candidate, typename = void> struct PlanTraits {
            static constexpr bool valid = false;
        };
    } // namespace detail

    template <typename Candidate>
    concept NormalizationComponent = detail::ComponentTraits<std::remove_cvref_t<Candidate>>::valid;

    template <typename... Components> struct NormalizationPlan final {
        using ComponentTypes  = utils::blocks::type_list<Components...>;
        using ValueBlocks     = detail::ConcatenateT<typename Components::ValueBlocks...>;
        using ResidualBlocks  = detail::ConcatenateT<typename Components::ResidualBlocks...>;
    };

    namespace detail {
        template <typename... Components>
        struct PlanTraits<NormalizationPlan<Components...>> {
            static constexpr bool valid = (ComponentTraits<Components>::valid && ...);
        };

        template <typename Values, typename Residuals> struct MakeIdentityPlan;

        template <typename... Values, typename... Residuals>
        struct MakeIdentityPlan<
            utils::blocks::type_list<Values...>,
            utils::blocks::type_list<Residuals...>> {
            using Type = NormalizationPlan<
                CoordinateComponent<CoordinateKind::value, utils::blocks::type_list<Values>, IdentityCoordinate>...,
                CoordinateComponent<
                    CoordinateKind::residual,
                    utils::blocks::type_list<Residuals>,
                    IdentityCoordinate>...>;
        };

        template <
            NormalizationPrescription Prescription,
            typename Values,
            typename Residuals>
        struct MakeRuntimePreparedPlan;

        template <
            NormalizationPrescription Prescription,
            typename... Values,
            typename... Residuals>
        struct MakeRuntimePreparedPlan<
            Prescription,
            utils::blocks::type_list<Values...>,
            utils::blocks::type_list<Residuals...>> {
            using Method = RuntimePreparedCoordinate<Prescription>;
            using Type = NormalizationPlan<
                CoordinateComponent<
                    CoordinateKind::value,
                    utils::blocks::type_list<Values>,
                    Method>...,
                CoordinateComponent<
                    CoordinateKind::residual,
                    utils::blocks::type_list<Residuals>,
                    Method>...>;
        };
    } // namespace detail

    template <typename Candidate>
    concept NormalizationPlanType = detail::PlanTraits<std::remove_cvref_t<Candidate>>::valid;

    template <typename Form>
        requires utils::blocks::block_form_is_valid_v<Form>
    using IdentityNormalizationPlanFor = typename detail::MakeIdentityPlan<
        typename Form::value_blocks,
        typename Form::residual_blocks>::Type;

    template <NormalizationPrescription Prescription, typename Form>
        requires utils::blocks::block_form_is_valid_v<Form>
    using RuntimePreparedNormalizationPlanFor =
        typename detail::MakeRuntimePreparedPlan<
            std::remove_cvref_t<Prescription>,
            typename Form::value_blocks,
            typename Form::residual_blocks>::Type;

    template <typename Form, typename Plan>
        requires utils::blocks::block_form_is_valid_v<Form>
    struct NormalizationCoverage final {
        using DeclaredValueBlocks = typename Plan::ValueBlocks;
        using DeclaredResidualBlocks = typename Plan::ResidualBlocks;

        using MissingValueBlocks = detail::ListDifferenceT<typename Form::value_blocks, DeclaredValueBlocks>;
        using UnexpectedValueBlocks = detail::ListDifferenceT<DeclaredValueBlocks, typename Form::value_blocks>;
        using RepeatedValueBlocks = detail::RepeatedTypesT<DeclaredValueBlocks>;

        using MissingResidualBlocks = detail::ListDifferenceT<typename Form::residual_blocks, DeclaredResidualBlocks>;
        using UnexpectedResidualBlocks = detail::ListDifferenceT<DeclaredResidualBlocks, typename Form::residual_blocks>;
        using RepeatedResidualBlocks = detail::RepeatedTypesT<DeclaredResidualBlocks>;

        static constexpr bool hasEveryValueBlock = MissingValueBlocks::size == 0;
        static constexpr bool hasOnlyValueBlocks = UnexpectedValueBlocks::size == 0;
        static constexpr bool hasUniqueValueOwners = RepeatedValueBlocks::size == 0;
        static constexpr bool hasEveryResidualBlock = MissingResidualBlocks::size == 0;
        static constexpr bool hasOnlyResidualBlocks = UnexpectedResidualBlocks::size == 0;
        static constexpr bool hasUniqueResidualOwners = RepeatedResidualBlocks::size == 0;

        static constexpr bool complete = hasEveryValueBlock && hasOnlyValueBlocks && hasUniqueValueOwners &&
                                         hasEveryResidualBlock && hasOnlyResidualBlocks &&
                                         hasUniqueResidualOwners;
    };

    template <typename Plan, typename Form>
    concept CompleteNormalizationFor = utils::blocks::block_form_is_valid_v<Form> &&
                                       NormalizationPlanType<Plan> &&
                                       NormalizationCoverage<Form, std::remove_cvref_t<Plan>>::complete;
} // namespace mean_field::normalization