#include <array>
#include <cmath>
#include <cstdint>
#include <initializer_list>
#include <limits>
#include <type_traits>
#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_surface_border_schur.hpp"

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

    mfem::DenseMatrix Matrix(int size, std::initializer_list<double> entries) {
        mfem::DenseMatrix result(size);
        auto entry = entries.begin();
        for (int row = 0; row < size; ++row)
            for (int column = 0; column < size; ++column)
                result(row, column) = *entry++;
        return result;
    }

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

    double Difference(const mfem::Vector &left, const mfem::Vector &right) {
        mfem::Vector difference(left);
        difference -= right;
        return difference.Norml2();
    }

    class Operation final : public mfem::Operator {
    public:
        explicit Operation(mfem::DenseMatrix matrix)
            : mfem::Operator(matrix.Height()), matrix(std::move(matrix)),
              residualScales(Height()), stateScales(Width()) {
            residualScales = 1.0;
            stateScales = 1.0;
        }
        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            ++applications;
            matrix.Mult(input, output);
        }
        bool IsPrepared() const { return prepared; }
        const Operation &GetProblem() const { return *this; }
        std::uint64_t GetPreparationGeneration() const { return generation; }
        void DenormalizeResidual(const mfem::Vector &input, mfem::Vector &output) const {
            for (int i = 0; i < input.Size(); ++i)
                output(i) = input(i) * residualScales(i);
        }
        void NormalizeState(const mfem::Vector &input, mfem::Vector &output) const {
            for (int i = 0; i < input.Size(); ++i)
                output(i) = input(i) / stateScales(i);
        }
        mfem::DenseMatrix matrix;
        mfem::Vector residualScales, stateScales;
        bool prepared{true};
        std::uint64_t generation{1};
        mutable std::uint64_t applications{0};
    };

    class Inverse : public mfem::Solver {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        Inverse(const Operation &operation, mfem::DenseMatrix matrix)
            : mfem::Solver(matrix.Height()), operation(operation), matrix(std::move(matrix)),
              generation(operation.generation) {}
        void SetOperator(const mfem::Operator &) override {}
        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            ++applications;
            matrix.Mult(input, output);
            if (nonfinite)
                output(0) = std::numeric_limits<double>::infinity();
        }
        bool IsCurrent() const { return generation == operation.generation; }
        const Operation &GetNormalizedJacobian() const { return operation; }
        const Operation &operation;
        mfem::DenseMatrix matrix;
        std::uint64_t generation;
        bool nonfinite{false};
        mutable std::uint64_t applications{0};
    };

    class FlexibleInverse : public Inverse {
    public:
        using Inverse::Inverse;
        static constexpr auto applicationContract = preconditioning::ApplicationContract::flexible;
    };
    static_assert(probe::SurfaceBorderSchurOperation<Operation>);
    static_assert(probe::SurfaceBorderSchurInverse<Inverse>);
    static_assert(!probe::SurfaceBorderSchurInverse<FlexibleInverse>);
    static_assert(!probe::SurfaceBorderSchurOperation<mfem::Operator>);

    using Schur = probe::PreparedSurfaceBorderSchur<Operation, Inverse>;
    static_assert(probe::SurfaceBorderSchurInverse<Schur>);
    static_assert(!std::is_copy_constructible_v<Schur>);

    // Fields occupy different equation/unknown positions. A contiguous or
    // same-offset assumption would silently factor the wrong block here.
    probe::SurfaceBorderSchurCoordinates Coordinates() { return {{1, 3}, {0, 2}}; }

    mfem::DenseMatrix CoupledMatrix() {
        return Matrix(4, {0.3, 4, -1, 1,
                          0, 2,  2, 0.4,
                        0.5, 2, 0.7, 3,
                          3, -0.2, 1, 0.9});
    }

    mfem::DenseMatrix FieldInverse(bool exact) {
        mfem::DenseMatrix result(4);
        result = 0.0;
        result(1, 0) = exact ? 0.3 : 0.2;
        result(1, 2) = exact ? -0.1 : 0.0;
        result(3, 0) = exact ? -0.2 : 0.0;
        result(3, 2) = exact ? 0.4 : 0.2;
        // These coarse contributions must be excluded, not accidentally used
        // to complete the field response twice.
        result(0, 0) = 11;
        result(2, 2) = -17;
        result(1, 1) = 23;
        result(3, 3) = -31;
        return result;
    }

    namespace blocks = utils::blocks;
    struct StructureDescriptor {
        using BackendType = preconditioning::backend::Diagonal;
        using CorrectionBlocks = blocks::type_list<blocks::density::mass::value,
            blocks::surface_deformation::parameters::value, blocks::enthalpy::specific::value>;
        using ResidualBlocks = blocks::type_list<blocks::enthalpy::specific::residual,
            blocks::surface_deformation::shape_equilibrium::residual, blocks::density::mass::residual>;
    };
    struct RootDescriptor {
        using StructureComponent = StructureDescriptor;
        using CorrectionBlocks = blocks::type_list<blocks::density::mass::value,
            blocks::surface_deformation::parameters::value, blocks::enthalpy::specific::value,
            blocks::fixed_total_mass::mass_normalization::value,
            blocks::fixed_central_density::central_value::value>;
        using ResidualBlocks = blocks::type_list<blocks::enthalpy::specific::residual,
            blocks::surface_deformation::shape_equilibrium::residual, blocks::density::mass::residual,
            blocks::fixed_total_mass::mass_normalization::residual,
            blocks::fixed_central_density::central_value::residual>;
        static constexpr std::size_t borderValueArity = 2, borderResidualArity = 2;
    };
    struct CoordinateMap {
        using Range = preconditioning::EquilibriumCoordinateRange;
        std::array<Range, 5> columns{{{5, 0, 2}, {1, 2, 3}, {0, 5, 1}, {4, 6, 1}, {7, 7, 1}}};
        std::array<Range, 5> rows{{{6, 0, 1}, {2, 1, 3}, {0, 4, 2}, {5, 6, 1}, {7, 7, 1}}};
        const auto &GetCorrectionRanges() const { return columns; }
        const auto &GetResidualRanges() const { return rows; }
        void PackResidual(const mfem::Vector &input, mfem::Vector &output) const {
            for (const auto &range : rows)
                for (int i = 0; i < range.size; ++i)
                    output(range.preconditionerOffset + i) = input(range.equilibriumOffset + i);
        }
        void UnpackCorrection(const mfem::Vector &input, mfem::Vector &output) const {
            for (const auto &range : columns)
                for (int i = 0; i < range.size; ++i)
                    output(range.equilibriumOffset + i) = input(range.preconditionerOffset + i);
        }
    };
    struct CoordinateBinding {
        RootDescriptor root;
        CoordinateMap coordinates;
        const auto &GetBlock() const { return root; }
        const auto &GetCoordinateMap() const { return coordinates; }
    };

    class PackedStructure final : public mfem::Solver {
    public:
        PackedStructure() : mfem::Solver(6), matrix(6) {
            matrix = 0.0;
            for (int i = 0; i < 6; ++i)
                matrix(i, i) = i + 2.0;
            matrix(0, 5) = 0.4;
            matrix(4, 1) = -0.6;
        }
        void SetOperator(const mfem::Operator &) override {}
        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            ++applications;
            matrix.Mult(input, output);
        }
        mfem::DenseMatrix matrix;
        mutable unsigned applications{0};
    };
    struct GroupedStructure {
        PackedStructure structure;
        const auto &GetStructurePreconditioner() const { return structure; }
    };
    class PhysicalStructureBinding final : public mfem::Solver, public CoordinateBinding {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        explicit PhysicalStructureBinding(const Operation &operation)
            : mfem::Solver(operation.Width()), operation(operation), generation(operation.generation) {}
        void SetOperator(const mfem::Operator &) override {}
        void Mult(const mfem::Vector &, mfem::Vector &) const override {
            throw std::logic_error("The adapter must not apply the scalar-bordered inverse.");
        }
        const auto &GetProblem() const { return operation; }
        const auto &GetGroupedPreconditioner() const { return grouped; }
        bool IsCurrent() const { return generation == operation.generation; }
        const Operation &operation;
        std::uint64_t generation;
        GroupedStructure grouped;
    };
}

