#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 {
    [[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(
    "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 operators::PreparedCentralDensityStellarEquilibriumReport preparation =
        problem.Prepare(projected.values, make_dependencies(), make_zero_rotation());
    CHECK(preparation.assembledResidual);

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