#include <concepts>
#include <type_traits>
#include <utility>

#include <catch2/catch_test_macros.hpp>

import mean_field;

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

    using PhysicalForm = blocks::surface_deformed_stellar_equilibrium_form;
    using PhaseForm = blocks::central_density_bordered_stellar_equilibrium_form;
    using PhysicalPlan = normalization::PhysicalRieszNormalizationPlanFor<PhysicalForm>;
    using PhasePlan = normalization::PhysicalRieszNormalizationPlanFor<PhaseForm>;

    struct ValueA final : blocks::value_block_base { };
    struct ValueB final : blocks::value_block_base { };
    struct ResidualA final : blocks::residual_block_base { };
    struct ResidualB final : blocks::residual_block_base { };
    struct ForeignValue final : blocks::value_block_base { };
    struct ForeignResidual final : blocks::residual_block_base { };

    using SmallForm = blocks::block_form<
        blocks::type_list<ValueA, ValueB>,
        blocks::type_list<ResidualA, ResidualB>>;
    using ValueAIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::value,
        blocks::type_list<ValueA>,
        normalization::IdentityCoordinate>;
    using ValueBIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::value,
        blocks::type_list<ValueB>,
        normalization::IdentityCoordinate>;
    using ResidualAIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::residual,
        blocks::type_list<ResidualA>,
        normalization::IdentityCoordinate>;
    using ResidualBIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::residual,
        blocks::type_list<ResidualB>,
        normalization::IdentityCoordinate>;
    using ForeignIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::value,
        blocks::type_list<ForeignValue>,
        normalization::IdentityCoordinate>;
    using ForeignResidualIdentity = normalization::CoordinateComponent<
        normalization::CoordinateKind::residual,
        blocks::type_list<ForeignResidual>,
        normalization::IdentityCoordinate>;

    using CompleteSmallPlan = normalization::NormalizationPlan<
        ValueAIdentity,
        ValueBIdentity,
        ResidualAIdentity,
        ResidualBIdentity>;
    using MissingSmallPlan = normalization::NormalizationPlan<
        ValueAIdentity,
        ResidualAIdentity,
        ResidualBIdentity>;
    using DuplicateSmallPlan = normalization::NormalizationPlan<
        ValueAIdentity,
        ValueAIdentity,
        ValueBIdentity,
        ResidualAIdentity,
        ResidualBIdentity>;
    using ForeignSmallPlan = normalization::NormalizationPlan<
        ValueAIdentity,
        ValueBIdentity,
        ForeignIdentity,
        ResidualAIdentity,
        ResidualBIdentity,
        ForeignResidualIdentity>;

    struct MalformedComponent final {
        using Blocks = blocks::type_list<ValueA>;
        using Method = normalization::IdentityCoordinate;
    };

    struct IncoherentComponent final {
        using Blocks                         = blocks::type_list<ValueA>;
        using Method                         = normalization::IdentityCoordinate;
        using ValueBlocks                    = blocks::type_list<ValueB>;
        using ResidualBlocks                 = blocks::type_list<>;
        static constexpr auto kind           = normalization::CoordinateKind::value;
    };

    using WrongDensityTopology = normalization::CoordinateComponent<
        normalization::CoordinateKind::value,
        blocks::type_list<blocks::density::mass::value>,
        normalization::PhysicalRieszCoordinate<
            normalization::RieszTopology::vector_volume_l2,
            normalization::PhysicalScaleKind::density>>;

    struct FutureInvariantValue final : blocks::value_block_base { };
    struct FutureInvariantResidual final : blocks::residual_block_base { };
    using UnregisteredFutureForm = blocks::block_form<
        blocks::type_list<FutureInvariantValue>,
        blocks::type_list<FutureInvariantResidual>>;

    using BaseModel = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::integral::FixedTotalMass>>;
    using RieszPolicy = normalization::PhysicalRieszDiagonal<>;
    using RieszDiscretization = mean_field::equilibrium::StellarDiscretizationFor<RieszPolicy>;
    using BaselineProblem = mean_field::equilibrium::StellarEquilibriumProblem<BaseModel>;
    using RieszProblem = mean_field::equilibrium::StellarEquilibriumProblem<BaseModel, RieszDiscretization>;
    using AngularModel = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::integral::FixedTotalMass,
        mean_field::integral::FixedAngularMomentum>>;
    using AngularForm = mean_field::operators::CompiledStellarEquilibriumForm<AngularModel>;

    template <typename Mapper>
    concept CanMakeRieszDiscretization = requires(
        mean_field::fem::FEM &finiteElements,
        Mapper &&mapper,
        RieszPolicy policy
    ) {
        mean_field::equilibrium::makeStellarDiscretization(
            finiteElements,
            std::forward<Mapper>(mapper),
            policy
        );
    };
} // namespace

