#include <array>
#include <cmath>
#include <concepts>
#include <limits>
#include <memory>
#include <type_traits>
#include <utility>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    struct StellarModelExtensionTracker final {
        int structureValidationCount{0};

        const mean_field::eos::Polytrope *structureEquationOfState{nullptr};
    };

    class StellarModelTestStructure final {
    public:
        explicit StellarModelTestStructure(std::shared_ptr<StellarModelExtensionTracker> tracker)
            : m_tracker(std::move(tracker)),
              m_equationOfState(
                  3.0,
                  0.25
              ) {
        }

        [[nodiscard]] const mean_field::eos::Polytrope &equationOfState() const noexcept {
            m_tracker->structureEquationOfState = &m_equationOfState;
            return m_equationOfState;
        }

        [[nodiscard]] double targetMass() const noexcept {
            return 2.5;
        }

        [[nodiscard]] mean_field::models::structure::StructureSeed
        makeInitialSeed(const mean_field::models::structure::StructureSeedRequest &request) const {
            mean_field::models::structure::StructureSeed seed;

            seed.radius.SetSize(2);
            seed.density.SetSize(2);
            seed.enthalpy.SetSize(2);

            seed.radius(0)       = 0.0;
            seed.radius(1)       = 1.0;

            seed.density(0)      = request.centralDensity;
            seed.density(1)      = 0.0;

            seed.enthalpy(0)     = 1.0;
            seed.enthalpy(1)     = 0.0;

            seed.stellarRadius   = 1.0;
            seed.centralDensity  = request.centralDensity;
            seed.centralEnthalpy = 1.0;

            return seed;
        }

        void validate() const {
            ++m_tracker->structureValidationCount;
        }

    private:
        std::shared_ptr<StellarModelExtensionTracker> m_tracker;
        mean_field::eos::Polytrope m_equationOfState;
    };

    class StructureWithoutSeed final {
    public:
        [[nodiscard]] const mean_field::eos::Polytrope &equationOfState() const noexcept;

        [[nodiscard]] double targetMass() const noexcept;

        void validate() const;
    };

    class SurfaceWithoutPhysicalQuantity final { };

    struct ModelSurfaceState final {
        double specificEnthalpy;

        [[nodiscard]] mean_field::eos::SpecificEnthalpyValue
        value(mean_field::eos::quantity::SpecificEnthalpy) const noexcept {
            return mean_field::eos::SpecificEnthalpyValue{specificEnthalpy};
        }
    };

    using PolytropicStellarModel = mean_field::models::StellarModel<mean_field::models::structure::PolytropicStructure>;

    using ExtensionStellarModel  = mean_field::models::StellarModel<StellarModelTestStructure>;
} // namespace

