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

#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 = 3109, .revision = 1},
            .density            = {.identity = 3119, .revision = 1},
            .surfaceDeformation = {.identity = 3121, .revision = 1},
            .gravityGradient    = {.identity = 3137, .revision = 1},
            .gravityPotential   = {.identity = 3163, .revision = 1},
            .enthalpy           = {.identity = 3167, .revision = 1},
            .bernoulliConstant  = {.identity = 3169, .revision = 1},
            .rotation           = {.identity = 3181, .revision = 1},
            .targetMass         = {.identity = 3187, .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(
    "Central Density Contribution Composes Through The Variadic Root",
    tags::central_density_phase_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 stellarModel{
        models::structure::PolytropicStructure{eos::Polytrope{3.0, 0.25}, 1.0},
        surface::ConstantPressureSurface{dimensions::PressureValue{0.0}}
    };
    operators::PreparedStellarEquilibriumOperator physicalOperator(f, *f.domainMapperStateless, stellarModel);
    auto equilibriumProblem = equilibrium::discretize(
        model::StellarModel(
            constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{1.0}}),
            integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
            surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}), eos::Polytrope({.n = 3.0, .K = 0.25})
        ),
        equilibrium::StellarDiscretization{f, *f.domainMapperStateless}
    );
    auto &borderedOperator = equilibriumProblem.GetPreparedOperator();

    STATIC_CHECK(
        std::same_as<
            typename std::remove_cvref_t<decltype(equilibriumProblem)>::PreparedOperatorType,
            operators::PreparedVariadicStellarEquilibriumOperator<
                typename std::remove_cvref_t<decltype(equilibriumProblem)>::ModelType>>
    );
    CHECK(
        equilibriumProblem.GetStellarModel().specification<constraint::FixedCentralDensity>().targetDensity() ==
        dimensions::DensityValue{1.0}
    );
    CHECK(equilibriumProblem.StateSize() == equilibriumProblem.EquationSize());

    CHECK(borderedOperator.Width() == physicalOperator.Width() + 1);
    CHECK(borderedOperator.Height() == physicalOperator.Height() + 1);
    CHECK(borderedOperator.GetRootManifest().valueBlocks().size() == 7);
    CHECK(borderedOperator.GetRootManifest().residualBlocks().size() == 7);
    CHECK(borderedOperator.GetRootManifest().constraints().size() == 3);
    CHECK(borderedOperator.GetRootManifest().specificationDescriptors().size() == 4);

    const auto &centralDescriptor =
        borderedOperator.GetRootManifest().specification<constraint::FixedCentralDensity>();
    CHECK(centralDescriptor.stableId == "FixedCentralDensity");
    CHECK(centralDescriptor.role == models::SpecificationRole::phase_condition);
    CHECK(centralDescriptor.columnPolicy == operators::RootColumnPolicy::solver_border);
    CHECK(centralDescriptor.target == 1.0);
    REQUIRE(centralDescriptor.carrierTarget.has_value());
    CHECK(*centralDescriptor.carrierTarget == 1.0);
    CHECK(centralDescriptor.targetUnits == "density");
    CHECK(centralDescriptor.residualUnits == "specific_enthalpy");

    mfem::Vector physicalState(physicalOperator.Width());
    physicalState                                                   = 0.0;
    const auto physicalStateView =
        physicalOperator.GetRootManifest().stateView(physicalState);
    physicalStateView.block(utils::blocks::density_field.mass_term) = 1.0;
    physicalStateView.block(utils::blocks::enthalpy_field.specific_term) = 1.0;

    mfem::Vector borderedState(borderedOperator.Width());
    borderedState                                                = 0.0;
    mfem::Vector(borderedState.GetData(), physicalState.Size())  = physicalState;

    const operators::StellarEquilibriumDependencies dependencies = make_dependencies();
    const physics::RigidRotation rotation                        = make_zero_rotation();
    physicalOperator.Prepare(physicalState, dependencies, rotation);
    const auto initialReport = equilibriumProblem.Prepare(borderedState, dependencies, rotation);
    CHECK(initialReport.physical.assembledResidual);
    CHECK(initialReport.specification<constraint::FixedCentralDensity>().constraint.assembledResidual);
    CHECK(initialReport.assembledResidual);

    mfem::Vector physicalResidual;
    mfem::Vector borderedResidual;
    physicalOperator.BuildResidual(physicalResidual);
    equilibriumProblem.BuildResidual(borderedResidual);
    const mfem::Vector borderedPhysicalResidual(borderedResidual.GetData(), physicalResidual.Size());
    CHECK(relative_difference(borderedPhysicalResidual, physicalResidual) < 2.0e-15);
    CHECK(borderedResidual(borderedResidual.Size() - 1) == 0.0);

    const operators::CentralDensityConstraintReport centralReport = borderedOperator.GetCentralDensityReport();
    CHECK(centralReport.targetDensity == 1.0);
    CHECK(centralReport.achievedDensity == 1.0);
    CHECK(centralReport.enthalpyResidual == 0.0);

    mfem::Vector physicalDirection(physicalOperator.Width());
    for (int index = 0; index < physicalDirection.Size(); ++index) {
        physicalDirection(index) = 0.01 * std::sin(0.37 * static_cast<double>(index + 1));
    }
    mfem::Vector borderedDirection(borderedOperator.Width());
    borderedDirection                                                   = 0.0;
    mfem::Vector(borderedDirection.GetData(), physicalDirection.Size()) = physicalDirection;

    mfem::Vector physicalAction;
    mfem::Vector borderedAction;
    physicalOperator.Mult(physicalDirection, physicalAction);
    equilibriumProblem.ApplyLinearization(borderedDirection, borderedAction);
    const mfem::Vector borderedPhysicalAction(borderedAction.GetData(), physicalAction.Size());
    CHECK(relative_difference(borderedPhysicalAction, physicalAction) < 2.0e-15);

    const auto borderedDirectionView     = borderedOperator.GetRootManifest().directionView(borderedDirection);
    const mfem::Vector enthalpyDirection = borderedDirectionView.block(utils::blocks::enthalpy_field.specific_term);
    double localCenterDirection          = 0.0;
    for (const int centerDof : borderedOperator.GetCentralDensityConstraint().GetCenterDof().reduced_dofs()) {
        localCenterDirection += enthalpyDirection(centerDof);
    }
    double globalCenterDirection = 0.0;
    MPI_Allreduce(&localCenterDirection, &globalCenterDirection, 1, MPI_DOUBLE, MPI_SUM, f.mesh->GetComm());
    CHECK(borderedAction(borderedAction.Size() - 1) == globalCenterDirection);

    const auto repeatedReport = borderedOperator.Prepare(borderedState, dependencies, rotation);
    CHECK_FALSE(repeatedReport.physical.DidAnyWork());
    CHECK_FALSE(repeatedReport.specification<constraint::FixedCentralDensity>().DidAnyWork());
    CHECK(repeatedReport.assembledResidual);

    borderedState(borderedState.Size() - 1) = 0.375;
    const auto borderReport                 = borderedOperator.Prepare(borderedState, dependencies, rotation);
    CHECK_FALSE(borderReport.physical.DidAnyWork());
    CHECK(borderReport.specification<constraint::FixedCentralDensity>().constraint.refreshedBorder);
    CHECK(borderReport.assembledResidual);

    mfem::Vector borderOnlyDirection(borderedOperator.Width());
    borderOnlyDirection                                 = 0.0;
    borderOnlyDirection(borderOnlyDirection.Size() - 1) = -0.625;
    const std::uint64_t preparationsBeforeMult = borderedOperator.GetCentralDensityConstraint().GetPreparationCount();
    borderedOperator.Mult(borderOnlyDirection, borderedAction);
    CHECK(borderedOperator.GetCentralDensityConstraint().GetPreparationCount() == preparationsBeforeMult);
    CHECK(borderedAction(borderedAction.Size() - 1) == 0.0);

    const auto actionView             = borderedOperator.GetRootManifest().residualView(borderedAction);
    const mfem::Vector enthalpyAction = actionView.block(utils::blocks::enthalpy_field.specific_term);
    double localBorderEntry           = 0.0;
    for (const int centerDof : borderedOperator.GetCentralDensityConstraint().GetCenterDof().reduced_dofs()) {
        localBorderEntry += enthalpyAction(centerDof);
    }
    double globalBorderEntry = 0.0;
    MPI_Allreduce(&localBorderEntry, &globalBorderEntry, 1, MPI_DOUBLE, MPI_SUM, f.mesh->GetComm());
    CHECK(globalBorderEntry == -0.625);
}