#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <utility>
#include <vector>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>

import mean_field;

namespace {
    using namespace mean_field::solver;

    class DenseOperation final : public mfem::Operator {
    public:
        explicit DenseOperation(mfem::DenseMatrix matrix)
            : mfem::Operator(matrix.Height(), matrix.Width()),
              m_matrix(std::move(matrix)) {
        }

        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            output.SetSize(Height());
            m_matrix.Mult(input, output);
            ++applications;
        }

        [[nodiscard]] const mfem::DenseMatrix &Matrix() const noexcept {
            return m_matrix;
        }

        mutable int applications{0};

    private:
        mfem::DenseMatrix m_matrix;
    };

    class DenseGravitySubsystem final {
    public:
        DenseGravitySubsystem(mfem::DenseMatrix jacobian, GravityBlockPartition partition)
            : m_operation(std::move(jacobian)),
              m_partition(std::move(partition)),
              m_gravityMatrix(m_partition.GravityResidualSize(), m_partition.GravityValueSize()) {
            mfem::Vector gravityDirection(m_partition.GravityValueSize());
            mfem::Vector fullDirection(m_partition.RootValueSize());
            mfem::Vector fullAction(m_partition.RootResidualSize());
            mfem::Vector gravityAction(m_partition.GravityResidualSize());
            for (int column = 0; column < gravityDirection.Size(); ++column) {
                gravityDirection = 0.0;
                gravityDirection(column) = 1.0;
                m_partition.InjectGravityValues(gravityDirection, fullDirection);
                m_operation.Mult(fullDirection, fullAction);
                m_partition.GatherGravityResiduals(fullAction, gravityAction);
                for (int row = 0; row < gravityAction.Size(); ++row) {
                    m_gravityMatrix(row, column) = gravityAction(row);
                }
            }
        }

        [[nodiscard]] const GravityBlockPartition &GetPartition() const noexcept { return m_partition; }
        [[nodiscard]] const DenseOperation &GetFullOperator() const noexcept { return m_operation; }
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept { return MPI_COMM_SELF; }
        [[nodiscard]] bool IsCurrent() const noexcept { return current; }

        [[nodiscard]] LinearSolveReport Solve(
            const mfem::Vector &rightHandSide,
            mfem::Vector &correction,
            const LinearSolveControl &control
        ) {
            CHECK(correction.Norml2() == 0.0);
            ++solves;
            mfem::DenseMatrixInverse inverse(m_gravityMatrix);
            inverse.Mult(rightHandSide, correction);
            mfem::Vector defect(rightHandSide.Size());
            m_gravityMatrix.Mult(correction, defect);
            defect -= rightHandSide;
            const double rhsNorm = rightHandSide.Norml2();
            const double trueNorm = forcedTrueNorm >= 0.0 ? forcedTrueNorm : defect.Norml2();
            return {
                .status = status,
                .control = control,
                .iterations = 1,
                .rightHandSideNorm = rhsNorm,
                .trueResidualNorm = trueNorm,
                .relativeTrueResidualNorm = rhsNorm > 0.0 ? trueNorm / rhsNorm : 0.0
            };
        }

        bool current{true};
        int solves{0};
        LinearSolveStatus status{LinearSolveStatus::converged};
        double forcedTrueNorm{-1.0};

    private:
        DenseOperation m_operation;
        GravityBlockPartition m_partition;
        mfem::DenseMatrix m_gravityMatrix;
    };

    template <mean_field::preconditioning::ApplicationContract Contract>
    class ExactFullInverse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = Contract;

        explicit ExactFullInverse(const mfem::DenseMatrix &matrix)
            : mfem::Solver(matrix.Width(), matrix.Height(), false), m_matrix(matrix) {
        }

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Width() || operation.Width() != Height()) {
                throw std::invalid_argument("The dense test inverse requires matching dimensions.");
            }
        }

        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            output.SetSize(Height());
            mfem::DenseMatrixInverse inverse(m_matrix);
            inverse.Mult(input, output);
        }

    private:
        mfem::DenseMatrix m_matrix;
    };

    struct MissingGravityContract { };
    static_assert(PreparedGravityReductionSubsystem<DenseGravitySubsystem>);
    static_assert(!PreparedGravityReductionSubsystem<MissingGravityContract>);
    using StationaryFullInverse = ExactFullInverse<mean_field::preconditioning::ApplicationContract::stationary_linear>;
    using FlexibleFullInverse = ExactFullInverse<mean_field::preconditioning::ApplicationContract::flexible>;
    using StationaryReducedInverse = GravityReducedPreconditioner<DenseGravitySubsystem, StationaryFullInverse>;
    using FlexibleReducedInverse = GravityReducedPreconditioner<DenseGravitySubsystem, FlexibleFullInverse>;
    static_assert(linearPreconditionerApplicationContract<StationaryReducedInverse> ==
                  mean_field::preconditioning::ApplicationContract::stationary_linear);
    static_assert(linearPreconditionerApplicationContract<FlexibleReducedInverse> ==
                  mean_field::preconditioning::ApplicationContract::flexible);

    constexpr LinearSolveControl innerControl{.relativeTolerance = 1.0e-12, .absoluteTolerance = 1.0e-14,
                                              .maximumIterations = 20};

    [[nodiscard]] GravityBlockPartition testPartition(const int size = 4) {
        // Different row and column permutations expose accidental reliance on
        // the physical root's present ordering, including its appended borders.
        return GravityBlockPartition(size, size, {2, 0}, {1, 3});
    }

    [[nodiscard]] mfem::Vector manifoldState(const double first, const double second) {
        mfem::Vector state(4);
        state(0) = std::sin(second) / (1.0 + first * first);
        state(1) = first;
        state(2) = first * second;
        state(3) = second;
        return state;
    }

    [[nodiscard]] mfem::Vector fullResidual(const mfem::Vector &state) {
        const double first = state(1);
        const double second = state(3);
        const double gravityFirst = state(2);
        const double gravitySecond = state(0);
        const auto equilibrium = manifoldState(0.5, 0.25);
        mfem::Vector residual(4);
        residual(0) = second + 2.0 * gravityFirst - gravitySecond -
                      (0.25 + 2.0 * equilibrium(2) - equilibrium(0));
        residual(1) = gravityFirst - first * second;
        residual(2) = first * first + gravityFirst + 3.0 * gravitySecond -
                      (0.25 + equilibrium(2) + 3.0 * equilibrium(0));
        residual(3) = (1.0 + first * first) * gravitySecond - std::sin(second);
        return residual;
    }

    [[nodiscard]] mfem::DenseMatrix fullJacobian(const mfem::Vector &state, const bool border = false) {
        const double first = state(1);
        const double second = state(3);
        mfem::DenseMatrix matrix(border ? 5 : 4);
        matrix = 0.0;
        matrix(0, 0) = -1.0;
        matrix(0, 2) = 2.0;
        matrix(0, 3) = 1.0;
        matrix(1, 1) = -second;
        matrix(1, 2) = 1.0;
        matrix(1, 3) = -first;
        matrix(2, 0) = 3.0;
        matrix(2, 1) = 2.0 * first;
        matrix(2, 2) = 1.0;
        matrix(3, 0) = 1.0 + first * first;
        matrix(3, 1) = 2.0 * first * state(0);
        matrix(3, 3) = -std::cos(second);
        if (border) {
            matrix(4, 4) = 2.0;
            matrix(4, 0) = 0.25;
            matrix(1, 4) = -0.3;
            matrix(2, 4) = 0.1;
            matrix(4, 3) = 0.4;
        }
        return matrix;
    }

    [[nodiscard]] double relativeError(const mfem::Vector &actual, const mfem::Vector &expected) {
        mfem::Vector difference(actual);
        difference -= expected;
        return difference.Norml2() / std::max(1.0, expected.Norml2());
    }

    [[nodiscard]] mfem::DenseMatrix explicitReduced(const auto &operation) {
        mfem::DenseMatrix matrix(operation.Height(), operation.Width());
        mfem::Vector basis(operation.Width());
        mfem::Vector action(operation.Height());
        for (int column = 0; column < operation.Width(); ++column) {
            basis = 0.0;
            basis(column) = 1.0;
            operation.Mult(basis, action);
            for (int row = 0; row < operation.Height(); ++row) {
                matrix(row, column) = action(row);
            }
        }
        return matrix;
    }
} // namespace