TEST_CASE(
    "Stellar Model Owns Structure Prescription And Surface Condition",
    tags::stellar_model_type_contract &tags::surface_condition_type_contract
) {
    STATIC_CHECK(mean_field::models::StructurePrescription<mean_field::models::structure::PolytropicStructure>);
    STATIC_CHECK(mean_field::models::StructurePrescription<StellarModelTestStructure>);
    STATIC_CHECK_FALSE(mean_field::models::StructurePrescription<StructureWithoutSeed>);

    STATIC_CHECK(
        mean_field::models::SurfaceCondition<mean_field::surface::ConstantPressureSurface, mean_field::eos::Polytrope>
    );
    STATIC_CHECK_FALSE(
        mean_field::models::SurfaceCondition<SurfaceWithoutPhysicalQuantity, mean_field::eos::Polytrope>
    );

    STATIC_CHECK_FALSE(std::derived_from<StellarModelTestStructure, mean_field::models::structure::StructureBase>);
    STATIC_CHECK_FALSE(
        std::derived_from<
            mean_field::models::structure::PolytropicStructure, mean_field::models::structure::StructureBase>
    );
    STATIC_CHECK(
        std::same_as<
            decltype(std::declval<const mean_field::models::structure::StructureBase &>().equationOfState()),
            mean_field::eos::EquationOfStateView>
    );

    STATIC_REQUIRE_FALSE(std::is_copy_constructible_v<PolytropicStellarModel>);

    STATIC_REQUIRE_FALSE(std::is_copy_assignable_v<PolytropicStellarModel>);

    STATIC_REQUIRE(std::is_nothrow_move_constructible_v<PolytropicStellarModel>);

    STATIC_REQUIRE(std::is_nothrow_move_assignable_v<PolytropicStellarModel>);

    mean_field::models::StellarModel model{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    STATIC_CHECK(std::same_as<decltype(model), PolytropicStellarModel>);
    STATIC_CHECK(
        std::same_as<
            decltype(model.structurePrescription()), const mean_field::models::structure::PolytropicStructure &>
    );
    STATIC_CHECK(
        std::same_as<decltype(model.surfaceCondition()), const mean_field::surface::ConstantPressureSurface &>
    );
    STATIC_CHECK(std::same_as<decltype(model.equationOfState()), const mean_field::eos::Polytrope &>);
    STATIC_CHECK(
        std::same_as<
            typename PolytropicStellarModel::SurfaceDeformationPrescriptionType,
            mean_field::deformation::NodalRadialSurface>
    );
    STATIC_CHECK(
        std::same_as<
            typename PolytropicStellarModel::StellarInteriorDeformationExtensionType,
            mean_field::deformation::PowerLawRadialInteriorExtension>
    );
    STATIC_CHECK(
        std::same_as<
            typename PolytropicStellarModel::VacuumDeformationExtensionType,
            mean_field::deformation::FixedInfinityRadialVacuumExtension>
    );

    CHECK(model.targetMass() == 1.0);
    CHECK(model.compiledSurfaceConstraint().targetPressure() == mean_field::eos::PressureValue{0.0});
    CHECK(&model.equationOfState() == &model.structurePrescription().equationOfState());
    CHECK(model.surfaceCondition().targetPressure() == mean_field::eos::PressureValue{0.0});
    CHECK(model.surfaceDeformationPrescription().descriptor().name == "NodalRadialSurface");
    CHECK(model.stellarInteriorDeformationExtension().radialPower() == 2.0);
    CHECK(model.vacuumDeformationExtension().descriptor().name == "FixedInfinityRadialVacuumExtension");
}

TEST_CASE(
    "Stellar Model Owns Explicit Surface Interior And Vacuum Deformation Policies",
    tags::stellar_model_deformation_ownership
) {
    mfem::Vector referenceCenter(3);
    referenceCenter(0) = 0.125;
    referenceCenter(1) = -0.25;
    referenceCenter(2) = 0.375;

    mean_field::models::StellarModel model{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}},
        mean_field::deformation::NodalRadialSurface{referenceCenter},
        mean_field::deformation::PowerLawRadialInteriorExtension{3.0},
        mean_field::deformation::FixedInfinityRadialVacuumExtension{}
    };

    CHECK(model.surfaceDeformationPrescription().referenceCenter()(0) == referenceCenter(0));
    CHECK(model.surfaceDeformationPrescription().referenceCenter()(1) == referenceCenter(1));
    CHECK(model.surfaceDeformationPrescription().referenceCenter()(2) == referenceCenter(2));
    CHECK(model.stellarInteriorDeformationExtension().radialPower() == 3.0);
    CHECK(
        model.vacuumDeformationExtension().descriptor().outerBoundaryBehavior ==
        mean_field::deformation::VacuumOuterBoundaryBehavior::FixedAtReferenceInfinity
    );

    const auto *surfaceDeformationAddress  = &model.surfaceDeformationPrescription();
    const auto *interiorDeformationAddress = &model.stellarInteriorDeformationExtension();
    const auto *vacuumDeformationAddress   = &model.vacuumDeformationExtension();
    auto movedModel                        = std::move(model);

    CHECK(&movedModel.surfaceDeformationPrescription() == surfaceDeformationAddress);
    CHECK(&movedModel.stellarInteriorDeformationExtension() == interiorDeformationAddress);
    CHECK(&movedModel.vacuumDeformationExtension() == vacuumDeformationAddress);
}

