module;

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

#include <mfem.hpp>

export module mean_field:model.stellar;

export import :deformation.domain_deformation;
export import :eos.runtime;
export import :model.structure.base;
export import :surface.compiler;

export namespace mean_field::models {
    namespace detail {
        template <typename Candidate>
        concept ConstEquationOfStateReference =
            std::is_lvalue_reference_v<Candidate> && std::is_const_v<std::remove_reference_t<Candidate>> &&
            eos::EquationOfStateModel<std::remove_cvref_t<Candidate>>;
    } // namespace detail

    template <typename Candidate>
    concept StructurePrescription = requires(
        const std::remove_cvref_t<Candidate> &structurePrescription,
        const structure::StructureSeedRequest &seedRequest
    ) {
        { structurePrescription.equationOfState() } noexcept -> detail::ConstEquationOfStateReference;
        { structurePrescription.targetMass() } noexcept -> std::same_as<double>;
        { structurePrescription.makeInitialSeed(seedRequest) } -> std::same_as<structure::StructureSeed>;
        { structurePrescription.validate() } -> std::same_as<void>;
    };

    template <StructurePrescription Candidate>
    using StructureEquationOfStateT =
        std::remove_cvref_t<decltype(std::declval<const std::remove_cvref_t<Candidate> &>().equationOfState())>;

    template <typename Candidate, typename EquationOfState>
    concept SurfaceCondition =
        surface::ConstantPressureSurfaceType<Candidate> &&
        surface::PressureSurfaceCompilable<surface::BarotropicSurfaceFormulation, std::remove_cvref_t<EquationOfState>>;

    template <
        StructurePrescription Structure,
        deformation::SurfaceDeformationPrescription SurfaceDeformation = deformation::NodalRadialSurface,
        deformation::InteriorDeformationExtension StellarInteriorDeformation =
            deformation::PowerLawRadialInteriorExtension,
        deformation::VacuumDeformationExtension VacuumDeformation = deformation::FixedInfinityRadialVacuumExtension>
        requires SurfaceCondition<surface::ConstantPressureSurface, StructureEquationOfStateT<Structure>>
    class StellarModel final {
    public:
        using StructurePrescriptionType               = Structure;
        using SurfaceConditionType                    = surface::ConstantPressureSurface;
        using SurfaceDeformationPrescriptionType      = SurfaceDeformation;
        using StellarInteriorDeformationExtensionType = StellarInteriorDeformation;
        using VacuumDeformationExtensionType          = VacuumDeformation;
        using EquationOfStateType                     = StructureEquationOfStateT<Structure>;
        using SurfaceConstraintType =
            surface::CompiledPressureSurfaceConstraintT<surface::BarotropicSurfaceFormulation, EquationOfStateType>;

        template <typename StructureArgument>
            requires std::same_as<
                std::remove_cvref_t<StructureArgument>,
                Structure>
        explicit StellarModel(
            StructureArgument &&structurePrescription,
            const surface::ConstantPressureSurface surfaceCondition
        )
            : StellarModel(
                  std::forward<StructureArgument>(structurePrescription),
                  surfaceCondition,
                  deformation::NodalRadialSurface{defaultReferenceCenter()},
                  deformation::PowerLawRadialInteriorExtension{},
                  deformation::FixedInfinityRadialVacuumExtension{}
              ) {
        }

