#include <algorithm>
#include <cmath>
#include <span>
#include <vector>

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

#include "serif/backend/mfem/model/equilibrium.hpp"
#include "serif/model/finite_elements.hpp"
#include "serif/runtime.hpp"
#include "serif/solver/seed/state.hpp"
#include "../sandbox/polytropic_model.hpp"

namespace {
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace blocks = serif::discretization::blocks;
    namespace deformation = serif::geometry::deformation;
    namespace normalization = serif::solver::normalization;
    using Density = blocks::density::mass;
    using Potential = blocks::gravity::potential;
    using Constant = blocks::barotropic_constant::mass_normalization;
    using Omega = blocks::rotation::angular_velocity;
    using Shape = blocks::surface_deformation::coefficients;
    using Kind = blocks::RepresentedValueKind;

    /** @brief Make quadrature refinement explicit without changing the represented stellar state. */
    [[nodiscard]] serif::operators::backend::mfem::VolumeOptions integration_order(int order) {
        serif::discretization::quadrature::RuleSet rules;
        rules.set_fallback_control({.fixed_order = order});
        return {.quadrature_policy = serif::discretization::quadrature::Policy{rules}};
    }

    /** @brief Freeze a residual before preparing another physical candidate. */
    template <typename System>
    [[nodiscard]] mfem::Vector residual_copy(const System& system) {
        mfem::Vector result(static_cast<int>(system.residual().size()));
        std::copy(system.residual().begin(), system.residual().end(), result.GetData());
        return result;
    }

    /** @brief Exercise all coefficients with a reproducible bounded direction, independent of mesh numbering. */
    [[nodiscard]] mfem::Vector direction_values(int size, double phase) {
        mfem::Vector result(size);
        for (int index = 0; index < size; ++index) { result[index] = std::sin(0.37 * index + phase); }
        return result;
    }

    /** @brief Compare each physical equation separately so gravity cannot mask an incorrect scalar or surface row. */
    template <typename Layout>
    void check_rows(const mfem::Vector& actual, const mfem::Vector& expected, const Layout& layout,
        double relative_tolerance, double absolute_tolerance) {
        REQUIRE(actual.Size() == expected.Size());
        int offset = 0;
        for (std::size_t row = 0; row < layout.residual_sizes.size(); ++row) {
            CAPTURE(row);
            double error_squared = 0.0, reference_squared = 0.0;
            for (std::size_t entry = 0; entry < layout.residual_sizes[row]; ++entry, ++offset) {
                const double error = actual[offset] - expected[offset];
                error_squared += error * error;
                reference_squared += expected[offset] * expected[offset];
            }
            CHECK(std::sqrt(error_squared) <= absolute_tolerance + relative_tolerance * std::sqrt(reference_squared));
        }
    }

    /** @brief Perturb solver coordinates through their actual dimensional scaling. */
    template <typename Layout>
    [[nodiscard]] model::StellarState<Layout> perturbed(const model::StellarState<Layout>& state,
        const normalization::DiagonalNormalization& coordinates, const mfem::Vector& direction, double step) {
        mfem::Vector physical(direction.Size());
        coordinates.denormalize_state({direction.GetData(), static_cast<std::size_t>(direction.Size())},
            {physical.GetData(), static_cast<std::size_t>(physical.Size())});
        model::StellarState<Layout> result(state);
        for (int index = 0; index < direction.Size(); ++index) { result.values()[index] += step * physical[index]; }
        return result;
    }

    /** @brief Verify strong physical boundary equations after the complete normalized transpose assembly. */
    template <typename Field, typename Layout, typename Selection>
    void check_boundary_transpose(const backend::NormalizedSystem& system, const Layout& layout, const Selection& selected) {
        const auto& coordinates = system.coordinates();
        const int row_offset = static_cast<int>(layout.template offset<Kind::residual, Field>());
        const int column_offset = static_cast<int>(layout.template offset<Kind::value, Field>());
        mfem::Vector dual(system.jacobian().Height()), expected(system.jacobian().Width()), actual;
        dual = 0.0;
        expected = 0.0;
        REQUIRE(selected.size() > 0);
        for (int row : selected.reduced_dofs()) {
            const double value = 0.2 + 0.01 * (row % 7);
            dual[row_offset + row] = value;
            expected[column_offset + row] = value * coordinates.residual_factors()[row_offset + row] /
                coordinates.state_factors()[column_offset + row];
        }
        system.jacobian().MultTranspose(dual, actual);
        actual -= expected;
        CHECK(actual.Norml2() <= 2e-12 * expected.Norml2());
    }
}