TEST_CASE(
    "Stellar Model Delegates Seed Construction To Its Structure",
    tags::barotrope &tags::unit &tags::model
) {
    mean_field::models::StellarModel model{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    const mean_field::models::structure::StructureSeed seed =
        model.makeInitialSeed({.centralDensity = 2.0, .radialSampleCount = 64});

    CHECK(seed.radius.Size() == 64);
    CHECK(seed.density.Size() == 64);
    CHECK(seed.enthalpy.Size() == 64);
    CHECK(seed.centralDensity == 2.0);
    CHECK(seed.stellarRadius > 0.0);
    CHECK(seed.density(0) == 2.0);
    CHECK(seed.density(63) == 0.0);
    CHECK(seed.enthalpy(63) == 0.0);
}

TEST_CASE(
    "Moving A Stellar Model Preserves Stable Prescription Addresses",
    tags::barotrope &tags::unit &tags::model &tags::surface_constraint_lifetime
) {
    mean_field::models::StellarModel originalModel{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    const mean_field::models::structure::PolytropicStructure *structureAddress = &originalModel.structurePrescription();

    const mean_field::surface::ConstantPressureSurface *surfaceAddress         = &originalModel.surfaceCondition();

    const mean_field::eos::Polytrope *equationOfStateAddress                   = &originalModel.equationOfState();

    const auto *compiledSurfaceConstraintAddress = &originalModel.compiledSurfaceConstraint();

    mean_field::models::StellarModel movedModel{std::move(originalModel)};

    CHECK(&movedModel.structurePrescription() == structureAddress);
    CHECK(&movedModel.surfaceCondition() == surfaceAddress);
    CHECK(&movedModel.equationOfState() == equationOfStateAddress);
    CHECK(&movedModel.compiledSurfaceConstraint() == compiledSurfaceConstraintAddress);
    CHECK(movedModel.targetMass() == 1.0);
}

TEST_CASE(
    "Stellar Model Compiles A Positive Constant Pressure Surface",
    tags::barotrope &tags::unit &tags::model
) {
    constexpr double targetPressure = 0.03125;

    mean_field::models::StellarModel model{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{targetPressure}}
    };

    const double requiredSpecificEnthalpy = mean_field::eos::evaluate<mean_field::eos::quantity::SpecificEnthalpy>(
                                                model.equationOfState(), mean_field::eos::PressureValue{targetPressure}
    )
                                                .value();

    CHECK(requiredSpecificEnthalpy > 0.0);
    CHECK(model.compiledSurfaceConstraint().targetPressure() == mean_field::eos::PressureValue{targetPressure});
    CHECK(model.compiledSurfaceConstraint().residual(ModelSurfaceState{requiredSpecificEnthalpy}) == 0.0);
}

TEST_CASE(
    "Stellar Model Supports Custom Structure Prescriptions And Surface Conditions",
    tags::barotrope &tags::unit &tags::model
) {
    const auto tracker = std::make_shared<StellarModelExtensionTracker>();

    mean_field::models::StellarModel model{
        StellarModelTestStructure{tracker},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.375}}
    };

    STATIC_CHECK(std::same_as<decltype(model), ExtensionStellarModel>);

    REQUIRE(tracker->structureValidationCount == 1);

    const mean_field::eos::Polytrope *ownedEquationOfState = &model.equationOfState();

    CHECK(tracker->structureEquationOfState == ownedEquationOfState);

    CHECK(model.targetMass() == 2.5);

    CHECK(model.compiledSurfaceConstraint().targetPressure() == mean_field::eos::PressureValue{0.375});
}