TEST_CASE("Normalization Plans Prove Exact Ownership Of Every Compiled Coordinate", "[normalization][type]") {
    STATIC_CHECK(normalization::NormalizationPlanType<PhysicalPlan>);
    STATIC_CHECK(normalization::CompleteNormalizationFor<PhysicalPlan, PhysicalForm>);
    STATIC_CHECK(normalization::CompleteNormalizationFor<PhasePlan, PhaseForm>);
    STATIC_CHECK(normalization::CompilableNormalizationFor<normalization::Unnormalized, PhysicalForm>);
    STATIC_CHECK(normalization::CompilableNormalizationFor<RieszPolicy, PhysicalForm>);
    STATIC_CHECK(normalization::CompilableNormalizationFor<RieszPolicy, PhaseForm>);
    STATIC_CHECK(normalization::CompilableNormalizationFor<RieszPolicy, AngularForm>);
    STATIC_CHECK(normalization::StellarSpecificationNormalizationContribution<
                 mean_field::integral::FixedAngularMomentum>::registered);

    STATIC_CHECK(normalization::CompleteNormalizationFor<CompleteSmallPlan, SmallForm>);
    STATIC_CHECK_FALSE(normalization::CompleteNormalizationFor<MissingSmallPlan, SmallForm>);
    STATIC_CHECK_FALSE(normalization::CompleteNormalizationFor<DuplicateSmallPlan, SmallForm>);
    STATIC_CHECK_FALSE(normalization::CompleteNormalizationFor<ForeignSmallPlan, SmallForm>);

    using Missing = normalization::NormalizationCoverage<SmallForm, MissingSmallPlan>;
    using Duplicate = normalization::NormalizationCoverage<SmallForm, DuplicateSmallPlan>;
    using Foreign = normalization::NormalizationCoverage<SmallForm, ForeignSmallPlan>;
    STATIC_CHECK(Missing::MissingValueBlocks::size == 1);
    STATIC_CHECK(blocks::contains_type_v<ValueB, typename Missing::MissingValueBlocks>);
    STATIC_CHECK(Duplicate::RepeatedValueBlocks::size == 1);
    STATIC_CHECK(blocks::contains_type_v<ValueA, typename Duplicate::RepeatedValueBlocks>);
    STATIC_CHECK(Foreign::UnexpectedValueBlocks::size == 1);
    STATIC_CHECK(Foreign::UnexpectedResidualBlocks::size == 1);
}

