#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <limits>

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

import mean_field;
import test_helpers;

namespace {
    [[nodiscard]] mean_field::field::FieldPointDofMap make_center_map() {
        mfem::Array<int> centerDof(1);
        centerDof[0] = 2;
        return {5, centerDof};
    }

    [[nodiscard]] double relative_error(
        const double actual,
        const double expected
    ) {
        return std::abs(actual - expected) / std::max({1.0, std::abs(actual), std::abs(expected)});
    }
} // namespace

TEST_CASE(
    "Fixed Central Density Compiles A Carrier Phase Row And Solver Border",
    tags::model_specification_type_contract
) {
    using namespace mean_field;
    using Request = models::CentralDensityLayoutRequest;
    using Form    = utils::blocks::central_density_bordered_stellar_equilibrium_form;

    STATIC_CHECK(models::ConstraintLayoutRequestType<Request>);
    STATIC_CHECK(models::CompiledConstraint<models::CompiledFixedCentralDensity>);
    STATIC_CHECK(Request::rowInjection == models::ConstraintRowInjection::solver_border);
    STATIC_CHECK(Request::valueArity == 1);
    STATIC_CHECK(Request::residualArity == 1);
    STATIC_CHECK(Request::valueBlock<Form>().index == Form::value_block_count - 1);
    STATIC_CHECK(Request::residualBlock<Form>().index == Form::residual_block_count - 1);
    STATIC_CHECK(std::same_as<typename models::CompiledFixedCentralDensity::CarrierField, field::Enthalpy>);
    STATIC_CHECK(std::same_as<typename models::CompiledFixedCentralDensity::BorderField, field::CentralDensityBorder>);

    const eos::Polytrope equationOfState{3.0, 0.25};
    const models::CompiledFixedCentralDensity compiled =
        models::compileConstraint(models::FixedCentralDensity{eos::DensityValue{8.0}}, equationOfState);

    CHECK(compiled.targetDensity() == eos::DensityValue{8.0});
    CHECK(compiled.targetEnthalpy() == eos::SpecificEnthalpyValue{2.0});
    CHECK(compiled.densityFromEnthalpy(eos::SpecificEnthalpyValue{2.5}) == eos::DensityValue{15.625});
}

TEST_CASE(
    "Prepared Central Density Phase Has Exact Residual Jacobian And Transpose Actions",
    tags::central_density_phase_unit
) {
    using namespace mean_field;

    const eos::Polytrope equationOfState{3.0, 0.25};
    const models::CompiledFixedCentralDensity compiled =
        models::compileConstraint(models::FixedCentralDensity{eos::DensityValue{8.0}}, equationOfState);
    operators::PreparedCentralDensityConstraint phase(make_center_map(), MPI_COMM_SELF);

    mfem::Vector enthalpy(5);
    enthalpy    = 0.0;
    enthalpy(2) = 2.5;
    const operators::CentralDensityDependencies dependencies{.enthalpy = {.identity = 17, .revision = 1}};

    const operators::PreparedCentralDensityReport initial = phase.Prepare(compiled, enthalpy, 0.3, dependencies);
    CHECK(initial.refreshedCentralEnthalpy);
    CHECK(initial.refreshedBorder);
    CHECK(initial.assembledResidual);

    mfem::Vector carrierResidual(5);
    mfem::Vector phaseResidual(1);
    carrierResidual = 1.0;
    phaseResidual   = 0.0;
    phase.AddResidual(carrierResidual, phaseResidual);
    CHECK(carrierResidual(2) == 1.3);
    CHECK(phaseResidual(0) == 0.5);

    mfem::Vector enthalpyVariation(5);
    enthalpyVariation                = 0.0;
    enthalpyVariation(2)             = -0.4;
    constexpr double borderVariation = 0.7;

    mfem::Vector carrierAction(5);
    mfem::Vector phaseAction(1);
    carrierAction = 0.0;
    phaseAction   = 0.0;
    phase.ApplyJacobian(
        {.enthalpyVariation = enthalpyVariation, .borderVariation = borderVariation},
        {.enthalpyAction = carrierAction, .phaseAction = phaseAction}
    );
    CHECK(carrierAction(2) == borderVariation);
    CHECK(phaseAction(0) == enthalpyVariation(2));

    // The phase residual is affine, so a larger centered-difference step
    // reduces cancellation without introducing truncation error.
    constexpr double epsilon = 1.0e-3;
    mfem::Vector plusEnthalpy(enthalpy);
    mfem::Vector minusEnthalpy(enthalpy);
    plusEnthalpy.Add(epsilon, enthalpyVariation);
    minusEnthalpy.Add(-epsilon, enthalpyVariation);

    auto plusDependencies = dependencies;
    ++plusDependencies.enthalpy.revision;
    phase.Prepare(compiled, plusEnthalpy, 0.3 + epsilon * borderVariation, plusDependencies);
    mfem::Vector plusCarrier(5);
    mfem::Vector plusPhase(1);
    plusCarrier = 0.0;
    plusPhase   = 0.0;
    phase.AddResidual(plusCarrier, plusPhase);

    auto minusDependencies = plusDependencies;
    ++minusDependencies.enthalpy.revision;
    phase.Prepare(compiled, minusEnthalpy, 0.3 - epsilon * borderVariation, minusDependencies);
    mfem::Vector minusCarrier(5);
    mfem::Vector minusPhase(1);
    minusCarrier = 0.0;
    minusPhase   = 0.0;
    phase.AddResidual(minusCarrier, minusPhase);

    plusCarrier -= minusCarrier;
    plusCarrier /= 2.0 * epsilon;
    const double phaseDifference = (plusPhase(0) - minusPhase(0)) / (2.0 * epsilon);
    plusCarrier -= carrierAction;
    CHECK(plusCarrier.Norml2() < 1.0e-10);
    const double phaseDifferenceError = relative_error(phaseDifference, phaseAction(0));
    INFO("Central-density phase action = " << phaseAction(0));
    INFO("Central-density centered difference = " << phaseDifference);
    INFO("Central-density centered-difference error = " << phaseDifferenceError);
    CHECK(phaseDifferenceError < 1.0e-10);

    auto restoredDependencies = minusDependencies;
    ++restoredDependencies.enthalpy.revision;
    phase.Prepare(compiled, enthalpy, 0.3, restoredDependencies);
    mfem::Vector carrierDual(5);
    carrierDual                = 0.0;
    carrierDual(2)             = -0.8;
    constexpr double phaseDual = 1.1;
    mfem::Vector enthalpyDual(5);
    mfem::Vector borderDual(1);
    enthalpyDual = 0.0;
    borderDual   = 0.0;
    phase.ApplyJacobianTranspose(
        {.enthalpyResidualDual = carrierDual, .phaseResidualDual = phaseDual},
        {.enthalpyDual = enthalpyDual, .borderDual = borderDual}
    );

    const double forwardPairing   = carrierAction * carrierDual + phaseAction(0) * phaseDual;
    const double transposePairing = enthalpyVariation * enthalpyDual + borderVariation * borderDual(0);
    CHECK(relative_error(transposePairing, forwardPairing) < 8.0 * std::numeric_limits<double>::epsilon());

    const operators::CentralDensityConstraintReport report = phase.GetConstraintReport();
    CHECK(report.targetDensity == 8.0);
    CHECK(report.achievedDensity == 15.625);
    CHECK(report.targetEnthalpy == 2.0);
    CHECK(report.achievedEnthalpy == 2.5);
    CHECK(report.enthalpyResidual == 0.5);
    CHECK(report.scaledResidual == 0.25);
}