module;

#include <concepts>
#include <tuple>
#include <type_traits>
#include <utility>

export module mean_field:preconditioning.plan;

export import :preconditioning.backend;
export import :utils.blocks;

export namespace mean_field::preconditioning {
    template <typename ResidualBlock, typename CorrectionBlock> struct Coupling final {
        using Residual   = ResidualBlock;
        using Correction = CorrectionBlock;
    };

    template <
        typename CorrectionBlockList,
        typename ResidualBlockList,
        typename RequiredCouplingList,
        typename Characteristics,
        typename Backend,
        typename PreparationRequirements =
            typename backend::Traits<std::remove_cvref_t<Backend>>::PreparationDependencies>
    struct ComponentDeclaration {
        using CorrectionBlocks        = CorrectionBlockList;
        using ResidualBlocks          = ResidualBlockList;
        using RequiredCouplings       = RequiredCouplingList;
        using OperatorDescription     = Characteristics;
        using BackendType             = Backend;
        using PreparationDependencies = PreparationRequirements;
    };

    template <typename CorrectionBlock, typename ResidualBlock>
        requires std::derived_from<CorrectionBlock, utils::blocks::value_block_base> &&
                 std::derived_from<ResidualBlock, utils::blocks::residual_block_base>
    struct IdentityBlock final {
        using CorrectionBlocks        = utils::blocks::type_list<CorrectionBlock>;
        using ResidualBlocks          = utils::blocks::type_list<ResidualBlock>;
        using RequiredCouplings       = utils::blocks::type_list<>;
        using OperatorDescription     = IdentityOperatorCharacteristics;
        using BackendType             = backend::Identity;
        using PreparationDependencies = NoPreparationDependencies;
    };

    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 Candidate>
        inline constexpr bool isTypeList = IsTypeList<std::remove_cvref_t<Candidate>>::value;

        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 Candidate> struct IsCoupling : std::false_type { };

        template <typename ResidualBlock, typename CorrectionBlock>
        struct IsCoupling<Coupling<ResidualBlock, CorrectionBlock>>
            : std::bool_constant<
                  std::derived_from<ResidualBlock, utils::blocks::residual_block_base> &&
                  std::derived_from<CorrectionBlock, utils::blocks::value_block_base>> { };

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

        template <typename... Couplings>
        struct IsCouplingList<utils::blocks::type_list<Couplings...>>
            : std::bool_constant<
                  (IsCoupling<Couplings>::value && ...) &&
                  utils::blocks::types_are_unique_v<utils::blocks::type_list<Couplings...>>> { };

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

        template <typename Candidate>
        struct ComponentTraits<
            Candidate,
            std::void_t<
                typename Candidate::CorrectionBlocks,
                typename Candidate::ResidualBlocks,
                typename Candidate::RequiredCouplings,
                typename Candidate::OperatorDescription,
                typename Candidate::BackendType,
                typename Candidate::PreparationDependencies>> {
            using CorrectionBlocks               = typename Candidate::CorrectionBlocks;
            using ResidualBlocks                 = typename Candidate::ResidualBlocks;
            using RequiredCouplings              = typename Candidate::RequiredCouplings;
            using OperatorDescription            = typename Candidate::OperatorDescription;
            using BackendType                    = typename Candidate::BackendType;
            using PreparationDependencies        = typename Candidate::PreparationDependencies;
            using BackendPreparationDependencies = typename backend::Traits<BackendType>::PreparationDependencies;

            static constexpr bool valid =
                IsUniqueDerivedBlockList<CorrectionBlocks, utils::blocks::value_block_base>::value &&
                IsUniqueDerivedBlockList<ResidualBlocks, utils::blocks::residual_block_base>::value &&
                IsCouplingList<RequiredCouplings>::value && OperatorCharacteristicsType<OperatorDescription> &&
                backend::Registered<BackendType> && backend::isCompatible<BackendType, OperatorDescription> &&
                PreparationDependenciesType<PreparationDependencies> &&
                ((PreparationDependencies::mask & BackendPreparationDependencies::mask) ==
                 BackendPreparationDependencies::mask);
        };

        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 Type> struct Append<utils::blocks::type_list<Types...>, Type> {
            using Result = utils::blocks::type_list<Types..., Type>;
        };

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

        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 <bool AllowsOverlap, typename... Components> class PlanStorage {
        public:
            using ComponentTypes    = utils::blocks::type_list<Components...>;
            using CorrectionBlocks  = ConcatenateT<typename ComponentTraits<Components>::CorrectionBlocks...>;
            using ResidualBlocks    = ConcatenateT<typename ComponentTraits<Components>::ResidualBlocks...>;
            using RequiredCouplings = ConcatenateT<typename ComponentTraits<Components>::RequiredCouplings...>;

            static constexpr bool allowsOverlappingOwnership = AllowsOverlap;
            static constexpr bool stationaryLinear =
                ((backend::applicationContract<typename ComponentTraits<Components>::BackendType> ==
                  ApplicationContract::stationary_linear) &&
                 ...);

            constexpr explicit PlanStorage(Components... components) : m_components(std::move(components)...) {
            }

            template <typename Component> [[nodiscard]] constexpr const Component &component() const noexcept {
                return std::get<Component>(m_components);
            }

            [[nodiscard]] constexpr const std::tuple<Components...> &components() const noexcept {
                return m_components;
            }

        private:
            std::tuple<Components...> m_components;
        };

