#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace preconditioning_runtime_test {
    namespace blocks   = mean_field::utils::blocks;

    using Form         = blocks::surface_deformed_stellar_equilibrium_form;
    using JacobianForm = blocks::surface_deformed_stellar_equilibrium_jacobian_form;
    using Layout       = blocks::form_layout<Form>;

    class Manifest final {
    public:
        Manifest()
            : m_layout(
                  std::array<
                      int,
                      Form::value_block_count>{
                      2,
                      3,
                      4,
                      5,
                      6,
                      1
                  },
                  std::array<
                      int,
                      Form::residual_block_count>{
                      4,
                      5,
                      2,
                      3,
                      6,
                      1
                  }
              ) {
        }

        [[nodiscard]] const Layout &layout() const noexcept {
            return m_layout;
        }

    private:
        Layout m_layout;
    };

    class Problem final {
    public:
        Problem() : m_linearization(m_manifest.layout().value_offsets().Last()) {
            m_snapshot.discretization               = {.identity = 11, .revision = 1};
            m_snapshot.geometry                     = {.identity = 12, .revision = 1};
            m_snapshot.equationOfStateIdentity      = &m_equationOfStateToken;
            m_snapshot.linearization.discretization = m_snapshot.discretization;
            m_snapshot.linearization.density        = {.identity = 21, .revision = 1};
        }

        void AdvanceDensity() noexcept {
            ++m_snapshot.linearization.density.revision;
        }

        void AdvanceGeometry() noexcept {
            ++m_snapshot.geometry.revision;
        }

        void AdvancePreparation() noexcept {
            ++m_snapshot.preparedOperatorGeneration;
        }

        void SetPrepared(const bool prepared) noexcept {
            m_prepared = prepared;
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_prepared;
        }

        [[nodiscard]] int StateSize() const noexcept {
            return m_manifest.layout().value_offsets().Last();
        }

        [[nodiscard]] int EquationSize() const noexcept {
            return m_manifest.layout().residual_offsets().Last();
        }

        [[nodiscard]] const Manifest &GetManifest() const noexcept {
            return m_manifest;
        }

        [[nodiscard]] const mfem::Operator &GetLinearizationOperator() const noexcept {
            return m_linearization;
        }

        [[nodiscard]] mean_field::preconditioning::StellarPreconditionerLifecycleSnapshot Snapshot() const {
            return m_snapshot;
        }

    private:
        Manifest m_manifest;
        mfem::IdentityOperator m_linearization;
        std::uint8_t m_equationOfStateToken{0};
        mean_field::preconditioning::StellarPreconditionerLifecycleSnapshot m_snapshot;
        bool m_prepared{true};
    };
} // namespace preconditioning_runtime_test

namespace unsupported_physical_preconditioner_test {
    class Constraint;
    class PreparedConstraint;

    class Constraint final {
    public:
        struct Parameters final {
            mean_field::dimensions::SpecificEnthalpyValue target;
        };

        using TargetValue       = mean_field::dimensions::SpecificEnthalpyValue;
        using ScalarDescription = mean_field::stellar::ScalarConstraint<
            mean_field::dimensions::quantity::SpecificEnthalpy,
            mean_field::dimensions::quantity::Dimensionless,
            mean_field::dimensions::quantity::SpecificEnthalpy,
            "test.unsupported_physical_edge.coordinate",
            "q_u",
            "test.unsupported_physical_edge.residual",
            "R_u">;
        using ModelDefinition = mean_field::constraint::ScalarPhaseCondition<
            Constraint,
            "UnsupportedPhysicalPreconditionerEdge",
            mean_field::stellar::
                Reads<mean_field::stellar::state::SpecificEnthalpy, mean_field::stellar::state::OwnGeneratedCoordinate>,
            mean_field::stellar::Changes<mean_field::stellar::equation::PoissonEquation>,
            ScalarDescription>;
        using EquilibriumPhysics = mean_field::operators::LocalSpecificationEquilibriumPhysics<PreparedConstraint>;

        explicit constexpr Constraint(const Parameters parameters) noexcept : m_target(parameters.target) {
        }

        [[nodiscard]] constexpr TargetValue target() const noexcept {
            return m_target;
        }

    private:
        TargetValue m_target;
    };

    struct PreparationReport final {
        bool stateChanged{true};
    };

    class PreparedConstraint final {
    public:
        using Report = PreparationReport;

        explicit PreparedConstraint(const Constraint &) noexcept {
        }

        template <typename StateView> [[nodiscard]] Report PrepareAfterPhysical(const StateView &) noexcept {
            m_isPrepared = true;
            return {};
        }

