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

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

import mean_field;
import test_helpers;

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

    using Form                = blocks::central_density_bordered_stellar_equilibrium_form;
    using GroupedComponent    = preconditioning::ComponentDeclaration<
        blocks::type_list<
            blocks::density::mass::value,
            blocks::surface_deformation::parameters::value,
            blocks::enthalpy::specific::value,
            blocks::gravity::gradient::value,
            blocks::gravity::poisson::value,
            blocks::fixed_total_mass::mass_normalization::value,
            blocks::fixed_central_density::central_value::value>,
        blocks::type_list<
            blocks::density::mass::residual,
            blocks::surface_deformation::shape_equilibrium::residual,
            blocks::enthalpy::specific::residual,
            blocks::gravity::gradient::residual,
            blocks::gravity::poisson::residual,
            blocks::fixed_total_mass::mass_normalization::residual,
            blocks::fixed_central_density::central_value::residual>,
        blocks::type_list<>,
        preconditioning::IdentityOperatorCharacteristics,
        preconditioning::backend::Identity>;
    using IncompleteComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<
            blocks::density::mass::value,
            blocks::surface_deformation::parameters::value,
            blocks::enthalpy::specific::value,
            blocks::gravity::gradient::value,
            blocks::gravity::poisson::value,
            blocks::fixed_total_mass::mass_normalization::value>,
        blocks::type_list<
            blocks::density::mass::residual,
            blocks::surface_deformation::shape_equilibrium::residual,
            blocks::enthalpy::specific::residual,
            blocks::gravity::gradient::residual,
            blocks::gravity::poisson::residual,
            blocks::fixed_total_mass::mass_normalization::residual>,
        blocks::type_list<>,
        preconditioning::IdentityOperatorCharacteristics,
        preconditioning::backend::Identity>;
    using LifetimeModel    = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::models::FixedTotalMass,
        mean_field::models::FixedCentralDensity>>;
    using LifetimeProblem  = mean_field::equilibrium::StellarEquilibriumProblem<LifetimeModel>;
    using LifetimeBlock    = decltype(preconditioning::makePreconditioner(std::declval<const LifetimeProblem &>()));
    using LifetimePrepared = preconditioning::PreparedStellarPreconditioner<LifetimeProblem, LifetimeBlock>;

    template <typename Problem, typename Block>
    concept CanPrepareStellarPreconditioner =
        requires(const Problem &problem, Block block) { preconditioning::prepare(problem, std::move(block)); };

    template <typename Problem, typename Block>
    concept CanPrepareStellarPreconditionerFromTemporary =
        requires(Block block) { preconditioning::prepare(std::declval<Problem &&>(), std::move(block)); };

    template <typename Problem, typename Block>
    concept CanPrepareStellarPreconditionerFromConstTemporary =
        requires(Block block) { preconditioning::prepare(std::declval<const Problem &&>(), std::move(block)); };

    [[nodiscard]] blocks::form_layout<Form> makeUnevenLayout() {
        return {
            std::array<int, Form::value_block_count>{2, 3, 4, 5, 6, 1, 1},
            std::array<int, Form::residual_block_count>{4, 5, 2, 3, 6, 1, 1}
        };
    }

    [[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()});
    }

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies makeDependencies() {
        return {
            .discretization     = {.identity = 9301, .revision = 1},
            .density            = {.identity = 9303, .revision = 1},
            .surfaceDeformation = {.identity = 9307, .revision = 1},
            .gravityGradient    = {.identity = 9311, .revision = 1},
            .gravityPotential   = {.identity = 9317, .revision = 1},
            .enthalpy           = {.identity = 9323, .revision = 1},
            .bernoulliConstant  = {.identity = 9329, .revision = 1},
            .rotation           = {.identity = 9331, .revision = 1},
            .targetMass         = {.identity = 9337, .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};
    }
} // namespace

