module;

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <memory>
#include <numbers>
#include <stdexcept>
#include <utility>

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

module mean_field;

import :solver.stellar_equilibrium;

namespace solver_internal_architecture_test {
    struct LifetimeProbe final {
        const void *problemIdentity{nullptr};
        bool backendDestroyed{false};
        bool dependenciesAliveAtBackendDestruction{false};
    };

    struct InspectingBackend final : mean_field::solver::LinearBackendConfigurationTag {
        static constexpr mean_field::preconditioning::ApplicationContract supportedPreconditionerContract =
            mean_field::preconditioning::ApplicationContract::flexible;

        std::shared_ptr<LifetimeProbe> probe;

        explicit InspectingBackend(std::shared_ptr<LifetimeProbe> lifetimeProbe = nullptr)
            : probe(std::move(lifetimeProbe)) {
        }
    };

    struct ThrowingBackend final : mean_field::solver::LinearBackendConfigurationTag {
        static constexpr mean_field::preconditioning::ApplicationContract supportedPreconditionerContract =
            mean_field::preconditioning::ApplicationContract::flexible;
    };

    struct ZeroMetric final { };

    [[nodiscard]] mean_field::solver::nonlinear::MetricEvaluation getMetric(
        const ZeroMetric &,
        const mfem::Vector &,
        MPI_Comm
    ) {
        return {.residualNorm = 0.0, .merit = 0.0};
    }

    template <typename Operator, typename Preconditioner> class PreparedBackend final {
    public:
        PreparedBackend(
            const Operator &operation,
            Preconditioner &preconditioner,
            const MPI_Comm communicator,
            std::shared_ptr<LifetimeProbe> probe
        )
            : m_operation(&operation),
              m_preconditioner(&preconditioner),
              m_communicator(communicator),
              m_rightHandSide(operation.Height()),
              m_correction(operation.Width()),
              m_probe(std::move(probe)) {
            m_rightHandSide = 0.0;
            m_correction    = 0.0;
            if (m_probe != nullptr) {
                m_probe->problemIdentity = std::addressof(operation.GetProblem());
            }
        }

        PreparedBackend(const PreparedBackend &)            = delete;
        PreparedBackend &operator=(const PreparedBackend &) = delete;
        PreparedBackend(PreparedBackend &&)                 = delete;
        PreparedBackend &operator=(PreparedBackend &&)      = delete;

        ~PreparedBackend() {
            if (m_probe != nullptr) {
                m_probe->backendDestroyed = true;
                m_probe->dependenciesAliveAtBackendDestruction =
                    m_operation != nullptr && m_preconditioner != nullptr && m_preconditioner->IsCurrent();
            }
        }

        [[nodiscard]] const Operator &GetOperator() const noexcept {
            return *m_operation;
        }

        [[nodiscard]] const Preconditioner &GetPreconditioner() const noexcept {
            return *m_preconditioner;
        }

        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept {
            return m_communicator;
        }

        [[nodiscard]] bool IsReady() const noexcept {
            return m_operation != nullptr && m_preconditioner != nullptr && m_communicator != MPI_COMM_NULL &&
                   m_rightHandSide.Size() == m_operation->Height() && m_correction.Size() == m_operation->Width();
        }

        [[nodiscard]] int RightHandSideSize() const noexcept {
            return m_rightHandSide.Size();
        }

        [[nodiscard]] int CorrectionSize() const noexcept {
            return m_correction.Size();
        }

        [[nodiscard]] mean_field::solver::LinearSolveReport Solve(
            const mfem::Vector &rightHandSide,
            mfem::Vector &correction,
            const mean_field::solver::LinearSolveControl &control
        ) {
            control.Validate();
            if (rightHandSide.Size() != RightHandSideSize() || correction.Size() != CorrectionSize()) {
                throw std::invalid_argument("The internal test backend requires preallocated compatible vectors.");
            }
            m_rightHandSide = rightHandSide;
            m_correction    = 0.0;
            correction      = m_correction;
            return {
                .status               = mean_field::solver::LinearSolveStatus::converged,
                .control              = control,
                .initialResidualNorm  = rightHandSide.Norml2(),
                .reportedResidualNorm = 0.0,
                .trueResidualNorm     = 0.0
            };
        }

    private:
        const Operator *m_operation;
        Preconditioner *m_preconditioner;
        MPI_Comm m_communicator;
        mfem::Vector m_rightHandSide;
        mfem::Vector m_correction;
        std::shared_ptr<LifetimeProbe> m_probe;
    };