        template <typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddResidual(
            mean_field::stellar::equation::OwnConstraint,
            Row &
        ) const noexcept {
            return mean_field::stellar::structuralZero;
        }

        template <typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddResidual(
            mean_field::stellar::equation::PoissonEquation,
            Row &
        ) const noexcept {
            return mean_field::stellar::structuralZero;
        }

        template <
            typename Direction,
            typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddJacobianAction(
            mean_field::stellar::Derivative<
                mean_field::stellar::equation::OwnConstraint,
                mean_field::stellar::state::SpecificEnthalpy>,
            const Direction &,
            Row &
        ) const noexcept {
            return mean_field::stellar::zeroDerivative;
        }

        template <
            typename Direction,
            typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddJacobianAction(
            mean_field::stellar::Derivative<
                mean_field::stellar::equation::OwnConstraint,
                mean_field::stellar::state::OwnGeneratedCoordinate>,
            const Direction &,
            Row &
        ) const noexcept {
            return mean_field::stellar::zeroDerivative;
        }

        template <
            typename Direction,
            typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddJacobianAction(
            mean_field::stellar::Derivative<
                mean_field::stellar::equation::PoissonEquation,
                mean_field::stellar::state::SpecificEnthalpy>,
            const Direction &,
            Row &
        ) const noexcept {
            return mean_field::stellar::zeroDerivative;
        }

        template <
            typename Direction,
            typename Row>
        [[nodiscard]] mean_field::stellar::StructuralZero AddJacobianAction(
            mean_field::stellar::Derivative<
                mean_field::stellar::equation::PoissonEquation,
                mean_field::stellar::state::OwnGeneratedCoordinate>,
            const Direction &,
            Row &
        ) const noexcept {
            return mean_field::stellar::zeroDerivative;
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_isPrepared;
        }

    private:
        bool m_isPrepared{false};
    };
} // namespace unsupported_physical_preconditioner_test

template <> struct mean_field::preconditioning::StellarEquilibriumProblemTraits<preconditioning_runtime_test::Problem> {
    using Problem                    = preconditioning_runtime_test::Problem;
    using Form                       = preconditioning_runtime_test::Form;
    using JacobianForm               = preconditioning_runtime_test::JacobianForm;
    using Manifest                   = preconditioning_runtime_test::Manifest;

    static constexpr bool registered = true;

    [[nodiscard]] static bool IsPrepared(const Problem &problem) noexcept {
        return problem.IsPrepared();
    }

    [[nodiscard]] static int StateSize(const Problem &problem) noexcept {
        return problem.StateSize();
    }

    [[nodiscard]] static int EquationSize(const Problem &problem) noexcept {
        return problem.EquationSize();
    }

    [[nodiscard]] static const Manifest &ManifestOf(const Problem &problem) noexcept {
        return problem.GetManifest();
    }

    [[nodiscard]] static const mfem::Operator &LinearizationOperator(const Problem &problem) noexcept {
        return problem.GetLinearizationOperator();
    }

    [[nodiscard]] static mean_field::preconditioning::StellarPreconditionerLifecycleSnapshot
    Snapshot(const Problem &problem) {
        return problem.Snapshot();
    }
};

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

    using ModelWithoutPhase           = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass>>;
    using CentralDensityModel         = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass,
        mean_field::models::FixedCentralDensity>>;
    using AngularMomentumModel        = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass,
        mean_field::models::FixedAngularMomentum>>;
    using ProvenZeroPhysicalEdgeModel = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass,
        unsupported_physical_preconditioner_test::Constraint>>;

    using ProblemWithoutPhase         = mean_field::equilibrium::StellarEquilibriumProblem<ModelWithoutPhase>;
    using CentralDensityProblem       = mean_field::equilibrium::StellarEquilibriumProblem<CentralDensityModel>;
    using AngularMomentumProblem      = mean_field::equilibrium::StellarEquilibriumProblem<AngularMomentumModel>;
    using ProvenZeroPhysicalEdgeProblem =
        mean_field::equilibrium::StellarEquilibriumProblem<ProvenZeroPhysicalEdgeModel>;
    using PlanWithoutPhase   = preconditioning::IdentityPreconditionerPlanFor<ProblemWithoutPhase>;
    using CentralDensityPlan = preconditioning::IdentityPreconditionerPlanFor<CentralDensityProblem>;

    template <typename Problem>
    concept CanMakeDefaultStellarStructureBlock =
        requires(const Problem &problem) { preconditioning::stellarStructureBlock(problem); };

    using RefreshingDensityIdentity = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::density::mass::value>,
        blocks::type_list<blocks::density::mass::residual>,
        blocks::type_list<>,
        preconditioning::IdentityOperatorCharacteristics,
        preconditioning::backend::Identity,
        preconditioning::PreparationDependencies<preconditioning::PreparationDependency::linearization>>;
    using SurfaceIdentity = preconditioning::IdentityBlock<
        blocks::surface_deformation::parameters::value,
        blocks::surface_deformation::shape_equilibrium::residual>;
    using GravityGradientIdentity =
        preconditioning::IdentityBlock<blocks::gravity::gradient::value, blocks::gravity::gradient::residual>;
    using GravityPotentialIdentity =
        preconditioning::IdentityBlock<blocks::gravity::poisson::value, blocks::gravity::poisson::residual>;
    using EnthalpyIdentity =
        preconditioning::IdentityBlock<blocks::enthalpy::specific::value, blocks::enthalpy::specific::residual>;
    using FixedMassIdentity = preconditioning::IdentityBlock<
        blocks::fixed_total_mass::mass_normalization::value,
        blocks::fixed_total_mass::mass_normalization::residual>;
    using SelectiveRefreshPlan = preconditioning::PreconditionerPlan<
        RefreshingDensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        FixedMassIdentity>;

    [[nodiscard]] constexpr SelectiveRefreshPlan makeSelectiveRefreshPlan() {
        return SelectiveRefreshPlan{RefreshingDensityIdentity{}, SurfaceIdentity{},  GravityGradientIdentity{},
                                    GravityPotentialIdentity{},  EnthalpyIdentity{}, FixedMassIdentity{}};
    }
} // namespace