TEST_CASE("Gravity reduction follows the curved gravity manifold", "[solver][gravity_reduction]") {
    const auto state = manifoldState(0.7, -0.4);
    DenseGravitySubsystem subsystem(fullJacobian(state), testPartition());
    GravityReducedOperator operation(subsystem, innerControl);
    mfem::Vector direction(2);
    direction(0) = 0.3;
    direction(1) = -0.2;
    mfem::Vector action(2);
    operation.Mult(direction, action);

    double previousError = 0.0;
    for (const double step : {1.0e-2, 1.0e-3, 1.0e-4}) {
        const auto plus = fullResidual(manifoldState(state(1) + step * direction(0),
                                                     state(3) + step * direction(1)));
        const auto minus = fullResidual(manifoldState(state(1) - step * direction(0),
                                                      state(3) - step * direction(1)));
        mfem::Vector difference(plus);
        difference -= minus;
        difference /= 2.0 * step;
        mfem::Vector reducedDifference(2);
        subsystem.GetPartition().GatherRetainedResiduals(difference, reducedDifference);
        const double error = relativeError(reducedDifference, action);
        if (previousError > 0.0) {
            CHECK(error < previousError / 80.0);
        }
        previousError = error;
    }
    CHECK(previousError < 5.0e-10);

    mfem::Vector fullDirection(4);
    REQUIRE(operation.TryLiftTangent(direction, fullDirection).Converged());
    mfem::Vector fullAction(4);
    subsystem.GetFullOperator().Mult(fullDirection, fullAction);
    mfem::Vector gravityAction(2), retainedAction(2);
    subsystem.GetPartition().GatherGravityResiduals(fullAction, gravityAction);
    subsystem.GetPartition().GatherRetainedResiduals(fullAction, retainedAction);
    CHECK(gravityAction.Norml2() < 1.0e-14);
    CHECK(relativeError(retainedAction, action) < 1.0e-14);

    mfem::Vector frozenDirection(4), frozenAction(4), frozenRetained(2);
    subsystem.GetPartition().InjectRetainedValues(direction, frozenDirection);
    subsystem.GetFullOperator().Mult(frozenDirection, frozenAction);
    subsystem.GetPartition().GatherRetainedResiduals(frozenAction, frozenRetained);
    CHECK(relativeError(frozenRetained, action) > 0.1);
}