TEST_CASE(
    "Move Assignment Preserves Stellar Model Prescription Addresses",
    tags::barotrope &tags::unit &tags::model
) {
    mean_field::models::StellarModel sourceModel{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.25},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    mean_field::models::StellarModel destinationModel{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{2.0, 0.5}, 4.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.02}}
    };

    const mean_field::models::structure::PolytropicStructure *sourceStructureAddress =
        &sourceModel.structurePrescription();

    const mean_field::surface::ConstantPressureSurface *sourceSurfaceAddress = &sourceModel.surfaceCondition();

    const mean_field::eos::Polytrope *sourceEquationOfStateAddress           = &sourceModel.equationOfState();
    const auto *sourceSurfaceDeformationAddress  = &sourceModel.surfaceDeformationPrescription();
    const auto *sourceInteriorDeformationAddress = &sourceModel.stellarInteriorDeformationExtension();
    const auto *sourceVacuumDeformationAddress   = &sourceModel.vacuumDeformationExtension();

    const mean_field::eos::PressureValue sourceTargetPressure =
        sourceModel.compiledSurfaceConstraint().targetPressure();

    destinationModel = std::move(sourceModel);

    CHECK(&destinationModel.structurePrescription() == sourceStructureAddress);

    CHECK(&destinationModel.surfaceCondition() == sourceSurfaceAddress);

    CHECK(&destinationModel.equationOfState() == sourceEquationOfStateAddress);

    CHECK(&destinationModel.surfaceDeformationPrescription() == sourceSurfaceDeformationAddress);

    CHECK(&destinationModel.stellarInteriorDeformationExtension() == sourceInteriorDeformationAddress);

    CHECK(&destinationModel.vacuumDeformationExtension() == sourceVacuumDeformationAddress);

    CHECK(destinationModel.targetMass() == 1.25);

    CHECK(destinationModel.compiledSurfaceConstraint().targetPressure() == sourceTargetPressure);
}

TEST_CASE(
    "Stellar Model View Supports Heterogeneous Typed Models",
    tags::stellar_model_runtime_view
) {
    const auto tracker = std::make_shared<StellarModelExtensionTracker>();

    const mean_field::models::StellarModel polytropicModel{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    const mean_field::models::StellarModel extensionModel{
        StellarModelTestStructure{tracker},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.375}}
    };

    STATIC_CHECK(std::is_trivially_copyable_v<mean_field::models::StellarModelView>);
    STATIC_CHECK_FALSE(std::constructible_from<mean_field::models::StellarModelView, PolytropicStellarModel &&>);

    const std::array views{
        mean_field::models::StellarModelView{polytropicModel}, mean_field::models::StellarModelView{extensionModel}
    };

    CHECK(views[0].targetMass() == 1.0);
    CHECK(views[1].targetMass() == 2.5);
    CHECK(views[1].surfaceCondition().targetPressure == 0.375);
    CHECK(views[0].surfaceDeformation().name == "NodalRadialSurface");
    CHECK(views[0].surfaceDeformation().motionKind == mean_field::deformation::SurfaceMotionKind::Radial);
    CHECK(views[0].stellarInteriorDeformation().name == "PowerLawRadialInteriorExtension");
    CHECK(
        views[0].vacuumDeformation().outerBoundaryBehavior ==
        mean_field::deformation::VacuumOuterBoundaryBehavior::FixedAtReferenceInfinity
    );
    REQUIRE(views[1].surfaceDependencies().stateFields.size() == 1);
    CHECK(
        views[1].surfaceDependencies().residualRowField ==
        mean_field::surface::surfaceFieldId<mean_field::field::Enthalpy>
    );

    const auto pressure =
        views[0].equationOfState().tryEvaluate<mean_field::eos::quantity::Pressure>(mean_field::eos::DensityValue{0.7});

    REQUIRE(pressure.has_value());
    CHECK(
        pressure->value() == mean_field::eos::evaluate<mean_field::eos::quantity::Pressure>(
                                 polytropicModel.equationOfState(), mean_field::eos::DensityValue{0.7}
                             )
                                 .value()
    );

    const mean_field::models::structure::StructureSeed seed =
        views[1].makeInitialSeed({.centralDensity = 1.75, .radialSampleCount = 2});

    CHECK(seed.centralDensity == 1.75);
    CHECK(seed.radius.Size() == 2);
}