TEST_CASE(
    "Typed Equilibrium Coordinate Maps Preserve Every Uneven Block Without Scaling",
    "[preconditioning][equilibrium_coordinates][unit]"
) {
    STATIC_CHECK(preconditioning::EquilibriumCoordinateComponentFor<GroupedComponent, Form>);
    STATIC_CHECK_FALSE(preconditioning::EquilibriumCoordinateComponentFor<IncompleteComponent, Form>);
    STATIC_CHECK(CanPrepareStellarPreconditioner<LifetimeProblem, LifetimeBlock>);
    STATIC_CHECK_FALSE(CanPrepareStellarPreconditionerFromTemporary<LifetimeProblem, LifetimeBlock>);
    STATIC_CHECK_FALSE(CanPrepareStellarPreconditionerFromConstTemporary<LifetimeProblem, LifetimeBlock>);
    STATIC_CHECK(std::constructible_from<LifetimePrepared, const LifetimeProblem &, LifetimeBlock>);
    STATIC_CHECK_FALSE(std::constructible_from<LifetimePrepared, LifetimeProblem &&, LifetimeBlock>);
    STATIC_CHECK_FALSE(std::constructible_from<LifetimePrepared, const LifetimeProblem &&, LifetimeBlock>);

    const auto layout = makeUnevenLayout();
    preconditioning::EquilibriumPreconditionerCoordinateMap<Form, GroupedComponent> coordinates(layout);

    REQUIRE(coordinates.EquilibriumStateSize() == 22);
    REQUIRE(coordinates.EquilibriumResidualSize() == 22);
    REQUIRE(coordinates.PreconditionerCorrectionSize() == 22);
    REQUIRE(coordinates.PreconditionerResidualSize() == 22);

    const auto &correctionRanges = coordinates.GetCorrectionRanges();
    CHECK(correctionRanges[0] == (preconditioning::EquilibriumCoordinateRange{0, 0, 2}));
    CHECK(correctionRanges[1] == (preconditioning::EquilibriumCoordinateRange{2, 2, 3}));
    CHECK(correctionRanges[2] == (preconditioning::EquilibriumCoordinateRange{14, 5, 6}));
    CHECK(correctionRanges[3] == (preconditioning::EquilibriumCoordinateRange{5, 11, 4}));
    CHECK(correctionRanges[4] == (preconditioning::EquilibriumCoordinateRange{9, 15, 5}));
    CHECK(correctionRanges[5] == (preconditioning::EquilibriumCoordinateRange{20, 20, 1}));
    CHECK(correctionRanges[6] == (preconditioning::EquilibriumCoordinateRange{21, 21, 1}));

    const auto &residualRanges = coordinates.GetResidualRanges();
    CHECK(residualRanges[0] == (preconditioning::EquilibriumCoordinateRange{9, 0, 2}));
    CHECK(residualRanges[1] == (preconditioning::EquilibriumCoordinateRange{11, 2, 3}));
    CHECK(residualRanges[2] == (preconditioning::EquilibriumCoordinateRange{14, 5, 6}));
    CHECK(residualRanges[3] == (preconditioning::EquilibriumCoordinateRange{0, 11, 4}));
    CHECK(residualRanges[4] == (preconditioning::EquilibriumCoordinateRange{4, 15, 5}));
    CHECK(residualRanges[5] == (preconditioning::EquilibriumCoordinateRange{20, 20, 1}));
    CHECK(residualRanges[6] == (preconditioning::EquilibriumCoordinateRange{21, 21, 1}));

    mfem::Vector equilibriumCorrection(22);
    mfem::Vector equilibriumResidual(22);
    for (int index = 0; index < 22; ++index) {
        equilibriumCorrection(index) = 100.0 + static_cast<double>(index);
        equilibriumResidual(index)   = -200.0 - static_cast<double>(index);
    }

    mfem::Vector groupedCorrection(22);
    mfem::Vector groupedResidual(22);
    const double *const groupedCorrectionStorage = groupedCorrection.GetData();
    const double *const groupedResidualStorage   = groupedResidual.GetData();
    coordinates.PackCorrection(equilibriumCorrection, groupedCorrection);
    coordinates.PackResidual(equilibriumResidual, groupedResidual);
    CHECK(groupedCorrection.GetData() == groupedCorrectionStorage);
    CHECK(groupedResidual.GetData() == groupedResidualStorage);
    CHECK(groupedCorrection(5) == equilibriumCorrection(14));
    CHECK(groupedCorrection(11) == equilibriumCorrection(5));
    CHECK(groupedResidual(0) == equilibriumResidual(9));
    CHECK(groupedResidual(11) == equilibriumResidual(0));

    mfem::Vector recoveredCorrection(22);
    mfem::Vector recoveredResidual(22);
    coordinates.UnpackCorrection(groupedCorrection, recoveredCorrection);
    coordinates.UnpackResidual(groupedResidual, recoveredResidual);
    CHECK(relativeError(recoveredCorrection, equilibriumCorrection) == 0.0);
    CHECK(relativeError(recoveredResidual, equilibriumResidual) == 0.0);

    const auto &statistics = coordinates.GetStatistics();
    CHECK(statistics.correctionPacks == 1);
    CHECK(statistics.correctionUnpacks == 1);
    CHECK(statistics.residualPacks == 1);
    CHECK(statistics.residualUnpacks == 1);

    mfem::Vector wrongSize(21);
    CHECK_THROWS_AS(coordinates.PackResidual(wrongSize, groupedResidual), std::invalid_argument);
    CHECK_THROWS_AS(coordinates.UnpackCorrection(wrongSize, recoveredCorrection), std::invalid_argument);
}

