#include <algorithm>
#include <cmath>
#include <numbers>
#include <ranges>
#include <type_traits>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>

#include "serif/model/finite_elements.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/runtime.hpp"
#include "serif/solver/seed/state.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/model/execution_validation.hpp"
#include "../sandbox/polytropic_model.hpp"

TEST_CASE("Production resources reproduce the established stellar coefficient representation", "[model][resources]") {
    namespace fixture = serif::tests::solver;
    namespace model = serif::model;
    namespace deformation = serif::geometry::deformation;

    serif::initialize();
    serif::initialize(); // Reusing the application runtime must be harmless.
    const auto star = fixture::make_model();
    const auto radial = serif::solver::seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto old_resources = fixture::make_resources(radial->scales.radius.value());
    CHECK(old_resources->geometry_space.FEColl() == old_resources->displacement_space.FEColl());
    CHECK(old_resources->geometry_space.FEColl() == old_resources->reference_displacement.FEColl());
    const auto previous = fixture::make_discretization(old_resources, 2);
    const auto current = model::make_stellar_discretization(star, {.radius = radial->scales.radius},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{
            .surface = {.maximum_degree = 2}, .stellar_interior = {}, .vacuum = {}});
    const auto old_state = serif::solver::seed::make_polytropic_state(star, previous, *radial);
    const auto new_state = serif::solver::seed::make_polytropic_state(star, current, *radial);
    REQUIRE(old_state);
    REQUIRE(new_state);
    CHECK(std::ranges::equal(old_state->values(), new_state->values()));
    const auto copied = current;
    CHECK(copied.owner == current.owner);
    CHECK(&copied.template binding<fixture::Density>().space() == &current.template binding<fixture::Density>().space());
}

TEST_CASE("Geometry construction needs only the fields of the selected equations", "[model][resources]") {
    namespace model = serif::model;
    namespace equations = serif::equations;
    namespace deformation = serif::geometry::deformation;

    const auto star = model::make_stellar_model(serif::eos::models::Polytrope{1.0, 1.0}, model::EquilibriumEquations{
        .interior_equations = equations::EquationRegistry{equations::PrimalPoissonBalance{1.0}}});
    const auto finite_elements = model::make_stellar_discretization(star, {},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{});
    const auto& reference = finite_elements.deformation.reference;
    CHECK(reference.scalar_space.FEColl() == reference.displacement_space.FEColl());
    CHECK(reference.scalar_space.GetVDim() == 1);
    CHECK(reference.displacement_space.GetVDim() == 3);
    CHECK_THROWS(model::make_stellar_discretization(star, {.radius = serif::dimensions::LengthValue{-1.0}}));
    CHECK_THROWS(model::make_stellar_discretization(star, {.infinity_radius_fraction = 1.0}));
}

TEST_CASE("The polytropic factory retains fixed physical targets in the scalar formulation", "[model][unit][primal]") {
    namespace model = serif::model;
    namespace equations = serif::equations;
    namespace blocks = serif::discretization::blocks;
    namespace dependencies = serif::utils::dependencies;

    const auto star = serif::examples::make_polytropic_model(2.0, 0.4, 1.7, 0.03, 0.8, 0.02);
    const auto definition = model::resolve_equilibrium_equations(star);
    using Description = std::remove_cvref_t<decltype(definition)>;
    using Model = std::remove_cvref_t<decltype(star)>;
    using Plan = model::EquilibriumPlan<typename Model::EquationOfState, Description>;
    static_assert(model::executable_model_topology<Model, Description, Plan>());
    static_assert(dependencies::key_count_v<typename Plan::Values> == 5);
    static_assert(dependencies::key_count_v<typename Plan::Residuals> == 5);
    static_assert(std::same_as<typename Plan::Surface::Formulation::CarrierTerm, blocks::density::mass>);
    CHECK(star.equation_of_state().polytropic_index() == 2.0);
    CHECK(star.equation_of_state().polytropic_constant() == 0.4);
    CHECK(star.specification<model::FixedTotalMass>().target().value() == 1.7);
    CHECK(star.specification<model::FixedAngularMomentum>().target().value() == 0.03);
    CHECK(definition.interior_equations.get<equations::MassNormalization>().residual_offset() == -1.7);
    CHECK(definition.interior_equations.get<equations::AngularMomentumConstraint>().residual_offset() == -0.03);
    CHECK(definition.interior_equations.get<equations::PrimalPoissonBalance>().four_pi_G() == Catch::Approx(4.0 * std::numbers::pi * 0.8));
    CHECK(definition.surface_condition.target_pressure().value() == 0.02);
}

