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

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "serif/solver/normalization/scales.hpp"
#include "serif/solver/normalization/physical_coordinates.hpp"
#include "serif/equations/barycenter.hpp"
#include "serif/utils/serialize/equations.hpp"

namespace {
    namespace equations = serif::equations;
    namespace dimensions = serif::dimensions;
    using Density = serif::discretization::blocks::density::mass;

    /** @brief Check the density partial and every explicit position partial by central differences. */
    template <typename Equation, typename Point>
    void check_partials(const Equation& equation, const Point& point) {
        const auto linearization = equation.evaluate(point);
        for (const double step : {1.0e-5, 1.0e-6}) {
            auto plus = point, minus = point;
            plus.fields.template get<Density>() += step;
            minus.fields.template get<Density>() -= step;
            const double difference = (equation.evaluate(plus).residual - equation.evaluate(minus).residual) / (2.0 * step);
            const double expected = linearization.field_derivatives.template get<Density>();
            CHECK(std::abs(difference - expected) < 2.0e-9 * std::max(1.0, std::abs(expected)));
        }
        for (int axis = 0; axis < 3; ++axis) {
            auto plus = point, minus = point;
            constexpr double step = 1.0e-5;
            plus.position[axis] += step;
            minus.position[axis] -= step;
            const double difference = (equation.evaluate(plus).residual - equation.evaluate(minus).residual) / (2.0 * step);
            CHECK(std::abs(difference - linearization.position_derivative[axis]) < 2.0e-9);
        }
    }
}

TEST_CASE("Barycenter constraints differentiate all axes about a displaced center and persist", "[equations][barycenter][unit]") {
    const auto check_axis = []<int Axis>() {
        using Center = equations::BarycenterConstraint<Axis>;
        const std::array<double, 3> origin{0.1, -0.2, 0.3};
        const std::array<double, 3> position{0.6, 0.3, -0.4};
        constexpr double density = 1.2;
        constexpr double mass = 2.7;
        equations::EquationPoint<typename Center::Operands> center;
        center.position = position;
        center.fields.template get<Density>() = density;
        const Center center_law{mass, origin};
        CHECK(center_law.evaluate(center).residual == Catch::Approx(density * (position[Axis] - origin[Axis]) / mass));
        check_partials(center_law, center);
        namespace serialization = serif::utils::serialize;
        const auto saved = serialization::encode(center_law);
        const auto restored = serialization::decode<Center>(saved);
        REQUIRE(restored);
        CHECK(serialization::encode(*restored) == saved);
        for (const double invalid : {std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
            auto bad_center = origin;
            bad_center[Axis] = invalid;
            CHECK_THROWS(Center(mass, bad_center));
        }
        for (const double invalid : {0.0, -0.1, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
            CHECK_THROWS(Center{invalid});
        }
    };
    check_axis.template operator()<0>();
    check_axis.template operator()<1>();
    check_axis.template operator()<2>();
}

TEST_CASE("Frame acceleration coordinates use the stellar acceleration scale", "[equations][barycenter][unit]") {
    namespace normalization = serif::solver::normalization;
    const auto scales = normalization::prepare_stellar_reference_scales(dimensions::MassValue{12.0}, dimensions::LengthValue{3.0},
        dimensions::QuantityValue<dimensions::GravitationalConstant>{2.0});
    REQUIRE(scales);
    CHECK(normalization::characteristic_scale(*scales, dimensions::Acceleration{}).value() == Catch::Approx(24.0 / 9.0));
}