#include "serif/tests/equations/fixture.hpp"

#include <catch2/catch_test_macros.hpp>
#include <catch2/generators/catch_generators.hpp>
#include <memory>
#include <stdexcept>
#include <utility>

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    namespace dependencies = utils::dependencies;
    namespace stages = backend::volume_keys;

    using Vector = ::mfem::Vector;
    using Mode = backend::PreparationMode;
    using Status = dependencies::PreparationStatus;
    using DensityValue = discretization::forms::Operand<blocks::density::mass>;
    using EnthalpyValue = discretization::forms::Operand<::serif::tests::equations::Enthalpy>;
}

TEST_CASE("Named volume caches retain independent equation configurations", "[equations][computation_graphs][unit]") {
    Fixture fixture;
    const eos::models::Polytrope first_eos{1.5, 0.25};
    const eos::models::Polytrope second_eos{3.0, 1.5};
    auto first_density = fixture.density();
    const auto first_enthalpy = fixture.enthalpy();
    const auto first_displacement = fixture.displacement(0.04);
    const auto second_density = fixture.density(1.0);
    const auto second_enthalpy = fixture.enthalpy(0.5);
    const auto second_displacement = fixture.displacement(-0.03);

    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const auto first_inputs = backend::make_volume_inputs<Closure>(
        discretization, fixture.state(first_density, first_enthalpy), first_displacement
    );
    const auto second_inputs = backend::make_volume_inputs<Closure>(
        discretization, fixture.state(second_density, second_enthalpy), second_displacement
    );
    auto first = backend::make_volume_equation_cache(Closure{first_eos}, discretization, first_inputs);
    auto second = backend::make_volume_equation_cache(Closure{second_eos}, discretization, second_inputs);
    CHECK_FALSE(first.ready());
    CHECK_FALSE(second.ready());

    backend::prepare_volume(first, discretization, fixture.state(first_density, first_enthalpy), first_displacement);
    backend::prepare_volume(second, discretization, fixture.state(second_density, second_enthalpy), second_displacement);
    Vector first_residual, second_residual;
    backend::volume_residual(first.get<stages::Linearization>(), first_residual);
    backend::volume_residual(second.get<stages::Linearization>(), second_residual);
    CHECK(relative_error(first_residual, manual_closure(fixture, first_eos, first_density, first_enthalpy, first_displacement)) < 2.0e-12);
    CHECK(relative_error(second_residual, manual_closure(fixture, second_eos, second_density, second_enthalpy, second_displacement)) < 2.0e-12);

    first_density += 0.07;
    backend::prepare_volume(first, discretization, fixture.state(first_density, first_enthalpy), first_displacement);
    backend::volume_residual(first.get<stages::Linearization>(), first_residual);
    CHECK(relative_error(first_residual, manual_closure(fixture, first_eos, first_density, first_enthalpy, first_displacement)) < 2.0e-12);

    Vector still_second;
    backend::volume_residual(second.get<stages::Linearization>(), still_second);
    CHECK(relative_error(still_second, second_residual) == 0.0);
    CHECK(second.preparation_count() == 1);
    CHECK(first.preparation_count() == 2);
}

