#include <limits>
#include <memory>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>

#include "serif/tests/equations/test_equations.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/equations/angular_momentum.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/solver/normalization/normalization.hpp"
#include "serif/utils/dependencies/dependencies.hpp"

namespace
{
    namespace normalization = serif::solver::normalization;
    namespace dependencies = serif::utils::dependencies;
    namespace dimensions = serif::dimensions;
    namespace blocks = serif::discretization::blocks;
    using Kind = blocks::RepresentedValueKind;
    using Density = blocks::density::mass;
    using Constant = blocks::barotropic_constant::mass_normalization;
    using Rotation = blocks::rotation::angular_velocity;
    using Mass = serif::equations::MassNormalization;
    using Momentum = serif::equations::AngularMomentumConstraint;

    /** @brief Make a test reference with deliberately different mass, enthalpy, pressure and rotation scales. */
    normalization::StellarReferenceScales reference_scales()
    {
        return normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{9.0}, dimensions::QuantityValue<dimensions::Length>{2.0}, dimensions::QuantityValue<dimensions::GravitationalConstant>{4.0}).value();
    }

    /** @brief Construct a cache from the ordinary computation tuple, as a coupled graph would. */
    template <typename Plan, typename Scales = normalization::StellarReferenceScales, typename Context, typename Layout>
    auto make_cache(const Context& context, const Layout& layout)
    {
        return std::apply([&](auto... computations)
        {
            return dependencies::make_computation_cache(context, std::move(computations)...);
        }, normalization::make_normalization_computations<Plan, Scales>(layout));
    }

    struct SurfaceRows
    {
    };

    struct TemperatureCoordinate
    {
        using Quantity = dimensions::Temperature;
    };

    struct TemperatureMetric
    {
        using Coordinates = dependencies::KeyList<TemperatureCoordinate>;
    };

    struct ThermalScales
    {
        dimensions::QuantityValue<dimensions::Temperature> temperature;

        /** @brief Track an external scale provider without adding normalization-specific revision machinery. */
        bool operator==(const ThermalScales&) const = default;
    };

    /** @brief Supply a quantity that cannot be inferred from stellar mass, radius, and gravity. */
    dimensions::QuantityValue<dimensions::Temperature> characteristic_scale(const ThermalScales& scales, dimensions::Temperature)
    {
        return scales.temperature;
    }

    /** @brief A third-party method whose configuration is an explicit tracked input, including move-only ownership. */
    struct UniformDiagonal : normalization::NormalizationPrescriptionTag
    {
        std::unique_ptr<double> multiplier;
        double residual_multiplier;
        std::size_t extra_count = 0;

        /** @brief Return the actual policy's factors for the exact requested role and coordinate. */
        template <Kind Role, typename Coordinate>
        std::expected<std::vector<double>, std::string> prepare_factors(normalization::keys::Factors<Role, Coordinate>, std::size_t count) const
        {
            return std::vector<double>(count + extra_count, Role == Kind::value ? *multiplier : residual_multiplier);
        }
    };
}

TEST_CASE("Normalization plans preserve distinct global quantities and obey layout rather than declaration order", "[unit][normalization][dependencies]")
{
    using Plan = normalization::NormalizationPlan<
        dependencies::Bindings<normalization::PhysicalValueRule<Rotation>, normalization::PhysicalValueRule<Constant>>,
        dependencies::Bindings<normalization::PhysicalResidualRule<Mass>, normalization::PhysicalResidualRule<Momentum>>>;
    const normalization::CoordinateLayout<dependencies::KeyList<Constant, Rotation>, dependencies::KeyList<Momentum, Mass>> layout{.state_sizes = {1, 1}, .residual_sizes = {1, 1}};
    auto scales = reference_scales();
    double irrelevant_candidate = 0.0;
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout), dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<Density>(irrelevant_candidate));
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    const normalization::DiagonalNormalization frozen = cache.get<normalization::keys::Normalization>();
    CHECK(frozen.state_factors()[0] == Catch::Approx(1.0 / scales.specific_energy.value()).epsilon(2.0e-15));
    CHECK(frozen.state_factors()[1] == Catch::Approx(1.0 / scales.angular_velocity.value()).epsilon(2.0e-15));
    CHECK(frozen.residual_factors()[0] == Catch::Approx(1.0 / scales.angular_momentum.value()).epsilon(2.0e-15));
    CHECK(frozen.residual_factors()[1] == Catch::Approx(1.0 / scales.mass.value()).epsilon(2.0e-15));

    irrelevant_candidate = std::numeric_limits<double>::quiet_NaN();
    auto report = cache.try_prepare(context);
    REQUIRE(report);
    CHECK(report->computed_count() == 0);
    scales.mass = dimensions::QuantityValue<dimensions::Mass>{18.0};
    report = cache.try_prepare(context);
    REQUIRE(report);
    CHECK(cache.get<normalization::keys::Normalization>().residual_factors()[1] == Catch::Approx(1.0 / 18.0));
    CHECK(frozen.residual_factors()[1] == Catch::Approx(1.0 / 9.0));
}

