#include <array>
#include <cmath>
#include <expected>
#include <limits>
#include <memory>
#include <numbers>
#include <stdexcept>
#include <string>

#include <catch2/catch_test_macros.hpp>

#include "mfem/fem/coefficient.hpp"
#include "serif/backend/mfem/mapping/compactification/kelvin.hpp"
#include "serif/dimensions/type_alias.hpp"
#include "serif/backend/mfem/mapping/evaluator.hpp"
#include "serif/utils/laneemden/laneemden.hpp"
#include "serif/utils/projection/projection.hpp"
#include "serif/tests/deformation/mesh.hpp"
#include "serif/backend/mfem/mapping/transformations.hpp"

namespace {
    namespace projection = serif::utils::projection;
    namespace backend = projection::backend::mfem;
    namespace laneemden = serif::utils::laneemden;
    namespace dimensions = serif::dimensions;
    namespace mapping = serif::mapping;

    /** @brief Resolve radius without sharing the production sampler's distance implementation. */
    [[nodiscard]] double radius_squared(const mfem::Vector& point, const std::array<double, 3>& center) {
        double squared = 0.0;
        for (int component = 0; component < 3; ++component) {
            squared += std::pow(point[component] - center[component], 2);
        }
        return squared;
    }

    /** @brief Check physical reconstruction through independent quadrature. */
    void check_linear_vector(mfem::FiniteElementSpace& space, const std::array<double, 3>& center) {
        const auto result = backend::project_vector(space, [](const double radius) { return 2.0 * radius; }, {.center = center});
        REQUIRE(result);
        mfem::VectorFunctionCoefficient exact(3, [center](const mfem::Vector& point, mfem::Vector& value) {
            value.SetSize(3);
            for (int component = 0; component < 3; ++component) {
                value[component] = 2.0 * (point[component] - center[component]);
            }
        });
        CHECK(result->ComputeL2Error(exact) < 2.0e-12);
    }
}

TEST_CASE("Radial scalar projection reproduces physical polynomials in H1 and L2", "[unit][projection]") {
    const std::array<double, 3> center{0.5, 0.25, 0.75};
    for (const mfem::Element::Type geometry : {mfem::Element::HEXAHEDRON, mfem::Element::TETRAHEDRON}) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, geometry);
        mfem::H1_FECollection h1(2, 3);
        mfem::L2_FECollection l2(2, 3);
        mfem::L2_FECollection integral_l2(2, 3, mfem::BasisType::GaussLegendre, mfem::FiniteElement::INTEGRAL);
        for (mfem::FiniteElementCollection* collection : std::array<mfem::FiniteElementCollection*, 3>{&h1, &l2, &integral_l2}) {
            mfem::FiniteElementSpace space(&mesh, collection);
            const auto result = backend::project_scalar(space, [](const double radius) {
                return 3.0 - radius * radius;
            }, {.center = center});
            REQUIRE(result);
            CHECK(result->FESpace() == &space);
            mfem::FunctionCoefficient exact([center](const mfem::Vector& point) {
                return 3.0 - radius_squared(point, center);
            });
            CHECK(result->ComputeL2Error(exact) < 2.0e-12);
        }
    }
}

TEST_CASE("Radial vectors preserve component order", "[unit][projection]") {
    for (const mfem::Element::Type geometry : {mfem::Element::HEXAHEDRON, mfem::Element::TETRAHEDRON}) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, geometry);
        mfem::H1_FECollection h1(1, 3);
        mfem::L2_FECollection l2(1, 3);
        for (const mfem::Ordering::Type order : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
            mfem::FiniteElementSpace continuous(&mesh, &h1, 3, order);
            mfem::FiniteElementSpace discontinuous(&mesh, &l2, 3, order);
            check_linear_vector(continuous, {0.5, 0.5, 0.5});
            check_linear_vector(continuous, {0.0, 0.0, 0.0});
            check_linear_vector(discontinuous, {0.5, 0.5, 0.5});
        }
    }
}

TEST_CASE("Radial projection samples curved physical mesh coordinates", "[unit][projection]") {
    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, mfem::Element::HEXAHEDRON);
    mesh.SetCurvature(2);
    mesh.Transform([](const mfem::Vector& point, mfem::Vector& image) {
        image = point;
        image[1] += 0.1 * point[0] * point[0];
        image[2] *= 1.0 + 0.1 * point[0];
    });
    mfem::H1_FECollection h1(4, 3);
    mfem::FiniteElementSpace space(&mesh, &h1);
    const auto result = backend::project_scalar(space, [](const double radius) { return radius * radius; });
    REQUIRE(result);
    mfem::FunctionCoefficient exact([](const mfem::Vector& point) { return point * point; });
    CHECK(result->ComputeL2Error(exact) < 3.0e-12);
}

