#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <numbers>
#include <span>
#include <vector>

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

#include "serif/backend/mfem/geometry/deformation/computations.hpp"
#include "serif/model/finite_elements.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace {
    namespace model = serif::model;
    namespace factory = model::backend::mfem;
    namespace deformation = serif::geometry::deformation;
    namespace geometry = deformation::backend::mfem;
    namespace domain = serif::discretization::domain;
    using Resources = serif::discretization::backend::mfem::Resources;
    using Schema = domain::schema::CoreEnvelopeVacuumDomainSchema;
    constexpr int stellar_surface = Schema::boundary_id<domain::StellarSurfaceBoundary>();
    constexpr int infinity_surface = Schema::boundary_id<domain::InfinitySurfaceBoundary>();

    struct SurfaceStatistics {
        int faces{};
        int equatorial_faces{};
        int polar_faces{};
        double largest_equatorial_area{};
        double smallest_polar_area{std::numeric_limits<double>::infinity()};
        std::vector<double> absolute_latitudes;
        std::vector<double> areas;
    };

    // Measure concentration in physical space, independently of the selector's
    // logical coordinates and refinement-depth bookkeeping.
    SurfaceStatistics surface_statistics(mfem::Mesh& mesh, std::array<double, 3> axis) {
        double scale = 0.0;
        for (const double value : axis) scale = std::max(scale, std::abs(value));
        for (double& value : axis) value /= scale;
        const double norm = std::hypot(axis[0], axis[1], axis[2]);
        for (double& value : axis) value /= norm;
        SurfaceStatistics result;
        mfem::Vector position(3);
        for (int boundary = 0; boundary < mesh.GetNBE(); ++boundary) {
            if (mesh.GetBdrAttribute(boundary) != stellar_surface) continue;
            ++result.faces;
            auto& transform = *mesh.GetBdrElementTransformation(boundary);
            transform.Transform(mfem::Geometries.GetCenter(mesh.GetBdrElementBaseGeometry(boundary)), position);
            double latitude = 0.0;
            for (int component = 0; component < 3; ++component) latitude += axis[component] * position[component];
            latitude = std::abs(latitude) / position.Norml2();
            double area = 0.0;
            const auto& rule = mfem::IntRules.Get(mesh.GetBdrElementBaseGeometry(boundary), 8);
            for (int sample = 0; sample < rule.GetNPoints(); ++sample) {
                const auto& point = rule.IntPoint(sample);
                transform.SetIntPoint(&point);
                area += point.weight * transform.Weight();
            }
            REQUIRE(std::isfinite(area));
            REQUIRE(area > 0.0);
            result.absolute_latitudes.push_back(latitude);
            result.areas.push_back(area);
            if (latitude < std::sin(15.0 * std::numbers::pi / 180.0)) {
                ++result.equatorial_faces;
                result.largest_equatorial_area = std::max(result.largest_equatorial_area, area);
            }
            if (latitude > 0.75) {
                ++result.polar_faces;
                result.smallest_polar_area = std::min(result.smallest_polar_area, area);
            }
        }
        std::ranges::sort(result.absolute_latitudes);
        std::ranges::sort(result.areas);
        return result;
    }

    void check_mesh_invariants(Resources& resources) {
        REQUIRE(resources.meshes.size() == 2);
        auto& physical = *resources.meshes[0];
        auto& logical = *resources.meshes[1];
        REQUIRE(physical.GetNE() == logical.GetNE());
        REQUIRE(physical.GetNV() == logical.GetNV());
        REQUIRE(physical.GetNBE() == logical.GetNBE());
        CHECK(physical.Nonconforming() == logical.Nonconforming());
        mfem::Array<int> physical_vertices, logical_vertices;
        for (int element = 0; element < physical.GetNE(); ++element) {
            CHECK(physical.GetAttribute(element) == logical.GetAttribute(element));
            physical.GetElementVertices(element, physical_vertices);
            logical.GetElementVertices(element, logical_vertices);
            REQUIRE(physical_vertices.Size() == logical_vertices.Size());
            CHECK(std::ranges::equal(std::span{physical_vertices.GetData(), static_cast<std::size_t>(physical_vertices.Size())},
                std::span{logical_vertices.GetData(), static_cast<std::size_t>(logical_vertices.Size())}));
            auto& transform = *physical.GetElementTransformation(element);
            const auto& rule = mfem::IntRules.Get(physical.GetElementBaseGeometry(element), 6);
            for (int sample = 0; sample < rule.GetNPoints(); ++sample) {
                transform.SetIntPoint(&rule.IntPoint(sample));
                const double determinant = transform.Jacobian().Det();
                REQUIRE(std::isfinite(determinant));
                REQUIRE(determinant > 0.0);
            }
        }
        REQUIRE(resources.compactification);
        auto& coordinate = *resources.compactification;
        auto& scalar_space = *coordinate.FESpace();
        mfem::Array<int> dofs;
        int stellar_faces = 0, infinity_faces = 0;
        for (int boundary = 0; boundary < physical.GetNBE(); ++boundary) {
            CHECK(physical.GetBdrAttribute(boundary) == logical.GetBdrAttribute(boundary));
            physical.GetBdrElementVertices(boundary, physical_vertices);
            logical.GetBdrElementVertices(boundary, logical_vertices);
            REQUIRE(physical_vertices.Size() == logical_vertices.Size());
            CHECK(std::ranges::equal(std::span{physical_vertices.GetData(), static_cast<std::size_t>(physical_vertices.Size())},
                std::span{logical_vertices.GetData(), static_cast<std::size_t>(logical_vertices.Size())}));
            const int attribute = physical.GetBdrAttribute(boundary);
            if (attribute != stellar_surface && attribute != infinity_surface) continue;
            if (attribute == stellar_surface) ++stellar_faces;
            else ++infinity_faces;
            scalar_space.GetBdrElementDofs(boundary, dofs);
            for (const int dof : dofs) {
                CHECK(coordinate[mfem::FiniteElementSpace::DecodeDof(dof)] ==
                    Catch::Approx(attribute == stellar_surface ? 0.0 : 1.0).margin(2.0e-13));
            }
        }
        REQUIRE(stellar_faces > 0);
        REQUIRE(infinity_faces > 0);
        // This retains the production guard: volume hanging nodes are allowed,
        // but every stellar trace coefficient must remain independent.
        const auto trace = geometry::select_surface_scalar_dofs(scalar_space, stellar_surface);
        INFO((trace ? "Conforming stellar trace" : trace.error().message));
        REQUIRE(trace);
    }
}

