#include <cmath>
#include <algorithm>
#include <array>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <numbers>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>
#include <format>
#include <memory>
#include <print>

#include "mfem/fem/intrules.hpp"
#include "serif/runtime.hpp"
#include "serif/backend/mfem/serialize/discretization.hpp"
#include "serif/backend/mfem/serialize/top.hpp"
#include "serif/utils/serialize/serialize.hpp"
#include "serif/equations/barotropic_density.hpp"
#include "serif/model/finite_elements.hpp"
#include "serif/solver/equilibrium.hpp"
#include "serif/solver/preconditioner/default_preconditioners.hpp"
#include "serif/backend/mfem/solver/preconditioner/inverse/sparse_direct.hpp"
#include "serif/solver/seed/state.hpp"
#include "serif/utils/laneemden/solution.hpp"
#include "polytropic_model.hpp"
#include "stellar_rayfem.hpp"
#include "serif/tests/solver/virial.hpp"
#include "libRayfem/init.hpp"

namespace {
    // Defaults: the n=2, J=0.1 star (Omega/Omega_K ~0.70, flattening ~0.22), the fastest rotator in the ESTER/TOP grid
    // that converges from the spherical seed. L=12 resolves its surface: all 33 reference modes match within 0.55 PLATO
    // sigma (L=8 leaves 3.65 sigma). --angular-momentum 0.05 --surface-degree 8 gives the original benchmark star.
    struct Arguments {
        bool visualize = true;
        bool visualize_requested = false;
        double angular_momentum = 0.1;
        double polytropic_index = 2.0;
        std::optional<double> polytropic_constant;
        bool surface_power_density = true;
        int surface_degree = 12;
        int mesh_refinement = 0;
        int surface_refinement = 1;
        int equatorial_refinement = 1;
        double equatorial_band_degrees = 10.0;
        int radial_surface_refinement = 3;
        /// Compactified infinity surface, in stellar radii. The vacuum map folds once the equatorial bulge reaches
        /// 1/(F - 1) stellar radii, so the fastest rotators need F below the library default of 4.
        double infinity_radius = 4.0;
        /// Converged archive (profile_solve --save-state) with this n, K and discretization to start Newton from.
        std::filesystem::path seed_state;
        /// Fixed volume quadrature order; empty uses SERiF's automatic rule (~35 on this mesh). Order 11 matches
        /// the automatic rule's ESTER/TOP agreement for the benchmark star at ~1/14 of the memory.
        std::optional<int> quadrature_order = 11;
        /// Potential-block inverse: sparse LU or three fixed BoomerAMG cycles (SERiF's default recipe). Empty (auto)
        /// uses LU when this build has a sparse_direct backend that works on this machine.
        std::optional<bool> potential_inverse_lu;
        std::array<double, 3> rotation_axis{0.0, 0.0, 1.0};
        /// Tidally locked companion on +x: its mass (G = M = 1) and separation; zero mass means a single star.
        double companion_mass = 0.0;
        double separation = 0.0;
        /// The star spins about its own axis at the rate J sets and feels only the companion's tide (asynchronous Roche).
        bool asynchronous = false;
        /// Continuation: solve first at separation_start, then step geometrically to separation, seeding each solve.
        double separation_start = 0.0;
        int separation_steps = 0;
        /// Converged binary surface radius on a (theta, phi) grid, for plotting.
        std::filesystem::path surface_table;
        bool mesh_only = false;
        int viewer_port = 5555;
        std::string viewer;
        std::filesystem::path top_export;
        serif::utils::serialize::TopExportOptions top_options;
    };

