#include <array>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <initializer_list>
#include <limits>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>

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

import mean_field;

#include "../../experiments/plateau_coupled_coarse.hpp"

namespace coarse_test {
    namespace probe = ::experiment::plateau_diagnostics;
    using namespace mean_field;

    [[nodiscard]] mfem::Vector Vector(std::initializer_list<double> values) {
        mfem::Vector result(static_cast<int>(values.size()));
        int index = 0;
        for (double value : values)
            result(index++) = value;
        return result;
    }

    [[nodiscard]] mfem::DenseMatrix Identity(int size) {
        mfem::DenseMatrix result(size);
        result = 0.0;
        for (int i = 0; i < size; ++i)
            result(i, i) = 1.0;
        return result;
    }

    [[nodiscard]] double Norm(const mfem::Vector &value) {
        double local = value * value, global = 0.0;
        MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
        return std::sqrt(global);
    }

    [[nodiscard]] double Difference(
        const mfem::Vector &left,
        const mfem::Vector &right
    ) {
        mfem::Vector difference(left);
        difference -= right;
        return Norm(difference);
    }

    class Operation final : public mfem::Operator {
    public:
        explicit Operation(mfem::DenseMatrix matrix)
            : mfem::Operator(
                  matrix.Height(),
                  matrix.Width()
              ),
              matrix(std::move(matrix)),
              factors(Width()) {
            factors = 1.0;
        }
        void Mult(
            const mfem::Vector &value,
            mfem::Vector &action
        ) const override {
            ++applications;
            matrix.Mult(value, action);
        }
        void NormalizeState(
            const mfem::Vector &physical,
            mfem::Vector &normalized
        ) const {
            ++normalizations;
            for (int i = 0; i < Width(); ++i)
                normalized(i) = physical(i) * factors(i);
        }
        bool IsPrepared() const {
            return true;
        }
        const Operation &GetProblem() const {
            return *this;
        }
        std::uint64_t GetPreparationGeneration() const {
            return generation;
        }
        void Prepare() {
            ++generation;
        }
        mfem::DenseMatrix matrix;
        mfem::Vector factors;
        mutable std::uint64_t applications{0}, normalizations{0};
        std::uint64_t generation{1};
    };

    class Inverse : public mfem::Solver {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        explicit Inverse(const Operation &operation)
            : mfem::Solver(operation.Width()),
              matrix(Identity(operation.Width())),
              operation(operation) {
            Refresh();
        }
        void SetOperator(const mfem::Operator &) override {
        }
        void Mult(
            const mfem::Vector &rhs,
            mfem::Vector &action
        ) const override {
            ++applications;
            matrix.Mult(rhs, action);
            if (produceNonfinite && action.Size())
                action(0) = std::numeric_limits<double>::infinity();
        }
        void Refresh() {
            generation = operation.generation;
        }
        bool IsCurrent() const {
            return generation == operation.generation;
        }
        mfem::DenseMatrix matrix;
        const Operation &operation;
        bool produceNonfinite{false};
        mutable std::uint64_t applications{0};
        std::uint64_t generation{0};
    };

    class FlexibleInverse final : public Inverse {
    public:
        using Inverse::Inverse;
        static constexpr auto applicationContract = preconditioning::ApplicationContract::flexible;
    };

    static_assert(probe::CoupledCoarseNormalizedOperator<Operation>);
    static_assert(probe::CoupledCoarseStationaryInverse<Inverse>);
    static_assert(!probe::CoupledCoarseStationaryInverse<FlexibleInverse>);
    static_assert(!probe::CoupledCoarseNormalizedOperator<mfem::Operator>);

