#include <array>
#include <memory>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace {
    using ScalarTable = mean_field::fem::ScalarReferenceTable;
    using VectorTable = mean_field::fem::VectorReferenceTable;
    using TableCache  = mean_field::fem::ReferenceTableCache;

    static_assert(std::is_same_v<
                  decltype(std::declval<const TableCache &>().GetScalarTable(
                      std::declval<const mfem::FiniteElement &>(),
                      std::declval<const mfem::IntegrationRule &>()
                  )),
                  std::shared_ptr<const ScalarTable>>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const ScalarTable &>().GetValues()),
                  const mfem::DenseMatrix &>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const ScalarTable &>().GetGradients(0)),
                  const mfem::DenseMatrix &>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const TableCache &>().GetVectorTable(
                      std::declval<const mfem::FiniteElement &>(),
                      std::declval<const mfem::IntegrationRule &>()
                  )),
                  std::shared_ptr<const VectorTable>>);
    static_assert(std::is_same_v<
                  decltype(std::declval<const VectorTable &>().GetValues(0)),
                  const mfem::DenseMatrix &>);

    void CheckMatrixExactly(
        const mfem::DenseMatrix &actual,
        const mfem::DenseMatrix &expected
    ) {
        REQUIRE(actual.Height() == expected.Height());
        REQUIRE(actual.Width() == expected.Width());
        for (int column = 0; column < actual.Width(); ++column) {
            for (int row = 0; row < actual.Height(); ++row) {
                CHECK(actual(row, column) == expected(row, column));
            }
        }
    }

    void CheckScalarTable(
        const ScalarTable &table,
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) {
        REQUIRE(table.GetPointCount() == rule.GetNPoints());
        REQUIRE(table.GetDofCount() == element.GetDof());
        REQUIRE(table.GetDimension() == element.GetDim());
        REQUIRE(table.GetValues().Height() == rule.GetNPoints());
        REQUIRE(table.GetValues().Width() == element.GetDof());

        mfem::Vector shape(element.GetDof());
        mfem::DenseMatrix gradient(element.GetDof(), element.GetDim());
        for (int point = 0; point < rule.GetNPoints(); ++point) {
            element.CalcShape(rule.IntPoint(point), shape);
            element.CalcDShape(rule.IntPoint(point), gradient);
            for (int dof = 0; dof < element.GetDof(); ++dof) {
                CHECK(table.GetValues()(point, dof) == shape(dof));
            }
            CheckMatrixExactly(table.GetGradients(point), gradient);
        }
    }

    void CheckVectorTable(
        const VectorTable &table,
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) {
        REQUIRE(table.GetPointCount() == rule.GetNPoints());
        REQUIRE(table.GetDofCount() == element.GetDof());
        REQUIRE(table.GetDimension() == element.GetRangeDim());
        mfem::DenseMatrix shape(element.GetDof(), element.GetRangeDim());
        for (int point = 0; point < rule.GetNPoints(); ++point) {
            element.CalcVShape(rule.IntPoint(point), shape);
            CheckMatrixExactly(table.GetValues(point), shape);
        }
    }

    mfem::IntegrationRule CopyRule(const mfem::IntegrationRule &source) {
        mfem::IntegrationRule copy(source.GetNPoints());
        copy.SetOrder(source.GetOrder());
        for (int point = 0; point < source.GetNPoints(); ++point) {
            copy.IntPoint(point) = source.IntPoint(point);
        }
        return copy;
    }
} // namespace

TEST_CASE(
    "Reference Table Cache Matches Scalar MFEM Values And Gradients",
    tags::unit &tags::quadrature
) {
    for (const int dimension : std::array{2, 3}) {
        const auto geometry = dimension == 2 ? mfem::Geometry::SQUARE : mfem::Geometry::CUBE;
        for (const int order : std::array{1, 3}) {
            CAPTURE(dimension, order);
            mfem::H1_FECollection h1(order, dimension);
            mfem::L2_FECollection l2(order - 1, dimension);
            TableCache cache;
            const mfem::IntegrationRule &rule = mfem::IntRules.Get(geometry, 2 * order + 1);
            for (const mfem::FiniteElement *element :
                 std::array{h1.FiniteElementForGeometry(geometry), l2.FiniteElementForGeometry(geometry)}) {
                REQUIRE(element != nullptr);
                const auto table = cache.GetScalarTable(*element, rule);
                REQUIRE(table != nullptr);
                CheckScalarTable(*table, *element, rule);
            }
        }
    }
}