TEST_CASE(
    "Stellar Model View Retains Stable Pointees When Its Owner Moves",
    tags::stellar_model_runtime_view
) {
    mean_field::models::StellarModel originalModel{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
    };

    const mean_field::models::StellarModelView view{originalModel};
    PolytropicStellarModel movedModel{std::move(originalModel)};

    const auto pressure =
        view.equationOfState().tryEvaluate<mean_field::eos::quantity::Pressure>(mean_field::eos::DensityValue{0.7});
    const mean_field::models::structure::StructureSeed seed =
        view.makeInitialSeed({.centralDensity = 1.0, .radialSampleCount = 8});

    REQUIRE(pressure.has_value());
    CHECK(view.targetMass() == movedModel.targetMass());
    CHECK(view.surfaceDeformation() == movedModel.surfaceDeformationPrescription().descriptor());
    CHECK(view.stellarInteriorDeformation() == movedModel.stellarInteriorDeformationExtension().descriptor());
    CHECK(view.vacuumDeformation() == movedModel.vacuumDeformationExtension().descriptor());
    CHECK(
        pressure->value() == mean_field::eos::evaluate<mean_field::eos::quantity::Pressure>(
                                 movedModel.equationOfState(), mean_field::eos::DensityValue{0.7}
                             )
                                 .value()
    );
    CHECK(seed.radius.Size() == 8);
}

TEST_CASE(
    "Stellar Model Compiles Its Deformation Policies Against The Finite Element Discretization",
    tags::stellar_model_deformation_compilation
) {
    mean_field::utils::Args args = test_utils::setup_args();
    mean_field::fem::FEM fem     = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());

    mfem::Vector referenceCenter(fem.mesh->SpaceDimension());
    referenceCenter = 0.0;
    const mean_field::models::StellarModel model{
        mean_field::models::structure::PolytropicStructure{mean_field::eos::Polytrope{3.0, 0.25}, 1.0},
        mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}},
        mean_field::deformation::NodalRadialSurface{referenceCenter},
        mean_field::deformation::PowerLawRadialInteriorExtension{3.0},
        mean_field::deformation::FixedInfinityRadialVacuumExtension{}
    };

    auto prepared = model.compileDomainDeformation(fem);
    STATIC_CHECK(mean_field::deformation::PreparedDomainDeformationOperator<decltype(prepared)>);
    CHECK(prepared.matchesCurrentDiscretization());
    CHECK(prepared.stellarInteriorExtension().radialPower() == 3.0);
    CHECK(prepared.discretizationDependencies().physicalMeshIdentity == fem.mesh.get());
    CHECK(prepared.discretizationDependencies().logicalReferenceMeshIdentity == fem.logicalReferenceMesh.get());
    CHECK(prepared.parameterCount() == prepared.surfaceDeformationPrescription().parameterCount());
    CHECK(prepared.volumeDisplacementSize() == fem.displacementFes->GetTrueVSize());

    STATIC_CHECK_FALSE(std::is_copy_constructible_v<mean_field::deformation::PreparedDomainDeformationRuntime>);
    STATIC_CHECK(std::is_nothrow_move_constructible_v<mean_field::deformation::PreparedDomainDeformationRuntime>);
    mean_field::deformation::PreparedDomainDeformationRuntime runtime{std::move(prepared)};
    STATIC_CHECK(mean_field::deformation::PreparedDomainDeformationOperator<decltype(runtime)>);
    mean_field::deformation::PreparedDomainDeformationRuntime movedRuntime{std::move(runtime)};

    CHECK(movedRuntime.matchesCurrentDiscretization());
    CHECK(movedRuntime.parameterCount() > 0);
    CHECK(movedRuntime.volumeDisplacementSize() == fem.displacementFes->GetTrueVSize());
    CHECK(movedRuntime.compositionReport().assignedScalarDofCount() == fem.surfaceDeformationFes->GetTrueVSize());

    mfem::Vector zeroParameters(movedRuntime.parameterCount());
    mfem::Vector volumeDisplacement(movedRuntime.volumeDisplacementSize());
    zeroParameters = 0.0;
    movedRuntime.buildVolumeDisplacement(zeroParameters, volumeDisplacement);
    CHECK(volumeDisplacement.Norml2() == 0.0);
}