#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <numbers>
#include <stdexcept>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace {
    namespace backend           = mean_field::preconditioning::backend;
    namespace blocks            = mean_field::utils::blocks;
    namespace preconditioning   = mean_field::preconditioning;

    using BaseModel             = mean_field::operators::StellarEquilibriumSpecificationModel;
    using CentralModel          = mean_field::operators::CentralDensityStellarEquilibriumSpecificationModel;
    using ReorderedCentralModel = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::models::FixedCentralDensity,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass,
        mean_field::eos::Polytrope>>;
    using BaseProblem           = mean_field::equilibrium::StellarEquilibriumProblem<BaseModel>;
    using CentralProblem        = mean_field::equilibrium::StellarEquilibriumProblem<CentralModel>;
    using BaseBorder            = preconditioning::CompiledSpecificationBorderFor<BaseModel>;
    using CentralBorder         = preconditioning::CompiledSpecificationBorderFor<CentralModel>;
    using BaseComponent = decltype(preconditioning::specificationBorderBlock(std::declval<const BaseProblem &>()));
    using CentralComponent =
        decltype(preconditioning::specificationBorderBlock(std::declval<const CentralProblem &>()));
    using BasePlan    = preconditioning::PreconditionerPlan<BaseComponent>;
    using CentralPlan = preconditioning::PreconditionerPlan<CentralComponent>;

    class KnownBorderCouplings final {
    public:
        explicit KnownBorderCouplings(const int borderSize)
            : m_borderSize(borderSize),
              m_structureToBorder(
                  borderSize,
                  StructureSize()
              ),
              m_borderToStructure(
                  StructureSize(),
                  borderSize
              ),
              m_borderDiagonal(borderSize) {
            if (borderSize <= 0) {
                throw std::invalid_argument("The known border must have positive size.");
            }
            for (int row = 0; row < borderSize; ++row) {
                for (int column = 0; column < StructureSize(); ++column) {
                    m_structureToBorder(row, column) = 0.04 * static_cast<double>((row + 1) * (column + 2));
                    m_borderToStructure(column, row) = -0.03 * static_cast<double>((column + 1) * (row + 2));
                }
                for (int column = 0; column < borderSize; ++column) {
                    m_borderDiagonal(row, column) =
                        row == column ? 2.0 + static_cast<double>(row) : 0.01 * static_cast<double>(row + column + 1);
                }
            }
        }

        [[nodiscard]] static constexpr int StructureSize() noexcept {
            return 3;
        }

        [[nodiscard]] int BorderSize() const noexcept {
            return m_borderSize;
        }

        void ApplyStructureToBorder(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const {
            m_structureToBorder.Mult(direction, action);
        }

        void ApplyBorderToStructure(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const {
            m_borderToStructure.Mult(direction, action);
        }

        void ApplyBorderToBorder(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const {
            m_borderDiagonal.Mult(direction, action);
        }

        void IncreaseBorderDiagonal(const double increment) {
            for (int index = 0; index < m_borderSize; ++index) {
                m_borderDiagonal(index, index) += increment;
            }
        }

        [[nodiscard]] const mfem::DenseMatrix &StructureToBorder() const noexcept {
            return m_structureToBorder;
        }

        [[nodiscard]] const mfem::DenseMatrix &BorderToStructure() const noexcept {
            return m_borderToStructure;
        }

        [[nodiscard]] const mfem::DenseMatrix &BorderDiagonal() const noexcept {
            return m_borderDiagonal;
        }

    private:
        int m_borderSize;
        mfem::DenseMatrix m_structureToBorder;
        mfem::DenseMatrix m_borderToStructure;
        mfem::DenseMatrix m_borderDiagonal;
    };

    [[nodiscard]] double relativeError(
        const mfem::Vector &left,
        const mfem::Vector &right
    ) {
        mfem::Vector difference(left);
        difference -= right;
        return difference.Norml2() / std::max({1.0, left.Norml2(), right.Norml2()});
    }

    template <
        preconditioning::ApplicationContract StructureInverseContract =
            preconditioning::ApplicationContract::stationary_linear>
    void verifyKnownBorderFactorization(const int borderSize) {
        mfem::Vector structureDiagonal(KnownBorderCouplings::StructureSize());
        structureDiagonal(0)  = 2.0;
        structureDiagonal(1)  = 3.0;
        structureDiagonal(2)  = 5.0;
        auto structureInverse = backend::prepare(backend::Diagonal{}, structureDiagonal);
        KnownBorderCouplings couplings(borderSize);
        using Factorization =
            preconditioning::SpecificationBorderFactorizationOperator<KnownBorderCouplings, StructureInverseContract>;
        Factorization factorization(structureInverse, couplings);
        constexpr bool cachesStructureResponse = Factorization::cachesStructureInverseBorderCoupling;

        const auto expectedSchurEntry          = [&](const int row, const int column) {
            double correction = 0.0;
            for (int inner = 0; inner < KnownBorderCouplings::StructureSize(); ++inner) {
                correction += couplings.StructureToBorder()(row, inner) * couplings.BorderToStructure()(inner, column) /
                              structureDiagonal(inner);
            }
            return couplings.BorderDiagonal()(row, column) - correction;
        };
        for (int row = 0; row < borderSize; ++row) {
            for (int column = 0; column < borderSize; ++column) {
                CHECK(
                    factorization.GetSchurComplement()(row, column) ==
                    Catch::Approx(expectedSchurEntry(row, column)).margin(2.0e-14)
                );
            }
        }

        const int completeSize = KnownBorderCouplings::StructureSize() + borderSize;
        mfem::DenseMatrix completeMatrix(completeSize);
        completeMatrix = 0.0;
        for (int index = 0; index < KnownBorderCouplings::StructureSize(); ++index) {
            completeMatrix(index, index) = structureDiagonal(index);
        }
        for (int row = 0; row < KnownBorderCouplings::StructureSize(); ++row) {
            for (int column = 0; column < borderSize; ++column) {
                completeMatrix(row, KnownBorderCouplings::StructureSize() + column) =
                    couplings.BorderToStructure()(row, column);
                completeMatrix(KnownBorderCouplings::StructureSize() + column, row) =
                    couplings.StructureToBorder()(column, row);
            }
        }
        for (int row = 0; row < borderSize; ++row) {
            for (int column = 0; column < borderSize; ++column) {
                completeMatrix(
                    KnownBorderCouplings::StructureSize() + row, KnownBorderCouplings::StructureSize() + column
                ) = couplings.BorderDiagonal()(row, column);
            }
        }

        mfem::Vector rightHandSide(completeSize);
        for (int index = 0; index < completeSize; ++index) {
            rightHandSide(index) = 0.25 + 0.17 * static_cast<double>(index + 1);
        }
        mfem::Vector actual(completeSize);
        mfem::Vector expected(completeSize);
        factorization.Mult(rightHandSide, actual);
        mfem::DenseMatrixInverse exactInverse(completeMatrix);
        exactInverse.Mult(rightHandSide, expected);
        CHECK(relativeError(actual, expected) <= 2.0e-13);

        const auto statisticsBeforeRefresh = factorization.GetStatistics();
        CHECK(statisticsBeforeRefresh.setups == 1);
        CHECK(statisticsBeforeRefresh.schurProbes == static_cast<std::uint64_t>(borderSize));
        CHECK(statisticsBeforeRefresh.applications == 1);
        CHECK(
            statisticsBeforeRefresh.structureInverseApplications ==
            static_cast<std::uint64_t>(borderSize + (cachesStructureResponse ? 1 : 2))
        );
        CHECK(
            statisticsBeforeRefresh.cachedStructureInverseBorderApplications ==
            static_cast<std::uint64_t>(cachesStructureResponse ? 1 : 0)
        );
        CHECK(statisticsBeforeRefresh.structureToBorderApplications == static_cast<std::uint64_t>(borderSize + 1));
        CHECK(
            statisticsBeforeRefresh.borderToStructureApplications ==
            static_cast<std::uint64_t>(borderSize + (cachesStructureResponse ? 0 : 1))
        );
        CHECK(statisticsBeforeRefresh.borderToBorderApplications == static_cast<std::uint64_t>(borderSize));
        CHECK(
            structureInverse.GetStatistics().applications ==
            static_cast<std::uint64_t>(borderSize + (cachesStructureResponse ? 1 : 2))
        );

        for (int index = 0; index < KnownBorderCouplings::StructureSize(); ++index) {
            structureDiagonal(index) += 0.25 * static_cast<double>(index + 1);
            completeMatrix(index, index) = structureDiagonal(index);
        }
        structureInverse.Refresh(structureDiagonal);
        couplings.IncreaseBorderDiagonal(0.5);
        for (int index = 0; index < borderSize; ++index) {
            completeMatrix(
                KnownBorderCouplings::StructureSize() + index, KnownBorderCouplings::StructureSize() + index
            ) += 0.5;
        }
        factorization.RefreshSchurComplement();
        CHECK(factorization.GetStatistics().setups == 2);
        CHECK(factorization.GetStatistics().schurProbes == static_cast<std::uint64_t>(2 * borderSize));
        CHECK(factorization.GetStatistics().borderToBorderApplications == static_cast<std::uint64_t>(2 * borderSize));
        CHECK(
            factorization.GetStatistics().structureInverseApplications ==
            static_cast<std::uint64_t>(2 * borderSize + (cachesStructureResponse ? 1 : 2))
        );
        CHECK(
            factorization.GetStatistics().cachedStructureInverseBorderApplications ==
            static_cast<std::uint64_t>(cachesStructureResponse ? 1 : 0)
        );
        CHECK(
            factorization.GetStatistics().structureToBorderApplications ==
            static_cast<std::uint64_t>(2 * borderSize + 1)
        );
        CHECK(
            factorization.GetStatistics().borderToStructureApplications ==
            static_cast<std::uint64_t>(2 * borderSize + (cachesStructureResponse ? 0 : 1))
        );
        for (int row = 0; row < borderSize; ++row) {
            for (int column = 0; column < borderSize; ++column) {
                CHECK(
                    factorization.GetSchurComplement()(row, column) ==
                    Catch::Approx(expectedSchurEntry(row, column)).margin(2.0e-14)
                );
            }
        }

        mfem::Vector refreshedActual(completeSize);
        mfem::Vector refreshedExpected(completeSize);
        factorization.Mult(rightHandSide, refreshedActual);
        mfem::DenseMatrixInverse refreshedExactInverse(completeMatrix);
        refreshedExactInverse.Mult(rightHandSide, refreshedExpected);
        CHECK(relativeError(refreshedActual, refreshedExpected) <= 2.0e-13);

        const auto statisticsAfterRefreshApplication = factorization.GetStatistics();
        CHECK(statisticsAfterRefreshApplication.applications == 2);
        CHECK(
            statisticsAfterRefreshApplication.structureInverseApplications ==
            static_cast<std::uint64_t>(2 * borderSize + (cachesStructureResponse ? 2 : 4))
        );
        CHECK(
            statisticsAfterRefreshApplication.cachedStructureInverseBorderApplications ==
            static_cast<std::uint64_t>(cachesStructureResponse ? 2 : 0)
        );
        CHECK(
            statisticsAfterRefreshApplication.structureToBorderApplications ==
            static_cast<std::uint64_t>(2 * borderSize + 2)
        );
        CHECK(
            statisticsAfterRefreshApplication.borderToStructureApplications ==
            static_cast<std::uint64_t>(2 * borderSize + (cachesStructureResponse ? 0 : 2))
        );
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies
    makeDependencies(const std::uint64_t revision = 1) {
        return {
            .discretization     = {.identity = 9201, .revision = 1},
            .density            = {.identity = 9203, .revision = revision},
            .surfaceDeformation = {.identity = 9207, .revision = revision},
            .gravityGradient    = {.identity = 9211, .revision = revision},
            .gravityPotential   = {.identity = 9217, .revision = revision},
            .enthalpy           = {.identity = 9223, .revision = revision},
            .bernoulliConstant  = {.identity = 9229, .revision = revision},
            .rotation           = {.identity = 9231, .revision = revision},
            .targetMass         = {.identity = 9237, .revision = 1}
        };
    }

    [[nodiscard]] mean_field::physics::RigidRotation zeroRotation() {
        mfem::Vector angularVelocity(3);
        mfem::Vector center(3);
        angularVelocity = 0.0;
        center          = 0.0;
        return {angularVelocity, center};
    }

    template <
        typename View,
        typename Term>
    void assignStateBlock(
        const View &view,
        const Term &term,
        const mfem::Vector &source,
        mfem::Vector &state
    ) {
        mfem::Vector destination = view.block(term);
        REQUIRE(destination.Size() == source.Size());
        destination = source;
        destination.SyncAliasMemory(state);
    }
} // namespace

TEST_CASE(
    "Model Specifications Compile Complete Canonical Preconditioning Borders",
    "[preconditioning][specification_border][unit][type_contract]"
) {
    using ExpectedBaseCorrections    = blocks::type_list<blocks::fixed_total_mass::mass_normalization::value>;
    using ExpectedBaseResiduals      = blocks::type_list<blocks::fixed_total_mass::mass_normalization::residual>;
    using ExpectedCentralCorrections = blocks::type_list<
        blocks::fixed_total_mass::mass_normalization::value, blocks::fixed_central_density::central_value::value>;
    using ExpectedCentralResiduals = blocks::type_list<
        blocks::fixed_total_mass::mass_normalization::residual, blocks::fixed_central_density::central_value::residual>;

    STATIC_CHECK(std::same_as<CentralModel, ReorderedCentralModel>);
    STATIC_CHECK(BaseBorder::valueArity == 1);
    STATIC_CHECK(BaseBorder::residualArity == 1);
    STATIC_CHECK(BaseBorder::specificationCount == 1);
    STATIC_CHECK(std::same_as<typename BaseBorder::CorrectionBlocks, ExpectedBaseCorrections>);
    STATIC_CHECK(std::same_as<typename BaseBorder::ResidualBlocks, ExpectedBaseResiduals>);
    STATIC_CHECK(BaseBorder::RequiredCouplings::size == 3);

    STATIC_CHECK(CentralBorder::valueArity == 2);
    STATIC_CHECK(CentralBorder::residualArity == 2);
    STATIC_CHECK(CentralBorder::specificationCount == 2);
    STATIC_CHECK(std::same_as<typename CentralBorder::CorrectionBlocks, ExpectedCentralCorrections>);
    STATIC_CHECK(std::same_as<typename CentralBorder::ResidualBlocks, ExpectedCentralResiduals>);
    STATIC_CHECK(CentralBorder::RequiredCouplings::size == 5);
    STATIC_CHECK(
        preconditioning::specificationBorderValueOffset<mean_field::models::FixedTotalMass, CentralModel> == 0
    );
    STATIC_CHECK(
        preconditioning::specificationBorderValueOffset<mean_field::models::FixedCentralDensity, CentralModel> == 1
    );
    STATIC_CHECK(
        preconditioning::specificationBorderResidualOffset<mean_field::models::FixedTotalMass, CentralModel> == 0
    );
    STATIC_CHECK(
        preconditioning::specificationBorderResidualOffset<mean_field::models::FixedCentralDensity, CentralModel> == 1
    );

    STATIC_CHECK(preconditioning::PreconditionerComponent<BaseComponent>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<CentralComponent>);
    STATIC_CHECK(BaseComponent::RequiredCouplings::size == 19);
    STATIC_CHECK(CentralComponent::RequiredCouplings::size == 21);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<BasePlan, typename BaseProblem::FormType>);
    STATIC_CHECK(
        preconditioning::CompatiblePreconditionerFor<
            BasePlan, typename BaseProblem::FormType, typename BaseProblem::JacobianFormType>
    );
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<CentralPlan, typename CentralProblem::FormType>);
    STATIC_CHECK(
        preconditioning::CompatiblePreconditionerFor<
            CentralPlan, typename CentralProblem::FormType, typename CentralProblem::JacobianFormType>
    );
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename CentralComponent::BackendType>);
}

