//region Includes
#include <memory>
#include <mfem.hpp>
#include <print>
#include <format>
#include <string>
#include <functional>
#include <utility>
#include <vector>

#include <cmath>
#include <expected>

#include <CLI/CLI.hpp>
#include <XAD/XAD.hpp>
//endregion

//region Test Utilities
constexpr std::string_view ANSI_GREEN = "\033[32m";
constexpr std::string_view ANSI_RED = "\033[31m";
constexpr std::string_view ANSI_YELLOW = "\033[33m";
constexpr std::string_view ANSI_BLUE = "\033[34m";
constexpr std::string_view ANSI_RESET = "\033[0m";

enum class TEST_RESULT_TYPE : uint8_t {
    SUCCESS,
    FAILURE,
    PARTIAL
};

std::string fmt_test_msg(const std::string_view test_name, const TEST_RESULT_TYPE type, size_t num_fails, size_t total) {
    std::string_view color;
    switch (type) {
        case TEST_RESULT_TYPE::SUCCESS:
            color = ANSI_GREEN;
            break;
        case TEST_RESULT_TYPE::FAILURE:
            color = ANSI_RED;
            break;
        case TEST_RESULT_TYPE::PARTIAL:
            color = ANSI_YELLOW;
            break;
        default:
            color = ANSI_RESET;
    }
    return std::format("{}[TEST: {}] {}/{}{}", color, test_name, total-num_fails, total, ANSI_RESET);
}
//endregion

//region Constants
/********************
 * Constants
 *********************/
constexpr double G = 1.0;
constexpr double MASS = 1.0;
constexpr double RADIUS = 1.0;
constexpr double CENTRAL_DENSITY = 1.0;

constexpr char HOST[10] = "localhost";
constexpr int PORT = 19916;

//endregion

//region Concepts and Typedefs
/********************
 * Concepts
 *********************/
template <typename T>
concept is_xad =
       std::is_same_v<T, xad::AReal<long double>>
    || std::is_same_v<T, xad::AReal<double>>
    || std::is_same_v<T, xad::AReal<float>>;

template <typename T>
concept is_real = std::is_floating_point_v<T> || is_xad<T>;

/********************
 * Type Defs
 *********************/
template <is_real T>
using EOS_P = std::function<T(T rho, T temp)>;
//endregion

//region User Argument Structs
/********************
 * User Args
 *********************/
struct potential {
    double tol;
    int max_iters;
};

struct rot {
    bool enabled;
    double omega;
};

struct Args {
    std::string mesh_file;
    potential p{};
    rot r{};
    bool verbose{};
    double index{};
    double mass{};
    double c{};

    int max_iters{};
    double tol{};
};
//endregion

//region Misc Structs
struct OblatePotential {
    bool use{false};
    double a{1};
    double c{1};
    double rho_0{1};
};

struct Bounds {
    double r_star_ref;
    double r_inf_ref;
};

enum BoundsError : uint8_t {
    CANNOT_FIND_VACUUM
};
//endregion

//region Domain Enums
enum class Domains : uint8_t {
    NONE = 0,
    CORE = 1 << 0,
    ENVELOPE = 1 << 1,
    VACUUM = 1 << 2,
    ALL = 0x7
};

inline Domains operator|(Domains lhs, Domains rhs) {
    return static_cast<Domains>(static_cast<uint8_t>(lhs) | static_cast<uint8_t>(rhs));
}

inline Domains operator&(Domains lhs, Domains rhs) {
    return static_cast<Domains>(static_cast<uint8_t>(lhs) & static_cast<uint8_t>(rhs));
}

const Domains STELLAR = Domains::CORE | Domains::ENVELOPE;
//endregion

//region Domain Mapper
/********************
 * Mappers
 *********************/
class DomainMapper {
public:
    DomainMapper(
        const double r_star_ref,
        const double r_inf_ref
    ) :
    m_d(nullptr),
    m_r_star_ref(r_star_ref),
    m_r_inf_ref(r_inf_ref) {}

    explicit DomainMapper(
        const mfem::GridFunction &d,
        const double r_star_ref,
        const double r_inf_ref
    ) :
    m_d(&d),
    m_dim(d.FESpace()->GetMesh()->Dimension()),
    m_r_star_ref(r_star_ref),
    m_r_inf_ref(r_inf_ref) {};

    [[nodiscard]] bool is_vacuum(const mfem::ElementTransformation &T) const {
        if (T.ElementType == mfem::ElementTransformation::ELEMENT) {
            return T.Attribute == m_vacuum_attr;
        } else if (T.ElementType == mfem::ElementTransformation::BDR_ELEMENT) {
            return T.Attribute == m_vacuum_attr - 1; // TODO: In a more robust code this should really be read from the stroid API to ensure that the vacuum boundary is really 1 - the vacuum material attribute
        }
        return false;
    }

    void SetDisplacement(const mfem::GridFunction &d) {
        if (m_dim != d.FESpace()->GetMesh()->Dimension()) {
            const std::string err_msg = std::format("Dimension mismatch: DomainMapper is initialized for dimension {}, but provided displacement field has dimension {}.", m_dim, d.FESpace()->GetMesh()->Dimension());
            throw std::invalid_argument(err_msg);
        }
        m_d = &d;
    }

    [[nodiscard]] bool IsIdentity() const {
        return (m_d == nullptr);
    }

    void ResetDisplacement() {
        m_d = nullptr;
    }

    void ComputeJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &J) const {
        mfem::DenseMatrix J_D(m_dim, m_dim);
        J.SetSize(m_dim, m_dim);
        J = 0.0;
        J_D = 0.0;
        if (IsIdentity()) {
            for (int i = 0; i < m_dim; ++i) {
                J_D(i, i) = 1.0; // Identity mapping
            }
        } else {
            m_d->GetVectorGradient(T, J_D);
            for (int i = 0; i < m_dim; ++i) {
                J_D(i, i) += 1.0; // Add identity to get the total Jacobian of the mapping
            }
        }

        if (is_vacuum(T)) {
            mfem::Vector x_ref(m_dim);
            T.Transform(T.GetIntPoint(), x_ref);

            mfem::Vector d_val(m_dim), x_disp(m_dim);
            if (IsIdentity()) {
                x_disp = x_ref;
            } else {
                m_d->GetVectorValue(T, T.GetIntPoint(), d_val);
                add(x_ref, d_val, x_disp);
            }

            mfem::DenseMatrix J_K(m_dim, m_dim);
            ComputeKelvinJacobian(x_ref, x_disp, J_K);
            mfem::Mult(J_K, J_D, J);
        } else {
            J = J_D;
        }
    }

    double ComputeDetJ(mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) const {
        if (IsIdentity() && !is_vacuum(T)) return 1.0; // If no mapping, the determinant of the Jacobian is 1
        mfem::DenseMatrix J;
        ComputeJacobian(T, J);
        return J.Det();
    }

    void ComputeMappedDiffusionTensor(mfem::ElementTransformation &T, mfem::DenseMatrix &D) const {
        mfem::DenseMatrix J(m_dim, m_dim), JInv(m_dim, m_dim);

        ComputeJacobian(T, J);
        const double detJ = J.Det();
        mfem::CalcInverse(J, JInv);
        D.SetSize(m_dim, m_dim);
        mfem::MultABt(JInv, JInv, D);
        D *= fabs(detJ);
    }

    void ComputeInverseJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &JInv) const {
        mfem::DenseMatrix J(m_dim, m_dim);
        ComputeJacobian(T, J);
        JInv.SetSize(m_dim, m_dim);
        mfem::CalcInverse(J, JInv);
    }

    void GetPhysicalPoint(mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip, mfem::Vector& x_phys) const {
        x_phys.SetSize(m_dim);
        mfem::Vector x_ref(m_dim);
        T.Transform(ip, x_ref);

        if (IsIdentity()) {
            x_phys = x_ref;
        } else {
            mfem::Vector d_val(m_dim);
            m_d->GetVectorValue(T, ip, d_val);
            add(x_ref, d_val, x_phys);
        }
        if (is_vacuum(T)) {
            ApplyKelvinMapping(x_ref, x_phys);
        }
    }

    [[nodiscard]] const mfem::GridFunction* GetDisplacement() const { return m_d; }


private:
    void ApplyKelvinMapping(const mfem::Vector& x_ref, mfem::Vector& x_phys) const {
        const double r_ref = x_ref.Norml2();
        double xi = (r_ref - m_r_star_ref) / (m_r_inf_ref - m_r_star_ref);
        xi = std::clamp(xi, 0.0, m_xi_clamp);
        const double factor = m_r_star_ref / (r_ref * (1 - xi));
        x_phys *= factor;
    }

    void ComputeKelvinJacobian(const mfem::Vector& x_ref, const mfem::Vector &x_disp, mfem::DenseMatrix& J_K) const {
        const double r_ref = x_ref.Norml2();
        const double delta_R = m_r_inf_ref - m_r_star_ref;

        double xi = (r_ref - m_r_star_ref) / delta_R;
        xi = std::clamp(xi, 0.0, m_xi_clamp);

        const double denom = 1.0 - xi;

        const double k = m_r_star_ref / (r_ref * denom);

        const double dk_dr = m_r_star_ref * (( 1.0 / (delta_R* r_ref * denom * denom)) - ( 1.0 / (r_ref * r_ref * denom)));

        J_K.SetSize(m_dim, m_dim);
        const double outer_factor = dk_dr / r_ref;

        for (int i = 0; i < m_dim; ++i) {
            for (int j = 0; j < m_dim; ++j) {
                J_K(i, j) = outer_factor * x_disp(i) * x_ref(j);
                if (i == j) {
                    J_K(i, j) += k;
                }
            }
        }
    }
