module;

#include <concepts>
#include <type_traits>

#include <mfem.hpp>

export module mean_field:deformation.interior_extension;

export import :deformation.descriptors;

export namespace mean_field::deformation {
    template <typename Candidate>
    concept PreparedInteriorDeformationExtension = requires(
        const std::remove_cvref_t<Candidate> &preparedExtension,
        const mfem::Vector &surfaceDisplacement,
        const mfem::Vector &surfaceDisplacementDirection,
        const mfem::Vector &interiorDisplacementDual,
        mfem::Vector &interiorDisplacement,
        mfem::Vector &surfaceDisplacementDual
    ) {
        { preparedExtension.descriptor() } noexcept -> std::same_as<InteriorDeformationExtensionDescriptor>;
        { preparedExtension.surfaceDisplacementSize() } noexcept -> std::same_as<int>;
        { preparedExtension.interiorDisplacementSize() } noexcept -> std::same_as<int>;
        { preparedExtension.scalarTrueDofCount() } noexcept -> std::same_as<int>;
        { preparedExtension.hasStellarSupport(0) } -> std::same_as<bool>;
        {
            preparedExtension.buildInteriorDisplacement(surfaceDisplacement, interiorDisplacement)
        } -> std::same_as<void>;
        {
            preparedExtension.applyJacobian(surfaceDisplacement, surfaceDisplacementDirection, interiorDisplacement)
        } -> std::same_as<void>;
        {
            preparedExtension.applyJacobianTranspose(
                surfaceDisplacement, interiorDisplacementDual, surfaceDisplacementDual
            )
        } -> std::same_as<void>;
        {
            preparedExtension.applyPullbackDerivative(
                surfaceDisplacement, surfaceDisplacementDirection, interiorDisplacementDual, surfaceDisplacementDual
            )
        } -> std::same_as<void>;
    };

    template <typename Candidate>
    concept InteriorDeformationExtension = requires(const std::remove_cvref_t<Candidate> &extension) {
        typename std::remove_cvref_t<Candidate>::PreparedType;
        requires PreparedInteriorDeformationExtension<typename std::remove_cvref_t<Candidate>::PreparedType>;
        { extension.descriptor() } noexcept -> std::same_as<InteriorDeformationExtensionDescriptor>;
        { extension.validate() } -> std::same_as<void>;
    };

    template <typename Extension, typename CompilationContext>
    concept InteriorDeformationExtensionCompilable =
        InteriorDeformationExtension<Extension> && requires(
                                                       const std::remove_cvref_t<Extension> &extension,
                                                       const std::remove_cvref_t<CompilationContext> &context
                                                   ) {
            {
                compileInteriorDeformationExtension(extension, context)
            } -> std::same_as<typename std::remove_cvref_t<Extension>::PreparedType>;
        };
} // namespace mean_field::deformation