#include <array>
#include <cmath>
#include <cstdint>
#include <stdexcept>
#include <vector>

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

import mean_field;

namespace gravity_compact_test {
    using namespace mean_field;

    template <typename Form> class Operation final : public mfem::Operator {
        template <std::size_t N> static auto Sizes() {
            std::array<int, N> sizes;
            sizes.fill(1);
            return sizes;
        }

    public:
        Operation()
            : mfem::Operator(
                  Form::residual_block_count,
                  Form::value_block_count
              ),
              layout(
                  Sizes<Form::value_block_count>(),
                  Sizes<Form::residual_block_count>()
              ),
              partition(solver::MakeGravityBlockPartition(layout)),
              state(Width()),
              rows(Height()),
              physicalDirection(Width()),
              physicalAction(Height()),
              gravityDirection(2),
              gravityAction(2) {
            for (int i = 0; i < Width(); ++i)
                state(i) = 0.3 + 0.7 * (i + 1);
            for (int i = 0; i < Height(); ++i)
                rows(i) = 1.2 + 0.4 * (i + 1);
        }
        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override {
            ++rootApplications;
            for (int i = 0; i < Width(); ++i)
                physicalDirection(i) = direction(i) / state(i);
            partition.GatherGravityValues(physicalDirection, gravityDirection);
            ApplyGravity(gravityDirection, gravityAction);
            partition.InjectGravityResiduals(gravityAction, physicalAction);
            for (int i = 0; i < Height(); ++i)
                action(i) = rows(i) * physicalAction(i);
        }
        void ApplyGravity(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const {
            action(0) = 3.0 * direction(0) + 2.0 * direction(1);
            action(1) = 2.0 * direction(0);
        }
        bool IsPrepared() const {
            return prepared;
        }
        const Operation &GetProblem() const {
            return *this;
        }
        std::uint64_t GetPreparationGeneration() const {
            return generation;
        }
        void Prepare() {
            prepared = true;
            ++generation;
        }
        utils::blocks::form_layout<Form> layout;
        solver::GravityBlockPartition partition;
        mfem::Vector state, rows;
        bool prepared{true};
        std::uint64_t generation{1};
        mutable int rootApplications{0}, compactApplications{0};

    private:
        mutable mfem::Vector physicalDirection, physicalAction, gravityDirection, gravityAction;
    };

    template <typename Form, bool Registered = true> class Inverse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        explicit Inverse(const Operation<Form> &operation)
            : mfem::Solver(operation.Width()),
              operation(operation),
              physicalRhs(Width()),
              physicalCorrection(Height()),
              compactRhs(2),
              compactCorrection(2) {
            Refresh();
        }
        void SetOperator(const mfem::Operator &) override {
        }
        void Mult(
            const mfem::Vector &rhs,
            mfem::Vector &action
        ) const override {
            ++rootApplications;
            for (int i = 0; i < Width(); ++i)
                physicalRhs(i) = rhs(i) / operation.rows(i);
            operation.partition.GatherGravityResiduals(physicalRhs, compactRhs);
            ApplyGravity(compactRhs, compactCorrection);
            operation.partition.InjectGravityValues(compactCorrection, physicalCorrection);
            for (int i = 0; i < Height(); ++i)
                action(i) = operation.state(i) * physicalCorrection(i);
        }
        void ApplyGravity(
            const mfem::Vector &rhs,
            mfem::Vector &action
        ) const {
            action = rhs;
        }
        bool IsCurrent() const {
            return operation.IsPrepared() && generation == operation.generation;
        }
        void Refresh() {
            generation     = operation.generation;
            compactCurrent = true;
        }
        const Operation<Form> &operation;
        bool compactCurrent{true};
        mutable int rootApplications{0}, compactApplications{0};

    private:
        std::uint64_t generation{0};
        mutable mfem::Vector physicalRhs, physicalCorrection, compactRhs, compactCorrection;
    };
} // namespace gravity_compact_test

namespace mean_field::solver {
    template <typename Form> struct GravityCompactOperatorAccess<gravity_compact_test::Operation<Form>> {
        using Operation                 = gravity_compact_test::Operation<Form>;
        static constexpr bool available = true;
        static GravityBlockPartition CanonicalPartition(const Operation &operation) {
            return operation.partition;
        }
        static const mfem::Vector &StateFactors(const Operation &operation) {
            return operation.state;
        }
        static const mfem::Vector &ResidualFactors(const Operation &operation) {
            return operation.rows;
        }
        static void ApplyPhysical(
            const Operation &operation,
            const mfem::Vector &direction,
            mfem::Vector &action
        ) {
            ++operation.compactApplications;
            operation.ApplyGravity(direction, action);
        }
    };
    template <typename Form>
    struct GravityCompactInverseAccess<gravity_compact_test::Operation<Form>, gravity_compact_test::Inverse<Form>> {
        using Operation                           = gravity_compact_test::Operation<Form>;
        using Inverse                             = gravity_compact_test::Inverse<Form>;
        static constexpr bool available           = true;
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        static bool Matches(
            const Operation &operation,
            const Inverse &inverse
        ) {
            return &operation == &inverse.operation;
        }
        static bool IsCurrent(const Inverse &inverse) {
            return inverse.compactCurrent;
        }
        static void ApplyPhysical(
            const Inverse &inverse,
            const mfem::Vector &rhs,
            mfem::Vector &action
        ) {
            ++inverse.compactApplications;
            inverse.ApplyGravity(rhs, action);
        }
    };
} // namespace mean_field::solver