private:
    const mfem::GridFunction *m_d;
    std::unique_ptr<mfem::GridFunction> m_internal_d;
    const int m_dim{3};
    const int m_vacuum_attr{3};
    const double m_r_star_ref{1.0};
    const double m_r_inf_ref{2.0};
    const double m_xi_clamp{0.999};
};

//endregion

//region State Types
/********************
 * State Types
 *********************/
struct FEM {
    std::unique_ptr<mfem::Mesh> mesh;
    std::unique_ptr<mfem::FiniteElementCollection> H1_fec;
    std::unique_ptr<mfem::FiniteElementSpace> H1_fes;
    std::unique_ptr<mfem::FiniteElementSpace> Vec_H1_fes;
    std::unique_ptr<DomainMapper> mapping;

    mfem::Array<int> block_offsets;
    std::unique_ptr<mfem::GridFunction> reference_x;

    mfem::Vector com;
    mfem::DenseMatrix Q;
    mfem::Array<int> ess_tdof_x;

    std::unique_ptr<mfem::Array<int>> star_marker;
    std::unique_ptr<mfem::Array<int>> vacuum_marker;
    std::unique_ptr<mfem::Array<int>> surface_marker;
    std::unique_ptr<mfem::Array<int>> inf_marker;

    [[nodiscard]] bool okay() const { return (mesh != nullptr) && (H1_fec != nullptr) && (H1_fes != nullptr) && (Vec_H1_fes != nullptr); }

    [[nodiscard]] bool has_mapping() const { return mapping != nullptr; }
};

struct CoupledState {
    std::unique_ptr<mfem::BlockVector> U;
    mfem::GridFunction rho;
    mfem::GridFunction d; // Stability depends on solving for the displacement vector not the nodal position vector, those live on a reference domain.

    explicit CoupledState(const FEM& fem) {
        U = std::make_unique<mfem::BlockVector>(fem.block_offsets);
        rho.MakeRef(fem.H1_fes.get(), U->GetBlock(0), 0);
        d.MakeRef(fem.Vec_H1_fes.get(), U->GetBlock(1), 0);
        *U = 0.0;
        U->GetBlock(2) = 1.0;
    }
};

//endregion

//region Function Definitions
/********************
 * Core Setup Functions
 *********************/
FEM setup_fem(const std::string& filename, bool verbose=true);

/********************
 * Utility Functions
 *********************/
void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title);
double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf, Domains domain = Domains::ALL);
mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho);
void get_physical_coordinates(const mfem::GridFunction& reference_pos, const mfem::GridFunction& displacement, mfem::GridFunction& physical_pos);
void populate_element_mask(const FEM& fem, Domains domain, mfem::Array<int>& mask);
std::expected<Bounds, BoundsError> DiscoverBounds(const mfem::Mesh *mesh, int vacuum_attr);

/********************
 * Physics Functions
 *********************/
double centrifugal_potential(const mfem::Vector& phys_x, double omega);
double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho);
double oblate_spheroid_surface_potential(const mfem::Vector& x, double a, double c, double total_mass);
std::unique_ptr<mfem::GridFunction> grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho);
std::unique_ptr<mfem::GridFunction> get_potential(const FEM& fem, const Args& args, const mfem::GridFunction& rho);
mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com);
double l2_multipole_potential(const FEM& fem, double total_mass, const mfem::Vector& phys_x);

/********************
 * Conservation Functions
 *********************/
void conserve_mass(const FEM& fem, mfem::GridFunction& rho, double target_mass);
//endregion

//region Mapping Coefficients
class MappedScalarCoefficient : public mfem::Coefficient {
public:
    enum class EVAL_POINTS : uint8_t {
        PHYSICAL,
        REFERENCE
    };

    MappedScalarCoefficient(
        const DomainMapper& map,
        mfem::Coefficient& coeff,
        const EVAL_POINTS eval_point=EVAL_POINTS::PHYSICAL
    ) :
    m_map(map),
    m_coeff(coeff),
    m_eval_point(eval_point) {};

    double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override {
        T.SetIntPoint(&ip);

        double detJ = m_map.ComputeDetJ(T, ip);
        double f_val = 0.0;

        switch (m_eval_point) {
            case EVAL_POINTS::PHYSICAL: {
                f_val = eval_at_point(m_coeff, T, ip);
                break;
            }
            case EVAL_POINTS::REFERENCE: {
                f_val = m_coeff.Eval(T, ip);
                break;
            }
        }
        return f_val * fabs(detJ);
    }
private:
    static double eval_at_point(mfem::Coefficient& c, mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) {
        return c.Eval(T, ip);
    }

private:
    const DomainMapper& m_map;
    mfem::Coefficient& m_coeff;
    EVAL_POINTS m_eval_point;

};

class MappedDiffusionCoefficient : public mfem::MatrixCoefficient {
public:
    MappedDiffusionCoefficient(
        const DomainMapper& map,
        mfem::Coefficient& sigma,
        const int dim
    ) :
    mfem::MatrixCoefficient(dim),
    m_map(map),
    m_scalar(&sigma),
    m_tensor(nullptr) {};

    MappedDiffusionCoefficient(
        const DomainMapper& map,
        mfem::MatrixCoefficient& sigma
    ) :
    mfem::MatrixCoefficient(sigma.GetHeight()),
    m_map(map),
    m_scalar(nullptr),
    m_tensor(&sigma) {};

    void Eval(mfem::DenseMatrix &K, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override {
        const int dim = height;
        T.SetIntPoint(&ip);

        mfem::DenseMatrix J(dim, dim), JInv(dim, dim);
        m_map.ComputeJacobian(T, J);
        const double detJ = J.Det();
        mfem::CalcInverse(J, JInv);

        if (m_scalar) {
            const double sig_val = m_scalar->Eval(T, ip);
            mfem::MultABt(JInv, JInv, K);
            K *= sig_val * fabs(detJ);
        } else {
            mfem::DenseMatrix sig_mat(dim, dim);
            m_tensor->Eval(sig_mat, T, ip);

            mfem::DenseMatrix temp(dim, dim);
            Mult(JInv, sig_mat, temp);

            MultABt(temp, JInv, K);
            K *= fabs(detJ);
        }
    }
private:
    const DomainMapper& m_map;
    mfem::Coefficient* m_scalar;
    mfem::MatrixCoefficient* m_tensor;
};

class MappedVectorCoefficient : public mfem::VectorCoefficient {
public:
    MappedVectorCoefficient(
        const DomainMapper& map,
        mfem::VectorCoefficient& coeff
    ) :
    mfem::VectorCoefficient(coeff.GetVDim()),
    m_map(map),
    m_coeff(coeff) {};

    void Eval(mfem::Vector& V, mfem::ElementTransformation& T, const mfem::IntegrationPoint& ip) override {
        const int dim = vdim;
        T.SetIntPoint(&ip);

        mfem::DenseMatrix JInv(dim, dim);
        m_map.ComputeInverseJacobian(T, JInv);
        double detJ = m_map.ComputeDetJ(T, ip);

        mfem::Vector C_phys(dim);
        m_coeff.Eval(C_phys, T, ip);

        V.SetSize(dim);
        JInv.Mult(C_phys, V);
        V *= fabs(detJ);
    }
private:
    const DomainMapper& m_map;
    mfem::VectorCoefficient& m_coeff;
};

class PhysicalPositionFunctionCoefficient : public mfem::Coefficient {
public:
    using Func = std::function<double(const mfem::Vector& x)>;

    PhysicalPositionFunctionCoefficient(
        const DomainMapper& map,
        Func f
    ) :
    m_f(std::move(f)),
    m_map(map) {};

    double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override {
        T.SetIntPoint(&ip);
        mfem::Vector x;
        m_map.GetPhysicalPoint(T, ip, x);
        return m_f(x);
    }
private:
    Func m_f;
    const DomainMapper& m_map;
};
//endregion

//region Integrators
/********************
 * Integrators
 *********************/
template <is_xad EOS_T>
class FluidIntegrator : public mfem::NonlinearFormIntegrator {
    using Scalar = EOS_T::value_type;
public:
    explicit FluidIntegrator(
        EOS_P<EOS_T> eos,
        const DomainMapper* map = nullptr
    ) :
    m_eos(std::move(eos)),
    m_map(map)
    {};

    void AssembleElementVector(
        const mfem::FiniteElement &el,
        mfem::ElementTransformation &Tr,
        const mfem::Vector &elfun,
        mfem::Vector &elvect
    ) override {
        const int dof = el.GetDof();
        elvect.SetSize(dof);
        elvect = 0.0;

        const mfem::IntegrationRule *ir = IntRule ? IntRule : &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 1);

        mfem::Vector shape(dof);
        for (int i = 0; i < ir->GetNPoints(); i++) {
            const mfem::IntegrationPoint& ip = ir->IntPoint(i);
            Tr.SetIntPoint(&ip);
            el.CalcShape(ip, shape);

            double u = shape * elfun;
            EOS_T rho = u;
            const double val = m_eos(rho, 0.0).value();
            double weight = ip.weight * Tr.Weight() * val;

            if (m_map) weight *= fabs(m_map->ComputeDetJ(Tr, ip));


            elvect.Add(weight, shape);
        }
    }
    void AssembleElementGrad(
        const mfem::FiniteElement &el,
        mfem::ElementTransformation &Tr,
        const mfem::Vector &elfun,
        mfem::DenseMatrix &elmat
    ) override {
        const int dof = el.GetDof();
        elmat.SetSize(dof);
        elmat = 0.0;

        const mfem::IntegrationRule *ir = IntRule ? IntRule : &mfem::IntRules.Get(el.GetGeomType(), 2 * el.GetOrder() + 1);

        mfem::Vector shape(dof);
        for (int i = 0; i < ir->GetNPoints(); i++) {
            const mfem::IntegrationPoint& ip = ir->IntPoint(i);
            Tr.SetIntPoint(&ip);
            el.CalcShape(ip, shape);

            double u = shape * elfun;

            double d_val_d_rho = 0.0;
            if (u > 1e-15) {
                xad::Tape<Scalar> tape;
                EOS_T x_r = u;
                EOS_T x_t = 0.0; // In future this is one area where we introduce a temp dependency

                tape.registerInput(x_r);
                EOS_T result = m_eos(x_r, x_t);
                tape.registerOutput(result);
                result.setAdjoint(1.0);
                tape.computeAdjoints();
                d_val_d_rho = x_r.getAdjoint();
            }

            double weight = ip.weight * Tr.Weight() * d_val_d_rho;
            if (m_map) weight *= fabs(m_map->ComputeDetJ(Tr, ip));

            mfem::AddMult_a_VVt(weight, shape, elmat);
        }
    }

