module;
#include <mfem.hpp>

export module mean_field:model.structure_profile;
export import :fem;

export namespace mean_field::models {
    constexpr double DEFAULT_COLATITUDE_DEGREES = 90;
    constexpr double DEFAULT_LONGITUDE_DEGREES  = 0;

    struct RadialDirection {
        double coLatitudeDegrees{DEFAULT_COLATITUDE_DEGREES};
        double longitudeDegrees{DEFAULT_LONGITUDE_DEGREES};

        [[nodiscard]] mfem::Vector toUnitCartesian() const;
    };

    struct RadialProfile {
        double coLatitudeDegrees;
        double longitudeDegrees;

        mfem::Vector radius;
        mfem::Vector param;
    };

    struct SliceProfile {
        double radius;
        std::array<double, 2> normal;

        mfem::GridFunction param;
    };

    class StructureProfile {
    public:
        explicit StructureProfile(
            const fem::FEM &fem,
            mfem::Vector state
        );

        // Profiles
        mfem::GridFunction pressureProfile();
        mfem::GridFunction densityProfile();
        mfem::GridFunction gravitationalPotentialProfile();
        mfem::GridFunction gravitationalFieldProfile();
        mfem::GridFunction enthalpyProfile();
        mfem::GridFunction entropyProfile();
        mfem::GridFunction temperatureProfile();
        mfem::GridFunction internalEnergyProfile();

        // Local Evaluation
        double pressureAt(const mfem::Vector &position);
        double densityAt(const mfem::Vector &position);
        double gravitationalPotentialAt(const mfem::Vector &position);
        mfem::Vector gravitationalFieldAt(const mfem::Vector &position);
        double enthalpyAt(const mfem::Vector &position);
        double entropyAt(const mfem::Vector &position);
        double temperatureAt(const mfem::Vector &position);
        double internalEnergyAt(const mfem::Vector &position);

        // Radial helpers
        mfem::Vector radius(RadialDirection direction = {});

        RadialProfile radialPressureProfile(RadialDirection direction = {});
        RadialProfile radialDensityProfile(RadialDirection direction = {});
        RadialProfile radialGravitationalPotentialProfile(RadialDirection direction = {});
        RadialProfile radialGravitationalFieldProfile(RadialDirection direction = {});
        RadialProfile radialEnthalpyProfile(RadialDirection direction = {});
        RadialProfile radialEntropyProfile(RadialDirection direction = {});
        RadialProfile radialTemperatureProfile(RadialDirection direction = {});
        RadialProfile radialInternalEnergyProfile(RadialDirection direction = {});

        // Ellipsoidal Slices
        SliceProfile slicePressureProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceDensityProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceGravitationalPotentialProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceGravitationalFieldProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceEnthalpyProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceEntropyProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceTemperatureProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );
        SliceProfile sliceInternalEnergyProfile(
            double radius,
            const std::array<
                double,
                2> &normal
        );

        // Integral Constraints
        double totalMass();
        double virialRatio();

        // Diagnostics
        bool isBound();

        // IO
        void radialToCSV(
            const std::string &filename,
            RadialDirection direction = {}
        );
        void radialToBIN(
            const std::string &filename,
            RadialDirection direction = {}
        );
        void toBIN(const std::string &filename);

    private:
        const fem::FEM &m_fem;
        const mfem::Vector m_state;
    };

    StructureProfile StructureProfileFromBIN(const std::string &filename);
} // namespace mean_field::models