TEST_CASE("Gravity condensation and lifting recover the monolithic Newton correction", "[solver][gravity_reduction]") {
    for (const bool border : {false, true}) {
        CAPTURE(border);
        const auto matrix = fullJacobian(manifoldState(0.7, -0.4), border);
        const int size = matrix.Height();
        DenseGravitySubsystem subsystem(matrix, testPartition(size));
        GravityReducedOperator operation(subsystem, innerControl);
        const auto reducedMatrix = explicitReduced(operation);
        mfem::DenseMatrixInverse reducedInverse(reducedMatrix);
        mfem::DenseMatrixInverse fullInverse(matrix);
        mfem::Vector fullRightHandSide(size), condensed(size - 2), retainedCorrection(size - 2);
        mfem::Vector fullCorrection(size), reference(size);
        for (int index = 0; index < size; ++index) {
            fullRightHandSide(index) = 0.4 - 0.17 * static_cast<double>(index);
        }
        REQUIRE(operation.TryCondenseRightHandSide(fullRightHandSide, condensed).Converged());
        reducedInverse.Mult(condensed, retainedCorrection);
        REQUIRE(operation.TryLiftCorrection(retainedCorrection, fullRightHandSide, fullCorrection).Converged());
        fullInverse.Mult(fullRightHandSide, reference);
        CHECK(relativeError(fullCorrection, reference) < 2.0e-13);
        mfem::Vector checked(size);
        matrix.Mult(fullCorrection, checked);
        CHECK(relativeError(checked, fullRightHandSide) < 2.0e-14);
    }
}

