#include <algorithm>
#include <cmath>
#include <memory>
#include <stdexcept>

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

import mean_field;
import test_helpers;

namespace {
    namespace blocks          = mean_field::utils::blocks;
    namespace preconditioning = mean_field::preconditioning;
    namespace backend         = mean_field::preconditioning::backend;

    using DiagonalMass        = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::mass_like,
        preconditioning::OperatorValueStructure::scalar,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_definite,
        preconditioning::OperatorRepresentation::diagonal,
        preconditioning::OperatorDistribution::local,
        preconditioning::OperatorFESpace::h1>;
    using MatrixFreeHdivMass = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::mass_like,
        preconditioning::OperatorValueStructure::vector,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_definite,
        preconditioning::OperatorRepresentation::matrix_free,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h_div>;
    using DenseBorder = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::dense_border,
        preconditioning::OperatorValueStructure::block,
        preconditioning::OperatorSymmetry::nonsymmetric,
        preconditioning::OperatorDefiniteness::indefinite,
        preconditioning::OperatorRepresentation::assembled_dense,
        preconditioning::OperatorDistribution::local>;
    using ScalarH1Elliptic = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::elliptic_like,
        preconditioning::OperatorValueStructure::scalar,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_definite,
        preconditioning::OperatorRepresentation::assembled_sparse,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h1>;
    using ConstantNullspaceH1Elliptic = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::elliptic_like,
        preconditioning::OperatorValueStructure::scalar,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_semidefinite,
        preconditioning::OperatorRepresentation::assembled_sparse,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h1,
        preconditioning::OperatorNullspace::constant_mode>;
    using HdivElliptic = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::elliptic_like,
        preconditioning::OperatorValueStructure::vector,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_definite,
        preconditioning::OperatorRepresentation::assembled_sparse,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h_div>;
    using NonsymmetricH1Elliptic = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::elliptic_like,
        preconditioning::OperatorValueStructure::scalar,
        preconditioning::OperatorSymmetry::nonsymmetric,
        preconditioning::OperatorDefiniteness::indefinite,
        preconditioning::OperatorRepresentation::assembled_sparse,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h1>;
    using SuppliedNullspaceH1Elliptic = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::elliptic_like,
        preconditioning::OperatorValueStructure::scalar,
        preconditioning::OperatorSymmetry::symmetric,
        preconditioning::OperatorDefiniteness::positive_semidefinite,
        preconditioning::OperatorRepresentation::assembled_sparse,
        preconditioning::OperatorDistribution::distributed_true_dof,
        preconditioning::OperatorFESpace::h1,
        preconditioning::OperatorNullspace::supplied_basis>;

    using FixedAMG          = backend::HypreBoomerAMG<backend::FixedCycles>;
    using AdaptiveAMG       = backend::HypreBoomerAMG<backend::SolveToTolerance>;

    using DiagonalComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::density::mass::value>,
        blocks::type_list<blocks::density::mass::residual>,
        blocks::type_list<>,
        DiagonalMass,
        backend::Diagonal>;
    using UnderdeclaredDiagonalComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::density::mass::value>,
        blocks::type_list<blocks::density::mass::residual>,
        blocks::type_list<>,
        DiagonalMass,
        backend::Diagonal,
        preconditioning::NoPreparationDependencies>;

    class TinyParallelH1Operator final {
    public:
        TinyParallelH1Operator()
            : m_serialMesh(mfem::Mesh::MakeCartesian1D(4)),
              m_parallelMesh(
                  MPI_COMM_WORLD,
                  m_serialMesh
              ),
              m_collection(
                  1,
                  1
              ),
              m_space(
                  &m_parallelMesh,
                  &m_collection
              ),
              m_form(&m_space) {
            m_form.AddDomainIntegrator(new mfem::DiffusionIntegrator());
            m_form.AddDomainIntegrator(new mfem::MassIntegrator());
            m_form.Assemble();
            m_form.Finalize();
            m_matrix.reset(m_form.ParallelAssemble());
        }

        [[nodiscard]] const mfem::HypreParMatrix &matrix() const {
            return *m_matrix;
        }