        template <
            bool ComponentsAreValid,
            typename DeclaredCorrectionBlocks,
            typename DeclaredResidualBlocks,
            typename DeclaredCouplings,
            typename... Components>
        struct CoherentPlanDeclaration : std::false_type { };

        template <
            typename DeclaredCorrectionBlocks,
            typename DeclaredResidualBlocks,
            typename DeclaredCouplings,
            typename... Components>
        struct CoherentPlanDeclaration<
            true,
            DeclaredCorrectionBlocks,
            DeclaredResidualBlocks,
            DeclaredCouplings,
            Components...>
            : std::bool_constant<
                  std::same_as<
                      DeclaredCorrectionBlocks,
                      ConcatenateT<typename ComponentTraits<Components>::CorrectionBlocks...>> &&
                  std::same_as<
                      DeclaredResidualBlocks,
                      ConcatenateT<typename ComponentTraits<Components>::ResidualBlocks...>> &&
                  std::same_as<
                      DeclaredCouplings,
                      ConcatenateT<typename ComponentTraits<Components>::RequiredCouplings...>>> { };

        template <
            typename ComponentList,
            typename DeclaredCorrectionBlocks,
            typename DeclaredResidualBlocks,
            typename DeclaredCouplings>
        struct PlanDeclarationIsCoherent : std::false_type { };

        template <
            typename... Components,
            typename DeclaredCorrectionBlocks,
            typename DeclaredResidualBlocks,
            typename DeclaredCouplings>
        struct PlanDeclarationIsCoherent<
            utils::blocks::type_list<Components...>,
            DeclaredCorrectionBlocks,
            DeclaredResidualBlocks,
            DeclaredCouplings>
            : CoherentPlanDeclaration<
                  (ComponentTraits<Components>::valid && ...),
                  DeclaredCorrectionBlocks,
                  DeclaredResidualBlocks,
                  DeclaredCouplings,
                  Components...> { };

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

        template <typename Candidate>
        struct PlanTraits<
            Candidate,
            std::void_t<
                typename Candidate::ComponentTypes,
                typename Candidate::CorrectionBlocks,
                typename Candidate::ResidualBlocks,
                typename Candidate::RequiredCouplings>> {
            static constexpr bool valid =
                isTypeList<typename Candidate::ComponentTypes> && isTypeList<typename Candidate::CorrectionBlocks> &&
                isTypeList<typename Candidate::ResidualBlocks> && isTypeList<typename Candidate::RequiredCouplings> &&
                PlanDeclarationIsCoherent<
                    typename Candidate::ComponentTypes,
                    typename Candidate::CorrectionBlocks,
                    typename Candidate::ResidualBlocks,
                    typename Candidate::RequiredCouplings>::value;
        };

        template <typename CouplingList, typename JacobianForm> struct CouplingsExistInJacobian;

        template <typename JacobianForm>
        struct CouplingsExistInJacobian<utils::blocks::type_list<>, JacobianForm> : std::true_type { };

