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

#include <catch2/catch_test_macros.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 And Surface Prescriptions",
    tags::stellar_model_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::SurfacePrescription<
            mean_field::surface::ConstantPressureSurface, mean_field::eos::Polytrope>
    );
    STATIC_CHECK_FALSE(
        mean_field::models::SurfacePrescription<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.surfacePrescription()), const mean_field::surface::ConstantPressureSurface &>
    );
    STATIC_CHECK(std::same_as<decltype(model.equationOfState()), const mean_field::eos::Polytrope &>);

    CHECK(model.targetMass() == 1.0);
    CHECK(model.compiledSurfaceConstraint().targetPressure() == mean_field::eos::PressureValue{0.0});
    CHECK(&model.equationOfState() == &model.structurePrescription().equationOfState());
    CHECK(model.surfacePrescription().targetPressure() == mean_field::eos::PressureValue{0.0});
}

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

    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.surfacePrescription() == 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 And Surface Prescriptions",
    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.surfacePrescription();

    const mean_field::eos::Polytrope *sourceEquationOfStateAddress           = &sourceModel.equationOfState();

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

    destinationModel = std::move(sourceModel);

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

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

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

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