TEST_CASE("Volume graph nodes expose coefficient and geometry reuse independently of workspace storage", "[equations][computation_graphs][unit]") {
    Fixture fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    auto displacement = fixture.displacement(0.04);
    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const auto inputs = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
    auto cache = backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    const auto initial = cache.report();
    CHECK(initial.status<stages::Basis<0>>() == Status::computed);
    CHECK(initial.status<stages::Geometry<0>>() == Status::computed);
    CHECK(initial.status<stages::ReferenceValues<0, DensityValue>>() == Status::computed);
    CHECK(initial.status<stages::ReferenceValues<0, EnthalpyValue>>() == Status::computed);
    CHECK(initial.status<stages::QuadratureLinearization<0>>() == Status::computed);
    CHECK(initial.status<stages::GeometryDerivatives<0>>() == Status::computed);
    const auto original_basis = cache.get<stages::Basis<0>>().front();
    REQUIRE(original_basis->reference_geometry.size() == static_cast<std::size_t>(original_basis->rule->GetNPoints()));
    const auto original_derivatives = cache.get<stages::GeometryDerivatives<0>>().front();
    CHECK(cache.get<stages::Linearization>().geometry_derivatives[0].front() == original_derivatives);

    density += 0.03;
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    const auto coefficients_changed = cache.report();
    CHECK(coefficients_changed.status<stages::Basis<0>>() == Status::reused);
    CHECK(coefficients_changed.status<stages::Geometry<0>>() == Status::reused);
    CHECK(coefficients_changed.status<stages::ReferenceValues<0, DensityValue>>() == Status::computed);
    CHECK(coefficients_changed.status<stages::ReferenceValues<0, EnthalpyValue>>() == Status::reused);
    CHECK(coefficients_changed.status<stages::QuadratureLinearization<0>>() == Status::computed);
    CHECK(coefficients_changed.status<stages::GeometryDerivatives<0>>() == Status::computed);
    CHECK(cache.get<stages::GeometryDerivatives<0>>().front() != original_derivatives);
    Vector residual;
    backend::volume_residual(cache.get<stages::Linearization>(), residual);
    CHECK(relative_error(residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);

    displacement *= 1.5;
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    const auto geometry_changed = cache.report();
    CHECK(geometry_changed.status<stages::Basis<0>>() == Status::reused);
    CHECK(cache.get<stages::Basis<0>>().front() == original_basis);
    CHECK(geometry_changed.status<stages::Geometry<0>>() == Status::computed);
    CHECK(geometry_changed.status<stages::ReferenceValues<0, DensityValue>>() == Status::reused);
    CHECK(geometry_changed.status<stages::ReferenceValues<0, EnthalpyValue>>() == Status::reused);
    CHECK(geometry_changed.status<stages::QuadratureLinearization<0>>() == Status::computed);
    CHECK(geometry_changed.status<stages::GeometryDerivatives<0>>() == Status::computed);
    backend::volume_residual(cache.get<stages::Linearization>(), residual);
    CHECK(relative_error(residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);

    // Equal coefficient copies and different application scratch describe the same numerical state.
    const Vector density_copy(density), enthalpy_copy(enthalpy), displacement_copy(displacement);
    backend::Workspace another_workspace;
    backend::prepare_volume(cache, discretization, fixture.state(density_copy, enthalpy_copy), displacement_copy);
    const auto unchanged = cache.report();
    CHECK(unchanged.computed_count() == 0);
    CHECK(unchanged.status<stages::Linearization>() == Status::reused);
    CHECK(unchanged.status<stages::GeometryDerivatives<0>>() == Status::reused);
    Vector repeated;
    backend::volume_residual(cache.get<stages::Linearization>(), repeated);
    CHECK(relative_error(repeated, residual) == 0.0);

    // Application reads accepted snapshots until the caller explicitly prepares new inputs.
    const auto direction = packed(fixture.density(0.7), fixture.enthalpy(0.2), fixture.displacement(0.01));
    const auto dual = deterministic(fixture.density_binding.size(), 0.3);
    Vector action, transpose, frozen_action, frozen_transpose;
    backend::apply_jacobian(cache.get<stages::Linearization>(), direction, fixture.workspace, action);
    backend::apply_jacobian_transpose(cache.get<stages::Linearization>(), dual, fixture.workspace, transpose);
    density += 0.09;
    displacement *= 1.1;
    backend::volume_residual(cache.get<stages::Linearization>(), repeated);
    backend::apply_jacobian(cache.get<stages::Linearization>(), direction, another_workspace, frozen_action);
    backend::apply_jacobian_transpose(cache.get<stages::Linearization>(), dual, another_workspace, frozen_transpose);
    CHECK(relative_error(repeated, residual) == 0.0);
    CHECK(relative_error(frozen_action, action) == 0.0);
    CHECK(relative_error(frozen_transpose, transpose) == 0.0);
}

TEST_CASE("Rejected volume graph candidates never expose partially updated operator results", "[equations][computation_graphs][unit]") {
    Fixture fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    auto density = fixture.density();
    auto enthalpy = fixture.enthalpy();
    auto displacement = fixture.displacement(0.04);
    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const auto inputs = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
    auto cache = backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);

    // Several upstream results change before the equation rejects this candidate.
    density += 0.05;
    displacement *= 1.5;
    enthalpy = -1.0;
    const auto rejected = backend::try_prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::EQUATION);
    CHECK_FALSE(cache.ready());
    CHECK_THROWS_AS(cache.get<stages::Linearization>(), std::logic_error);
    CHECK_THROWS_AS(cache.get<stages::Basis<0>>(), std::logic_error);
    CHECK_THROWS_AS(cache.get<stages::GeometryDerivatives<0>>(), std::logic_error);
    Vector result;
    CHECK_THROWS_AS(backend::volume_residual(cache.get<stages::Linearization>(), result), std::logic_error);

    enthalpy = fixture.enthalpy(0.5);
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    backend::volume_residual(cache.get<stages::Linearization>(), result);
    CHECK(relative_error(result, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);
    CHECK(cache.preparation_count() == 2);

    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement, Mode::rebuild);
    Vector rebuilt;
    backend::volume_residual(cache.get<stages::Linearization>(), rebuilt);
    CHECK(relative_error(result, rebuilt) < 2.0e-15);
}

