#include <cmath>
#include <concepts>
#include <cstdint>
#include <numbers>
#include <stdexcept>
#include <type_traits>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    template <typename... Specifications>
    using ProjectionModelWith = mean_field::model::StellarModel<
        mean_field::models::SpecificationSet<Specifications...>>;

    class UnregisteredProjectionConstraint final {
    public:
        struct Parameters final { };
        using ModelDefinition = mean_field::constraint::PhaseCondition<
            UnregisteredProjectionConstraint,
            "UnregisteredProjectionConstraint",
            mean_field::models::DependsOn<mean_field::models::stellar::state::SpecificEnthalpy>,
            mean_field::models::Affects<mean_field::models::stellar::equation::HydrostaticBalance>>;

        explicit constexpr UnregisteredProjectionConstraint(Parameters) noexcept {
        }
    };

    class ExplicitNoChangeProjectionConstraint final {
    public:
        struct Parameters final { };
        using ModelDefinition = mean_field::constraint::PhaseCondition<
            ExplicitNoChangeProjectionConstraint,
            "ExplicitNoChangeProjectionConstraint",
            mean_field::models::DependsOn<mean_field::models::stellar::state::SpecificEnthalpy>,
            mean_field::models::Affects<mean_field::models::stellar::equation::HydrostaticBalance>>;
        using RadialProjection = mean_field::seed::projection::Use<mean_field::seed::projection::NoStateChange>;

        explicit constexpr ExplicitNoChangeProjectionConstraint(Parameters) noexcept {
        }
    };

    struct IncompleteProjectionPhysics final {
        static constexpr bool registered         = true;
        static constexpr bool providesRadialMass = false;

        template <typename Model>
        static constexpr bool supports = true;
    };

    class IncompleteProjectionConstraint final {
    public:
        struct Parameters final { };
        using ModelDefinition = mean_field::constraint::PhaseCondition<
            IncompleteProjectionConstraint,
            "IncompleteProjectionConstraint",
            mean_field::models::DependsOn<mean_field::models::stellar::state::SpecificEnthalpy>,
            mean_field::models::Affects<mean_field::models::stellar::equation::HydrostaticBalance>>;
        using RadialProjection = mean_field::seed::projection::Use<IncompleteProjectionPhysics>;

        explicit constexpr IncompleteProjectionConstraint(Parameters) noexcept {
        }
    };

    class UnregisteredProjectionEquationOfState final {
    public:
        struct Parameters final { };
        using ModelDefinition = mean_field::eos::ConstitutiveLaw<
            UnregisteredProjectionEquationOfState,
            "UnregisteredProjectionEquationOfState">;

        explicit constexpr UnregisteredProjectionEquationOfState(Parameters) noexcept {
        }
    };

    class UnregisteredProjectionSurface final {
    public:
        struct Parameters final { };
        using ModelDefinition = mean_field::surface::BoundaryCondition<
            UnregisteredProjectionSurface,
            "UnregisteredProjectionSurface">;

        explicit constexpr UnregisteredProjectionSurface(Parameters) noexcept {
        }
    };

    struct SecondRadialMassProjectionPhysics final {
        static constexpr bool registered         = true;
        static constexpr bool providesRadialMass = true;

        template <typename Model>
        static constexpr bool supports = true;

        template <typename Specification>
        [[nodiscard]] static mean_field::dimensions::MassValue targetMass(const Specification &specification) {
            return specification.targetMass();
        }

        template <typename Specification, typename Model>
        static void validate(
            const Specification &,
            const Model &,
            const mean_field::seed::RadialProfile &,
            const mean_field::seed::StellarEquilibriumProjectionOptions &
        ) noexcept {
        }

        template <typename Specification, typename Model>
        static void initialize(
            const Specification &,
            const Model &,
            const mean_field::seed::RadialProjectionScales &,
            mean_field::seed::RadialProjectionState &,
            mfem::Vector coordinate
        ) {
            if (coordinate.Size() == 1) {
                coordinate(0) = 0.0;
            }
        }
    };

    class SecondRadialMassConstraint final {
    public:
        struct Parameters final {
            mean_field::dimensions::MassValue mass;
        };
        using ModelDefinition = mean_field::integral::FixedWithMultiplier<
            SecondRadialMassConstraint,
            "SecondRadialMassConstraint">;
        using RadialProjection = mean_field::seed::projection::Use<SecondRadialMassProjectionPhysics>;

        explicit constexpr SecondRadialMassConstraint(Parameters parameters) noexcept : m_mass(parameters.mass) {
        }

        [[nodiscard]] constexpr mean_field::dimensions::MassValue targetMass() const noexcept {
            return m_mass;
        }

    private:
        mean_field::dimensions::MassValue m_mass;
    };

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies make_dependencies() {
        return {
            .discretization     = {.identity = 7001, .revision = 1},
            .density            = {.identity = 7003, .revision = 1},
            .surfaceDeformation = {.identity = 7009, .revision = 1},
            .gravityGradient    = {.identity = 7013, .revision = 1},
            .gravityPotential   = {.identity = 7019, .revision = 1},
            .enthalpy           = {.identity = 7027, .revision = 1},
            .bernoulliConstant  = {.identity = 7039, .revision = 1},
            .rotation           = {.identity = 7043, .revision = 1},
            .targetMass         = {.identity = 7057, .revision = 1}
        };
    }

    [[nodiscard]] mean_field::physics::RigidRotation make_zero_rotation() {
        mfem::Vector angularVelocity(3);
        mfem::Vector center(3);
        angularVelocity = 0.0;
        center          = 0.0;
        return {angularVelocity, center};
    }

    template <typename Vector> void check_finite(const Vector &values) {
        for (int index = 0; index < values.Size(); ++index) {
            REQUIRE(std::isfinite(values(index)));
        }
    }
} // namespace

