#include "serif/utils/error/errors.hpp"
#include <cmath>
#include <limits>
#include <memory>
#include <utility>

#include "serif/equations/barotropic_density.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/backend/mfem/geometry/surface/pressure_boundary.hpp"
#include "serif/geometry/surface/evaluation.hpp"
#include "serif/tests/equations/volume_equation_fixture.hpp"
#include "serif/tests/equations/gravity_port.hpp"

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    using namespace serif::tests::equations::gravity;
    namespace physics = serif::geometry::surface;
    namespace surface = physics::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    using Enthalpy = ::serif::tests::equations::Enthalpy;
    using SurfaceState = dependencies::Bindings<dependencies::Binding<dimensions::MassDensity, Density>>;
    using Formulation = physics::PressureSurfaceFormulation<eos::models::Polytrope, SurfaceState>;
    using EnthalpyState = dependencies::Bindings<dependencies::Binding<dimensions::SpecificEnthalpy, Enthalpy>>;
    using EnthalpyFormulation = physics::PressureSurfaceFormulation<eos::models::Polytrope, EnthalpyState>;
    using Balance = serif::equations::BarotropicDensityBalance<eos::models::Polytrope>;
    using Vector = ::mfem::Vector;

    /** @brief Bind the carrier without introducing unrelated field requirements. */
    auto enthalpy_state(GravityFixture& f, const Vector& h) {
        return backend::make_term_coefficients(backend::bind(f.enthalpy_binding, h));
    }

    /** @brief Retain every density-balance column when testing strong boundary replacement. */
    auto balance_state(GravityFixture& f, const Vector& rho, const Vector& phi, const Vector& C, const Vector& omega) {
        return backend::make_term_coefficients(backend::bind(f.density_binding, rho), backend::bind(f.potential_binding, phi),
                              backend::bind(f.constant_binding, C), backend::bind(f.angular_velocity_binding, omega));
    }

    /** @brief Track both polytropic parameters alongside the prescribed pressure and current coefficients. */
    template <typename Boundary, typename Coefficients>
    auto pressure_inputs(const std::shared_ptr<const Boundary>& boundary, const eos::models::Polytrope& model, dimensions::PressureValue pressure, const Coefficients& coefficients) {
        const auto equation_of_state = dependencies::borrow<physics::keys::EquationOfState>(model, dependencies::track_by([](const eos::models::Polytrope& eos) {
            return std::pair{eos.polytropic_index(), eos.polytropic_constant()};
        }));
        return surface::make_pressure_boundary_inputs(boundary, equation_of_state, physics::Isobaric{pressure}, coefficients);
    }

    /** @brief Constrain the Cartesian fixture's stellar x=0 face, leaving its other faces free. */
    void tag_cartesian_boundary(GravityFixture& f) {
        for (int boundary = 0; boundary < f.mesh.GetNBE(); ++boundary) {
            auto& transform = *f.mesh.GetBdrElementTransformation(boundary);
            Vector center(3);
            transform.Transform(::mfem::Geometries.GetCenter(transform.GetGeometryType()), center);
            f.mesh.GetBdrElement(boundary)->SetAttribute(std::abs(center[0]) < 1.0e-14 ? 1 : 2);
        }
        f.mesh.SetAttributes();
    }

    /** @brief Select essential DOFs independently with MFEM, without the boundary trace compiler. */
    Vector boundary_mask(GravityFixture& f) {
        ::mfem::Array<int> attributes(f.mesh.bdr_attributes.Max()), marker;
        attributes = 0;
        attributes[0] = 1;
        f.density_space.GetEssentialVDofs(attributes, marker);
        Vector full(marker.Size());
        for (int dof = 0; dof < marker.Size(); ++dof) {
            full[dof] = marker[dof] ? 1.0 : 0.0;
        }
        return f.density_binding.gather(full);
    }

    /** @brief Preserve MeanField prepared_stellar_equilibrium.cpp:810 profiles and surface-row tolerances. */
    void check_rows(GravityFixture& f) {
        CAPTURE(f.mesh.GetNE(), f.mesh.GetElementGeometry(0), f.displacement_space.GetOrdering());
        f.set_compactification();
        const eos::models::Polytrope eos{3.0, 0.25};
        const dimensions::PressureValue pressure{0.03125};
        // For n=3 and K=1/4, rho(P) = (P/K)^(3/4), independently of the EOS implementation.
        const double required_rho = std::pow(pressure.value() / 0.25, 0.75);
        const auto rho = f.project(f.density_binding, [](const Vector& x) {
            return 0.82 + 0.08 * std::sin(0.62 * x[0] + 0.53) + 0.045 * std::cos(0.57 * x[1] - 0.53) + 0.02 * x[2] * x[2];
        });
        const auto drho = f.project(f.density_binding, [](const Vector& x) {
            return 0.21 * std::sin(0.74 * x[0] + 0.97) - 0.14 * std::cos(0.64 * x[1] - 0.97) + 0.075 * x[2];
        });
        const auto phi = f.project(f.potential_binding, [](const Vector& x) {
            return 0.24 + 0.09 * std::sin(0.67 * x[0] + 0.47) - 0.06 * std::cos(0.53 * x[1] - 0.47) + 0.035 * x[2];
        });
        const auto dphi = f.project(f.potential_binding, [](const Vector& x) {
            return 0.17 * std::sin(0.81 * x[0] + 0.89) + 0.11 * std::cos(0.69 * x[1] - 0.89) - 0.07 * x[2];
        });
        const auto C = constant(1, 1.07), dC = constant(1, -0.37);
        const auto omega = constant(1, 0.82), domega = constant(1, 0.13);
        // Volume displacement is already supported; surface-parameter extension is deliberately not used.
        const auto u = f.displacement(0.08), du = f.displacement(-0.11);
        const auto state = balance_state(f, rho, phi, C, omega);
        const auto direction = balance_state(f, drho, dphi, dC, domega);
        VolumeEquationFixture balance{Balance{eos, pressure}, f.density_binding, f.displacement_binding,
                                          f.compactification, f.mapper, reference_policy()};
        const auto boundary = surface::make_pressure_boundary<Formulation>(state);
        const auto inputs = pressure_inputs(boundary, eos, pressure, state);
        auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(surface::make_pressure_boundary_graph<Formulation>(), inputs);
        const auto mask = boundary_mask(f);
        REQUIRE(boundary->selected_rows->size() > 0);
        REQUIRE(boundary->selected_rows->size() < f.density_binding.size());
        CHECK(boundary->selected_rows->size() == static_cast<int>(mask.Sum()));
        int previous = -1;
        for (int dof : boundary->selected_rows->reduced_dofs()) {
            CHECK(dof > previous);
            CHECK(mask[dof] == 1.0);
            previous = dof;
        }

        balance.prepare(state, u, f.workspace);
        REQUIRE(cache.try_prepare(inputs));
        Vector residual, action;
        balance.residual(f.workspace, residual);
        balance.apply_jacobian(direction, du, f.workspace, action);
        const auto interior_residual = residual, interior_action = action;
        surface::apply_residual_rows(*boundary, cache.get<physics::keys::Linearization>(), residual);
        surface::apply_jacobian_rows(*boundary, cache.get<physics::keys::Linearization>(), direction, action);
        for (int row = 0; row < rho.Size(); ++row) {
            CAPTURE(row);
            if (mask[row]) {
                CHECK(std::abs(residual[row] - (rho[row] - required_rho)) < 1.0e-13);
                CHECK(std::abs(action[row] - drho[row]) < 1.0e-13);
            } else {
                CHECK(residual[row] == interior_residual[row]);
                CHECK(action[row] == interior_action[row]);
            }
        }

        const auto zero_rho = constant(rho.Size(), 0.0);
        const auto other_columns = balance_state(f, zero_rho, dphi, dC, domega);
        Vector other_action;
        balance.apply_jacobian(other_columns, du, f.workspace, other_action);
        const auto before_replacement = other_action;
        surface::apply_jacobian_rows(*boundary, cache.get<physics::keys::Linearization>(), other_columns, other_action);
        double removed_action = 0.0;
        for (int row = 0; row < rho.Size(); ++row) {
            if (mask[row]) {
                CHECK(other_action[row] == 0.0);
                removed_action += std::abs(before_replacement[row]);
            } else {
                CHECK(other_action[row] == before_replacement[row]);
            }
        }
        CHECK(removed_action > 0.0);

        // The constrained density-balance row retains the old complete-root difference step and threshold.
        constexpr double step = 1.0e-5;
        const auto shifted_residual = [&](double scale) {
            auto trial_rho = rho, trial_phi = phi, trial_C = C, trial_omega = omega, trial_u = u;
            trial_rho.Add(scale, drho); trial_phi.Add(scale, dphi); trial_C.Add(scale, dC);
            trial_omega.Add(scale, domega); trial_u.Add(scale, du);
            const auto trial = balance_state(f, trial_rho, trial_phi, trial_C, trial_omega);
            balance.prepare(trial, trial_u, f.workspace);
            REQUIRE(cache.try_prepare(pressure_inputs(boundary, eos, pressure, trial)));
            Vector result;
            balance.residual(f.workspace, result);
            surface::apply_residual_rows(*boundary, cache.get<physics::keys::Linearization>(), result);
            return result;
        };
        auto difference = shifted_residual(step);
        difference -= shifted_residual(-step);
        difference /= 2.0 * step;
        auto error = action;
        error -= difference;
        CHECK(error.Norml2() / std::max({action.Norml2(), difference.Norml2(), 100.0 * std::numeric_limits<double>::epsilon()}) < 2.0e-6);
    }

    /** @brief A second EOS prevents assumptions of a polytrope or h(0)=0. */
    struct OffsetEOS {
        using Relations = eos::relations::RelationCatalog<eos::relations::SpecificEnthalpyFromPressure>;
        double offset{0.3};
        /** @brief Give zero pressure a nonzero enthalpy that changes with the EOS parameters. */
        dimensions::SpecificEnthalpyValue evaluate(eos::relations::SpecificEnthalpyFromPressure, dimensions::PressureValue pressure) const {
            return dimensions::SpecificEnthalpyValue{offset + 2.0 * pressure.value()};
        }
    };
}