TEST_CASE("The factory rejects only the unresolved nonrotating n3 vacuum scale at equation resolution", "[model][unit][primal]") {
    namespace model = serif::model;
    CHECK_NOTHROW(serif::eos::models::Polytrope{3.0, 0.4});
    const auto unresolved = serif::examples::make_polytropic_model(3.0, 0.4, 1.0, 0.0, 1.0, 0.0);
    CHECK(unresolved.equation_of_state().polytropic_index() == 3.0);
    CHECK_THROWS_WITH(model::resolve_equilibrium_equations(unresolved),
        Catch::Matchers::ContainsSubstring("compatible mass leaves radius free"));
    const auto rotating = serif::examples::make_polytropic_model(3.0, 0.4, 1.0, 0.01, 1.0, 0.0);
    const auto confined = serif::examples::make_polytropic_model(3.0, 0.4, 1.0, 0.0, 1.0, 0.01);
    const auto below = serif::examples::make_polytropic_model(2.999, 0.4, 1.0, 0.0, 1.0, 0.0);
    CHECK_NOTHROW(model::resolve_equilibrium_equations(rotating));
    CHECK_NOTHROW(model::resolve_equilibrium_equations(confined));
    CHECK_NOTHROW(model::resolve_equilibrium_equations(below));
    // Admission of these declarations does not certify existence or uniqueness for their chosen targets.
}

TEST_CASE("The scalar factory seed contains only primal fields and fixes its density trace", "[model][resources][primal]") {
    namespace model = serif::model;
    namespace blocks = serif::discretization::blocks;
    namespace dependencies = serif::utils::dependencies;
    namespace seed = serif::solver::seed;
    namespace deformation = serif::geometry::deformation;
    using Density = blocks::density::mass;
    using Potential = blocks::gravity::potential;
    using SurfaceLevel = blocks::barotropic_constant::mass_normalization;

    serif::initialize();
    const auto star = serif::examples::make_polytropic_model();
    const auto radial = seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto finite_elements = model::make_stellar_discretization(star, {.radius = radial->scales.radius},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{
            .surface = {.maximum_degree = 2}, .stellar_interior = {}, .vacuum = {}});
    const auto state = seed::make_polytropic_state(star, finite_elements, *radial);
    if (!state) {
        INFO(state.error());
    }
    REQUIRE(state);
    using State = std::remove_cvref_t<decltype(*state)>;
    static_assert(dependencies::key_count_v<typename State::Layout::Values> == 5);
    CHECK(std::ranges::all_of(state->values(), [](double value) { return std::isfinite(value); }));
    CHECK(state->template get<Density>().size() == static_cast<std::size_t>(finite_elements.template binding<Density>().size()));
    CHECK(state->template get<Potential>().size() == static_cast<std::size_t>(finite_elements.template binding<Potential>().size()));
    CHECK(state->template get<blocks::surface_deformation::coefficients>().size() == 9);
    CHECK(state->template get<SurfaceLevel>()[0] == Catch::Approx(-1.0).epsilon(3.0e-7));
    CHECK(state->template get<blocks::rotation::angular_velocity>()[0] == 0.0);
    CHECK(std::ranges::max(state->template get<Density>()) > 0.7);
    const auto surface_rows = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::StellarSurfaceBoundary>(finite_elements.template binding<Density>());
    REQUIRE(surface_rows.size() > 0);
    for (const int dof : surface_rows.reduced_dofs()) {
        CHECK(state->template get<Density>()[dof] == 0.0);
    }
    const auto infinity_rows = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::InfinitySurfaceBoundary>(finite_elements.template binding<Potential>());
    REQUIRE(infinity_rows.size() > 0);
    for (const int dof : infinity_rows.reduced_dofs()) {
        CHECK(state->template get<Potential>()[dof] == 0.0);
    }
    CHECK(star.specification<model::FixedTotalMass>().target().value() == 1.0);
    CHECK(star.equation_of_state().polytropic_constant() == 2.0 / std::numbers::pi);
}