TEST_CASE(
    "Radial Projection Capabilities Are Inferred From Every Model Specification",
    tags::stellar_seed_projection_type_contract
) {
    using namespace mean_field;
    using BaseModel = ProjectionModelWith<eos::Polytrope, surface::Isobaric, integral::FixedTotalMass>;
    using CentralModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        constraint::FixedCentralDensity>;
    using AngularModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        integral::FixedAngularMomentum>;
    using ExplicitExtensionModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        ExplicitNoChangeProjectionConstraint>;
    using MissingConstraintRuleModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        UnregisteredProjectionConstraint>;
    using IncompleteConstraintRuleModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        IncompleteProjectionConstraint>;
    using MissingEquationOfStateRuleModel = ProjectionModelWith<
        UnregisteredProjectionEquationOfState,
        surface::Isobaric,
        integral::FixedTotalMass>;
    using MissingSurfaceRuleModel = ProjectionModelWith<
        eos::Polytrope,
        UnregisteredProjectionSurface,
        integral::FixedTotalMass>;
    using MissingMassProviderModel = ProjectionModelWith<eos::Polytrope, surface::Isobaric>;
    using AmbiguousMassProviderModel = ProjectionModelWith<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        SecondRadialMassConstraint>;

    STATIC_CHECK(seed::RadialProfileProjectableModel<BaseModel>);
    STATIC_CHECK(seed::RadialProfileProjectableModel<CentralModel>);
    STATIC_CHECK(seed::RadialProfileProjectableModel<AngularModel>);
    STATIC_CHECK(seed::RadialProfileProjectableModel<ExplicitExtensionModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<MissingConstraintRuleModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<IncompleteConstraintRuleModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<MissingEquationOfStateRuleModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<MissingSurfaceRuleModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<MissingMassProviderModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<AmbiguousMassProviderModel>);
    STATIC_CHECK_FALSE(seed::RadialProfileProjectableModel<int>);
}

TEST_CASE(
    "Projected Equilibrium States Preserve Their Compiled Stellar Model Type",
    tags::stellar_seed_projection_type_contract
) {
    using namespace mean_field;
    using BaseModel =
        model::StellarModel<models::SpecificationSet<eos::Polytrope, surface::Isobaric, integral::FixedTotalMass>>;
    using CentralDensityModel = model::StellarModel<models::SpecificationSet<
        eos::Polytrope, surface::Isobaric, integral::FixedTotalMass, constraint::FixedCentralDensity>>;
    using BaseState           = seed::ProjectedEquilibriumState<BaseModel>;
    using CentralDensityState = seed::ProjectedEquilibriumState<CentralDensityModel>;

    STATIC_CHECK_FALSE(std::same_as<BaseState, CentralDensityState>);
    STATIC_CHECK(std::same_as<typename BaseState::ModelType, BaseModel>);
    STATIC_CHECK(std::same_as<typename CentralDensityState::ModelType, CentralDensityModel>);
}