TEST_CASE("Surface Border Schur Structure Adapter Preserves Nonunit Normalization And Typed Packing",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    mfem::DenseMatrix matrix(8);
    matrix = 0.0;
    Operation operation(std::move(matrix));
    operation.residualScales = Vector({2, 3, 5, 7, 11, 13, 17, 19});
    operation.stateScales = Vector({23, 29, 31, 37, 41, 43, 47, 53});
    PhysicalStructureBinding physical(operation);
    probe::NormalizedUnborderedStructureInverse adapter(operation, physical);
    STATIC_CHECK(probe::SurfaceBorderSchurInverse<decltype(adapter)>);
    auto rhs = Vector({1, 2, 3, 4, 5, 6, 7, 8});
    const auto expected = Vector({42.0 / 23, 112.0 / 29, 275.0 / 31, 3.0 / 37,
                                  0, 240.4 / 43, 45.0 / 47, 0});
    mfem::Vector observed(8);
    adapter.Mult(rhs, observed);
    CHECK(Difference(observed, expected) < 1.e-13);
    CHECK(physical.grouped.structure.applications == 1);
    adapter.Mult(rhs, rhs);
    CHECK(Difference(rhs, expected) < 1.e-13);
    physical.grouped.structure.iterative_mode = true;
    CHECK_FALSE(adapter.IsCurrent());
    physical.grouped.structure.iterative_mode = false;
    ++operation.generation;
    physical.generation = operation.generation;
    CHECK_FALSE(adapter.IsCurrent());
    CHECK_THROWS_AS(adapter.Mult(rhs, observed), std::logic_error);
}