    Arguments parse_arguments(int argc, char** argv) {
        Arguments args;
        for (int i = 1; i < argc; ++i) {
            const std::string_view option(argv[i]);
            if (option == "--visualize") { args.visualize = args.visualize_requested = true; continue; }
            if (option == "--no-visualize") { args.visualize = false; continue; }
            if (option == "--surface-power-density") { args.surface_power_density = true; continue; }
            if (option == "--ordinary-density") { args.surface_power_density = false; continue; }
            if (option == "--mesh-only") { args.mesh_only = true; continue; }
            if (option == "--asynchronous") { args.asynchronous = true; continue; }
            if (option == "--rotation-axis") {
                for (double& component : args.rotation_axis) {
                    if (++i == argc) throw std::invalid_argument("--rotation-axis requires three components: X Y Z");
                    const std::string value(argv[i]);
                    std::size_t consumed = 0;
                    component = std::stod(value, &consumed);
                    if (consumed != value.size() || !std::isfinite(component))
                        throw std::invalid_argument("Rotation-axis components must be finite numbers");
                }
                continue;
            }
            if (++i == argc) throw std::invalid_argument("Missing value for " + std::string(option));
            const std::string value(argv[i]);
            std::size_t consumed = 0;
            if (option == "--export-top") { args.top_export = value; continue; }
            if (option == "--surface-table") { args.surface_table = value; continue; }
            if (option == "--seed-state") { args.seed_state = value; continue; }
            if (option == "--quadrature-order" && value == "auto") { args.quadrature_order.reset(); continue; }
            if (option == "--potential-inverse") {
                if (value != "lu" && value != "amg" && value != "auto") throw std::invalid_argument("--potential-inverse takes auto, lu or amg");
                if (value == "auto") args.potential_inverse_lu.reset();
                else args.potential_inverse_lu = value == "lu";
                continue;
            }
            if (option == "--viewer") { args.viewer = value; args.visualize = args.visualize_requested = true; continue; }
            if (option == "--angular-momentum") args.angular_momentum = std::stod(value, &consumed);
            else if (option == "--polytropic-index") args.polytropic_index = std::stod(value, &consumed);
            else if (option == "--K") args.polytropic_constant = std::stod(value, &consumed);
            else if (option == "--surface-degree") args.surface_degree = std::stoi(value, &consumed);
            else if (option == "--mesh-refinement") args.mesh_refinement = std::stoi(value, &consumed);
            else if (option == "--surface-refinement") args.surface_refinement = std::stoi(value, &consumed);
            else if (option == "--equatorial-refinement") args.equatorial_refinement = std::stoi(value, &consumed);
            else if (option == "--equatorial-band-degrees") args.equatorial_band_degrees = std::stod(value, &consumed);
            else if (option == "--infinity-radius") args.infinity_radius = std::stod(value, &consumed);
            else if (option == "--radial-surface-refinement") args.radial_surface_refinement = std::stoi(value, &consumed);
            else if (option == "--quadrature-order") args.quadrature_order = std::stoi(value, &consumed);
            else if (option == "--top-radial-points") args.top_options.radial_points = std::stoi(value, &consumed);
            else if (option == "--top-degree") args.top_options.maximum_degree = std::stoi(value, &consumed);
            else if (option == "--top-angular-points") args.top_options.angular_points = std::stoi(value, &consumed);
            else if (option == "--top-azimuthal-points") args.top_options.azimuthal_points = std::stoi(value, &consumed);
            else if (option == "--viewer-port") args.viewer_port = std::stoi(value, &consumed);
            else if (option == "--companion-mass") args.companion_mass = std::stod(value, &consumed);
            else if (option == "--separation") args.separation = std::stod(value, &consumed);
            else if (option == "--separation-start") args.separation_start = std::stod(value, &consumed);
            else if (option == "--separation-steps") args.separation_steps = std::stoi(value, &consumed);
            else throw std::invalid_argument("Unknown option: " + std::string(option));
            if (consumed != value.size()) throw std::invalid_argument("Invalid value for " + std::string(option));
        }
        if (!std::isfinite(args.polytropic_index) || args.polytropic_index < 1.0 || args.polytropic_index >= 3.0)
            throw std::invalid_argument("Use a polytropic index in [1,3) for this vacuum benchmark");
        if (!std::isfinite(args.angular_momentum) || args.surface_degree < 0 ||
            args.viewer_port < 1 || args.viewer_port > 65534)
            throw std::invalid_argument("Use finite angular momentum, nonnegative surface degree and viewer port 1..65534");
        if (args.polytropic_constant && (!std::isfinite(*args.polytropic_constant) || *args.polytropic_constant <= 0.0))
            throw std::invalid_argument("K must be finite and positive");
        if (args.quadrature_order && *args.quadrature_order < 0)
            throw std::invalid_argument("The volume quadrature order must be nonnegative, or auto");
        if (args.mesh_refinement < 0 || args.surface_refinement < 0 || args.equatorial_refinement < 0 || args.radial_surface_refinement < 0)
            throw std::invalid_argument("Mesh, surface, equatorial and radial surface refinement counts must be nonnegative");
        if (!std::isfinite(args.equatorial_band_degrees) || args.equatorial_band_degrees <= 0.0 || args.equatorial_band_degrees > 35.0)
            throw std::invalid_argument("The equatorial band half-angle must satisfy 0 < degrees <= 35");
        if (!std::isfinite(args.infinity_radius) || args.infinity_radius <= 1.0)
            throw std::invalid_argument("The infinity radius must be a finite number of stellar radii above 1");
        if (std::ranges::all_of(args.rotation_axis, [](double x) { return x == 0.0; }))
            throw std::invalid_argument("The rotation axis must be nonzero");
        if (!std::isfinite(args.companion_mass) || !std::isfinite(args.separation) || args.companion_mass < 0.0 ||
            (args.companion_mass > 0.0) != (args.separation > 0.0))
            throw std::invalid_argument("--companion-mass and --separation go together: a positive mass at a separation of several stellar radii");
        if (args.companion_mass > 0.0 && (args.rotation_axis != std::array<double, 3>{0.0, 0.0, 1.0} || !args.top_export.empty() || !args.seed_state.empty()))
            throw std::invalid_argument("A companion orbits about the default z axis; TOP export and seed states are single-star only");
        if (args.separation_steps < 0 || (args.separation_steps > 0 && !(args.companion_mass > 0.0 && std::isfinite(args.separation_start) && args.separation_start > 0.0)) ||
            (!args.surface_table.empty() && args.companion_mass <= 0.0))
            throw std::invalid_argument("--separation-steps needs a companion and a positive --separation-start; --surface-table needs a companion");
        if (args.asynchronous && args.companion_mass <= 0.0)
            throw std::invalid_argument("--asynchronous needs a companion");
        if (args.mesh_only && args.visualize_requested)
            throw std::invalid_argument("--mesh-only reports mesh statistics without solving; omit --visualize");
        if (args.mesh_only) args.visualize = false;
        return args;
    }