TEST_CASE("Volume caches reject another discretization of the same type and recover", "[equations][computation_graphs][unit]") {
    Fixture fixture;
    Fixture other_fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    const auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    const auto displacement = fixture.displacement(0.04);
    const auto other_density = other_fixture.density(0.5);
    const auto other_enthalpy = other_fixture.enthalpy(0.7);
    const auto other_displacement = other_fixture.displacement(-0.03);
    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const backend::VolumeDiscretization other_discretization{
        .residual_binding = other_fixture.density_binding,
        .displacement_binding = other_fixture.displacement_binding,
        .compactification = other_fixture.compactification,
        .mapper = other_fixture.mapper
    };
    const auto inputs = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
    auto cache = backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    CHECK(cache.preparation_count() == 1);

    CHECK_THROWS_AS(
        backend::try_prepare_volume(cache, other_discretization, other_fixture.state(other_density, other_enthalpy), other_displacement),
        std::invalid_argument
    );
    CHECK_FALSE(cache.ready());
    CHECK_THROWS_AS(cache.get<stages::Linearization>(), std::logic_error);
    CHECK_THROWS_AS(cache.get<stages::Basis<0>>(), std::logic_error);
    CHECK_THROWS_AS(cache.get<stages::GeometryDerivatives<0>>(), std::logic_error);
    CHECK(cache.preparation_count() == 1);

    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    Vector residual;
    backend::volume_residual(cache.get<stages::Linearization>(), residual);
    CHECK(relative_error(residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);
    CHECK(cache.preparation_count() == 2);
}

TEST_CASE("The volume graph refreshes immutable reference geometry after mesh coordinate updates", "[equations][computation_graphs][unit]") {
    Fixture fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    const auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    const auto displacement = fixture.displacement(0.04);
    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const auto inputs = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
    auto cache = backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    const auto original = cache.get<stages::Basis<0>>().front();
    const double original_weight = original->reference_geometry.front().weight;

    fixture.mesh.Transform([](const ::mfem::Vector& position, ::mfem::Vector& transformed) {
        transformed = position;
        transformed[0] *= 1.2;
    });
    backend::prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    CHECK(cache.report().status<stages::Basis<0>>() == Status::computed);
    const auto changed = cache.get<stages::Basis<0>>().front();
    CHECK(changed != original);
    CHECK(changed->reference_geometry.front().weight > original_weight);
    CHECK(original->reference_geometry.front().weight == original_weight);
    Vector residual;
    backend::volume_residual(cache.get<stages::Linearization>(), residual);
    CHECK(relative_error(residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);

    fixture.mesh.Transform([](const ::mfem::Vector& position, ::mfem::Vector& transformed) {
        transformed = position;
        transformed[0] = 0.0;
    });
    const auto rejected = backend::try_prepare_volume(cache, discretization, fixture.state(density, enthalpy), displacement);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::MAPPING);
    CHECK_FALSE(cache.ready());
    CHECK(original->reference_geometry.front().weight == original_weight);
}

TEST_CASE("Volume graph source tracking rejects stale reused input contexts", "[equations][computation_graphs][unit]") {
    const bool initially_prepared = GENERATE(false, true);
    Fixture fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    const auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    const auto displacement = fixture.displacement(0.04);
    const backend::VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper
    };
    const auto inputs = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
    auto cache = backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
    if (initially_prepared) {
        REQUIRE(cache.try_prepare(inputs));
    }
    SECTION("Coordinate changes require fresh source metadata") {
        fixture.mesh.NodesUpdated();
        CHECK_THROWS_AS(cache.try_prepare(inputs), std::logic_error);
        CHECK_FALSE(cache.ready());
        const auto current = backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
        REQUIRE(cache.try_prepare(current));
        Vector residual;
        backend::volume_residual(cache.get<stages::Linearization>(), residual);
        CHECK(relative_error(residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);
    }
    SECTION("Domain edits invalidate the borrowed bindings") {
        fixture.mesh.GetElement(0)->SetAttribute(3);
        CHECK_THROWS_AS(cache.try_prepare(inputs), std::logic_error);
        CHECK_FALSE(cache.ready());
    }
    SECTION("Topology edits invalidate the borrowed bindings") {
        fixture.mesh.UniformRefinement();
        CHECK_THROWS_AS(cache.try_prepare(inputs), std::logic_error);
        CHECK_FALSE(cache.ready());
    }
}