#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <memory>
#include <vector>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/backend/mfem/geometry/surface/pressure_boundary.hpp"
#include "serif/backend/mfem/geometry/surface/pressure_field_block.hpp"
#include "serif/tests/surface/boundary_fixture.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace test = serif::tests::surface;
    namespace fixture = test::boundary;
    namespace physics = serif::geometry::surface;
    namespace boundary = physics::backend::mfem;
    namespace operators = serif::operators::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    using Temperature = fixture::ContinuousScalar<test::quantities::Temperature, 3>;
    using Vector = ::mfem::Vector;

    /** @brief Give the volume block nonzero off-diagonal couplings that constrained rows must completely remove. */
    [[nodiscard]] std::shared_ptr<const ::mfem::SparseMatrix> volume_matrix(int rows, int columns)
    {
        auto matrix = std::make_shared<::mfem::SparseMatrix>(rows, columns);
        for (int row = 0; row < rows; ++row)
        {
            matrix->Add(row, row % columns, 2.0);
            matrix->Add(row, (row + 1) % columns, -0.25);
        }
        matrix->Finalize();
        return matrix;
    }

    /** @brief Compare each physical row rather than a norm that could hide a wrong small boundary entry. */
    void check_rows(const Vector& actual, const Vector& expected)
    {
        REQUIRE(actual.Size() == expected.Size());
        for (int row = 0; row < actual.Size(); ++row)
        {
            CHECK(actual[row] == Catch::Approx(expected[row]).margin(3e-13));
        }
    }

    /** @brief Exercise different nodal bases at equal dimensions so an interpolated diagonal cannot be mistaken for a row sum. */
    template <typename Formulation>
    void check_pressure_field()
    {
        fixture::Fixture mesh;
        ::mfem::H1_FECollection uniform_collection(3, 3, ::mfem::BasisType::ClosedUniform);
        ::mfem::FiniteElementSpace uniform_space(&mesh.mesh, &uniform_collection);
        serif::discretization::backend::mfem::TermBinding<Temperature, fixture::Schema> temperature_binding(uniform_space);
        REQUIRE(temperature_binding.size() == mesh.density.size());
        Vector density = fixture::sequence(mesh.density.size(), 0.8, 0.2);
        Vector temperature = fixture::sequence(temperature_binding.size(), 1.7, 0.4);
        const Vector direction = fixture::sequence(temperature.Size(), 0.2, 0.1);
        Vector zero_density(density.Size());
        zero_density = 0.0;
        const auto state = operators::make_term_coefficients(operators::bind(mesh.density, density), operators::bind(temperature_binding, temperature));
        const auto directions = operators::make_term_coefficients(operators::bind(mesh.density, zero_density), operators::bind(temperature_binding, direction));
        test::GasEOS model;
        const auto selected = boundary::make_pressure_boundary<Formulation>(state);
        /** @brief The ideal-gas fixture has no variable parameters; observational counters do not change its physics. */
        const auto eos_tracking = [](const test::GasEOS&)
        {
            return 0;
        };
        const auto inputs = boundary::make_pressure_boundary_inputs(selected, dependencies::borrow<physics::keys::EquationOfState>(model, dependencies::track_by(eos_tracking)), physics::IsobaricSurface{test::quantities::PressureValue{0.8}}, state);
        auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(boundary::make_pressure_boundary_graph<Formulation>(), inputs);
        REQUIRE(cache.try_prepare(inputs));
        std::vector<physics::SurfaceLinearization<Formulation>> accepted = cache.template get<physics::keys::Linearization>();
        const auto volume = volume_matrix(density.Size(), temperature.Size());
        const operators::LinearizedBlock block = boundary::pressure_field_block<Temperature>(operators::linearized_block(volume), *selected, accepted);
        REQUIRE(block.has_diagonal());
        REQUIRE(block.has_assembly());
        const auto assembled = block.assemble();
        const std::vector<double> diagonal = block.diagonal();
        Vector actual(density.Size()), expected(density.Size()), assembled_action(density.Size());
        volume->Mult(direction, expected);
        boundary::apply_jacobian_rows(*selected, accepted, directions, expected);
        block.action()->Mult(direction, actual);
        assembled->Mult(direction, assembled_action);
        check_rows(actual, expected);
        check_rows(assembled_action, expected);

        const ::mfem::SparseMatrix& interpolation = selected->traces.template get<Temperature>().interpolation();
        bool interpolating_diagonal = false;
        for (int index = 0; index < selected->selected_rows->size(); ++index)
        {
            const int row = selected->selected_rows->reduced_dofs()[index];
            const double derivative = accepted[index].field_derivatives.template get<Temperature>();
            CHECK(diagonal[row] == Catch::Approx(derivative * interpolation(index, row)).margin(3e-14));
            if (std::abs(interpolation(index, row) - 1.0) > 1e-3)
            {
                interpolating_diagonal = true;
            }
        }
        REQUIRE(interpolating_diagonal);
        for (int row = 0; row < assembled->Height(); ++row)
        {
            CHECK(diagonal[row] == Catch::Approx((*assembled)(row, row)).margin(3e-14));
        }
        Vector alias(direction);
        block.action()->Mult(alias, alias);
        check_rows(alias, expected);
        const Vector dual = fixture::sequence(density.Size(), 0.4, 0.8);
        Vector transpose(temperature.Size()), assembled_transpose(temperature.Size());
        block.action()->MultTranspose(dual, transpose);
        assembled->MultTranspose(dual, assembled_transpose);
        check_rows(transpose, assembled_transpose);
        CHECK(dual * actual == Catch::Approx(direction * transpose).margin(2e-12));
        alias = dual;
        block.action()->MultTranspose(alias, alias);
        check_rows(alias, assembled_transpose);

        // Later accepted candidates cannot mutate the frozen row coefficients of an earlier physical block.
        for (auto& point : accepted)
        {
            point.field_derivatives.template get<Temperature>() *= 2.0;
        }
        const auto refreshed = boundary::pressure_field_block<Temperature>(operators::linearized_block(volume), *selected, accepted);
        refreshed.action()->Mult(direction, assembled_action);
        block.action()->Mult(direction, actual);
        check_rows(actual, expected);
        for (int row : selected->selected_rows->reduced_dofs())
        {
            CHECK(assembled_action[row] == Catch::Approx(2.0 * expected[row]).margin(3e-13));
        }

        const auto absent = boundary::pressure_field_block<test::Entropy>(operators::linearized_block(volume), *selected, accepted);
        absent.action()->Mult(direction, actual);
        volume->Mult(direction, expected);
        const auto absent_matrix = absent.assemble();
        const auto absent_diagonal = absent.diagonal();
        for (int row : selected->selected_rows->reduced_dofs())
        {
            expected[row] = 0.0;
            CHECK(absent_diagonal[row] == 0.0);
        }
        check_rows(actual, expected);
        absent_matrix->Mult(direction, actual);
        check_rows(actual, expected);

        accepted.pop_back();
        CHECK_THROWS(boundary::pressure_field_block<Temperature>(operators::linearized_block(volume), *selected, accepted));
    }
} // namespace

