#include <array>
#include <cmath>
#include <concepts>
#include <cstdint>
#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;

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

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override {
            if (!m_prepared) {
                throw std::logic_error("Test operator is unprepared.");
            }
            m_matrix.Mult(direction, action);
        }
        void Prepare(mfem::DenseMatrix matrix) {
            m_matrix   = std::move(matrix);
            m_prepared = true;
            ++m_generation;
        }
        void Invalidate() {
            m_prepared = false;
        }
        [[nodiscard]] bool IsPrepared() const {
            return m_prepared;
        }
        [[nodiscard]] const PreparedDenseOperator &GetProblem() const {
            return *this;
        }
        [[nodiscard]] std::uint64_t GetPreparationGeneration() const {
            return m_generation;
        }

    private:
        mfem::DenseMatrix m_matrix;
        std::uint64_t m_generation{1};
        bool m_prepared{true};
    };

    class PreparedPermutationInverse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;

        explicit PreparedPermutationInverse(
            const PreparedDenseOperator &operation,
            std::vector<int> residualForValue = {}
        )
            : mfem::Solver(operation.Width()),
              m_operation(operation),
              m_residualForValue(std::move(residualForValue)) {
            if (m_residualForValue.empty()) {
                for (int index = 0; index < Width(); ++index) {
                    m_residualForValue.push_back(index);
                }
            }
            Refresh();
        }
        void SetOperator(const mfem::Operator &) override {
        }
        void Mult(
            const mfem::Vector &rhs,
            mfem::Vector &action
        ) const override {
            if (!IsCurrent()) {
                throw std::logic_error("Test inverse is stale.");
            }
            for (int index = 0; index < Width(); ++index) {
                action(index) = rhs(m_residualForValue[static_cast<std::size_t>(index)]);
            }
        }
        void Refresh() {
            m_generation = m_operation.GetPreparationGeneration();
        }
        [[nodiscard]] bool IsCurrent() const {
            return m_operation.IsPrepared() && m_generation == m_operation.GetPreparationGeneration();
        }

    private:
        const PreparedDenseOperator &m_operation;
        std::vector<int> m_residualForValue;
        std::uint64_t m_generation{0};
    };

    [[nodiscard]] mfem::DenseMatrix Diagonal(const std::initializer_list<double> entries) {
        mfem::DenseMatrix result(static_cast<int>(entries.size()));
        result    = 0.0;
        int index = 0;
        for (const double entry : entries) {
            result(index, index) = entry;
            ++index;
        }
        return result;
    }

    [[nodiscard]] mfem::DenseMatrix CurvedJacobian(
        const double q,
        const double z
    ) {
        mfem::DenseMatrix result(2);
        result(0, 0) = 2.0 * q * z - 3.0 * q * q;
        result(0, 1) = 1.0 + q * q;
        result(1, 0) = 1.0;
        result(1, 1) = 2.0;
        return result;
    }

    [[nodiscard]] mfem::Vector CurvedResidual(
        const double q,
        const double z
    ) {
        mfem::Vector result(2);
        result(0) = (1.0 + q * q) * z - q * q * q;
        result(1) = q + 2.0 * z;
        return result;
    }

    using PlainForm   = utils::blocks::surface_deformed_stellar_equilibrium_form;
    using CentralForm = utils::blocks::central_density_bordered_stellar_equilibrium_form;
    static_assert(solver::GravityPartitionForm<PlainForm>);
    static_assert(solver::GravityPartitionForm<CentralForm>);
    static_assert(!solver::GravityPartitionForm<int>);
    static_assert(solver::GravitySubsystemOperation<PreparedDenseOperator>);
    static_assert(solver::GravitySubsystemPreconditioner<PreparedPermutationInverse>);
    static_assert(!solver::GravitySubsystemOperation<mfem::Operator>);
    static_assert(!solver::GravitySubsystemPreconditioner<mfem::Solver>);
} // namespace

