#include <algorithm>
#include <cmath>
#include <concepts>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace {
    using BaseModel           = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::integral::FixedTotalMass>>;

    using CentralDensityModel = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::integral::FixedTotalMass,
        mean_field::constraint::FixedCentralDensity>>;

    using IncompleteModel =
        mean_field::model::StellarModel<mean_field::models::SpecificationSet<mean_field::eos::Polytrope>>;

    template <typename Candidate>
    concept HasLegacyNumericalModelAdapter = requires { typename Candidate::NumericalModelAdapter; };

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies make_dependencies() {
        return {
            .discretization     = {.identity = 4001, .revision = 1},
            .density            = {.identity = 4003, .revision = 1},
            .surfaceDeformation = {.identity = 4007, .revision = 1},
            .gravityGradient    = {.identity = 4013, .revision = 1},
            .gravityPotential   = {.identity = 4019, .revision = 1},
            .enthalpy           = {.identity = 4021, .revision = 1},
            .bernoulliConstant  = {.identity = 4027, .revision = 1},
            .rotation           = {.identity = 4049, .revision = 1},
            .targetMass         = {.identity = 4051, .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};
    }

    [[nodiscard]] double relative_difference(
        const mfem::Vector &left,
        const mfem::Vector &right
    ) {
        mfem::Vector difference(left);
        difference -= right;
        return difference.Norml2() / std::max({1.0, left.Norml2(), right.Norml2()});
    }
} // namespace

TEST_CASE(
    "Stellar Model Selects A Compile-Time Equilibrium Problem Type",
    tags::stellar_equilibrium_problem_type_contract
) {
    using namespace mean_field;

    using BaseProblem           = equilibrium::StellarEquilibriumProblem<BaseModel>;
    using CentralDensityProblem = equilibrium::StellarEquilibriumProblem<CentralDensityModel>;

    STATIC_CHECK(equilibrium::StellarEquilibriumModel<BaseModel>);
    STATIC_CHECK(equilibrium::StellarEquilibriumModel<CentralDensityModel>);
    STATIC_CHECK_FALSE(equilibrium::StellarEquilibriumModel<IncompleteModel>);
    STATIC_CHECK_FALSE(std::same_as<BaseProblem, CentralDensityProblem>);
    STATIC_CHECK(BaseProblem::symbolicallySquare);
    STATIC_CHECK(CentralDensityProblem::symbolicallySquare);
    STATIC_CHECK_FALSE(BaseProblem::hasFixedCentralDensity);
    STATIC_CHECK(CentralDensityProblem::hasFixedCentralDensity);
    STATIC_CHECK_FALSE(HasLegacyNumericalModelAdapter<BaseProblem>);
    STATIC_CHECK_FALSE(HasLegacyNumericalModelAdapter<CentralDensityProblem>);
    STATIC_CHECK(std::same_as<BaseProblem, equilibrium::StellarEquilibriumSystem<BaseModel>>);
    STATIC_CHECK(
        std::same_as<typename BaseProblem::PreparedOperatorType, operators::PreparedStellarEquilibriumOperator>
    );
    STATIC_CHECK(
        std::same_as<
            typename CentralDensityProblem::PreparedOperatorType,
            operators::PreparedCentralDensityStellarEquilibriumOperator>
    );
    STATIC_CHECK(
        std::same_as<
            typename BaseProblem::CompiledSurfaceConstraintType,
            surface::CompiledPressureSurfaceConstraintT<
                typename BaseProblem::ThermodynamicEquationsType::PressureSurfaceFormulation, eos::Polytrope>>
    );
    STATIC_CHECK(material::CompiledThermodynamicEquations<typename BaseProblem::ThermodynamicEquationsType>);
}

TEST_CASE(
    "Discretized Stellar Equilibrium Problem Is Exactly Equivalent To The Legacy Construction Path",
    tags::stellar_equilibrium_problem_integration
) {
    using namespace mean_field;

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

    models::StellarModel legacyModel{
        models::structure::PolytropicStructure{eos::Polytrope{3.0, 0.25}, 1.25},
        surface::ConstantPressureSurface{eos::PressureValue{0.0}}
    };
    operators::PreparedStellarEquilibriumOperator legacyOperator(f, *f.domainMapperStateless, legacyModel);

    const equilibrium::StellarDiscretization discretization{f, *f.domainMapperStateless};
    auto equilibriumProblem = equilibrium::discretize(
        model::StellarModel(
            integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.25}}),
            surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}), eos::Polytrope({.n = 3.0, .K = 0.25})
        ),
        discretization
    );
    auto &modelDrivenOperator = equilibriumProblem.GetPreparedOperator();

    CHECK(equilibriumProblem.StateSize() == legacyOperator.Width());
    CHECK(equilibriumProblem.EquationSize() == legacyOperator.Height());
    CHECK(equilibriumProblem.StateSize() == equilibriumProblem.EquationSize());
    CHECK(&equilibriumProblem.GetDiscretization().finiteElementModel() == &f);
    CHECK(&equilibriumProblem.GetDiscretization().domainMapper() == f.domainMapperStateless.get());
    CHECK(equilibriumProblem.GetDiscretization().isCurrent());
    CHECK(modelDrivenOperator.GetTargetMass() == 1.25);
    CHECK(modelDrivenOperator.GetSurfaceConstraintOperator().GetPhysicalCondition().targetPressure == 0.0);
    CHECK(equilibriumProblem.GetCompiledSurfaceConstraint().targetPressure() == dimensions::PressureValue{0.0});
    CHECK(modelDrivenOperator.GetDomainDeformation().matchesCurrentDiscretization());
    CHECK(&equilibriumProblem.GetLinearizationOperator() == &modelDrivenOperator);
    CHECK(equilibriumProblem.GetManifest().constraints()[0].target == 1.25);

    mfem::Vector state(legacyOperator.Width());
    state                                                        = 0.0;
    const auto stateView                                         = legacyOperator.GetRootStateView(state);
    stateView.block(utils::blocks::density_field.mass_term)      = 1.0;
    stateView.block(utils::blocks::enthalpy_field.specific_term) = 1.0;

    const operators::StellarEquilibriumDependencies dependencies = make_dependencies();
    const physics::RigidRotation rotation                        = make_zero_rotation();
    legacyOperator.Prepare(state, dependencies, rotation);
    equilibriumProblem.Prepare(state, dependencies, rotation);

    mfem::Vector legacyResidual;
    mfem::Vector modelDrivenResidual;
    legacyOperator.BuildResidual(legacyResidual);
    equilibriumProblem.BuildResidual(modelDrivenResidual);
    CHECK(relative_difference(modelDrivenResidual, legacyResidual) < 2.0e-15);

    mfem::Vector direction(state.Size());
    for (int index = 0; index < direction.Size(); ++index) {
        direction(index) = 0.01 * std::sin(0.31 * static_cast<double>(index + 1));
    }
    mfem::Vector legacyAction;
    mfem::Vector modelDrivenAction;
    legacyOperator.Mult(direction, legacyAction);
    equilibriumProblem.ApplyLinearization(direction, modelDrivenAction);
    CHECK(relative_difference(modelDrivenAction, legacyAction) < 2.0e-15);
}