#include <print>
#include <vector>
#include "serif/model/finite_elements.hpp"
#include "serif/types.hpp"
#include "serif/geometry/deformation/finite_elements.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"
#include "polytropic_model.hpp"

/**
 * @brief Deform a stellar surface, extend its motion, and reuse the reference calculations for a second shape.
 *
 * The physical model chooses radial displacement lengths in normalized real harmonics, a quadratic interior response,
 * and fixed reference infinity. The owning discretization binds those choices to paired physical and logical meshes. Deformation reference computations wait until try_prepare. Coefficients are borrowed; finite_elements and shape_coefficients must outlive the inputs.
 * Changing model options requires rebuilding the lightweight inputs; changing the coefficient vector does not.
 *
 * The same computations can be appended to an equation graph with make_deformation_computations. They do not own another
 * preparation cache. SurfaceMetric is an optional separate reference calculation, independent of the volume extension.
 *
 * Low-level derivative actions use spans and an explicitly allocated workspace. They do not assemble a composite Jacobian.
 * For nonlinear surface or extension models the same interface includes the derivative of the force pullback.
 * This MFEM reference compiler currently supports conforming serial three-dimensional nodal spaces with byNODES displacement.
 */
int main() {
    namespace deformation = serif::geometry::deformation;
    namespace dependencies = serif::utils::dependencies;
    namespace keys = deformation::keys;

    const deformation::DomainDeformation model{
        .surface = deformation::SphericalHarmonicSurface{.maximum_degree = 2},
        .stellar_interior = deformation::PowerLawRadialExtension{.radial_power = 2.0},
        .vacuum = deformation::FixedInfinityRadialExtension{}
    };
    std::vector<double> shape_coefficients(9, 0.0);
    shape_coefficients[0] = 0.02;
    const auto finite_elements = serif::model::make_stellar_discretization(serif::examples::make_polytropic_model(),
        {.geometry_order = 2}, model);
    const auto inputs = deformation::make_deformation_inputs(model, finite_elements.deformation.reference, shape_coefficients);
    auto cache = dependencies::make_computation_cache<deformation::DeformationFailure>(
        deformation::make_deformation_graph<deformation::SphericalHarmonicSurface>(), inputs);
    const auto prepared = cache.try_prepare(inputs);
    if (!prepared) {
        std::println(stderr, "Stellar deformation failed: {}", prepared.error().error.message);
        return 1;
    }
    std::println("{} shape coefficients produce {} volume displacement coefficients.", shape_coefficients.size(), cache.get<keys::VolumeDisplacement>().size());

    const auto accepted_shape = cache.get<keys::Linearization>();
    auto workspace = deformation::make_deformation_workspace(*accepted_shape.map);
    std::vector<double> shape_direction(shape_coefficients.size(), 0.0);
    // Increasing degree, then signed order places the axisymmetric quadrupole Y(2, 0) at coordinate six.
    shape_direction[6] = 1.0;
    serif::Vector displacement_direction(static_cast<int>(accepted_shape.map->output_size()));
    if (const auto action = deformation::apply_jacobian(accepted_shape, shape_direction, displacement_direction, workspace); !action) {
        std::println(stderr, "Shape tangent failed: {}", action.error().message);
        return 1;
    }

    shape_coefficients[6] = 0.01;
    const auto updated = cache.try_prepare(inputs);
    if (!updated) {
        std::println(stderr, "Updated stellar deformation failed: {}", updated.error().error.message);
        return 1;
    }
    const bool reference_reused = updated->status<keys::RadialInterpolation>() == dependencies::PreparationStatus::reused;
    std::println("Reference radial interpolation reused after changing shape: {}.", reference_reused);
    return reference_reused ? 0 : 1;
}