TEST_CASE(
    "Reference Table Cache Matches RT Reference Values And Shares Equal Rules",
    tags::unit &tags::quadrature
) {
    for (const int dimension : std::array{2, 3}) {
        const auto geometry = dimension == 2 ? mfem::Geometry::SQUARE : mfem::Geometry::CUBE;
        for (const int order : std::array{0, 2}) {
            CAPTURE(dimension, order);
            mfem::RT_FECollection standard(order, dimension);
            mfem::RT_FECollection integrated(
                order, dimension, mfem::BasisType::GaussLobatto, mfem::BasisType::IntegratedGLL
            );
            const mfem::FiniteElement &standardElement   = *standard.FiniteElementForGeometry(geometry);
            const mfem::FiniteElement &integratedElement = *integrated.FiniteElementForGeometry(geometry);
            const mfem::IntegrationRule &rule            = mfem::IntRules.Get(geometry, 2 * order + 3);
            const mfem::IntegrationRule copiedRule       = CopyRule(rule);
            const TableCache cache;
            const auto standardTable   = cache.GetVectorTable(standardElement, rule);
            const auto integratedTable = cache.GetVectorTable(integratedElement, rule);
            REQUIRE(standardTable != nullptr);
            REQUIRE(integratedTable != nullptr);
            CHECK(cache.GetVectorTable(standardElement, copiedRule).get() == standardTable.get());
            CHECK(cache.GetVectorTable(integratedElement, copiedRule).get() == integratedTable.get());
            CHECK(standardTable.get() != integratedTable.get());
            CheckVectorTable(*standardTable, standardElement, rule);
            CheckVectorTable(*integratedTable, integratedElement, rule);
        }
    }
}

TEST_CASE(
    "Reference Table Cache Shares Equal Rules And Distinguishes Rule Contents",
    tags::unit &tags::quadrature
) {
    mfem::H1_FECollection collection(3, 3);
    const mfem::FiniteElement &element      = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
    const mfem::IntegrationRule &rule       = mfem::IntRules.Get(mfem::Geometry::CUBE, 7);
    mfem::IntegrationRule copiedRule        = CopyRule(rule);
    mfem::IntegrationRule movedPointRule    = CopyRule(rule);
    mfem::IntegrationRule changedWeightRule = CopyRule(rule);
    movedPointRule.IntPoint(0).x += 0.03125;
    changedWeightRule.IntPoint(0).weight *= 1.25;

    const TableCache cache;
    const auto original = cache.GetScalarTable(element, rule);
    const mfem::DenseMatrix originalValues(original->GetValues());
    const auto copy          = cache.GetScalarTable(element, copiedRule);
    const auto movedPoint    = cache.GetScalarTable(element, movedPointRule);
    const auto changedWeight = cache.GetScalarTable(element, changedWeightRule);

    CHECK(copy.get() == original.get());
    CHECK(movedPointRule.GetOrder() == rule.GetOrder());
    CHECK(changedWeightRule.GetOrder() == rule.GetOrder());
    CHECK(movedPoint.get() != original.get());
    CHECK(changedWeight.get() != original.get());
    CHECK(changedWeight.get() != movedPoint.get());
    CheckScalarTable(*movedPoint, element, movedPointRule);
    CheckScalarTable(*changedWeight, element, changedWeightRule);
    CheckMatrixExactly(original->GetValues(), originalValues);

    // Rule identity is its contents, not its address, including after mutation.
    copiedRule.IntPoint(0).x = movedPointRule.IntPoint(0).x;
    CHECK(cache.GetScalarTable(element, copiedRule).get() == movedPoint.get());
    CHECK(cache.GetScalarTable(element, rule).get() == original.get());
}