TEST_CASE(
    "Dense Specification Borders Cache Stationary Structure Responses And Reproduce Exact Block Factorizations",
    "[preconditioning][specification_border][unit][factorization]"
) {
    SECTION("one generated scalar") {
        verifyKnownBorderFactorization(1);
    }
    SECTION("two generated scalars") {
        verifyKnownBorderFactorization(2);
    }
    SECTION("four generated scalars") {
        verifyKnownBorderFactorization(4);
    }
}

TEST_CASE(
    "Flexible Specification Borders Preserve Per-Application Structure Solves",
    "[preconditioning][specification_border][unit][factorization]"
) {
    verifyKnownBorderFactorization<preconditioning::ApplicationContract::flexible>(2);
}

TEST_CASE(
    "Generated Specification Border Actions Match The Authoritative Stellar Jacobian",
    "[preconditioning][specification_border][integration]"
) {
    using namespace mean_field;
    const utils::Args arguments = test_utils::setup_args();
    fem::FEM finiteElements     = fem::setup_fem(arguments.mesh_file, arguments, 0);
    REQUIRE(finiteElements.okay());

    constexpr double radius         = utils::RADIUS;
    constexpr double mass           = utils::MASS;
    const double polytropicConstant = 2.0 * utils::G * radius * radius / std::numbers::pi_v<double>;
    const double centralDensity     = std::numbers::pi_v<double> * mass / (4.0 * radius * radius * radius);
    const auto stellarModel         = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{mass}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
    );
    auto problem   = equilibrium::discretize(stellarModel, finiteElements);
    auto projected = seed::makeProjectedEquilibriumState(problem, seed::LaneEmden({.radialSampleCount = 512}));
    problem.Prepare(projected.values, makeDependencies(), zeroRotation());

    preconditioning::SpecificationBorderJacobianOperator coupling(problem);
    REQUIRE(coupling.BorderSize() == 2);
    REQUIRE(coupling.StructureSize() + coupling.BorderSize() == problem.StateSize());
    const auto &offsets = coupling.GetStructureOffsets();

    mfem::Vector groupedDirection(coupling.Width());
    for (int index = 0; index < groupedDirection.Size(); ++index) {
        groupedDirection(index) = 0.015 * std::sin(0.23 * static_cast<double>(index + 1));
    }
    const auto groupedBlock = [&](const int block) {
        return mfem::Vector(groupedDirection.GetData() + offsets[block], offsets[block + 1] - offsets[block]);
    };

    mfem::Vector structureOnlyRoot(problem.StateSize());
    structureOnlyRoot        = 0.0;
    const auto structureView = problem.GetManifest().directionView(structureOnlyRoot);
    assignStateBlock(structureView, blocks::density_field.mass_term, groupedBlock(0), structureOnlyRoot);
    assignStateBlock(
        structureView, blocks::surface_deformation_field.parameters_term, groupedBlock(1), structureOnlyRoot
    );
    assignStateBlock(structureView, blocks::enthalpy_field.specific_term, groupedBlock(2), structureOnlyRoot);
    assignStateBlock(structureView, blocks::gravity_field.gradient_term, groupedBlock(3), structureOnlyRoot);
    assignStateBlock(structureView, blocks::gravity_field.poisson_term, groupedBlock(4), structureOnlyRoot);

    mfem::Vector borderOnlyRoot(problem.StateSize());
    borderOnlyRoot        = 0.0;
    const auto borderView = problem.GetManifest().directionView(borderOnlyRoot);
    mfem::Vector massDirection(groupedDirection.GetData() + coupling.StructureSize(), 1);
    mfem::Vector centralDirection(groupedDirection.GetData() + coupling.StructureSize() + 1, 1);
    assignStateBlock(
        borderView, blocks::fixed_total_mass_constraint.mass_normalization_term, massDirection, borderOnlyRoot
    );
    assignStateBlock(
        borderView, blocks::fixed_central_density_phase.central_value_term, centralDirection, borderOnlyRoot
    );

    mfem::Vector structureOnlyAction;
    mfem::Vector borderOnlyAction;
    problem.ApplyLinearization(structureOnlyRoot, structureOnlyAction);
    problem.ApplyLinearization(borderOnlyRoot, borderOnlyAction);
    auto structureOnlyResidual = problem.GetManifest().residualView(structureOnlyAction);
    auto borderOnlyResidual    = problem.GetManifest().residualView(borderOnlyAction);

    mfem::Vector expected(coupling.Height());
    expected = 0.0;
    expected.SetVector(borderOnlyResidual.block(blocks::density_field.mass_term), offsets[0]);
    expected.SetVector(borderOnlyResidual.block(blocks::surface_deformation_field.shape_equilibrium_term), offsets[1]);
    expected.SetVector(borderOnlyResidual.block(blocks::enthalpy_field.specific_term), offsets[2]);
    expected.SetVector(borderOnlyResidual.block(blocks::gravity_field.gradient_term), offsets[3]);
    expected.SetVector(borderOnlyResidual.block(blocks::gravity_field.poisson_term), offsets[4]);
    expected.SetVector(
        structureOnlyResidual.block(blocks::fixed_total_mass_constraint.mass_normalization_term),
        coupling.StructureSize()
    );
    expected.SetVector(
        structureOnlyResidual.block(blocks::fixed_central_density_phase.central_value_term),
        coupling.StructureSize() + 1
    );
    mfem::Vector borderDiagonal(2);
    borderDiagonal(0) = borderOnlyResidual.block(blocks::fixed_total_mass_constraint.mass_normalization_term)(0);
    borderDiagonal(1) = borderOnlyResidual.block(blocks::fixed_central_density_phase.central_value_term)(0);
    mfem::Vector expectedBorder(expected, coupling.StructureSize(), coupling.BorderSize());
    expectedBorder += borderDiagonal;
    expectedBorder.SyncAliasMemory(expected);

    mfem::Vector actual(coupling.Height());
    coupling.Mult(groupedDirection, actual);
    CHECK(relativeError(actual, expected) <= 2.0e-12);

    auto component  = preconditioning::makePreconditioner(problem);
    using Component = decltype(component);
    STATIC_CHECK(std::same_as<Component, CentralComponent>);
    auto prepared               = preconditioning::prepare(problem, component);
    using GroupedPreconditioner = typename decltype(prepared)::GroupedPreconditioner;
    using PreparedFactorization = typename GroupedPreconditioner::Factorization;
    STATIC_CHECK(PreparedFactorization::cachesStructureInverseBorderCoupling);
    mfem::Vector rightHandSide(prepared.Width());
    for (int index = 0; index < rightHandSide.Size(); ++index) {
        rightHandSide(index) = std::cos(0.11 * static_cast<double>(index + 1));
    }
    mfem::Vector correction(prepared.Height());
    prepared.Mult(rightHandSide, correction);
    for (int index = 0; index < correction.Size(); ++index) {
        REQUIRE(std::isfinite(correction(index)));
    }
    const auto &factorizationStatistics = prepared.GetGroupedPreconditioner().GetFactorization().GetStatistics();
    CHECK(factorizationStatistics.setups == 1);
    CHECK(factorizationStatistics.schurProbes == 2);
    CHECK(factorizationStatistics.applications == 1);
    CHECK(factorizationStatistics.structureInverseApplications == 3);
    CHECK(factorizationStatistics.cachedStructureInverseBorderApplications == 1);
    CHECK(factorizationStatistics.borderToStructureApplications == 2);
    const auto unchanged = prepared.Refresh();
    CHECK_FALSE(unchanged.DidAnyWork());
    CHECK(prepared.IsCurrent());
}