    [[nodiscard]] bool has_mapping() const { return m_map != nullptr; }

    void set_mapping(const DomainMapper* map) { m_map = map; }

    void clear_mapping() { m_map = nullptr; }
private:
    EOS_P<EOS_T> m_eos;
    const DomainMapper* m_map{nullptr};
};

//endregion

//region Coefficients
/********************
 * Coefficient Defs
 *********************/
template <is_xad EOS_T>
class PressureBoundaryForce : public mfem::VectorCoefficient {
public:
    PressureBoundaryForce(
        const int dim,
        const FEM& fem,
        const mfem::GridFunction& rho,
        const EOS_P<EOS_T>& eos,
        const double P_fit
    ) : VectorCoefficient(dim), m_fem(fem), m_rho(rho), m_eos(eos), m_P_fit(P_fit) {};

    void Eval(mfem::Vector &V, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override {
        V.SetSize(vdim);
        V = 0.0;

        double rho = m_rho.GetValue(T, ip);

        EOS_T x_rho = rho;
        const double P_curr = m_eos(x_rho, 0.0).value();
        const double delta_P = P_curr - m_P_fit;

        mfem::Vector phys(vdim);
        T.Transform(ip, phys);
        mfem::Vector normal(vdim);
        mfem::CalcOrtho(T.Jacobian(), normal);

        for (int i = 0; i < vdim; ++i) {
            V(i) = delta_P * normal(i);
        }
    }
private:
    const FEM& m_fem;
    const mfem::GridFunction& m_rho;
    const EOS_P<EOS_T>& m_eos;
    double m_P_fit;
};

template <is_xad EOS_T>
class MappingDetJacobianCoefficient : public mfem::Coefficient {
public:
    MappingDetJacobianCoefficient(
        const mfem::GridFunction& d,
        const mfem::FiniteElementSpace& fes
    ) :
    m_d(d),
    m_fes(fes) {}

    double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override {
        const int dim = T.GetSpaceDim();
        mfem::DenseMatrix grad_d(dim, dim);
        m_d.GetVectorGradient(T, grad_d);
        for (int i = 0; i < dim; ++i) {
            grad_d(i, i) += 1.0;
        }
        return grad_d.Det();
    }
private:
    const mfem::GridFunction& m_d; // Displacement vector
    const mfem::FiniteElementSpace &m_fes;
};
//endregion

//region Operators
/********************
 * Operator Defs
 *********************/
template <is_xad EOS_T>
class VectorOperator : public mfem::Operator {
public:
    VectorOperator() = default;

    VectorOperator(
        const mfem::Vector& v,
        const bool is_col
    )  :
    Operator(is_col ? v.Size() : 1, is_col ? 1 : v.Size()),
    m_v(v),
    m_is_col(is_col),
    m_is_initialized(true) {}

    void SetVector(const mfem::Vector& v, const bool is_col) {
        if (v.Size() != m_v.Size()) {
            m_v.SetSize(v.Size());
        }
        m_v = v;
        m_is_col = is_col;
        height = is_col ? v.Size() : 1;
        width = is_col ? 1 : v.Size();
        m_is_initialized = true;
    }

    void Mult(const mfem::Vector &x, mfem::Vector &y) const override {
        if (!m_is_initialized) throw std::runtime_error("VectorOperator Not initialized");
        if (m_is_col) {
            y.SetSize(m_v.Size());
            y = 0.0;
            y.Add(x(0), m_v);
        } else {
            y.SetSize(1);
            y(0) = m_v * x;
        }
    }
private:
    mfem::Vector m_v;
    bool m_is_col{false};
    bool m_is_initialized{false};
};

template <is_xad EOS_T>
class PressureDensityCoupling : public mfem::Operator {
public:
    PressureDensityCoupling(
        FEM& fem,
        const mfem::GridFunction& rho,
        const EOS_P<EOS_T>& eos_pressure)
    : Operator(fem.Vec_H1_fes->GetTrueVSize(), fem.H1_fes->GetTrueVSize()),
    m_fem(fem),
    m_rho(rho),
    m_eos(eos_pressure) {
        Assemble();
    }

    void Assemble() {
        const int dim = m_fem.mesh->Dimension();
        const int scalar_size = m_fem.H1_fes->GetTrueVSize();
        const int vector_size  = m_fem.Vec_H1_fes->GetTrueVSize();

        m_mat = std::make_unique<mfem::SparseMatrix>(vector_size, scalar_size);

        for (int be = 0; be < m_fem.mesh->GetNBE(); ++be) {
            auto* ftr = m_fem.mesh->GetBdrFaceTransformations(be);
            if (!ftr) continue;

            const int elem = ftr->Elem1No;

            const auto& scalar_fe = *m_fem.H1_fes->GetFE(elem);

            const int sdof = scalar_fe.GetDof();

            mfem::Array<int> scalar_dofs, vector_dofs;
            m_fem.H1_fes->GetElementDofs(elem, scalar_dofs);
            m_fem.Vec_H1_fes->GetElementDofs(elem, vector_dofs);

            mfem::DenseMatrix elmat(vector_dofs.Size(), scalar_dofs.Size());
            elmat = 0.0;

            const auto& face_ir = mfem::IntRules.Get(ftr->GetGeometryType(), 2 * scalar_fe.GetOrder() + 2);

            mfem::Vector shape(sdof);
            mfem::Vector normal(dim);

            for (int q = 0; q < face_ir.GetNPoints(); ++q) {
                const auto& face_ip = face_ir.IntPoint(q);

                ftr->SetAllIntPoints(&face_ip);

                const mfem::IntegrationPoint& vol_ip = ftr->GetElement1IntPoint();

                scalar_fe.CalcShape(vol_ip, shape);

                mfem::CalcOrtho(ftr->Face->Jacobian(), normal);

                mfem::ElementTransformation& vol_trans = ftr->GetElement1Transformation();
                vol_trans.SetIntPoint(&vol_ip);
                double rho_val = m_rho.GetValue(vol_trans, vol_ip);

                double dPdrho = 0.0;
                if (rho_val > 1e-15) {
                    using Scalar = EOS_T::value_type;
                    xad::Tape<Scalar> tape;

                    EOS_T x_rho = rho_val;
                    tape.registerInput(x_rho);
                    EOS_T P = m_eos(x_rho, EOS_T(0.0));
                    tape.registerOutput(P);

                    P.setAdjoint(1.0);

                    tape.computeAdjoints();
                    dPdrho = x_rho.getAdjoint();
                }

                const double w = face_ip.weight;

                for (int k = 0; k < sdof; ++k) {
                    for (int j = 0; j < sdof; ++j) {
                        double base = -dPdrho * shape(j) * shape(k) * w;
                        for (int c = 0; c < dim; ++c) {
                            elmat(k + c * sdof, j) += base * normal(c);
                        }
                    }
                }
            }
            m_mat->AddSubMatrix(vector_dofs, scalar_dofs, elmat);
        }
        m_mat->Finalize();
    }

    void Mult(const mfem::Vector &x, mfem::Vector &y) const override {
        m_mat->Mult(x, y);
    }

    [[nodiscard]] mfem::SparseMatrix& SpMat() const { return *m_mat; }

private:
    FEM& m_fem;
    const mfem::GridFunction& m_rho;
    const EOS_P<EOS_T>& m_eos;
    std::unique_ptr<mfem::SparseMatrix> m_mat;
};

template <is_xad EOS_T>
class MassDisplacementCoupling : public mfem::Operator {
public:
    MassDisplacementCoupling(
        FEM& fem,
        const mfem::GridFunction& rho,
        const bool is_col
    ) :
    Operator(
        is_col ? fem.Vec_H1_fes->GetTrueVSize() : 1,
        is_col ? 1 : fem.Vec_H1_fes->GetTrueVSize()
    ),
    m_fem(fem),
    m_rho(rho),
    m_is_col(is_col){
        m_D.SetSize(m_fem.Vec_H1_fes->GetTrueVSize());
        Assemble();
    }

