#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 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

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::operators::StellarEquilibriumSpecificationModel;
    using CentralDensityModel       = mean_field::operators::CentralDensityStellarEquilibriumSpecificationModel;

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

    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(
    "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.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
    );
}