TEST_CASE("Surface pressure resolves the prescribed quantity through the EOS", "[equations][pressure_surface][unit]") {
    // MeanField constant_surface_compilation.cpp: pressure-only contract and exact polytropic residual/derivative.
    const eos::models::Polytrope eos{3.0, 0.25};
    const dimensions::PressureValue pressure{0.03125};
    const physics::Isobaric condition{pressure};
    CHECK(condition.target_pressure().value() == pressure.value());
    physics::SurfacePoint<Formulation> point{};
    point.fields.get<Density>() = eos::evaluate<dimensions::MassDensity>(eos, pressure).value();
    const auto result = physics::evaluate_pressure_surface<Formulation>(eos, condition, point);
    REQUIRE(result);
    CHECK(result->residual == 0.0);
    CHECK(result->field_derivatives.get<Density>() * -0.19 == -0.19);
    CHECK(result->position_derivative == discretization::forms::VectorValue{});
    for (double invalid : {-0.1, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
        CHECK_THROWS_AS((physics::Isobaric{dimensions::PressureValue{invalid}}), std::invalid_argument);
    }
    OffsetEOS offset_eos;
    using OffsetFormulation = physics::PressureSurfaceFormulation<OffsetEOS, EnthalpyState>;
    physics::SurfacePoint<OffsetFormulation> offset_point{};
    offset_point.fields.get<Enthalpy>() = 0.3;
    const physics::Isobaric zero_pressure{dimensions::PressureValue{0.0}};
    const auto original = physics::evaluate_pressure_surface<OffsetFormulation>(offset_eos, zero_pressure, offset_point);
    REQUIRE(original);
    CHECK(original->residual == 0.0);
    offset_eos.offset = 1.2;
    const auto updated = physics::evaluate_pressure_surface<OffsetFormulation>(offset_eos, zero_pressure, offset_point);
    REQUIRE(updated);
    CHECK(updated->residual == 0.3 - 1.2);
    CHECK(original->residual == 0.0);
    offset_eos.offset = std::numeric_limits<double>::infinity();
    const auto rejected = physics::evaluate_pressure_surface<OffsetFormulation>(offset_eos, condition, offset_point);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().reason == physics::SurfaceFailureReason::non_finite_eos_response);
}