    void Assemble() const {
        const int dim = m_fem.mesh->Dimension();
        m_D = 0.0;
        for (int elemID = 0; elemID < m_fem.mesh->GetNE(); ++elemID) {
            auto* trans = m_fem.mesh->GetElementTransformation(elemID);
            const auto& fe = *m_fem.Vec_H1_fes->GetFE(elemID);
            const int dof = fe.GetDof();

            mfem::Array<int> vdofs;
            m_fem.Vec_H1_fes->GetElementDofs(elemID, vdofs);

            const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fe.GetOrder() + 1);

            mfem::DenseMatrix dshape(dof, dim);
            mfem::Vector elvec(dof * dim);
            elvec = 0.0;

            for (int q = 0; q < ir.GetNPoints(); ++q) {
                const auto& ip = ir.IntPoint(q);
                trans->SetIntPoint(&ip);

                fe.CalcPhysDShape(*trans, dshape);
                double rho_val = m_rho.GetValue(elemID, ip);
                double ref_weight = trans->Weight() * ip.weight;

                mfem::DenseMatrix J_map(dim, dim), J_inv(dim, dim);
                m_fem.mapping->ComputeJacobian(*trans, J_map);
                double detJ = J_map.Det();
                mfem::CalcInverse(J_map, J_inv);

                for (int k = 0; k < dof; ++k) {
                    for (int c = 0; c < dim; ++c) {
                        double trace_contrib = 0.0;
                        for (int j = 0; j < dim; ++j) {
                            trace_contrib += J_inv(j, c) * dshape(k, j);
                        }
                        elvec(k + c * dof) += rho_val * fabs(detJ) * trace_contrib * ref_weight;
                    }
                }
            }
            m_D.AddElementVector(vdofs, elvec);
        }
    }

    [[nodiscard]] mfem::Vector& GetVec() const {
        return m_D;
    }

    void Mult(const mfem::Vector &x, mfem::Vector &y) const override {
        if (m_is_col) {
            y.SetSize(m_D.Size());
            y = 0.0;
            y.Add(x(0), m_D);
        } else {
            y.SetSize(1);
            y(0) = m_D * x;
        }
    }


private:
    const FEM& m_fem;
    const mfem::GridFunction& m_rho;
    const bool m_is_col;

    mutable mfem::Vector m_D;
};

template <is_xad EOS_T>
class ResidualOperator : public mfem::Operator {
public:
    ResidualOperator(
        FEM& fem,
        const Args& args,
        const EOS_P<EOS_T>& eos_enthalpy,
        const EOS_P<EOS_T>& eos_pressure,
        const double alpha
    ) :
    Operator(fem.block_offsets.Last()),
    m_fem(fem),
    m_args(args),
    m_eos_enthalpy(eos_enthalpy),
    m_eos_pressure(eos_pressure),
    m_alpha(std::make_unique<mfem::ConstantCoefficient>(alpha)),
    m_fluid_nlf(m_fem.H1_fes.get()),
    m_reference_stiffness(m_fem.Vec_H1_fes.get())
    {
        m_fluid_nlf.AddDomainIntegrator(new FluidIntegrator<EOS_T>(m_eos_enthalpy, m_fem.mapping.get()));

        m_reference_stiffness.AddDomainIntegrator(new mfem::VectorMassIntegrator(*m_alpha));
        m_reference_stiffness.AddDomainIntegrator(new mfem::VectorDiffusionIntegrator());
        m_reference_stiffness.Assemble();
        m_reference_stiffness.Finalize();
        std::println("ResidualOperator initialized with XAD-enabled FluidIntegrator and reference stiffness.");
    };


    void Mult(const mfem::Vector &u, mfem::Vector &r) const override {
        mfem::GridFunction rho, d;
        rho.MakeRef(m_fem.H1_fes.get(), u.GetData() + m_fem.block_offsets[0]);
        d.MakeRef(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]);
        double lambda = u(m_fem.block_offsets[2]);

        m_fem.mapping->SetDisplacement(d);

        m_fem.com = get_com(m_fem, rho);
        m_fem.Q = compute_quadrupole_moment_tensor(m_fem, rho, m_fem.com);

        mfem::GridFunction r_rho, r_d;
        r_rho.MakeRef(m_fem.H1_fes.get(), r.GetData() + m_fem.block_offsets[0]);
        r_d.MakeRef(m_fem.Vec_H1_fes.get(), r.GetData() + m_fem.block_offsets[1]);

        double &r_lambda = r(m_fem.block_offsets[2]);

        m_fluid_nlf.Mult(rho, r_rho);

        auto phi = get_potential(m_fem, m_args, rho);
        mfem::GridFunctionCoefficient phi_c(phi.get());
        MappedScalarCoefficient mapped_phi_c(*m_fem.mapping, phi_c);
        mfem::LinearForm phi_lf(m_fem.H1_fes.get());
        phi_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_phi_c));
        phi_lf.Assemble();
        r_rho += phi_lf;

        mfem::ConstantCoefficient lambda_c(lambda);
        MappedScalarCoefficient mapped_lambda_c(*m_fem.mapping, lambda_c);
        mfem::LinearForm mass_grad_lf(m_fem.H1_fes.get());
        mass_grad_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_lambda_c));
        mass_grad_lf.Assemble();

        // ReSharper disable once CppDFAUnusedValue
        r_rho += mass_grad_lf;

        m_reference_stiffness.Mult(d, r_d);

        PressureBoundaryForce<EOS_T> pbf(
            m_fem.H1_fes->GetMesh()->Dimension(),
            m_fem,
            rho,
            m_eos_pressure,
            m_args.c
        );
        mfem::LinearForm pbf_lf(m_fem.Vec_H1_fes.get());
        pbf_lf.AddBoundaryIntegrator(new mfem::VectorBoundaryLFIntegrator(pbf));
        pbf_lf.Assemble();

        r_d -= pbf_lf;

        for (int i = 0; i < m_fem.ess_tdof_x.Size(); ++i) {
            r_d(m_fem.ess_tdof_x[i]) = 0.0;
        }

        double current_mass = domain_integrate_grid_function(m_fem, rho);
        r_lambda = current_mass - m_args.mass;
    }

    [[nodiscard]] Operator& GetGradient(const mfem::Vector &u) const override {
        mfem::GridFunction rho, d;
        rho.MakeRef(m_fem.H1_fes.get(), u.GetData() + m_fem.block_offsets[0]);
        d.MakeRef(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]);

        m_fem.mapping->SetDisplacement(d);

        m_approx_jacobian = std::make_unique<mfem::BlockOperator>(m_fem.block_offsets);

        const mfem::GridFunction grad(m_fem.Vec_H1_fes.get(), u.GetData() + m_fem.block_offsets[1]);
        // m_fem.mesh->SetNodes(grad);

        m_approx_jacobian->SetBlock(0, 0, &m_fluid_nlf.GetGradient(rho));
        m_approx_jacobian->SetBlock(1, 1, &m_reference_stiffness);

        B_vec.SetSize(m_fem.H1_fes->GetTrueVSize());
        B_vec = 0.0;

        mfem::ConstantCoefficient one(1.0);
        MappedScalarCoefficient mapped_b_vec(*m_fem.mapping, one);
        mfem::LinearForm b_lf(m_fem.H1_fes.get());
        b_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_b_vec));
        b_lf.Assemble();
        B_vec += b_lf;

        m_B_vec_op_col.SetVector(B_vec, true);
        m_B_vec_op_row.SetVector(B_vec, false);

        m_approx_jacobian->SetBlock(0, 2, &m_B_vec_op_col);
        m_approx_jacobian->SetBlock(2, 0, &m_B_vec_op_row);

        m_C = std::make_unique<PressureDensityCoupling<EOS_T>>(m_fem, rho, m_eos_pressure);
        m_D = std::make_unique<MassDisplacementCoupling<EOS_T>>(m_fem, rho, false);

        m_approx_jacobian->SetBlock(1, 0, m_C.get());
        m_approx_jacobian->SetBlock(2, 1, m_D.get());

        return *m_approx_jacobian;

    }
private:
    FEM& m_fem;
    const Args& m_args;
    const EOS_P<EOS_T>& m_eos_enthalpy;
    const EOS_P<EOS_T>& m_eos_pressure;
    const std::unique_ptr<mfem::ConstantCoefficient> m_alpha;

    mutable mfem::NonlinearForm m_fluid_nlf;
    mutable mfem::BilinearForm m_reference_stiffness;
    mutable std::unique_ptr<mfem::BlockOperator> m_approx_jacobian = nullptr;
    mutable mfem::Vector B_vec;
    mutable VectorOperator<EOS_T> m_B_vec_op_col;
    mutable VectorOperator<EOS_T> m_B_vec_op_row;
    mutable std::unique_ptr<PressureDensityCoupling<EOS_T>> m_C;
    mutable std::unique_ptr<MassDisplacementCoupling<EOS_T>> m_D;
};
//endregion

//region Utility Functions
FEM setup_fem(const std::string& filename, const bool verbose) {
    FEM fem;
    fem.mesh = std::make_unique<mfem::Mesh>(filename, 0, 0);
    fem.mesh->EnsureNodes();
    int dim = fem.mesh->Dimension();
    fem.H1_fec   = std::make_unique<mfem::H1_FECollection>(2, dim);
    fem.H1_fes   = std::make_unique<mfem::FiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get());
    fem.Vec_H1_fes = std::make_unique<mfem::FiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get(), dim, mfem::Ordering::byNODES);

    auto [r_star_ref, r_inf_ref] = DiscoverBounds(fem.mesh.get(), 3)
    .or_else([](const BoundsError& err)->std::expected<Bounds, BoundsError> {
        throw std::runtime_error("Unable to determine vacuum domain reference boundary...");
    }).value();

    fem.mapping = std::make_unique<DomainMapper>(r_star_ref, r_inf_ref);

    fem.block_offsets.SetSize(4);
    fem.block_offsets[0] = 0;
    fem.block_offsets[1] = fem.H1_fes->GetTrueVSize();
    fem.block_offsets[2] = fem.block_offsets[1] + fem.Vec_H1_fes->GetTrueVSize();
    fem.block_offsets[3] = fem.block_offsets[2] + 1;

    // Setup Domain Markers to constrain integration
    fem.star_marker = std::make_unique<mfem::Array<int>>(fem.mesh->attributes.Max());
    fem.star_marker->operator=(0);
    fem.star_marker.operator->()[0] = 1; // core
    fem.star_marker.operator->()[1] = 1; // envelope

    fem.vacuum_marker = std::make_unique<mfem::Array<int>>(fem.mesh->attributes.Max());
    fem.vacuum_marker->operator=(0);
    fem.vacuum_marker.operator->()[2] = 1; // vacuum

    fem.surface_marker = std::make_unique<mfem::Array<int>>(fem.mesh->bdr_attributes.Max());
    fem.surface_marker->operator=(0);
    fem.surface_marker.operator->()[0] = 1; // surface

    fem.vacuum_marker = std::make_unique<mfem::Array<int>>(fem.mesh->bdr_attributes.Max());
    fem.vacuum_marker->operator=(0);
    fem.vacuum_marker.operator->()[1] = 1; // Infinity

    // Initial COM should be 0, initial mass distribution should be uniform
    fem.com.SetSize(dim); fem.com = 0.0;
    fem.Q.SetSize(dim, dim); fem.Q = 0.0;

    fem.reference_x = std::make_unique<mfem::GridFunction>(fem.Vec_H1_fes.get());
    fem.mesh->GetNodes(*fem.reference_x);

    fem.ess_tdof_x.SetSize(0); // No essential boundary conditions for the displacement, the null space here is handled with a weak penalty term
    return fem;
}

