module;

#include <mfem.hpp>

export module mean_field:deformation.nodal_radial_surface;

export import :deformation.surface_prescription;
export import :field.mfem;

export namespace mean_field::deformation {
    class PreparedNodalRadialSurface;

    class SurfaceDeformationCompilationContext final {
    public:
        SurfaceDeformationCompilationContext(
            mfem::ParFiniteElementSpace &scalarFiniteElementSpace,
            field::ScalarBoundaryDofMap surfaceDofMap
        );

        [[nodiscard]] mfem::ParFiniteElementSpace &scalarFiniteElementSpace() const noexcept;

        [[nodiscard]] const field::ScalarBoundaryDofMap &surfaceDofMap() const noexcept;

    private:
        mfem::ParFiniteElementSpace *m_scalarFiniteElementSpace;
        field::ScalarBoundaryDofMap m_surfaceDofMap;
    };

    class NodalRadialSurface final {
    public:
        using PreparedType = PreparedNodalRadialSurface;

        explicit NodalRadialSurface(mfem::Vector referenceCenter);

        [[nodiscard]] const mfem::Vector &referenceCenter() const noexcept;

        [[nodiscard]] SurfaceDeformationDescriptor descriptor() const noexcept;

        void validate() const;

    private:
        mfem::Vector m_referenceCenter;
    };

    class PreparedNodalRadialSurface final {
    public:
        [[nodiscard]] SurfaceDeformationDescriptor descriptor() const noexcept;

        [[nodiscard]] int parameterCount() const noexcept;

        [[nodiscard]] long long globalParameterCount() const noexcept;

        [[nodiscard]] long long globalParameterOffset() const noexcept;

        [[nodiscard]] int spatialDimension() const noexcept;

        [[nodiscard]] int surfaceDisplacementSize() const noexcept;

        [[nodiscard]] long long globalSurfaceDisplacementSize() const noexcept;

        [[nodiscard]] long long globalSurfaceDisplacementOffset() const noexcept;

        [[nodiscard]] int surfaceDisplacementDof(
            int parameterDof,
            int component
        ) const;

        [[nodiscard]] double radialDirection(
            int parameterDof,
            int component
        ) const;

        [[nodiscard]] double referenceRadius(int parameterDof) const;

        [[nodiscard]] const mfem::Vector &referenceCenter() const noexcept;

        [[nodiscard]] const field::ScalarBoundaryDofMap &surfaceDofMap() const noexcept;

        void buildSurfaceDisplacement(
            const mfem::Vector &parameters,
            mfem::Vector &surfaceDisplacement
        ) const;

        void applyJacobian(
            const mfem::Vector &parameters,
            const mfem::Vector &parameterDirection,
            mfem::Vector &surfaceDisplacementDirection
        ) const;

        void applyJacobianTranspose(
            const mfem::Vector &parameters,
            const mfem::Vector &surfaceDisplacementDual,
            mfem::Vector &parameterDual
        ) const;

        void applyPullbackDerivative(
            const mfem::Vector &parameters,
            const mfem::Vector &parameterDirection,
            const mfem::Vector &surfaceDisplacementDual,
            mfem::Vector &parameterDualAction
        ) const;

    private:
        friend PreparedNodalRadialSurface compileSurfaceDeformationPrescription(
            const NodalRadialSurface &prescription,
            const SurfaceDeformationCompilationContext &context
        );

        PreparedNodalRadialSurface(
            SurfaceDeformationDescriptor descriptor,
            mfem::Vector referenceCenter,
            field::ScalarBoundaryDofMap surfaceDofMap,
            mfem::Vector radialDirections,
            mfem::Vector referenceRadii
        );

        void requireParameterSize(const mfem::Vector &parameters) const;

        void requireSurfaceDisplacementSize(const mfem::Vector &surfaceDisplacement) const;

        SurfaceDeformationDescriptor m_descriptor;
        mfem::Vector m_referenceCenter;
        field::ScalarBoundaryDofMap m_surfaceDofMap;
        mfem::Vector m_radialDirections;
        mfem::Vector m_referenceRadii;
    };

    [[nodiscard]] PreparedNodalRadialSurface compileSurfaceDeformationPrescription(
        const NodalRadialSurface &prescription,
        const SurfaceDeformationCompilationContext &context
    );

    static_assert(SurfaceDeformationPrescription<NodalRadialSurface>);
    static_assert(PreparedSurfaceDeformationPrescription<PreparedNodalRadialSurface>);
    static_assert(SurfaceDeformationCompilable<
                  NodalRadialSurface,
                  SurfaceDeformationCompilationContext>);
} // namespace mean_field::deformation