TEST_CASE("Physical Riesz Methods Reject Incompatible Or Unregistered Field Topologies", "[normalization][type]") {
    STATIC_CHECK_FALSE(normalization::NormalizationComponent<MalformedComponent>);
    STATIC_CHECK_FALSE(normalization::NormalizationComponent<IncoherentComponent>);
    STATIC_CHECK_FALSE(normalization::NormalizationComponent<WrongDensityTopology>);
    STATIC_CHECK_FALSE(normalization::CompilableNormalizationFor<RieszPolicy, UnregisteredFutureForm>);
    STATIC_CHECK(normalization::CompilableNormalizationFor<normalization::Unnormalized, UnregisteredFutureForm>);

    using Density = normalization::PhysicalRieszBlockTraits<blocks::density::mass::value>;
    using Gravity = normalization::PhysicalRieszBlockTraits<blocks::gravity::gradient::value>;
    using Surface = normalization::PhysicalRieszBlockTraits<blocks::surface_deformation::parameters::value>;
    using EnthalpyResidual = normalization::PhysicalRieszBlockTraits<blocks::enthalpy::specific::residual>;
    using MassResidual = normalization::PhysicalRieszBlockTraits<
        blocks::fixed_total_mass::mass_normalization::residual>;

    STATIC_CHECK(Density::Method::topology == normalization::RieszTopology::scalar_volume_l2);
    STATIC_CHECK(Gravity::Method::topology == normalization::RieszTopology::vector_volume_l2);
    STATIC_CHECK(Surface::Method::topology == normalization::RieszTopology::scalar_boundary_l2);
    STATIC_CHECK(
        EnthalpyResidual::Method::topology == normalization::RieszTopology::hybrid_scalar_volume_point_rows
    );
    STATIC_CHECK(MassResidual::Method::topology == normalization::RieszTopology::global_scalar);
    STATIC_CHECK(MassResidual::Method::scale == normalization::PhysicalScaleKind::mass);
}

TEST_CASE("Normalization Is Part Of The Compile-Time Discretization And Problem Type", "[normalization][type]") {
    STATIC_CHECK(mean_field::equilibrium::StellarDiscretizationType<RieszDiscretization>);
    STATIC_CHECK_FALSE(std::same_as<RieszDiscretization, mean_field::equilibrium::StellarDiscretization>);
    STATIC_CHECK_FALSE(std::same_as<RieszProblem, BaselineProblem>);
    STATIC_CHECK(std::same_as<typename BaselineProblem::NormalizationPrescriptionType, normalization::Unnormalized>);
    STATIC_CHECK(std::same_as<typename RieszProblem::NormalizationPrescriptionType, RieszPolicy>);
    STATIC_CHECK(mean_field::equilibrium::DiscretizedStellarEquilibriumProblem<RieszProblem>);
    STATIC_CHECK(std::constructible_from<
                 RieszDiscretization,
                 mean_field::fem::FEM &,
                 const mean_field::mapping::DomainMapper &,
                 RieszPolicy>);
    STATIC_CHECK_FALSE(std::constructible_from<
                       RieszDiscretization,
                       mean_field::fem::FEM &,
                       mean_field::mapping::DomainMapper &&,
                       RieszPolicy>);
    STATIC_CHECK_FALSE(std::constructible_from<
                       RieszDiscretization,
                       mean_field::fem::FEM &,
                       const mean_field::mapping::DomainMapper &&,
                       RieszPolicy>);
    STATIC_CHECK(std::constructible_from<
                 mean_field::equilibrium::StellarDiscretization,
                 mean_field::fem::FEM &,
                 const mean_field::mapping::DomainMapper &>);
    STATIC_CHECK_FALSE(std::constructible_from<
                       mean_field::equilibrium::StellarDiscretization,
                       mean_field::fem::FEM &,
                       mean_field::mapping::DomainMapper &&>);
    STATIC_CHECK_FALSE(std::constructible_from<
                       mean_field::equilibrium::StellarDiscretization,
                       mean_field::fem::FEM &,
                       const mean_field::mapping::DomainMapper &&>);
    STATIC_CHECK(CanMakeRieszDiscretization<
                 mean_field::mapping::DomainMapper &>);
    STATIC_CHECK_FALSE(CanMakeRieszDiscretization<
                       mean_field::mapping::DomainMapper>);
    STATIC_CHECK_FALSE(CanMakeRieszDiscretization<
                       const mean_field::mapping::DomainMapper>);

    using SmallLayout = blocks::form_layout<SmallForm>;
    using SmallBuilder = normalization::DiagonalNormalizationBuilder<SmallForm>;
    STATIC_CHECK(std::constructible_from<SmallBuilder, const SmallLayout &>);
    STATIC_CHECK_FALSE(std::constructible_from<SmallBuilder, SmallLayout &&>);
    STATIC_CHECK_FALSE(std::constructible_from<SmallBuilder, const SmallLayout &&>);
}