    private:
        mfem::Mesh m_serialMesh;
        mfem::ParMesh m_parallelMesh;
        mfem::H1_FECollection m_collection;
        mfem::ParFiniteElementSpace m_space;
        mfem::ParBilinearForm m_form;
        std::unique_ptr<mfem::HypreParMatrix> m_matrix;
    };

    class KnownMatrixFreeSPDOperator final : public mfem::Operator {
    public:
        KnownMatrixFreeSPDOperator()
            : mfem::Operator(3),
              m_matrix(3) {
            m_matrix       = 0.0;
            m_matrix(0, 0) = 4.0;
            m_matrix(0, 1) = 1.0;
            m_matrix(1, 0) = 1.0;
            m_matrix(1, 1) = 3.0;
            m_matrix(1, 2) = 0.5;
            m_matrix(2, 1) = 0.5;
            m_matrix(2, 2) = 2.0;
        }

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

        void AssembleDiagonal(mfem::Vector &diagonal) const override {
            diagonal.SetSize(Height());
            for (int index = 0; index < Height(); ++index) {
                diagonal(index) = m_matrix(index, index);
            }
        }

    private:
        mfem::DenseMatrix m_matrix;
    };

    [[nodiscard]] double relativeResidual(
        const mfem::HypreParMatrix &matrix,
        const mfem::Vector &rightHandSide,
        const mfem::Vector &action
    ) {
        mfem::Vector residual(rightHandSide.Size());
        matrix.Mult(action, residual);
        residual -= rightHandSide;
        return gravity_prepared_test_utils::global_norm(residual, matrix.GetComm()) /
               gravity_prepared_test_utils::global_norm(rightHandSide, matrix.GetComm());
    }
} // namespace

TEST_CASE(
    "Preconditioning Backends Expose Complete Compile-Time Capabilities",
    tags::preconditioning_backend_unit
) {
    STATIC_CHECK(backend::Compatible<backend::Diagonal, DiagonalMass>);
    STATIC_CHECK(backend::Compatible<backend::Diagonal, MatrixFreeHdivMass>);
    STATIC_CHECK(backend::Compatible<backend::MatrixFreeChebyshev, MatrixFreeHdivMass>);
    STATIC_CHECK(backend::Compatible<backend::DenseDirect, DenseBorder>);
    STATIC_CHECK(backend::Compatible<FixedAMG, ScalarH1Elliptic>);
    STATIC_CHECK(backend::Compatible<FixedAMG, ConstantNullspaceH1Elliptic>);
    STATIC_CHECK_FALSE(backend::Compatible<FixedAMG, HdivElliptic>);
    STATIC_CHECK_FALSE(backend::Compatible<FixedAMG, NonsymmetricH1Elliptic>);
    STATIC_CHECK_FALSE(backend::Compatible<FixedAMG, SuppliedNullspaceH1Elliptic>);

    STATIC_CHECK(backend::ArnoldiAdmissible<FixedAMG>);
    STATIC_CHECK(backend::ArnoldiAdmissible<backend::MatrixFreeChebyshev>);
    STATIC_CHECK_FALSE(backend::ArnoldiAdmissible<AdaptiveAMG>);
    STATIC_CHECK(backend::requiresAssembledSparseSurrogate<FixedAMG>);
    STATIC_CHECK_FALSE(backend::requiresAssembledSparseSurrogate<backend::Diagonal>);
    STATIC_CHECK(backend::Traits<backend::Diagonal>::supportsSerialExecution);
    STATIC_CHECK(backend::Traits<backend::Diagonal>::supportsDistributedExecution);
    STATIC_CHECK(backend::Traits<backend::DenseDirect>::supportsSerialExecution);
    STATIC_CHECK_FALSE(backend::Traits<backend::DenseDirect>::supportsDistributedExecution);
    STATIC_CHECK_FALSE(backend::Traits<FixedAMG>::supportsSerialExecution);
    STATIC_CHECK(backend::Traits<FixedAMG>::supportsDistributedExecution);

    STATIC_CHECK(preconditioning::PreconditionerComponent<DiagonalComponent>);
    STATIC_CHECK(
        DiagonalComponent::PreparationDependencies::contains(preconditioning::PreparationDependency::linearization)
    );
    STATIC_CHECK_FALSE(preconditioning::PreconditionerComponent<UnderdeclaredDiagonalComponent>);
}

