module;

#include <array>
#include <cmath>
#include <compare>
#include <concepts>
#include <cstddef>
#include <format>
#include <span>
#include <stdexcept>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>

export module mean_field:model.specifications;

export import :eos.polytrope;
export import :surface.constant;

export namespace mean_field::models {
    enum class SpecificationRole {
        constitutive_law,
        boundary_condition,
        invariant,
        phase_condition,
        gauge_choice,
        rotation_law
    };

    struct SpecificationKey final {
        SpecificationRole role;
        std::size_t ordinal;

        constexpr auto operator<=>(const SpecificationKey &) const = default;
    };

    struct SpecificationDescriptor final {
        std::string_view name;
        SpecificationRole role;
        SpecificationKey key;
        std::size_t generatedValueArity;
        std::size_t generatedResidualArity;

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

    enum class EquilibriumSystemCompilation {
        complete_equilibrium_system,
        equation_contributions_only,

        // Transitional spellings retained while internal solver code is
        // migrated to physics-facing equilibrium-system terminology.
        isolated_root = complete_equilibrium_system,
        assembly_only = equation_contributions_only
    };

    using ModelCompilationClass = EquilibriumSystemCompilation;

    struct RuntimeSpecificationDescriptor final {
        SpecificationDescriptor specification;
        std::size_t canonicalIndex;
        bool hasRootCompiler;

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

    template <typename Candidate> struct SpecificationTraits;

    template <typename Candidate>
    concept ModelSpecification = requires {
        typename std::remove_cvref_t<Candidate>::Parameters;
        { SpecificationTraits<std::remove_cvref_t<Candidate>>::name } -> std::convertible_to<std::string_view>;
        { SpecificationTraits<std::remove_cvref_t<Candidate>>::role } -> std::convertible_to<SpecificationRole>;
        { SpecificationTraits<std::remove_cvref_t<Candidate>>::key } -> std::convertible_to<SpecificationKey>;
    } && std::constructible_from<std::remove_cvref_t<Candidate>, typename std::remove_cvref_t<Candidate>::Parameters>;

    class FixedTotalMass final {
    public:
        struct Parameters final {
            dimensions::MassValue Mtotal;
        };

        using TargetValue = dimensions::MassValue;

        explicit FixedTotalMass(const Parameters parameters) : FixedTotalMass(parameters.Mtotal) {
        }

        explicit FixedTotalMass(const TargetValue targetMass) : m_targetMass(targetMass) {
            if (!std::isfinite(targetMass.value()) || targetMass.value() <= 0.0) {
                throw std::invalid_argument(
                    std::format(
                        "The fixed total mass must be finite and positive. Instead M = {} was provided.",
                        targetMass.value()
                    )
                );
            }
        }

        [[nodiscard]] TargetValue targetMass() const noexcept {
            return m_targetMass;
        }

    private:
        TargetValue m_targetMass;
    };

    class FixedCentralDensity final {
    public:
        struct Parameters final {
            dimensions::DensityValue RhoC;
        };

        using TargetValue = dimensions::DensityValue;

        explicit FixedCentralDensity(const Parameters parameters) : FixedCentralDensity(parameters.RhoC) {
        }

        explicit FixedCentralDensity(const TargetValue targetDensity) : m_targetDensity(targetDensity) {
            if (!std::isfinite(targetDensity.value()) || targetDensity.value() <= 0.0) {
                throw std::invalid_argument(
                    std::format(
                        "The fixed central density must be finite and positive. Instead rho_c = {} was provided.",
                        targetDensity.value()
                    )
                );
            }
        }

        [[nodiscard]] TargetValue targetDensity() const noexcept {
            return m_targetDensity;
        }

    private:
        TargetValue m_targetDensity;
    };

    template <> struct SpecificationTraits<eos::Polytrope> {
        static constexpr std::string_view name  = "Polytrope";
        static constexpr SpecificationRole role = SpecificationRole::constitutive_law;
        static constexpr SpecificationKey key{role, 0};
    };

    template <> struct SpecificationTraits<surface::ConstantPressureSurface> {
        static constexpr std::string_view name  = "IsobaricSurface";
        static constexpr SpecificationRole role = SpecificationRole::boundary_condition;
        static constexpr SpecificationKey key{role, 0};
    };

    template <> struct SpecificationTraits<FixedTotalMass> {
        static constexpr std::string_view name  = "FixedTotalMass";
        static constexpr SpecificationRole role = SpecificationRole::invariant;
        static constexpr SpecificationKey key{role, 0};
    };

    template <> struct SpecificationTraits<FixedCentralDensity> {
        static constexpr std::string_view name  = "FixedCentralDensity";
        static constexpr SpecificationRole role = SpecificationRole::phase_condition;
        static constexpr SpecificationKey key{role, 0};
    };

