#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include "serif/tests/surface/contexts.hpp"

namespace {
    using namespace serif::tests::surface;
    using Status = dependencies::PreparationStatus;
    using Response = physics::keys::EOSResponse;
    using Linearization = physics::keys::Linearization;
    using OffsetFormulation = physics::PressureSurfaceFormulation<OffsetEOS, BarotropicState>;
}

TEST_CASE("Surface graph construction performs no physics and pressure-only response is broadcast", "[surface][unit]") {
    OffsetEOS model;
    const auto context = offset_context(model, 0.4, {0.7, 1.1, 1.4});
    auto graph = physics::make_pressure_surface_graph<OffsetEOS, BarotropicState>();
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(std::move(graph), context);
    CHECK(model.evaluations == 0);
    CHECK_FALSE(cache.ready());
    CHECK_THROWS_AS(cache.get<Linearization>(), std::logic_error);

    const auto prepared = cache.try_prepare(context);
    REQUIRE(prepared);
    CHECK(prepared->computed_count() == 2);
    CHECK(model.evaluations == 1);
    REQUIRE(cache.get<Response>().size() == 1);
    const auto& rows = cache.get<Linearization>();
    REQUIRE(rows.size() == 3);
    for (std::size_t index = 0; index < rows.size(); ++index) {
        const double expected = context.get<Enthalpy>().values[index] - (0.3 + 2.0 * 0.4);
        CHECK(rows[index].residual == expected);
        CHECK(rows[index].field_derivatives.get<Enthalpy>() == 1.0);
    }
}

TEST_CASE("Surface source changes refresh only the calculations that depend on them", "[surface][unit]") {
    OffsetEOS model;
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<OffsetFormulation>(), offset_context(model, 0.0, {0.3, 0.7}));
    REQUIRE(cache.try_prepare(offset_context(model, 0.0, {0.3, 0.7})));
    CHECK(cache.get<Linearization>()[0].residual == 0.0);

    auto report = cache.try_prepare(offset_context(model, 0.0, {0.3, 0.7}));
    REQUIRE(report);
    CHECK(report->computed_count() == 0);
    CHECK(model.evaluations == 1);

    report = cache.try_prepare(offset_context(model, 0.0, {0.4, 0.9}));
    REQUIRE(report);
    CHECK(report->status<Response>() == Status::reused);
    CHECK(report->status<Linearization>() == Status::computed);
    CHECK(model.evaluations == 1);

    report = cache.try_prepare(offset_context(model, 0.2, {0.4, 0.9}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(model.evaluations == 2);

    model.offset = 0.6;
    report = cache.try_prepare(offset_context(model, 0.2, {0.4, 0.9}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(cache.get<Linearization>()[0].residual == 0.4 - (0.6 + 2.0 * 0.2));

    OffsetEOS replacement{.offset = model.offset};
    report = cache.try_prepare(offset_context(replacement, 0.2, {0.4, 0.9}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(replacement.evaluations == 1);

    report = cache.try_prepare(offset_context(replacement, 0.2, {0.4, 0.9}), {}, dependencies::PreparationMode::rebuild);
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
}

TEST_CASE("Surface composition response caches values and all partial derivatives together", "[surface][unit]") {
    GeneralEOS model;
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<GeneralFormulation>(), general_context(model, 0.4, {1.7, 1.9}, {0.2, 0.3}, {0.1, 0.2}));
    REQUIRE(cache.try_prepare(general_context(model, 0.4, {1.7, 1.9}, {0.2, 0.3}, {0.1, 0.2})));
    CHECK(model.evaluations == 2);
    CHECK(model.partials == 4);
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Entropy>() == -3.0);
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Composition>() == -5.0);

    auto report = cache.try_prepare(general_context(model, 0.4, {1.8, 2.0}, {0.2, 0.3}, {0.1, 0.2}));
    REQUIRE(report);
    CHECK(report->status<Response>() == Status::reused);
    CHECK(model.partials == 4);

    report = cache.try_prepare(general_context(model, 0.4, {1.8, 2.0}, {0.4, 0.3}, {0.1, 0.2}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(model.evaluations == 4);

    report = cache.try_prepare(general_context(model, 0.4, {1.8, 2.0}, {0.4, 0.3}, {0.3, 0.2}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(model.evaluations == 6);

    model.scale = 2.0;
    REQUIRE(cache.try_prepare(general_context(model, 0.4, {1.8, 2.0}, {0.4, 0.3}, {0.3, 0.2})));
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Entropy>() == -6.0);
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Composition>() == -10.0);
}

TEST_CASE("Density surface graph refreshes nonlinear temperature derivatives", "[surface][unit]") {
    GasEOS model;
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<GasFormulation>(), gas_context(model, 0.8, {0.4, 0.6}, {1.0, 2.0}));
    REQUIRE(cache.try_prepare(gas_context(model, 0.8, {0.4, 0.6}, {1.0, 2.0})));
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Temperature>() == 0.8);
    CHECK(cache.get<Linearization>()[1].field_derivatives.get<Temperature>() == 0.2);

    const auto report = cache.try_prepare(gas_context(model, 0.8, {0.4, 0.6}, {2.0, 4.0}));
    REQUIRE(report);
    CHECK(report->computed_count() == 2);
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Temperature>() == 0.2);
    CHECK(cache.get<Linearization>()[1].field_derivatives.get<Temperature>() == 0.05);
}

TEST_CASE("Direct pressure surface reuses EOS work when prescribed pressure changes", "[surface][unit]") {
    GasEOS model;
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<DirectGasFormulation>(), gas_context(model, 0.8, {0.4, 0.6}, {1.0, 2.0}));
    REQUIRE(cache.try_prepare(gas_context(model, 0.8, {0.4, 0.6}, {1.0, 2.0})));
    const auto report = cache.try_prepare(gas_context(model, 0.9, {0.4, 0.6}, {1.0, 2.0}));
    REQUIRE(report);
    CHECK(report->status<Response>() == Status::reused);
    CHECK(report->status<Linearization>() == Status::computed);
    CHECK(model.evaluations == 2);
    CHECK(cache.get<Linearization>()[0].residual == 0.4 * 1.0 - 0.9);
    CHECK(cache.get<Linearization>()[1].field_derivatives.get<Density>() == 2.0);
    CHECK(cache.get<Linearization>()[1].field_derivatives.get<Temperature>() == 0.6);
}