    template <
        typename Operator,
        typename Preconditioner>
    [[nodiscard]] auto prepareLinearBackend(
        InspectingBackend configuration,
        const Operator &operation,
        Preconditioner &preconditioner,
        const MPI_Comm communicator
    ) {
        return PreparedBackend<Operator, Preconditioner>{
            operation, preconditioner, communicator, std::move(configuration.probe)
        };
    }

    template <
        typename Operator,
        typename Preconditioner>
    [[nodiscard]] auto prepareLinearBackend(
        ThrowingBackend,
        const Operator &,
        Preconditioner &,
        MPI_Comm
    )
        -> PreparedBackend<
            Operator,
            Preconditioner> {
        throw std::runtime_error("The internal test backend rejected restart preparation.");
    }

    [[nodiscard]] mean_field::fem::FEM makeFiniteElements() {
        mean_field::utils::Args arguments;
        arguments.mesh_file = "sandbox.smesh";
        arguments.p.rtol    = 1.0e-12;
        arguments.p.atol    = 1.0e-12;
        return mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    }

    [[nodiscard]] double matchingCentralDensity() {
        constexpr double stellarRadius = mean_field::utils::RADIUS;
        constexpr double targetMass    = mean_field::utils::MASS;
        return std::numbers::pi_v<double> * targetMass / (4.0 * stellarRadius * stellarRadius * stellarRadius);
    }

    [[nodiscard]] auto makeModel() {
        using namespace mean_field;
        constexpr double stellarRadius  = utils::RADIUS;
        constexpr double targetMass     = utils::MASS;
        const double polytropicConstant = 2.0 * utils::G * stellarRadius * stellarRadius / std::numbers::pi_v<double>;
        return model::StellarModel(
            eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
            surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
            integral::FixedTotalMass({.Mtotal = dimensions::MassValue{targetMass}}),
            constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{matchingCentralDensity()}})
        );
    }

    [[nodiscard]] auto makeGeneratedRotationModel() {
        using namespace mean_field;
        constexpr double stellarRadius  = utils::RADIUS;
        constexpr double targetMass     = utils::MASS;
        const double polytropicConstant = 2.0 * utils::G * stellarRadius * stellarRadius / std::numbers::pi_v<double>;
        return model::StellarModel(
            eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
            surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
            integral::FixedTotalMass({.Mtotal = dimensions::MassValue{targetMass}}),
            integral::FixedAngularMomentum({.Jtotal = dimensions::AngularMomentumValue{0.05}}),
            constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{matchingCentralDensity()}})
        );
    }

    void checkZeroRotation(const mean_field::physics::RigidRotation &rotation) {
        REQUIRE(rotation.angular_velocity().Size() == 3);
        REQUIRE(rotation.center().Size() == 3);
        for (int component = 0; component < 3; ++component) {
            CHECK(rotation.angular_velocity()(component) == 0.0);
            CHECK(rotation.center()(component) == 0.0);
        }
    }
} // namespace solver_internal_architecture_test