void view_mesh(const std::string& host, int port, const mfem::Mesh& mesh, const mfem::GridFunction& gf, const std::string& title) {
    mfem::socketstream sol_sock(host.c_str(), port);
    if (!sol_sock.is_open()) return;
    sol_sock << "solution\n" << mesh << gf;
    sol_sock << "window_title '" << title << "'\n" << std::flush;
}

double domain_integrate_grid_function(const FEM& fem, const mfem::GridFunction& gf, Domains domain) {
    mfem::LinearForm lf(fem.H1_fes.get());
    mfem::GridFunctionCoefficient gf_c(&gf);
    double integral;
    mfem::Array<int> elem_markers;
    populate_element_mask(fem, domain, elem_markers);

    if (fem.has_mapping()) {
        MappedScalarCoefficient mapped_gf_c(*fem.mapping, gf_c);
        lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_gf_c), elem_markers);
        lf.Assemble();
        integral = lf.Sum();
    } else {
        lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(gf_c), elem_markers);
        lf.Assemble();
        integral = lf.Sum();
    }
    return integral;
}

mfem::Vector get_com(const FEM& fem, const mfem::GridFunction &rho) {
    const int dim = fem.mesh->Dimension();
    mfem::Vector com(dim);
    com = 0.0;
    double total_mass = 0.0;

    for (int i = 0; i < fem.H1_fes->GetNE(); ++i) {
        mfem::ElementTransformation *trans = fem.H1_fes->GetElementTransformation(i);
        const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), fem.H1_fes->GetOrder(0) + trans->OrderW());

        for (int j = 0; j < ir.GetNPoints(); ++j) {
            const mfem::IntegrationPoint &ip = ir.IntPoint(j);
            trans->SetIntPoint(&ip);

            double weight = trans->Weight() * ip.weight;
            if (fem.has_mapping()) {
                weight *= fem.mapping->ComputeDetJ(*trans, ip);
            }
            double rho_val = rho.GetValue(i, ip);

            mfem::Vector phys_point(dim);
            if (fem.has_mapping()) {
                fem.mapping->GetPhysicalPoint(*trans, ip, phys_point);
            } else {
                trans->Transform(ip, phys_point);
            }

            const double mass_term = rho_val * weight;
            total_mass += mass_term;

            for (int d = 0; d < dim; ++d) {
                com(d) += phys_point(d) * mass_term;
            }
        }
    }
    com /= total_mass;
    return com;
}

void get_physical_coordinates(const mfem::GridFunction& reference_pos, const mfem::GridFunction& displacement, mfem::GridFunction& physical_pos) {
    add(reference_pos, displacement, physical_pos);
}

void populate_element_mask(const FEM &fem, Domains domain, mfem::Array<int> &mask) {
    mask.SetSize(fem.mesh->attributes.Max());
    mask = 0;

    if ((domain & Domains::CORE) == Domains::CORE) {
        mask[0] = 1;
    }

    if ((domain & Domains::ENVELOPE) == Domains::ENVELOPE) {
        mask[1] = 1;
    }

    if ((domain & Domains::VACUUM) == Domains::VACUUM) {
        mask[2] = 1;
    }
}

std::expected<Bounds, BoundsError> DiscoverBounds(const mfem::Mesh *mesh, const int vacuum_attr) {
    double min_r = std::numeric_limits<double>::max();
    double max_r = -std::numeric_limits<double>::max();

    bool found_vacuum = false;

    for (int i = 0; i < mesh->GetNE(); ++i) {
        if (mesh->GetAttribute(i) == vacuum_attr) {
            found_vacuum = true;
            mfem::Array<int> vertices;
            mesh->GetElementVertices(i, vertices);
            for (const int v: vertices) {
                const double* coords = mesh->GetVertex(v);
                double r = std::sqrt(coords[0] * coords[0] + coords[1] * coords[1] + coords[2] * coords[2]);
                min_r = std::min(min_r, r);
                max_r = std::max(max_r, r);
            }
        }
    }
    if (found_vacuum) {
        return Bounds(min_r, max_r);
    }
    return std::unexpected(BoundsError::CANNOT_FIND_VACUUM);
}

void conserve_mass(const FEM& fem, mfem::GridFunction& rho, const double target_mass) {
    if (const double current_mass = domain_integrate_grid_function(fem, rho); current_mass > 1e-15) rho *= (target_mass / current_mass);
}
//endregion

//region Physics Functions
double centrifugal_potential(const mfem::Vector& phys_x, double omega) {
    const double s2 = std::pow(phys_x(0), 2) + std::pow(phys_x(1), 2);
    return -0.5 * s2 * std::pow(omega, 2);
}

double get_moment_of_inertia(const FEM& fem, const mfem::GridFunction& rho) {
    auto s2_func = [](const mfem::Vector& x) {
        return std::pow(x(0), 2) + std::pow(x(1), 2);
    };

    std::unique_ptr<mfem::Coefficient> s2_coeff;
    if (fem.has_mapping()) {
        s2_coeff = std::make_unique<PhysicalPositionFunctionCoefficient>(*fem.mapping, s2_func);
    } else {
        s2_coeff = std::make_unique<mfem::FunctionCoefficient>(s2_func);
    }

    mfem::GridFunctionCoefficient rho_coeff(&rho);
    mfem::ProductCoefficient I_integrand(rho_coeff, *s2_coeff);

    mfem::LinearForm I_lf(fem.H1_fes.get());

    double I = 0.0;
    if (fem.has_mapping()) {
        MappedScalarCoefficient mapped_integrand(*fem.mapping, I_integrand);
        I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(mapped_integrand));
        I_lf.Assemble();
        I = I_lf.Sum();
    } else {
        I_lf.AddDomainIntegrator(new mfem::DomainLFIntegrator(I_integrand));
        I_lf.Assemble();
        I = I_lf.Sum();
    }

    return I;
}
//endregion

//region Analytic External Potential for an oblate spheroid (doi.10.16.j.pss.2018.01.005)

inline double sq(const double x) { return x * x; }

double eccentricity(const double a, const double c) {
    return std::sqrt(1.0 - sq(c)/sq(a));
}

double equatorial_radius(const mfem::Vector& x) {
    return std::sqrt(sq(x(0)) + sq(x(1)));
}

double beta_ext(const mfem::Vector& x, const double a, const double c) {
    const double e = eccentricity(a, c);
    const double z = x(2);
    const double r = equatorial_radius(x);
    const double ae_sq = sq(a*e);

    if (std::abs(r) < 1e-12) {
        return std::atan2(a*e, std::abs(x(2)));
    }

    const double fraction = (sq(z) + ae_sq) / sq(r);
    const double first_term = -fraction;

    const double radical_one = sq(1.0 + fraction);
    const double radical_two = 4.0 * ae_sq / sq(r);
    const double second_term = std::sqrt(std::max(0.0, radical_one - radical_two));

    const double cos_term = first_term + second_term;
    return 0.5 * std::acos(std::clamp(cos_term, -1.0, 1.0));
}

double oblate_spheroid_surface_potential(const mfem::Vector& x, const double a, const double c, const double total_mass) {
    const double beta = beta_ext(x, a, c);
    const double r = equatorial_radius(x);
    const double e = eccentricity(a, c);
    const double z = x(2);

    const double G_M = G * total_mass;
    const double e3a3 = std::pow(e, 3) * std::pow(a, 3);

    const double t1 = (3.0 * G_M * beta) / (2.0 * e * a);

    const double t2aa = (3.0 * G_M * sq(r)) / (2.0 * e3a3);
    const double t2ab = beta - std::sin(beta) * std::cos(beta);

    const double t2ba = (3.0 * G_M * sq(z)) / e3a3;
    const double t2bb = std::tan(beta) - beta;

    return -t1 + 0.5 * (t2aa * t2ab + t2ba * t2bb);
}
//endregion End Analytic Potential