TEST_CASE(
    "Diagonal Backend Exactly Inverts A Known Diagonal Operator",
    tags::preconditioning_backend_unit
) {
    mfem::Vector diagonal(3);
    diagonal(0)   = 2.0;
    diagonal(1)   = -4.0;
    diagonal(2)   = 0.5;

    auto prepared = backend::prepare(backend::Diagonal{}, diagonal);
    mfem::Vector rightHandSide(3);
    rightHandSide(0) = 4.0;
    rightHandSide(1) = 8.0;
    rightHandSide(2) = -1.0;
    mfem::Vector action(3);
    prepared.Mult(rightHandSide, action);

    CHECK(action(0) == Catch::Approx(2.0));
    CHECK(action(1) == Catch::Approx(-2.0));
    CHECK(action(2) == Catch::Approx(-2.0));
    CHECK(prepared.GetStatistics().setups == 1);
    CHECK(prepared.GetStatistics().applications == 1);
    CHECK(prepared.GetStatistics().innerIterations == 0);

    diagonal(1) = 0.0;
    CHECK_THROWS_AS(prepared.Refresh(diagonal), std::invalid_argument);
}

TEST_CASE(
    "Dense Direct Backend Exactly Solves And Refreshes A Known Border",
    tags::preconditioning_backend_unit
) {
    mfem::DenseMatrix matrix(2);
    matrix(0, 0)  = 4.0;
    matrix(0, 1)  = 1.0;
    matrix(1, 0)  = 2.0;
    matrix(1, 1)  = 3.0;

    auto prepared = backend::prepare(backend::DenseDirect{}, matrix);
    mfem::Vector rightHandSide(2);
    rightHandSide(0) = 7.0;
    rightHandSide(1) = 1.0;
    mfem::Vector action(2);
    prepared.Mult(rightHandSide, action);

    CHECK(action(0) == Catch::Approx(2.0).margin(1.0e-14));
    CHECK(action(1) == Catch::Approx(-1.0).margin(1.0e-14));

    matrix       = 0.0;
    matrix(0, 0) = 2.0;
    matrix(1, 1) = 4.0;
    prepared.Refresh(matrix);
    prepared.Mult(rightHandSide, action);

    CHECK(action(0) == Catch::Approx(3.5).margin(1.0e-14));
    CHECK(action(1) == Catch::Approx(0.25).margin(1.0e-14));
    CHECK(prepared.GetStatistics().setups == 2);
    CHECK(prepared.GetStatistics().applications == 2);
}

TEST_CASE(
    "Matrix-Free Chebyshev Backend Is A Fixed Linear Positive Approximate Inverse",
    tags::preconditioning_backend_unit
) {
    const KnownMatrixFreeSPDOperator operation;
    auto prepared =
        backend::prepare(backend::MatrixFreeChebyshev{.order = 3, .powerIterations = 20}, operation, MPI_COMM_WORLD);

    mfem::Vector first(3);
    first(0) = 1.0;
    first(1) = -2.0;
    first(2) = 0.25;
    mfem::Vector second(3);
    second(0) = -0.5;
    second(1) = 0.75;
    second(2) = 3.0;
    mfem::Vector combination(first);
    combination *= 1.7;
    combination.Add(-0.4, second);

    mfem::Vector firstAction(3);
    mfem::Vector secondAction(3);
    mfem::Vector combinationAction(3);
    prepared.Mult(first, firstAction);
    prepared.Mult(second, secondAction);
    prepared.Mult(combination, combinationAction);

    mfem::Vector expected(firstAction);
    expected *= 1.7;
    expected.Add(-0.4, secondAction);
    mfem::Vector linearityError(combinationAction);
    linearityError -= expected;
    CHECK(linearityError.Norml2() <= 2.0e-12 * std::max(1.0, expected.Norml2()));
    CHECK((first * firstAction) > 0.0);
    CHECK(prepared.GetStatistics().setups == 1);
    CHECK(prepared.GetStatistics().applications == 3);
    CHECK(prepared.GetStatistics().innerIterations == 9);
    CHECK(prepared.GetStatistics().lastInnerIterations == 3);
    CHECK_THROWS_AS(
        backend::prepare(backend::MatrixFreeChebyshev{.order = 0}, operation, MPI_COMM_WORLD), std::invalid_argument
    );
    CHECK_THROWS_AS(
        backend::prepare(backend::MatrixFreeChebyshev{.order = 6}, operation, MPI_COMM_WORLD), std::invalid_argument
    );
    CHECK_THROWS_AS(
        backend::prepare(backend::MatrixFreeChebyshev{.powerTolerance = 1.0}, operation, MPI_COMM_WORLD),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        backend::prepare(backend::MatrixFreeChebyshev{.powerSeed = 0}, operation, MPI_COMM_WORLD), std::invalid_argument
    );
}