TEST_CASE(
    "Gravity partition preserves noncontiguous coordinate and equation order",
    "[solver][gravity_subsystem]"
) {
    using Catch::Approx;
    using namespace mean_field;
    solver::GravityBlockPartition partition(5, 5, {3, 1}, {0, 4});
    mfem::Vector full(5), compact(2), retained(3), assembled(5);
    for (int index = 0; index < 5; ++index) {
        full(index) = 10.0 + index;
    }
    partition.GatherGravityValues(full, compact);
    CHECK(compact(0) == 13.0);
    CHECK(compact(1) == 11.0);
    partition.InjectGravityResiduals(compact, assembled);
    CHECK(assembled(0) == 13.0);
    CHECK(assembled(4) == 11.0);
    CHECK(assembled(1) == 0.0);
    partition.GatherRetainedValues(full, retained);
    CHECK(retained(0) == 10.0);
    CHECK(retained(1) == 12.0);
    CHECK(retained(2) == 14.0);
    partition.InjectRetainedValues(retained, assembled);
    partition.AddGravityValues(compact, assembled);
    assembled -= full;
    CHECK(assembled.Norml2() == 0.0);
    CHECK_THROWS_AS(solver::GravityBlockPartition(3, 3, {0, 0}, {0, 1}), std::invalid_argument);
    CHECK_THROWS_AS(solver::GravityBlockPartition(3, 3, {3}, {0}), std::invalid_argument);
    CHECK_THROWS_AS(solver::GravityBlockPartition(3, 3, {0}, {0, 1}), std::invalid_argument);
    mfem::Vector aliased(full.GetData() + 1, 2);
    CHECK_THROWS_AS(partition.InjectGravityValues(aliased, full), std::invalid_argument);
}

TEST_CASE(
    "Gravity partition compiles from forms with and without a central border",
    "[solver][gravity_subsystem]"
) {
    using namespace mean_field;
    const auto check = []<typename Form>() {
        std::array<int, Form::value_block_count> valueSizes{};
        std::array<int, Form::residual_block_count> residualSizes{};
        valueSizes.fill(1);
        residualSizes.fill(1);
        utils::blocks::form_layout<Form> layout(valueSizes, residualSizes);
        const auto partition = solver::MakeGravityBlockPartition<Form>(layout);
        CHECK(partition.GravityValueSize() == 2);
        CHECK(partition.GravityResidualSize() == 2);
        CHECK(partition.RetainedValueSize() == Form::value_block_count - 2);
        CHECK(partition.RetainedResidualSize() == Form::residual_block_count - 2);
    };
    check.template operator()<PlainForm>();
    check.template operator()<CentralForm>();
}

