#include <algorithm>
#include <array>
#include <cmath>
#include <limits>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace rotational_displacement_force_test_utils {
    using CoupledForm = mean_field::utils::blocks::barotropic_equilibrium_form;

    constexpr auto densityValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::density_field.mass_term);

    constexpr auto displacementValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::displacement_field.geometry_term
    );

    constexpr auto gravityGradientValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::gravity_field.gradient_term);

    constexpr auto gravityPotentialValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::gravity_field.poisson_term);

    constexpr auto enthalpyValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::enthalpy_field.specific_term
    );

    constexpr auto barotropicConstantValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    constexpr auto gravityGradientResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::gravity_field.gradient_term
    );

    constexpr auto gravityPotentialResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::gravity_field.poisson_term
    );

    constexpr auto densityResidual =
        mean_field::utils::blocks::get_residual_block<CoupledForm>(mean_field::utils::blocks::density_field.mass_term);

    constexpr auto displacementResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::displacement_field.geometry_term
    );

    constexpr auto enthalpyResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::enthalpy_field.specific_term
    );

    constexpr auto massResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    [[nodiscard]] mean_field::operators::RotationalDisplacementForceLayout make_layout(const mean_field::fem::FEM &f) {
        const std::array<int, CoupledForm::value_block_count> valueSizes{
            f.densityFes->GetTrueVSize(),          f.displacementFes->GetTrueVSize(), f.gravityFluxFes->GetTrueVSize(),
            f.gravityPotentialFes->GetTrueVSize(), f.enthalpyFes->GetTrueVSize(),     1
        };

        const std::array<int, CoupledForm::residual_block_count> residualSizes{
            f.gravityFluxFes->GetTrueVSize(),  f.gravityPotentialFes->GetTrueVSize(), f.densityFes->GetTrueVSize(),
            f.displacementFes->GetTrueVSize(), f.enthalpyFes->GetTrueVSize(),         1
        };

        return {valueSizes, residualSizes};
    }

    [[nodiscard]] mfem::Vector make_density(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction densityField(f.densityFes.get());

        mfem::FunctionCoefficient densityCoefficient([phase](const mfem::Vector &position) {
            return 0.88 + 0.06 * std::sin(0.7 * position(0) + phase) + 0.04 * std::cos(0.6 * position(1) - phase) +
                   0.025 * position(2) * position(2);
        });

        densityField.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector make_density_direction(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction densityField(f.densityFes.get());

        mfem::FunctionCoefficient densityCoefficient([phase](const mfem::Vector &position) {
            return 0.17 * std::sin(0.9 * position(0) + phase) - 0.12 * std::cos(0.8 * position(1) - phase) +
                   0.07 * position(2);
        });

        densityField.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector make_displacement_direction(const mean_field::fem::FEM &f) {
        mfem::Vector direction    = gravity_prepared_test_utils::make_displacement(f, 0.91);

        const mfem::Vector second = gravity_prepared_test_utils::make_displacement(f, 0.27);

        direction -= second;
        return direction;
    }

    [[nodiscard]] mean_field::physics::RigidRotation make_rotation(const double scale = 1.0) {
        mfem::Vector angularVelocity(3);
        angularVelocity(0) = scale * 0.17;
        angularVelocity(1) = scale * -0.09;
        angularVelocity(2) = scale * 0.62;

        mfem::Vector center(3);
        center(0) = 0.04;
        center(1) = -0.03;
        center(2) = 0.02;

        return mean_field::physics::RigidRotation(angularVelocity, center);
    }

    [[nodiscard]] mean_field::operators::context::rotational_displacement_force::RotationalDisplacementForceDependencies
    make_dependencies() {
        return {
            .discretization = {.identity = 211, .revision = 3},
            .density        = {.identity = 223, .revision = 5},
            .displacement   = {.identity = 227, .revision = 7},
            .rotation       = {.identity = 229, .revision = 11}
        };
    }

    [[nodiscard]] mfem::Vector make_vacuum_only_density(const mean_field::fem::FEM &f) {
        mfem::ParGridFunction densityField(f.densityFes.get());
        densityField              = 0.0;

        const int vacuumAttribute = f.domainMapperStateless->GetVacuumElementAttribute();

        mfem::Array<int> densityDofs;
        int localVacuumElements = 0;

        for (int elementId = 0; elementId < f.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation = f.mesh->GetElementTransformation(elementId);

            REQUIRE(transformation != nullptr);

            if (transformation->Attribute != vacuumAttribute) {
                continue;
            }

            f.densityFes->GetElementDofs(elementId, densityDofs);

            mfem::Vector elementDensity(densityDofs.Size());
            elementDensity = 1.0;
            densityField.SetSubVector(densityDofs, elementDensity);
            ++localVacuumElements;
        }

        int globalVacuumElements = 0;

        MPI_Allreduce(&localVacuumElements, &globalVacuumElements, 1, MPI_INT, MPI_SUM, f.mesh->GetComm());

        REQUIRE(globalVacuumElements > 0);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] double global_norm(
        const mfem::Vector &vector,
        MPI_Comm communicator
    ) {
        const double localSquaredNorm = vector * vector;
        double globalSquaredNorm      = 0.0;

        MPI_Allreduce(&localSquaredNorm, &globalSquaredNorm, 1, MPI_DOUBLE, MPI_SUM, communicator);

        return std::sqrt(globalSquaredNorm);
    }

    [[nodiscard]] double global_dot(
        const mfem::Vector &left,
        const mfem::Vector &right,
        MPI_Comm communicator
    ) {
        REQUIRE(left.Size() == right.Size());

        const double localDot = left * right;
        double globalDot      = 0.0;

        MPI_Allreduce(&localDot, &globalDot, 1, MPI_DOUBLE, MPI_SUM, communicator);

        return globalDot;
    }

    [[nodiscard]] double relative_difference(
        const mfem::Vector &computed,
        const mfem::Vector &reference,
        MPI_Comm communicator
    ) {
        REQUIRE(computed.Size() == reference.Size());

        mfem::Vector difference(computed);
        difference -= reference;

        return global_norm(difference, communicator) /
               std::max(global_norm(reference, communicator), std::numeric_limits<double>::epsilon());
    }

    [[nodiscard]] mfem::Vector centered_difference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &baseDensity,
        const mfem::Vector &densityDirection,
        const mfem::Vector &baseDisplacement,
        const mfem::Vector &displacementDirection,
        const double step
    ) {
        mfem::Vector plusDensity(baseDensity);
        plusDensity.Add(step, densityDirection);

        mfem::Vector minusDensity(baseDensity);
        minusDensity.Add(-step, densityDirection);

        mfem::Vector plusDisplacement(baseDisplacement);
        plusDisplacement.Add(step, displacementDirection);

        mfem::Vector minusDisplacement(baseDisplacement);
        minusDisplacement.Add(-step, displacementDirection);

        mfem::Vector plusResidual;
        mfem::Vector minusResidual;

        mean_field::operators::kernels::apply_rotational_displacement_force_residual(
            f, *f.domainMapperStateless, rotation, plusDensity, plusDisplacement, plusResidual
        );

        mean_field::operators::kernels::apply_rotational_displacement_force_residual(
            f, *f.domainMapperStateless, rotation, minusDensity, minusDisplacement, minusResidual
        );

        plusResidual -= minusResidual;
        plusResidual /= 2.0 * step;
        return plusResidual;
    }

    template <int index>
    [[nodiscard]] mfem::Vector copy_residual_block(
        const mfem::Vector &action,
        const mean_field::operators::RotationalDisplacementForceLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block
    ) {
        mfem::Vector result(layout.size(block));
        const int offset = layout.offset(block);

        for (int entry = 0; entry < result.Size(); ++entry) {
            result(entry) = action(offset + entry);
        }

        return result;
    }
} // namespace rotational_displacement_force_test_utils

