#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <tuple>
#include <variant>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include "serif/backend/mfem/discretization/boundary_trace.hpp"
#include "serif/backend/mfem/geometry/surface/pressure_boundary.hpp"
#include "serif/tests/equations/gravity_fixture.hpp"
#include "serif/tests/equations/volume_equation_fixture.hpp"
#include "serif/tests/surface/boundary_fixture.hpp"

namespace {
    namespace test = serif::tests::equations;
    namespace gravity = test::gravity;
    namespace blocks = serif::discretization::blocks;
    namespace physics = serif::geometry::surface;
    namespace surface = physics::backend::mfem;
    namespace volume = serif::operators::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    using Vector = mfem::Vector;
    using Enthalpy = ::serif::tests::equations::Enthalpy;

    /** @brief Identify boundary coefficients that depend on several native independent coordinates. */
    [[nodiscard]] std::vector<int> hanging_boundary_dofs(const mfem::FiniteElementSpace& space, const mfem::Array<int>& attributes) {
        const mfem::SparseMatrix* prolongation = space.GetConformingProlongation();
        REQUIRE(prolongation != nullptr);
        mfem::Array<int> boundary;
        space.GetEssentialVDofs(attributes, boundary);
        std::vector<int> hanging;
        for (int local = 0; local < boundary.Size(); ++local) {
            if (boundary[local] == 0) {
                continue;
            }
            int contributors = 0;
            for (int entry = prolongation->GetI()[local]; entry < prolongation->GetI()[local + 1]; ++entry) {
                contributors += prolongation->GetData()[entry] != 0.0;
            }
            if (contributors > 1) {
                hanging.push_back(local);
            }
        }
        return hanging;
    }

    /** @brief Evaluate an MFEM field at native restriction representatives, without backend trace or reduction helpers. */
    [[nodiscard]] Vector native_boundary_samples(const mfem::FiniteElementSpace& row_space,
        const mfem::Array<int>& independent_rows, const mfem::GridFunction& field) {
        const mfem::SparseMatrix* restriction = row_space.GetConformingRestriction();
        REQUIRE(restriction != nullptr);
        std::vector<int> selected(row_space.GetVSize(), -1);
        for (int row = 0; row < independent_rows.Size(); ++row) {
            const int independent = independent_rows[row];
            const int entry = restriction->GetI()[independent];
            REQUIRE(restriction->GetI()[independent + 1] == entry + 1);
            REQUIRE(restriction->GetData()[entry] == 1.0);
            selected[restriction->GetJ()[entry]] = row;
        }
        Vector samples(independent_rows.Size());
        samples = std::numeric_limits<double>::quiet_NaN();
        mfem::Array<int> dofs;
        for (int element = 0; element < row_space.GetNE(); ++element) {
            row_space.GetElementDofs(element, dofs);
            const mfem::IntegrationRule& nodes = row_space.GetFE(element)->GetNodes();
            for (int local = 0; local < dofs.Size(); ++local) {
                const int row = selected[mfem::FiniteElementSpace::DecodeDof(dofs[local])];
                if (row < 0 || std::isfinite(samples[row])) {
                    continue;
                }
                samples[row] = field.GetValue(element, nodes.IntPoint(local));
            }
        }
        for (int row = 0; row < samples.Size(); ++row) {
            REQUIRE(std::isfinite(samples[row]));
        }
        return samples;
    }
} // namespace

