module;

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

export module mean_field:model.stellar;

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 SurfacePrescription =
        surface::ConstantPressureSurfaceType<Candidate> &&
        surface::PressureSurfaceCompilable<surface::BarotropicSurfaceFormulation, std::remove_cvref_t<EquationOfState>>;

    template <StructurePrescription Structure>
        requires SurfacePrescription<surface::ConstantPressureSurface, StructureEquationOfStateT<Structure>>
    class StellarModel final {
    public:
        using StructurePrescriptionType = Structure;
        using SurfacePrescriptionType   = surface::ConstantPressureSurface;
        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 surfacePrescription
        )
            : m_structurePrescription(
                  std::make_unique<Structure>(std::forward<StructureArgument>(structurePrescription))
              ),
              m_surfacePrescription(std::make_unique<surface::ConstantPressureSurface>(surfacePrescription)),
              m_compiledSurfaceConstraint(
                  std::make_unique<SurfaceConstraintType>(validateAndCompileSurface(
                      *m_structurePrescription,
                      *m_surfacePrescription
                  ))
              ) {
        }

        ~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 &surfacePrescription() const noexcept {
            return *m_surfacePrescription;
        }

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

    private:
        [[nodiscard]] static SurfaceConstraintType validateAndCompileSurface(
            const Structure &structurePrescription,
            const surface::ConstantPressureSurface &surfacePrescription
        ) {
            structurePrescription.validate();

            return surface::compilePressureSurfaceConstraint<surface::BarotropicSurfaceFormulation>(
                surfacePrescription, structurePrescription.equationOfState()
            );
        }

        std::unique_ptr<Structure> m_structurePrescription;
        std::unique_ptr<surface::ConstantPressureSurface> m_surfacePrescription;
        std::unique_ptr<SurfaceConstraintType> m_compiledSurfaceConstraint;
    };

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

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

        template <typename Structure> struct IsStellarModel<StellarModel<Structure>> : 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()) {
        }

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

    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;
    };
} // namespace mean_field::models