TEST_CASE(
    "Rotational Displacement Force Query Includes Density Test And Linear "
    "Position",
    tags::centrifugal &tags::quadrature &tags::unit
) {
    using DisplacementField           = mean_field::field::Field<mean_field::field::Displacement>;

    constexpr int geometryWeightOrder = 4;

    constexpr mean_field::quadrature::Query query =
        DisplacementField::make_query<mean_field::field::Displacement::Form::CentrifugalForce>(
            mean_field::quadrature::QuadratureRole::discretization, geometryWeightOrder, std::array<int, 1>{1},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

    /* density: 2, displacement test: 3, position: 1, geometry: 4 */
    constexpr int expectedBaseOrder = 2 + 3 + 1 + 4;

    STATIC_REQUIRE(query.term == mean_field::quadrature::Term::centrifugal);
    STATIC_REQUIRE(query.domain == mean_field::utils::DOMAINS::STELLAR);
    STATIC_REQUIRE(query.mapping == mean_field::quadrature::MappingKind::general);
    STATIC_REQUIRE(query.base_order.has_value());
    STATIC_REQUIRE(*query.base_order == expectedBaseOrder);
}

TEST_CASE(
    "Rotational Displacement Force Uses Negative Rotation-Potential "
    "Gradient And Excludes Vacuum",
    tags::centrifugal &tags::kernels &tags::integration &tags::accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const mfem::Vector density = rotational_displacement_force_test_utils::make_density(f, 0.31);

    mfem::Vector displacement(f.displacementFes->GetTrueVSize());
    displacement                                      = 0.0;

    const mean_field::physics::RigidRotation rotation = rotational_displacement_force_test_utils::make_rotation();

    mfem::Vector residual;

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, rotation, density, displacement, residual
    );

    mfem::ParGridFunction gradientTestField(f.displacementFes.get());

    auto gradientFunction = [&rotation](const mfem::Vector &position, mfem::Vector &value) {
        rotation.potential_gradient(position, value);
    };

    mfem::VectorFunctionCoefficient gradientCoefficient(3, gradientFunction);

    gradientTestField.ProjectCoefficient(gradientCoefficient);

    mfem::Vector gradientTestDirection;
    gradientTestField.GetTrueDofs(gradientTestDirection);

    const double signedWork =
        rotational_displacement_force_test_utils::global_dot(residual, gradientTestDirection, f.mesh->GetComm());

    INFO("Rotation-force work against grad(Psi) = " << signedWork);
    CHECK(signedWork < 0.0);

    const mfem::Vector vacuumDensity = rotational_displacement_force_test_utils::make_vacuum_only_density(f);

    mfem::Vector vacuumResidual;

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, rotation, vacuumDensity, displacement, vacuumResidual
    );

    CHECK(rotational_displacement_force_test_utils::global_norm(vacuumResidual, f.mesh->GetComm()) == 0.0);

    mfem::Vector zeroAngularVelocity(3);
    mfem::Vector zeroCenter(3);
    zeroAngularVelocity = 0.0;
    zeroCenter          = 0.0;

    const mean_field::physics::RigidRotation zeroRotation(zeroAngularVelocity, zeroCenter);

    mfem::Vector zeroRotationResidual;

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, zeroRotation, density, displacement, zeroRotationResidual
    );

    CHECK(rotational_displacement_force_test_utils::global_norm(zeroRotationResidual, f.mesh->GetComm()) == 0.0);
}