TEST_CASE(
    "Reference Table Cache Distinguishes Scalar Basis Variants",
    tags::unit &tags::quadrature
) {
    constexpr int dimension = 3;
    constexpr int order     = 3;
    mfem::H1_FECollection nodalH1(order, dimension, mfem::BasisType::GaussLobatto);
    mfem::H1_FECollection positiveH1(order, dimension, mfem::BasisType::Positive);
    mfem::L2_FECollection openL2(order, dimension, mfem::BasisType::GaussLegendre);
    mfem::L2_FECollection closedL2(order, dimension, mfem::BasisType::GaussLobatto);
    const mfem::IntegrationRule &rule = mfem::IntRules.Get(mfem::Geometry::CUBE, 5);
    const TableCache cache;
    std::array<std::shared_ptr<const ScalarTable>, 4> tables;
    const std::array<const mfem::FiniteElement *, 4> elements{
        nodalH1.FiniteElementForGeometry(mfem::Geometry::CUBE),
        positiveH1.FiniteElementForGeometry(mfem::Geometry::CUBE),
        openL2.FiniteElementForGeometry(mfem::Geometry::CUBE), closedL2.FiniteElementForGeometry(mfem::Geometry::CUBE)
    };
    for (std::size_t index = 0; index < elements.size(); ++index) {
        REQUIRE(elements[index] != nullptr);
        REQUIRE(elements[index]->GetOrder() == order);
        REQUIRE(elements[index]->GetDof() == elements[0]->GetDof());
        tables[index] = cache.GetScalarTable(*elements[index], rule);
        CheckScalarTable(*tables[index], *elements[index], rule);
        for (std::size_t previous = 0; previous < index; ++previous) {
            CHECK(tables[index].get() != tables[previous].get());
        }
    }
}

TEST_CASE(
    "Reference Table Cache Published Scalar Storage Outlives Its Cache",
    tags::unit &tags::quadrature
) {
    std::shared_ptr<const ScalarTable> retained;
    mfem::DenseMatrix expectedValues;
    mfem::DenseMatrix expectedGradient;
    {
        // FE objects remain immutable and alive throughout the cache lifetime.
        mfem::H1_FECollection collection(3, 2);
        const mfem::FiniteElement &element = *collection.FiniteElementForGeometry(mfem::Geometry::SQUARE);
        const mfem::IntegrationRule rule   = CopyRule(mfem::IntRules.Get(mfem::Geometry::SQUARE, 7));
        const TableCache cache;
        retained = cache.GetScalarTable(element, rule);
        CheckScalarTable(*retained, element, rule);
        expectedValues   = retained->GetValues();
        expectedGradient = retained->GetGradients(0);
    }
    REQUIRE(retained != nullptr);
    CheckMatrixExactly(retained->GetValues(), expectedValues);
    CheckMatrixExactly(retained->GetGradients(0), expectedGradient);
}

TEST_CASE(
    "Reference Table Cache Published RT Storage Outlives Its Cache",
    tags::unit &tags::quadrature
) {
    std::shared_ptr<const VectorTable> retained;
    mfem::DenseMatrix firstExpected;
    mfem::DenseMatrix lastExpected;
    {
        mfem::RT_FECollection collection(2, 3, mfem::BasisType::GaussLobatto, mfem::BasisType::IntegratedGLL);
        const mfem::FiniteElement &element = *collection.FiniteElementForGeometry(mfem::Geometry::CUBE);
        const mfem::IntegrationRule rule   = CopyRule(mfem::IntRules.Get(mfem::Geometry::CUBE, 7));
        const TableCache cache;
        retained = cache.GetVectorTable(element, rule);
        CheckVectorTable(*retained, element, rule);
        firstExpected = retained->GetValues(0);
        lastExpected  = retained->GetValues(rule.GetNPoints() - 1);
    }
    REQUIRE(retained != nullptr);
    CheckMatrixExactly(retained->GetValues(0), firstExpected);
    CheckMatrixExactly(retained->GetValues(retained->GetPointCount() - 1), lastExpected);
}