module;

#include <utility>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

export module mean_field:deformation.radial_extensions;

export import :deformation.interior_extension;
export import :deformation.vacuum_extension;
export import :field.mfem;
export import :utils.domain;

export namespace mean_field::deformation {
    class RadialDeformationExtensionCompilationContext final {
    public:
        RadialDeformationExtensionCompilationContext(
            mfem::ParFiniteElementSpace &scalarFiniteElementSpace,
            mfem::ParFiniteElementSpace &vectorFiniteElementSpace,
            mfem::ParMesh &logicalReferenceMesh,
            field::ScalarBoundaryDofMap stellarSurfaceDofMap,
            field::ScalarBoundaryDofMap infinitySurfaceDofMap,
            mfem::Array<int> stellarMaterialMarker,
            mfem::Array<int> vacuumMaterialMarker,
            int stellarSurfaceBoundaryAttribute,
            int infinitySurfaceBoundaryAttribute
        );

        [[nodiscard]] int spatialDimension() const noexcept;
        [[nodiscard]] int surfaceDisplacementSize() const noexcept;
        [[nodiscard]] int volumeDisplacementSize() const noexcept;
        [[nodiscard]] int scalarTrueDofCount() const noexcept;
        [[nodiscard]] double logicalRadius(int scalarTrueDof) const;
        [[nodiscard]] double stellarSurfaceLogicalRadius() const noexcept;
        [[nodiscard]] double infinitySurfaceLogicalRadius() const noexcept;
        [[nodiscard]] int surfaceInterpolationEntryCount(int scalarTrueDof) const;
        [[nodiscard]] int surfaceGlobalCoordinate(
            int scalarTrueDof,
            int interpolationEntry
        ) const;
        [[nodiscard]] double surfaceInterpolationWeight(
            int scalarTrueDof,
            int interpolationEntry
        ) const;

    private:
        friend class PreparedPowerLawRadialInteriorExtension;
        friend class PreparedFixedInfinityRadialVacuumExtension;

        int m_spatialDimension{0};
        int m_scalarTrueDofCount{0};
        int m_volumeDisplacementSize{0};
        int m_surfaceDisplacementSize{0};
        int m_globalSurfaceDisplacementSize{0};
        int m_globalSurfaceDisplacementOffset{0};
        MPI_Comm m_communicator{MPI_COMM_NULL};
        mfem::Array<int> m_stellarSupport;
        mfem::Array<int> m_vacuumSupport;
        mfem::Vector m_logicalRadius;
        double m_stellarSurfaceLogicalRadius{0.0};
        double m_infinitySurfaceLogicalRadius{0.0};
        std::vector<int> m_surfaceInterpolationRowOffsets;
        std::vector<int> m_surfaceInterpolationGlobalCoordinates;
        std::vector<double> m_surfaceInterpolationWeights;
        std::vector<int> m_surfaceDisplacementCounts;
        std::vector<int> m_surfaceDisplacementOffsets;
    };

    template <utils::domain::IsSchema SchemaT = utils::domain::CoreEnvelopeVacuumDomainSchema>
        requires(
            SchemaT::template contains_domain<utils::domain::Stellar>() &&
            SchemaT::template contains_domain<utils::domain::Vacuum>() &&
            SchemaT::template contains_boundary<utils::domain::StellarSurface>() &&
            SchemaT::template contains_boundary<utils::domain::InfinitySurface>()
        )
    [[nodiscard]] RadialDeformationExtensionCompilationContext makeRadialDeformationExtensionCompilationContext(
        mfem::ParFiniteElementSpace &scalarFiniteElementSpace,
        mfem::ParFiniteElementSpace &vectorFiniteElementSpace,
        mfem::ParMesh &logicalReferenceMesh
    ) {
        const mfem::Mesh *mesh = scalarFiniteElementSpace.GetMesh();
        MFEM_VERIFY(mesh != nullptr, "Radial deformation extension compilation requires an MFEM mesh.");

        return RadialDeformationExtensionCompilationContext(
            scalarFiniteElementSpace, vectorFiniteElementSpace, logicalReferenceMesh,
            field::make_scalar_boundary_dof_map<utils::domain::StellarSurface, SchemaT>(scalarFiniteElementSpace),
            field::make_scalar_boundary_dof_map<utils::domain::InfinitySurface, SchemaT>(scalarFiniteElementSpace),
            utils::domain::make_attribute_marker<utils::domain::Stellar, SchemaT>(*mesh),
            utils::domain::make_attribute_marker<utils::domain::Vacuum, SchemaT>(*mesh),
            SchemaT::template boundary_attribute<utils::domain::StellarSurface>(),
            SchemaT::template boundary_attribute<utils::domain::InfinitySurface>()
        );
    }