    /**
     * @brief K that gives the nonrotating polytrope of index n a radius of 1 when G = M = 1.
     * @note Lane-Emden scaling: R = xi_1 alpha and M = 4 pi alpha^3 rho_c (-xi_1^2 theta'(xi_1)), with
     * alpha^2 = (n+1) K rho_c^(1/n-1) / (4 pi G). For n = 1 this is 2/pi.
     */
    double unit_radius_polytropic_constant(const double n) {
        // K is a model constant, so integrate well below the seed solver's default step.
        const auto lane_emden = serif::utils::laneemden::solve(n, {.step = 1.0e-4, .maximum_steps = 1000000});
        if (!lane_emden || !lane_emden->surface)
            throw std::runtime_error("Cannot find the Lane-Emden surface for n = " + std::to_string(n));
        const double xi = lane_emden->surface->xi;
        const double alpha = 1.0 / xi;
        const double central_density = 1.0 / (4.0 * std::numbers::pi * alpha * alpha * alpha * (-xi * xi * lane_emden->surface->derivative));
        return 4.0 * std::numbers::pi * alpha * alpha * std::pow(central_density, 1.0 - 1.0 / n) / (n + 1.0);
    }

    /** @brief Resolve --potential-inverse auto against the compiled sparse_direct backend. */
    bool use_lu(const Arguments& args) {
        return args.potential_inverse_lu.value_or(serif::solver::preconditioner::backend::mfem::sparse_direct_available());
    }

    void print_options(const Arguments& args, const double K) {
        std::println("complete_sandbox options (G = M = 1; K gives a nonrotating radius of 1 unless --K is set):");
        std::println("  polytropic index n          {}", args.polytropic_index);
        std::println("  polytropic constant K       {:.17g}{}", K, args.polytropic_constant ? " (--K)" : " (computed from Lane-Emden)");
        if (args.companion_mass > 0.0)
            std::println("  companion                   mass {} at separation {} on +x, {}{}", args.companion_mass, args.separation,
                args.asynchronous ? "tide only (asynchronous)" : "tidally locked (Omega from force balance)",
                args.separation_steps ? std::format("; continued from {} in {} steps", args.separation_start, args.separation_steps) : std::string{});
        if (args.companion_mass <= 0.0 || args.asynchronous)
            std::println("  angular momentum J          {}", args.angular_momentum);
        std::println("  rotation axis               ({}, {}, {})", args.rotation_axis[0], args.rotation_axis[1], args.rotation_axis[2]);
        std::println("  surface harmonic degree     {}", args.surface_degree);
        std::println("  density representation      {}", args.surface_power_density ? "SurfacePower t^n on surface-adjacent cells" : "ordinary");
        std::println("  mesh refinement             {}", args.mesh_refinement);
        std::println("  surface refinement          {}", args.surface_refinement);
        std::println("  equatorial refinement       {} (band half-angle {} degrees)", args.equatorial_refinement, args.equatorial_band_degrees);
        std::println("  radial surface refinement   {}", args.radial_surface_refinement);
        std::println("  volume quadrature           {}", args.quadrature_order
            ? std::format("fixed order {} ({} points per hexahedron)", *args.quadrature_order,
                ::mfem::IntRules.Get(::mfem::Geometry::CUBE, *args.quadrature_order).GetNPoints())
            : std::string{"automatic (SERiF's per-form rule)"});
        namespace inverse_backend = serif::solver::preconditioner::backend::mfem;
        std::println("  potential inverse           {}{}", use_lu(args) ? std::format("sparse LU ({})", inverse_backend::sparse_direct_backend()) : std::string{"3 fixed BoomerAMG cycles"},
            args.potential_inverse_lu ? "" : " (auto)");
        std::println("  visualization               {}", args.visualize ? std::format("RayFEM on port {}", args.viewer_port) : "off");
        std::println("  TOP export                  {}", args.top_export.empty() ? "off" : std::format("{} ({} radial, degree {}, {} angular, {} azimuthal)",
            args.top_export.string(), args.top_options.radial_points, args.top_options.maximum_degree,
            args.top_options.angular_points, args.top_options.azimuthal_points));
        std::fflush(stdout);
    }
}

/**
 * @brief One RayFEM connection shared by every continuation step, so all their frames form one sequence.
 * @note Every step uses the same reference mesh, so it is sent once and frame indices keep increasing across steps.
 */
struct LiveStream {
    RayFEM::Client client;
    std::uint32_t next_frame = 0;
    bool mesh_sent = false;
    bool enabled = true;
    bool continuation = false;
    int step = 0;

    /** @brief Start or reuse the viewer on the requested port. */
    static std::unique_ptr<LiveStream> open(const Arguments& args) {
        RayFEM::WindowParams window{.port = args.viewer_port, .executable = args.viewer};
    #ifdef SERIF_RAYFEM_EXECUTABLE
        const char* viewer_override = std::getenv("RAYFEM_EXECUTABLE");
        if (window.executable.empty() && (!viewer_override || !*viewer_override))
            window.executable = SERIF_RAYFEM_EXECUTABLE;
    #endif
        const auto viewer = RayFEM::init(window);
        if (!viewer) throw std::runtime_error("Cannot initialize RayFEM: " + viewer.error);
        if (viewer.status == RayFEM::InitStatus::Started)
            std::println("Started RayFEM: {}", viewer.executable);
        else
            std::println("Reusing RayFEM on port {}", args.viewer_port);
        auto stream = std::unique_ptr<LiveStream>(new LiveStream{RayFEM::Client("127.0.0.1", args.viewer_port)});
        // Frames go out from a background thread: a send returns once the frame is copied,
        // so Newton never waits for the viewer or the network.
        stream->client.startQueue();
        return stream;
    }
};