TEST_CASE(
    "Gravity subsystem solves the actual restricted operator and rejects "
    "stale preparation",
    "[solver][gravity_subsystem]"
) {
    using Catch::Approx;
    using namespace mean_field;
    PreparedDenseOperator operation(Diagonal({2.0, 3.0, 7.0}));
    PreparedPermutationInverse inverse(operation);
    solver::PreparedGravitySubsystem subsystem(
        operation, inverse, solver::GravityBlockPartition(3, 3, {1, 0}, {1, 0}), MPI_COMM_WORLD
    );
    static_assert(solver::GravitySubsystemRuntime<decltype(subsystem)>);
    mfem::Vector rhs(2), correction(2), action(2);
    rhs(0)     = 9.0;
    rhs(1)     = -8.0;
    correction = 0.0;
    const solver::LinearSolveControl control{.relativeTolerance = 1.0e-12, .maximumIterations = 8};
    const auto report = subsystem.Solve(rhs, correction, control);
    REQUIRE(report.Converged());
    CHECK(correction(0) == Approx(3.0).margin(1.0e-12));
    CHECK(correction(1) == Approx(-4.0).margin(1.0e-12));
    subsystem.GetOperator().Mult(correction, action);
    action -= rhs;
    double localSquared  = action * action;
    double globalSquared = 0.0;
    MPI_Allreduce(&localSquared, &globalSquared, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    CHECK(report.trueResidualNorm == Approx(std::sqrt(globalSquared)).margin(1.0e-14));

    operation.Prepare(Diagonal({4.0, 6.0, 7.0}));
    CHECK_FALSE(subsystem.IsCurrent());
    CHECK_THROWS_AS(subsystem.Solve(rhs, correction, control), std::logic_error);
    CHECK_THROWS_AS(subsystem.Refresh(), std::logic_error);
    inverse.Refresh();
    CHECK_FALSE(subsystem.IsCurrent());
    subsystem.Refresh();
    REQUIRE(subsystem.Solve(rhs, correction, control).Converged());
    CHECK(correction(0) == Approx(1.5).margin(1.0e-12));
    CHECK(correction(1) == Approx(-2.0).margin(1.0e-12));
    operation.Invalidate();
    CHECK_FALSE(subsystem.IsCurrent());
    CHECK_THROWS_AS(subsystem.Refresh(), std::logic_error);
}

TEST_CASE(
    "Gravity inner solve reports an unresolved true defect instead of "
    "certifying a correction",
    "[solver][gravity_subsystem]"
) {
    using namespace mean_field;
    PreparedDenseOperator operation(Diagonal({2.0, 3.0, 7.0}));
    PreparedPermutationInverse inverse(operation);
    solver::PreparedGravitySubsystem subsystem(
        operation, inverse, solver::GravityBlockPartition(3, 3, {0, 1}, {0, 1}), MPI_COMM_WORLD
    );
    mfem::Vector rhs(2), correction(2);
    rhs        = 1.0;
    correction = 0.0;
    const solver::LinearSolveControl control{.relativeTolerance = 1.0e-12, .maximumIterations = 1};
    const auto report = subsystem.Solve(rhs, correction, control);
    CHECK_FALSE(report.Converged());
    CHECK(report.status == solver::LinearSolveStatus::maximum_iterations);
    CHECK(report.trueResidualNorm > control.ConvergenceThreshold(report.rightHandSideNorm));
    CHECK(report.relativeTrueResidualNorm > 0.1);
    CHECK(subsystem.GetSolveCount() == 1);
    CHECK(subsystem.GetIterationCount() == static_cast<std::uint64_t>(report.iterations));
    CHECK(subsystem.GetFailureCount() == 1);
    CHECK(subsystem.GetSolveSeconds() == report.solveSeconds);
    REQUIRE(subsystem.GetLastSolveReport().has_value());
    CHECK(subsystem.GetLastSolveReport()->trueResidualNorm == report.trueResidualNorm);
    rhs(0) = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS_AS(subsystem.Solve(rhs, correction, control), std::invalid_argument);
    CHECK(subsystem.GetSolveCount() == 2);
    CHECK(subsystem.GetFailureCount() == 2);
}

TEST_CASE(
    "Distributed gravity solves admit empty local partitions but reject "
    "a globally empty system",
    "[solver][gravity_subsystem]"
) {
    using Catch::Approx;
    using namespace mean_field;
    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    const int localSize = rank == 0 ? 1 : 0;
    mfem::DenseMatrix matrix(localSize);
    if (localSize != 0) {
        matrix(0, 0) = 4.0;
    }
    PreparedDenseOperator operation(matrix);
    PreparedPermutationInverse inverse(operation);
    std::vector<int> indices = localSize != 0 ? std::vector<int>{0} : std::vector<int>{};
    solver::PreparedGravitySubsystem subsystem(
        operation, inverse, solver::GravityBlockPartition(localSize, localSize, indices, indices), MPI_COMM_WORLD
    );
    mfem::Vector rhs(localSize), correction(localSize);
    rhs               = 12.0;
    correction        = 0.0;
    const auto report = subsystem.Solve(rhs, correction, {});
    REQUIRE(report.Converged());
    CHECK(report.rightHandSideNorm == 12.0);
    if (localSize != 0) {
        CHECK(correction(0) == Approx(3.0));
    }

    auto backend = solver::linear::prepareLinearBackend(
        solver::linear::FGMRES{{.restartLength = 13, .printLevel = -1}}, operation, inverse, MPI_COMM_WORLD
    );
    CHECK(backend.GetConfiguration().GetOptions().restartLength == 13);
    CHECK(backend.GetConfiguration().GetOptions().printLevel == -1);

    PreparedDenseOperator emptyOperation(mfem::DenseMatrix(0));
    PreparedPermutationInverse emptyInverse(emptyOperation);
    CHECK_THROWS_AS(
        (solver::PreparedGravitySubsystem(
            emptyOperation, emptyInverse, solver::GravityBlockPartition(0, 0, {}, {}), MPI_COMM_WORLD
        )),
        std::invalid_argument
    );
}

TEST_CASE(
    "Affine gravity receives no curvature correction even away from "
    "gravity equilibrium",
    "[solver][gravity_subsystem]"
) {
    using namespace mean_field;
    mfem::DenseMatrix jacobian(2);
    jacobian(0, 0) = 2.0;
    jacobian(0, 1) = 3.0;
    jacobian(1, 0) = 1.0;
    jacobian(1, 1) = 1.0;
    PreparedDenseOperator operation(jacobian);
    PreparedPermutationInverse inverse(operation, {1, 0});
    solver::PreparedGravitySubsystem subsystem(
        operation, inverse, solver::GravityBlockPartition(2, 2, {1}, {0}), MPI_COMM_WORLD
    );
    solver::PreparedGravityCurvatureCorrection correction(subsystem);
    mfem::Vector baseResidual(2), direction(2), trialResidual(2), gravityStep(1), prediction(1);
    baseResidual(0) = 0.75;
    baseResidual(1) = 0.75;
    direction(0)    = 0.125;
    direction(1)    = 0.5;
    correction.CaptureBase(baseResidual, direction);
    const auto baseGeneration = correction.GetBaseGeneration();
    for (const double alpha : {1.0, 0.5, 0.25, 0.125, 0.0}) {
        operation.Prepare(jacobian);
        inverse.Refresh();
        subsystem.Refresh();
        trialResidual(0) = 0.75 + alpha * 1.75;
        trialResidual(1) = 0.75 + alpha * 0.625;
        correction.BuildLinearPrediction(alpha, prediction);
        CHECK(prediction(0) == trialResidual(0));
        const auto report = correction.SolveCorrection(alpha, trialResidual, gravityStep, {});
        REQUIRE(report.Converged());
        CHECK(report.rightHandSideNorm == 0.0);
        CHECK(gravityStep.Norml2() == 0.0);
        CHECK(correction.GetBaseGeneration() == baseGeneration);
    }
    correction.Invalidate();
    CHECK_FALSE(correction.HasBase());
    CHECK_THROWS_AS(correction.BuildLinearPrediction(1.0, prediction), std::logic_error);
    CHECK_THROWS_AS(correction.SolveCorrection(1.0, trialResidual, gravityStep, {}), std::logic_error);
}

TEST_CASE(
    "Gravity curvature correction preserves the frozen tangent and is "
    "second order",
    "[solver][gravity_subsystem]"
) {
    using Catch::Approx;
    using namespace mean_field;
    constexpr double baseQ = 0.4, baseZ = 1.2, deltaQ = 0.3, deltaZ = -0.2;
    PreparedDenseOperator operation(CurvedJacobian(baseQ, baseZ));
    PreparedPermutationInverse inverse(operation, {1, 0});
    solver::PreparedGravitySubsystem subsystem(
        operation, inverse, solver::GravityBlockPartition(2, 2, {1}, {0}), MPI_COMM_WORLD
    );
    solver::PreparedGravityCurvatureCorrection correction(subsystem);
    const auto baseResidual = CurvedResidual(baseQ, baseZ);
    mfem::Vector direction(2), action(2), gravityStep(1), candidate(2);
    direction(0) = deltaQ;
    direction(1) = deltaZ;
    operation.Mult(direction, action);
    correction.CaptureBase(baseResidual, direction);
    double previousStep = 0.0;
    for (const double alpha : {0.02, 0.01, 0.005, 0.0025}) {
        candidate(0) = baseQ + alpha * deltaQ;
        candidate(1) = baseZ + alpha * deltaZ;
        operation.Prepare(CurvedJacobian(candidate(0), candidate(1)));
        inverse.Refresh();
        subsystem.Refresh();
        const auto trialResidual = CurvedResidual(candidate(0), candidate(1));
        const auto report        = correction.SolveCorrection(
            alpha, trialResidual, gravityStep, {.relativeTolerance = 1.0e-12, .maximumIterations = 4}
        );
        REQUIRE(report.Converged());
        subsystem.GetPartition().AddGravityValues(gravityStep, candidate);
        const auto correctedResidual = CurvedResidual(candidate(0), candidate(1));
        CHECK(correctedResidual(0) == Approx(baseResidual(0) + alpha * action(0)).margin(2.0e-15));
        CHECK(candidate(0) == baseQ + alpha * deltaQ);
        if (previousStep != 0.0) {
            CHECK(std::abs(gravityStep(0) / previousStep - 0.25) < 0.004);
        }
        previousStep = gravityStep(0);
    }
    CHECK(baseResidual(0) == CurvedResidual(baseQ, baseZ)(0));
    CHECK_THROWS_AS(correction.SolveCorrection(-1.0, baseResidual, gravityStep, {}), std::logic_error);
}