#include <array>
#include <cmath>
#include <cstdio>
#include <print>
#include <tuple>
#include <utility>
#include <vector>

#include "serif/equations/mass_normalization.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/solver/normalization/normalization.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace normalization = serif::solver::normalization;
    namespace dependencies = serif::utils::dependencies;
    namespace dimensions = serif::dimensions;

    /** @brief Convert illustrative physical coefficients and assembled residuals into one solve's fixed coordinates. */
    void show_coordinate_conversions(const normalization::DiagonalNormalization& coordinates, const std::array<double, 3>& physical_state, const std::array<double, 3>& physical_residual)
    {
        // These orders follow the layout below: [potential, surface potential level] and [Poisson balance, mass constraint].
        std::array<double, 3> normalized_state{};
        std::array<double, 3> normalized_residual{};
        coordinates.normalize_state(physical_state, normalized_state);
        coordinates.normalize_residual(physical_residual, normalized_residual);

        std::println("Physical potential coefficients: [{:.6g}, {:.6g}]; surface potential level: {:.6g}", physical_state[0], physical_state[1], physical_state[2]);
        std::println("Normalized potential coefficients: [{:.6g}, {:.6g}]; surface potential level: {:.6g}", normalized_state[0], normalized_state[1], normalized_state[2]);
        std::println("Physical assembled Poisson residual: [{:.6g}, {:.6g}]; mass residual: {:.6g}", physical_residual[0], physical_residual[1], physical_residual[2]);
        std::println("Normalized Poisson residual: [{:.6g}, {:.6g}]; mass residual: {:.6g}", normalized_residual[0], normalized_residual[1], normalized_residual[2]);
        std::println("Residual norm in these coordinates: {:.6g}", std::sqrt(coordinates.residual_norm_squared(physical_residual)));

        // Physical equations still consume physical coefficients. Convert a solver's coordinates back before evaluation.
        std::array<double, 3> recovered_state{};
        coordinates.denormalize_state(normalized_state, recovered_state);
        std::println("Recovered physical potential coefficients: [{:.6g}, {:.6g}]; surface potential level: {:.6g}", recovered_state[0], recovered_state[1], recovered_state[2]);
    }
}

/**
 * @brief Prepare, apply, reuse, and explicitly replace the normalization for a small selection of stellar coordinates.
 *
 * Potential and the surface potential level use GM/R, Poisson balance uses GM/R^3, and the mass constraint uses M.
 * PhysicalValueRule reads the unknown's Quantity; PhysicalResidualRule reads the equation's ResidualQuantity.
 * Those quantities need not agree, even when an equation closes a particular unknown.
 *
 * The illustrative metric contains two diagonal entries m_i = integral psi_i^2 dV_ref in the field's coefficient order.
 * A configured stellar discretization builds reference metrics in its computational graph automatically.
 * complete_sandbox demonstrates that path; this small example supplies a diagonal explicitly.
 * State factors are sqrt(m_i)/state_scale; assembled weak-form residual factors are 1/(sqrt(m_i)*residual_scale).
 * Algebraic coordinates use a unit metric. This is a diagonal metric, not the full mass-matrix norm.
 *
 * The example starts with supplied coefficients and assembled residuals; it does not assemble a complete stellar model.
 * Reference units must be consistent with those physical quantities and the reference geometry.
 */