    using OrdinaryModel                  = model::StellarModel<models::SpecificationSet<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        integral::FixedAngularMomentum>>;
    using CentralModel                   = model::StellarModel<models::SpecificationSet<
        eos::Polytrope,
        surface::Isobaric,
        integral::FixedTotalMass,
        integral::FixedAngularMomentum,
        constraint::FixedCentralDensity>>;
    template <class Model> using Problem = equilibrium::StellarEquilibriumProblem<Model>;
    template <class Model>
    using NormalizedOperation = normalization::NormalizedStellarEquilibriumOperator<Problem<Model>>;
    template <class Model>
    using PhysicalInverse =
        preconditioning::PreparedStellarInverseType<preconditioning::DefaultStellarPreconditioner, Problem<Model>>;
    template <class Model>
    using NormalizedInverse = normalization::NormalizedStellarPreconditioner<Problem<Model>, PhysicalInverse<Model>>;
    template <class Model>
    using Coarse = probe::PreparedCoupledShapeCoarse<NormalizedOperation<Model>, NormalizedInverse<Model>>;

    // The production problem erases its prepared deformation. Recompile only
    // the existing surface prescription, then verify its actual runtime action.
    // Explicit instantiations below keep both complete model paths checked.
    template <class Model>
    concept ReusesRadialCompiler = requires(const Problem<Model> &problem) {
        {
            probe::coupled_shape_detail::MakeDefaultRadialSurface(problem)
        } -> std::same_as<deformation::PreparedNodalRadialSurface>;
        requires probe::coupled_shape_detail::SupportsShapeModes<Problem<Model>, PhysicalInverse<Model>>;
    };
    static_assert(ReusesRadialCompiler<OrdinaryModel>);
    static_assert(ReusesRadialCompiler<CentralModel>);
    static_assert(solver::LinearPreconditionerApplicationContractAvailable<Coarse<OrdinaryModel>>);
    static_assert(solver::LinearPreconditionerApplicationContractAvailable<Coarse<CentralModel>>);
} // namespace coarse_test

template class experiment::plateau_diagnostics::PreparedCoupledShapeCoarse<
    coarse_test::NormalizedOperation<coarse_test::OrdinaryModel>,
    coarse_test::NormalizedInverse<coarse_test::OrdinaryModel>>;
template class experiment::plateau_diagnostics::PreparedCoupledShapeCoarse<
    coarse_test::NormalizedOperation<coarse_test::CentralModel>,
    coarse_test::NormalizedInverse<coarse_test::CentralModel>>;

template std::vector<mfem::Vector> experiment::plateau_diagnostics::MakeCoupledShapeModes(
    const coarse_test::Problem<coarse_test::OrdinaryModel> &,
    const coarse_test::PhysicalInverse<coarse_test::OrdinaryModel> &
);
template std::vector<mfem::Vector> experiment::plateau_diagnostics::MakeCoupledShapeModes(
    const coarse_test::Problem<coarse_test::CentralModel> &,
    const coarse_test::PhysicalInverse<coarse_test::CentralModel> &
);