//region potentials
std::unique_ptr<mfem::GridFunction> grav_potential(const FEM& fem, const Args& args, const mfem::GridFunction& rho) {
    auto phi = std::make_unique<mfem::GridFunction>(fem.H1_fes.get());

    mfem::Array<int> ess_bdr(fem.mesh->bdr_attributes.Max());
    ess_bdr = 0;
    assert(ess_bdr.Size() == 2);
    ess_bdr[1] = 1; // The only boundary condition is that the potential goes to 0 at infinity

    // mfem::ConstantCoefficient zero(0.0);
    // phi->ProjectBdrCoefficient(zero, ess_bdr); // Set the potential to 0 at infinity;

    double total_mass = domain_integrate_grid_function(fem, rho, STELLAR);
    auto bdr_func = [&fem, total_mass](const mfem::Vector& x) {
        double r = x.Norml2();
        double theta = std::atan2(x(1), x(0));
        double phi = std::acos(x(2) / r);
        double val = l2_multipole_potential(fem, total_mass, x);
        // std::println("{},{},{},{},{}", r, theta, phi, val, -G*total_mass / r);
        return l2_multipole_potential(fem, total_mass, x);
    };

    std::unique_ptr<mfem::Coefficient> phi_bdr_coeff;
    if (fem.has_mapping()) {
        phi_bdr_coeff = std::make_unique<PhysicalPositionFunctionCoefficient>(*fem.mapping, bdr_func);
    } else {
        phi_bdr_coeff = std::make_unique<mfem::FunctionCoefficient>(bdr_func);
    }
    phi->ProjectBdrCoefficient(*phi_bdr_coeff, ess_bdr);


    mfem::Array<int> ess_tdof_list;
    fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list);

    mfem::Array<int> stellar_marker;
    populate_element_mask(fem, STELLAR, stellar_marker);

    auto laplacian = std::make_unique<mfem::BilinearForm>(fem.H1_fes.get());

    mfem::ConstantCoefficient one_coeff(1.0);
    std::unique_ptr<MappedDiffusionCoefficient> mapped_diff_coeff;
    if (fem.has_mapping()) {
        mapped_diff_coeff = std::make_unique<MappedDiffusionCoefficient>(*fem.mapping, one_coeff, fem.mesh->Dimension());
        laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator(*mapped_diff_coeff)); // The laplacian is global
    } else {
        laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator());
    }

    laplacian->Assemble();
    laplacian->Finalize();

    mfem::LinearForm b(fem.H1_fes.get());
    mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G);
    mfem::GridFunctionCoefficient rho_coeff(&rho);
    mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G);

    auto mapped_rhs = std::make_unique<MappedScalarCoefficient>(*fem.mapping, rhs_coeff, MappedScalarCoefficient::EVAL_POINTS::REFERENCE);
    if (fem.has_mapping()) {
        b.AddDomainIntegrator(new mfem::DomainLFIntegrator(*mapped_rhs), stellar_marker); // The mass contribution is local to the stellar domain
    } else {
        b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff), stellar_marker);
    }

    b.Assemble();

    mfem::OperatorPtr A;
    mfem::Vector B, X;
    laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B);

    mfem::GSSmoother prec;
    mfem::CGSolver cg;
    cg.SetPreconditioner(prec);
    cg.SetOperator(*A);
    cg.SetRelTol(args.p.tol);
    cg.SetMaxIter(args.p.max_iters);
    cg.Mult(B, X);

    laplacian->RecoverFEMSolution(X, b, *phi);
    return phi;
}

// std::unique_ptr<mfem::GridFunction> grav_potential(const FEM& fem, const Args &args, const mfem::GridFunction& rho, std::optional<OblatePotential> oblate) {
//     auto phi = std::make_unique<mfem::GridFunction>(fem.H1_fes.get());
//
//     mfem::Array<int> ess_bdr(fem.mesh->bdr_attributes.Max());
//     ess_bdr = 1;
//
//     mfem::GridFunctionCoefficient rho_coeff(&rho);
//     double total_mass = domain_integrate_grid_function(fem, rho);
//
//     auto grav_potential = [&fem, &total_mass, &oblate](const mfem::Vector& x) {
//         if (!oblate.has_value() or oblate->use == false) {
//             return l2_multipole_potential(fem, total_mass, x);
//         }
//         return oblate_spheroid_surface_potential(x, oblate->a, oblate->c, total_mass);
//     };
//
//     std::unique_ptr<mfem::Coefficient> phi_bdr_coeff;
//     if (fem.has_mapping()) {
//         phi_bdr_coeff = std::make_unique<PhysicalPositionFunctionCoefficient>(*fem.mapping, grav_potential);
//     } else {
//         phi_bdr_coeff = std::make_unique<mfem::FunctionCoefficient>(grav_potential);
//     }
//     phi->ProjectBdrCoefficient(*phi_bdr_coeff, ess_bdr);
//
//     mfem::Array<int> ess_tdof_list;
//     fem.H1_fes->GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
//
//     auto laplacian = std::make_unique<mfem::BilinearForm>(fem.H1_fes.get());
//     // ReSharper disable once CppTooWideScope
//     std::unique_ptr<mfem::ConstantCoefficient> one_coeff;
//     // ReSharper disable once CppTooWideScope
//     std::unique_ptr<MappedDiffusionCoefficient> mapped_diff_coeff;
//     if (fem.has_mapping()) {
//         one_coeff = std::make_unique<mfem::ConstantCoefficient>(1.0);
//         mapped_diff_coeff = std::make_unique<MappedDiffusionCoefficient>(*fem.mapping, *one_coeff, fem.mesh->Dimension());
//         laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator(*mapped_diff_coeff));
//     } else {
//         laplacian->AddDomainIntegrator(new mfem::DiffusionIntegrator());
//     }
//     laplacian->Assemble();
//     laplacian->Finalize();
//
//     mfem::ConstantCoefficient four_pi_G(-4.0 * M_PI * G);
//     mfem::ProductCoefficient rhs_coeff(rho_coeff, four_pi_G);
//     mfem::LinearForm b(fem.H1_fes.get());
//
//     // ReSharper disable once CppTooWideScope
//     std::unique_ptr<MappedScalarCoefficient> mapped_rhs;
//     if (fem.has_mapping()) {
//         mapped_rhs = std::make_unique<MappedScalarCoefficient>(*fem.mapping, rhs_coeff, MappedScalarCoefficient::EVAL_POINTS::REFERENCE);
//         b.AddDomainIntegrator(new mfem::DomainLFIntegrator(*mapped_rhs));
//     } else {
//         b.AddDomainIntegrator(new mfem::DomainLFIntegrator(rhs_coeff));
//     }
//     b.Assemble();
//
//     mfem::OperatorPtr A;
//     mfem::Vector B, X;
//     laplacian->FormLinearSystem(ess_tdof_list, *phi, b, A, X, B);
//
//     mfem::GSSmoother prec;
//     mfem::CGSolver cg;
//     cg.SetPreconditioner(prec);
//     cg.SetOperator(*A);
//     cg.SetRelTol(args.p.tol);
//     cg.SetMaxIter(args.p.max_iters);
//     cg.SetPrintLevel(0);
//
//     cg.Mult(B, X);
//     laplacian->RecoverFEMSolution(X, b, *phi);
//     return phi;
// }

std::unique_ptr<mfem::GridFunction> get_potential(const FEM &fem, const Args &args, const mfem::GridFunction &rho) {
    auto phi = grav_potential(fem, args, rho);

    if (args.r.enabled) {
        auto rot = [&args](const mfem::Vector& x) {
            return centrifugal_potential(x, args.r.omega);
        };

        std::unique_ptr<mfem::Coefficient> centrifugal_coeff;
        if (fem.has_mapping()) {
            centrifugal_coeff = std::make_unique<PhysicalPositionFunctionCoefficient>(*fem.mapping, rot);
        } else {
            centrifugal_coeff = std::make_unique<mfem::FunctionCoefficient>(rot);
        }
        mfem::GridFunction centrifugal_gf(fem.H1_fes.get());
        centrifugal_gf.ProjectCoefficient(*centrifugal_coeff);
        (*phi) += centrifugal_gf;
    }

    return phi;
}

mfem::DenseMatrix compute_quadrupole_moment_tensor(const FEM& fem, const mfem::GridFunction& rho, const mfem::Vector& com) {
    const int dim = fem.mesh->Dimension();
    mfem::DenseMatrix Q(dim, dim);
    Q = 0.0;

    for (int i = 0; i < fem.H1_fes->GetNE(); ++i) {
        mfem::ElementTransformation *trans = fem.mesh->GetElementTransformation(i);
        const mfem::IntegrationRule &ir = mfem::IntRules.Get(trans->GetGeometryType(), 2 * fem.H1_fes->GetOrder(0) + trans->OrderW());
        for (int j = 0; j < ir.GetNPoints(); ++j) {
            const mfem::IntegrationPoint &ip = ir.IntPoint(j);
            trans->SetIntPoint(&ip);

            double weight = trans->Weight() * ip.weight;

            if (fem.has_mapping()) {
                weight *= fem.mapping->ComputeDetJ(*trans, ip);
            }

            const double rho_val = rho.GetValue(i, ip);

            mfem::Vector phys_point(dim);
            if (fem.has_mapping()) {
                fem.mapping->GetPhysicalPoint(*trans, ip, phys_point);
            } else {
                trans->Transform(ip, phys_point);
            }

            mfem::Vector x_prime(dim);
            double r_sq = 0.0;

            for (int d = 0; d < dim; ++d) {
                x_prime(d) = phys_point(d) - com(d);
                r_sq += x_prime(d) * x_prime(d);
            }

            for (int m = 0; m < dim; ++m) {
                for (int n = 0; n < dim; ++n) {
                    const double delta = (m == n) ? 1.0 : 0.0;
                    const double contrib = 3.0 * x_prime(m) * x_prime(n) - delta * r_sq;
                    Q(m, n) += rho_val * contrib * weight;
                }
            }
        }
    }
    return Q;
}

double l2_multipole_potential(const FEM &fem, const double total_mass, const mfem::Vector &phys_x) {
    const double r = phys_x.Norml2();
    if (r < 1e-12) return 0.0;

    const int dim = fem.mesh->Dimension();

    mfem::Vector n(phys_x);
    n /= r;

    double l2_mult_factor = 0.0;
    for (int i = 0; i < dim; ++i) {
        for (int j = 0; j < dim; ++j) {
            l2_mult_factor += fem.Q(i, j) * n(i) * n(j);
        }
    }

    const double l2_contrib = (G / (2.0 * std::pow(r, 3))) * l2_mult_factor;

    const double l0_contrib = -G * total_mass / r;

    // l1 contribution is zero for a system centered on its COM
    return l0_contrib + l2_contrib;
}
//endregion

