module;

#include <concepts>
#include <type_traits>

export module mean_field:material.thermodynamic_equations;

export import :eos.polytrope;
export import :surface.compiler;
export import :utils.blocks;

export namespace mean_field::material {
    /**
     * A thermodynamic field identifies the physical quantity represented by
     * its discrete degree of freedom. The quantity belongs to the field, not
     * to an EOS-specific aggregate description.
     */
    template <typename Candidate>
    concept ThermodynamicField = surface::SurfaceFieldType<Candidate> && requires {
        typename std::remove_cvref_t<Candidate>::PhysicalQuantity;
        requires eos::ThermodynamicQuantityType<typename std::remove_cvref_t<Candidate>::PhysicalQuantity>;
    };

    /**
     * Declares the algebraic blocks owned by one thermodynamic governing
     * equation. Its physical quantity is inferred from Field.
     */
    template <ThermodynamicField Field, typename CorrectionBlock, typename ResidualBlock>
        requires std::derived_from<CorrectionBlock, utils::blocks::value_block_base> &&
                 std::derived_from<ResidualBlock, utils::blocks::residual_block_base>
    struct ThermodynamicEquation final {
        using FieldType        = Field;
        using PhysicalQuantity = typename Field::PhysicalQuantity;
        using Correction       = CorrectionBlock;
        using Residual         = ResidualBlock;
    };

    /**
     * Registry of thermodynamic governing equations available to a problem.
     * The problem form, rather than the registry, selects the active subset.
     */
    template <typename... Equations> struct ThermodynamicEquationCatalog final {
        static constexpr int size = sizeof...(Equations);
    };

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

        template <typename Field, typename CorrectionBlock, typename ResidualBlock>
        struct IsThermodynamicEquation<ThermodynamicEquation<Field, CorrectionBlock, ResidualBlock>> : std::true_type {
        };

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

        template <typename... Equations>
        struct IsThermodynamicEquationCatalog<ThermodynamicEquationCatalog<Equations...>>
            : std::bool_constant<(sizeof...(Equations) > 0) && (IsThermodynamicEquation<Equations>::value && ...)> { };

        template <typename Catalog> struct CatalogEntriesAreUnique : std::false_type { };

        template <typename... Equations>
        struct CatalogEntriesAreUnique<ThermodynamicEquationCatalog<Equations...>>
            : std::bool_constant<
                  utils::blocks::types_are_unique_v<utils::blocks::type_list<typename Equations::FieldType...>> &&
                  utils::blocks::types_are_unique_v<
                      utils::blocks::type_list<typename Equations::PhysicalQuantity...>> &&
                  utils::blocks::types_are_unique_v<utils::blocks::type_list<typename Equations::Correction...>> &&
                  utils::blocks::types_are_unique_v<utils::blocks::type_list<typename Equations::Residual...>>> { };

        template <typename Catalog, typename Field> struct EquationForField;

        template <typename Field, typename First, typename... Remaining>
        struct EquationForField<ThermodynamicEquationCatalog<First, Remaining...>, Field>
            : std::conditional_t<
                  std::same_as<Field, typename First::FieldType>,
                  std::type_identity<First>,
                  EquationForField<ThermodynamicEquationCatalog<Remaining...>, Field>> { };

        template <typename Field> struct EquationForField<ThermodynamicEquationCatalog<>, Field> {
            using type = void;
        };

        template <typename Catalog, typename Field> struct EquationFieldCount;

        template <typename Field, typename... Equations>
        struct EquationFieldCount<ThermodynamicEquationCatalog<Equations...>, Field>
            : std::integral_constant<
                  int,
                  (int{0} + ... + (std::same_as<Field, typename Equations::FieldType> ? 1 : 0))> { };

        template <typename Catalog, typename Correction> struct EquationForCorrection;

        template <typename Correction, typename First, typename... Remaining>
        struct EquationForCorrection<ThermodynamicEquationCatalog<First, Remaining...>, Correction>
            : std::conditional_t<
                  std::same_as<Correction, typename First::Correction>,
                  std::type_identity<First>,
                  EquationForCorrection<ThermodynamicEquationCatalog<Remaining...>, Correction>> { };

        template <typename Correction> struct EquationForCorrection<ThermodynamicEquationCatalog<>, Correction> {
            using type = void;
        };

