module;

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

#include <mpi.h>

export module mean_field:equilibrium.stellar_discretization;

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> class StellarDiscretizationFor final {
    public:
        using NormalizationPrescriptionType = std::remove_cvref_t<Normalization>;

        explicit StellarDiscretizationFor(fem::FEM &&finiteElementModel)
            requires std::same_as<
                NormalizationPrescriptionType,
                normalization::Unnormalized>
            : StellarDiscretizationFor(
                  std::move(finiteElementModel),
                  normalization::Unnormalized{}
              ) {
        }

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

        StellarDiscretizationFor(
            fem::FEM &&finiteElementModel,
            NormalizationPrescriptionType normalizationPrescription
        )
            : m_finiteElementModel(TakeOwnership(std::move(finiteElementModel))),
              m_normalizationPrescription(std::move(normalizationPrescription)) {
        }

        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]] 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 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;
    };

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

    using StellarDiscretization = StellarDiscretizationFor<normalization::Unnormalized>;

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

    template <normalization::NormalizationPrescription Normalization>
    struct IsStellarDiscretization<StellarDiscretizationFor<Normalization>> : 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)
        };
    }

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