//region Tests
void test_mesh_load(const FEM& fem) {
    size_t failed = 0;
    if (not fem.okay()) ++failed;
    const int dim = fem.mesh->Dimension();
    if (dim != 3) ++failed;

    const int n_scalar = fem.H1_fes->GetTrueVSize();
    const int n_vector = fem.Vec_H1_fes->GetTrueVSize();
    if (n_vector != dim * n_scalar) ++failed;

    if (fem.block_offsets[0] != 0) ++failed;
    if (fem.block_offsets[1] != n_scalar) ++failed;
    if (fem.block_offsets[2] != n_scalar + n_vector) ++failed;
    if (fem.block_offsets[3] != n_scalar + n_vector + 1) ++failed;

    constexpr size_t num_tests = 6;
    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }
    std::string test_msg = fmt_test_msg("Mesh Load Test", result_type, failed, num_tests);
    std::println("{}", test_msg);

    assert(dim == 3);
    assert(n_vector == (n_scalar * dim));
    assert (fem.block_offsets[0] == 0);
    assert (fem.block_offsets[1] == n_scalar);
    assert (fem.block_offsets[2] == n_scalar + n_vector);
    assert (fem.block_offsets[3] == n_scalar + n_vector + 1);
}

void test_ref_coord_storage(const FEM& fem) {
    size_t failed = 0;
    if (not fem.mapping->IsIdentity()) ++failed;

    const size_t num_elemIDs = std::min(30, fem.mesh->GetNE());
    for (int elemID = 0; elemID < num_elemIDs; ++elemID) {
        auto* trans = fem.mesh->GetElementTransformation(elemID);
        const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2);
        const auto& ip = ir.IntPoint(0);
        trans->SetIntPoint(&ip);

        mfem::Vector x_ref, x_phys;
        trans->Transform(ip, x_ref);
        fem.mapping->GetPhysicalPoint(*trans, ip, x_phys);
        x_ref -= x_phys;

        if (not (x_ref.Norml2() < 1e-12)) ++failed;
    }

    const size_t num_tests = num_elemIDs + 1;
    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::string test_msg = fmt_test_msg("Mesh Ref Coord", result_type, failed, num_tests);
    std::println("{}", test_msg);
}

void test_reference_volume_integral(const FEM& fem) {
    size_t failed = 0;

    mfem::GridFunction ones(fem.H1_fes.get());
    ones = 1.0;

    double vol = domain_integrate_grid_function(fem, ones, STELLAR);
    double expected = (4.0/3.0) * M_PI * std::pow(RADIUS, 3.0);
    double rel_err = std::abs(vol - expected) / expected;

    if (rel_err > 1e-2) ++failed;

    constexpr size_t num_tests = 1;
    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    }

    std::println("{}", fmt_test_msg("Reference Volume Integral", result_type, failed, num_tests));

    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\tFAILURE: Volume: {}, Expected: {}, Error (rel): {}", vol, expected, rel_err);
    }
}

void test_spherically_symmetric_com(const FEM& fem) {
    mfem::GridFunction rho(fem.H1_fes.get());
    rho = 1.0;

    mfem::Vector com = get_com(fem, rho);
    size_t failed = 0;

    const size_t dim = fem.mesh->Dimension();
    const size_t num_tests = dim;

    for (int dimID = 0; dimID < num_tests; ++dimID) {
        if (std::abs(com(dimID)) > 1e-12) ++failed;
    }

    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::println("{}", fmt_test_msg("Uniform COM", result_type, failed, num_tests));

    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\t COM=<{:+0.3E}, {:+0.3E}, {:+0.3E}>", com(0), com(1), com(2));
    }
}

void test_com_variance_to_displacement(const FEM& fem) {
    size_t failed = 0;
    mfem::GridFunction linear_displacement(fem.Vec_H1_fes.get());
    linear_displacement = 10.0; // This will linearly displace the domain by 10 unit along all axes

    fem.mapping->SetDisplacement(linear_displacement);

    mfem::GridFunction rho(fem.H1_fes.get());
    rho = 1.0;

    mfem::Vector mapped_com = get_com(fem, rho);

    const size_t dim = fem.mesh->Dimension();
    const size_t num_tests = dim;
    for (int dimID = 0; dimID < num_tests; ++dimID) {
        if (10 - std::abs(mapped_com(dimID)) > 1e-12) ++failed;
    }

    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::println("{}", fmt_test_msg("COM variance to displacement", result_type, failed, num_tests));

    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\tFAILURE COM=<{:+0.2E}, {:+0.2E}, {:+0.2E}>", mapped_com(0), mapped_com(1), mapped_com(2));
    }

    fem.mapping->ResetDisplacement();
}

void test_volume_invariance_to_displacement(const FEM& fem) {
    size_t failed = 0;
    mfem::GridFunction linear_displacement(fem.Vec_H1_fes.get());
    linear_displacement = 10.0; // This will linearly displace the domain by 10 unit along all axes

    fem.mapping->SetDisplacement(linear_displacement);

    mfem::GridFunction ones(fem.H1_fes.get());
    ones = 1.0;
    double mapped_vol = domain_integrate_grid_function(fem, ones, STELLAR);
    double expected = (4.0/3.0) * M_PI * std::pow(RADIUS, 3.0);
    double rel_err = std::abs(mapped_vol - expected) / expected;

    if (rel_err > 1e-2) ++failed;
    constexpr size_t num_tests = 1;
    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    }

    std::println("{}", fmt_test_msg("Invariance of volume against translation", result_type, failed, num_tests));

    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\tFAILURE: Volume: {}, Expected: {}", mapped_vol, expected);
    }
    fem.mapping->ResetDisplacement();
}

void test_volume_ellipsoid_deformation(const FEM& fem) {
    size_t failed = 0;
    size_t num_tests = 0;

    constexpr double a = 2.0; // x-axis
    constexpr double b = 0.5; // y-axis
    constexpr double c = 1.5; // z-axis
    constexpr double expected_vol = (4.0 / 3.0) * M_PI * a * b * c;

    mfem::GridFunction ellipsoid_displacement(fem.Vec_H1_fes.get());
    {
        const int dim = fem.mesh->Dimension();
        mfem::VectorFunctionCoefficient disp_coeff(dim, [&](const mfem::Vector& x, mfem::Vector& d) {
            d.SetSize(x.Size());
            d(0) = (a - 1.0) * x(0);
            d(1) = (b - 1.0) * x(1);
            d(2) = (c - 1.0) * x(2);
        });
        ellipsoid_displacement.ProjectCoefficient(disp_coeff);
    }
    fem.mapping->SetDisplacement(ellipsoid_displacement);

    {
        ++num_tests;
        mfem::GridFunction ones(fem.H1_fes.get());
        ones = 1.0;
        const double mapped_vol = domain_integrate_grid_function(fem, ones, STELLAR);
        const double rel_err = std::abs(mapped_vol - expected_vol) / expected_vol;
        if (rel_err > 1e-3) {
            ++failed;
            std::println("\tFAILURE Test 1: Mapped volume = {:.6f}, expected = {:.6f}, rel_err = {:.2e}",
                         mapped_vol, expected_vol, rel_err);
        }
    }

    {
        ++num_tests;
        const double expected_x2_integral = std::pow(a, 3) * b * c * (4.0 * M_PI / 15.0);

        mfem::GridFunction x_ref_sq(fem.H1_fes.get());
        mfem::FunctionCoefficient x_sq_coeff([](const mfem::Vector& x) {
            return x(0) * x(0);
        });
        x_ref_sq.ProjectCoefficient(x_sq_coeff);

        mfem::GridFunction x_phys_sq(fem.H1_fes.get());
        PhysicalPositionFunctionCoefficient x_phys_sq_coeff(*fem.mapping,
            [](const mfem::Vector& x_phys) {
                return x_phys(0) * x_phys(0);
            }
        );
        x_phys_sq.ProjectCoefficient(x_phys_sq_coeff);

        const double mapped_x2_integral = domain_integrate_grid_function(fem, x_phys_sq, STELLAR);
        const double rel_err = std::abs(mapped_x2_integral - expected_x2_integral) / expected_x2_integral;
        if (rel_err > 1e-3) {
            ++failed;
            std::println("\tFAILURE Test 2: integral x_phys^2 = {:.6f}, expected = {:.6f}, rel_err = {:.2e}",
                         mapped_x2_integral, expected_x2_integral, rel_err);
        }
    }

    {
        ++num_tests;
        constexpr double expected_detJ = a * b * c;
        double max_detJ_err = 0.0;
        for (int e = 0; e < std::min(5, fem.mesh->GetNE()); ++e) {
            auto* trans = fem.mesh->GetElementTransformation(e);
            const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2);
            for (int q = 0; q < ir.GetNPoints(); ++q) {
                const auto& ip = ir.IntPoint(q);
                trans->SetIntPoint(&ip);
                const double detJ = fem.mapping->ComputeDetJ(*trans, ip);
                max_detJ_err = std::max(max_detJ_err, std::abs(detJ - expected_detJ));
            }
        }
        if (max_detJ_err > 1e-10) {
            ++failed;
            std::println("\tFAILURE Test 3: max pointwise |det(J) - a*b*c| = {:.2e}", max_detJ_err);
        }
    }

    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::println("{}", fmt_test_msg("Volume under ellipsoidal deformation", result_type, failed, num_tests));

    fem.mapping->ResetDisplacement();
}