TEST_CASE(
    "Stellar Identity Plans Follow The Compiled Equilibrium Problem Type",
    tags::preconditioning_runtime_unit
) {
    STATIC_CHECK(mean_field::equilibrium::DiscretizedStellarEquilibriumProblem<ProblemWithoutPhase>);
    STATIC_CHECK(mean_field::equilibrium::DiscretizedStellarEquilibriumProblem<CentralDensityProblem>);
    STATIC_CHECK(preconditioning::StellarPreconditionerProblem<ProblemWithoutPhase>);
    STATIC_CHECK(preconditioning::StellarPreconditionerProblem<CentralDensityProblem>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<PlanWithoutPhase, typename ProblemWithoutPhase::FormType>);
    STATIC_CHECK(
        preconditioning::CompletePreconditionerFor<CentralDensityPlan, typename CentralDensityProblem::FormType>
    );
    STATIC_CHECK(PlanWithoutPhase::ComponentTypes::size == 6);
    STATIC_CHECK(CentralDensityPlan::ComponentTypes::size == 7);
}

TEST_CASE(
    "Default Stellar Structure Availability Distinguishes Proven Zeros From Unhandled Physical Edges",
    "[preconditioning][stellar_structure][type_contract][compiler]"
) {
    using BaseCompilation = mean_field::operators::CompiledStellarEquilibriumSystem<ModelWithoutPhase>;
    using BaseSupport     = preconditioning::DefaultStellarStructurePhysicalTopologySupport<ModelWithoutPhase>;
    using ProvenZeroSupport =
        preconditioning::DefaultStellarStructurePhysicalTopologySupport<ProvenZeroPhysicalEdgeModel>;
    using TrustedFixedMassEdge = mean_field::operators::StellarEquilibriumJacobianCoupling<
        blocks::enthalpy::specific::residual, blocks::density::mass::value>;
    using ProvenZeroPoissonEnthalpyEdge = mean_field::operators::StellarEquilibriumJacobianCoupling<
        blocks::gravity::poisson::residual, blocks::enthalpy::specific::value>;

    // FixedTotalMass contributes h <- rho outside the generic five-field base
    // graph. It remains supported because that specification is explicitly
    // embedded in the trusted numerical core, not because of a model-pack
    // special case.
    STATIC_CHECK_FALSE(
        mean_field::utils::blocks::contains_type_v<
            TrustedFixedMassEdge, typename BaseCompilation::BaseJacobianCouplings>
    );
    STATIC_CHECK(
        mean_field::utils::blocks::contains_type_v<
            TrustedFixedMassEdge, typename BaseCompilation::ContributionJacobianCouplings>
    );
    STATIC_CHECK(BaseSupport::UnsupportedCouplings::size == 0);
    STATIC_CHECK(preconditioning::DefaultStellarStructurePhysicalTopologySupportedFor<ModelWithoutPhase>);
    STATIC_CHECK(preconditioning::DefaultStellarStructurePhysicalTopologySupportedFor<CentralDensityModel>);
    STATIC_CHECK(preconditioning::DefaultStellarStructurePhysicalTopologySupportedFor<AngularMomentumModel>);
    STATIC_CHECK(preconditioning::StellarStructurePreconditionerProblem<ProblemWithoutPhase>);
    STATIC_CHECK(CanMakeDefaultStellarStructureBlock<ProblemWithoutPhase>);
    STATIC_CHECK(preconditioning::DefaultStellarPreconditionerAvailableFor<ProblemWithoutPhase>);

    // The mock's novel Poisson <- enthalpy edge is absent from the structure
    // backend, but its exact nested provider returns StructuralZero.  That is
    // a compile-time proof that no preconditioner term is missing; generated-
    // coordinate edges are handled independently by the inferred border.
    STATIC_CHECK(mean_field::equilibrium::StellarEquilibriumModel<ProvenZeroPhysicalEdgeModel>);
    STATIC_CHECK(mean_field::equilibrium::DiscretizedStellarEquilibriumProblem<ProvenZeroPhysicalEdgeProblem>);
    STATIC_CHECK(ProvenZeroSupport::UnsupportedCouplings::size == 0);
    STATIC_CHECK(
        mean_field::utils::blocks::contains_type_v<
            ProvenZeroPoissonEnthalpyEdge, typename mean_field::operators::CompiledStellarEquilibriumSystem<
                                               ProvenZeroPhysicalEdgeModel>::ContributionJacobianCouplings>
    );
    STATIC_CHECK(preconditioning::DefaultStellarStructurePhysicalTopologySupportedFor<ProvenZeroPhysicalEdgeModel>);
    STATIC_CHECK(preconditioning::StellarStructurePreconditionerProblem<ProvenZeroPhysicalEdgeProblem>);
    STATIC_CHECK(CanMakeDefaultStellarStructureBlock<ProvenZeroPhysicalEdgeProblem>);
    STATIC_CHECK(preconditioning::DefaultStellarPreconditionerAvailableFor<ProvenZeroPhysicalEdgeProblem>);
}