TEST_CASE("Weighted moving stars retain complete normalized derivatives and physical boundary rows", "[model][integration][reconstruction][reconstruction_system]") {
    serif::initialize();
    const auto star = serif::examples::make_polytropic_model(4.0 / 3.0, 0.6, 1.0, 0.04);
    const auto radial = serif::solver::seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto finite_elements = model::make_stellar_discretization(star,
        {.radius = radial->scales.radius, .geometry_order = 3},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{.surface = {.maximum_degree = 2}},
        serif::discretization::reconstruction::SurfacePower<Density>{.density_exponent = 4.0 / 3.0});
    auto seeded = serif::solver::seed::make_polytropic_state(star, finite_elements, *radial);
    REQUIRE(seeded);
    auto state = *seeded;
    state.get<Omega>()[0] = 0.08;
    REQUIRE(state.get<Shape>().size() == 9);
    state.get<Shape>()[0] = 0.01;
    state.get<Shape>()[6] = -0.015;
    const auto scales = normalization::prepare_stellar_reference_scales(radial->scales.mass, radial->scales.radius);
    REQUIRE(scales);
    const auto configuration = model::make_discretization(model::DiscretizationOptions{
        .model = star, .finite_elements = finite_elements,
        .normalization = normalization::NormalizationOptions{.scheme = *scales}});
    auto evaluator = backend::bind_stellar_model(configuration, state, integration_order(12));
    const auto evaluated = evaluator.evaluate(state);
    REQUIRE(evaluated);
    const auto frozen = *evaluated;
    const auto& layout = state.layout();
    const auto& coordinates = frozen->coordinates();
    const mfem::Vector original_residual = residual_copy(*frozen);
    const mfem::Vector seed = direction_values(frozen->jacobian().Width(), 0.43);
    const mfem::Vector dual = direction_values(frozen->jacobian().Height(), 1.17);
    mfem::Vector transpose;
    frozen->jacobian().MultTranspose(dual, transpose);

    int offset = 0;
    for (std::size_t column = 0; column <= layout.state_sizes.size(); ++column) {
        CAPTURE(column);
        mfem::Vector direction(seed);
        if (column < layout.state_sizes.size()) {
            direction = 0.0;
            std::copy_n(seed.GetData() + offset, layout.state_sizes[column], direction.GetData() + offset);
            offset += static_cast<int>(layout.state_sizes[column]);
        }
        mfem::Vector action;
        frozen->jacobian().Mult(direction, action);
        REQUIRE(action.Norml2() > 1e-10);
        CHECK(action * dual == Catch::Approx(direction * transpose).epsilon(2e-11).margin(2e-11));
        const std::vector<double> steps = column < layout.state_sizes.size() ? std::vector<double>{2e-6} : std::vector<double>{2e-6, 1e-6};
        for (const double step : steps) {
            CAPTURE(step);
            const auto plus = evaluator.evaluate(perturbed(state, coordinates, direction, step));
            REQUIRE(plus);
            mfem::Vector difference = residual_copy(**plus);
            const auto minus = evaluator.evaluate(perturbed(state, coordinates, direction, -step));
            REQUIRE(minus);
            difference -= residual_copy(**minus);
            difference /= 2.0 * step;
            check_rows(action, difference, layout, 5e-6, 2e-8);
        }
    }
    const auto density_surface = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::StellarSurfaceBoundary>(finite_elements.binding<Density>());
    const auto infinity = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::InfinitySurfaceBoundary>(finite_elements.binding<Potential>());
    check_boundary_transpose<Density>(*frozen, layout, density_surface);
    check_boundary_transpose<Potential>(*frozen, layout, infinity);
    const auto restored = evaluator.evaluate(state);
    REQUIRE(restored);
    check_rows(residual_copy(**restored), original_residual, layout, 0.0, 0.0);
    mfem::Vector repeated;
    frozen->jacobian().MultTranspose(dual, repeated);
    repeated -= transpose;
    CHECK(repeated.Norml2() == 0.0);
}

TEST_CASE("Weighted moving-star mass and angular momentum converge under independent quadrature refinement", "[model][integration][reconstruction][reconstruction_quadrature]") {
    serif::initialize();
    const auto star = serif::examples::make_polytropic_model(4.0 / 3.0, 0.6, 1.0, 0.04);
    const auto radial = serif::solver::seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto finite_elements = model::make_stellar_discretization(star,
        {.radius = radial->scales.radius, .geometry_order = 3},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{.surface = {.maximum_degree = 2}},
        serif::discretization::reconstruction::SurfacePower<Density>{.density_exponent = 4.0 / 3.0});
    auto state = serif::solver::seed::make_polytropic_state(star, finite_elements, *radial);
    REQUIRE(state);
    state->get<Omega>()[0] = 0.08;
    state->get<Shape>()[6] = -0.015;
    std::vector<mfem::Vector> residuals;
    for (int order : {12, 20, 28}) {
        auto evaluator = backend::bind_stellar_model(star, finite_elements, *state, integration_order(order));
        const auto evaluated = evaluator.evaluate(*state);
        REQUIRE(evaluated);
        residuals.push_back(residual_copy(**evaluated));
    }
    const auto check_functional = [&]<typename Row>() {
        const int offset = static_cast<int>(state->layout().template offset<Kind::residual, Row>());
        const double coarse_error = std::abs(residuals[0][offset] - residuals[2][offset]);
        const double fine_error = std::abs(residuals[1][offset] - residuals[2][offset]);
        CAPTURE(offset, coarse_error, fine_error);
        CHECK(fine_error <= 0.7 * coarse_error + 1e-11);
        CHECK(fine_error < 1e-4);
    };
    check_functional.template operator()<Constant>();
    check_functional.template operator()<Omega>();
}