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

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>

import mean_field;

namespace {
    using mean_field::solver::PolynomialSignStatus;
    constexpr std::array cubicNodes{0.0, 1.0/3.0, 2.0/3.0, 1.0};
    constexpr std::array quadraticNodes{0.0, 0.5, 1.0};

    template <class Function>
    std::vector<double> Values(std::span<const double> nodes, int dimension, Function function) {
        int count = 1;
        for (int d = 0; d < dimension; ++d) count *= static_cast<int>(nodes.size());
        std::vector<double> values(count);
        for (int i = 0; i < count; ++i) {
            int index = i;
            std::array<double,3> point{};
            for (int d = 0; d < dimension; ++d) {
                point[d] = nodes[index%nodes.size()]; index /= static_cast<int>(nodes.size());
            }
            values[i] = function(point);
        }
        return values;
    }

    struct ParallelFixture {
        mfem::Mesh serial{mfem::Mesh::MakeCartesian3D(2,2,1,mfem::Element::HEXAHEDRON,1.0,1.0,1.0)};
        mfem::ParMesh mesh{MPI_COMM_WORLD,serial};
        mfem::H1_FECollection collection{3,3};
        mfem::ParFiniteElementSpace space{&mesh,&collection};
        std::vector<int> elements;
        mean_field::field::FieldDofMap dofMap;
        ParallelFixture() {
            elements.resize(space.GetNE()); std::iota(elements.begin(),elements.end(),0);
            mfem::Array<int> trueDofs(space.GetTrueVSize());
            for (int i = 0; i < trueDofs.Size(); ++i) trueDofs[i] = i;
            dofMap = mean_field::field::FieldDofMap(space.GetTrueVSize(),trueDofs);
        }
        template <class Function> mfem::Vector Project(Function function) {
            mfem::FunctionCoefficient coefficient(function);
            mfem::ParGridFunction grid(&space); grid.ProjectCoefficient(coefficient);
            mfem::Vector values; grid.GetTrueDofs(values); return values;
        }
    };
}

TEST_CASE("Bernstein bounds detect a nodal-positive cubic with an interior negative pocket",
          "[solver][enthalpy_admissibility][unit]") {
    mean_field::solver::TensorBernsteinPolynomial<1> polynomial(cubicNodes);
    const auto values = Values(cubicNodes,1,[](auto p) {
        return ((p[0]-0.17)*(p[0]-0.17)-0.002)*(1.0+0.2*p[0]);
    });
    REQUIRE(std::all_of(values.begin(),values.end(),[](double value) { return value > 0.0; }));
    const auto report = polynomial.Check(values);
    REQUIRE(report.status == PolynomialSignStatus::negative);
    const double x = report.witnessReferencePoint[0];
    CHECK(((x-0.17)*(x-0.17)-0.002)*(1.0+0.2*x) < 0.0);
}

TEST_CASE("Negative coarse Bernstein coefficients cause refinement rather than a false rejection",
          "[solver][enthalpy_admissibility][unit]") {
    mean_field::solver::TensorBernsteinPolynomial<1> polynomial(quadraticNodes);
    const auto values = Values(quadraticNodes,1,[](auto p) { return (p[0]-0.5)*(p[0]-0.5)+0.01; });
    CHECK(polynomial.Check(values,{.maximumSubdivisionDepth=0}).status == PolynomialSignStatus::unresolved);
    const auto refined = polynomial.Check(values);
    CHECK(refined.status == PolynomialSignStatus::certified_nonnegative);
    CHECK(refined.visitedSubcells > 1);
    CHECK(refined.lowerBound >= 0.0);
}

TEST_CASE("Tensor certification retains structural zero faces without an admissibility epsilon",
          "[solver][enthalpy_admissibility][unit]") {
    mean_field::solver::TensorBernsteinPolynomial<3> polynomial(cubicNodes);
    const auto values = Values(cubicNodes,3,[](auto p) {
        return p[0]*(1.0-p[0])*p[1]*(1.0-p[1])*p[2]*(1.0-p[2]);
    });
    const auto report = polynomial.Check(values);
    CHECK(report.status == PolynomialSignStatus::certified_nonnegative);
    CHECK(report.lowerBound == 0.0);
    const std::vector<double> zero(polynomial.Size(),0.0);
    CHECK(polynomial.Check(zero).status == PolynomialSignStatus::certified_nonnegative);
}

