#include <concepts>
#include <type_traits>

#include <catch2/catch_test_macros.hpp>

import mean_field;
import test_helpers;

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

    using Form                = blocks::surface_deformed_stellar_equilibrium_form;
    using JacobianForm        = blocks::surface_deformed_stellar_equilibrium_jacobian_form;
    using CentralForm         = blocks::central_density_bordered_stellar_equilibrium_form;
    using CentralJacobianForm = blocks::central_density_bordered_stellar_equilibrium_jacobian_form;

    using DensityIdentity =
        preconditioning::IdentityBlock<blocks::density::mass::value, blocks::density::mass::residual>;
    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 MassIdentity = preconditioning::IdentityBlock<
        blocks::fixed_total_mass::mass_normalization::value,
        blocks::fixed_total_mass::mass_normalization::residual>;
    using CentralDensityIdentity = preconditioning::IdentityBlock<
        blocks::fixed_central_density::central_value::value,
        blocks::fixed_central_density::central_value::residual>;

    using IdentityPlan = preconditioning::PreconditionerPlan<
        DensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        MassIdentity>;

    using CentralIdentityPlan = preconditioning::PreconditionerPlan<
        DensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        MassIdentity,
        CentralDensityIdentity>;

    using IncompleteCentralPlan = IdentityPlan;

    struct AlternateDensityIdentity final : preconditioning::ComponentDeclaration<
                                                blocks::type_list<blocks::density::mass::value>,
                                                blocks::type_list<blocks::density::mass::residual>,
                                                blocks::type_list<>,
                                                preconditioning::IdentityOperatorCharacteristics,
                                                preconditioning::backend::Identity> { };

    using DuplicateOwnershipPlan = preconditioning::PreconditionerPlan<
        DensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        MassIdentity,
        CentralDensityIdentity,
        AlternateDensityIdentity>;

    using ExplicitOverlapPlan = preconditioning::OverlappingPreconditionerPlan<
        DensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        MassIdentity,
        CentralDensityIdentity,
        AlternateDensityIdentity>;

    struct ExtraCorrection final : blocks::value_block_base { };
    struct ExtraResidual final : blocks::residual_block_base { };
    using ExtraIdentity           = preconditioning::IdentityBlock<ExtraCorrection, ExtraResidual>;
    using UnexpectedOwnershipPlan = preconditioning::PreconditionerPlan<
        DensityIdentity,
        SurfaceIdentity,
        GravityGradientIdentity,
        GravityPotentialIdentity,
        EnthalpyIdentity,
        MassIdentity,
        ExtraIdentity>;

    using CoupledGravity = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::gravity::gradient::value, blocks::gravity::poisson::value>,
        blocks::type_list<blocks::gravity::gradient::residual, blocks::gravity::poisson::residual>,
        blocks::type_list<
            preconditioning::Coupling<blocks::gravity::gradient::residual, blocks::gravity::gradient::value>,
            preconditioning::Coupling<blocks::gravity::gradient::residual, blocks::gravity::poisson::value>,
            preconditioning::Coupling<blocks::gravity::poisson::residual, blocks::gravity::gradient::value>>,
        preconditioning::IdentityOperatorCharacteristics,
        preconditioning::backend::Identity>;

    using ValidCoupledPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, CoupledGravity, EnthalpyIdentity, MassIdentity>;

    using InvalidGravityCoupling = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::gravity::gradient::value, blocks::gravity::poisson::value>,
        blocks::type_list<blocks::gravity::gradient::residual, blocks::gravity::poisson::residual>,
        blocks::type_list<preconditioning::Coupling<blocks::gravity::gradient::residual, blocks::density::mass::value>>,
        preconditioning::IdentityOperatorCharacteristics,
        preconditioning::backend::Identity>;

    using InvalidCoupledPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, InvalidGravityCoupling, EnthalpyIdentity, MassIdentity>;

    using IncompatibleBackendComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<blocks::gravity::gradient::value>,
        blocks::type_list<blocks::gravity::gradient::residual>,
        blocks::type_list<>,
        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>,
        preconditioning::backend::HypreBoomerAMG<>>;

    struct IncoherentPlanDeclaration final {
        using ComponentTypes    = blocks::type_list<DensityIdentity>;
        using CorrectionBlocks  = blocks::type_list<blocks::gravity::gradient::value>;
        using ResidualBlocks    = blocks::type_list<blocks::density::mass::residual>;
        using RequiredCouplings = blocks::type_list<>;
    };
} // namespace