TEST_CASE(
    "Prepared Rotational Displacement Force Reprepares Selectively",
    tags::centrifugal &tags::prepared &tags::integration
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    mfem::Vector density                        = rotational_displacement_force_test_utils::make_density(f, 0.37);

    const mfem::Vector displacement             = gravity_prepared_test_utils::make_displacement(f, 0.53);

    mean_field::physics::RigidRotation rotation = rotational_displacement_force_test_utils::make_rotation(0.81);

    auto dependencies                           = rotational_displacement_force_test_utils::make_dependencies();

    mean_field::operators::PreparedRotationalDisplacementForceOperator preparedOperator(f, *f.domainMapperStateless);

    const auto initialReport =
        preparedOperator.Prepare({.density = density, .displacement = displacement}, dependencies, rotation);

    REQUIRE(initialReport.DidAnyWork());
    REQUIRE(initialReport.updatedRotation);
    REQUIRE(initialReport.preparedResidual);
    REQUIRE(preparedOperator.IsPrepared());

    mfem::Vector preparedResidual;
    mfem::Vector kernelResidual;

    preparedOperator.BuildResidual(preparedResidual);

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, rotation, density, displacement, kernelResidual
    );

    CHECK(
        rotational_displacement_force_test_utils::relative_difference(
            preparedResidual, kernelResidual, f.mesh->GetComm()
        ) < 2.0e-12
    );

    CHECK_FALSE(preparedOperator.Prepare({.density = density, .displacement = displacement}, dependencies, rotation)
                    .DidAnyWork());

    density = rotational_displacement_force_test_utils::make_density(f, 0.79);

    ++dependencies.density.revision;

    const auto densityReport =
        preparedOperator.Prepare({.density = density, .displacement = displacement}, dependencies, rotation);

    CHECK(densityReport.preparedResidual);
    CHECK_FALSE(densityReport.updatedRotation);

    rotation = rotational_displacement_force_test_utils::make_rotation(1.23);

    ++dependencies.rotation.revision;

    const auto rotationReport =
        preparedOperator.Prepare({.density = density, .displacement = displacement}, dependencies, rotation);

    CHECK(rotationReport.updatedRotation);
    CHECK(rotationReport.preparedResidual);
    CHECK(preparedOperator.GetResidualPreparationCount() == 3);
    CHECK(preparedOperator.GetResidualApplicationCount() == 1);
}

