#include <cmath>
#include <memory>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/general/communication.hpp"
#include "serif/backend/mfem/solver/seed/state.hpp"
#include "serif/solver/equilibrium.hpp"
#include "serif/solver/seed/polytrope.hpp"
#include "serif/tests/nonconforming.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/tests/solver/virial.hpp"

namespace {
    namespace solver = serif::solver;
    namespace fixture = serif::tests::solver;
}

TEST_CASE("A fixed STROID nonconforming mesh supports a complete coupled stellar solve", "[solver][nonconforming][nonconforming_stellar_solve]") {
    if (!mfem::Mpi::IsInitialized()) {
        mfem::Mpi::Init();
    }
    const auto model = fixture::make_model();
    const auto seed = solver::seed::make_polytropic_seed(model);
    REQUIRE(seed);
    CHECK(seed->scales.radius.value() == Catch::Approx(1.0).margin(2.0e-10));
    const auto resources = std::make_shared<fixture::StellarResources>(serif::tests::nonconforming::mesh(1, 2, 6));
    REQUIRE(resources->potential_space.GetTrueVSize() < resources->potential_space.GetVSize());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    const auto initial = solver::seed::backend::mfem::make_polytropic_state(model, finite_elements, *seed);
    REQUIRE(initial);
    const auto scales = solver::normalization::prepare_stellar_reference_scales(
        serif::dimensions::MassValue{1.0}, serif::dimensions::LengthValue{1.0});
    REQUIRE(scales);
    const auto configuration = serif::model::make_discretization(serif::model::DiscretizationOptions{
        .model = model, .finite_elements = finite_elements,
        .normalization = solver::normalization::NormalizationOptions{.scheme = *scales}});
    const solver::SolverOptions options{
        .method = {.relative_tolerance = 0.0, .absolute_tolerance = 1.0e-10, .maximum_iterations = 20,
            .linear_solve = {.relative_tolerance = 1.0e-3, .maximum_iterations = 200}}};
    const auto result = solver::solve_equilibrium(configuration, *initial, options);
    INFO(result.message());
    REQUIRE(result.converged());
    REQUIRE(result.diagnostics().final_residual_norm);
    CHECK(*result.diagnostics().final_residual_norm <= 1.0e-10);
    CHECK(result.diagnostics().accepted_iterations > 0);

    const auto& state = result.state();
    const auto trace = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::StellarSurfaceBoundary>(resources->density);
    const auto density = state.template get<fixture::Density>();
    for (const int dof : trace.reduced_dofs()) {
        CHECK(density[dof] == 0.0);
    }
    const double omega = state.template get<fixture::Omega>()[0];
    const auto virial = fixture::vacuum_virial(finite_elements, state, model.equation_of_state(),
        serif::equations::RigidRotation{{0.0, 0.0, omega}});
    if (!virial) {
        INFO(virial.error());
    }
    REQUIRE(virial);
    const auto& physical = virial->samples.back();
    INFO("Full virial = " << physical.energy_virial() << ", quadrature allowance = " << virial->quadrature_allowance
        << ", passes unchanged 1e-6 gate = " << virial->passes());
    CHECK(physical.mass == Catch::Approx(1.0).margin(2.0e-9));
    CHECK(std::abs(omega) < 1.0e-12);
    CHECK(std::isfinite(physical.energy_virial()));
    CHECK(virial->quadrature_allowance < 1.0e-7);
}