TEST_CASE("One reference volume metric supplies independent primal and dual coordinates", "[unit][normalization][dependencies]")
{
    using Metric = normalization::ReferenceVolumeMetric<Density>;
    using Method = normalization::PhysicalRieszCoordinate<dimensions::MassDensity, Metric>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<dependencies::Binding<Density, Method>>, dependencies::Bindings<dependencies::Binding<Density, Method>>>;
    const normalization::CoordinateLayout<dependencies::KeyList<Density>, dependencies::KeyList<Density>> layout{.state_sizes = {3}, .residual_sizes = {3}};
    const auto scales = reference_scales();
    std::vector<double> metric{1.0e-20, 4.0, 9.0e20};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout), dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<Metric>(metric));
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    const auto& map = cache.get<normalization::keys::Normalization>();
    for (std::size_t index = 0; index < metric.size(); ++index)
    {
        CHECK(map.state_factors()[index] == Catch::Approx(std::sqrt(metric[index]) / scales.density.value()).epsilon(3.0e-15));
        CHECK(map.residual_factors()[index] == Catch::Approx(1.0 / (std::sqrt(metric[index]) * scales.density.value())).epsilon(3.0e-15));
    }
    metric[1] = 0.0;
    const auto failure = cache.try_prepare(context);
    REQUIRE_FALSE(failure);
    const auto originating_failure = normalization::make_riesz_factors(Kind::value, metric, scales.density.value());
    REQUIRE_FALSE(originating_failure);
    CHECK(failure.error().error == originating_failure.error());
    CHECK_THAT(failure.error().error, Catch::Matchers::ContainsSubstring("NORMALIZATION_ERROR"));
    const std::string& diagnostic = failure.error().error;
    REQUIRE(diagnostic.find("code No.:") != std::string::npos);
    CHECK(diagnostic.find("code No.:") == diagnostic.rfind("code No.:"));
    CHECK_FALSE(cache.ready());
    CHECK_THROWS_AS(cache.get<normalization::keys::Normalization>(), std::logic_error);
    metric[1] = 16.0;
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[1] == Catch::Approx(4.0 / scales.density.value()));
    metric.push_back(1.0);
    const auto wrong_size = cache.try_prepare(context);
    REQUIRE_FALSE(wrong_size);
    CHECK_THAT(wrong_size.error().error, Catch::Matchers::ContainsSubstring("NORMALIZATION_VECTOR_SIZE_MISMATCH"));
    metric.pop_back();
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[1] == Catch::Approx(4.0 / scales.density.value()));
}

TEST_CASE("Pressure point rows replace both the volume metric and its physical scale", "[unit][normalization][dependencies]")
{
    using Enthalpy = ::serif::tests::equations::Enthalpy;
    using Metric = normalization::ReferenceVolumeMetric<Enthalpy>;
    using Volume = normalization::PhysicalRieszCoordinate<dimensions::SpecificEnthalpy, Metric>;
    using Hybrid = normalization::WithPointRows<Volume, dimensions::Pressure, SurfaceRows>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<>, dependencies::Bindings<dependencies::Binding<Enthalpy, Hybrid>>>;
    const normalization::CoordinateLayout<dependencies::KeyList<>, dependencies::KeyList<Enthalpy>> layout{.state_sizes = {}, .residual_sizes = {4}};
    const auto scales = reference_scales();
    const std::vector<double> metric{4.0, 0.0, 16.0, 0.0};
    std::vector<int> rows{1, 3};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout), dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<Metric>(metric), dependencies::borrow<SurfaceRows>(rows));
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    const auto& map = cache.get<normalization::keys::Normalization>();
    CHECK(map.state_size() == 0);
    CHECK(map.residual_factors()[0] == Catch::Approx(1.0 / (2.0 * scales.specific_energy.value())));
    CHECK(map.residual_factors()[1] == Catch::Approx(1.0 / scales.pressure.value()));
    CHECK(map.residual_factors()[2] == Catch::Approx(1.0 / (4.0 * scales.specific_energy.value())));
    CHECK(map.residual_factors()[3] == Catch::Approx(1.0 / scales.pressure.value()));
    for (const std::vector<int>& invalid : {std::vector<int>{1, 1}, {1, 4}, {-1, 3}, {1}})
    {
        rows = invalid;
        CHECK_FALSE(cache.try_prepare(context));
    }
    rows = {1, 3};
    REQUIRE(cache.try_prepare(context));
}