TEST_CASE("A small negative pocket is not hidden by relative or absolute positivity tolerances",
          "[solver][enthalpy_admissibility][unit]") {
    mean_field::solver::TensorBernsteinPolynomial<1> polynomial(quadraticNodes);
    const auto values = Values(quadraticNodes,1,[](auto p) { return (p[0]-0.37)*(p[0]-0.37)-1.0e-10; });
    CHECK(polynomial.Check(values).status == PolynomialSignStatus::negative);
    auto invalid = values; invalid[1] = std::numeric_limits<double>::quiet_NaN();
    CHECK(polynomial.Check(invalid).status == PolynomialSignStatus::non_finite);
    invalid[1] = std::numeric_limits<double>::infinity();
    CHECK(polynomial.Check(invalid).status == PolynomialSignStatus::non_finite);
    CHECK_THROWS_AS(polynomial.Check(std::span<const double>{}),std::invalid_argument);
}

TEST_CASE("Certification supports runtime polynomial degree and honest unresolved results at tangential zeros",
          "[solver][enthalpy_admissibility][unit]") {
    const std::array nodes{0.0,0.2,0.4,0.6,0.8,1.0};
    mean_field::solver::TensorBernsteinPolynomial<2> polynomial(nodes);
    const auto values = Values(nodes,2,[](auto p) { return 0.2+std::pow(p[0],5)+p[1]; });
    CHECK(polynomial.Order() == 5);
    CHECK(polynomial.Check(values).status == PolynomialSignStatus::certified_nonnegative);
    mean_field::solver::TensorBernsteinPolynomial<1> quadratic(quadraticNodes);
    const auto tangent = Values(quadraticNodes,1,[](auto p) { return (p[0]-0.37)*(p[0]-0.37); });
    CHECK(quadratic.Check(tangent,{.maximumSubdivisionDepth=6}).status == PolynomialSignStatus::unresolved);
}

TEST_CASE("Feasible-anchor fractions describe a sufficient convex polynomial cone",
          "[solver][enthalpy_admissibility][unit]") {
    mean_field::solver::TensorBernsteinPolynomial<1> polynomial(cubicNodes);
    const auto anchor = Values(cubicNodes,1,[](auto p) { return p[0]*(1.0-p[0]); });
    const auto candidate = Values(cubicNodes,1,[](auto p) { return p[0]*(1.0-p[0])*(1.0-6.0*p[0]); });
    const double fraction = polynomial.MinimumRepairFraction(anchor,candidate);
    REQUIRE(std::isfinite(fraction)); REQUIRE(fraction > 0.0); REQUIRE(fraction <= 1.0);
    std::vector<double> repaired(anchor.size());
    const double safe = fraction+1.0e-10*(1.0-fraction);
    for (std::size_t i = 0; i < repaired.size(); ++i) repaired[i] = (1.0-safe)*candidate[i]+safe*anchor[i];
    CHECK(polynomial.Check(repaired).status == PolynomialSignStatus::certified_nonnegative);
    CHECK(repaired.front() == 0.0); CHECK(repaired.back() == 0.0);
}

