module;

#include <array>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <expected>
#include <memory>
#include <span>
#include <string>
#include <string_view>
#include <tuple>
#include <type_traits>
#include <utility>

export module mean_field:eos.runtime;
export import :eos.evaluation;

export namespace mean_field::eos {
    class ThermodynamicQuantityId final {
    public:
        explicit constexpr ThermodynamicQuantityId(const std::string_view name) noexcept : m_name(name) {
        }

        [[nodiscard]] constexpr std::string_view name() const noexcept {
            return m_name;
        }

        [[nodiscard]] friend constexpr bool operator==(
            const ThermodynamicQuantityId &,
            const ThermodynamicQuantityId &
        ) noexcept = default;

    private:
        std::string_view m_name;
    };

    template <typename Quantity>
    concept RuntimeIdentifiedThermodynamicQuantity = ThermodynamicQuantityType<Quantity> && requires {
        { Quantity::identifier } -> std::convertible_to<std::string_view>;
    } && (std::string_view{Quantity::identifier}.size() > 0);

    template <RuntimeIdentifiedThermodynamicQuantity Quantity>
    inline constexpr ThermodynamicQuantityId thermodynamicQuantityId{std::string_view{Quantity::identifier}};

    struct RuntimeQuantityValue final {
        ThermodynamicQuantityId quantity;
        double value;
    };

    struct RuntimeRelationDescriptor final {
        ThermodynamicQuantityId outputQuantity;
        std::span<const ThermodynamicQuantityId> inputQuantities;
        std::uint64_t partialDerivativeMask;

        [[nodiscard]] constexpr bool hasPartialDerivative(const std::size_t inputIndex) const noexcept {
            return inputIndex < inputQuantities.size() &&
                   (partialDerivativeMask & (std::uint64_t{1} << inputIndex)) != 0;
        }
    };

    namespace detail {
        template <typename RelationType> struct HasRuntimeQuantityIdentifiers : std::false_type { };

        template <typename Output, typename... Inputs>
        struct HasRuntimeQuantityIdentifiers<Relation<Output, Inputs...>>
            : std::bool_constant<
                  RuntimeIdentifiedThermodynamicQuantity<Output> &&
                  (RuntimeIdentifiedThermodynamicQuantity<Inputs> && ...)> { };

        template <typename RelationType> struct RuntimeRelationQuantities;

        template <typename Output, typename... Inputs> struct RuntimeRelationQuantities<Relation<Output, Inputs...>> {
            using Type = std::tuple<Output, Inputs...>;
        };

        template <typename... Relations>
        using RuntimeCatalogQuantityTuple =
            decltype(std::tuple_cat(std::declval<typename RuntimeRelationQuantities<Relations>::Type>()...));

        template <
            typename FirstQuantity,
            typename SecondQuantity>
        [[nodiscard]] consteval bool runtimeQuantityIdentifiersAreCompatible() {
            if constexpr (std::same_as<FirstQuantity, SecondQuantity>) {
                return true;
            } else {
                return thermodynamicQuantityId<FirstQuantity> != thermodynamicQuantityId<SecondQuantity>;
            }
        }

        template <
            typename QuantityTuple,
            std::size_t First,
            std::size_t... Offsets>
        [[nodiscard]] consteval bool runtimeQuantityIdentifierIsUnambiguous(std::index_sequence<Offsets...>) {
            return (
                runtimeQuantityIdentifiersAreCompatible<
                    std::tuple_element_t<First, QuantityTuple>,
                    std::tuple_element_t<First + 1 + Offsets, QuantityTuple>>() &&
                ...
            );
        }

        template <
            typename QuantityTuple,
            std::size_t... Indices>
        [[nodiscard]] consteval bool runtimeQuantityIdentifiersAreUnambiguous(std::index_sequence<Indices...>) {
            return (
                runtimeQuantityIdentifierIsUnambiguous<QuantityTuple, Indices>(
                    std::make_index_sequence<std::tuple_size_v<QuantityTuple> - Indices - 1>{}
                ) &&
                ...
            );
        }

        template <bool QuantitiesAreIdentified, typename... Relations>
        struct RuntimeRelationsAreSupported : std::false_type { };

        template <typename... Relations>
        struct RuntimeRelationsAreSupported<true, Relations...>
            : std::bool_constant<runtimeQuantityIdentifiersAreUnambiguous<RuntimeCatalogQuantityTuple<Relations...>>(
                  std::make_index_sequence<std::tuple_size_v<RuntimeCatalogQuantityTuple<Relations...>>>{}
              )> { };

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