TEST_CASE("Surface cache requested results and moves retain the ordinary dependency lifecycle", "[surface][unit]") {
    OffsetEOS model;
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<OffsetFormulation>(), offset_context(model, 0.0, {0.5}));
    const auto report = cache.try_prepare(offset_context(model, 0.0, {std::numeric_limits<double>::quiet_NaN()}), dependencies::RequestedResults<Response>{});
    REQUIRE(report);
    CHECK(report->computed_count() == 1);
    CHECK_FALSE(cache.contains<Linearization>());
    CHECK_THROWS_AS(cache.get<Linearization>(), std::logic_error);
    REQUIRE(cache.try_prepare(offset_context(model, 0.0, {0.5})));

    auto moved = std::move(cache);
    CHECK_FALSE(cache.ready());
    REQUIRE(moved.ready());
    CHECK(moved.get<Linearization>()[0].residual == 0.5 - 0.3);
}

TEST_CASE("Surface batch with aliased quantities gathers each term once and sums every derivative", "[surface][unit]") {
    using SharedState = dependencies::Bindings<
        dependencies::Binding<quantities::SpecificEnthalpy, Enthalpy>,
        dependencies::Binding<quantities::Entropy, Enthalpy>,
        dependencies::Binding<ElectronFraction, Enthalpy>
    >;
    GeneralEOS model;
    const auto context = dependencies::make_dependency_context(
        dependencies::borrow<physics::keys::EquationOfState>(model, dependencies::track_by([](const GeneralEOS& eos) {
            return eos.scale;
        })),
        physics::pressure_condition(physics::IsobaricSurface{quantities::PressureValue{0.4}}),
        physics::boundary_values<Enthalpy>({0.2, 0.3})
    );
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(physics::make_pressure_surface_graph<GeneralEOS, SharedState>(), context);
    REQUIRE(cache.try_prepare(context));
    REQUIRE(cache.get<Linearization>().size() == 2);
    CHECK(cache.get<Linearization>()[0].field_derivatives.get<Enthalpy>() == -7.0);
    CHECK(cache.get<Linearization>()[1].field_derivatives.get<Enthalpy>() == -7.0);
}