#include <cstdio>
#include <memory>
#include <print>
#include <utility>

#include "serif/types.hpp"
#include "serif/operators/linearized_block.hpp"
#include "serif/discretization/blocks/fields.hpp"
#include "serif/solver/preconditioner/preconditioner.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace {
    /** @brief Make the complete accepted Jacobian of this illustrative two-coordinate example. */
    [[nodiscard]] std::shared_ptr<const serif::DenseMatrix> material_jacobian(double density_response) {
        auto matrix = std::make_shared<serif::DenseMatrix>(2);
        (*matrix)(0, 0) = density_response;
        (*matrix)(0, 1) = 3.0;
        (*matrix)(1, 0) = 0.0;
        (*matrix)(1, 1) = 4.0;
        return matrix;
    }

    /** @brief Publish one frozen physical block; real discretized equations supply volume_field_block and pressure_field_block instead. */
    [[nodiscard]] serif::operators::LinearizedBlock scalar_block(double entry) {
        auto matrix = std::make_shared<serif::DenseMatrix>(1);
        (*matrix)(0, 0) = entry;
        return serif::operators::linearized_block(std::shared_ptr<const serif::DenseMatrix>(std::move(matrix)));
    }
} // namespace

/** @brief Demonstrate public preparation, reuse, MFEM application, and accepted-factor ownership with a transparent 2x2 system. */
int main() {
    namespace preconditioner = serif::solver::preconditioner;
    namespace dependencies = serif::utils::dependencies;
    namespace inverse = preconditioner::strategies::inverse;
    namespace coupling = preconditioner::strategies::coupling;
    using Density = serif::discretization::blocks::density::mass;
    using Potential = serif::discretization::blocks::gravity::potential;

    std::println("Illustrative 2x2 material block; these coefficients do not define a stellar model.");
    // Correct potential first, then subtract its response from the density equation.
    const auto density = preconditioner::make_preconditioner<Density>(inverse::Diagonal{});
    const auto potential = preconditioner::make_preconditioner<Potential>(inverse::Diagonal{});
    const auto choice = preconditioner::couple(density, potential, coupling::SecondThenFirst{});
    using Recipe = decltype(choice);
    using DensityBlock = preconditioner::keys::Block<dependencies::KeyList<Density>, dependencies::KeyList<Density>>;

    // Physical blocks and the full Jacobian come from one accepted state, in (density, potential) order.
    auto accepted_jacobian = material_jacobian(2.0);
    auto density_block = scalar_block((*accepted_jacobian)(0, 0));
    const auto context = dependencies::make_dependency_context(
        dependencies::borrow<DensityBlock>(density_block),
        preconditioner::physical_block<Potential, Potential>(scalar_block((*accepted_jacobian)(1, 1))),
        preconditioner::physical_block<Density, Potential>(scalar_block((*accepted_jacobian)(0, 1))));

    // Describing a graph performs no numerical setup; preparation explicitly accepts its inputs.
    const auto graph = preconditioner::make_preconditioner_graph(choice);
    auto cache = dependencies::make_computation_cache<preconditioner::Failure>(graph, context);
    const auto prepared = cache.try_prepare(context);
    if (!prepared) {
        std::println(stderr, "Cannot prepare material preconditioner: {}", prepared.error().error.message);
        return 1;
    }
    std::println("Initial preparation: {} computations", prepared->computed_count());

    // Sharing fixes the numerical factors for this solve, even after the cache accepts another state.
    const auto frozen = cache.share<Recipe::Result>();
    preconditioner::Solver original_solver(*frozen, accepted_jacobian);
    original_solver.SetOperator(*accepted_jacobian);
    serif::Vector residual(2), correction(2);
    residual[0] = 10.0;
    residual[1] = 8.0;
    correction = 99.0;
    original_solver.Mult(residual, correction);
    std::println("J = [[2, 3], [0, 4]], residual = [10, 8]: correction = [{:.3f}, {:.3f}]", correction[0], correction[1]);

    const auto unchanged = cache.try_prepare(context);
    if (!unchanged) {
        std::println(stderr, "Cannot reuse material preconditioner: {}", unchanged.error().error.message);
        return 1;
    }
    std::println("Unchanged accepted state: {} computations", unchanged->computed_count());

    // Replacing an accepted block changes its identity; the shared dependency engine refreshes only its consumers.
    accepted_jacobian = material_jacobian(4.0);
    density_block = scalar_block((*accepted_jacobian)(0, 0));
    const auto refreshed = cache.try_prepare(context);
    if (!refreshed) {
        std::println(stderr, "Cannot refresh material preconditioner: {}", refreshed.error().error.message);
        return 1;
    }
    const auto current = cache.share<Recipe::Result>();
    preconditioner::Solver current_solver(*current, accepted_jacobian);
    current_solver.SetOperator(*accepted_jacobian);
    current_solver.Mult(residual, correction);
    std::println("Density response 2 -> 4: {} computations; correction = [{:.3f}, {:.3f}]", refreshed->computed_count(), correction[0], correction[1]);

    // The old adapter still owns both its original Jacobian and original inverse; never rebind it to the new Jacobian.
    original_solver.Mult(residual, correction);
    std::println("Original frozen solve remains: [{:.3f}, {:.3f}]", correction[0], correction[1]);
    return 0;
}