TEST_CASE("Pressure constraints replace density-balance boundary rows at the original accuracy", "[equations][pressure_surface][integration]") {
    for (const auto geometry : {::mfem::Element::TETRAHEDRON, ::mfem::Element::HEXAHEDRON}) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            GravityFixture f(geometry, ordering);
            tag_cartesian_boundary(f);
            check_rows(f);
        }
    }
    auto fixture = gravity_port::stellar_fixture();
    check_rows(*fixture);
}

TEST_CASE("Surface pressure preserves exact analytic closure density-balance and mass rows", "[equations][pressure_surface][integration]") {
    // MeanField prepared_stellar_equilibrium.cpp:979. Retain all physical constants and numerical thresholds.
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    const eos::models::Polytrope eos{1.0, 0.25};
    const dimensions::SpecificEnthalpyValue enthalpy{0.60};
    const double density = eos::evaluate<dimensions::MassDensity>(eos, enthalpy).value();
    const auto pressure = eos::evaluate<dimensions::Pressure>(eos, enthalpy);
    double volume = 0.0;
    for (int element = 0; element < f.mesh.GetNE(); ++element) {
        if (!f.enthalpy_binding.supports(element)) {
            continue;
        }
        auto& transform = *f.mesh.GetElementTransformation(element);
        const auto& rule = ::mfem::IntRules.Get(f.mesh.GetElementGeometry(element), reference_order);
        for (int q = 0; q < rule.GetNPoints(); ++q) {
            const auto& ip = rule.IntPoint(q);
            transform.SetIntPoint(&ip);
            volume += ip.weight * transform.Weight();
        }
    }
    const double target_mass = density * volume;
    const auto rho = constant(f.density_binding.size(), density), h = constant(f.enthalpy_binding.size(), enthalpy.value());
    // A uniform state at the surface density has A(rho) = A(rho_surface), so the balance closes at Us = phi.
    const auto phi = constant(f.potential_binding.size(), 0.20), C = constant(1, 0.20);
    const auto omega = constant(1, 0.0), u = constant(f.displacement_binding.size(), 0.0);
    const auto state = backend::make_term_coefficients(backend::bind(f.density_binding, rho), backend::bind(f.enthalpy_binding, h),
                                      backend::bind(f.potential_binding, phi), backend::bind(f.constant_binding, C),
                                      backend::bind(f.angular_velocity_binding, omega));
    VolumeEquationFixture closure{Closure{eos}, f.density_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    VolumeEquationFixture balance{Balance{eos, pressure}, f.density_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    VolumeEquationFixture mass{serif::equations::MassNormalization{target_mass}, f.constant_binding,
                                   f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    const auto boundary = surface::make_pressure_boundary<Formulation>(state);
    const auto inputs = pressure_inputs(boundary, eos, pressure, state);
    auto cache = dependencies::make_computation_cache<physics::SurfaceFailure>(surface::make_pressure_boundary_graph<Formulation>(), inputs);
    closure.prepare(state, u, f.workspace);
    balance.prepare(state, u, f.workspace);
    mass.prepare(state, u, f.workspace);
    REQUIRE(cache.try_prepare(inputs));
    Vector result;
    closure.residual(f.workspace, result);
    CHECK(result.Norml2() < 2.0e-12);
    balance.residual(f.workspace, result);
    surface::apply_residual_rows(*boundary, cache.get<physics::keys::Linearization>(), result);
    CHECK(result.Norml2() < 2.0e-12);
    mass.residual(f.workspace, result);
    REQUIRE(result.Size() == 1);
    CHECK(std::abs(result[0]) < 2.0e-11 * target_mass);
}

TEST_CASE("Boundary selection requires supported tagged nodal coefficients", "[equations][pressure_surface][unit]") {
    GravityFixture f;
    tag_cartesian_boundary(f);
    const Vector h = constant(f.enthalpy_binding.size(), 0.7);
    const auto state = enthalpy_state(f, h);
    SECTION("Absent boundary") {
        for (int boundary = 0; boundary < f.mesh.GetNBE(); ++boundary) {
            f.mesh.GetBdrElement(boundary)->SetAttribute(2);
        }
        CHECK_THROWS_AS(surface::make_pressure_boundary<EnthalpyFormulation>(state), std::invalid_argument);
    }
    SECTION("Boundary outside the stellar support") {
        for (int boundary = 0; boundary < f.mesh.GetNBE(); ++boundary) {
            f.mesh.GetBdrElement(boundary)->SetAttribute(1);
        }
        CHECK_THROWS_AS(surface::make_pressure_boundary<EnthalpyFormulation>(state), std::invalid_argument);
    }
    SECTION("Positive H1 coefficients are not nodal values") {
        ::mfem::H1_FECollection positive{Enthalpy::FunctionalSpace::order, 3, ::mfem::BasisType::Positive};
        ::mfem::FiniteElementSpace space{&f.mesh, &positive};
        discretization::backend::mfem::TermBinding<Enthalpy, Schema> binding{space};
        const Vector values = constant(binding.size(), 0.7);
        const auto positive_state = backend::make_term_coefficients(backend::bind(binding, values));
        CHECK_THROWS_AS(surface::make_pressure_boundary<EnthalpyFormulation>(positive_state), std::invalid_argument);
    }
    SECTION("Indices must be sorted unique and in range") {
        using discretization::backend::mfem::DofSelection;
        for (const ::mfem::Array<int>& indices : {::mfem::Array<int>{-1}, ::mfem::Array<int>{0, 0},
                                               ::mfem::Array<int>{1, 0}, ::mfem::Array<int>{f.enthalpy_binding.size()}}) {
            CHECK_THROWS_AS((DofSelection{f.enthalpy_binding, indices}), std::invalid_argument);
        }
        CHECK_THROWS_AS(f.enthalpy_binding.reduced_dof(-1), std::out_of_range);
        CHECK_THROWS_AS(f.enthalpy_binding.reduced_dof(f.enthalpy_binding.full_size()), std::out_of_range);
    }
}