#include <algorithm>
#include <limits>
#include <memory>
#include <tuple>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/backend/mfem/geometry/surface/pressure_boundary.hpp"
#include "serif/backend/mfem/model/computations/pressure_boundary.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 surface = serif::geometry::surface;
    namespace boundary = surface::backend::mfem;
    namespace physical = serif::operators::backend::mfem;
    namespace model = serif::model;
    namespace dependencies = serif::utils::dependencies;
    namespace blocks = serif::discretization::blocks;
    using Density = fixture::BoundaryDensity;
    using Temperature = test::Temperature;
    using Vector = mfem::Vector;
    struct OtherEquation {
    };

    /** @brief Keep an independent nonzero interior action so replacement must remove every old constrained-row column. */
    std::shared_ptr<const mfem::SparseMatrix> base_matrix(const int rows, const int columns) {
        auto matrix = std::make_shared<mfem::SparseMatrix>(rows, columns);
        for (int row = 0; row < rows; ++row) {
            matrix->Add(row, row % columns, 7.0);
        }
        matrix->Finalize();
        return matrix;
    }

    /** @brief Check every row against independent MFEM interpolation and the analytic ideal-gas boundary relation. */
    template <typename Formulation, bool SolveTemperature>
    void check_pressure_graph() {
        using Values = std::conditional_t<SolveTemperature, dependencies::KeyList<Temperature, Density>, dependencies::KeyList<Density>>;
        using Layout = blocks::SystemLayout<Values, dependencies::KeyList<OtherEquation, Density>>;
        fixture::Fixture mesh;
        Vector density = fixture::sequence(mesh.density.size(), 0.8, 0.2);
        Vector temperature = fixture::sequence(mesh.temperature.size(), 1.7, 0.4);
        const auto state = physical::make_term_coefficients(physical::bind(mesh.density, density), physical::bind(mesh.temperature, temperature));
        const auto compiled = boundary::make_pressure_boundary<Formulation>(state);
        Layout layout;
        if constexpr (SolveTemperature) {
            layout.state_sizes = {static_cast<std::size_t>(temperature.Size()), static_cast<std::size_t>(density.Size())};
        }
        else {
            layout.state_sizes = {static_cast<std::size_t>(density.Size())};
        }
        layout.residual_sizes = {2, static_cast<std::size_t>(density.Size())};
        test::GasEOS eos;
        const auto inputs = boundary::make_pressure_boundary_inputs(compiled, dependencies::borrow<surface::keys::EquationOfState>(eos, dependencies::track_by([](const test::GasEOS&) {
                                                                                                                                                                   return 0;
                                                                                                                                                               })),
                                                                    surface::IsobaricSurface{test::quantities::PressureValue{0.8}}, state);
        auto cache = std::apply([&](auto... node) {
                                    return dependencies::make_computation_cache<surface::SurfaceFailure>(inputs, std::move(node)...);
                                },
                                std::tuple_cat(boundary::make_pressure_boundary_computations<Formulation>(), std::tuple{model::computations::pressure_boundary_contribution<Formulation>(layout)}));
        REQUIRE(cache.try_prepare(inputs));
        const auto frozen = cache.template share<model::computations::PressureBoundaryContribution>();
        REQUIRE(frozen->replacements.size() == 1);
        const auto& replacement = frozen->replacements[0];
        REQUIRE(replacement.jacobian.size() == (SolveTemperature ? 2 : 1));
        REQUIRE(replacement.residual.size() == static_cast<std::size_t>(compiled->selected_rows->size()));
        const int rows = density.Size() + 2;
        const int columns = static_cast<int>(layout.template total_size<blocks::RepresentedValueKind::value>());
        const auto base = base_matrix(rows, columns);
        const model::backend::mfem::PhysicalSystem accepted(rows, columns,
                                                            {{.row_offset = 0, .values = std::vector<double>(rows, 9.0)}, {.row_offset = 0, .values = std::vector<double>(rows, 2.0)}},
                                                            {{.row_offset = 0, .column_offset = 0, .block = physical::linearized_block(base)}}, frozen->replacements);
        const Vector density_samples = fixture::reference_samples(mesh.density, *compiled->selected_rows, mesh.density, density);
        const Vector temperature_samples = fixture::reference_samples(mesh.density, *compiled->selected_rows, mesh.temperature, temperature);
        Vector direction = fixture::sequence(columns, 0.3, 0.7);
        const int density_offset = static_cast<int>(layout.template offset<blocks::RepresentedValueKind::value, Density>());
        Vector density_direction(direction.GetData() + density_offset, density.Size());
        const Vector density_variation = fixture::reference_samples(mesh.density, *compiled->selected_rows, mesh.density, density_direction);
        Vector temperature_variation(compiled->selected_rows->size());
        temperature_variation = 0.0;
        if constexpr (SolveTemperature) {
            Vector temperature_direction(direction.GetData(), temperature.Size());
            temperature_variation = fixture::reference_samples(mesh.density, *compiled->selected_rows, mesh.temperature, temperature_direction);
        }
        Vector expected(rows), actual(rows), expected_residual(rows);
        base->Mult(direction, expected);
        expected_residual = 11.0;
        for (int index = 0; index < compiled->selected_rows->size(); ++index) {
            const int row = 2 + compiled->selected_rows->reduced_dofs()[index];
            CHECK(replacement.selected_rows[index] == row);
            const double rho = density_samples[index];
            const double thermal = temperature_samples[index];
            if constexpr (Formulation::resolves_carrier) {
                expected_residual[row] = rho - 0.8 / thermal;
                expected[row] = density_variation[index] + 0.8 * temperature_variation[index] / (thermal * thermal);
            }
            else {
                expected_residual[row] = rho * thermal - 0.8;
                expected[row] = thermal * density_variation[index] + rho * temperature_variation[index];
            }
        }
        accepted.jacobian().action()->Mult(direction, actual);
        for (int row = 0; row < rows; ++row) {
            CHECK(accepted.residual()[row] == Catch::Approx(expected_residual[row]).margin(3.0e-13));
            CHECK(actual[row] == Catch::Approx(expected[row]).margin(3.0e-13));
        }
        const Vector dual = fixture::sequence(rows, 0.5, 0.2);
        Vector transpose(columns);
        accepted.jacobian().action()->MultTranspose(dual, transpose);
        CHECK((dual * actual) == Catch::Approx(direction * transpose).margin(2.0e-11));
        REQUIRE(cache.try_prepare(inputs));
        CHECK(cache.report().computed_count() == 0);
        CHECK(cache.template share<model::computations::PressureBoundaryContribution>() == frozen);
        density += 0.04;
        REQUIRE(cache.try_prepare(inputs));
        CHECK(cache.template share<model::computations::PressureBoundaryContribution>() != frozen);
        accepted.jacobian().action()->Mult(direction, actual);
        for (int row = 0; row < rows; ++row) {
            CHECK(actual[row] == Catch::Approx(expected[row]).margin(3.0e-13));
        }
    }
} // namespace

TEST_CASE("Model pressure replacements preserve general EOS fields and independent system ordering", "[unit][model][pressure_graph]") {
    SECTION("Density carries an isobaric condition with solved temperature") {
        check_pressure_graph<fixture::DensityFormulation<>, true>();
    }
    SECTION("Direct pressure contributes interpolated columns for every solved EOS field") {
        check_pressure_graph<fixture::PressureFormulation<>, true>();
    }
    SECTION("Prescribed temperature is evaluated without introducing a Newton column") {
        check_pressure_graph<fixture::PressureFormulation<>, false>();
    }
}

TEST_CASE("Pressure replacement rejects layouts that cannot be represented by MFEM indices", "[unit][model][pressure_graph]") {
    using Layout = blocks::SystemLayout<dependencies::KeyList<Density>, dependencies::KeyList<Density>>;
    const Layout layout{.state_sizes = {static_cast<std::size_t>(std::numeric_limits<int>::max()) + 1}, .residual_sizes = {1}};
    CHECK_THROWS(model::computations::pressure_boundary_contribution<fixture::PressureFormulation<>>(layout));
}