module;

#include <concepts>
#include <memory>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <mfem/linalg/vector.hpp>
#include <mpi.h>

export module mean_field:equilibrium.stellar_discretization;

export import :deformation.domain_deformation;
export import :fem;
export import :mapping.domain_mapper;
export import :normalization.physical_riesz;

namespace mean_field::equilibrium::detail {
    struct StellarEquilibriumProblemFactory;
}

export namespace mean_field::equilibrium {
    /*
     * The complete numerical discretization used by a stellar equilibrium
     * problem. Moving the FEM into stable heap storage lets the problem and a
     * completed Structure transfer unique ownership without invalidating the
     * references retained by prepared operators. The mapper is part of that
     * owned FEM and therefore has the same lifetime.
     */
    template <
        normalization::NormalizationPrescription Normalization,
        deformation::SurfaceDeformationPrescription SurfaceDeformation = deformation::NodalRadialSurface,
        deformation::InteriorDeformationExtension InteriorDeformation  = deformation::PowerLawRadialInteriorExtension,
        deformation::VacuumDeformationExtension VacuumDeformation = deformation::FixedInfinityRadialVacuumExtension>
        requires deformation::SurfaceDeformationCompilable<
                     SurfaceDeformation,
                     deformation::SurfaceDeformationCompilationContext> &&
                 deformation::InteriorDeformationExtensionCompilable<
                     InteriorDeformation,
                     deformation::RadialDeformationExtensionCompilationContext> &&
                 deformation::VacuumDeformationExtensionCompilable<
                     VacuumDeformation,
                     deformation::RadialDeformationExtensionCompilationContext>
    class StellarDiscretizationFor final {
    public:
        using NormalizationPrescriptionType      = std::remove_cvref_t<Normalization>;
        using SurfaceDeformationPrescriptionType = SurfaceDeformation;
        using InteriorDeformationExtensionType   = InteriorDeformation;
        using VacuumDeformationExtensionType     = VacuumDeformation;

        explicit StellarDiscretizationFor(fem::FEM &&finiteElementModel)
            requires std::same_as<
                         NormalizationPrescriptionType,
                         normalization::Unnormalized> &&
                     std::same_as<
                         SurfaceDeformation,
                         deformation::NodalRadialSurface> &&
                     std::same_as<
                         InteriorDeformation,
                         deformation::PowerLawRadialInteriorExtension> &&
                     std::same_as<
                         VacuumDeformation,
                         deformation::FixedInfinityRadialVacuumExtension>
            : StellarDiscretizationFor(
                  std::move(finiteElementModel),
                  normalization::Unnormalized{}
              ) {
        }

        explicit StellarDiscretizationFor(fem::FEM &)
            requires std::same_as<
                         NormalizationPrescriptionType,
                         normalization::Unnormalized>
        = delete;

        StellarDiscretizationFor(
            fem::FEM &&finiteElementModel,
            NormalizationPrescriptionType normalizationPrescription
        )
            requires std::same_as<
                         SurfaceDeformation,
                         deformation::NodalRadialSurface> &&
                     std::same_as<
                         InteriorDeformation,
                         deformation::PowerLawRadialInteriorExtension> &&
                     std::same_as<
                         VacuumDeformation,
                         deformation::FixedInfinityRadialVacuumExtension>
            : StellarDiscretizationFor(
                  std::move(finiteElementModel),
                  std::move(normalizationPrescription),
                  MakeDefaultSurface(finiteElementModel),
                  InteriorDeformation{},
                  VacuumDeformation{}
              ) {
        }

        StellarDiscretizationFor(
            fem::FEM &&finiteElementModel,
            NormalizationPrescriptionType normalizationPrescription,
            SurfaceDeformation surfaceDeformation
        )
            requires std::default_initializable<InteriorDeformation> && std::default_initializable<VacuumDeformation>
            : StellarDiscretizationFor(
                  std::move(finiteElementModel),
                  std::move(normalizationPrescription),
                  std::move(surfaceDeformation),
                  InteriorDeformation{},
                  VacuumDeformation{}
              ) {
        }