        template <typename... Relations>
        struct RuntimeCatalogIsSupported<RelationCatalog<Relations...>>
            : RuntimeRelationsAreSupported<(HasRuntimeQuantityIdentifiers<Relations>::value && ...), Relations...> { };
    } // namespace detail

    template <typename Candidate>
    concept RuntimeEquationOfStateModel =
        EquationOfStateModel<Candidate> &&
        detail::RuntimeCatalogIsSupported<typename std::remove_cvref_t<Candidate>::Relations>::value;

    namespace detail {
        template <typename EquationOfState, typename RelationType> struct RuntimeRelationStorage;

        template <typename EquationOfState, typename Output, typename... Inputs>
        struct RuntimeRelationStorage<EquationOfState, Relation<Output, Inputs...>> {
            using RelationType = Relation<Output, Inputs...>;

            static_assert(
                sizeof...(Inputs) <= 64,
                "Runtime EOS relation descriptors support at most 64 inputs."
            );

            inline static constexpr std::array<ThermodynamicQuantityId, sizeof...(Inputs)> inputQuantityIds{
                thermodynamicQuantityId<Inputs>...
            };

            template <std::size_t... Indices>
            [[nodiscard]] static consteval std::uint64_t makePartialDerivativeMask(std::index_sequence<Indices...>) {
                using InputTuple = std::tuple<Inputs...>;

                return (
                    std::uint64_t{0} | ... |
                    (SupportsPartialDerivative<EquationOfState, RelationType, std::tuple_element_t<Indices, InputTuple>>
                         ? (std::uint64_t{1} << Indices)
                         : std::uint64_t{0})
                );
            }

            inline static constexpr std::uint64_t partialDerivativeMask =
                makePartialDerivativeMask(std::index_sequence_for<Inputs...>{});

            inline static constexpr RuntimeRelationDescriptor descriptor{
                thermodynamicQuantityId<Output>, std::span<const ThermodynamicQuantityId>{inputQuantityIds},
                partialDerivativeMask
            };
        };

        template <typename EquationOfState, typename Catalog> struct RuntimeCatalogStorage;

        template <typename EquationOfState, typename... Relations>
        struct RuntimeCatalogStorage<EquationOfState, RelationCatalog<Relations...>> {
            inline static constexpr std::array descriptors{
                RuntimeRelationStorage<EquationOfState, Relations>::descriptor...
            };
        };

        [[nodiscard]] inline std::expected<
            double,
            EvaluationError>
        runtimeEvaluationFailure(
            const EvaluationErrorCode code,
            std::string message
        ) {
            return std::unexpected<EvaluationError>{EvaluationError{code, std::move(message)}};
        }

        template <
            typename EquationOfState,
            typename Output,
            typename... Inputs>
        [[nodiscard]] std::expected<
            double,
            EvaluationError>
        evaluateRuntimeRelation(
            const EquationOfState &equationOfState,
            Relation<
                Output,
                Inputs...>,
            const std::span<const RuntimeQuantityValue> inputValues
        ) {
            const auto invoke = [&]<std::size_t... Indices>(std::index_sequence<Indices...>) {
                return eos::evaluate<Output>(equationOfState, QuantityValue<Inputs>{inputValues[Indices].value}...)
                    .value();
            };

            try {
                return invoke(std::index_sequence_for<Inputs...>{});
            } catch (const EvaluationError &error) {
                return std::unexpected<EvaluationError>{error};
            }
        }

        template <
            typename InputQuantity,
            typename EquationOfState,
            typename Output,
            typename... Inputs>
        [[nodiscard]] bool tryRuntimePartialDerivative(
            const EquationOfState &equationOfState,
            Relation<
                Output,
                Inputs...> relation,
            const ThermodynamicQuantityId withRespectTo,
            const std::span<const RuntimeQuantityValue> inputValues,
            std::expected<
                double,
                EvaluationError> &result
        ) {
            if (withRespectTo != thermodynamicQuantityId<InputQuantity>) {
                return false;
            }

            if constexpr (SupportsPartialDerivative<EquationOfState, Relation<Output, Inputs...>, InputQuantity>) {
                const auto invoke = [&]<std::size_t... Indices>(std::index_sequence<Indices...>) {
                    return eos::partialDerivative<Output, InputQuantity>(
                               equationOfState, QuantityValue<Inputs>{inputValues[Indices].value}...
                    )
                        .value();
                };

                try {
                    result = invoke(std::index_sequence_for<Inputs...>{});
                } catch (const EvaluationError &error) {
                    result = std::unexpected<EvaluationError>{error};
                }
            } else {
                result = runtimeEvaluationFailure(
                    EvaluationErrorCode::unsupported_derivative,
                    "The requested EOS partial derivative is not available."
                );
            }

            return true;
        }