    template <typename... Types> struct ModelTypeList final {
        static constexpr std::size_t size = sizeof...(Types);
    };

    template <typename Query, typename List> struct ModelTypeListContains;

    template <typename Query, typename... Types>
    struct ModelTypeListContains<Query, ModelTypeList<Types...>>
        : std::bool_constant<(std::same_as<Query, Types> || ...)> { };

    template <typename Query, typename List>
    inline constexpr bool modelTypeListContains = ModelTypeListContains<Query, List>::value;

    template <ModelSpecification Specification> struct ResidualFor final {
        using SpecificationType                  = Specification;

        static constexpr std::size_t scalarArity = 1;
    };

    template <ModelSpecification Specification> struct MultiplierFor final {
        using SpecificationType                  = Specification;

        static constexpr std::size_t scalarArity = 1;
    };

    template <ModelSpecification Specification> struct BorderFor final {
        using SpecificationType                  = Specification;

        static constexpr std::size_t scalarArity = 1;
    };

    template <ModelSpecification Specification> struct SpecificationContribution {
        using GeneratedValues                 = ModelTypeList<>;
        using GeneratedResiduals              = ModelTypeList<>;

        static constexpr bool isDefined       = false;
        static constexpr bool hasRootCompiler = false;
    };

    template <> struct SpecificationContribution<eos::Polytrope> {
        using GeneratedValues                 = ModelTypeList<>;
        using GeneratedResiduals              = ModelTypeList<>;

        static constexpr bool isDefined       = true;
        static constexpr bool hasRootCompiler = true;
    };

    template <> struct SpecificationContribution<surface::ConstantPressureSurface> {
        using GeneratedValues                 = ModelTypeList<>;
        using GeneratedResiduals              = ModelTypeList<>;

        static constexpr bool isDefined       = true;
        static constexpr bool hasRootCompiler = true;
    };

    template <> struct SpecificationContribution<FixedTotalMass> {
        using GeneratedValues                 = ModelTypeList<MultiplierFor<FixedTotalMass>>;
        using GeneratedResiduals              = ModelTypeList<ResidualFor<FixedTotalMass>>;

        static constexpr bool isDefined       = true;
        static constexpr bool hasRootCompiler = true;
    };

    template <> struct SpecificationContribution<FixedCentralDensity> {
        using GeneratedValues                 = ModelTypeList<BorderFor<FixedCentralDensity>>;
        using GeneratedResiduals              = ModelTypeList<ResidualFor<FixedCentralDensity>>;

        static constexpr bool isDefined       = true;
        static constexpr bool hasRootCompiler = true;
    };

    template <typename Candidate>
    concept ResolvedModelSpecification =
        ModelSpecification<Candidate> && SpecificationContribution<std::remove_cvref_t<Candidate>>::isDefined;

    namespace detail {
        template <typename... Specifications> struct SpecificationSetStorage final {
            static constexpr std::size_t size = sizeof...(Specifications);
        };

        template <typename... Lists> struct ConcatenateModelTypeLists;

        template <> struct ConcatenateModelTypeLists<> {
            using Type = ModelTypeList<>;
        };

        template <typename... Types> struct ConcatenateModelTypeLists<ModelTypeList<Types...>> {
            using Type = ModelTypeList<Types...>;
        };

        template <typename... First, typename... Second, typename... Remaining>
        struct ConcatenateModelTypeLists<ModelTypeList<First...>, ModelTypeList<Second...>, Remaining...> {
            using Type = typename ConcatenateModelTypeLists<ModelTypeList<First..., Second...>, Remaining...>::Type;
        };

        template <ModelSpecification Specification, typename Set> struct InsertSpecification;

        template <ModelSpecification Specification>
        struct InsertSpecification<Specification, SpecificationSetStorage<>> {
            using Type = SpecificationSetStorage<Specification>;
        };

        template <ModelSpecification Specification, ModelSpecification Head, ModelSpecification... Tail>
        struct InsertSpecification<Specification, SpecificationSetStorage<Head, Tail...>> {
        private:
            using InsertedTail = typename InsertSpecification<Specification, SpecificationSetStorage<Tail...>>::Type;

            template <typename First, typename Rest> struct PrependSpecification;

            template <typename First, ModelSpecification... Rest>
            struct PrependSpecification<First, SpecificationSetStorage<Rest...>> {
                using Type = SpecificationSetStorage<First, Rest...>;
            };