TEST_CASE("Radial projection pulls physical values back through stellar mapping", "[unit][projection]") {
    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, mfem::Element::HEXAHEDRON);
    mfem::H1_FECollection h1(2, 3);
    mfem::FiniteElementSpace displacement_space(&mesh, &h1, 3);
    mfem::FiniteElementSpace scalar_space(&mesh, &h1);
    mfem::GridFunction displacement(&displacement_space);
    mfem::GridFunction compactification(&scalar_space);
    const std::array<double, 3> stretch{1.2, 0.8, 1.1};
    mfem::VectorFunctionCoefficient affine(3, [stretch](const mfem::Vector& point, mfem::Vector& value) {
        value.SetSize(3);
        for (int component = 0; component < 3; ++component) {
            value[component] = (stretch[component] - 1.0) * point[component];
        }
    });
    displacement.ProjectCoefficient(affine);
    compactification = 0.0;
    mapping::DomainMapper mapper(
        {.dimension = 3, .vacuum_element_attribute_ID = 3},
        std::make_unique<mapping::compactification::KelvinCompactification>(
            serif::utils::user_options::KelvinCompactificationOptions{.r_star_ref = 1.0, .r_inf_ref = 4.0}
        )
    );
    mapping::GridFunctionMappingEvaluator evaluator(mapper, displacement, compactification);
    CHECK(evaluator.mesh() == &mesh);
    // Equal element counts do not make geometry from another mesh interchangeable.
    mfem::Mesh other_mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, mfem::Element::HEXAHEDRON);
    mfem::FiniteElementSpace other_scalar(&other_mesh, &h1), other_vector(&other_mesh, &h1, 3);
    CHECK_THROWS_AS(backend::project_scalar(other_scalar, [](const double r) { return r; }, {.mapping = &evaluator}), std::invalid_argument);
    CHECK_THROWS_AS(backend::project_vector(other_vector, [](const double r) { return r; }, {.mapping = &evaluator}), std::invalid_argument);
    const auto scalar = backend::project_scalar(scalar_space, [](const double radius) { return radius * radius; }, {.mapping = &evaluator});
    REQUIRE(scalar);
    mfem::FunctionCoefficient exact_scalar([stretch](const mfem::Vector& point) {
        double value = 0.0;
        for (int component = 0; component < 3; ++component) {
            value += std::pow(stretch[component] * point[component], 2);
        }
        return value;
    });
    CHECK(scalar->ComputeL2Error(exact_scalar) < 3.0e-12);

    const auto vector = backend::project_vector(displacement_space, [](const double radius) { return 2.0 * radius; }, {.mapping = &evaluator});
    REQUIRE(vector);
    // Vector H1 values are physical: 2 x_phys = 2 A x_mesh.
    mfem::VectorFunctionCoefficient exact_vector(3, [stretch](const mfem::Vector& point, mfem::Vector& value) {
        value.SetSize(3);
        for (int component = 0; component < 3; ++component) {
            value[component] = 2.0 * stretch[component] * point[component];
        }
    });
    CHECK(vector->ComputeL2Error(exact_vector) < 3.0e-12);
    const double determinant = stretch[0] * stretch[1] * stretch[2];

    mfem::L2_FECollection integral_l2(2, 3, mfem::BasisType::GaussLegendre, mfem::FiniteElement::INTEGRAL);
    mfem::FiniteElementSpace density_space(&mesh, &integral_l2);
    const auto density = backend::project_scalar(density_space, [](const double) { return 2.0; }, {.mapping = &evaluator});
    REQUIRE(density);
    mfem::ConstantCoefficient exact_density(2.0 * determinant);
    CHECK(density->ComputeL2Error(exact_density) < 3.0e-12);

    displacement *= -6.0;
    evaluator.refresh();
    CHECK_FALSE(backend::project_scalar(scalar_space, [](const double radius) { return radius; }, {.mapping = &evaluator}));

    displacement = 0.0;
    compactification = 1.0;
    for (int element = 0; element < mesh.GetNE(); ++element) {
        mesh.GetElement(element)->SetAttribute(3);
    }
    mesh.SetAttributes();
    evaluator.refresh();
    const auto infinite = backend::project_scalar(scalar_space, [](const double) { return 0.0; }, {.mapping = &evaluator});
    REQUIRE_FALSE(infinite);
    CHECK(infinite.error().find("compactified infinity") != std::string::npos);
}