TEST_CASE(
    "Prepared Stellar Identity Preconditioning Is Bitwise Equivalent To The P0 Baseline",
    tags::preconditioning_runtime_unit
) {
    preconditioning_runtime_test::Problem problem;
    auto plan           = preconditioning::makeIdentityPlan(problem);
    using Plan          = decltype(plan);
    using Problem       = preconditioning_runtime_test::Problem;
    auto preconditioner = preconditioning::prepare(problem, std::move(plan));

    STATIC_CHECK(preconditioning::PreparedPreconditionerPlanFor<Plan, Problem>);
    CHECK(preconditioner.Height() == problem.StateSize());
    CHECK(preconditioner.Width() == problem.EquationSize());
    CHECK(preconditioner.IsCurrent());
    CHECK(&preconditioner.GetLinearizationOperator() == &problem.GetLinearizationOperator());

    preconditioner.SetOperator(problem.GetLinearizationOperator());

    mfem::Vector residual(problem.EquationSize());
    mfem::Vector correction(problem.StateSize());
    for (int index = 0; index < residual.Size(); ++index) {
        residual(index) = static_cast<double>(index) - 10.25;
    }
    correction                                  = -1.0;

    const mfem::real_t *const correctionStorage = correction.GetData();
    const auto statisticsBefore                 = preconditioner.GetStatistics();
    preconditioner.Mult(residual, correction);
    const auto statisticsAfter = preconditioner.GetStatistics();

    CHECK(correction.GetData() == correctionStorage);
    CHECK(std::memcmp(correction.GetData(), residual.GetData(), sizeof(mfem::real_t) * residual.Size()) == 0);

    const mfem::Vector densityCorrection = preconditioner.GetCorrectionBlock<blocks::density::mass::value>(correction);
    const mfem::Vector densityResidual   = preconditioner.GetResidualBlock<blocks::density::mass::residual>(residual);
    CHECK(densityCorrection.Size() == 2);
    CHECK(densityCorrection.GetData() == correction.GetData());
    CHECK(densityResidual.Size() == 2);
    CHECK(densityResidual.GetData() == residual.GetData() + 9);
    CHECK(statisticsAfter.setups == statisticsBefore.setups);
    CHECK(statisticsAfter.refreshes == statisticsBefore.refreshes);
    CHECK(statisticsAfter.componentSetups == 6);
    CHECK(statisticsAfter.applications == statisticsBefore.applications + 1);
    CHECK(statisticsAfter.backendApplications == statisticsBefore.backendApplications + 1);
    CHECK(statisticsAfter.innerIterations == 0);
    CHECK(statisticsAfter.operatorBindings == 1);
}