        template <
            typename StructureArgument,
            typename SurfaceDeformationArgument,
            typename StellarInteriorDeformationArgument,
            typename VacuumDeformationArgument>
            requires std::same_as<
                         std::remove_cvref_t<StructureArgument>,
                         Structure> &&
                         std::same_as<
                             std::remove_cvref_t<SurfaceDeformationArgument>,
                             SurfaceDeformation> &&
                         std::same_as<
                             std::remove_cvref_t<StellarInteriorDeformationArgument>,
                             StellarInteriorDeformation> &&
                         std::same_as<
                             std::remove_cvref_t<VacuumDeformationArgument>,
                             VacuumDeformation>
        explicit StellarModel(
            StructureArgument &&structurePrescription,
            const surface::ConstantPressureSurface surfaceCondition,
            SurfaceDeformationArgument &&surfaceDeformation,
            StellarInteriorDeformationArgument &&stellarInteriorDeformation,
            VacuumDeformationArgument &&vacuumDeformation
        )
            : m_structurePrescription(
                  std::make_unique<Structure>(std::forward<StructureArgument>(structurePrescription))
              ),
              m_surfaceCondition(std::make_unique<surface::ConstantPressureSurface>(surfaceCondition)),
              m_surfaceDeformation(
                  std::make_unique<SurfaceDeformation>(std::forward<SurfaceDeformationArgument>(surfaceDeformation))
              ),
              m_stellarInteriorDeformation(
                  std::make_unique<StellarInteriorDeformation>(
                      std::forward<StellarInteriorDeformationArgument>(stellarInteriorDeformation)
                  )
              ),
              m_vacuumDeformation(
                  std::make_unique<VacuumDeformation>(std::forward<VacuumDeformationArgument>(vacuumDeformation))
              ),
              m_compiledSurfaceConstraint(std::make_unique<SurfaceConstraintType>(validateAndCompileConfiguration())) {
        }

        ~StellarModel()                                   = default;

        StellarModel(const StellarModel &)                = delete;
        StellarModel &operator=(const StellarModel &)     = delete;
        StellarModel(StellarModel &&) noexcept            = default;
        StellarModel &operator=(StellarModel &&) noexcept = default;

        [[nodiscard]] const Structure &structurePrescription() const noexcept {
            return *m_structurePrescription;
        }

        [[nodiscard]] const surface::ConstantPressureSurface &surfaceCondition() const noexcept {
            return *m_surfaceCondition;
        }

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

        [[nodiscard]] const StellarInteriorDeformation &stellarInteriorDeformationExtension() const noexcept {
            return *m_stellarInteriorDeformation;
        }

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

        [[nodiscard]] const EquationOfStateType &equationOfState() const noexcept {
            return m_structurePrescription->equationOfState();
        }

        [[nodiscard]] double targetMass() const noexcept {
            return m_structurePrescription->targetMass();
        }

        [[nodiscard]] structure::StructureSeed makeInitialSeed(const structure::StructureSeedRequest &request) const {
            return m_structurePrescription->makeInitialSeed(request);
        }

        [[nodiscard]] const SurfaceConstraintType &compiledSurfaceConstraint() const noexcept {
            return *m_compiledSurfaceConstraint;
        }

        [[nodiscard]] auto compileDomainDeformation(fem::FEM &finiteElementModel) const {
            return deformation::compileDomainDeformation(
                surfaceDeformationPrescription(), stellarInteriorDeformationExtension(), vacuumDeformationExtension(),
                finiteElementModel
            );
        }

    private:
        [[nodiscard]] static mfem::Vector defaultReferenceCenter() {
            mfem::Vector center(3);
            center = 0.0;
            return center;
        }

        [[nodiscard]] SurfaceConstraintType validateAndCompileConfiguration() const {
            m_structurePrescription->validate();
            m_surfaceDeformation->validate();
            m_stellarInteriorDeformation->validate();
            m_vacuumDeformation->validate();

            return surface::compilePressureSurfaceConstraint<surface::BarotropicSurfaceFormulation>(
                *m_surfaceCondition, m_structurePrescription->equationOfState()
            );
        }

        std::unique_ptr<Structure> m_structurePrescription;
        std::unique_ptr<surface::ConstantPressureSurface> m_surfaceCondition;
        std::unique_ptr<SurfaceDeformation> m_surfaceDeformation;
        std::unique_ptr<StellarInteriorDeformation> m_stellarInteriorDeformation;
        std::unique_ptr<VacuumDeformation> m_vacuumDeformation;
        std::unique_ptr<SurfaceConstraintType> m_compiledSurfaceConstraint;
    };

    template <typename Structure>
    StellarModel(
        Structure &&,
        surface::ConstantPressureSurface
    )
        -> StellarModel<
            std::remove_cvref_t<Structure>,
            deformation::NodalRadialSurface,
            deformation::PowerLawRadialInteriorExtension,
            deformation::FixedInfinityRadialVacuumExtension>;