TEST_CASE("Surface refinement adds levels beyond the existing global mesh", "[model][resources][surface_refinement]") {
    for (const int base : {0, 1}) {
        CAPTURE(base);
        model::StellarMeshOptions options{.refinement = base, .geometry_order = 2};
        auto original = factory::make_stellar_resources(options);
        const auto before = surface_statistics(*original.meshes[0], {0.0, 0.0, 1.0});
        options.surface_refinement.levels = 1;
        auto refined = factory::make_stellar_resources(options);
        const auto after = surface_statistics(*refined.meshes[0], {0.0, 0.0, 1.0});
        CHECK(after.faces == 4 * before.faces);
        check_mesh_invariants(refined);

        if (base == 1) {
            options.surface_refinement.equatorial_levels = 1;
            auto combined = factory::make_stellar_resources(options);
            const auto combined_stats = surface_statistics(*combined.meshes[0], {0.0, 0.0, 1.0});
            CHECK(combined_stats.faces > after.faces);
            CHECK(combined_stats.polar_faces == after.polar_faces);
            check_mesh_invariants(combined);
        }
    }
}

TEST_CASE("Equatorial rings concentrate resolution around the requested physical axis", "[model][resources][surface_refinement]") {
    model::StellarMeshOptions options{.geometry_order = 2};
    auto original = factory::make_stellar_resources(options);
    const auto before = surface_statistics(*original.meshes[0], {0.0, 0.0, 1.0});
    options.surface_refinement.equatorial_levels = 3;
    auto native = factory::make_stellar_resources(options);
    const auto expected = surface_statistics(*native.meshes[0], options.surface_refinement.axis);
    CHECK(expected.faces > before.faces);
    CHECK(expected.equatorial_faces > before.equatorial_faces);
    CHECK(expected.polar_faces == before.polar_faces);
    CHECK(expected.largest_equatorial_area < 0.5 * expected.smallest_polar_area);
    check_mesh_invariants(native);

    const double largest = std::numeric_limits<double>::max();
    const double smallest = std::numeric_limits<double>::min();
    for (const std::array<double, 3> axis : {
        std::array{1.0, 0.0, 0.0}, std::array{1.0, 1.0, 1.0},
        std::array{-1.0, -1.0, -1.0}, std::array{0.0, 0.0, -1.0},
        std::array{largest, largest, largest}, std::array{smallest, smallest, smallest}}) {
        CAPTURE(axis);
        options.surface_refinement.axis = axis;
        auto rotated = factory::make_stellar_resources(options);
        const auto actual = surface_statistics(*rotated.meshes[0], axis);
        REQUIRE(actual.faces == expected.faces);
        CHECK(actual.equatorial_faces == expected.equatorial_faces);
        CHECK(actual.polar_faces == expected.polar_faces);
        for (std::size_t index = 0; index < actual.areas.size(); ++index) {
            CHECK(actual.absolute_latitudes[index] == Catch::Approx(expected.absolute_latitudes[index]).margin(2.0e-12));
            CHECK(actual.areas[index] == Catch::Approx(expected.areas[index]).epsilon(2.0e-12));
        }
        check_mesh_invariants(rotated);
    }
}