        public:
            using Type = std::conditional_t<
                (SpecificationTraits<Specification>::key < SpecificationTraits<Head>::key),
                SpecificationSetStorage<Specification, Head, Tail...>,
                typename PrependSpecification<Head, InsertedTail>::Type>;
        };

        template <typename Set, ModelSpecification... Specifications> struct CanonicalizeSpecifications;

        template <typename Set> struct CanonicalizeSpecifications<Set> {
            using Type = Set;
        };

        template <typename Set, ModelSpecification Head, ModelSpecification... Tail>
        struct CanonicalizeSpecifications<Set, Head, Tail...> {
            using Inserted = typename InsertSpecification<Head, Set>::Type;
            using Type     = typename CanonicalizeSpecifications<Inserted, Tail...>::Type;
        };

        template <ModelSpecification... Specifications>
        using CanonicalSpecificationSet =
            typename CanonicalizeSpecifications<SpecificationSetStorage<>, Specifications...>::Type;

        template <
            ModelSpecification Head,
            ModelSpecification... Tail>
        consteval bool specificationKeyIsUnique() {
            return ((SpecificationTraits<Head>::key != SpecificationTraits<Tail>::key) && ...);
        }

        template <ModelSpecification... Specifications> struct SpecificationKeysAreUnique;

        template <> struct SpecificationKeysAreUnique<> : std::true_type { };

        template <ModelSpecification Head, ModelSpecification... Tail>
        struct SpecificationKeysAreUnique<Head, Tail...>
            : std::bool_constant<
                  specificationKeyIsUnique<Head, Tail...>() && SpecificationKeysAreUnique<Tail...>::value> { };

        template <SpecificationRole Role, ModelSpecification... Specifications>
        inline constexpr std::size_t specificationRoleCount =
            (std::size_t{0} + ... + (SpecificationTraits<Specifications>::role == Role ? 1 : 0));

        template <typename List> struct ModelTypeListScalarArity;

        template <typename... Types>
        struct ModelTypeListScalarArity<ModelTypeList<Types...>>
            : std::integral_constant<std::size_t, (std::size_t{0} + ... + Types::scalarArity)> { };

        template <typename Query, typename... Types>
        inline constexpr bool isOneOf = (std::same_as<Query, Types> || ...);

        template <typename Query, typename... Types>
        inline constexpr std::size_t typeCount =
            (std::size_t{0} + ... +
             (std::same_as<Query, std::remove_cvref_t<Types>> ? std::size_t{1} : std::size_t{0}));

        template <typename CanonicalSet, typename... Arguments> struct ArgumentsMatchCanonicalSpecifications;

        template <ModelSpecification... CanonicalSpecifications, typename... Arguments>
        struct ArgumentsMatchCanonicalSpecifications<SpecificationSetStorage<CanonicalSpecifications...>, Arguments...>
            : std::bool_constant<
                  sizeof...(CanonicalSpecifications) == sizeof...(Arguments) &&
                  (isOneOf<std::remove_cvref_t<Arguments>, CanonicalSpecifications...> && ...) &&
                  ((typeCount<CanonicalSpecifications, Arguments...> == 1) && ...)> { };
    } // namespace detail

    template <ModelSpecification... Specifications>
    inline constexpr bool specificationKeysAreUnique = detail::SpecificationKeysAreUnique<Specifications...>::value;

    template <typename... Specifications>
    concept ValidModelSpecificationPack =
        (ResolvedModelSpecification<Specifications> && ...) && specificationKeysAreUnique<Specifications...> &&
        detail::specificationRoleCount<SpecificationRole::constitutive_law, Specifications...> == 1;

    template <ModelSpecification... Specifications>
        requires specificationKeysAreUnique<Specifications...>
    using SpecificationSet = detail::CanonicalSpecificationSet<Specifications...>;

    template <typename SpecificationSet> struct SpecificationOperatorSignature;

    template <ModelSpecification... Specifications>
    struct SpecificationOperatorSignature<detail::SpecificationSetStorage<Specifications...>> final {
        using GeneratedValues = typename detail::ConcatenateModelTypeLists<
            typename SpecificationContribution<Specifications>::GeneratedValues...>::Type;

        using GeneratedResiduals = typename detail::ConcatenateModelTypeLists<
            typename SpecificationContribution<Specifications>::GeneratedResiduals...>::Type;

        static constexpr std::size_t generatedValueArity = detail::ModelTypeListScalarArity<GeneratedValues>::value;

        static constexpr std::size_t generatedResidualArity =
            detail::ModelTypeListScalarArity<GeneratedResiduals>::value;

        static constexpr bool symbolicallySquare = generatedValueArity == generatedResidualArity;
    };

