module;

#include <optional>
#include <span>
#include <vector>

#include <mfem/linalg/densemat.hpp>

export module mean_field:deformation.spherical_harmonic_surface;

export import :deformation.nodal_radial_surface;

export namespace mean_field::deformation {
    struct SphericalHarmonicSurfaceOptions final {
        int maximumDegree{4};
        // A cutoff retains both real sine/cosine partners. Zero restricts the
        // shape to axisymmetry; an omitted cutoff retains every order.
        std::optional<int> maximumAzimuthalOrder{};

        void validate() const;
    };

    struct SphericalHarmonicMode final {
        int degree;
        int order;

        constexpr bool operator==(const SphericalHarmonicMode &) const = default;
    };

    class PreparedSphericalHarmonicSurface;

    class SphericalHarmonicSurface final {
    public:
        using PreparedType = PreparedSphericalHarmonicSurface;

        explicit SphericalHarmonicSurface(SphericalHarmonicSurfaceOptions options = {});
        explicit SphericalHarmonicSurface(
            mfem::Vector referenceCenter,
            SphericalHarmonicSurfaceOptions options = {}
        );

        [[nodiscard]] const mfem::Vector &referenceCenter() const noexcept;
        [[nodiscard]] const SphericalHarmonicSurfaceOptions &options() const noexcept;
        [[nodiscard]] int modeCount() const noexcept;
        [[nodiscard]] SurfaceDeformationDescriptor descriptor() const noexcept;
        void validate() const;

    private:
        mfem::Vector m_referenceCenter;
        SphericalHarmonicSurfaceOptions m_options;
    };

    // Coefficients have units of radial length. The fixed reference functions
    // satisfy integral(Y_lm^2 dOmega)=1, with Condon--Shortley phase, negative
    // orders denoting sine and positive orders cosine. FE interpolation of
    // their nodal samples need not retain continuous spherical orthogonality.
    class PreparedSphericalHarmonicSurface final {
    public:
        static constexpr bool hasNodalRadialParameters     = false;
        // Trace entries follow dimension * canonical boundary-node index + component.
        static constexpr bool hasCanonicalNodalVectorTrace = true;

        [[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 surfaceNodeCount() const noexcept;
        [[nodiscard]] int surfaceDisplacementSize() const noexcept;
        [[nodiscard]] long long globalSurfaceDisplacementSize() const noexcept;
        [[nodiscard]] long long globalSurfaceDisplacementOffset() const noexcept;
        [[nodiscard]] int surfaceDisplacementDof(
            int surfaceNode,
            int component
        ) const;
        [[nodiscard]] double radialDirection(
            int surfaceNode,
            int component
        ) const;
        [[nodiscard]] double referenceRadius(int surfaceNode) const;
        [[nodiscard]] const mfem::Vector &referenceCenter() const noexcept;
        [[nodiscard]] const field::ScalarBoundaryDofMap &surfaceDofMap() const noexcept;
        [[nodiscard]] std::span<const SphericalHarmonicMode> modes() const noexcept;
        [[nodiscard]] const mfem::DenseMatrix &nodalBasis() 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 PreparedSphericalHarmonicSurface compileSurfaceDeformationPrescription(
            const SphericalHarmonicSurface &,
            const SurfaceDeformationCompilationContext &
        );

        PreparedSphericalHarmonicSurface(
            PreparedNodalRadialSurface radialSurface,
            std::vector<SphericalHarmonicMode> modes,
            mfem::DenseMatrix nodalBasis
        );
        void requireParameters(const mfem::Vector &parameters) const;
        void requireSurfaceSize(const mfem::Vector &surfaceDisplacement) const;
        void applyEmbedding(
            const mfem::Vector &parameters,
            mfem::Vector &surfaceDisplacement
        ) const;

        PreparedNodalRadialSurface m_radialSurface;
        std::vector<SphericalHarmonicMode> m_modes;
        mfem::DenseMatrix m_nodalBasis;
    };

    [[nodiscard]] PreparedSphericalHarmonicSurface compileSurfaceDeformationPrescription(
        const SphericalHarmonicSurface &prescription,
        const SurfaceDeformationCompilationContext &context
    );

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