TEST_CASE("The projected full inverse preconditions the reduced Schur operator", "[solver][gravity_reduction]") {
    const auto matrix = fullJacobian(manifoldState(0.7, -0.4), true);
    DenseGravitySubsystem subsystem(matrix, testPartition(5));
    GravityReducedOperator operation(subsystem, innerControl);
    StationaryFullInverse fullInverse(matrix);
    GravityReducedPreconditioner inverse(operation, fullInverse);
    inverse.SetOperator(operation);
    mfem::Vector rightHandSide(3), correction(3), actual(3);
    rightHandSide(0) = 0.4;
    rightHandSide(1) = -0.3;
    rightHandSide(2) = 0.8;
    inverse.Mult(rightHandSide, correction);
    operation.Mult(correction, actual);
    CHECK(relativeError(actual, rightHandSide) < 2.0e-14);
    CHECK_THROWS_AS(inverse.SetOperator(subsystem.GetFullOperator()), std::invalid_argument);
}

TEST_CASE("Reduced and monolithic Newton converge to the same curved-system equilibrium", "[solver][gravity_reduction]") {
    // This system has a second equilibrium near (0.713, 0.171). The old
    // (0.6, 0.2) seed lay near the intervening fold and selected that root.
    // Test local Newton equivalence in the known root's basin, rather than
    // assuming uniqueness or comparing a fixed iteration count at two rates.
    auto reducedState = manifoldState(0.52, 0.24);
    auto monolithicState = reducedState;
    for (int iteration = 0; iteration < 8; ++iteration) {
        if (fullResidual(reducedState).Norml2() < 2.0e-14 &&
            fullResidual(monolithicState).Norml2() < 2.0e-14) {
            break;
        }
        DenseGravitySubsystem subsystem(fullJacobian(reducedState), testPartition());
        GravityReducedOperator operation(subsystem, innerControl);
        auto rhs = fullResidual(reducedState);
        rhs *= -1.0;
        mfem::Vector condensed(2), reducedCorrection(2);
        REQUIRE(operation.TryCondenseRightHandSide(rhs, condensed).Converged());
        const auto reducedMatrix = explicitReduced(operation);
        mfem::DenseMatrixInverse inverse(reducedMatrix);
        inverse.Mult(condensed, reducedCorrection);
        reducedState = manifoldState(reducedState(1) + reducedCorrection(0),
                                     reducedState(3) + reducedCorrection(1));

        const auto matrix = fullJacobian(monolithicState);
        mfem::DenseMatrixInverse fullInverse(matrix);
        auto monolithicRhs = fullResidual(monolithicState);
        monolithicRhs *= -1.0;
        mfem::Vector fullCorrection(4);
        fullInverse.Mult(monolithicRhs, fullCorrection);
        monolithicState += fullCorrection;
        CHECK(reducedState(1) > 0.45);
        CHECK(reducedState(1) < 0.55);
        CHECK(monolithicState(1) > 0.45);
        CHECK(monolithicState(1) < 0.55);
    }
    const auto expected = manifoldState(0.5, 0.25);
    CHECK(fullResidual(reducedState).Norml2() < 2.0e-14);
    CHECK(fullResidual(monolithicState).Norml2() < 2.0e-14);
    CHECK(relativeError(reducedState, expected) < 2.0e-13);
    CHECK(relativeError(monolithicState, expected) < 2.0e-13);
}