    template <ResolvedModelSpecification Specification>
    [[nodiscard]] consteval SpecificationDescriptor specificationDescriptor() {
        using Contribution = SpecificationContribution<Specification>;

        return {
            .name                   = SpecificationTraits<Specification>::name,
            .role                   = SpecificationTraits<Specification>::role,
            .key                    = SpecificationTraits<Specification>::key,
            .generatedValueArity    = detail::ModelTypeListScalarArity<typename Contribution::GeneratedValues>::value,
            .generatedResidualArity = detail::ModelTypeListScalarArity<typename Contribution::GeneratedResiduals>::value
        };
    }

    namespace detail {
        template <typename Specifications> class SpecifiedModel;

        template <ModelSpecification... Specifications>
        class SpecifiedModel<SpecificationSetStorage<Specifications...>> final {
        public:
            using SpecificationTypes                 = SpecificationSetStorage<Specifications...>;
            using OperatorSignature                  = SpecificationOperatorSignature<SpecificationTypes>;

            static constexpr bool symbolicallySquare = OperatorSignature::symbolicallySquare;
            static constexpr bool hasCompleteRootCompiler =
                (SpecificationContribution<Specifications>::hasRootCompiler && ...);
            static constexpr EquilibriumSystemCompilation compilationClass =
                symbolicallySquare && hasCompleteRootCompiler
                    ? EquilibriumSystemCompilation::complete_equilibrium_system
                    : EquilibriumSystemCompilation::equation_contributions_only;

            template <typename... Arguments>
                requires ArgumentsMatchCanonicalSpecifications<
                    SpecificationTypes,
                    Arguments...>::value
            explicit SpecifiedModel(Arguments &&...arguments)
                : m_specifications(
                      std::get<Specifications>(
                          std::tuple<std::remove_cvref_t<Arguments>...>{std::forward<Arguments>(arguments)...}
                      )...
                  ) {
            }

            template <ModelSpecification Specification>
                requires isOneOf<
                    Specification,
                    Specifications...>
            [[nodiscard]] const Specification &specification() const noexcept {
                return std::get<Specification>(m_specifications);
            }

            template <ModelSpecification Specification>
            static constexpr bool containsSpecification = isOneOf<Specification, Specifications...>;

            [[nodiscard]] static constexpr std::span<const RuntimeSpecificationDescriptor>
            runtimeSpecificationDescriptors() noexcept {
                return runtimeDescriptors;
            }

        private:
            inline static constexpr std::array<RuntimeSpecificationDescriptor, sizeof...(Specifications)>
                runtimeDescriptors = [] {
                    std::array<RuntimeSpecificationDescriptor, sizeof...(Specifications)> descriptors{};
                    std::size_t index = 0;
                    ((descriptors[index] =
                          {.specification   = specificationDescriptor<Specifications>(),
                           .canonicalIndex  = index,
                           .hasRootCompiler = SpecificationContribution<Specifications>::hasRootCompiler},
                      ++index),
                     ...);
                    return descriptors;
                }();

            std::tuple<Specifications...> m_specifications;
        };
    } // namespace detail

    template <ModelSpecification... Specifications>
        requires ValidModelSpecificationPack<Specifications...> &&
                     SpecificationOperatorSignature<SpecificationSet<Specifications...>>::symbolicallySquare
    using Model = detail::SpecifiedModel<SpecificationSet<Specifications...>>;

    template <typename Candidate>
    concept SpecifiedModelType = requires {
        typename std::remove_cvref_t<Candidate>::SpecificationTypes;
        typename std::remove_cvref_t<Candidate>::OperatorSignature;
        requires std::remove_cvref_t<Candidate>::symbolicallySquare;
        { std::remove_cvref_t<Candidate>::compilationClass } -> std::convertible_to<ModelCompilationClass>;
        {
            std::remove_cvref_t<Candidate>::runtimeSpecificationDescriptors()
        } -> std::same_as<std::span<const RuntimeSpecificationDescriptor>>;
    };

    static_assert(ModelSpecification<eos::Polytrope>);
    static_assert(ModelSpecification<surface::ConstantPressureSurface>);
    static_assert(ModelSpecification<FixedTotalMass>);
    static_assert(ModelSpecification<FixedCentralDensity>);
    static_assert(ResolvedModelSpecification<eos::Polytrope>);
    static_assert(ResolvedModelSpecification<surface::ConstantPressureSurface>);
    static_assert(ResolvedModelSpecification<FixedTotalMass>);
    static_assert(ResolvedModelSpecification<FixedCentralDensity>);
} // namespace mean_field::models

export namespace mean_field::integral {
    using FixedTotalMass = models::FixedTotalMass;
}

export namespace mean_field::constraint {
    using FixedCentralDensity = models::FixedCentralDensity;
}