TEST_CASE("Lane Emden physical profiles project density enthalpy potential and gravity", "[unit][projection][laneemden]") {
    const auto solution = laneemden::solve(0.0);
    REQUIRE(solution);
    const auto scales = laneemden::scale_from_mass(*solution,
        dimensions::QuantityValue<dimensions::Mass>{3.0},
        dimensions::QuantityValue<dimensions::MassDensity>{2.0},
        dimensions::QuantityValue<dimensions::GravitationalConstant>{1.5});
    REQUIRE(scales);
    const double radius = scales->radius.value();
    const double central_enthalpy = scales->central_enthalpy.value();
    const double potential_at_surface = -1.5 * 3.0 / radius;
    const double extent = 0.5 * radius;
    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, mfem::Element::HEXAHEDRON, extent, extent, extent);
    mfem::H1_FECollection h1(2, 3);
    mfem::L2_FECollection l2(2, 3);
    mfem::FiniteElementSpace enthalpy_space(&mesh, &h1), density_space(&mesh, &l2), gravity_space(&mesh, &h1, 3);
    const auto field = [&](const double physical_radius) {
        return laneemden::evaluate_physical(*solution, *scales, dimensions::QuantityValue<dimensions::Length>{physical_radius});
    };
    const auto density = backend::project_scalar(density_space, [&](const double r) -> projection::RadialValue {
        return field(r).transform([](const auto& value) { return value.density.value(); });
    });
    const auto enthalpy = backend::project_scalar(enthalpy_space, [&](const double r) -> projection::RadialValue {
        return field(r).transform([](const auto& value) { return value.enthalpy.value(); });
    });
    const auto potential = backend::project_scalar(density_space, [&](const double r) -> projection::RadialValue {
        return field(r).transform([](const auto& value) { return value.potential.value(); });
    });
    const auto gravity = backend::project_vector(gravity_space, [&](const double r) -> projection::RadialValue {
        return field(r).transform([](const auto& value) { return value.gravity.value(); });
    });
    REQUIRE(density);
    REQUIRE(enthalpy);
    REQUIRE(potential);
    REQUIRE(gravity);
    mfem::ConstantCoefficient exact_density(2.0);
    mfem::FunctionCoefficient exact_enthalpy([=](const mfem::Vector& point) {
        return central_enthalpy * (1.0 - (point * point) / (radius * radius));
    });
    mfem::FunctionCoefficient exact_potential([=](const mfem::Vector& point) {
        return potential_at_surface - central_enthalpy * (1.0 - (point * point) / (radius * radius));
    });
    mfem::VectorFunctionCoefficient exact_gravity(3, [](const mfem::Vector& point, mfem::Vector& value) {
        value = point;
        value *= 4.0 * std::numbers::pi * 1.5 * 2.0 / 3.0;
    });
    CHECK(density->ComputeL2Error(exact_density) < 1.0e-11);
    CHECK(enthalpy->ComputeL2Error(exact_enthalpy) < 1.0e-9);
    CHECK(potential->ComputeL2Error(exact_potential) < 1.0e-9);
    CHECK(gravity->ComputeL2Error(exact_gravity) < 1.0e-9);
}

TEST_CASE("Projection retains profile exterior choices and rejects invalid samples atomically", "[unit][projection]") {
    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 2, 2, mfem::Element::HEXAHEDRON);
    mfem::H1_FECollection h1(1, 3);
    mfem::FiniteElementSpace scalar(&mesh, &h1), vector(&mesh, &h1, 3);
    const auto exterior = backend::project_scalar(scalar, [](const double radius) {
        return radius > 0.25 ? -7.0 : 3.0;
    });
    REQUIRE(exterior);
    CHECK(exterior->Min() == -7.0);
    CHECK(exterior->Max() == 3.0);
    for (const double invalid : {std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()}) {
        CHECK_FALSE(backend::project_scalar(scalar, [invalid](const double) { return invalid; }));
        CHECK_FALSE(backend::project_vector(vector, [invalid](const double) { return invalid; }));
        CHECK_FALSE(backend::project_scalar(scalar, [](const double r) { return r; }, {.center = {invalid, 0.0, 0.0}}));
    }
    CHECK_FALSE(backend::project_vector(vector, [](const double) { return 1.0; }));
    int successful_samples = 0;
    const auto failure = backend::project_scalar(scalar, [&](const double r) -> projection::RadialValue {
        if (r > 0.5) {
            return std::unexpected("profile is only available inside radius 0.5");
        }
        ++successful_samples;
        return 2.0;
    });
    REQUIRE_FALSE(failure);
    CHECK(successful_samples > 0);
    CHECK(failure.error().find("profile is only available") != std::string::npos);
    const auto recovered = backend::project_scalar(scalar, [](const double) { return 4.0; });
    REQUIRE(recovered);
    CHECK(recovered->Min() == 4.0);
    CHECK(recovered->Max() == 4.0);
    CHECK_THROWS_AS(backend::project_scalar(vector, [](const double r) { return r; }), std::invalid_argument);
    CHECK_THROWS_AS(backend::project_vector(scalar, [](const double r) { return r; }), std::invalid_argument);
}