int main()
{
    using Potential = serif::discretization::blocks::gravity::potential;
    using SurfaceLevel = serif::discretization::blocks::barotropic_constant::mass_normalization;
    using PoissonBalance = serif::equations::PrimalPoissonBalance;
    using MassConstraint = serif::equations::MassNormalization;
    using PotentialMetric = normalization::ReferenceVolumeMetric<Potential>;

    // The plan chooses physical scaling. The layout independently chooses the order and sizes of packed blocks.
    using Plan = normalization::NormalizationPlan<
        dependencies::Bindings<normalization::PhysicalValueRule<Potential>, normalization::PhysicalValueRule<SurfaceLevel>>,
        dependencies::Bindings<normalization::PhysicalResidualRule<PoissonBalance>, normalization::PhysicalResidualRule<MassConstraint>>>;
    const normalization::CoordinateLayout<dependencies::KeyList<Potential, SurfaceLevel>, dependencies::KeyList<PoissonBalance, MassConstraint>> layout{
        .state_sizes = {2, 1},
        .residual_sizes = {2, 1}
    };

    const dimensions::QuantityValue<dimensions::Length> reference_radius{2.0};
    const dimensions::QuantityValue<dimensions::GravitationalConstant> gravitational_constant{1.0};
    auto reference = normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{8.0}, reference_radius, gravitational_constant);
    if (!reference)
    {
        std::println(stderr, "Cannot prepare reference scales: {}", reference.error());
        return 1;
    }
    normalization::StellarReferenceScales reference_scales = *reference;
    const std::vector<double> reference_metric{4.0, 9.0};

    // Borrowed inputs must outlive the cache. Trial coefficients are absent: changing a candidate cannot change its units.
    const auto inputs = dependencies::make_dependency_context(
        dependencies::borrow<normalization::keys::Layout>(layout),
        dependencies::borrow<normalization::keys::Scales>(reference_scales),
        dependencies::borrow<PotentialMetric>(reference_metric));

    // Factories only declare work. Both coordinate roles share this ordinary dependency cache.
    auto computations = normalization::make_normalization_computations<Plan>(layout);
    auto cache = std::apply([&](auto... nodes)
    {
        return dependencies::make_computation_cache(inputs, std::move(nodes)...);
    }, std::move(computations));

    const auto prepared = cache.try_prepare(inputs);
    if (!prepared)
    {
        std::println(stderr, "Cannot prepare reference normalization: {}", prepared.error().error);
        return 1;
    }

    const normalization::StellarReferenceScales& scales = reference_scales;
    std::println("Reference mass: {:.6g}; radius: {:.6g}; G: {:.6g}", scales.mass.value(), scales.radius.value(), scales.gravitational_constant.value());
    std::println("Potential and surface-potential scale: {:.6g}; Poisson scale: {:.6g}; mass scale: {:.6g}", scales.specific_energy.value(), scales.inverse_time_squared.value(), scales.mass.value());

    // This owning copy fixes the coordinates for a solve, even if the cache later prepares another reference.
    const normalization::DiagonalNormalization solve_coordinates = cache.get<normalization::keys::Normalization>();
    const std::array physical_state{-8.0, -12.0, 6.0};
    const std::array physical_residual{0.2, -0.3, 0.4};
    show_coordinate_conversions(solve_coordinates, physical_state, physical_residual);

    const auto unchanged = cache.try_prepare(inputs);
    if (!unchanged)
    {
        std::println(stderr, "Cannot reuse reference normalization: {}", unchanged.error().error);
        return 1;
    }
    std::println("Unchanged reference: {} computations repeated.", unchanged->computed_count());

    // After this solve, explicitly choose a new reference mass for the next solve. The accepted copy remains valid.
    reference = normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{16.0}, reference_radius, gravitational_constant);
    if (!reference)
    {
        std::println(stderr, "Cannot prepare the next reference scales: {}", reference.error());
        return 1;
    }
    reference_scales = *reference;
    const auto updated = cache.try_prepare(inputs);
    if (!updated)
    {
        std::println(stderr, "Cannot prepare the next reference normalization: {}", updated.error().error);
        return 1;
    }
    const normalization::DiagonalNormalization next_solve_coordinates = cache.get<normalization::keys::Normalization>();
    std::array<double, 3> old_reference_state{};
    std::array<double, 3> new_reference_state{};
    solve_coordinates.normalize_state(physical_state, old_reference_state);
    next_solve_coordinates.normalize_state(physical_state, new_reference_state);
    std::println("After doubling reference mass, surface potential level in original coordinates: {:.6g}; next solve's coordinates: {:.6g}", old_reference_state[2], new_reference_state[2]);
    return 0;
}