TEST_CASE("Surface Border Schur Coordinates Include All Native Modes And Generated Central Border",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    const auto coordinates = probe::MakeSurfaceBorderSchurCoordinates(CoordinateBinding{});
    CHECK(coordinates.correctionIndices == std::vector<int>{1, 2, 3, 4, 7});
    CHECK(coordinates.residualIndices == std::vector<int>{2, 3, 4, 5, 7});
}

TEST_CASE("Surface Border Schur With Exact Field Inverse Solves Full Nonsymmetric System",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(CoupledMatrix());
    Inverse inverse(operation, FieldInverse(true));
    Schur schur(operation, inverse, Coordinates(), MPI_COMM_WORLD);
    REQUIRE(schur.IsCurrent());
    CHECK(schur.GetStatistics().setupOperatorApplications == 4);
    CHECK(schur.GetStatistics().setupInverseApplications == 2);
    const auto rhs = Vector({1, -2, 3, 4});
    mfem::Vector action(4), expected(4), defect(4);
    mfem::DenseMatrixInverse direct(operation.matrix);
    direct.Mult(rhs, expected);
    schur.Mult(rhs, action);
    CHECK(Difference(action, expected) < 1.e-13);
    operation.Mult(action, defect);
    CHECK(Difference(defect, rhs) < 1.e-13);
    CHECK(schur.GetStatistics().operatorApplications == 1);
    CHECK(schur.GetStatistics().inverseApplications == 1);
    for (int column = 0; column < 2; ++column) {
        const auto &completed = schur.GetCompletedColumns()[column];
        operation.Mult(completed, defect);
        CHECK(std::abs(defect(0)) < 1.e-13);
        CHECK(std::abs(defect(2)) < 1.e-13);
        CHECK(completed(Coordinates().correctionIndices[column]) == 1.0);
    }
}