TEST_CASE(
    "Rotational Displacement Force Jacobian Matches Both Columns And "
    "Centered Differences",
    tags::centrifugal &tags::prepared &tags::jacobian &tags::accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const mfem::Vector density          = rotational_displacement_force_test_utils::make_density(f, 0.43);

    const mfem::Vector densityDirection = rotational_displacement_force_test_utils::make_density_direction(f, 0.59);

    const mfem::Vector displacement     = gravity_prepared_test_utils::make_displacement(f, 0.61);

    const mfem::Vector displacementDirection = rotational_displacement_force_test_utils::make_displacement_direction(f);

    const mean_field::physics::RigidRotation rotation = rotational_displacement_force_test_utils::make_rotation(0.93);

    mean_field::operators::PreparedRotationalDisplacementForceOperator preparedOperator(f, *f.domainMapperStateless);

    preparedOperator.Prepare(
        {.density = density, .displacement = displacement},
        rotational_displacement_force_test_utils::make_dependencies(), rotation
    );

    mfem::Vector densityAction;
    mfem::Vector displacementAction;
    mfem::Vector completeAction;

    preparedOperator.ApplyDensityJacobianAction(densityDirection, densityAction);

    preparedOperator.ApplyDisplacementJacobianAction(displacementDirection, displacementAction);

    preparedOperator.ApplyCompleteJacobianAction(densityDirection, displacementDirection, completeAction);

    mfem::Vector summedColumns(densityAction);
    summedColumns += displacementAction;

    CHECK(
        rotational_displacement_force_test_utils::relative_difference(
            completeAction, summedColumns, f.mesh->GetComm()
        ) < 2.0e-12
    );

    mfem::Vector zeroDensity(densityDirection.Size());
    mfem::Vector zeroDisplacement(displacementDirection.Size());
    zeroDensity                          = 0.0;
    zeroDisplacement                     = 0.0;

    constexpr double step                = 1.0e-5;

    const mfem::Vector densityDifference = rotational_displacement_force_test_utils::centered_difference(
        f, rotation, density, densityDirection, displacement, zeroDisplacement, step
    );

    const mfem::Vector displacementDifference = rotational_displacement_force_test_utils::centered_difference(
        f, rotation, density, zeroDensity, displacement, displacementDirection, step
    );

    const mfem::Vector completeDifference = rotational_displacement_force_test_utils::centered_difference(
        f, rotation, density, densityDirection, displacement, displacementDirection, step
    );

    const double densityError = rotational_displacement_force_test_utils::relative_difference(
        densityAction, densityDifference, f.mesh->GetComm()
    );

    const double displacementError = rotational_displacement_force_test_utils::relative_difference(
        displacementAction, displacementDifference, f.mesh->GetComm()
    );

    const double completeError = rotational_displacement_force_test_utils::relative_difference(
        completeAction, completeDifference, f.mesh->GetComm()
    );

    INFO("Density-column centered-difference error = " << densityError);
    INFO("Displacement-column centered-difference error = " << displacementError);
    INFO("Complete centered-difference error = " << completeError);

    CHECK(densityError < 2.0e-9);
    CHECK(displacementError < 3.0e-8);
    CHECK(completeError < 4.0e-8);
}

