module;

#include <cstddef>
#include <tuple>
#include <type_traits>

export module mean_field:surface.compiler;

export import :surface.compiled;

export namespace mean_field::surface {
    namespace detail {
        template <typename RelationType, typename Formulation, typename EquationOfState>
        struct PressureSurfaceRelationMatches : std::false_type { };

        template <typename OutputQuantity, typename... InputQuantities, typename Formulation, typename EquationOfState>
        struct PressureSurfaceRelationMatches<
            eos::Relation<OutputQuantity, InputQuantities...>,
            Formulation,
            EquationOfState>
            : std::bool_constant<
                  std::same_as<OutputQuantity, typename Formulation::CarrierQuantity> &&
                  (std::same_as<eos::quantity::Pressure, InputQuantities> || ...) &&
                  ((std::same_as<eos::quantity::Pressure, InputQuantities> ||
                    (surfaceBindingCount<typename Formulation::StateBindings, InputQuantities> == 1 &&
                     eos::SupportsPartialDerivative<
                         EquationOfState,
                         eos::Relation<OutputQuantity, InputQuantities...>,
                         InputQuantities>)) &&
                   ...)> { };

        template <typename RelationType, typename Bindings, typename EquationOfState>
        struct PressureSurfaceRelationMatchesBindings : std::false_type { };

        template <typename OutputQuantity, typename... InputQuantities, typename Bindings, typename EquationOfState>
        struct PressureSurfaceRelationMatchesBindings<
            eos::Relation<OutputQuantity, InputQuantities...>,
            Bindings,
            EquationOfState>
            : std::bool_constant<
                  (surfaceBindingCount<Bindings, OutputQuantity> == 1) &&
                  (std::same_as<eos::quantity::Pressure, InputQuantities> || ...) &&
                  ((std::same_as<eos::quantity::Pressure, InputQuantities> ||
                    (surfaceBindingCount<Bindings, InputQuantities> == 1 &&
                     eos::SupportsPartialDerivative<
                         EquationOfState,
                         eos::Relation<OutputQuantity, InputQuantities...>,
                         InputQuantities>)) &&
                   ...)> { };

        template <typename Catalog, typename Formulation, typename EquationOfState>
        struct MatchingPressureSurfaceRelations;

        template <typename... Relations, typename Formulation, typename EquationOfState>
        struct MatchingPressureSurfaceRelations<eos::RelationCatalog<Relations...>, Formulation, EquationOfState> {
            using Tuple                        = decltype(std::tuple_cat(
                std::conditional_t<
                    PressureSurfaceRelationMatches<Relations, Formulation, EquationOfState>::value,
                    std::tuple<Relations>,
                    std::tuple<>>{}...
            ));

            static constexpr std::size_t count = std::tuple_size_v<Tuple>;
        };

        template <typename Catalog, typename Bindings, typename EquationOfState>
        struct MatchingPressureSurfaceRelationsForBindings;

        template <typename... Relations, typename Bindings, typename EquationOfState>
        struct MatchingPressureSurfaceRelationsForBindings<
            eos::RelationCatalog<Relations...>,
            Bindings,
            EquationOfState> {
            using Tuple                        = decltype(std::tuple_cat(
                std::conditional_t<
                    PressureSurfaceRelationMatchesBindings<Relations, Bindings, EquationOfState>::value,
                    std::tuple<Relations>,
                    std::tuple<>>{}...
            ));

            static constexpr std::size_t count = std::tuple_size_v<Tuple>;
        };

        template <std::size_t Count, typename Tuple> struct UniquePressureSurfaceRelation {
            using Type = void;
        };

        template <typename Tuple> struct UniquePressureSurfaceRelation<1, Tuple> {
            using Type = std::tuple_element_t<0, Tuple>;
        };

        template <std::size_t Count, typename Tuple, typename Bindings> struct UniquePressureSurfaceFormulation {
            using Type = void;
        };

        template <typename Tuple, typename Bindings> struct UniquePressureSurfaceFormulation<1, Tuple, Bindings> {
            using Relation        = std::tuple_element_t<0, Tuple>;
            using CarrierQuantity = eos::RelationOutputT<Relation>;
            using CarrierField    = SurfaceFieldForQuantityT<Bindings, CarrierQuantity>;
            using Type            = SurfaceConstraintFormulation<CarrierQuantity, CarrierField, Bindings>;
        };

        template <typename Dependencies, typename Field> struct AppendSurfaceDependency;

        template <typename RowField, typename... StateFields, typename Field>
        struct AppendSurfaceDependency<SurfaceConstraintDependencies<RowField, StateFields...>, Field> {
            using Type = SurfaceConstraintDependencies<RowField, StateFields..., Field>;
        };

        template <typename Dependencies, typename InputQuantity, typename Bindings>
        struct AppendPressureSurfaceInputDependency {
            using Type =
                typename AppendSurfaceDependency<Dependencies, SurfaceFieldForQuantityT<Bindings, InputQuantity>>::Type;
        };

        template <typename Dependencies, typename Bindings>
        struct AppendPressureSurfaceInputDependency<Dependencies, eos::quantity::Pressure, Bindings> {
            using Type = Dependencies;
        };