namespace {
    template <typename Form> void CheckCanonicalApplications() {
        using namespace mean_field;
        gravity_compact_test::Operation<Form> operation;
        gravity_compact_test::Inverse<Form> inverse(operation);
        solver::PreparedGravitySubsystem compact(operation, inverse, operation.partition, MPI_COMM_WORLD);
        solver::PreparedGravitySubsystem reference(
            operation, inverse, operation.partition, MPI_COMM_WORLD, solver::linear::FGMRES{},
            solver::RootRestrictedGravityApplications{}
        );
        REQUIRE(compact.UsesCompactOperator());
        REQUIRE(compact.UsesCompactPreconditioner());
        REQUIRE_FALSE(reference.UsesCompactOperator());
        REQUIRE_FALSE(reference.UsesCompactPreconditioner());
        mfem::Vector direction(2), direct(2), restricted(2), correction(2);
        for (int trial = 0; trial < 4; ++trial) {
            direction(0) = std::sin(0.3 + trial);
            direction(1) = std::cos(0.7 + 2 * trial);
            compact.GetOperator().Mult(direction, direct);
            reference.GetOperator().Mult(direction, restricted);
            direct -= restricted;
            CHECK(direct.Norml2() < 1.0e-13);
            compact.GetPreconditioner().Mult(direction, direct);
            reference.GetPreconditioner().Mult(direction, restricted);
            direct -= restricted;
            CHECK(direct.Norml2() < 1.0e-13);
        }
        CHECK(operation.compactApplications == 4);
        CHECK(operation.rootApplications == 4);
        CHECK(inverse.compactApplications == 4);
        CHECK(inverse.rootApplications == 4);
        correction  = 0.0;
        auto report = compact.Solve(direction, correction, {.relativeTolerance = 1.0e-12, .maximumIterations = 4});
        REQUIRE(report.Converged());
        reference.GetOperator().Mult(correction, restricted);
        restricted -= direction;
        CHECK(restricted.Norml2() < 1.0e-11);
        CHECK(report.operatorApplications > 0);
        CHECK(report.inversePreconditionerApplications > 0);
        CHECK(report.operatorSeconds >= 0.0);
        CHECK(report.inversePreconditionerSeconds >= 0.0);

        // A normalization refresh invalidates the accepted linearization; the
        // compact factors must be reacquired when its new generation is ready.
        operation.state *= 2.0;
        operation.rows *= 0.25;
        operation.Prepare();
        CHECK_FALSE(compact.IsCurrent());
        CHECK_THROWS_AS(compact.GetOperator().Mult(direction, direct), std::logic_error);
        CHECK_THROWS_AS(compact.Refresh(), std::logic_error);
        inverse.Refresh();
        compact.Refresh();
        reference.Refresh();
        compact.GetOperator().Mult(direction, direct);
        reference.GetOperator().Mult(direction, restricted);
        direct -= restricted;
        CHECK(direct.Norml2() < 1.0e-13);
        inverse.compactCurrent = false;
        CHECK_FALSE(compact.IsCurrent());
        CHECK_THROWS_AS(compact.GetPreconditioner().Mult(direction, direct), std::logic_error);
    }
} // namespace

TEST_CASE(
    "Compact gravity preserves normalized block actions with optional borders",
    "[solver][gravity_subsystem][gravity_compact]"
) {
    CheckCanonicalApplications<mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form>();
    CheckCanonicalApplications<mean_field::utils::blocks::central_density_bordered_stellar_equilibrium_form>();
}

TEST_CASE(
    "Compact gravity selection preserves unknown inverses and collectively rejects noncanonical maps",
    "[solver][gravity_subsystem][gravity_compact]"
) {
    using namespace mean_field;
    using Form = utils::blocks::central_density_bordered_stellar_equilibrium_form;
    gravity_compact_test::Operation<Form> operation;
    gravity_compact_test::Inverse<Form, false> customInverse(operation);
    static_assert(!solver::GravityCompactInverseAccess<decltype(operation), decltype(customInverse)>::available);
    solver::PreparedGravitySubsystem unknown(operation, customInverse, operation.partition, MPI_COMM_WORLD);
    CHECK(unknown.UsesCompactOperator());
    CHECK_FALSE(unknown.UsesCompactPreconditioner());

    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    auto selected = operation.partition;
    if (rank == 0)
        selected = solver::GravityBlockPartition(operation.Width(), operation.Height(), {1, 0}, {0, 1});
    gravity_compact_test::Inverse<Form> inverse(operation);
    solver::PreparedGravitySubsystem fallback(operation, inverse, selected, MPI_COMM_WORLD);
    solver::PreparedGravitySubsystem reference(
        operation, inverse, selected, MPI_COMM_WORLD, solver::linear::FGMRES{},
        solver::RootRestrictedGravityApplications{}
    );
    CHECK_FALSE(fallback.UsesCompactOperator());
    CHECK_FALSE(fallback.UsesCompactPreconditioner());
    mfem::Vector direction(2), actual(2), expected(2);
    direction(0) = 0.43;
    direction(1) = -0.91;
    fallback.GetOperator().Mult(direction, actual);
    reference.GetOperator().Mult(direction, expected);
    actual -= expected;
    CHECK(actual.Norml2() == 0.0);
    fallback.GetPreconditioner().Mult(direction, actual);
    reference.GetPreconditioner().Mult(direction, expected);
    actual -= expected;
    CHECK(actual.Norml2() == 0.0);
}