TEST_CASE(
    "Preconditioning Backends Advertise Compile-Time Operator Compatibility",
    tags::preconditioning_type_contract
) {
    using ScalarElliptic = 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 VectorElliptic = 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 LocalDenseBorder = preconditioning::OperatorCharacteristics<
        preconditioning::OperatorCategory::dense_border, preconditioning::OperatorValueStructure::block,
        preconditioning::OperatorSymmetry::nonsymmetric, preconditioning::OperatorDefiniteness::indefinite,
        preconditioning::OperatorRepresentation::assembled_dense, preconditioning::OperatorDistribution::local>;

    using FixedAMG = preconditioning::backend::HypreBoomerAMG<preconditioning::backend::FixedCycles>;

    STATIC_CHECK(preconditioning::backend::Registered<FixedAMG>);
    STATIC_CHECK(preconditioning::backend::Compatible<FixedAMG, ScalarElliptic>);
    STATIC_CHECK_FALSE(preconditioning::backend::Compatible<FixedAMG, VectorElliptic>);
    STATIC_CHECK(preconditioning::backend::Compatible<preconditioning::backend::DenseDirect, LocalDenseBorder>);
    STATIC_CHECK_FALSE(preconditioning::PreconditionerComponent<IncompatibleBackendComponent>);
    STATIC_CHECK(
        preconditioning::backend::applicationContract<FixedAMG> ==
        preconditioning::ApplicationContract::stationary_linear
    );
}

TEST_CASE(
    "Preconditioner Plans Prove Complete Unique Ownership Of Every Equilibrium Block",
    tags::preconditioning_type_contract
) {
    STATIC_CHECK(preconditioning::PreconditionerComponent<DensityIdentity>);
    STATIC_CHECK(preconditioning::PreconditionerPlanType<IdentityPlan>);
    STATIC_CHECK_FALSE(preconditioning::PreconditionerPlanType<IncoherentPlanDeclaration>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<IdentityPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<IdentityPlan, Form, JacobianForm>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<CentralIdentityPlan, CentralForm>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<CentralIdentityPlan, CentralForm, CentralJacobianForm>);
    STATIC_CHECK(preconditioning::StationaryLinearPreconditionerPlan<CentralIdentityPlan>);
}

TEST_CASE(
    "Preconditioner Coverage Reports Missing Generated Borders And Rejects Accidental Overlap",
    tags::preconditioning_type_contract
) {
    using IncompleteCoverage = preconditioning::PreconditionerCoverage<CentralForm, IncompleteCentralPlan>;
    using DuplicateCoverage  = preconditioning::PreconditionerCoverage<CentralForm, DuplicateOwnershipPlan>;

    STATIC_CHECK_FALSE(preconditioning::CompletePreconditionerFor<IncompleteCentralPlan, CentralForm>);
    STATIC_CHECK(IncompleteCoverage::MissingCorrectionBlocks::size == 1);
    STATIC_CHECK(IncompleteCoverage::MissingResidualBlocks::size == 1);
    STATIC_CHECK(
        blocks::contains_type_v<
            blocks::fixed_central_density::central_value::value, IncompleteCoverage::MissingCorrectionBlocks>
    );
    STATIC_CHECK(
        blocks::contains_type_v<
            blocks::fixed_central_density::central_value::residual, IncompleteCoverage::MissingResidualBlocks>
    );

    STATIC_CHECK_FALSE(preconditioning::CompletePreconditionerFor<DuplicateOwnershipPlan, CentralForm>);
    STATIC_CHECK(DuplicateCoverage::RepeatedCorrectionBlocks::size == 1);
    STATIC_CHECK(DuplicateCoverage::RepeatedResidualBlocks::size == 1);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<ExplicitOverlapPlan, CentralForm>);
}

TEST_CASE(
    "Preconditioner Coverage Rejects Blocks Outside The Compiled Stellar Form",
    tags::preconditioning_type_contract
) {
    using Coverage = preconditioning::PreconditionerCoverage<Form, UnexpectedOwnershipPlan>;

    STATIC_CHECK_FALSE(preconditioning::CompletePreconditionerFor<UnexpectedOwnershipPlan, Form>);
    STATIC_CHECK(Coverage::UnexpectedCorrectionBlocks::size == 1);
    STATIC_CHECK(Coverage::UnexpectedResidualBlocks::size == 1);
    STATIC_CHECK(blocks::contains_type_v<ExtraCorrection, Coverage::UnexpectedCorrectionBlocks>);
    STATIC_CHECK(blocks::contains_type_v<ExtraResidual, Coverage::UnexpectedResidualBlocks>);
}

TEST_CASE(
    "Preconditioner Component Dependencies Must Exist In The Compiled Jacobian Graph",
    tags::preconditioning_type_contract
) {
    STATIC_CHECK(preconditioning::PreconditionerComponent<CoupledGravity>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<ValidCoupledPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<ValidCoupledPlan, Form, JacobianForm>);

    STATIC_CHECK(preconditioning::CompletePreconditionerFor<InvalidCoupledPlan, Form>);
    STATIC_CHECK_FALSE(preconditioning::requiredCouplingsExist<InvalidCoupledPlan, JacobianForm>);
    STATIC_CHECK_FALSE(preconditioning::CompatiblePreconditionerFor<InvalidCoupledPlan, Form, JacobianForm>);
}