TEST_CASE("Projected n one polytropes converge to their analytic radial fields", "[unit][projection][laneemden]") {
    const auto solution = laneemden::solve(1.0, {.step = 0.002});
    REQUIRE(solution);
    const auto scales = laneemden::scale_from_mass(*solution,
        dimensions::QuantityValue<dimensions::Mass>{2.0},
        dimensions::QuantityValue<dimensions::MassDensity>{1.0},
        dimensions::QuantityValue<dimensions::GravitationalConstant>{1.0});
    REQUIRE(scales);
    const double radial_scale = scales->radial_scale;
    double coarse_error = 0.0;
    for (const int count : {2, 4}) {
        const double extent = scales->radius.value() / std::sqrt(3.0);
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(count, count, count, mfem::Element::HEXAHEDRON, extent, extent, extent);
        mfem::H1_FECollection h1(2, 3);
        mfem::FiniteElementSpace space(&mesh, &h1);
        const auto density = backend::project_scalar(space, [&](const double radius) -> projection::RadialValue {
            return laneemden::evaluate_physical(*solution, *scales, dimensions::QuantityValue<dimensions::Length>{radius})
                .transform([](const auto& value) { return value.density.value(); });
        });
        REQUIRE(density);
        mfem::FunctionCoefficient exact([radial_scale](const mfem::Vector& point) {
            const double xi = point.Norml2() / radial_scale;
            if (xi == 0.0) {
                return 1.0;
            }
            return std::sin(xi) / xi;
        });
        const double error = density->ComputeL2Error(exact);
        if (count == 2) {
            coarse_error = error;
            CHECK(coarse_error > 1.0e-7);
        } else {
            CHECK(error < coarse_error / 6.0);
            CHECK(error < 2.0e-4 * std::pow(extent, 1.5));
        }
    }
}

TEST_CASE("Radial projection completes the compactified stellar mesh with explicit finite limits", "[unit][projection][laneemden][infinity]") {
    auto generated = serif::tests::deformation::make_mesh(0);
    mfem::H1_FECollection h1(2, 3);
    mfem::FiniteElementSpace scalar_space(generated.mesh.get(), &h1);
    mfem::FiniteElementSpace vector_space(generated.mesh.get(), &h1, 3);
    mfem::GridFunction displacement(&vector_space);
    displacement = 0.0;
    mapping::DomainMapper mapper(
        {.dimension = 3, .vacuum_element_attribute_ID = 3},
        std::make_unique<mapping::compactification::KelvinCompactification>(
            serif::utils::user_options::KelvinCompactificationOptions{.r_star_ref = 1.0, .r_inf_ref = 4.0}
        )
    );
    mapping::GridFunctionMappingEvaluator evaluator(mapper, displacement, *generated.exterior_coordinate->values);
    const auto solution = laneemden::solve(1.0, {.step = 0.002});
    REQUIRE(solution);
    // R = 1/2 leaves the entire curved exterior mesh in the analytic vacuum region.
    const auto scales = laneemden::scale_from_mass(*solution,
        dimensions::MassValue{2.0}, dimensions::DensityValue{4.0 * std::numbers::pi},
        dimensions::QuantityValue<dimensions::GravitationalConstant>{0.7});
    REQUIRE(scales);
    const double gravitational_mass = 0.7 * scales->mass.value();
    const auto fields = [&](double radius) {
        return laneemden::evaluate_physical(*solution, *scales, dimensions::QuantityValue<dimensions::Length>{radius});
    };
    int scalar_limits = 0;
    int vector_limits = 0;
    const auto scalar_limit = [&](const mfem::Vector&) -> projection::RadialValue {
        ++scalar_limits;
        return 0.0;
    };
    const auto vector_limit = [&](const mfem::Vector& point) -> std::expected<mfem::Vector, std::string> {
        ++vector_limits;
        mfem::Vector value(point);
        value *= gravitational_mass / std::pow(point.Norml2(), 3);
        return value;
    };
    const auto potential = backend::project_scalar(scalar_space,
        [&](double radius) -> projection::RadialValue {
            return fields(radius).transform([](const laneemden::PhysicalSample& sample) { return sample.potential.value(); });
        }, {.mapping = &evaluator, .scalar_at_infinity = scalar_limit});
    const auto gravity = backend::project_vector(vector_space,
        [&](double radius) -> projection::RadialValue {
            return fields(radius).transform([](const laneemden::PhysicalSample& sample) { return sample.gravity.value(); });
        }, {.mapping = &evaluator, .vector_at_infinity = vector_limit});
    REQUIRE(potential);
    REQUIRE(gravity);
    CHECK(scalar_limits > 0);
    CHECK(vector_limits > 0);
    mfem::Array<int> infinity_marker(generated.mesh->bdr_attributes.Max());
    infinity_marker = 0;
    infinity_marker[serif::tests::deformation::PairedMesh::infinity_attribute - 1] = 1;
    mfem::Array<int> infinity_dofs;
    scalar_space.GetEssentialTrueDofs(infinity_marker, infinity_dofs);
    REQUIRE(infinity_dofs.Size() > 0);
    for (const int dof : infinity_dofs) {
        CHECK((*potential)[dof] == 0.0);
    }

    CHECK(gravity->Norml2() > 0.0);
}