        template <
            typename EquationOfState,
            typename Output,
            typename... Inputs>
        [[nodiscard]] std::expected<
            double,
            EvaluationError>
        evaluateRuntimePartialDerivative(
            const EquationOfState &equationOfState,
            Relation<
                Output,
                Inputs...> relation,
            const ThermodynamicQuantityId withRespectTo,
            const std::span<const RuntimeQuantityValue> inputValues
        ) {
            std::expected<double, EvaluationError> result = runtimeEvaluationFailure(
                EvaluationErrorCode::unsupported_derivative,
                "The requested quantity is not an input to the EOS relation."
            );

            const bool matched =
                (tryRuntimePartialDerivative<Inputs>(equationOfState, relation, withRespectTo, inputValues, result) ||
                 ...);

            static_cast<void>(matched);
            return result;
        }

        template <
            typename EquationOfState,
            typename RelationType>
        [[nodiscard]] bool runtimeRelationMatches(
            const ThermodynamicQuantityId outputQuantity,
            const std::span<const RuntimeQuantityValue> inputValues
        ) {
            const RuntimeRelationDescriptor &descriptor =
                RuntimeRelationStorage<EquationOfState, RelationType>::descriptor;

            if (descriptor.outputQuantity != outputQuantity ||
                descriptor.inputQuantities.size() != inputValues.size()) {
                return false;
            }

            for (std::size_t index = 0; index < inputValues.size(); ++index) {
                if (descriptor.inputQuantities[index] != inputValues[index].quantity) {
                    return false;
                }
            }

            return true;
        }

        template <typename EquationOfState, typename Catalog> struct RuntimeCatalogDispatch;

        template <typename EquationOfState, typename... Relations>
        struct RuntimeCatalogDispatch<EquationOfState, RelationCatalog<Relations...>> {
            [[nodiscard]] static std::expected<
                double,
                EvaluationError>
            evaluate(
                const void *object,
                const ThermodynamicQuantityId outputQuantity,
                const std::span<const RuntimeQuantityValue> inputValues
            ) {
                const auto &equationOfState                   = *static_cast<const EquationOfState *>(object);

                std::expected<double, EvaluationError> result = runtimeEvaluationFailure(
                    EvaluationErrorCode::unsupported_relation, "The requested EOS relation is not available."
                );

                const bool matched =
                    ((runtimeRelationMatches<EquationOfState, Relations>(outputQuantity, inputValues)
                          ? (result = evaluateRuntimeRelation(equationOfState, Relations{}, inputValues), true)
                          : false) ||
                     ...);

                static_cast<void>(matched);
                return result;
            }

            [[nodiscard]] static std::expected<
                double,
                EvaluationError>
            partialDerivative(
                const void *object,
                const ThermodynamicQuantityId outputQuantity,
                const ThermodynamicQuantityId withRespectTo,
                const std::span<const RuntimeQuantityValue> inputValues
            ) {
                const auto &equationOfState                   = *static_cast<const EquationOfState *>(object);

                std::expected<double, EvaluationError> result = runtimeEvaluationFailure(
                    EvaluationErrorCode::unsupported_relation, "The requested EOS relation is not available."
                );

                const bool matched =
                    ((runtimeRelationMatches<EquationOfState, Relations>(outputQuantity, inputValues)
                          ? (result = evaluateRuntimePartialDerivative(
                                 equationOfState, Relations{}, withRespectTo, inputValues
                             ),
                             true)
                          : false) ||
                     ...);

                static_cast<void>(matched);
                return result;
            }
        };

        template <RuntimeEquationOfStateModel EquationOfState>
        using RuntimeAdapter = RuntimeCatalogDispatch<EquationOfState, typename EquationOfState::Relations>;

        template <RuntimeEquationOfStateModel EquationOfState>
        [[nodiscard]] constexpr std::span<const RuntimeRelationDescriptor> runtimeRelationDescriptors() noexcept {
            return RuntimeCatalogStorage<EquationOfState, typename EquationOfState::Relations>::descriptors;
        }
    } // namespace detail

    class EquationOfStateView final {
    public:
        template <RuntimeEquationOfStateModel EquationOfState>
        explicit EquationOfStateView(EquationOfState &equationOfState) noexcept
            : m_object(std::addressof(equationOfState)),
              m_relations(detail::runtimeRelationDescriptors<std::remove_cv_t<EquationOfState>>()),
              m_evaluate(&detail::RuntimeAdapter<std::remove_cv_t<EquationOfState>>::evaluate),
              m_partialDerivative(&detail::RuntimeAdapter<std::remove_cv_t<EquationOfState>>::partialDerivative) {
        }