        template <typename Head, typename Catalog> struct PrependEquation;

        template <typename Head, typename... Equations>
        struct PrependEquation<Head, ThermodynamicEquationCatalog<Equations...>> {
            using Type = ThermodynamicEquationCatalog<Head, Equations...>;
        };

        template <typename ValueBlocks, typename AvailableEquations> struct SelectActiveEquations;

        template <typename AvailableEquations>
        struct SelectActiveEquations<utils::blocks::type_list<>, AvailableEquations> {
            using Type = ThermodynamicEquationCatalog<>;
        };

        template <typename FirstValue, typename... RemainingValues, typename AvailableEquations>
        struct SelectActiveEquations<utils::blocks::type_list<FirstValue, RemainingValues...>, AvailableEquations> {
        private:
            using Tail =
                typename SelectActiveEquations<utils::blocks::type_list<RemainingValues...>, AvailableEquations>::Type;
            using Match = typename EquationForCorrection<AvailableEquations, FirstValue>::type;

        public:
            using Type =
                std::conditional_t<std::same_as<Match, void>, Tail, typename PrependEquation<Match, Tail>::Type>;
        };

        template <typename AvailableEquations, typename Form> struct CatalogMatchesForm : std::false_type { };

        template <typename... Equations, typename... Values, typename... Residuals>
        struct CatalogMatchesForm<
            ThermodynamicEquationCatalog<Equations...>,
            utils::blocks::block_form<utils::blocks::type_list<Values...>, utils::blocks::type_list<Residuals...>>>
            : std::bool_constant<
                  ((utils::blocks::
                        contains_type_v<typename Equations::Correction, utils::blocks::type_list<Values...>> ==
                    utils::blocks::
                        contains_type_v<typename Equations::Residual, utils::blocks::type_list<Residuals...>>) &&
                   ...)> { };

        template <typename Equations> struct SurfaceBindingsForEquations;

        template <typename... Equations>
        struct SurfaceBindingsForEquations<ThermodynamicEquationCatalog<Equations...>> {
            using Type = surface::SurfaceStateBindings<
                surface::SurfaceStateBinding<typename Equations::PhysicalQuantity, typename Equations::FieldType>...>;
        };

        template <typename SurfaceFields, typename Equations>
        struct SurfaceFieldsBelongToEquations : std::false_type { };

        template <typename... Fields, typename Equations>
        struct SurfaceFieldsBelongToEquations<field::TypeList<Fields...>, Equations>
            : std::bool_constant<((EquationFieldCount<Equations, Fields>::value == 1) && ...)> { };

        template <typename EquationOfState, typename Form, typename AvailableEquations, typename = void>
        struct ThermodynamicCompilationIsAvailable : std::false_type { };

        template <typename EquationOfState, typename Form, typename AvailableEquations>
        struct ThermodynamicCompilationIsAvailable<
            EquationOfState,
            Form,
            AvailableEquations,
            std::enable_if_t<
                eos::EquationOfStateModel<EquationOfState> && utils::blocks::block_form_is_valid_v<Form> &&
                IsThermodynamicEquationCatalog<AvailableEquations>::value &&
                CatalogEntriesAreUnique<AvailableEquations>::value &&
                CatalogMatchesForm<AvailableEquations, Form>::value>> {
        private:
            using ActiveEquations =
                typename SelectActiveEquations<typename Form::value_blocks, AvailableEquations>::Type;
            using StateBindings = typename SurfaceBindingsForEquations<ActiveEquations>::Type;

        public:
            static constexpr bool value = (ActiveEquations::size > 0) &&
                                          surface::PressureSurfaceFormulationCompilable<StateBindings, EquationOfState>;
        };

        template <typename Candidate, typename = void> struct IsCompiledThermodynamicEquations : std::false_type { };

