#include <array>
#include <memory>
#include <span>
#include <vector>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/model/deformation.hpp"
#include "serif/backend/mfem/model/system.hpp"

namespace {
    namespace deformation = serif::geometry::deformation;
    namespace backend = serif::model::backend::mfem;
    using Catch::Matchers::WithinAbs;

    /** @brief A nonlinear shape with independent mixed curvature and deliberately move-only scratch. */
    struct Shape {
        using InputSpace = deformation::spaces::ShapeParameters;
        using OutputSpace = deformation::spaces::VolumeDisplacement;
        using Workspace = std::unique_ptr<int>;
        static constexpr bool is_linear = false;

        /** @brief Two independent shape coefficients determine three displacement coefficients. */
        [[nodiscard]] std::size_t input_size() const {
            return 2;
        }
        /** @brief Keep the physical force dimension different from the shape dimension. */
        [[nodiscard]] std::size_t output_size() const {
            return 3;
        }
        /** @brief Ensure model adapters do not impose an undocumented copyable-workspace requirement. */
        [[nodiscard]] Workspace make_workspace() const {
            return std::make_unique<int>(0);
        }

        /** @brief Evaluate u(q) = (q0^2 + q1, q0*q1, q1^2). */
        [[nodiscard]] deformation::DeformationStatus build_displacement(std::span<const double> q, std::span<double> u, Workspace&) const {
            u[0] = q[0] * q[0] + q[1];
            u[1] = q[0] * q[1];
            u[2] = q[1] * q[1];
            return {};
        }

        /** @brief Apply the analytic displacement derivative. */
        [[nodiscard]] deformation::DeformationStatus apply_jacobian(std::span<const double> q, std::span<const double> direction, std::span<double> output, Workspace&) const {
            output[0] = 2.0 * q[0] * direction[0] + direction[1];
            output[1] = q[1] * direction[0] + q[0] * direction[1];
            output[2] = 2.0 * q[1] * direction[1];
            return {};
        }

        /** @brief Pull a physical force back into the two shape coordinates. */
        [[nodiscard]] deformation::DeformationStatus apply_jacobian_transpose(std::span<const double> q, std::span<const double> force, std::span<double> output, Workspace&) const {
            output[0] = 2.0 * q[0] * force[0] + q[1] * force[1];
            output[1] = force[0] + q[0] * force[1] + 2.0 * q[1] * force[2];
            return {};
        }

        /** @brief Retain both pure and mixed second derivatives in the shape Newton action. */
        [[nodiscard]] deformation::DeformationStatus apply_pullback_derivative(std::span<const double>, std::span<const double> direction, std::span<const double> force, std::span<double> output, Workspace&) const {
            output[0] = 2.0 * force[0] * direction[0] + force[1] * direction[1];
            output[1] = force[1] * direction[0] + 2.0 * force[2] * direction[1];
            return {};
        }
    };

    /** @brief Use a nonsymmetric physical force derivative so the full transpose cannot be replaced by the forward action. */
    [[nodiscard]] std::shared_ptr<const ::mfem::DenseMatrix> force_derivative() {
        auto matrix = std::make_shared<::mfem::DenseMatrix>(3);
        const std::array<std::array<double, 3>, 3> rows{{{2, -1, 0.5}, {0.2, 3, -0.7}, {1, 0.3, 4}}};
        for (int row = 0; row < 3; ++row) {
            for (int column = 0; column < 3; ++column) {
                (*matrix)(row, column) = rows[row][column];
            }
        }
        return matrix;
    }

    /** @brief Assemble E^T force and E^T K E using only the public model adapters. */
    [[nodiscard]] std::shared_ptr<const backend::PhysicalSystem> evaluate(const std::array<double, 2>& coordinates) {
        const auto map = std::make_shared<const Shape>();
        const auto state = std::make_shared<const deformation::DeformationLinearization<Shape>>(deformation::DeformationLinearization<Shape>{map, {coordinates[0], coordinates[1]}});
        auto workspace = map->make_workspace();
        ::mfem::Vector displacement(3), force(3), residual;
        REQUIRE(map->build_displacement(coordinates, std::span<double>(displacement.GetData(), 3), workspace));
        const auto matrix = force_derivative();
        matrix->Mult(displacement, force);
        force[0] += 0.7;
        force[1] -= 0.2;
        force[2] += 0.5;
        const auto tangent = backend::deformation_tangent(state);
        const auto transpose = serif::operators::backend::mfem::transpose_block(tangent);
        transpose.action()->Mult(force, residual);
        const auto force_block = serif::operators::backend::mfem::linearized_block(matrix);
        const auto material = serif::operators::backend::mfem::compose_blocks(transpose, serif::operators::backend::mfem::compose_blocks(force_block, tangent));
        return backend::make_physical_system(2, 2, {{0, std::vector<double>(residual.GetData(), residual.GetData() + 2)}}, {{0, 0, material}});
    }
} // namespace

TEST_CASE("Model shape tangent composes with a nonsymmetric force derivative as an exact transpose pair", "[model][runtime][deformation][unit]") {
    const std::array<double, 2> coordinates{0.4, -0.3};
    const auto accepted = evaluate(coordinates);
    ::mfem::Vector direction(2), action, dual(2), transpose;
    direction[0] = 0.7;
    direction[1] = -0.2;
    dual[0] = -0.8;
    dual[1] = 1.3;
    accepted->jacobian().action()->Mult(direction, action);
    accepted->jacobian().action()->MultTranspose(dual, transpose);
    CHECK_THAT(action * dual, WithinAbs(direction * transpose, 2.0e-13));
    ::mfem::Vector repeated;
    accepted->jacobian().action()->Mult(direction, repeated);
    for (int row = 0; row < 2; ++row) {
        CHECK(repeated[row] == action[row]);
    }
}