#include <concepts>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    struct DeformationCompilationContext final { };

    constexpr mean_field::deformation::SurfaceDeformationDescriptor surfaceDescriptor{
        .name                               = "TestRadialSurface",
        .spatialDimension                   = 3,
        .motionKind                         = mean_field::deformation::SurfaceMotionKind::Radial,
        .linearOnReferenceGeometry          = true,
        .requiresStarShapedReferenceSurface = true,
        .hasExactDerivativeTranspose        = true,
        .hasExactPullbackDerivative         = true,
        .translationTreatment = mean_field::deformation::GeometricGaugeTreatment::ExcludedByParameterization,
        .orientationTreatment = mean_field::deformation::GeometricGaugeTreatment::ExcludedByParameterization
    };

    constexpr mean_field::deformation::InteriorDeformationExtensionDescriptor interiorDescriptor{
        .name                        = "TestRadialInteriorExtension",
        .spatialDimension            = 3,
        .linearOnReferenceGeometry   = true,
        .requiresRadialFoliation     = true,
        .requiresAuxiliarySolve      = false,
        .hasExactDerivativeTranspose = true,
        .hasExactPullbackDerivative  = true,
        .centerBehavior              = mean_field::deformation::InteriorCenterBehavior::FixedAtReferenceCenter
    };

    constexpr mean_field::deformation::VacuumDeformationExtensionDescriptor vacuumDescriptor{
        .name                        = "TestFixedInfinityExtension",
        .spatialDimension            = 3,
        .linearOnReferenceGeometry   = true,
        .requiresRadialFoliation     = true,
        .requiresAuxiliarySolve      = false,
        .hasExactDerivativeTranspose = true,
        .hasExactPullbackDerivative  = true,
        .outerBoundaryBehavior       = mean_field::deformation::VacuumOuterBoundaryBehavior::FixedAtReferenceInfinity
    };

    template <
        bool HasBuild              = true,
        bool HasJacobian           = true,
        bool HasJacobianTranspose  = true,
        bool HasPullbackDerivative = true>
    class PreparedSurfaceFixture final {
    public:
        [[nodiscard]] mean_field::deformation::SurfaceDeformationDescriptor descriptor() const noexcept {
            return surfaceDescriptor;
        }

        [[nodiscard]] int parameterCount() const noexcept {
            return 4;
        }

        [[nodiscard]] int surfaceDisplacementSize() const noexcept {
            return 12;
        }

        void buildSurfaceDisplacement(
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasBuild
        {
        }

        void applyJacobian(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobian
        {
        }

        void applyJacobianTranspose(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobianTranspose
        {
        }

        void applyPullbackDerivative(
            const mfem::Vector &,
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasPullbackDerivative
        {
        }
    };

    template <typename Prepared, bool HasDescriptor = true, bool HasValidation = true, bool HasCompilation = true>
    class SurfacePrescriptionFixture final {
    public:
        using PreparedType = Prepared;

        [[nodiscard]] mean_field::deformation::SurfaceDeformationDescriptor descriptor() const noexcept
            requires HasDescriptor
        {
            return surfaceDescriptor;
        }

        void validate() const
            requires HasValidation
        {
        }

        static constexpr bool hasCompilation = HasCompilation;
    };

    template <
        typename Prepared,
        bool HasDescriptor,
        bool HasValidation,
        bool HasCompilation>
        requires HasCompilation
    Prepared compileSurfaceDeformationPrescription(
        const SurfacePrescriptionFixture<
            Prepared,
            HasDescriptor,
            HasValidation,
            HasCompilation> &,
        const DeformationCompilationContext &
    ) {
        return {};
    }

    template <
        bool HasBuild              = true,
        bool HasJacobian           = true,
        bool HasJacobianTranspose  = true,
        bool HasPullbackDerivative = true,
        bool HasScalarDofCount     = true,
        bool HasSupportQuery       = true>
    class PreparedInteriorExtensionFixture final {
    public:
        [[nodiscard]] mean_field::deformation::InteriorDeformationExtensionDescriptor descriptor() const noexcept {
            return interiorDescriptor;
        }

        [[nodiscard]] int surfaceDisplacementSize() const noexcept {
            return 12;
        }

        [[nodiscard]] int interiorDisplacementSize() const noexcept {
            return 24;
        }

        [[nodiscard]] int scalarTrueDofCount() const noexcept
            requires HasScalarDofCount
        {
            return 8;
        }

        [[nodiscard]] bool hasStellarSupport(const int) const
            requires HasSupportQuery
        {
            return true;
        }

        void buildInteriorDisplacement(
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasBuild
        {
        }

        void applyJacobian(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobian
        {
        }

        void applyJacobianTranspose(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobianTranspose
        {
        }

        void applyPullbackDerivative(
            const mfem::Vector &,
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasPullbackDerivative
        {
        }
    };

    template <typename Prepared, bool HasDescriptor = true, bool HasValidation = true, bool HasCompilation = true>
    class InteriorExtensionFixture final {
    public:
        using PreparedType = Prepared;

        [[nodiscard]] mean_field::deformation::InteriorDeformationExtensionDescriptor descriptor() const noexcept
            requires HasDescriptor
        {
            return interiorDescriptor;
        }

        void validate() const
            requires HasValidation
        {
        }

        static constexpr bool hasCompilation = HasCompilation;
    };

    template <
        typename Prepared,
        bool HasDescriptor,
        bool HasValidation,
        bool HasCompilation>
        requires HasCompilation
    Prepared compileInteriorDeformationExtension(
        const InteriorExtensionFixture<
            Prepared,
            HasDescriptor,
            HasValidation,
            HasCompilation> &,
        const DeformationCompilationContext &
    ) {
        return {};
    }

    template <
        bool HasBuild              = true,
        bool HasJacobian           = true,
        bool HasJacobianTranspose  = true,
        bool HasPullbackDerivative = true,
        bool HasScalarDofCount     = true,
        bool HasSupportQuery       = true>
    class PreparedVacuumExtensionFixture final {
    public:
        [[nodiscard]] mean_field::deformation::VacuumDeformationExtensionDescriptor descriptor() const noexcept {
            return vacuumDescriptor;
        }

        [[nodiscard]] int surfaceDisplacementSize() const noexcept {
            return 12;
        }

        [[nodiscard]] int vacuumDisplacementSize() const noexcept {
            return 30;
        }

        [[nodiscard]] int scalarTrueDofCount() const noexcept
            requires HasScalarDofCount
        {
            return 10;
        }

        [[nodiscard]] bool hasVacuumSupport(const int) const
            requires HasSupportQuery
        {
            return true;
        }

        void buildVacuumDisplacement(
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasBuild
        {
        }

        void applyJacobian(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobian
        {
        }

        void applyJacobianTranspose(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobianTranspose
        {
        }

        void applyPullbackDerivative(
            const mfem::Vector &,
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasPullbackDerivative
        {
        }
    };

    template <typename Prepared, bool HasDescriptor = true, bool HasValidation = true, bool HasCompilation = true>
    class VacuumExtensionFixture final {
    public:
        using PreparedType = Prepared;

        [[nodiscard]] mean_field::deformation::VacuumDeformationExtensionDescriptor descriptor() const noexcept
            requires HasDescriptor
        {
            return vacuumDescriptor;
        }

        void validate() const
            requires HasValidation
        {
        }

        static constexpr bool hasCompilation = HasCompilation;
    };

    template <
        typename Prepared,
        bool HasDescriptor,
        bool HasValidation,
        bool HasCompilation>
        requires HasCompilation
    Prepared compileVacuumDeformationExtension(
        const VacuumExtensionFixture<
            Prepared,
            HasDescriptor,
            HasValidation,
            HasCompilation> &,
        const DeformationCompilationContext &
    ) {
        return {};
    }

    using CompletePreparedSurface           = PreparedSurfaceFixture<>;
    using CompleteSurfacePrescription       = SurfacePrescriptionFixture<CompletePreparedSurface>;

    using CompletePreparedInteriorExtension = PreparedInteriorExtensionFixture<>;
    using CompleteInteriorExtension         = InteriorExtensionFixture<CompletePreparedInteriorExtension>;

    using CompletePreparedVacuumExtension   = PreparedVacuumExtensionFixture<>;
    using CompleteVacuumExtension           = VacuumExtensionFixture<CompletePreparedVacuumExtension>;

    constexpr mean_field::deformation::DomainDeformationDescriptor domainDescriptor{
        .surfaceDeformation          = surfaceDescriptor,
        .stellarInteriorExtension    = interiorDescriptor,
        .vacuumExtension             = vacuumDescriptor,
        .linearOnReferenceGeometry   = true,
        .requiresAuxiliarySolve      = false,
        .hasExactDerivativeTranspose = true,
        .hasExactPullbackDerivative  = true
    };

    template <
        bool HasBuild              = true,
        bool HasJacobian           = true,
        bool HasJacobianTranspose  = true,
        bool HasPullbackDerivative = true>
    class PreparedDomainDeformationFixture final {
    public:
        [[nodiscard]] mean_field::deformation::DomainDeformationDescriptor descriptor() const noexcept {
            return domainDescriptor;
        }

        [[nodiscard]] int parameterCount() const noexcept {
            return 4;
        }

        [[nodiscard]] int surfaceDisplacementSize() const noexcept {
            return 12;
        }

        [[nodiscard]] int volumeDisplacementSize() const noexcept {
            return 24;
        }

        void buildVolumeDisplacement(
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasBuild
        {
        }

        void applyJacobian(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobian
        {
        }

        void applyJacobianTranspose(
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasJacobianTranspose
        {
        }

        void applyPullbackDerivative(
            const mfem::Vector &,
            const mfem::Vector &,
            const mfem::Vector &,
            mfem::Vector &
        ) const
            requires HasPullbackDerivative
        {
        }
    };
} // namespace

TEST_CASE(
    "Surface Deformation Contracts Require Complete Forward And Pullback Operations",
    tags::surface_deformation_type_contract
) {
    STATIC_CHECK(mean_field::deformation::PreparedSurfaceDeformationPrescription<CompletePreparedSurface>);
    STATIC_CHECK(mean_field::deformation::SurfaceDeformationPrescription<CompleteSurfacePrescription>);
    STATIC_CHECK(
        mean_field::deformation::SurfaceDeformationCompilable<
            CompleteSurfacePrescription, DeformationCompilationContext>
    );

    STATIC_CHECK_FALSE(mean_field::deformation::PreparedSurfaceDeformationPrescription<PreparedSurfaceFixture<false>>);
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedSurfaceDeformationPrescription<PreparedSurfaceFixture<true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedSurfaceDeformationPrescription<PreparedSurfaceFixture<true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedSurfaceDeformationPrescription<PreparedSurfaceFixture<true, true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::SurfaceDeformationPrescription<
            SurfacePrescriptionFixture<CompletePreparedSurface, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::SurfaceDeformationPrescription<
            SurfacePrescriptionFixture<CompletePreparedSurface, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::SurfaceDeformationCompilable<
            SurfacePrescriptionFixture<CompletePreparedSurface, true, true, false>, DeformationCompilationContext>
    );
}

TEST_CASE(
    "Interior Extension Contracts Require Complete Forward And Pullback Operations",
    tags::interior_deformation_extension_type_contract
) {
    STATIC_CHECK(mean_field::deformation::PreparedInteriorDeformationExtension<CompletePreparedInteriorExtension>);
    STATIC_CHECK(mean_field::deformation::InteriorDeformationExtension<CompleteInteriorExtension>);
    STATIC_CHECK(
        mean_field::deformation::InteriorDeformationExtensionCompilable<
            CompleteInteriorExtension, DeformationCompilationContext>
    );

    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<PreparedInteriorExtensionFixture<false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<PreparedInteriorExtensionFixture<true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<
            PreparedInteriorExtensionFixture<true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<
            PreparedInteriorExtensionFixture<true, true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<
            PreparedInteriorExtensionFixture<true, true, true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedInteriorDeformationExtension<
            PreparedInteriorExtensionFixture<true, true, true, true, true, false>>
    );
    STATIC_CHECK_FALSE(mean_field::deformation::PreparedInteriorDeformationExtension<CompletePreparedVacuumExtension>);
    STATIC_CHECK_FALSE(
        mean_field::deformation::InteriorDeformationExtension<
            InteriorExtensionFixture<CompletePreparedInteriorExtension, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::InteriorDeformationExtensionCompilable<
            InteriorExtensionFixture<CompletePreparedInteriorExtension, true, true, false>,
            DeformationCompilationContext>
    );
}

TEST_CASE(
    "Vacuum Extension Contracts Require Complete Forward And Pullback Operations",
    tags::vacuum_deformation_extension_type_contract
) {
    STATIC_CHECK(mean_field::deformation::PreparedVacuumDeformationExtension<CompletePreparedVacuumExtension>);
    STATIC_CHECK(mean_field::deformation::VacuumDeformationExtension<CompleteVacuumExtension>);
    STATIC_CHECK(
        mean_field::deformation::VacuumDeformationExtensionCompilable<
            CompleteVacuumExtension, DeformationCompilationContext>
    );

    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<PreparedVacuumExtensionFixture<false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<PreparedVacuumExtensionFixture<true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<PreparedVacuumExtensionFixture<true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<
            PreparedVacuumExtensionFixture<true, true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<
            PreparedVacuumExtensionFixture<true, true, true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedVacuumDeformationExtension<
            PreparedVacuumExtensionFixture<true, true, true, true, true, false>>
    );
    STATIC_CHECK_FALSE(mean_field::deformation::PreparedVacuumDeformationExtension<CompletePreparedInteriorExtension>);
    STATIC_CHECK_FALSE(
        mean_field::deformation::VacuumDeformationExtension<
            VacuumExtensionFixture<CompletePreparedVacuumExtension, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::VacuumDeformationExtensionCompilable<
            VacuumExtensionFixture<CompletePreparedVacuumExtension, true, true, false>, DeformationCompilationContext>
    );
}

TEST_CASE(
    "Prepared Domain Deformation Contracts Require Complete Lift And Pullback Operations",
    tags::domain_deformation_type_contract
) {
    STATIC_CHECK(mean_field::deformation::PreparedDomainDeformationOperator<PreparedDomainDeformationFixture<>>);
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedDomainDeformationOperator<PreparedDomainDeformationFixture<false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedDomainDeformationOperator<PreparedDomainDeformationFixture<true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedDomainDeformationOperator<PreparedDomainDeformationFixture<true, true, false>>
    );
    STATIC_CHECK_FALSE(
        mean_field::deformation::PreparedDomainDeformationOperator<
            PreparedDomainDeformationFixture<true, true, true, false>>
    );
}

TEST_CASE(
    "Deformation Descriptors Report Geometry And Exact Linearization Capabilities",
    tags::deformation_type_contract
) {
    STATIC_CHECK(surfaceDescriptor.isValid());
    STATIC_CHECK(surfaceDescriptor.supportsExactNewtonLinearization());
    STATIC_CHECK(interiorDescriptor.isValid());
    STATIC_CHECK(interiorDescriptor.supportsExactNewtonLinearization());
    STATIC_CHECK(vacuumDescriptor.isValid());
    STATIC_CHECK(vacuumDescriptor.supportsExactNewtonLinearization());
    STATIC_CHECK(domainDescriptor.isValid());
    STATIC_CHECK(domainDescriptor.supportsExactNewtonLinearization());

    constexpr auto invalidInteriorDescriptor = [] {
        auto descriptor           = interiorDescriptor;
        descriptor.centerBehavior = mean_field::deformation::InteriorCenterBehavior::Unspecified;
        return descriptor;
    }();

    constexpr auto invalidVacuumDescriptor = [] {
        auto descriptor                  = vacuumDescriptor;
        descriptor.outerBoundaryBehavior = mean_field::deformation::VacuumOuterBoundaryBehavior::Unspecified;
        return descriptor;
    }();

    STATIC_CHECK_FALSE(invalidInteriorDescriptor.isValid());
    STATIC_CHECK_FALSE(invalidVacuumDescriptor.isValid());

    CHECK(surfaceDescriptor.motionKind == mean_field::deformation::SurfaceMotionKind::Radial);
    CHECK(interiorDescriptor.centerBehavior == mean_field::deformation::InteriorCenterBehavior::FixedAtReferenceCenter);
    CHECK(
        vacuumDescriptor.outerBoundaryBehavior ==
        mean_field::deformation::VacuumOuterBoundaryBehavior::FixedAtReferenceInfinity
    );
}