TEST_CASE(
    "Stellar Preconditioner Refresh Is Explicit And Dependency Aware",
    tags::preconditioning_runtime_unit
) {
    preconditioning_runtime_test::Problem problem;
    auto preconditioner = preconditioning::prepare(problem, preconditioning::makeIdentityPlan(problem));
    mfem::Vector residual(problem.EquationSize());
    mfem::Vector correction(problem.StateSize());
    residual            = 1.0;
    correction          = 0.0;

    const auto noChange = preconditioner.Refresh();
    CHECK_FALSE(noChange.changes.Any());
    CHECK_FALSE(noChange.DidAnyWork());
    CHECK(preconditioner.GetStatistics().noOpRefreshes == 1);

    problem.AdvanceDensity();
    CHECK_FALSE(preconditioner.IsCurrent());
    CHECK_THROWS_AS(preconditioner.Mult(residual, correction), std::logic_error);

    const auto linearizationRefresh = preconditioner.Refresh();
    CHECK(linearizationRefresh.changes.linearization);
    CHECK_FALSE(linearizationRefresh.changes.discretization);
    CHECK_FALSE(linearizationRefresh.changes.geometry);
    CHECK_FALSE(linearizationRefresh.DidAnyWork());
    CHECK(preconditioner.IsCurrent());
    CHECK(preconditioner.GetStatistics().refreshes == 1);
    CHECK(preconditioner.GetStatistics().componentRefreshes == 0);

    problem.AdvanceGeometry();
    const auto geometryRefresh = preconditioner.Refresh();
    CHECK(geometryRefresh.changes.geometry);
    CHECK_FALSE(geometryRefresh.changes.linearization);
    CHECK(preconditioner.GetStatistics().refreshes == 2);

    problem.AdvancePreparation();
    CHECK_FALSE(preconditioner.IsCurrent());
    const auto preparationRefresh = preconditioner.Refresh();
    CHECK(preparationRefresh.changes.linearization);
    CHECK_FALSE(preparationRefresh.changes.geometry);
    CHECK(preconditioner.IsCurrent());
    CHECK(preconditioner.GetStatistics().refreshes == 3);

    problem.SetPrepared(false);
    CHECK_FALSE(preconditioner.IsCurrent());
    CHECK_THROWS_AS(preconditioner.Refresh(), std::logic_error);
}

TEST_CASE(
    "Stellar Preconditioner Refresh Touches Only Components With Changed Dependencies",
    tags::preconditioning_runtime_unit
) {
    preconditioning_runtime_test::Problem problem;
    auto preconditioner = preconditioning::prepare(problem, makeSelectiveRefreshPlan());

    problem.AdvanceGeometry();
    const auto geometryRefresh = preconditioner.Refresh();
    CHECK(geometryRefresh.changes.geometry);
    CHECK_FALSE(geometryRefresh.changes.linearization);
    CHECK_FALSE(geometryRefresh.DidAnyWork());
    CHECK(geometryRefresh.refreshedComponents == 0);

    problem.AdvanceDensity();
    const auto linearizationRefresh = preconditioner.Refresh();
    CHECK_FALSE(linearizationRefresh.changes.geometry);
    CHECK(linearizationRefresh.changes.linearization);
    CHECK(linearizationRefresh.DidAnyWork());
    CHECK(linearizationRefresh.refreshedComponents == 1);
    CHECK(preconditioner.GetStatistics().componentRefreshes == 1);
}

TEST_CASE(
    "Stellar Preconditioner Application Requires Preallocated Compatible Vectors",
    tags::preconditioning_runtime_unit
) {
    preconditioning_runtime_test::Problem problem;
    auto preconditioner = preconditioning::prepare(problem, preconditioning::makeIdentityPlan(problem));
    mfem::Vector residual(problem.EquationSize());
    mfem::Vector missingCorrection;
    mfem::IdentityOperator wrongOperator(problem.StateSize() - 1);

    CHECK_THROWS_AS(preconditioner.Mult(residual, missingCorrection), std::invalid_argument);
    CHECK_THROWS_AS(preconditioner.SetOperator(wrongOperator), std::invalid_argument);

    preconditioning_runtime_test::Problem unpreparedProblem;
    unpreparedProblem.SetPrepared(false);
    CHECK_THROWS_AS(
        preconditioning::prepare(unpreparedProblem, preconditioning::makeIdentityPlan(unpreparedProblem)),
        std::logic_error
    );
}