TEST_CASE(
    "BoomerAMG Fixed Cycles Are Linear While Adaptive Application Meets Its Tolerance",
    tags::preconditioning_backend_unit
) {
    TinyParallelH1Operator problem;
    const mfem::HypreParMatrix &matrix = problem.matrix();

    mfem::Vector exact(matrix.Width());
    exact = 1.0;
    mfem::Vector rightHandSide(matrix.Height());
    matrix.Mult(exact, rightHandSide);

    auto fixed = backend::prepare(FixedAMG{backend::FixedCycles{.cycles = 2}}, matrix);
    mfem::Vector fixedAction(matrix.Width());
    fixedAction = 0.0;
    fixed.Mult(rightHandSide, fixedAction);

    CHECK(relativeResidual(matrix, rightHandSide, fixedAction) < 1.0);

    mfem::Vector secondExact(matrix.Width());
    for (int index = 0; index < secondExact.Size(); ++index) {
        secondExact(index) = 0.25 + static_cast<double>(index);
    }
    mfem::Vector secondRightHandSide(matrix.Height());
    matrix.Mult(secondExact, secondRightHandSide);
    mfem::Vector secondAction(matrix.Width());
    secondAction = 0.0;
    fixed.Mult(secondRightHandSide, secondAction);

    mfem::Vector combinedRightHandSide(rightHandSide);
    combinedRightHandSide *= 0.7;
    combinedRightHandSide.Add(-0.2, secondRightHandSide);
    mfem::Vector combinedAction(matrix.Width());
    combinedAction = 0.0;
    fixed.Mult(combinedRightHandSide, combinedAction);

    mfem::Vector expectedCombinedAction(fixedAction);
    expectedCombinedAction *= 0.7;
    expectedCombinedAction.Add(-0.2, secondAction);
    combinedAction -= expectedCombinedAction;
    CHECK(
        gravity_prepared_test_utils::global_norm(combinedAction, matrix.GetComm()) <
        1.0e-11 * gravity_prepared_test_utils::global_norm(expectedCombinedAction, matrix.GetComm())
    );

    CHECK(fixed.GetStatistics().setups == 1);
    CHECK(fixed.GetStatistics().applications == 3);
    CHECK(fixed.GetStatistics().lastInnerIterations >= 1);
    CHECK(fixed.GetStatistics().lastInnerIterations <= 2);

    auto adaptive = backend::prepare(
        AdaptiveAMG{backend::SolveToTolerance{.relativeTolerance = 1.0e-10, .maximumCycles = 50}}, matrix
    );
    mfem::Vector adaptiveAction(matrix.Width());
    adaptiveAction = 0.0;
    adaptive.Mult(rightHandSide, adaptiveAction);

    CHECK(relativeResidual(matrix, rightHandSide, adaptiveAction) < 1.0e-8);
    CHECK(adaptive.GetStatistics().setups == 1);
    CHECK(adaptive.GetStatistics().applications == 1);
    CHECK(adaptive.GetStatistics().lastInnerIterations >= 1);
    CHECK(adaptive.GetStatistics().lastInnerIterations <= 50);
}