TEST_CASE(
    "Filtered shape coarse inverse is exact on its source space and linear",
    "[solver][coupled_shape_coarse]"
) {
    using namespace coarse_test;
    mfem::DenseMatrix matrix(3);
    matrix       = 0.0;
    matrix(0, 0) = 3.0;
    matrix(0, 1) = 0.7;
    matrix(0, 2) = 0.2;
    matrix(1, 0) = -0.2;
    matrix(1, 1) = 2.0;
    matrix(1, 2) = 0.4;
    matrix(2, 0) = 0.5;
    matrix(2, 2) = 1.3;
    Operation operation(matrix);
    operation.factors = Vector({0.2, 3.0, 10.0});
    Inverse inverse(operation);
    inverse.matrix(0, 0) = 0.4;
    inverse.matrix(1, 1) = 0.5;
    inverse.matrix(2, 2) = 0.8;
    std::vector<mfem::Vector> modes{Vector({1.0, 0.2, -0.1}), Vector({0.3, 1.0, 0.4}), Vector({1.0, 0.2, -0.1})};
    probe::PreparedCoupledShapeCoarse coarse(operation, inverse, std::span<const mfem::Vector>(modes), MPI_COMM_WORLD);
    REQUIRE(coarse.UsesCoarseCorrection());
    REQUIRE(coarse.GetSourceBasis().size() == 2);
    CHECK(coarse.GetStatistics().activeModes == 2);
    CHECK(coarse.GetStatistics().sourceDependentModes == 1);
    CHECK(coarse.GetStatistics().setupOperatorApplications == 5);
    CHECK(coarse.GetStatistics().setupInverseApplications == 2);
    CHECK(operation.normalizations == 3);
    CHECK(coarse.GetStatistics().persistentVectorBytes == 18 * sizeof(mfem::real_t));

    mfem::Vector value(3), rhs(3), actual(3), expected(3);
    operation.NormalizeState(modes[0], value);
    operation.Mult(value, rhs);
    rhs /= Norm(rhs);
    CHECK(Difference(coarse.GetSourceBasis()[0], rhs) < 2.e-14);
    for (const auto &source : coarse.GetSourceBasis()) {
        inverse.Mult(source, expected);
        operation.Mult(expected, rhs);
        const auto before = operation.applications;
        coarse.Mult(rhs, actual);
        CHECK(operation.applications == before);
        CHECK(Difference(actual, expected) < 2.e-12 * std::max(1.0, Norm(expected)));
    }
    const auto before        = operation.applications;
    const auto inverseBefore = inverse.applications;
    rhs                      = 0.0;
    actual                   = 17.0;
    coarse.Mult(rhs, actual);
    CHECK(Norm(actual) == 0.0);
    CHECK(operation.applications == before);
    CHECK(inverse.applications == inverseBefore + 1);

    const auto first = Vector({0.2, -0.8, 1.3}), second = Vector({-0.4, 0.1, 0.7});
    mfem::Vector firstAction(3), secondAction(3), combination(first);
    combination *= 0.7;
    combination.Add(-1.2, second);
    coarse.Mult(first, firstAction);
    coarse.Mult(second, secondAction);
    coarse.Mult(combination, actual);
    expected = firstAction;
    expected *= 0.7;
    expected.Add(-1.2, secondAction);
    CHECK(Difference(actual, expected) < 2.e-12);
    coarse.Mult(combination, combination);
    CHECK(Difference(combination, actual) < 2.e-14);
    CHECK(operation.applications == before);
}

TEST_CASE(
    "Filtered coarse overlap rejects a near-right-angle even at rank one",
    "[solver][coupled_shape_coarse]"
) {
    using namespace coarse_test;
    for (const double cosine : {0.25, 1.e-12, 0.0}) {
        const double sine = std::sqrt(1.0 - cosine * cosine);
        mfem::DenseMatrix rotation(2);
        rotation(0, 0) = rotation(1, 1) = cosine;
        rotation(0, 1)                  = -sine;
        rotation(1, 0)                  = sine;
        Operation operation(rotation);
        Inverse inverse(operation);
        std::array<mfem::Vector, 1> modes{Vector({1.0, 0.0})};
        probe::PreparedCoupledShapeCoarse coarse(
            operation, inverse, std::span<const mfem::Vector>(modes), MPI_COMM_WORLD
        );
        CHECK(coarse.GetStatistics().independentModes == 1);
        if (cosine == 0.25) {
            CHECK(coarse.UsesCoarseCorrection());
            CHECK(coarse.GetStatistics().minimumOverlapLowerBound == Catch::Approx(cosine).epsilon(1.e-12));
        } else {
            CHECK_FALSE(coarse.UsesCoarseCorrection());
            CHECK(coarse.GetStatistics().fallback == probe::CoupledShapeCoarseFallback::insufficient_subspace_overlap);
            if (cosine != 0.0) {
                CHECK(coarse.GetStatistics().reciprocalOverlapCondition == Catch::Approx(1.0));
                CHECK(coarse.GetStatistics().minimumOverlapLowerBound < 1.e-8);
            }
            const auto rhs = Vector({0.3, 0.7});
            mfem::Vector result(2);
            const auto before = operation.applications;
            coarse.Mult(rhs, result);
            CHECK(Difference(result, rhs) == 0.0);
            CHECK(operation.applications == before);
            CHECK(coarse.GetStatistics().inverseApplications == 1);
        }
    }
}