TEST_CASE("Surface Border Schur With Approximate Fields Enforces Every Coarse Equation",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(CoupledMatrix());
    Inverse inverse(operation, FieldInverse(false));
    Schur schur(operation, inverse, Coordinates(), MPI_COMM_WORLD);
    REQUIRE(schur.IsCurrent());
    const auto rhs = Vector({1, -2, 3, 4});
    mfem::Vector action(4), image(4);
    schur.Mult(rhs, action);
    operation.Mult(action, image);
    for (int index : Coordinates().residualIndices)
        CHECK(std::abs(image(index) - rhs(index)) < 1.e-13);
    CHECK(Difference(image, rhs) > 1.e-3);
    // Verify S = D - A_CF B A_FC independently in this explicit packing.
    const auto &s = schur.GetSchurMatrix();
    CHECK(s(0, 0) == Catch::Approx(-0.16));
    CHECK(s(0, 1) == Catch::Approx(2.344));
    CHECK(s(1, 0) == Catch::Approx(2.922));
    CHECK(s(1, 1) == Catch::Approx(0.834));
}

TEST_CASE("Surface Border Schur Is Stationary Linear And Supports Aliasing",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(CoupledMatrix());
    Inverse inverse(operation, FieldInverse(false));
    Schur schur(operation, inverse, Coordinates(), MPI_COMM_WORLD);
    const auto first = Vector({1, 2, -1, 3}), second = Vector({-2, 4, 0.5, 1});
    mfem::Vector a(4), b(4), combined(first), expected(4), observed(4);
    schur.Mult(first, a);
    schur.Mult(second, b);
    combined *= 2.3;
    combined.Add(-0.7, second);
    expected = a;
    expected *= 2.3;
    expected.Add(-0.7, b);
    schur.Mult(combined, observed);
    CHECK(Difference(observed, expected) < 1.e-13);
    schur.Mult(combined, combined);
    CHECK(Difference(combined, expected) < 1.e-13);
    schur.Mult(first, observed);
    CHECK(Difference(observed, a) == 0.0);
}

TEST_CASE("Surface Border Schur Reports Singular And Ill Conditioned Coarse Blocks Without Regularization",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(Matrix(3, {1, 0, 0, 0, 0, 0, 0, 0, 1}));
    Inverse inverse(operation, Matrix(3, {1, 0, 0, 0, 0, 0, 0, 0, 0}));
    Schur singular(operation, inverse, {{1, 2}, {1, 2}}, MPI_COMM_WORLD);
    CHECK(singular.GetStatistics().status == probe::SurfaceBorderSchurStatus::singular);
    CHECK_FALSE(singular.IsCurrent());
    auto rhs = Vector({1, 2, 3});
    mfem::Vector action(3);
    CHECK_THROWS_AS(singular.Mult(rhs, action), std::logic_error);
    operation.matrix(1, 1) = 1.e-16;
    ++operation.generation;
    inverse.generation = operation.generation;
    singular.Refresh();
    CHECK(singular.GetStatistics().status == probe::SurfaceBorderSchurStatus::ill_conditioned);
    CHECK(singular.GetStatistics().reciprocalCondition == Catch::Approx(1.e-16));
    CHECK_FALSE(singular.IsCurrent());
    CHECK_THROWS_AS(singular.SetOperator(operation), std::logic_error);
    CHECK(singular.GetSchurMatrix()(0, 0) == 1.e-16);
}