TEST_CASE(
    "Report Views Retain Context-Owned Accepted State After Solver Destruction",
    "[solver][architecture][ownership][report][internal]"
) {
    using namespace mean_field;
    using namespace solver_internal_architecture_test;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());
    const MPI_Comm expectedCommunicator = finiteElements.mesh->GetComm();
    auto probe                          = std::make_shared<LifetimeProbe>();
    {
        auto context = solver::makeContext(
            makeModel(), equilibrium::StellarDiscretization{std::move(finiteElements)},
            preconditioning::makePreconditioner(), InspectingBackend{probe}
        );
        REQUIRE(context.isReady());
        REQUIRE(probe->problemIdentity != nullptr);

        auto report = [&] {
            auto borrowingSolver =
                solver::make(context, solver::nonlinear::Newton(solver::nonlinear::NewtonOptions{}, ZeroMetric{}));
            return borrowingSolver.evaluate();
        }();
        CHECK_FALSE(context.hasActiveSolver());
        CHECK_FALSE(probe->backendDestroyed);
        REQUIRE(report.converged());

        auto structure = report.structureView();
        REQUIRE(structure.valid());
        CHECK(
            structure.model().template specification<constraint::FixedCentralDensity>().targetDensity() ==
            dimensions::DensityValue{matchingCentralDensity()}
        );
        const auto state = structure.state();
        REQUIRE_FALSE(state.empty());
        CHECK(std::ranges::all_of(state, [](const mfem::real_t value) { return std::isfinite(value); }));
        REQUIRE_FALSE(structure.stateDescriptors().empty());
        REQUIRE_FALSE(structure.stateBlock(utils::blocks::density_field.mass_term).empty());

        int communicatorComparison = MPI_UNEQUAL;
        REQUIRE(
            MPI_Comm_compare(structure.communicator(), expectedCommunicator, &communicatorComparison) == MPI_SUCCESS
        );
        CHECK((communicatorComparison == MPI_IDENT || communicatorComparison == MPI_CONGRUENT));
        REQUIRE(structure.prescribedRotation().has_value());
        checkZeroRotation(*structure.prescribedRotation());
        checkZeroRotation(structure.rotation());
        CHECK_THROWS_AS(structure.capture(), std::logic_error);
    }

    CHECK(probe->backendDestroyed);
    CHECK(probe->dependenciesAliveAtBackendDestruction);

    auto rejectedFiniteElements = makeFiniteElements();
    CHECK_THROWS_AS(
        solver::makeContext(
            makeModel(), equilibrium::StellarDiscretization{std::move(rejectedFiniteElements)},
            preconditioning::makePreconditioner(), ThrowingBackend{}
        ),
        std::runtime_error
    );
}

TEST_CASE(
    "Generated Rotation Is Reported Without A Prescribed Rotation Payload",
    "[solver][architecture][ownership][rotation][report][internal]"
) {
    using namespace mean_field;
    using namespace solver_internal_architecture_test;

    auto finiteElements = makeFiniteElements();
    REQUIRE(finiteElements.okay());
    auto probe   = std::make_shared<LifetimeProbe>();
    auto context = solver::makeContext(
        makeGeneratedRotationModel(), equilibrium::StellarDiscretization{std::move(finiteElements)},
        preconditioning::makePreconditioner(), InspectingBackend{probe}
    );

    REQUIRE(context.isReady());
    auto borrowingSolver =
        solver::make(context, solver::nonlinear::Newton(solver::nonlinear::NewtonOptions{}, ZeroMetric{}));
    auto report    = borrowingSolver.evaluate();
    auto structure = report.structureView();

    CHECK_FALSE(probe->backendDestroyed);
    REQUIRE(structure.valid());
    CHECK_FALSE(structure.prescribedRotation().has_value());
    const auto rotation = structure.rotation();
    REQUIRE(rotation.angular_velocity().Size() == 3);
    REQUIRE(rotation.center().Size() == 3);
    CHECK(rotation.angular_velocity()(0) == 0.0);
    CHECK(rotation.angular_velocity()(1) == 0.0);
    CHECK(rotation.angular_velocity()(2) > 0.0);
    for (int component = 0; component < 3; ++component) {
        CHECK(std::isfinite(rotation.angular_velocity()(component)));
        CHECK(rotation.center()(component) == 0.0);
    }
}