        template <typename Residual, typename Correction, typename... Remaining, typename JacobianForm>
        struct CouplingsExistInJacobian<
            utils::blocks::type_list<Coupling<Residual, Correction>, Remaining...>,
            JacobianForm>
            : std::bool_constant<
                  utils::blocks::has_jacobian_coupling_v<Residual, Correction, JacobianForm> &&
                  CouplingsExistInJacobian<utils::blocks::type_list<Remaining...>, JacobianForm>::value> { };

        template <typename Plan, typename JacobianForm, bool = PlanTraits<std::remove_cvref_t<Plan>>::valid>
        struct RequiredCouplingsExist : std::false_type { };

        template <typename Plan, typename JacobianForm>
        struct RequiredCouplingsExist<Plan, JacobianForm, true>
            : CouplingsExistInJacobian<typename std::remove_cvref_t<Plan>::RequiredCouplings, JacobianForm> { };
    } // namespace detail

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

    template <PreconditionerComponent... Components>
    class PreconditionerPlan final : public detail::PlanStorage<false, Components...> {
        using Base = detail::PlanStorage<false, Components...>;

    public:
        using Base::Base;
    };

    template <typename... Components> PreconditionerPlan(Components...) -> PreconditionerPlan<Components...>;

    template <PreconditionerComponent... Components>
    class OverlappingPreconditionerPlan final : public detail::PlanStorage<true, Components...> {
        using Base = detail::PlanStorage<true, Components...>;

    public:
        using Base::Base;
    };

    template <typename... Components>
    OverlappingPreconditionerPlan(Components...) -> OverlappingPreconditionerPlan<Components...>;

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

    template <typename Form, typename Plan>
        requires utils::blocks::block_form_is_valid_v<Form> && PreconditionerPlanType<Plan>
    struct PreconditionerCoverage final {
        using DeclaredCorrectionBlocks = typename Plan::CorrectionBlocks;
        using DeclaredResidualBlocks   = typename Plan::ResidualBlocks;

        using MissingCorrectionBlocks  = detail::ListDifferenceT<typename Form::value_blocks, DeclaredCorrectionBlocks>;
        using UnexpectedCorrectionBlocks =
            detail::ListDifferenceT<DeclaredCorrectionBlocks, typename Form::value_blocks>;
        using RepeatedCorrectionBlocks = detail::RepeatedTypesT<DeclaredCorrectionBlocks>;

        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 hasEveryCorrectionBlock = MissingCorrectionBlocks::size == 0;
        static constexpr bool hasOnlyCorrectionBlocks = UnexpectedCorrectionBlocks::size == 0;
        static constexpr bool hasUniqueCorrectionOwners =
            Plan::allowsOverlappingOwnership || RepeatedCorrectionBlocks::size == 0;

        static constexpr bool hasEveryResidualBlock = MissingResidualBlocks::size == 0;
        static constexpr bool hasOnlyResidualBlocks = UnexpectedResidualBlocks::size == 0;
        static constexpr bool hasUniqueResidualOwners =
            Plan::allowsOverlappingOwnership || RepeatedResidualBlocks::size == 0;

        static constexpr bool complete = hasEveryCorrectionBlock && hasOnlyCorrectionBlocks &&
                                         hasUniqueCorrectionOwners && hasEveryResidualBlock && hasOnlyResidualBlocks &&
                                         hasUniqueResidualOwners;
    };

    template <typename Plan, typename Form>
    concept CompletePreconditionerFor = utils::blocks::block_form_is_valid_v<Form> && PreconditionerPlanType<Plan> &&
                                        PreconditionerCoverage<Form, std::remove_cvref_t<Plan>>::complete;

    template <typename Plan, typename JacobianForm>
    inline constexpr bool requiredCouplingsExist = detail::RequiredCouplingsExist<Plan, JacobianForm>::value;

    template <typename Plan, typename Form, typename JacobianForm>
    concept CompatiblePreconditionerFor =
        CompletePreconditionerFor<Plan, Form> && utils::blocks::valid_jacobian_form<Form, JacobianForm> &&
        requiredCouplingsExist<Plan, JacobianForm>;

    template <typename Plan>
    concept StationaryLinearPreconditionerPlan =
        PreconditionerPlanType<Plan> && std::remove_cvref_t<Plan>::stationaryLinear;
} // namespace mean_field::preconditioning