        template <typename Candidate>
        struct IsCompiledThermodynamicEquations<
            Candidate,
            std::void_t<
                typename Candidate::EquationOfStateType,
                typename Candidate::Equations,
                typename Candidate::StateBindings,
                typename Candidate::PressureSurfaceFormulation>>
            : std::bool_constant<
                  eos::EquationOfStateModel<typename Candidate::EquationOfStateType> &&
                  IsThermodynamicEquationCatalog<typename Candidate::Equations>::value &&
                  CatalogEntriesAreUnique<typename Candidate::Equations>::value &&
                  surface::ValidSurfaceStateBindings<typename Candidate::StateBindings> &&
                  std::same_as<
                      typename Candidate::StateBindings,
                      typename SurfaceBindingsForEquations<typename Candidate::Equations>::Type> &&
                  surface::SurfaceConstraintFormulationType<typename Candidate::PressureSurfaceFormulation> &&
                  std::same_as<
                      typename Candidate::PressureSurfaceFormulation::StateBindings,
                      typename Candidate::StateBindings>> { };
    } // namespace detail

    template <typename Candidate>
    concept ThermodynamicEquationType = detail::IsThermodynamicEquation<std::remove_cvref_t<Candidate>>::value;

    template <typename Candidate>
    concept ValidThermodynamicEquationCatalog =
        detail::IsThermodynamicEquationCatalog<std::remove_cvref_t<Candidate>>::value &&
        detail::CatalogEntriesAreUnique<std::remove_cvref_t<Candidate>>::value;

    template <typename Candidate>
    concept CompiledThermodynamicEquations =
        detail::IsCompiledThermodynamicEquations<std::remove_cvref_t<Candidate>>::value;

    template <ValidThermodynamicEquationCatalog Equations, surface::SurfaceFieldType Field>
        requires(detail::EquationFieldCount<Equations, Field>::value == 1)
    using ThermodynamicEquationForFieldT = typename detail::EquationForField<Equations, Field>::type;

    template <typename Fields, ValidThermodynamicEquationCatalog Equations>
    inline constexpr bool fieldsBelongToThermodynamicEquations =
        detail::SurfaceFieldsBelongToEquations<Fields, Equations>::value;

    template <typename EquationOfState, typename Form, typename AvailableEquations>
    concept ThermodynamicEquationsCompilable = detail::ThermodynamicCompilationIsAvailable<
        std::remove_cvref_t<EquationOfState>,
        std::remove_cvref_t<Form>,
        std::remove_cvref_t<AvailableEquations>>::value;

    template <eos::EquationOfStateModel EquationOfState, ValidThermodynamicEquationCatalog ActiveEquations>
        requires surface::PressureSurfaceFormulationCompilable<
            typename detail::SurfaceBindingsForEquations<ActiveEquations>::Type,
            EquationOfState>
    struct ThermodynamicEquationSet final {
        using EquationOfStateType = EquationOfState;
        using Equations           = ActiveEquations;
        using StateBindings       = typename detail::SurfaceBindingsForEquations<Equations>::Type;
        using PressureSurfaceFormulation =
            surface::CompiledPressureSurfaceFormulationT<StateBindings, EquationOfStateType>;
    };

    template <
        eos::EquationOfStateModel EquationOfState,
        typename Form,
        ValidThermodynamicEquationCatalog AvailableEquations>
        requires ThermodynamicEquationsCompilable<EquationOfState, Form, AvailableEquations>
    struct CompileThermodynamicEquations final {
        using Equations = typename detail::SelectActiveEquations<typename Form::value_blocks, AvailableEquations>::Type;
        using Type      = ThermodynamicEquationSet<EquationOfState, Equations>;
    };

    template <typename EquationOfState, typename Form, typename AvailableEquations>
        requires ThermodynamicEquationsCompilable<EquationOfState, Form, AvailableEquations>
    using CompiledThermodynamicEquationsT = typename CompileThermodynamicEquations<
        std::remove_cvref_t<EquationOfState>,
        std::remove_cvref_t<Form>,
        std::remove_cvref_t<AvailableEquations>>::Type;

    using StellarEquilibriumThermodynamicEquations = ThermodynamicEquationCatalog<
        ThermodynamicEquation<
            field::Density,
            utils::blocks::density::mass::value,
            utils::blocks::density::mass::residual>,
        ThermodynamicEquation<
            field::Enthalpy,
            utils::blocks::enthalpy::specific::value,
            utils::blocks::enthalpy::specific::residual>>;

    static_assert(ValidThermodynamicEquationCatalog<StellarEquilibriumThermodynamicEquations>);
} // namespace mean_field::material