TEST_CASE(
    "Coarse setup drops a null filtered image without altering equations",
    "[solver][coupled_shape_coarse]"
) {
    using namespace coarse_test;
    mfem::DenseMatrix matrix(2);
    matrix       = 0.0;
    matrix(0, 1) = 1.0;
    Operation operation(matrix);
    Inverse inverse(operation);
    std::array<mfem::Vector, 1> modes{Vector({0.0, 1.0})};
    probe::PreparedCoupledShapeCoarse coarse(operation, inverse, std::span<const mfem::Vector>(modes), MPI_COMM_WORLD);
    CHECK_FALSE(coarse.UsesCoarseCorrection());
    CHECK(coarse.GetStatistics().imageDependentModes == 1);
    CHECK(coarse.GetStatistics().fallback == probe::CoupledShapeCoarseFallback::no_independent_modes);
    CHECK(operation.matrix(0, 0) == 0.0);
    CHECK(operation.matrix(0, 1) == 1.0);
    CHECK(operation.matrix(1, 0) == 0.0);
    CHECK(operation.matrix(1, 1) == 0.0);
}

TEST_CASE(
    "Coarse inverse rejects stale preparation and preserves failed outputs",
    "[solver][coupled_shape_coarse]"
) {
    using namespace coarse_test;
    Operation operation(Identity(2));
    Inverse inverse(operation);
    std::array<mfem::Vector, 1> modes{Vector({1.0, 0.5})};
    probe::PreparedCoupledShapeCoarse coarse(operation, inverse, std::span<const mfem::Vector>(modes), MPI_COMM_WORLD);
    const auto rhs = Vector({0.3, 0.2});
    mfem::Vector result(2);
    result = 9.0;
    operation.Prepare();
    CHECK_FALSE(coarse.IsCurrent());
    CHECK_THROWS_AS(coarse.Mult(rhs, result), std::logic_error);
    CHECK(result(0) == 9.0);
    CHECK(result(1) == 9.0);
    inverse.Refresh();
    CHECK_FALSE(coarse.IsCurrent());
    coarse.Refresh(modes);
    CHECK(coarse.IsCurrent());
    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    inverse.produceNonfinite = rank == 0;
    CHECK_THROWS_AS(coarse.Mult(rhs, result), std::runtime_error);
    CHECK(result(0) == 9.0);
    CHECK(result(1) == 9.0);
    inverse.produceNonfinite = false;
    coarse.Mult(rhs, result);
    CHECK(Difference(result, rhs) < 1.e-14);
    CHECK(coarse.GetStatistics().failedApplications == 2);
    CHECK(coarse.GetStatistics().setups == 2);
}

TEST_CASE(
    "Coarse preparation supports locally empty ranks and rejects a globally empty system",
    "[solver][coupled_shape_coarse]"
) {
    using namespace coarse_test;
    int rank = 0;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    const int localSize = rank == 0 ? 2 : 0;
    Operation operation(Identity(localSize));
    Inverse inverse(operation);
    std::array<mfem::Vector, 1> modes{mfem::Vector(localSize)};
    modes[0] = 1.0;
    probe::PreparedCoupledShapeCoarse coarse(operation, inverse, std::span<const mfem::Vector>(modes), MPI_COMM_WORLD);
    REQUIRE(coarse.UsesCoarseCorrection());
    mfem::Vector result(localSize);
    coarse.Mult(modes[0], result);
    CHECK(Difference(result, modes[0]) < 1.e-14);
    Operation empty(Identity(0));
    Inverse emptyInverse(empty);
    std::array<mfem::Vector, 1> emptyMode{mfem::Vector(0)};
    CHECK_THROWS_AS(
        (probe::PreparedCoupledShapeCoarse(
            empty, emptyInverse, std::span<const mfem::Vector>(emptyMode), MPI_COMM_WORLD
        )),
        std::invalid_argument
    );
}