TEST_CASE("Nonconforming boundary interpolation retains weighted traces across scalar orders", "[unit][equations][pressure_boundary][nonconforming]") {
    namespace finite_elements = serif::discretization::backend::mfem;
    namespace boundary = serif::tests::surface::boundary;
    using StellarSurface = serif::discretization::domain::StellarSurfaceBoundary;

    mfem::Mesh mesh = boundary::make_boundary_mesh();
    mesh.EnsureNCMesh();
    const mfem::Array<int> refinement{0};
    mesh.GeneralRefinement(refinement);
    mfem::H1_FECollection cubic_collection(3, 3), quadratic_collection(2, 3);
    mfem::FiniteElementSpace row_space(&mesh, &cubic_collection), input_space(&mesh, &quadratic_collection);
    finite_elements::TermBinding<Enthalpy, boundary::Schema> row_binding(row_space);
    finite_elements::TermBinding<boundary::QuadraticComposition, boundary::Schema> input_binding(input_space);
    REQUIRE(row_space.GetTrueVSize() < row_space.GetVSize());
    REQUIRE(input_space.GetTrueVSize() < input_space.GetVSize());
    // Every cell is material, so supported coordinates have the native true ordering.
    REQUIRE(row_binding.size() == row_space.GetTrueVSize());
    REQUIRE(input_binding.size() == input_space.GetTrueVSize());
    const mfem::Array<int> attributes{1, 0};
    REQUIRE_FALSE(hanging_boundary_dofs(row_space, attributes).empty());
    const std::vector<int> input_hanging = hanging_boundary_dofs(input_space, attributes);
    REQUIRE_FALSE(input_hanging.empty());

    const auto selected = std::make_shared<const finite_elements::DofSelection>(
        finite_elements::select_boundary_dofs<StellarSurface>(row_binding));
    mfem::Array<int> native_selected;
    row_space.GetEssentialTrueDofs(attributes, native_selected);
    REQUIRE(selected->size() == native_selected.Size());
    for (int row = 0; row < selected->size(); ++row) {
        CHECK(selected->reduced_dofs()[row] == native_selected[row]);
    }
    const auto trace = finite_elements::make_boundary_trace(selected, input_binding);
    const Vector values = boundary::sequence(input_binding.size(), 0.8, 0.37);
    mfem::GridFunction field(&input_space);
    field.SetFromTrueDofs(values);
    const Vector expected = native_boundary_samples(row_space, native_selected, field);
    Vector actual(selected->size());
    trace.interpolation().Mult(values, actual);
    CHECK(test::relative_error(actual, expected) < 2.0e-12);

    // This deliberately broken local copy proves that weighted boundary inputs affect the tested rows.
    mfem::GridFunction missing_hanging_values(&input_space);
    missing_hanging_values = field;
    for (const int local : input_hanging) {
        missing_hanging_values[local] = 0.0;
    }
    Vector omitted = native_boundary_samples(row_space, native_selected, missing_hanging_values);
    omitted -= expected;
    REQUIRE(omitted.Norml2() > 1.0e-3);

    const Vector dual = boundary::sequence(selected->size(), -0.2, 0.81);
    Vector transpose(input_binding.size());
    trace.interpolation().MultTranspose(dual, transpose);
    CHECK(std::abs(values * transpose - actual * dual) < 2.0e-12 * std::max(1.0, std::abs(actual * dual)));

    Vector zero_trace(row_binding.size());
    zero_trace = 1.0;
    for (const int row : selected->reduced_dofs()) {
        zero_trace[row] = 0.0;
    }
    mfem::GridFunction constrained(&row_space);
    constrained.SetFromTrueDofs(zero_trace);
    double unselected_boundary_maximum = 0.0;
    for (int face = 0; face < mesh.GetNBE(); ++face) {
        mfem::ElementTransformation& transformation = *mesh.GetBdrElementTransformation(face);
        const mfem::IntegrationRule& rule = mfem::IntRules.Get(transformation.GetGeometryType(), 8);
        for (int sample = 0; sample < rule.GetNPoints(); ++sample) {
            const double value = constrained.GetValue(transformation, rule.IntPoint(sample));
            if (mesh.GetBdrAttribute(face) == 1) {
                CHECK(std::abs(value) < 2.0e-13);
            } else {
                unselected_boundary_maximum = std::max(unselected_boundary_maximum, std::abs(value));
            }
        }
    }
    CHECK(unselected_boundary_maximum > 0.5);
}