        StellarDiscretizationFor(
            fem::FEM &&finiteElementModel,
            NormalizationPrescriptionType normalizationPrescription,
            SurfaceDeformation surfaceDeformation,
            InteriorDeformation interiorDeformation,
            VacuumDeformation vacuumDeformation
        )
            : m_finiteElementModel(TakeOwnership(std::move(finiteElementModel))),
              m_normalizationPrescription(std::move(normalizationPrescription)),
              m_surfaceDeformation(std::move(surfaceDeformation)),
              m_interiorDeformation(std::move(interiorDeformation)),
              m_vacuumDeformation(std::move(vacuumDeformation)) {
            m_surfaceDeformation.validate();
            m_interiorDeformation.validate();
            m_vacuumDeformation.validate();
        }

        StellarDiscretizationFor(
            fem::FEM &,
            NormalizationPrescriptionType
        )                                                                     = delete;

        StellarDiscretizationFor(const StellarDiscretizationFor &)            = delete;
        StellarDiscretizationFor &operator=(const StellarDiscretizationFor &) = delete;
        StellarDiscretizationFor(StellarDiscretizationFor &&)                 = default;
        StellarDiscretizationFor &operator=(StellarDiscretizationFor &&)      = delete;

        [[nodiscard]] const mapping::DomainMapper &domainMapper() const & {
            const auto &finiteElementModel = RequireFiniteElementModel();
            if (finiteElementModel.domainMapperStateless == nullptr) {
                throw std::logic_error("The stellar discretization has no domain mapper.");
            }
            return *finiteElementModel.domainMapperStateless;
        }

        [[nodiscard]] const mapping::DomainMapper &domainMapper() const && = delete;

        [[nodiscard]] MPI_Comm communicator() const & {
            const auto &finiteElementModel = RequireFiniteElementModel();
            if (finiteElementModel.mesh == nullptr) {
                throw std::logic_error("The stellar discretization has no parallel mesh.");
            }
            return finiteElementModel.mesh->GetComm();
        }

        [[nodiscard]] MPI_Comm communicator() const && = delete;

        [[nodiscard]] const NormalizationPrescriptionType &normalizationPrescription() const & noexcept {
            return m_normalizationPrescription;
        }

        [[nodiscard]] const NormalizationPrescriptionType &normalizationPrescription() const && = delete;

        [[nodiscard]] const SurfaceDeformation &surfaceDeformationPrescription() const & noexcept {
            return m_surfaceDeformation;
        }
        [[nodiscard]] const SurfaceDeformation &surfaceDeformationPrescription() const && = delete;

        [[nodiscard]] const InteriorDeformation &interiorDeformationExtension() const & noexcept {
            return m_interiorDeformation;
        }
        [[nodiscard]] const InteriorDeformation &interiorDeformationExtension() const && = delete;

        [[nodiscard]] const VacuumDeformation &vacuumDeformationExtension() const & noexcept {
            return m_vacuumDeformation;
        }
        [[nodiscard]] const VacuumDeformation &vacuumDeformationExtension() const && = delete;

        [[nodiscard]] bool isCurrent() const noexcept {
            return m_finiteElementModel != nullptr && m_finiteElementModel->okay();
        }

    private:
        friend struct detail::StellarEquilibriumProblemFactory;

        [[nodiscard]] fem::FEM &MutableFiniteElementModelForAssembly() & {
            return const_cast<fem::FEM &>(RequireFiniteElementModel());
        }

        [[nodiscard]] const fem::FEM &RequireFiniteElementModel() const {
            if (m_finiteElementModel == nullptr) {
                throw std::logic_error("A moved-from stellar discretization has no finite-element model.");
            }
            return *m_finiteElementModel;
        }

        [[nodiscard]] static deformation::NodalRadialSurface MakeDefaultSurface(const fem::FEM &finiteElementModel) {
            if (!finiteElementModel.okay()) {
                throw std::invalid_argument("A default surface requires a complete finite-element model.");
            }
            mfem::Vector center(finiteElementModel.mesh->SpaceDimension());
            center = 0.0;
            return deformation::NodalRadialSurface{std::move(center)};
        }

        [[nodiscard]] static std::unique_ptr<fem::FEM> TakeOwnership(fem::FEM &&finiteElementModel) {
            if (!finiteElementModel.okay()) {
                throw std::invalid_argument("A stellar discretization requires a complete finite-element model.");
            }
            return std::make_unique<fem::FEM>(std::move(finiteElementModel));
        }