        template <typename Dependencies, typename Bindings, typename... InputQuantities>
        struct AppendPressureSurfaceInputDependencies;

        template <typename Dependencies, typename Bindings>
        struct AppendPressureSurfaceInputDependencies<Dependencies, Bindings> {
            using Type = Dependencies;
        };

        template <typename Dependencies, typename Bindings, typename FirstInput, typename... RemainingInputs>
        struct AppendPressureSurfaceInputDependencies<Dependencies, Bindings, FirstInput, RemainingInputs...> {
            using WithFirst = typename AppendPressureSurfaceInputDependency<Dependencies, FirstInput, Bindings>::Type;
            using Type = typename AppendPressureSurfaceInputDependencies<WithFirst, Bindings, RemainingInputs...>::Type;
        };

        template <typename RelationType, typename Formulation> struct PressureSurfaceDependenciesForRelation;

        template <typename OutputQuantity, typename... InputQuantities, typename Formulation>
        struct PressureSurfaceDependenciesForRelation<eos::Relation<OutputQuantity, InputQuantities...>, Formulation> {
            using InitialDependencies =
                SurfaceConstraintDependencies<typename Formulation::CarrierField, typename Formulation::CarrierField>;
            using Type = typename AppendPressureSurfaceInputDependencies<
                InitialDependencies,
                typename Formulation::StateBindings,
                InputQuantities...>::Type;
        };

        template <SurfaceConstraintFormulationType Formulation, eos::EquationOfStateModel EquationOfState>
        struct PressureSurfaceCompilation {
            using Matches =
                MatchingPressureSurfaceRelations<typename EquationOfState::Relations, Formulation, EquationOfState>;
            using Relation = typename UniquePressureSurfaceRelation<Matches::count, typename Matches::Tuple>::Type;
        };

        template <ValidSurfaceStateBindings Bindings, eos::EquationOfStateModel EquationOfState>
        struct PressureSurfaceFormulationCompilation {
            using Matches = MatchingPressureSurfaceRelationsForBindings<
                typename EquationOfState::Relations,
                Bindings,
                EquationOfState>;
            using Formulation =
                typename UniquePressureSurfaceFormulation<Matches::count, typename Matches::Tuple, Bindings>::Type;
        };

        template <SurfaceConstraintFormulationType Formulation, eos::EquationOfStateModel EquationOfState>
            requires(PressureSurfaceCompilation<Formulation, EquationOfState>::Matches::count == 1)
        struct CompiledPressureSurfaceConstraintType {
            using Compilation  = PressureSurfaceCompilation<Formulation, EquationOfState>;
            using Relation     = typename Compilation::Relation;
            using Dependencies = typename PressureSurfaceDependenciesForRelation<Relation, Formulation>::Type;
            using Type = CompiledPressureSurfaceConstraint<EquationOfState, Formulation, Relation, Dependencies>;
        };
    } // namespace detail

    template <typename Formulation, typename EquationOfState>
    concept PressureSurfaceCompilable =
        SurfaceConstraintFormulationType<Formulation> && eos::EquationOfStateModel<EquationOfState> &&
        (detail::PressureSurfaceCompilation<std::remove_cvref_t<Formulation>, std::remove_cvref_t<EquationOfState>>::
             Matches::count == 1);

    template <typename Bindings, typename EquationOfState>
    concept PressureSurfaceFormulationCompilable =
        ValidSurfaceStateBindings<Bindings> && eos::EquationOfStateModel<EquationOfState> &&
        (detail::PressureSurfaceFormulationCompilation<
             std::remove_cvref_t<Bindings>,
             std::remove_cvref_t<EquationOfState>>::Matches::count == 1);

    template <ValidSurfaceStateBindings Bindings, eos::EquationOfStateModel EquationOfState>
        requires PressureSurfaceFormulationCompilable<Bindings, EquationOfState>
    using CompiledPressureSurfaceFormulationT = typename detail::PressureSurfaceFormulationCompilation<
        std::remove_cvref_t<Bindings>,
        std::remove_cvref_t<EquationOfState>>::Formulation;

    template <SurfaceConstraintFormulationType Formulation, eos::EquationOfStateModel EquationOfState>
        requires PressureSurfaceCompilable<Formulation, EquationOfState>
    using CompiledPressureSurfaceConstraintT = typename detail::CompiledPressureSurfaceConstraintType<
        std::remove_cvref_t<Formulation>,
        std::remove_cvref_t<EquationOfState>>::Type;

    template <
        SurfaceConstraintFormulationType Formulation,
        eos::EquationOfStateModel EquationOfState>
        requires PressureSurfaceCompilable<
            Formulation,
            EquationOfState>
    [[nodiscard]] CompiledPressureSurfaceConstraintT<
        Formulation,
        EquationOfState>
    compilePressureSurfaceConstraint(
        const ConstantPressureSurface condition,
        const EquationOfState &equationOfState
    ) noexcept {
        return CompiledPressureSurfaceConstraintT<Formulation, EquationOfState>{condition, equationOfState};
    }
} // namespace mean_field::surface