        [[nodiscard]] std::span<const RuntimeRelationDescriptor> relations() const noexcept {
            return m_relations;
        }

        [[nodiscard]] bool supports(
            const ThermodynamicQuantityId outputQuantity,
            const std::span<const ThermodynamicQuantityId> inputQuantities
        ) const noexcept {
            return findRelation(outputQuantity, inputQuantities) != nullptr;
        }

        template <
            RuntimeIdentifiedThermodynamicQuantity OutputQuantity,
            RuntimeIdentifiedThermodynamicQuantity... InputQuantities>
        [[nodiscard]] bool supports() const noexcept {
            constexpr std::array<ThermodynamicQuantityId, sizeof...(InputQuantities)> inputs{
                thermodynamicQuantityId<InputQuantities>...
            };

            return supports(thermodynamicQuantityId<OutputQuantity>, std::span<const ThermodynamicQuantityId>{inputs});
        }

        [[nodiscard]] std::expected<
            RuntimeQuantityValue,
            EvaluationError>
        tryEvaluate(
            const ThermodynamicQuantityId outputQuantity,
            const std::span<const RuntimeQuantityValue> inputValues
        ) const {
            const auto validation = validateRelationRequest(outputQuantity, inputValues);

            if (!validation.has_value()) {
                return std::unexpected<EvaluationError>{validation.error()};
            }

            auto result = m_evaluate(m_object, outputQuantity, inputValues);
            if (!result.has_value()) {
                return std::unexpected<EvaluationError>{result.error()};
            }

            return RuntimeQuantityValue{outputQuantity, *result};
        }

        template <
            RuntimeIdentifiedThermodynamicQuantity OutputQuantity,
            QuantityValueType... InputValues>
        [[nodiscard]] std::expected<
            QuantityValue<OutputQuantity>,
            EvaluationError>
        tryEvaluate(const InputValues... inputValues) const {
            constexpr bool inputsHaveRuntimeIdentifiers =
                (RuntimeIdentifiedThermodynamicQuantity<QuantityOfT<InputValues>> && ...);

            static_assert(inputsHaveRuntimeIdentifiers, "Every runtime EOS input quantity needs a stable identifier.");

            const std::array<RuntimeQuantityValue, sizeof...(InputValues)> runtimeInputs{
                RuntimeQuantityValue{thermodynamicQuantityId<QuantityOfT<InputValues>>, inputValues.value()}...
            };

            auto result = tryEvaluate(
                thermodynamicQuantityId<OutputQuantity>, std::span<const RuntimeQuantityValue>{runtimeInputs}
            );

            if (!result.has_value()) {
                return std::unexpected<EvaluationError>{result.error()};
            }

            return QuantityValue<OutputQuantity>{result->value};
        }

        [[nodiscard]] std::expected<
            double,
            EvaluationError>
        tryPartialDerivative(
            const ThermodynamicQuantityId outputQuantity,
            const ThermodynamicQuantityId withRespectTo,
            const std::span<const RuntimeQuantityValue> inputValues
        ) const {
            const auto validation = validateRelationRequest(outputQuantity, inputValues);

            if (!validation.has_value()) {
                return std::unexpected<EvaluationError>{validation.error()};
            }

            const RuntimeRelationDescriptor &descriptor = **validation;
            bool derivativeAvailable                    = false;

            for (std::size_t index = 0; index < descriptor.inputQuantities.size(); ++index) {
                if (descriptor.inputQuantities[index] == withRespectTo) {
                    derivativeAvailable = descriptor.hasPartialDerivative(index);
                    break;
                }
            }

            if (!derivativeAvailable) {
                return runtimeFailure<double>(
                    EvaluationErrorCode::unsupported_derivative,
                    "The requested EOS partial derivative is not available."
                );
            }

            return m_partialDerivative(m_object, outputQuantity, withRespectTo, inputValues);
        }