/**
 * @brief Build, solve and report one star; Binary adds the tidally locked companion and changes the model's type.
 * @note seed (same discretization) replaces the spherical start; converged receives the solved coordinates.
 * With a stream, accepted iterates are sent to RayFEM, tagged with the stream's continuation step and the separation.
 */
template <bool Binary>
int solve_star(const Arguments& args, const double K, std::vector<double>* converged, const std::vector<double>* seed,
    LiveStream* stream) {
    namespace model = serif::model;
    namespace equations = serif::equations;
    namespace dimensions = serif::dimensions;
    namespace solver = serif::solver;
    namespace normalization = solver::normalization;

    // 1. Describe the star with M=G=1, the requested n and J, and K (by default a nonrotating radius of 1).
    // Deferred equations read the same owned prescriptions used by continuation.
    // A tidally locked companion sits on +x: rotation is about the barycenter, and Omega's row fixes the star's center of mass.
    const double companion_mass = args.companion_mass, separation = args.separation;
    const auto stellar_equations = [=](const auto& physics) {
        const auto eos = model::borrow_eos(physics.equation_of_state());
        const auto& rotation = physics.template specification<model::FixedAngularMomentum>();
        const double mass = physics.template specification<model::FixedTotalMass>().target().value();
        std::optional<equations::PointMassCompanion> companion;
        auto center = rotation.center();  // the star stays at the specification's center; a binary rotates about the barycenter
        if (companion_mass > 0.0) companion.emplace(companion_mass, equations::RigidRotation::Vector{separation, 0.0, 0.0}, !Binary);
        if constexpr (Binary) center[0] += companion_mass * separation / (mass + companion_mass);
        auto omega_row = [&] {
            if constexpr (Binary) return equations::BarycenterConstraint<0, serif::discretization::blocks::rotation::angular_velocity>{mass};
            else return equations::AngularMomentumConstraint{rotation.target().value(), rotation.axis(), rotation.center()};
        }();
        return model::EquilibriumEquations{
            .interior_equations = equations::EquationRegistry{
                equations::BarotropicDensityBalance{eos, dimensions::PressureValue{0.0}, rotation.axis(), center, companion},
                equations::PrimalPoissonBalance{4.0 * std::numbers::pi},
                equations::MassNormalization{mass},
                std::move(omega_row)
            },
            .shape_force = equations::SurfacePotentialBalance{rotation.axis(), center, companion},
            .surface_condition = serif::geometry::surface::Isobaric{dimensions::PressureValue{0.0}}
        };
    };

    const double n = args.polytropic_index;
    const auto stellar_model = model::make_stellar_model(serif::eos::models::Polytrope{n, K}, stellar_equations,
        model::FixedTotalMass{dimensions::MassValue{1.0}}, model::FixedAngularMomentum{model::FixedAngularMomentum::Parameters{
            .angular_momentum = dimensions::AngularMomentumValue{Binary ? 0.0 : args.angular_momentum}, .axis = args.rotation_axis}});
    const auto& rotation = stellar_model.template specification<model::FixedAngularMomentum>();

    // 2. Obtain a radial initial guess while preserving the EOS and the target mass.
    const auto radial = solver::seed::make_polytropic_seed(stellar_model);
    if (!radial) {
        std::println(stderr, "Cannot prepare the radial star: {}", radial.error());
        return 1;
    }

    // 3. Derive field spaces from their term types and own all numerical resources.
    // Surface harmonics remain an explicit approximation choice.
    namespace deformation = serif::geometry::deformation;
    namespace reconstruction = serif::discretization::reconstruction;
    using Density = serif::discretization::blocks::density::mass;
    const deformation::DomainDeformation<deformation::SphericalHarmonicSurface> moving_boundary{
        .surface = {.maximum_degree = args.surface_degree}, .stellar_interior = {.radial_power = 2.0}, .vacuum = {}};
    const auto finite_elements = model::make_stellar_discretization(stellar_model,
        {.radius = radial->scales.radius, .infinity_radius_fraction = args.infinity_radius, .refinement = args.mesh_refinement,
            .surface_refinement = {.levels = args.surface_refinement, .equatorial_levels = args.equatorial_refinement,
                .equatorial_half_angle_degrees = args.equatorial_band_degrees, .radial_levels = args.radial_surface_refinement,
                .axis = rotation.axis()}}, moving_boundary,
        reconstruction::SurfacePower<Density>{.density_exponent = args.surface_power_density ? n : 1.0});
    const auto& mesh = *finite_elements.owner->resources.meshes.front();
    using Schema = serif::discretization::domain::schema::CoreEnvelopeVacuumDomainSchema;
    const int surface_attribute = Schema::boundary_id<serif::discretization::domain::StellarSurfaceBoundary>();
    int surface_faces = 0;
    for (int boundary = 0; boundary < mesh.GetNBE(); ++boundary)
        surface_faces += mesh.GetBdrAttribute(boundary) == surface_attribute;
    std::println("Mesh: {} elements, {} stellar surface faces; rotation axis = ({}, {}, {})",
        mesh.GetNE(), surface_faces, rotation.axis()[0], rotation.axis()[1], rotation.axis()[2]);
    if (args.mesh_only) return 0;
    auto initial_state = solver::seed::make_polytropic_state(stellar_model, finite_elements, *radial);
    if (!initial_state) {
        std::println(stderr, "Cannot project the radial star: {}", initial_state.error());
        return 1;
    }
    // Start the orbit at Kepler's rate for two point masses.
    const double kepler_omega = std::sqrt((1.0 + companion_mass) / (separation * separation * separation));
    if constexpr (Binary) initial_state->template get<serif::discretization::blocks::rotation::angular_velocity>()[0] = kepler_omega;
    if (seed) std::ranges::copy(*seed, initial_state->values().begin());

    if (!args.seed_state.empty()) {
        // Continuation: replace the spherical seed with a converged star of the same n, K and discretization.
        namespace archive = serif::utils::serialize;
        using State = std::remove_cvref_t<decltype(*initial_state)>;
        const auto saved = archive::Reader::open(args.seed_state);
        if (!saved) throw std::runtime_error(saved.error().message);
        const auto loaded = saved->read<State>("state");
        const auto old_n = saved->read<double>("n"), old_K = saved->read<double>("K"), old_J = saved->read<double>("J");
        const auto converged = saved->read<bool>("converged");
        const auto geometry = saved->node("finite_elements");
        if (!loaded || !old_n || !old_K || !old_J || !converged || !*converged || *old_n != n || *old_K != K || !geometry ||
            *geometry != serif::backend::mfem::serialize::save_discretization(finite_elements) || loaded->layout() != initial_state->layout())
            throw std::runtime_error("--seed-state must be a converged archive with the same n, K, discretization and layout");
        std::ranges::copy(loaded->values(), initial_state->values().begin());
        std::println("Seeded from {} (converged at J = {})", args.seed_state.string(), *old_J);
    }

    // Deduced here: clang does not deduce these class templates inside the generic lambda below.
    const auto normalization_options = normalization::NormalizationOptions{
        .scheme = normalization::LaneEmdenApproximation{.n = n, .central_density = radial->scales.central_density}};
    // 4-6. Solve and report with the chosen potential inverse; each choice is its own preconditioner type.
    const auto solve_and_report = [&](const auto potential_inverse) -> int {
        // 4. Freeze sensible numerical scales independently of the physical prescriptions.
        auto configuration = model::make_discretization(model::DiscretizationOptions{
            .model = stellar_model,
            .finite_elements = finite_elements,
            .normalization = normalization::make_normalization(normalization_options),
            .preconditioner = solver::preconditioner::make_primal_preconditioner(initial_state->layout(), potential_inverse)
        });

        // 5. Select numerical work limits and observe accepted stars without managing evaluators or caches.
        const auto observe = [](const solver::SolveEvent& event, const auto&, const auto&) {
            if (event.kind == solver::EventKind::initial) {
                std::println("Initial normalized residual = {}", event.metric->residual_norm);
            } else if (event.kind == solver::EventKind::accepted) {
                std::println("Newton step {}: normalized residual = {}, accepted step = {}", event.iteration + 1, event.metric->residual_norm, event.step);
            }
            std::fflush(stdout);
        };
        const solver::SolverOptions options{
            .method = {.relative_tolerance = 1.0e-8, .absolute_tolerance = 1.0e-10, .maximum_iterations = 30,
                .linear_solve = {.relative_tolerance = 0.03, .maximum_iterations = 200}},
            .linear_solver = {.restart_length = 160}
        };
        solver::VolumeOptions volume_options{};
        if (args.quadrature_order) {
            namespace quadrature = serif::discretization::quadrature;
            quadrature::RuleSet rules;
            rules.set_role_control(quadrature::QuadratureRole::discretization, {.fixed_order = *args.quadrature_order});
            volume_options.quadrature_policy = quadrature::Policy{std::move(rules)};
        }
        const auto solution = [&] {
            if (!stream)
                return solver::solve_equilibrium(std::move(configuration), std::move(*initial_state), options, observe, volume_options);

            auto problem = solver::backend::mfem::make_observed_stellar_problem(std::move(configuration), *initial_state,
                options.strategies, volume_options, serif::utils::dependencies::DependencyContext<>{},
                serif::utils::dependencies::NoGraphObservers{});
            serif::visualization::backend::mfem::StellarFrame frame(*problem.binding().discretization);
            auto& client = stream->client;
            constexpr std::uint64_t mesh_id = 1;
            // Frames are Newton iterates of one equilibrium, not moments in time.
            RayFEM::FrameSchema schema;
            schema.values = {
                {"Newton step", RayFEM::ValueType::Integer, true, "", "Accepted Newton steps so far; 0 is the initial state"},
                {"residual", RayFEM::ValueType::Real, false, "", "Normalized nonlinear residual"},
                {"step length", RayFEM::ValueType::Real, false, "", "Accepted damped Newton step length (1 is a full step)"},
                {"Omega", RayFEM::ValueType::Real, false, "", "Angular velocity of the rigid rotation"},
                {"virial ratio", RayFEM::ValueType::Real, false, "", "(2T + 3 integral of P)/|W|; 1 in equilibrium"},
                {"virial error", RayFEM::ValueType::Real, false, "", "(2T + W + 3 integral of P)/|W| at the finest quadrature; 0 in equilibrium"},
                {"mass", RayFEM::ValueType::Real, false, "M", "Integrated mass"},
                {"virial", RayFEM::ValueType::Text, false, "", "Why the vacuum virial could not be measured"}};
            schema.tags = {{"initial", "The seeded state before any Newton step"},
                           {"resumed", "A state resumed from a checkpoint"},
                           {"accepted", "An accepted Newton iterate"},
                           {"converged", "Meets the solver's residual tolerances"}};
            schema.has_time = false;
            schema.primary = "Newton step";
            if (stream->continuation) {
                // Newton steps restart in each continuation step, so frames are ordered by step.
                schema.values.insert(schema.values.begin(), {
                    {"continuation step", RayFEM::ValueType::Integer, true, "", "Continuation step; 0 is the widest separation"},
                    {"separation", RayFEM::ValueType::Real, true, "R0", "Orbital separation of the companion"}});
                schema.primary = "continuation step";
            }
            // RayFEM orients the view and measures flattening about this axis; a binary rotates about the barycenter.
            auto axis_origin = rotation.center();
            if constexpr (Binary) {
                const double mass = stellar_model.template specification<model::FixedTotalMass>().target().value();
                axis_origin[0] += companion_mass * separation / (mass + companion_mass);
            }
            const auto& spin = rotation.axis();
            schema.axes = {{"rotation", {spin[0], spin[1], spin[2]}, {axis_origin[0], axis_origin[1], axis_origin[2]},
                            Binary ? "Orbital rotation axis through the barycenter" : "Rigid rotation axis"}};
            // A locked binary's barycenter moves with the separation, so each step re-registers the schema.
            if (stream->enabled && !client.registerSchema(schema)) throw std::runtime_error(client.getLastError());
            if (stream->enabled && !stream->mesh_sent) {
                if (!client.sendMesh(mesh_id, frame.mesh())) throw std::runtime_error(client.getLastError());
                stream->mesh_sent = true;
            }
            double initial_residual = 0;
            const auto live_observe = [&](const solver::SolveEvent& event, const auto& state, const auto& system) {
                observe(event, state, system);
                if (!stream->enabled || (event.kind != solver::EventKind::initial && event.kind != solver::EventKind::resumed &&
                                   event.kind != solver::EventKind::accepted)) return;
                try {
                    // StandardNewton has just evaluated this state. The cache is unsafe at finished
                    // after a rejected trial, so only initial/resumed/accepted events are streamed.
                    frame.update(state, problem.evaluator().cache());
                    const std::uint32_t frame_index = stream->next_frame;
                    RayFEM::FrameMetadata metadata{.frame_index = frame_index, .sim_time = 0.0};
                    const double residual = event.metric->residual_norm;
                    if (event.kind != solver::EventKind::accepted) initial_residual = residual;
                    const int newton_step = event.kind == solver::EventKind::accepted ? event.iteration + 1 : 0;
                    if (stream->continuation)
                        std::snprintf(metadata.label, sizeof(metadata.label), "step %d a=%.4g Newton %d |R|=%.3g",
                            stream->step, separation, newton_step, residual);
                    else
                        std::snprintf(metadata.label, sizeof(metadata.label), "n=%.3g J=%.3g Newton %u |R|=%.3g",
                            n, args.angular_momentum, frame_index, residual);
                    const double omega = state.template get<serif::discretization::blocks::rotation::angular_velocity>()[0];
                    metadata.values = {{"Newton step", std::int64_t{newton_step}}, {"residual", residual}, {"Omega", omega}};
                    if (stream->continuation) {
                        metadata.values.push_back({"continuation step", std::int64_t{stream->step}});
                        metadata.values.push_back({"separation", separation});
                    }
                    if (event.kind == solver::EventKind::accepted) metadata.values.push_back({"step length", event.step});
                    // The physical check, independent of the algebraic residual.
                    const auto& axis = rotation.axis();
                    const auto virial = companion_mass > 0.0 ? std::expected<serif::tests::solver::VacuumVirialAssessment, std::string>{
                        std::unexpect, "the vacuum virial omits the companion's potential"} :
                        serif::tests::solver::vacuum_virial(*problem.binding().discretization, state, stellar_model.equation_of_state(),
                            equations::RigidRotation{{omega*axis[0], omega*axis[1], omega*axis[2]}, rotation.center()});
                    if (virial) {
                        const auto& sample = virial->samples.back();
                        metadata.values.push_back({"virial ratio", (2.0*sample.kinetic_energy + 3.0*sample.pressure_integral)/std::abs(sample.potential_energy)});
                        metadata.values.push_back({"virial error", sample.energy_virial()});
                        metadata.values.push_back({"mass", sample.mass});
                    } else metadata.values.push_back({"virial", virial.error().substr(0, 1000)});
                    metadata.tags = {event.kind == solver::EventKind::initial ? "initial" : event.kind == solver::EventKind::resumed ? "resumed" : "accepted"};
                    if (residual <= options.method.absolute_tolerance || residual <= options.method.relative_tolerance*initial_residual)
                        metadata.tags.push_back("converged");
                    const auto send_start = std::chrono::steady_clock::now();
                    if (!client.sendFrame(mesh_id, frame.fields(), metadata, &frame.displacement()))
                        throw std::runtime_error(client.getLastError());
                    const std::chrono::duration<double, std::milli> send_time = std::chrono::steady_clock::now() - send_start;
                    std::println("RayFEM frame {} (queued in {:.2f} ms): interior elements {}, max nodal displacement {}, virial {}",
                        frame_index, send_time.count(), frame.mesh().GetNE(), frame.displacement().Normlinf(),
                        virial ? std::format("ratio {:.12g} error {:.3g}", (2.0*virial->samples.back().kinetic_energy +
                            3.0*virial->samples.back().pressure_integral)/std::abs(virial->samples.back().potential_energy),
                            virial->samples.back().energy_virial()) : "unavailable: " + virial.error());
                    ++stream->next_frame;
                    std::fflush(stdout);
                } catch (const std::exception& error) {
                    stream->enabled = false;
                    std::println(stderr, "Live visualization disabled: {}", error.what());
                }
            };
            auto result = solver::solve_newton(problem, std::move(*initial_state), options, live_observe);
            // The last frames, the converged state among them, may still be queued.
            if (!client.flush(60000)) std::println(stderr, "RayFEM: queued frames were still unsent after 60 s");
            if (const auto queue = client.queueStats(); queue.failed)
                std::println(stderr, "RayFEM: {} queued packets were not sent", queue.failed);
            return result;
        }();
        if (!solution.converged()) {
            std::println(stderr, "Stellar solve stopped: {}", solution.message());
            if (solution.diagnostics().last_linear_solve) {
                const auto& linear = *solution.diagnostics().last_linear_solve;
                std::println(stderr, "Last linear solve: {} iterations, true relative residual {}", linear.iterations, linear.relative_true_residual_norm);
            }
            return 1;
        }

        // The result retains the physical definition, accepted coefficients and their matching system.
        using Omega = serif::discretization::blocks::rotation::angular_velocity;
        std::println("Converged after {} steps: Omega = {}, normalized residual = {}", solution.diagnostics().accepted_iterations,
            solution.state().template get<Omega>()[0], *solution.diagnostics().final_residual_norm);
        if (companion_mass > 0.0) {
            // Radii from the star's center toward the companion (+x), away from it, along the orbit (+y) and along the axis.
            namespace deformation_backend = serif::geometry::deformation::backend::mfem;
            const auto& state = solution.state();
            const auto inputs = deformation_backend::make_deformation_inputs(finite_elements.deformation.model,
                finite_elements.deformation.reference, state.template get<serif::discretization::blocks::surface_deformation::coefficients>());
            using Surface = std::remove_cvref_t<decltype(finite_elements.deformation.model.surface)>;
            auto cache = serif::utils::dependencies::make_computation_cache<serif::geometry::deformation::DeformationFailure>(
                deformation_backend::make_deformation_graph<Surface>(), inputs);
            if (!cache.try_prepare(inputs, serif::utils::dependencies::KeyList<serif::geometry::deformation::keys::VolumeDisplacement>{}))
                throw std::runtime_error("Cannot reconstruct the converged surface");
            serif::backend::mfem::serialize::detail::TopModelSampler sampler(finite_elements.template binding<Density>(),
                state.template get<Density>(), finite_elements.deformation.reference,
                cache.template get<serif::geometry::deformation::keys::VolumeDisplacement>(), rotation.axis(), {0.0, 0.0, 0.0});
            const auto radius = [&](double mu, double phi) { return sampler.surface_radius(mu, phi).value(); };
            if constexpr (Binary)
                std::println("Orbit: Omega / Omega_Kepler(point masses) = {:.9f}", solution.state().template get<Omega>()[0] / kepler_omega);
            else
                std::println("Spin: Omega_spin / Omega_orbit(Kepler, point masses) = {:.6f}", solution.state().template get<Omega>()[0] / kepler_omega);
            std::println("Radii: toward companion {:.6f}, away {:.6f}, along orbit {:.6f}, polar {:.6f}",
                radius(0.0, 0.0), radius(0.0, std::numbers::pi), radius(0.0, 0.5 * std::numbers::pi), radius(1.0, 0.0));
            if (!args.surface_table.empty()) {
                std::ofstream table(args.surface_table);
                table << "# theta_deg phi_deg radius (star center at the origin, companion at +x, separation " << separation << ")\n";
                for (int i = 0; i <= 90; ++i) {
                    for (int j = 0; j < 90; ++j) {
                        const double theta = 2.0 * i, phi = 4.0 * j;
                        table << theta << ' ' << phi << ' ' << radius(std::cos(theta * std::numbers::pi / 180.0), phi * std::numbers::pi / 180.0) << '\n';
                    }
                }
                std::println("Surface table written to {}", args.surface_table.string());
            }
        }
        if (converged) converged->assign(solution.state().values().begin(), solution.state().values().end());
        if (!args.top_export.empty()) {
            const auto exported = serif::utils::serialize::export_top(args.top_export, finite_elements, solution, args.top_options);
            if (!exported) {
                std::println(stderr, "TOP export failed at {}: {}", exported.error().path, exported.error().message);
                return 1;
            }
            std::println("TOP model written to {}: {} radial points, degree {}, angular frequency unit {}",
                exported->directory.string(), exported->radial_points, exported->maximum_degree, exported->angular_frequency_unit);
            std::println("TOP projection: surface error {}, enthalpy angular truncation {}, surface symmetry error {}, enthalpy symmetry error {}",
                exported->maximum_relative_surface_projection_error, exported->maximum_relative_enthalpy_projection_error,
                exported->maximum_relative_surface_symmetry_error,
                exported->maximum_relative_enthalpy_symmetry_error);
        }
        return 0;
    };
    if (use_lu(args)) return solve_and_report(solver::preconditioner::strategies::inverse::SparseDirect{});
    return solve_and_report(solver::preconditioner::strategies::inverse::FixedAMG{.cycles = 3});
}

