#include <concepts>
#include <type_traits>

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    [[nodiscard]] mean_field::fem::FEM makeFiniteElements() {
        const mean_field::utils::Args arguments = test_utils::setup_args();
        return mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    }

    template <typename Problem>
    concept PreparesWithGeneratedRotation = requires(
        Problem &problem,
        const mfem::Vector &state,
        const mean_field::operators::StellarEquilibriumDependencies &dependencies
    ) {
        problem.Prepare(state, dependencies);
    };

    template <typename Problem>
    concept PreparesWithPrescribedRotation = requires(
        Problem &problem,
        const mfem::Vector &state,
        const mean_field::operators::StellarEquilibriumDependencies &dependencies,
        const mean_field::physics::RigidRotation &rotation
    ) {
        problem.Prepare(state, dependencies, rotation);
    };
} // namespace

TEST_CASE(
    "Simple User API Builds A Stellar Model With Its Default Preconditioner",
    "[user-api][simple]"
) {
    using namespace mean_field;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());

    auto model = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}})
    );

    auto problem        = equilibrium::discretize(model, finiteElements);
    auto preconditioner = preconditioning::makePreconditioner(problem);
    const auto &gravity = preconditioner.structureComponent().gravityComponent();

    STATIC_CHECK(preconditioning::PreconditionerComponent<decltype(preconditioner)>);
    STATIC_CHECK(
        std::same_as<
            typename std::remove_cvref_t<decltype(gravity)>::MassBackend, preconditioning::backend::MatrixFreeChebyshev>
    );
    CHECK(problem.StateSize() == problem.EquationSize());
    CHECK(decltype(preconditioner)::borderValueArity == 1);
    CHECK(gravity.massInverseBackend().order == 5);
    CHECK(gravity.potentialSchurBackend().application.cycles == 3);
}

TEST_CASE(
    "Fixed Angular Momentum User API Generates Rotation And Its Composable Solver Border",
    "[user-api][fixed-angular-momentum][type]"
) {
    using namespace mean_field;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());

    auto model = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = 0.25}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        integral::FixedAngularMomentum({
            .Jtotal = dimensions::AngularMomentumValue{0.2},
            .axis = {0.0, 0.0, 2.0}
        })
    );

    auto problem = equilibrium::discretize(model, finiteElements);
    auto preconditioner = preconditioning::makePreconditioner(problem);
    using Problem = std::remove_cvref_t<decltype(problem)>;
    using Preconditioner = std::remove_cvref_t<decltype(preconditioner)>;

    STATIC_CHECK(Problem::hasFixedAngularMomentum);
    STATIC_CHECK_FALSE(Problem::hasFixedCentralDensity);
    STATIC_CHECK(PreparesWithGeneratedRotation<Problem>);
    STATIC_CHECK_FALSE(PreparesWithPrescribedRotation<Problem>);
    STATIC_CHECK(Problem::FormType::value_block_count == 7);
    STATIC_CHECK(Problem::FormType::residual_block_count == 7);
    STATIC_CHECK(preconditioning::PreconditionerComponent<Preconditioner>);
    STATIC_CHECK(Preconditioner::borderValueArity == 2);
    STATIC_CHECK(Preconditioner::borderResidualArity == 2);

    CHECK(problem.StateSize() == problem.EquationSize());
    CHECK(problem.StateSize() == problem.GetPhysicalOperator().Width() + 1);
    REQUIRE(problem.GetManifest().constraints().size() == 3);
    CHECK(problem.GetManifest().template specification<models::FixedAngularMomentum>().stableId ==
          "FixedAngularMomentum");
    CHECK(problem.GetManifest().valueBlocks().back().symbol == "Omega");
    CHECK(problem.GetManifest().residualBlocks().back().symbol == "R_J");
}

TEST_CASE(
    "Intermediate User API Selects A Coupled Stellar Factorization",
    "[user-api][intermediate]"
) {
    using namespace mean_field;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());

    auto model = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{1.0}})
    );
    auto problem  = equilibrium::discretize(model, finiteElements);

    auto material = preconditioning::materialSurfaceBlock(problem);
    auto gravity  = preconditioning::GravityFieldBlock(
        preconditioning::backend::Diagonal{},
        preconditioning::backend::HypreBoomerAMG(preconditioning::backend::FixedCycles{.cycles = 2}),
        preconditioning::GravityApproximateLDU{}
    );
    auto structure = preconditioning::stellarStructureBlock(
        problem, material, gravity, preconditioning::ApproximateStellarBlockLDU{}
    );
    auto preconditioner = preconditioning::specificationBorderBlock(problem, structure);

    STATIC_CHECK(preconditioning::PreconditionerComponent<decltype(preconditioner)>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename decltype(preconditioner)::BackendType>);
    CHECK(decltype(preconditioner)::borderValueArity == 2);
}

TEST_CASE(
    "Advanced User API Composes Explicit Blocks Backends And Application Modes",
    "[user-api][advanced]"
) {
    using namespace mean_field;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());

    auto model = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{1.0}})
    );
    auto problem  = equilibrium::discretize(model, finiteElements);

    auto material = preconditioning::materialSurfaceBlock(
        problem, preconditioning::backend::Diagonal{}, preconditioning::backend::Diagonal{},
        preconditioning::MaterialThenSurfaceTriangular{}, {.relativeFloor = 1.0e-10, .absoluteFloor = 1.0e-13}
    );
    auto gravity = preconditioning::GravityFieldBlock(
        preconditioning::backend::Diagonal{},
        preconditioning::backend::HypreBoomerAMG(
            preconditioning::backend::SolveToTolerance{.relativeTolerance = 1.0e-8, .maximumCycles = 20}
        ),
        preconditioning::GravityUpperTriangular{}
    );
    auto structure = preconditioning::stellarStructureBlock(
        problem, material, gravity, preconditioning::MaterialThenGravityTriangular{}
    );
    auto preconditioner =
        preconditioning::specificationBorderBlock(problem, structure, preconditioning::backend::DenseDirect{});

    STATIC_CHECK(preconditioning::PreconditionerComponent<decltype(preconditioner)>);
    STATIC_CHECK_FALSE(preconditioning::backend::ArnoldiAdmissible<typename decltype(preconditioner)::BackendType>);
    CHECK(preconditioner.structureComponent().materialSurfaceComponent().diagonalOptions().relativeFloor == 1.0e-10);
    CHECK(
        preconditioner.structureComponent().gravityComponent().potentialSchurBackend().application.maximumCycles == 20
    );
}