module;

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

export module mean_field:equilibrium.stellar_discretization;

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

export namespace mean_field::equilibrium {
    /*
     * An explicit, non-owning view of the numerical discretization used by a
     * stellar equilibrium problem. The referenced FEM and mapper must outlive
     * every problem and structure that uses this view.
     *
     * Ownership cannot move here yet because FEM currently also contains
     * mutable field workspaces. Separating those workspaces is a prerequisite
     * for shared discretization ownership by solved Structure objects.
     */
    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(
                  finiteElementModel,
                  RequireDomainMapper(finiteElementModel),
                  normalization::Unnormalized{}
              ) {
        }

        StellarDiscretizationFor(
            fem::FEM &finiteElementModel,
            const mapping::DomainMapper &domainMapper
        )
            requires std::same_as<NormalizationPrescriptionType, normalization::Unnormalized>
            : StellarDiscretizationFor(
                  finiteElementModel,
                  domainMapper,
                  normalization::Unnormalized{}
              ) {
        }

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

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

        StellarDiscretizationFor(
            fem::FEM &finiteElementModel,
            NormalizationPrescriptionType normalizationPrescription
        )
            : StellarDiscretizationFor(
                  finiteElementModel,
                  RequireDomainMapper(finiteElementModel),
                  std::move(normalizationPrescription)
              ) {
        }

        StellarDiscretizationFor(
            fem::FEM &finiteElementModel,
            const mapping::DomainMapper &domainMapper,
            NormalizationPrescriptionType normalizationPrescription
        )
            : m_finiteElementModel(std::addressof(finiteElementModel)),
              m_domainMapper(std::addressof(domainMapper)),
              m_normalizationPrescription(std::move(normalizationPrescription)) {
            if (!finiteElementModel.okay()) {
                throw std::invalid_argument("A stellar discretization requires a complete finite-element model.");
            }
        }

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

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

        [[nodiscard]] fem::FEM &finiteElementModel() const noexcept {
            return *m_finiteElementModel;
        }

        [[nodiscard]] const mapping::DomainMapper &domainMapper() const noexcept {
            return *m_domainMapper;
        }

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

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

    private:
        [[nodiscard]] static const mapping::DomainMapper &RequireDomainMapper(const fem::FEM &finiteElementModel) {
            if (finiteElementModel.domainMapperStateless == nullptr) {
                throw std::invalid_argument("A stellar discretization requires a domain mapper.");
            }
            return *finiteElementModel.domainMapperStateless;
        }

        fem::FEM *m_finiteElementModel;
        const mapping::DomainMapper *m_domainMapper;
        NormalizationPrescriptionType m_normalizationPrescription;
    };

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

    template <normalization::NormalizationPrescription Normalization>
    StellarDiscretizationFor(fem::FEM &, const mapping::DomainMapper &, 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>>{
            finiteElementModel,
            std::move(normalizationPrescription)
        };
    }

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

    template <normalization::NormalizationPrescription Normalization>
    StellarDiscretizationFor<std::remove_cvref_t<Normalization>>
    makeStellarDiscretization(
        fem::FEM &,
        mapping::DomainMapper &&,
        Normalization
    ) = delete;

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