    template <
        typename Structure,
        typename SurfaceDeformation,
        typename StellarInteriorDeformation,
        typename VacuumDeformation>
    StellarModel(
        Structure &&,
        surface::ConstantPressureSurface,
        SurfaceDeformation &&,
        StellarInteriorDeformation &&,
        VacuumDeformation &&
    )
        -> StellarModel<
            std::remove_cvref_t<Structure>,
            std::remove_cvref_t<SurfaceDeformation>,
            std::remove_cvref_t<StellarInteriorDeformation>,
            std::remove_cvref_t<VacuumDeformation>>;

    namespace detail {
        template <typename Candidate> struct IsStellarModel : std::false_type { };

        template <
            typename Structure,
            typename SurfaceDeformation,
            typename StellarInteriorDeformation,
            typename VacuumDeformation>
        struct IsStellarModel<
            StellarModel<Structure, SurfaceDeformation, StellarInteriorDeformation, VacuumDeformation>>
            : std::true_type { };
    } // namespace detail

    template <typename Candidate>
    concept StellarModelType = detail::IsStellarModel<std::remove_cvref_t<Candidate>>::value;

    class StellarModelView final {
    public:
        template <typename Model>
            requires StellarModelType<Model> &&
                         eos::RuntimeEquationOfStateModel<typename std::remove_cvref_t<Model>::EquationOfStateType>
        explicit StellarModelView(Model &model) noexcept
            : m_equationOfState(model.equationOfState()),
              m_structurePrescription(std::addressof(model.structurePrescription())),
              m_makeInitialSeed(&makeInitialSeedFor<typename std::remove_cvref_t<Model>::StructurePrescriptionType>),
              m_targetMass(model.targetMass()),
              m_surfaceCondition(model.compiledSurfaceConstraint().descriptor()),
              m_surfaceDependencies(model.compiledSurfaceConstraint().runtimeDependencies()),
              m_surfaceDeformation(model.surfaceDeformationPrescription().descriptor()),
              m_stellarInteriorDeformation(model.stellarInteriorDeformationExtension().descriptor()),
              m_vacuumDeformation(model.vacuumDeformationExtension().descriptor()) {
        }

        [[nodiscard]] eos::EquationOfStateView equationOfState() const noexcept {
            return m_equationOfState;
        }

        [[nodiscard]] double targetMass() const noexcept {
            return m_targetMass;
        }

        [[nodiscard]] structure::StructureSeed makeInitialSeed(const structure::StructureSeedRequest &request) const {
            return m_makeInitialSeed(m_structurePrescription, request);
        }

        [[nodiscard]] surface::PressureSurfaceDescriptor surfaceCondition() const noexcept {
            return m_surfaceCondition;
        }

        [[nodiscard]] surface::RuntimeSurfaceConstraintDependencies surfaceDependencies() const noexcept {
            return m_surfaceDependencies;
        }

        [[nodiscard]] deformation::SurfaceDeformationDescriptor surfaceDeformation() const noexcept {
            return m_surfaceDeformation;
        }

        [[nodiscard]] deformation::InteriorDeformationExtensionDescriptor stellarInteriorDeformation() const noexcept {
            return m_stellarInteriorDeformation;
        }

        [[nodiscard]] deformation::VacuumDeformationExtensionDescriptor vacuumDeformation() const noexcept {
            return m_vacuumDeformation;
        }

    private:
        using MakeInitialSeedFunction = structure::StructureSeed (*)(
            const void *,
            const structure::StructureSeedRequest &
        );

        template <StructurePrescription Structure>
        [[nodiscard]] static structure::StructureSeed makeInitialSeedFor(
            const void *structurePrescription,
            const structure::StructureSeedRequest &request
        ) {
            return static_cast<const Structure *>(structurePrescription)->makeInitialSeed(request);
        }

        eos::EquationOfStateView m_equationOfState;
        const void *m_structurePrescription;
        MakeInitialSeedFunction m_makeInitialSeed;
        double m_targetMass;
        surface::PressureSurfaceDescriptor m_surfaceCondition;
        surface::RuntimeSurfaceConstraintDependencies m_surfaceDependencies;
        deformation::SurfaceDeformationDescriptor m_surfaceDeformation;
        deformation::InteriorDeformationExtensionDescriptor m_stellarInteriorDeformation;
        deformation::VacuumDeformationExtensionDescriptor m_vacuumDeformation;
    };
} // namespace mean_field::models