TEST_CASE("Selected pressure blocks retain general EOS interpolation in actions diagonals and sparse rows", "[unit][preconditioner][surface][operators]")
{
    SECTION("Density carries the surface relation without enthalpy")
    {
        check_pressure_field<fixture::DensityFormulation<Temperature>>();
    }
    SECTION("Direct pressure derivatives include every thermodynamic field")
    {
        check_pressure_field<fixture::PressureFormulation<Temperature>>();
    }
}

TEST_CASE("Generic physical row replacement preserves optional capabilities and rejects invalid rows", "[unit][preconditioner][surface][operators]")
{
    const auto original = volume_matrix(3, 3);
    auto replacement = std::make_shared<::mfem::SparseMatrix>(1, 3);
    replacement->Add(0, 0, 0.2);
    replacement->Add(0, 1, -0.7);
    replacement->Finalize();
    const operators::LinearizedBlock action_only(original);
    const auto replaced = operators::replace_block_rows(action_only, {1}, replacement);
    CHECK_FALSE(replaced.has_diagonal());
    CHECK_FALSE(replaced.has_assembly());
    const auto full = operators::replace_block_rows(operators::linearized_block(original), {1}, replacement);
    CHECK(full.diagonal()[1] == -0.7);
    CHECK_THROWS(operators::replace_block_rows(action_only, {-1}, replacement));
    CHECK_THROWS(operators::replace_block_rows(action_only, {3}, replacement));
    CHECK_THROWS(operators::replace_block_rows(action_only, {}, replacement));
    auto repeated = std::make_shared<::mfem::SparseMatrix>(2, 3);
    repeated->Finalize();
    CHECK_THROWS(operators::replace_block_rows(action_only, {1, 1}, repeated));
    auto nonfinite = std::make_shared<::mfem::SparseMatrix>(1, 3);
    nonfinite->Add(0, 0, std::numeric_limits<double>::infinity());
    nonfinite->Finalize();
    CHECK_THROWS(operators::replace_block_rows(action_only, {1}, nonfinite));
}