TEST_CASE("Surface Border Schur Handles Pivoting And Uniform Matrix Scale",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    for (double scale : {1.e-100, 1.0, 1.e100}) {
        auto matrix = Matrix(3, {2, 0, 0, 0, 0, 2, 0, 3, 1});
        matrix *= scale;
        Operation operation(std::move(matrix));
        auto base = Matrix(3, {0.5, 0, 0, 0, 0, 0, 0, 0, 0});
        base *= 1.0 / scale;
        Inverse inverse(operation, std::move(base));
        Schur schur(operation, inverse, {{1, 2}, {1, 2}}, MPI_COMM_WORLD);
        REQUIRE(schur.IsCurrent());
        auto rhs = Vector({1, 2, 4});
        rhs *= scale;
        mfem::Vector action(3);
        schur.Mult(rhs, action);
        CHECK(Difference(action, Vector({0.5, 1, 1})) < 1.e-13);
        CHECK(schur.GetStatistics().reciprocalCondition == Catch::Approx(0.5));
    }
}

TEST_CASE("Surface Border Schur Requires Unbordered Fields For A Zero Scalar Diagonal",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(Matrix(2, {1, 1, 1, 0}));
    Inverse alreadyBordered(operation, Matrix(2, {0, 1, 1, -1}));
    Schur wrongRestriction(operation, alreadyBordered, {{1}, {1}}, MPI_COMM_WORLD);
    CHECK(wrongRestriction.GetStatistics().status == probe::SurfaceBorderSchurStatus::singular);
    Inverse unbordered(operation, Matrix(2, {1, 0, 0, 0}));
    Schur schur(operation, unbordered, {{1}, {1}}, MPI_COMM_WORLD);
    REQUIRE(schur.IsCurrent());
    auto rhs = Vector({3, 2});
    mfem::Vector action(2);
    schur.Mult(rhs, action);
    CHECK(Difference(action, Vector({2, 1})) < 1.e-13);
}

TEST_CASE("Surface Border Schur Rejects Invalid Partitions And Stale Or Nonfinite Dependencies",
          "[solver][surface_border_schur][serial]") {
    using namespace surface_border_schur_test;
    Operation operation(CoupledMatrix());
    Inverse inverse(operation, FieldInverse(true));
    const auto reject = [&](probe::SurfaceBorderSchurCoordinates coordinates) {
        CHECK_THROWS_AS(Schur(operation, inverse, std::move(coordinates), MPI_COMM_WORLD), std::invalid_argument);
    };
    reject({{}, {}});
    reject({{1}, {0, 2}});
    reject({{1, 1}, {0, 2}});
    reject({{1, 3}, {0, 0}});
    reject({{1, 4}, {0, 2}});
    reject({{0, 1, 2, 3}, {0, 1, 2, 3}});
    CHECK_THROWS_AS(Schur(operation, inverse, Coordinates(), MPI_COMM_NULL), std::invalid_argument);
    Schur schur(operation, inverse, Coordinates(), MPI_COMM_WORLD);
    auto rhs = Vector({1, 2, 3, 4});
    mfem::Vector action(4);
    Operation other(CoupledMatrix());
    CHECK_THROWS_AS(schur.SetOperator(other), std::invalid_argument);
    CHECK_THROWS_AS(Schur(other, inverse, Coordinates(), MPI_COMM_WORLD), std::invalid_argument);
    ++operation.generation;
    CHECK_FALSE(schur.IsCurrent());
    CHECK_THROWS_AS(schur.Mult(rhs, action), std::logic_error);
    inverse.generation = operation.generation;
    CHECK_FALSE(schur.IsCurrent());
    schur.Refresh();
    REQUIRE(schur.IsCurrent());
    inverse.iterative_mode = true;
    CHECK_FALSE(schur.IsCurrent());
    CHECK_THROWS_AS(schur.Mult(rhs, action), std::logic_error);
    inverse.iterative_mode = false;
    inverse.nonfinite = true;
    CHECK_THROWS_AS(schur.Mult(rhs, action), std::runtime_error);
    inverse.nonfinite = false;
    rhs(0) = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS_AS(schur.Mult(rhs, action), std::runtime_error);
}