TEST_CASE(
    "Lane Emden Projection Builds A Complete Compiled Stellar Equilibrium State",
    tags::stellar_seed_projection
) {
    using namespace mean_field;
    using Catch::Approx;

    utils::Args args            = test_utils::setup_args();
    fem::FEM finiteElementModel = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElementModel.okay());

    constexpr double stellarRadius  = utils::RADIUS;
    constexpr double targetMass     = utils::MASS;
    const double polytropicConstant = 2.0 * utils::G * stellarRadius * stellarRadius / std::numbers::pi_v<double>;
    const double centralDensity =
        std::numbers::pi_v<double> * targetMass / (4.0 * stellarRadius * stellarRadius * stellarRadius);

    const auto stellarModel = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{targetMass}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
    );
    auto problem = equilibrium::discretize(stellarModel, finiteElementModel);

    STATIC_CHECK(seed::RadialSeedStrategyFor<seed::LaneEmden, decltype(stellarModel)>);
    STATIC_CHECK(
        std::same_as<
            decltype(seed::makeProjectedEquilibriumState(problem, seed::LaneEmden{})),
            seed::ProjectedEquilibriumState<typename std::remove_cvref_t<decltype(problem)>::ModelType>>
    );

    const auto projected = seed::makeProjectedEquilibriumState(problem, seed::LaneEmden({.radialSampleCount = 4096}));
    REQUIRE(projected.values.Size() == problem.StateSize());
    check_finite(projected.values);

    const auto stateView                = problem.GetManifest().stateView(projected.values);
    const mfem::Vector density          = stateView.block(utils::blocks::density_field.mass_term);
    const mfem::Vector surface          = stateView.block(utils::blocks::surface_deformation_field.parameters_term);
    const mfem::Vector gravityGradient  = stateView.block(utils::blocks::gravity_field.gradient_term);
    const mfem::Vector gravityPotential = stateView.block(utils::blocks::gravity_field.poisson_term);
    const mfem::Vector enthalpy         = stateView.block(utils::blocks::enthalpy_field.specific_term);
    const mfem::Vector fixedMassCoordinate =
        stateView.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);
    const mfem::Vector centralDensityBorder =
        stateView.block(utils::blocks::fixed_central_density_phase.central_value_term);

    CHECK(density.Norml2() > 0.0);
    CHECK(gravityGradient.Norml2() > 0.0);
    CHECK(gravityPotential.Norml2() > 0.0);
    CHECK(enthalpy.Norml2() > 0.0);
    CHECK(surface.Normlinf() == 0.0);
    REQUIRE(fixedMassCoordinate.Size() == 1);
    CHECK(fixedMassCoordinate(0) == Approx(-utils::G * targetMass / stellarRadius).margin(2.0e-7));
    REQUIRE(centralDensityBorder.Size() == 1);
    CHECK(centralDensityBorder(0) == 0.0);

    const auto preparation = problem.Prepare(projected.values, make_dependencies(), make_zero_rotation());
    CHECK(preparation.assembledResidual);
    CHECK(preparation.template specification<models::FixedCentralDensity>().constraint.DidAnyWork());

    mfem::Vector residual;
    problem.BuildResidual(residual);
    REQUIRE(residual.Size() == problem.EquationSize());
    check_finite(residual);

    const operators::RootConstraintReport massReport = problem.GetPreparedOperator().GetFixedMassReport();
    CHECK(std::abs(massReport.scaledResidual) < 5.0e-4);
    const operators::CentralDensityConstraintReport centralDensityReport =
        problem.GetPreparedOperator().GetCentralDensityReport();
    CHECK(centralDensityReport.targetDensity == Approx(centralDensity));
    CHECK(std::abs(centralDensityReport.enthalpyResidual) < 1.0e-10);

    const auto residualView             = problem.GetManifest().residualView(residual);
    const mfem::Vector enthalpyResidual = residualView.block(utils::blocks::enthalpy_field.specific_term);
    const auto &surfaceRows             = problem.GetPressureSurfaceRows();
    for (const int surfaceRow : surfaceRows.reduced_dofs()) {
        CHECK(enthalpy(surfaceRow) == 0.0);
        CHECK(enthalpyResidual(surfaceRow) == 0.0);
    }
}

TEST_CASE(
    "Lane Emden Projection Rejects A Seed Whose Surface Does Not Match The Reference Discretization",
    tags::stellar_seed_projection
) {
    using namespace mean_field;

    utils::Args args            = test_utils::setup_args();
    fem::FEM finiteElementModel = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElementModel.okay());

    const auto stellarModel = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{1.0}})
    );
    auto problem = equilibrium::discretize(stellarModel, finiteElementModel);

    const seed::RadialProfile mismatchedProfile =
        seed::generateRadialProfile(problem.GetStellarModel(), seed::LaneEmden({.radialSampleCount = 64}));
    CHECK_THROWS_AS(seed::projectRadialProfile(problem, mismatchedProfile), std::invalid_argument);
}

TEST_CASE(
    "Lane Emden Projection Rejects A Nonzero Isobaric Surface",
    tags::stellar_seed_projection
) {
    using namespace mean_field;

    utils::Args args            = test_utils::setup_args();
    fem::FEM finiteElementModel = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElementModel.okay());

    const double polytropicConstant = 2.0 * utils::G / std::numbers::pi_v<double>;
    const double centralDensity     = std::numbers::pi_v<double> / 4.0;
    const auto stellarModel         = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.01}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
    );
    auto problem = equilibrium::discretize(stellarModel, finiteElementModel);
    const seed::RadialProfile profile =
        seed::generateRadialProfile(problem.GetStellarModel(), seed::LaneEmden({.radialSampleCount = 64}));

    CHECK_THROWS_AS(seed::projectRadialProfile(problem, profile), std::invalid_argument);
}