    class PreparedPowerLawRadialInteriorExtension;

    class PowerLawRadialInteriorExtension final {
    public:
        using PreparedType = PreparedPowerLawRadialInteriorExtension;

        explicit PowerLawRadialInteriorExtension(double radialPower = 2.0);
        [[nodiscard]] double radialPower() const noexcept;
        [[nodiscard]] InteriorDeformationExtensionDescriptor descriptor() const noexcept;
        void validate() const;

    private:
        double m_radialPower;
    };

    class PreparedPowerLawRadialInteriorExtension final {
    public:
        [[nodiscard]] InteriorDeformationExtensionDescriptor descriptor() const noexcept;
        [[nodiscard]] int surfaceDisplacementSize() const noexcept;
        [[nodiscard]] int interiorDisplacementSize() const noexcept;
        [[nodiscard]] int scalarTrueDofCount() const noexcept;
        [[nodiscard]] double radialPower() const noexcept;
        [[nodiscard]] bool hasStellarSupport(int scalarTrueDof) const;
        [[nodiscard]] double radialWeight(int scalarTrueDof) const;
        [[nodiscard]] int surfaceInterpolationEntryCount(int scalarTrueDof) const;
        [[nodiscard]] int surfaceGlobalCoordinate(
            int scalarTrueDof,
            int interpolationEntry
        ) const;
        [[nodiscard]] double surfaceInterpolationWeight(
            int scalarTrueDof,
            int interpolationEntry
        ) const;

        void buildInteriorDisplacement(
            const mfem::Vector &surfaceDisplacement,
            mfem::Vector &interiorDisplacement
        ) const;
        void applyJacobian(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &surfaceDisplacementDirection,
            mfem::Vector &interiorDisplacementDirection
        ) const;
        void applyJacobianTranspose(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &interiorDisplacementDual,
            mfem::Vector &surfaceDisplacementDual
        ) const;
        void applyPullbackDerivative(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &surfaceDisplacementDirection,
            const mfem::Vector &interiorDisplacementDual,
            mfem::Vector &surfaceDisplacementDualAction
        ) const;

    private:
        friend PreparedPowerLawRadialInteriorExtension compileInteriorDeformationExtension(
            const PowerLawRadialInteriorExtension &extension,
            const RadialDeformationExtensionCompilationContext &context
        );

        PreparedPowerLawRadialInteriorExtension(
            const PowerLawRadialInteriorExtension &extension,
            const RadialDeformationExtensionCompilationContext &context
        );
        void requireSurfaceSize(const mfem::Vector &surfaceDisplacement) const;
        void requireInteriorSize(const mfem::Vector &interiorDisplacement) const;
        void applyForward(
            const mfem::Vector &surfaceDisplacement,
            mfem::Vector &interiorDisplacement
        ) const;
        void applyTranspose(
            const mfem::Vector &interiorDisplacementDual,
            mfem::Vector &surfaceDisplacementDual
        ) const;

        InteriorDeformationExtensionDescriptor m_descriptor;
        double m_radialPower{2.0};
        int m_surfaceDisplacementSize{0};
        int m_interiorDisplacementSize{0};
        int m_spatialDimension{0};
        int m_globalSurfaceDisplacementSize{0};
        int m_globalSurfaceDisplacementOffset{0};
        MPI_Comm m_communicator{MPI_COMM_NULL};
        mfem::Array<int> m_stellarSupport;
        mfem::Vector m_radialWeights;
        std::vector<int> m_surfaceInterpolationRowOffsets;
        std::vector<int> m_surfaceInterpolationGlobalCoordinates;
        std::vector<double> m_surfaceInterpolationWeights;
        std::vector<int> m_surfaceDisplacementCounts;
        std::vector<int> m_surfaceDisplacementOffsets;
        mutable mfem::Vector m_globalSurfaceDisplacementWorkspace;
        mutable mfem::Vector m_localGlobalSurfaceDualWorkspace;
        mutable mfem::Vector m_globalSurfaceDualWorkspace;
    };

    [[nodiscard]] PreparedPowerLawRadialInteriorExtension compileInteriorDeformationExtension(
        const PowerLawRadialInteriorExtension &extension,
        const RadialDeformationExtensionCompilationContext &context
    );

    class PreparedFixedInfinityRadialVacuumExtension;

    class FixedInfinityRadialVacuumExtension final {
    public:
        using PreparedType = PreparedFixedInfinityRadialVacuumExtension;

        [[nodiscard]] VacuumDeformationExtensionDescriptor descriptor() const noexcept;
        void validate() const;
    };