void test_uniform_potential(FEM& fem, const Args& args) {
    const double analytic_vol = (4.0/3.0) * M_PI * std::pow(RADIUS, 3);
    const double rho0 = MASS / analytic_vol;

    mfem::GridFunction rho_uniform(fem.H1_fes.get());
    rho_uniform = rho0;

    fem.com = get_com(fem, rho_uniform);
    fem.Q = compute_quadrupole_moment_tensor(fem, rho_uniform, fem.com);

    const auto phi = grav_potential(fem, args, rho_uniform);

    double max_abs_err = 0.0;
    double max_rel_err = 0.0;
    constexpr double tol = 1e-3;

    size_t failed = 0;
    size_t num_tests = 0;
    const size_t num_elemIDs = std::min(30, fem.mesh->GetNE());
    for (int elemID = 0; elemID < num_elemIDs; ++elemID) {
        num_tests++;
        auto* trans = fem.mesh->GetElementTransformation(elemID);
        const auto& ir = mfem::IntRules.Get(trans->GetGeometryType(), 2);
        const auto& ip = ir.IntPoint(0);
        trans->SetIntPoint(&ip);

        mfem::Vector x_ref, x_phys;
        trans->Transform(ip, x_ref);
        fem.mapping->GetPhysicalPoint(*trans, ip, x_phys);

        const double r = x_phys.Norml2();

        if (r < 1e-9) continue;

        const double phi_analytic = (-G * MASS / (2.0 * std::pow(RADIUS, 3.0))) * (3.0*RADIUS*RADIUS - r*r);

        const double phi_fem = phi->GetValue(elemID, ip);

        const double abs_err = std::abs(phi_fem  - phi_analytic);
        const double rel_err = abs_err / std::abs(phi_analytic);

        max_abs_err = std::max(max_abs_err, abs_err);
        max_rel_err = std::max(max_rel_err, rel_err);

        if (rel_err > tol) ++failed;
    }

    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::println("{}", fmt_test_msg("Test Uniform Potential", result_type, failed, num_tests));

    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\tFAILURE: max abs error: {:+0.2E}, max rel error: {:+0.2E}", max_abs_err, max_rel_err);
    }
}

void test_ellipsoidal_potential(FEM& fem, const Args& args) {
    constexpr double a = 1.0 * RADIUS;
    constexpr double b = a; // oblate
    constexpr double c = 0.99 * RADIUS;

    constexpr double expected_vol = (4.0 / 3.0) * M_PI * a * b * c;
    constexpr double rho0 = MASS / expected_vol;

    mfem::GridFunction ellipsoidal_disp(fem.Vec_H1_fes.get());
    mfem::VectorFunctionCoefficient disp_coeff(3, [&](const mfem::Vector& x, mfem::Vector& d) {
        d.SetSize(3);
        d(0) = (a/RADIUS - 1.0) * x(0);
        d(1) = (b/RADIUS - 1.0) * x(1);
        d(2) = (c/RADIUS - 1.0) * x(2);
    });
    ellipsoidal_disp.ProjectCoefficient(disp_coeff);
    fem.mapping->SetDisplacement(ellipsoidal_disp);

    mfem::GridFunction rho(fem.H1_fes.get());
    rho = rho0;

    fem.com = get_com(fem, rho);
    fem.Q = compute_quadrupole_moment_tensor(fem, rho, fem.com);

    // OblatePotential oblate{.use=true, .a=a, .c=c,.rho_0=rho0};
    const auto phi = grav_potential(fem, args, rho);

    constexpr double e_sq = 1.0 - (c * c)/(a*a);
    const double e = std::sqrt(e_sq);

    const double I_const = (2.0 * std::sqrt(1.0 - e_sq) / e) * std::asin(e);
    const double A_R = (std::sqrt(1.0-e_sq) / std::pow(e, 3.0)) * std::asin(e) - (1.0 - e_sq)/e_sq;
    const double A_z = (2.0 / e_sq) * (1.0 - (std::sqrt(1.0-e_sq) / e) * std::asin(e));

    size_t failed = 0;
    size_t num_tests = 0;
    double max_rel_err = 0.0;
    double total_err = 0.0;
    const size_t check_count = std::min(50, fem.mesh->GetNE());

    for (int elemID = 0; elemID < check_count; ++elemID) {
        auto* trans = fem.mesh->GetElementTransformation(elemID);
        const auto& ip = mfem::IntRules.Get(trans->GetGeometryType(), 2).IntPoint(0);
        trans->SetIntPoint(&ip);

        mfem::Vector x_phys;
        fem.mapping->GetPhysicalPoint(*trans, ip, x_phys);

        const double R2 = x_phys(0)*x_phys(0) + x_phys(1)*x_phys(1);
        const double z2 = x_phys(2)*x_phys(2);
        const double phi_analytic = -M_PI * G * rho0 * (a*a*I_const - A_R * R2 - A_z * z2);

        const double phi_fem = phi->GetValue(elemID, ip);
        const double rel_err = std::abs(phi_fem - phi_analytic) / std::abs(phi_analytic);
        max_rel_err = std::max(max_rel_err, rel_err);
        total_err += rel_err;
        num_tests++;
        if (rel_err > 1e-3) ++failed;
    }

    auto result_type = TEST_RESULT_TYPE::FAILURE;
    if (failed == 0) {
        result_type = TEST_RESULT_TYPE::SUCCESS;
    } else if (failed < num_tests) {
        result_type = TEST_RESULT_TYPE::PARTIAL;
    }

    std::println("{}", fmt_test_msg("Test Ellipsoidal Potential", result_type, failed, num_tests));
    if (result_type == TEST_RESULT_TYPE::FAILURE) {
        std::println("\tFAILURE: max rel error: {:+0.2E}, mean rel error: {:+0.2E}", max_rel_err, total_err/static_cast<double>(num_tests));
    }

}
//endregion

int main(int argc, char** argv) {
    Args args;

    CLI::App app{"Mapped Coordinates XAD-Enabled Non-Linear Solver"};
    app.add_option("-m,--mesh", args.mesh_file)->required()->check(CLI::ExistingFile);
    app.add_option("--max-iters", args.max_iters)->default_val(20);
    app.add_option("--index", args.index)->default_val(1);
    app.add_option("--mass", args.mass)->default_val(MASS);
    app.add_option("--surface-pressure-scalar", args.c)->default_val(1e-4);

    args.r.enabled = false;
    args.p.tol = 1e-6;
    args.p.max_iters = 1000;

    CLI11_PARSE(app, argc, argv);
    FEM fem = setup_fem(args.mesh_file, true);
    // fem.mesh->UniformRefinement();


    test_mesh_load(fem);
    test_ref_coord_storage(fem);
    test_reference_volume_integral(fem);
    test_spherically_symmetric_com(fem);

    test_com_variance_to_displacement(fem);
    test_volume_invariance_to_displacement(fem);
    test_volume_ellipsoid_deformation(fem);
    // test_kelvin_jacobian(fem);
    test_uniform_potential(fem, args);
    test_ellipsoidal_potential(fem, args); // Note that this test is currently predicated on an analytic solution for the surface boundary potential

    CoupledState state(fem);

    // typedef xad::AReal<double> adouble;
    //
    // EOS_P<adouble> eos_enthalpy = [index = args.index](const adouble& rho, const adouble& temp) -> adouble {
    //     if (rho.value() < 1e-15) return {0.0};
    //     return (index + 1.0) * pow(rho, 1.0 / index);
    // };
    //
    // EOS_P<adouble> eos_pressure = [index = args.index](const adouble& rho, const adouble& temp) -> adouble {
    //     if (rho.value() < 1e-15) return {0.0};
    //     return pow(rho, 1.0 + (1.0 / index));
    // };
    //
    // auto init_rho_func = [&](const mfem::Vector& pt) {
    //     const double r = pt.Norml2();
    //     return (r < RADIUS) ? (1.0 - r/RADIUS) : 0.0;
    // };
    // mfem::FunctionCoefficient init_rho_coeff(init_rho_func);
    // state.rho.ProjectCoefficient(init_rho_coeff);
    // conserve_mass(fem, state.rho, args.mass);
    // view_mesh(HOST, PORT, *fem.mesh, state.rho, "Initial Density Solution with XAD AD");
    // view_mesh(HOST, PORT, *fem.mesh, state.d, "Initial Position Solution with XAD AD");
    //
    // std::println("Starting Coupled Block Solver with XAD AD...");
    //
    // ResidualOperator coupled_operator(fem, args, eos_enthalpy, eos_pressure);
    //
    // mfem::NewtonSolver newton_solver;
    // newton_solver.SetOperator(coupled_operator);
    // newton_solver.SetMaxIter(500);
    // newton_solver.SetRelTol(1e-6);
    // newton_solver.SetPrintLevel(1);
    //
    // mfem::GMRESSolver jfnk_solver;
    // jfnk_solver.SetMaxIter(100);
    // jfnk_solver.SetRelTol(1e-4);
    // newton_solver.SetSolver(jfnk_solver);
    // newton_solver.SetAdaptiveLinRtol(2, 0.5, 0.9, 0.9, 1.1);
    //
    // mfem::Vector zero_rhs(fem.block_offsets.Last());
    // zero_rhs = 0.0;
    // newton_solver.Mult(zero_rhs, *state.U);
    //
    // mfem::GridFunction rho(fem.H1_fes.get(), state.U->GetBlock(0), 0);
    // mfem::GridFunction x(fem.Vec_H1_fes.get(), state.U->GetBlock(1), 0);
    // view_mesh(HOST, PORT, *fem.mesh, rho, "Final Density Solution with XAD AD");
    // view_mesh(HOST, PORT, *fem.mesh, x, "Final Position Solution with XAD AD");
    //
    // std::println("Solver finished using XAD machine-precision gradients.");
    // return 0;

}