TEST_CASE("Prepared admissibility preserves shared true DOFs and rejects an invalid starting state",
          "[solver][enthalpy_admissibility][integration][parallel]") {
    using namespace mean_field::solver;
    ParallelFixture fixture;
    PreparedEnthalpyAdmissibility prepared(fixture.space,fixture.dofMap,fixture.elements);
    const auto base = fixture.Project([](const mfem::Vector &) { return 1.0; });
    const auto anchor = fixture.Project([](const mfem::Vector &p) { return 1.0+0.2*p(1); });
    const auto direction = fixture.Project([](const mfem::Vector &p) { return -2.0-p(0)+0.5*p(2); });
    const mfem::Vector savedBase(base), savedAnchor(anchor), savedDirection(direction);
    const auto repair = prepared.MinimumRepairFraction(base,anchor,direction);
    REQUIRE(repair.status == EnthalpyDirectionRepairStatus::feasible);
    REQUIRE(repair.minimumSafeFraction > 0.0);
    mfem::Vector repaired;
    const double beta = repair.minimumSafeFraction+1.0e-8*(1.0-repair.minimumSafeFraction);
    REQUIRE(prepared.ApplyRepairFraction(base,anchor,direction,beta,repaired).status ==
            PolynomialSignStatus::certified_nonnegative);
    mfem::Vector candidate(base); candidate += repaired;
    CHECK(prepared.Check(candidate).status == PolynomialSignStatus::certified_nonnegative);
    for (int i = 0; i < repaired.Size(); ++i)
        CHECK(repaired(i) == Catch::Approx((1.0-beta)*direction(i)+beta*(anchor(i)-base(i))));
    for (int i = 0; i < base.Size(); ++i) {
        CHECK(base(i) == savedBase(i)); CHECK(anchor(i) == savedAnchor(i));
        CHECK(direction(i) == savedDirection(i));
    }
    const auto repeated = prepared.MinimumRepairFraction(base,anchor,direction);
    CHECK(repeated.minimumSafeFraction == repair.minimumSafeFraction);
    mfem::Vector invalid(base); invalid *= -1.0;
    CHECK(prepared.MinimumRepairFraction(invalid,anchor,direction).status ==
          EnthalpyDirectionRepairStatus::invalid_base);
    int rank = 0; MPI_Comm_rank(MPI_COMM_WORLD,&rank);
    mfem::Vector nonfinite(base);
    if (rank == 0 && nonfinite.Size()) nonfinite(0) = std::numeric_limits<double>::quiet_NaN();
    CHECK(prepared.Check(nonfinite).status == PolynomialSignStatus::non_finite);
    const auto saved = repaired;
    CHECK(prepared.ApplyRepairFraction(base,anchor,direction,0.0,repaired).status == PolynomialSignStatus::negative);
    for (int i = 0; i < saved.Size(); ++i) CHECK(repaired(i) == saved(i));
}

TEST_CASE("Feasible direction repair enforces exact homogeneous trace and rejects incompatible traces",
          "[solver][enthalpy_admissibility][integration][parallel]") {
    using namespace mean_field::solver;
    ParallelFixture fixture;
    mfem::Array<int> marker(fixture.mesh.bdr_attributes.Max()); marker = 1;
    mfem::Array<int> boundary; fixture.space.GetEssentialTrueDofs(marker,boundary);
    PreparedEnthalpyAdmissibility prepared(fixture.space,fixture.dofMap,fixture.elements,
        std::span<const int>(boundary.GetData(),boundary.Size()));
    const auto anchor = fixture.Project([](const mfem::Vector &p) {
        return p(0)*(1.0-p(0))*p(1)*(1.0-p(1))*p(2)*(1.0-p(2));
    });
    const auto base = anchor;
    mfem::Vector direction(anchor); direction *= -2.0;
    for (int row : boundary) direction(row) = -0.25;
    const auto report = prepared.MinimumRepairFraction(base,anchor,direction);
    REQUIRE(report.status == EnthalpyDirectionRepairStatus::feasible);
    mfem::Vector repaired;
    REQUIRE(prepared.ApplyRepairFraction(base,anchor,direction,1.0,repaired).status ==
            PolynomialSignStatus::certified_nonnegative);
    for (int row : boundary) CHECK(repaired(row) == 0.0);
    auto nonfiniteDirection = direction;
    for (int row : boundary) nonfiniteDirection(row) = std::numeric_limits<double>::quiet_NaN();
    CHECK(prepared.MinimumRepairFraction(base,anchor,nonfiniteDirection).status ==
          EnthalpyDirectionRepairStatus::invalid_direction);
    auto invalidTrace = base;
    for (int row : boundary) invalidTrace(row) = 1.0e-20;
    CHECK(prepared.MinimumRepairFraction(invalidTrace,anchor,direction).status ==
          EnthalpyDirectionRepairStatus::invalid_base);
}