/** @brief Solve one star, or continue a binary inward in separation, seeding each solve with the previous star. */
template <bool Binary>
int run(const Arguments& args, const double K) {
    const auto stream = args.visualize ? LiveStream::open(args) : nullptr;
    if (args.separation_steps > 0) {
        if (stream) stream->continuation = true;
        std::vector<double> state;
        for (int step = 0; step <= args.separation_steps; ++step) {
            Arguments stage = args;
            stage.separation = args.separation_start * std::pow(args.separation / args.separation_start, static_cast<double>(step) / args.separation_steps);
            if (step < args.separation_steps) stage.surface_table.clear();
            if (stream) stream->step = step;
            std::println("Continuation step {} of {}: separation {}", step, args.separation_steps, stage.separation);
            const std::vector<double> seed = state;
            if (const int status = solve_star<Binary>(stage, K, &state, step ? &seed : nullptr, stream.get()); status != 0) return status;
        }
        return 0;
    }
    return solve_star<Binary>(args, K, nullptr, nullptr, stream.get());
}

/** @brief Show the complete route from physical equations to an owning converged stellar solution. */
int main(int argc, char** argv) try {
    if (argc == 2 && std::string_view(argv[1]) == "--help") {
        std::println("Usage: complete_sandbox [--angular-momentum J] [--surface-degree L]\n"
                     "                        [--polytropic-index N] [--K K] [--surface-power-density | --ordinary-density]\n"
                     "                        [--mesh-refinement N] [--surface-refinement N]\n"
                     "                        [--equatorial-refinement N] [--equatorial-band-degrees D]\n"
                     "                        [--radial-surface-refinement N] [--quadrature-order N | auto]\n"
                     "                        [--potential-inverse auto | lu | amg]\n"
                     "                        [--rotation-axis X Y Z] [--mesh-only] [--infinity-radius F] [--seed-state FILE]\n"
                     "                        [--visualize | --no-visualize] [--viewer PATH] [--viewer-port PORT]\n"
                     "                        [--export-top DIRECTORY] [--top-radial-points N] [--top-degree L]\n"
                     "                        [--top-angular-points N] [--top-azimuthal-points N]\n"
                     "                        [--companion-mass MC --separation A [--separation-start A0 --separation-steps N]\n"
                     "                        [--surface-table FILE]]\n"
                     "Defaults: the n=2, J=0.1 rotator (Omega/Omega_K ~0.70) compared with ESTER/TOP: n=2, J=0.1, L=12, SurfacePower density,\n"
                     "surface refinement 1, equatorial refinement 1 with a 10 degree band, radial surface refinement 3,\n"
                     "axis=(0,0,1), streamed live to RayFEM. G=M=1 and, unless --K is given, K makes the nonrotating radius 1.\n"
                     "--angular-momentum 0.05 --surface-degree 8 gives the slower n=2, J=0.05 benchmark star.\n"
                     "--surface-power-density (default) selects t^n reconstruction on surface-adjacent density elements;\n"
                     "--ordinary-density selects the plain finite-element density.\n"
                     "Radial surface refinement halves only the stellar-side surface shell, radially, per level.\n"
                     "--quadrature-order fixes the volume quadrature order (default 11); auto uses SERiF's automatic rule.\n"
                     "--potential-inverse lu factors the potential block exactly (build option sparse_direct); amg uses SERiF's\n"
                     "3-cycle BoomerAMG; auto (default) uses lu when this build and machine support it, else amg.\n"
                     "Surface degree controls harmonics; refinement controls mesh cells.\n"
                     "Surface refinement adds isotropic interface layers; equatorial refinement adds latitude splits.\n"
                     "Equatorial rings conservatively cover the requested band (0 < D <= 35 degrees).\n"
                     "--infinity-radius places compactified infinity at F stellar radii (default 4); the vacuum map folds once the\n"
                     "equatorial bulge reaches 1/(F-1), so stars flatter than ~0.33 need F < 4.\n"
                     "--seed-state starts Newton from a converged profile_solve --save-state archive with the same n, K and\n"
                     "discretization (any J): continuation for fast rotators that do not converge from the spherical seed.\n"
                     "--mesh-only prints mesh counts and exits (no viewer). Viewer defaults: enabled, port 5555.\n"
                     "Live frames show initial and accepted Newton states, not physical time.\n"
                     "--export-top writes a converged polytrope to a new directory for TOP.\n"
                     "--companion-mass MC --separation A adds a tidally locked point-mass companion on +x: rotation is about the\n"
                     "barycenter and Omega follows from force balance (the star's center of mass stays at the origin); J is unused.\n"
                     "--separation-start A0 --separation-steps N solves first at A0 (from the spherical seed), then at N separations\n"
                     "stepping geometrically to A, each seeded by the previous star; every step streams to one RayFEM sequence,\n"
                     "tagged with its continuation step and separation. Close or massive\n"
                     "companions need this: from a sphere Newton stalls well inside the Roche lobe. Use --infinity-radius 3 near it.\n"
                     "--asynchronous instead spins the star about its own axis at the rate --angular-momentum sets, with only the\n"
                     "companion's tide acting (its uniform pull removed): the asynchronous Roche approximation.\n"
                     "--surface-table FILE writes the converged binary's surface radius on a (theta, phi) grid (phi = 0 faces the companion).\n"
                     "TOP defaults: 64 radial points, maximum even degree 16, 48 angular and 8 azimuthal samples.");
        return 0;
    }
    const auto args = parse_arguments(argc, argv);
    serif::initialize();
    const double K = args.polytropic_constant.value_or(unit_radius_polytropic_constant(args.polytropic_index));
    print_options(args, K);
    return args.companion_mass > 0.0 && !args.asynchronous ? run<true>(args, K) : run<false>(args, K);
} catch (const std::exception& error) {
    std::println(stderr, "{}", error.what());
    return 1;
}