TEST_CASE("Compactified projection validates explicit limits before publishing", "[unit][projection][infinity]") {
    mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::HEXAHEDRON);
    mesh.GetElement(0)->SetAttribute(3);
    mesh.SetAttributes();
    mfem::H1_FECollection h1(1, 3);
    mfem::FiniteElementSpace scalar_space(&mesh, &h1), vector_space(&mesh, &h1, 3);
    mfem::GridFunction displacement(&vector_space), compactification(&scalar_space);
    displacement = 0.0;
    compactification = 1.0;
    mapping::DomainMapper mapper(
        {.dimension = 3, .vacuum_element_attribute_ID = 3},
        std::make_unique<mapping::compactification::KelvinCompactification>(
            serif::utils::user_options::KelvinCompactificationOptions{.r_star_ref = 1.0, .r_inf_ref = 4.0}
        )
    );
    mapping::GridFunctionMappingEvaluator evaluator(mapper, displacement, compactification);
    int ordinary_samples = 0;
    const auto profile = [&](double) {
        ++ordinary_samples;
        return 1.0;
    };
    CHECK_FALSE(backend::project_scalar(scalar_space, profile, {.mapping = &evaluator}));
    CHECK_FALSE(backend::project_vector(vector_space, profile, {.mapping = &evaluator}));
    for (double invalid : {std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()}) {
        CHECK_FALSE(backend::project_scalar(scalar_space, profile, {.mapping = &evaluator,
            .scalar_at_infinity = [invalid](const mfem::Vector&) -> projection::RadialValue { return invalid; }}));
        CHECK_FALSE(backend::project_vector(vector_space, profile, {.mapping = &evaluator,
            .vector_at_infinity = [invalid](const mfem::Vector&) -> std::expected<mfem::Vector, std::string> {
                mfem::Vector value(3);
                value = invalid;
                return value;
            }}));
    }
    const auto rejected_scalar = backend::project_scalar(scalar_space, profile, {.mapping = &evaluator,
        .scalar_at_infinity = [](const mfem::Vector&) -> projection::RadialValue {
            return std::unexpected("scalar asymptote unavailable");
        }});
    const auto rejected_vector = backend::project_vector(vector_space, profile, {.mapping = &evaluator,
        .vector_at_infinity = [](const mfem::Vector&) -> std::expected<mfem::Vector, std::string> {
            return std::unexpected("vector asymptote unavailable");
        }});
    REQUIRE_FALSE(rejected_scalar);
    REQUIRE_FALSE(rejected_vector);
    CHECK(rejected_scalar.error().find("scalar asymptote unavailable") != std::string::npos);
    CHECK(rejected_vector.error().find("vector asymptote unavailable") != std::string::npos);
    CHECK_FALSE(backend::project_vector(vector_space, profile, {.mapping = &evaluator,
        .vector_at_infinity = [](const mfem::Vector&) -> std::expected<mfem::Vector, std::string> {
            mfem::Vector value(2);
            value = 0.0;
            return value;
        }}));
    const auto recovered = backend::project_scalar(scalar_space, profile, {.mapping = &evaluator,
        .scalar_at_infinity = [](const mfem::Vector& point) -> projection::RadialValue { return point[0]; }});
    REQUIRE(recovered);
    CHECK(recovered->Min() == 0.0);
    CHECK(recovered->Max() == 1.0);
    CHECK(ordinary_samples == 0);
}