TEST_CASE(
    "Prepared Rotational Displacement Force MFEM Adapter Routes Only R-d",
    tags::centrifugal &tags::prepared &tags::mfem_operators &tags::unit
) {
    mean_field::utils::Args args = test_utils::setup_args();

    mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const mfem::Vector density          = rotational_displacement_force_test_utils::make_density(f, 0.47);

    const mfem::Vector densityDirection = rotational_displacement_force_test_utils::make_density_direction(f, 0.63);

    const mfem::Vector displacement     = gravity_prepared_test_utils::make_displacement(f, 0.57);

    const mfem::Vector displacementDirection = rotational_displacement_force_test_utils::make_displacement_direction(f);

    const mean_field::physics::RigidRotation rotation = rotational_displacement_force_test_utils::make_rotation(0.87);

    mean_field::operators::PreparedRotationalDisplacementForceOperator preparedOperator(f, *f.domainMapperStateless);

    preparedOperator.Prepare(
        {.density = density, .displacement = displacement},
        rotational_displacement_force_test_utils::make_dependencies(), rotation
    );

    const auto layout = rotational_displacement_force_test_utils::make_layout(f);

    mean_field::operators::PreparedRotationalDisplacementForceJacobianOperator adapter(layout, preparedOperator);

    mfem::BlockVector direction(layout.value_offsets());
    direction                                                                             = 0.0;

    direction.GetBlock(rotational_displacement_force_test_utils::densityValue)            = densityDirection;

    direction.GetBlock(rotational_displacement_force_test_utils::displacementValue)       = displacementDirection;

    direction.GetBlock(rotational_displacement_force_test_utils::gravityGradientValue)    = 0.23;

    direction.GetBlock(rotational_displacement_force_test_utils::gravityPotentialValue)   = -0.31;

    direction.GetBlock(rotational_displacement_force_test_utils::enthalpyValue)           = 0.37;

    direction.GetBlock(rotational_displacement_force_test_utils::barotropicConstantValue) = -0.41;

    mfem::Vector action;
    adapter.Mult(direction, action);

    mfem::Vector expectedDisplacementAction;

    preparedOperator.ApplyCompleteJacobianAction(densityDirection, displacementDirection, expectedDisplacementAction);

    const mfem::Vector actualDisplacementAction = rotational_displacement_force_test_utils::copy_residual_block(
        action, layout, rotational_displacement_force_test_utils::displacementResidual
    );

    CHECK(
        rotational_displacement_force_test_utils::relative_difference(
            actualDisplacementAction, expectedDisplacementAction, f.mesh->GetComm()
        ) < 2.0e-12
    );

    const std::array<mfem::Vector, 5> zeroRows{
        rotational_displacement_force_test_utils::copy_residual_block(
            action, layout, rotational_displacement_force_test_utils::gravityGradientResidual
        ),
        rotational_displacement_force_test_utils::copy_residual_block(
            action, layout, rotational_displacement_force_test_utils::gravityPotentialResidual
        ),
        rotational_displacement_force_test_utils::copy_residual_block(
            action, layout, rotational_displacement_force_test_utils::densityResidual
        ),
        rotational_displacement_force_test_utils::copy_residual_block(
            action, layout, rotational_displacement_force_test_utils::enthalpyResidual
        ),
        rotational_displacement_force_test_utils::copy_residual_block(
            action, layout, rotational_displacement_force_test_utils::massResidual
        )
    };

    for (const mfem::Vector &row : zeroRows) {
        CHECK(rotational_displacement_force_test_utils::global_norm(row, f.mesh->GetComm()) == 0.0);
    }
}