TEST_CASE("Refined rotated meshes retain one surface and volume deformation", "[model][resources][surface_refinement]") {
    const auto star = serif::tests::solver::make_model();
    const auto finite_elements = model::make_stellar_discretization(star,
        {.geometry_order = 3, .surface_refinement = {.levels = 1, .equatorial_levels = 2, .axis = {1.0, 1.0, 1.0}}},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{
            .surface = {.maximum_degree = 2}, .stellar_interior = {}, .vacuum = {}});
    std::vector<double> shape(9);
    for (std::size_t mode = 0; mode < shape.size(); ++mode) shape[mode] = 0.002 * std::sin(0.37 * (mode + 1));
    const auto inputs = geometry::make_deformation_inputs(finite_elements.deformation.model, finite_elements.deformation.reference, shape);
    auto cache = serif::utils::dependencies::make_computation_cache<deformation::DeformationFailure>(
        geometry::make_deformation_graph<deformation::SphericalHarmonicSurface>(), inputs);
    const auto prepared = cache.try_prepare(inputs);
    INFO((prepared ? "Deformation graph prepared" : prepared.error().error.message));
    REQUIRE(prepared);
    const auto& reference = *cache.get<deformation::keys::ReferenceSurface>();
    const auto& surface = cache.get<deformation::keys::SurfaceDisplacement>();
    const auto& volume = cache.get<deformation::keys::VolumeDisplacement>();
    const auto& space = finite_elements.deformation.reference.displacement_space;
    REQUIRE(std::ranges::any_of(surface, [](double value) { return std::abs(value) > 1.0e-6; }));
    for (std::size_t node = 0; node < reference.scalar_dofs.size(); ++node) {
        for (int component = 0; component < 3; ++component) {
            CHECK(volume[reference.scalar_dofs[node] + component * space.GetNDofs()] ==
                Catch::Approx(surface[3 * node + component]).margin(2.0e-12));
        }
    }
}

TEST_CASE("Surface refinement rejects undefined levels bands and axes", "[model][resources][surface_refinement]") {
    model::StellarMeshOptions options{.geometry_order = 2};
    SECTION("negative surface levels") { options.surface_refinement.levels = -1; }
    SECTION("negative equatorial levels") { options.surface_refinement.equatorial_levels = -1; }
    SECTION("zero axis") { options.surface_refinement.axis = {0.0, 0.0, 0.0}; }
    SECTION("nonfinite axis") { options.surface_refinement.axis = {0.0, std::numeric_limits<double>::quiet_NaN(), 1.0}; }
    SECTION("nonfinite band") { options.surface_refinement.equatorial_half_angle_degrees = std::numeric_limits<double>::quiet_NaN(); }
    SECTION("zero band") { options.surface_refinement.equatorial_half_angle_degrees = 0.0; }
    SECTION("band outside the supported side patches") { options.surface_refinement.equatorial_half_angle_degrees = 35.1; }
    CHECK_THROWS(factory::make_stellar_resources(options));
}