TEST_CASE("Reduced gravity actions fail without exposing partial results", "[solver][gravity_reduction]") {
    DenseGravitySubsystem subsystem(fullJacobian(manifoldState(0.7, -0.4)), testPartition());
    GravityReducedOperator operation(subsystem, innerControl);
    mfem::Vector direction(2), output(2), fullOutput(4), fullRhs(4);
    direction = 0.3;
    output = 7.0;
    fullOutput = 9.0;
    fullRhs = 0.4;
    subsystem.status = LinearSolveStatus::maximum_iterations;
    CHECK_FALSE(operation.TryMult(direction, output).Converged());
    CHECK(output(0) == 7.0);
    CHECK(output(1) == 7.0);
    CHECK_THROWS_AS(operation.Mult(direction, output), GravityReductionFailure);
    CHECK_FALSE(operation.TryCondenseRightHandSide(fullRhs, output).Converged());
    CHECK_FALSE(operation.TryLiftCorrection(direction, fullRhs, fullOutput).Converged());
    CHECK_FALSE(operation.TryLiftTangent(direction, fullOutput).Converged());
    for (int index = 0; index < fullOutput.Size(); ++index) {
        CHECK(fullOutput(index) == 9.0);
    }
    subsystem.status = LinearSolveStatus::converged;
    subsystem.forcedTrueNorm = 1.0;
    CHECK(operation.TryMult(direction, output).status == LinearSolveStatus::backend_failure);
    subsystem.forcedTrueNorm = std::numeric_limits<double>::infinity();
    CHECK(operation.TryMult(direction, output).status == LinearSolveStatus::non_finite);
    CHECK(output(0) == 7.0);
}

TEST_CASE("Reduced gravity actions enforce preparation and dimension contracts", "[solver][gravity_reduction]") {
    DenseGravitySubsystem subsystem(fullJacobian(manifoldState(0.7, -0.4)), testPartition());
    GravityReducedOperator operation(subsystem, innerControl);
    StationaryFullInverse fullInverse(subsystem.GetFullOperator().Matrix());
    GravityReducedPreconditioner inverse(operation, fullInverse);
    mfem::Vector direction(2), output(2), badDirection(3), badRhs(3), fullOutput(4);
    direction = 0.2;
    badDirection = 0.2;
    badRhs = 0.2;
    CHECK_THROWS_AS(operation.TryMult(badDirection, output), std::invalid_argument);
    CHECK_THROWS_AS(operation.TryCondenseRightHandSide(badRhs, output), std::invalid_argument);
    CHECK_THROWS_AS(operation.TryLiftCorrection(direction, badRhs, fullOutput), std::invalid_argument);
    subsystem.current = false;
    const int applications = subsystem.GetFullOperator().applications;
    CHECK_THROWS_AS(operation.TryMult(direction, output), std::logic_error);
    CHECK_THROWS_AS(operation.TryLiftTangent(direction, fullOutput), std::logic_error);
    CHECK_THROWS_AS(inverse.Mult(direction, output), std::logic_error);
    CHECK(subsystem.GetFullOperator().applications == applications);
}

TEST_CASE("Reduced gravity action is repeatable, linear, and alias safe with an exact inner solve", "[solver][gravity_reduction]") {
    DenseGravitySubsystem subsystem(fullJacobian(manifoldState(0.7, -0.4)), testPartition());
    GravityReducedOperator operation(subsystem, innerControl);
    mfem::Vector first(2), second(2), sum(2), firstAction(2), secondAction(2), combinedAction(2);
    first(0) = 0.3;
    first(1) = -0.2;
    second(0) = -0.1;
    second(1) = 0.5;
    operation.Mult(first, firstAction);
    operation.Mult(second, secondAction);
    sum = first;
    sum.Add(1.7, second);
    operation.Mult(sum, combinedAction);
    firstAction.Add(1.7, secondAction);
    CHECK(relativeError(combinedAction, firstAction) < 1.0e-14);
    operation.Mult(sum, sum);
    CHECK(relativeError(sum, combinedAction) < 1.0e-14);
    operation.Mult(first, firstAction);
    operation.Mult(first, secondAction);
    CHECK(relativeError(firstAction, secondAction) == 0.0);
}