TEST_CASE(
    "Prepared Stellar Preconditioning Matches An Explicit Canonical Coordinate Transformation",
    "[preconditioning][equilibrium_coordinates][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);
    auto model                      = 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(model, std::move(finiteElements));
    auto projected = seed::makeProjectedEquilibriumState(problem, seed::LaneEmden({.radialSampleCount = 512}));
    problem.Prepare(projected.values, makeDependencies(), zeroRotation());

    auto component = preconditioning::makePreconditioner(problem);
    auto prepared  = preconditioning::prepare(problem, component);
    prepared.SetOperator(problem.GetLinearizationOperator());

    mfem::Vector equilibriumResidual(problem.EquationSize());
    for (int index = 0; index < equilibriumResidual.Size(); ++index) {
        equilibriumResidual(index) = 0.25 * std::cos(0.19 * static_cast<double>(index + 1));
    }

    mfem::Vector groupedResidual(problem.EquationSize());
    mfem::Vector groupedCorrection(problem.StateSize());
    mfem::Vector expected(problem.StateSize());
    prepared.GetCoordinateMap().PackResidual(equilibriumResidual, groupedResidual);
    prepared.GetGroupedPreconditioner().Mult(groupedResidual, groupedCorrection);
    prepared.GetCoordinateMap().UnpackCorrection(groupedCorrection, expected);

    mfem::Vector actual(problem.StateSize());
    const double *const actionStorage = actual.GetData();
    prepared.Mult(equilibriumResidual, actual);
    CHECK(actual.GetData() == actionStorage);
    CHECK(relativeError(actual, expected) <= 2.0e-12);

    const auto &statistics = prepared.GetStatistics();
    CHECK(statistics.applications == 1);
    CHECK(statistics.residualCoordinateMappings == 1);
    CHECK(statistics.correctionCoordinateMappings == 1);
    CHECK(prepared.GetCoordinateMap().GetStatistics().residualPacks == 2);
    CHECK(prepared.GetCoordinateMap().GetStatistics().correctionUnpacks == 2);

    const auto unchanged = prepared.Refresh();
    CHECK_FALSE(unchanged.DidAnyWork());
    CHECK(prepared.IsCurrent());
}