TEST_CASE("Custom normalization uses its exact move-only prescription and tracks configuration changes", "[unit][normalization][dependencies]")
{
    using Method = normalization::RuntimePreparedCoordinate<UniformDiagonal>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<dependencies::Binding<Density, Method>>, dependencies::Bindings<dependencies::Binding<Density, Method>>>;
    normalization::CoordinateLayout<dependencies::KeyList<Density>, dependencies::KeyList<Density>> layout{.state_sizes = {2}, .residual_sizes = {3}};
    UniformDiagonal policy{.multiplier = std::make_unique<double>(2.0), .residual_multiplier = 5.0};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout), dependencies::borrow<normalization::keys::Prescription<UniformDiagonal>>(policy, dependencies::track_by([](const UniformDiagonal& method)
    {
        return std::tuple{*method.multiplier, method.residual_multiplier, method.extra_count};
    })));
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_size() == 2);
    CHECK(cache.get<normalization::keys::Normalization>().residual_size() == 3);
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[1] == 2.0);
    CHECK(cache.get<normalization::keys::Normalization>().residual_factors()[2] == 5.0);
    *policy.multiplier = 3.0;
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[1] == 3.0);
    policy.residual_multiplier = -1.0;
    CHECK_FALSE(cache.try_prepare(context));
    policy.residual_multiplier = 5.0;
    policy.extra_count = 1;
    CHECK_FALSE(cache.try_prepare(context));
    policy.extra_count = 0;
    layout.state_sizes[0] = std::numeric_limits<std::size_t>::max();
    CHECK_FALSE(cache.try_prepare(context));
}

TEST_CASE("External physical scales share ordinary metric preparation and invalidation", "[unit][normalization][dependencies]")
{
    using Method = normalization::PhysicalRieszCoordinate<dimensions::Temperature, TemperatureMetric>;
    using Rule = dependencies::Binding<TemperatureCoordinate, Method>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<Rule>, dependencies::Bindings<Rule>>;
    const normalization::CoordinateLayout<dependencies::KeyList<TemperatureCoordinate>, dependencies::KeyList<TemperatureCoordinate>> layout{.state_sizes = {2}, .residual_sizes = {2}};
    ThermalScales scales{.temperature = dimensions::QuantityValue<dimensions::Temperature>{6000.0}};
    const std::vector<double> metric{4.0, 9.0};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout), dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<TemperatureMetric>(metric));
    auto cache = make_cache<Plan, ThermalScales>(context, layout);
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[0] == Catch::Approx(2.0 / 6000.0));
    CHECK(cache.get<normalization::keys::Normalization>().residual_factors()[1] == Catch::Approx(1.0 / 18000.0));
    REQUIRE(cache.try_prepare(context)->computed_count() == 0);
    scales.temperature = dimensions::QuantityValue<dimensions::Temperature>{3000.0};
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[0] == Catch::Approx(2.0 / 3000.0));
    scales.temperature = dimensions::QuantityValue<dimensions::Temperature>{-1.0};
    CHECK_FALSE(cache.try_prepare(context));
}

TEST_CASE("Oversized normalization layouts fail before allocating identity factors", "[unit][normalization][dependencies]")
{
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<dependencies::Binding<Density, normalization::IdentityCoordinate>>, dependencies::Bindings<>>;
    const normalization::CoordinateLayout<dependencies::KeyList<Density>, dependencies::KeyList<>> layout{.state_sizes = {std::numeric_limits<std::size_t>::max()}, .residual_sizes = {}};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout));
    auto cache = make_cache<Plan>(context, layout);
    CHECK_FALSE(cache.try_prepare(context));
}

TEST_CASE("Identity preparation needs no physical sources and validates global coordinate arity", "[unit][normalization][dependencies]")
{
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<dependencies::Binding<Constant, normalization::IdentityCoordinate>>, dependencies::Bindings<>>;
    normalization::CoordinateLayout<dependencies::KeyList<Constant>, dependencies::KeyList<>> layout{.state_sizes = {1}, .residual_sizes = {}};
    const auto context = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout));
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_factors()[0] == 1.0);
    layout.state_sizes[0] = 2;
    CHECK_FALSE(cache.try_prepare(context));
}

TEST_CASE("Empty normalization plans preserve empty independently packed coordinate spaces", "[unit][normalization][dependencies]")
{
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<>, dependencies::Bindings<>>;
    const normalization::CoordinateLayout<dependencies::KeyList<>, dependencies::KeyList<>> layout{.state_sizes = {}, .residual_sizes = {}};
    const auto context = dependencies::make_dependency_context();
    auto cache = make_cache<Plan>(context, layout);
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.get<normalization::keys::Normalization>().state_size() == 0);
    CHECK(cache.get<normalization::keys::Normalization>().residual_size() == 0);
}