TEST_CASE("Prepared gravity assessment retains the scale of cancelling field contributions",
          "[solver][gravity_prepared_residual]") {
    using mean_field::solver::detail::AssessGravityPreparedResidual;
    const std::array contributions{1.e6, 1.e6, 1.e-8, 2.0, 3.0, 1.e-8, 1.e-8, 1.e-8};
    const double affineDisagreement = 1.e-8;
    const double measuredInnerResidual = 5.e-15;
    const auto result = AssessGravityPreparedResidual(contributions, affineDisagreement,
        affineDisagreement + measuredInnerResidual, measuredInnerResidual);
    REQUIRE(result.admissible);
    CHECK(result.arithmeticAllowance > affineDisagreement);
    CHECK(result.targetThreshold == result.arithmeticAllowance + measuredInnerResidual);

    const std::array residualOnly{1.e-8, 1.e-8, 1.e-8};
    CHECK_FALSE(AssessGravityPreparedResidual(residualOnly, affineDisagreement,
        affineDisagreement + measuredInnerResidual, measuredInnerResidual).admissible);
    CHECK_FALSE(AssessGravityPreparedResidual(contributions, 2.0 * result.arithmeticAllowance,
        measuredInnerResidual, measuredInnerResidual).admissible);
    CHECK_FALSE(AssessGravityPreparedResidual(contributions, 0.0,
        2.0 * result.targetThreshold, measuredInnerResidual).admissible);
}

TEST_CASE("Prepared gravity assessment scales with equation units without an absolute floor",
          "[solver][gravity_prepared_residual]") {
    using mean_field::solver::detail::AssessGravityPreparedResidual;
    const std::array contributions{1.e8, 1.e8, 5.0, 3.0, 2.0};
    const auto reference = AssessGravityPreparedResidual(contributions, 1.e-7, 1.5e-7, 1.e-8);
    REQUIRE(reference.admissible);
    for (const double factor : {1.e-12, 1.e12}) {
        auto scaled = contributions;
        for (double &value : scaled) value *= factor;
        const auto result = AssessGravityPreparedResidual(scaled, factor * 1.e-7, factor * 1.5e-7, factor * 1.e-8);
        CHECK(result.admissible == reference.admissible);
        CHECK(result.arithmeticAllowance / factor == Catch::Approx(reference.arithmeticAllowance).epsilon(1.e-14));
        CHECK(result.targetThreshold / factor == Catch::Approx(reference.targetThreshold).epsilon(1.e-14));
    }
    const std::array zero{0.0};
    const auto exactZero = AssessGravityPreparedResidual(zero, 0.0, 0.0, 0.0);
    CHECK(exactZero.admissible);
    CHECK(exactZero.arithmeticAllowance == 0.0);
    CHECK_FALSE(AssessGravityPreparedResidual(zero, 1.e-30, 0.0, 0.0).admissible);
}

TEST_CASE("Prepared gravity assessment rejects nonfinite or invalid measurements",
          "[solver][gravity_prepared_residual]") {
    using mean_field::solver::detail::AssessGravityPreparedResidual;
    const std::array contributions{1.0, 2.0};
    for (const double invalid : {-1.0, std::numeric_limits<double>::infinity(),
                                 std::numeric_limits<double>::quiet_NaN()}) {
        CHECK_FALSE(AssessGravityPreparedResidual(contributions, invalid, 0.0, 0.0).admissible);
        CHECK_FALSE(AssessGravityPreparedResidual(contributions, 0.0, invalid, 0.0).admissible);
        CHECK_FALSE(AssessGravityPreparedResidual(contributions, 0.0, 0.0, invalid).admissible);
        const std::array invalidContributions{1.0, invalid};
        CHECK_FALSE(AssessGravityPreparedResidual(invalidContributions, 0.0, 0.0, 0.0).admissible);
    }
}