        std::unique_ptr<fem::FEM> m_finiteElementModel;
        NormalizationPrescriptionType m_normalizationPrescription;
        SurfaceDeformation m_surfaceDeformation;
        InteriorDeformation m_interiorDeformation;
        VacuumDeformation m_vacuumDeformation;
    };

    template <normalization::NormalizationPrescription Normalization>
    StellarDiscretizationFor(
        fem::FEM &&,
        Normalization
    ) -> StellarDiscretizationFor<std::remove_cvref_t<Normalization>>;

    template <
        typename Normalization,
        typename Surface>
    StellarDiscretizationFor(
        fem::FEM &&,
        Normalization,
        Surface
    )
        -> StellarDiscretizationFor<
            std::remove_cvref_t<Normalization>,
            std::remove_cvref_t<Surface>>;

    template <
        typename Normalization,
        typename Surface,
        typename Interior,
        typename Vacuum>
    StellarDiscretizationFor(
        fem::FEM &&,
        Normalization,
        Surface,
        Interior,
        Vacuum
    )
        -> StellarDiscretizationFor<
            std::remove_cvref_t<Normalization>,
            std::remove_cvref_t<Surface>,
            std::remove_cvref_t<Interior>,
            std::remove_cvref_t<Vacuum>>;

    using StellarDiscretization = StellarDiscretizationFor<normalization::Unnormalized>;

    template <typename Candidate> struct IsStellarDiscretization : std::false_type { };

    template <typename Normalization, typename Surface, typename Interior, typename Vacuum>
    struct IsStellarDiscretization<StellarDiscretizationFor<Normalization, Surface, Interior, Vacuum>>
        : std::true_type { };

    template <typename Candidate>
    concept StellarDiscretizationType = IsStellarDiscretization<std::remove_cvref_t<Candidate>>::value;

    template <normalization::NormalizationPrescription Normalization>
    [[nodiscard]] auto makeStellarDiscretization(
        fem::FEM &&finiteElementModel,
        Normalization normalizationPrescription
    ) {
        return StellarDiscretizationFor<std::remove_cvref_t<Normalization>>{
            std::move(finiteElementModel), std::move(normalizationPrescription)
        };
    }

    // Geometry is a numerical prescription. Keeping it on the discretization
    // leaves the model's EOS, surface pressure and integral constraints intact.
    template <
        normalization::NormalizationPrescription Normalization,
        deformation::SurfaceDeformationPrescription Surface>
        requires deformation::SurfaceDeformationCompilable<
            Surface,
            deformation::SurfaceDeformationCompilationContext>
    [[nodiscard]] auto makeStellarDiscretization(
        fem::FEM &&finiteElementModel,
        Normalization normalizationPrescription,
        Surface surface
    ) {
        return StellarDiscretizationFor<Normalization, Surface>{
            std::move(finiteElementModel), std::move(normalizationPrescription), std::move(surface)
        };
    }

    template <
        normalization::NormalizationPrescription Normalization,
        deformation::SurfaceDeformationPrescription Surface,
        deformation::InteriorDeformationExtension Interior,
        deformation::VacuumDeformationExtension Vacuum>
        requires deformation::SurfaceDeformationCompilable<
                     Surface,
                     deformation::SurfaceDeformationCompilationContext> &&
                 deformation::InteriorDeformationExtensionCompilable<
                     Interior,
                     deformation::RadialDeformationExtensionCompilationContext> &&
                 deformation::VacuumDeformationExtensionCompilable<
                     Vacuum,
                     deformation::RadialDeformationExtensionCompilationContext>
    [[nodiscard]] auto makeStellarDiscretization(
        fem::FEM &&finiteElementModel,
        Normalization normalizationPrescription,
        Surface surface,
        Interior interior,
        Vacuum vacuum
    ) {
        return StellarDiscretizationFor<Normalization, Surface, Interior, Vacuum>{
            std::move(finiteElementModel), std::move(normalizationPrescription), std::move(surface),
            std::move(interior), std::move(vacuum)
        };
    }

    template <normalization::NormalizationPrescription Normalization>
    StellarDiscretizationFor<std::remove_cvref_t<Normalization>> makeStellarDiscretization(
        fem::FEM &,
        Normalization
    ) = delete;
} // namespace mean_field::equilibrium