    class PreparedFixedInfinityRadialVacuumExtension final {
    public:
        [[nodiscard]] VacuumDeformationExtensionDescriptor descriptor() const noexcept;
        [[nodiscard]] int surfaceDisplacementSize() const noexcept;
        [[nodiscard]] int vacuumDisplacementSize() const noexcept;
        [[nodiscard]] int scalarTrueDofCount() const noexcept;
        [[nodiscard]] bool hasVacuumSupport(int scalarTrueDof) const;
        [[nodiscard]] double radialWeight(int scalarTrueDof) const;
        [[nodiscard]] int surfaceInterpolationEntryCount(int scalarTrueDof) const;
        [[nodiscard]] int surfaceGlobalCoordinate(
            int scalarTrueDof,
            int interpolationEntry
        ) const;
        [[nodiscard]] double surfaceInterpolationWeight(
            int scalarTrueDof,
            int interpolationEntry
        ) const;

        void buildVacuumDisplacement(
            const mfem::Vector &surfaceDisplacement,
            mfem::Vector &vacuumDisplacement
        ) const;
        void applyJacobian(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &surfaceDisplacementDirection,
            mfem::Vector &vacuumDisplacementDirection
        ) const;
        void applyJacobianTranspose(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &vacuumDisplacementDual,
            mfem::Vector &surfaceDisplacementDual
        ) const;
        void applyPullbackDerivative(
            const mfem::Vector &surfaceDisplacement,
            const mfem::Vector &surfaceDisplacementDirection,
            const mfem::Vector &vacuumDisplacementDual,
            mfem::Vector &surfaceDisplacementDualAction
        ) const;

    private:
        friend PreparedFixedInfinityRadialVacuumExtension compileVacuumDeformationExtension(
            const FixedInfinityRadialVacuumExtension &extension,
            const RadialDeformationExtensionCompilationContext &context
        );

        PreparedFixedInfinityRadialVacuumExtension(
            const FixedInfinityRadialVacuumExtension &extension,
            const RadialDeformationExtensionCompilationContext &context
        );
        void requireSurfaceSize(const mfem::Vector &surfaceDisplacement) const;
        void requireVacuumSize(const mfem::Vector &vacuumDisplacement) const;
        void applyForward(
            const mfem::Vector &surfaceDisplacement,
            mfem::Vector &vacuumDisplacement
        ) const;
        void applyTranspose(
            const mfem::Vector &vacuumDisplacementDual,
            mfem::Vector &surfaceDisplacementDual
        ) const;

        VacuumDeformationExtensionDescriptor m_descriptor;
        int m_surfaceDisplacementSize{0};
        int m_vacuumDisplacementSize{0};
        int m_spatialDimension{0};
        int m_globalSurfaceDisplacementSize{0};
        int m_globalSurfaceDisplacementOffset{0};
        MPI_Comm m_communicator{MPI_COMM_NULL};
        mfem::Array<int> m_vacuumSupport;
        mfem::Vector m_radialWeights;
        std::vector<int> m_surfaceInterpolationRowOffsets;
        std::vector<int> m_surfaceInterpolationGlobalCoordinates;
        std::vector<double> m_surfaceInterpolationWeights;
        std::vector<int> m_surfaceDisplacementCounts;
        std::vector<int> m_surfaceDisplacementOffsets;
        mutable mfem::Vector m_globalSurfaceDisplacementWorkspace;
        mutable mfem::Vector m_localGlobalSurfaceDualWorkspace;
        mutable mfem::Vector m_globalSurfaceDualWorkspace;
    };

    [[nodiscard]] PreparedFixedInfinityRadialVacuumExtension compileVacuumDeformationExtension(
        const FixedInfinityRadialVacuumExtension &extension,
        const RadialDeformationExtensionCompilationContext &context
    );

    static_assert(InteriorDeformationExtension<PowerLawRadialInteriorExtension>);
    static_assert(PreparedInteriorDeformationExtension<PreparedPowerLawRadialInteriorExtension>);
    static_assert(InteriorDeformationExtensionCompilable<
                  PowerLawRadialInteriorExtension,
                  RadialDeformationExtensionCompilationContext>);
    static_assert(VacuumDeformationExtension<FixedInfinityRadialVacuumExtension>);
    static_assert(PreparedVacuumDeformationExtension<PreparedFixedInfinityRadialVacuumExtension>);
    static_assert(VacuumDeformationExtensionCompilable<
                  FixedInfinityRadialVacuumExtension,
                  RadialDeformationExtensionCompilationContext>);
} // namespace mean_field::deformation