        template <
            RuntimeIdentifiedThermodynamicQuantity OutputQuantity,
            RuntimeIdentifiedThermodynamicQuantity InputQuantity,
            QuantityValueType... InputValues>
        [[nodiscard]] std::expected<
            PartialDerivative<
                OutputQuantity,
                InputQuantity>,
            EvaluationError>
        tryPartialDerivative(const InputValues... inputValues) const {
            constexpr bool inputsHaveRuntimeIdentifiers =
                (RuntimeIdentifiedThermodynamicQuantity<QuantityOfT<InputValues>> && ...);

            static_assert(inputsHaveRuntimeIdentifiers, "Every runtime EOS input quantity needs a stable identifier.");

            const std::array<RuntimeQuantityValue, sizeof...(InputValues)> runtimeInputs{
                RuntimeQuantityValue{thermodynamicQuantityId<QuantityOfT<InputValues>>, inputValues.value()}...
            };

            auto result = tryPartialDerivative(
                thermodynamicQuantityId<OutputQuantity>, thermodynamicQuantityId<InputQuantity>,
                std::span<const RuntimeQuantityValue>{runtimeInputs}
            );

            if (!result.has_value()) {
                return std::unexpected<EvaluationError>{result.error()};
            }

            return PartialDerivative<OutputQuantity, InputQuantity>{*result};
        }

    private:
        using RuntimeEvaluateFunction = std::expected<
            double,
            EvaluationError> (*)(
            const void *,
            ThermodynamicQuantityId,
            std::span<const RuntimeQuantityValue>
        );

        using RuntimePartialDerivativeFunction = std::expected<
            double,
            EvaluationError> (*)(
            const void *,
            ThermodynamicQuantityId,
            ThermodynamicQuantityId,
            std::span<const RuntimeQuantityValue>
        );

        [[nodiscard]] const RuntimeRelationDescriptor *findRelation(
            const ThermodynamicQuantityId outputQuantity,
            const std::span<const ThermodynamicQuantityId> inputQuantities
        ) const noexcept {
            for (const RuntimeRelationDescriptor &descriptor : m_relations) {
                if (descriptor.outputQuantity != outputQuantity ||
                    descriptor.inputQuantities.size() != inputQuantities.size()) {
                    continue;
                }

                bool matches = true;
                for (std::size_t index = 0; index < inputQuantities.size(); ++index) {
                    if (descriptor.inputQuantities[index] != inputQuantities[index]) {
                        matches = false;
                        break;
                    }
                }

                if (matches) {
                    return std::addressof(descriptor);
                }
            }

            return nullptr;
        }

        [[nodiscard]] std::expected<
            const RuntimeRelationDescriptor *,
            EvaluationError>
        validateRelationRequest(
            const ThermodynamicQuantityId outputQuantity,
            const std::span<const RuntimeQuantityValue> inputValues
        ) const {
            bool outputAvailable     = false;
            bool inputCountAvailable = false;

            for (const RuntimeRelationDescriptor &descriptor : m_relations) {
                if (descriptor.outputQuantity != outputQuantity) {
                    continue;
                }

                outputAvailable = true;

                if (descriptor.inputQuantities.size() != inputValues.size()) {
                    continue;
                }

                inputCountAvailable = true;
                bool matches        = true;

                for (std::size_t index = 0; index < inputValues.size(); ++index) {
                    if (descriptor.inputQuantities[index] != inputValues[index].quantity) {
                        matches = false;
                        break;
                    }
                }

                if (matches) {
                    return std::addressof(descriptor);
                }
            }

            if (!outputAvailable) {
                return runtimeFailure<const RuntimeRelationDescriptor *>(
                    EvaluationErrorCode::unsupported_relation,
                    "The EOS does not provide a relation for output quantity '" + std::string{outputQuantity.name()} +
                        "'."
                );
            }

            if (!inputCountAvailable) {
                return runtimeFailure<const RuntimeRelationDescriptor *>(
                    EvaluationErrorCode::wrong_input_count, "No EOS relation for output quantity '" +
                                                                std::string{outputQuantity.name()} +
                                                                "' accepts the supplied number of inputs."
                );
            }

            return runtimeFailure<const RuntimeRelationDescriptor *>(
                EvaluationErrorCode::wrong_input_quantity, "No EOS relation for output quantity '" +
                                                               std::string{outputQuantity.name()} +
                                                               "' accepts the supplied input quantities."
            );
        }

        template <typename Value>
        [[nodiscard]] static std::expected<
            Value,
            EvaluationError>
        runtimeFailure(
            const EvaluationErrorCode code,
            std::string message
        ) {
            return std::unexpected<EvaluationError>{EvaluationError{code, std::move(message)}};
        }

        const void *m_object;
        std::span<const RuntimeRelationDescriptor> m_relations;
        RuntimeEvaluateFunction m_evaluate;
        RuntimePartialDerivativeFunction m_partialDerivative;
    };
} // namespace mean_field::eos