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

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace pressure_force_kernel_test_utils {
    [[nodiscard]] mfem::Vector make_deterministic_vector(
        const int size,
        const double phase
    ) {
        mfem::Vector vector(size);

        for (int index = 0; index < size; ++index) {
            const double position = static_cast<double>(index + 1);

            vector(index) = 0.71 + 0.19 * std::sin(0.31 * position + phase) +
                            0.08 * std::cos(0.17 * position - 0.5 * phase);
        }

        return vector;
    }

    [[nodiscard]] mfem::Vector
    make_zero_displacement(const mean_field::fem::FEM &f) {
        mfem::Vector displacementTrue(f.displacementFes->GetTrueVSize());
        displacementTrue = 0.0;
        return displacementTrue;
    }

    [[nodiscard]] mfem::Vector
    make_vacuum_only_enthalpy(const mean_field::fem::FEM &f) {
        mfem::Vector enthalpyTrue =
            make_deterministic_vector(f.enthalpyFes->GetTrueVSize(), 0.43);

        mfem::Array<int> stellarElementMask;
        mean_field::utils::populate_element_mask(
            f.mesh.get(), mean_field::utils::DOMAINS::STELLAR,
            stellarElementMask
        );

        mfem::Array<int> stellarEnthalpyTrueDofs;
        mean_field::utils::populate_domain_tdofs(
            f.enthalpyFes.get(), stellarElementMask, stellarEnthalpyTrueDofs
        );

        for (int listIndex = 0; listIndex < stellarEnthalpyTrueDofs.Size();
             ++listIndex) {
            const int trueDof = stellarEnthalpyTrueDofs[listIndex];

            MFEM_VERIFY(
                trueDof >= 0 && trueDof < enthalpyTrue.Size(),
                "The stellar enthalpy true-DOF mask contains an "
                "invalid index."
            );

            enthalpyTrue(trueDof) = 0.0;
        }

        return enthalpyTrue;
    }

    [[nodiscard]] mfem::Vector
    make_positive_asymmetric_enthalpy(const mean_field::fem::FEM &f) {
        mfem::FunctionCoefficient coefficient([](const mfem::Vector &position) {
            return 1.10 + 0.07 * position(0) - 0.04 * position(1) +
                   0.03 * position(2);
        });

        mfem::ParGridFunction enthalpyField(f.enthalpyFes.get());

        enthalpyField.ProjectCoefficient(coefficient);

        mfem::Vector enthalpyTrue;
        enthalpyField.GetTrueDofs(enthalpyTrue);

        return enthalpyTrue;
    }

    [[nodiscard]] mfem::Vector make_component_test_field(
        const mean_field::fem::FEM &f,
        const int component,
        const int coordinate
    ) {
        const int dimension = f.mesh->Dimension();

        MFEM_VERIFY(
            component >= 0 && component < dimension,
            "The requested vector component is invalid."
        );

        MFEM_VERIFY(
            coordinate >= -1 && coordinate < dimension,
            "The requested coordinate is invalid."
        );

        /*
         * coordinate == -1 gives the rigid translation e_component.
         *
         * Otherwise this gives
         *
         *     w = x_coordinate e_component.
         */
        mfem::VectorFunctionCoefficient coefficient(
            dimension,
            [component, coordinate,
             dimension](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(dimension);
                value            = 0.0;

                value(component) = coordinate < 0 ? 1.0 : position(coordinate);
            }
        );

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

        field.ProjectCoefficient(coefficient);

        mfem::Vector fieldTrue;
        field.GetTrueDofs(fieldTrue);

        return fieldTrue;
    }

    [[nodiscard]] double global_dot(
        const mfem::Vector &left,
        const mfem::Vector &right,
        MPI_Comm communicator
    ) {
        MFEM_VERIFY(
            left.Size() == right.Size(),
            "The global dot-product vectors have different sizes."
        );

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

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

        return globalDot;
    }
} // namespace pressure_force_kernel_test_utils

TEST_CASE(
    "Pressure Force Residual Vanishes For Zero Enthalpy",
    tags::barotrope &tags::pressure &tags::kernels &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());

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 0.25);

    mfem::Vector enthalpyTrue(f.enthalpyFes->GetTrueVSize());
    enthalpyTrue = 0.0;

    const mfem::Vector displacementTrue =
        pressure_force_kernel_test_utils::make_zero_displacement(f);

    mfem::Vector residualTrue;

    mean_field::operators::kernels::apply_pressure_force_residual(
        f, *f.domainMapperStateless, barotrope, enthalpyTrue, displacementTrue,
        residualTrue
    );

    REQUIRE(residualTrue.Size() == f.displacementFes->GetTrueVSize());

    const double residualNorm = gravity_prepared_test_utils::global_norm(
        residualTrue, f.mesh->GetComm()
    );

    CHECK(residualNorm == 0.0);
}

TEST_CASE(
    "Pressure Force Residual Excludes Vacuum Enthalpy Exactly",
    tags::barotrope &tags::pressure &tags::kernels &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());

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 0.25);

    const mfem::Vector enthalpyTrue =
        pressure_force_kernel_test_utils::make_vacuum_only_enthalpy(f);

    const double enthalpyNorm = gravity_prepared_test_utils::global_norm(
        enthalpyTrue, f.mesh->GetComm()
    );

    /*
     * Ensure this is a real exclusion test rather than another
     * all-zero-input test.
     */
    REQUIRE(enthalpyNorm > 0.0);

    const mfem::Vector displacementTrue =
        pressure_force_kernel_test_utils::make_zero_displacement(f);

    mfem::Vector residualTrue;

    mean_field::operators::kernels::apply_pressure_force_residual(
        f, *f.domainMapperStateless, barotrope, enthalpyTrue, displacementTrue,
        residualTrue
    );

    REQUIRE(residualTrue.Size() == f.displacementFes->GetTrueVSize());

    const double residualNorm = gravity_prepared_test_utils::global_norm(
        residualTrue, f.mesh->GetComm()
    );

    CHECK(residualNorm == 0.0);
}

TEST_CASE(
    "Pressure Force Residual Is Nonzero For Positive Stellar Pressure",
    tags::barotrope &tags::pressure &tags::kernels &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());

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 0.25);

    /*
     * With n = 3 and K = 1/4:
     *
     *     P(1) = 1/4.
     */
    mfem::Vector enthalpyTrue(f.enthalpyFes->GetTrueVSize());
    enthalpyTrue = 1.0;

    const mfem::Vector displacementTrue =
        pressure_force_kernel_test_utils::make_zero_displacement(f);

    mfem::Vector residualTrue;

    mean_field::operators::kernels::apply_pressure_force_residual(
        f, *f.domainMapperStateless, barotrope, enthalpyTrue, displacementTrue,
        residualTrue
    );

    const double residualNorm = gravity_prepared_test_utils::global_norm(
        residualTrue, f.mesh->GetComm()
    );

    INFO("Positive-pressure residual norm = " << residualNorm);

    CHECK(std::isfinite(residualNorm));

    CHECK(residualNorm > 100.0 * std::numeric_limits<double>::epsilon());
}

TEST_CASE(
    "Pressure Force Residual Does No Work Against Rigid Translations",
    tags::barotrope &tags::pressure &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());

    REQUIRE(f.displacementFes->GetOrdering() == mfem::Ordering::byNODES);

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 0.25);

    const mfem::Vector enthalpyTrue =
        pressure_force_kernel_test_utils::make_positive_asymmetric_enthalpy(f);

    const mfem::Vector displacementTrue =
        pressure_force_kernel_test_utils::make_zero_displacement(f);

    mfem::Vector residualTrue;

    mean_field::operators::kernels::apply_pressure_force_residual(
        f, *f.domainMapperStateless, barotrope, enthalpyTrue, displacementTrue,
        residualTrue
    );

    const double residualNorm = gravity_prepared_test_utils::global_norm(
        residualTrue, f.mesh->GetComm()
    );

    REQUIRE(residualNorm > 0.0);

    const int dimension = f.mesh->Dimension();

    for (int component = 0; component < dimension; ++component) {
        const mfem::Vector translationTrue =
            pressure_force_kernel_test_utils::make_component_test_field(
                f, component, -1
            );

        const double translationNorm = gravity_prepared_test_utils::global_norm(
            translationTrue, f.mesh->GetComm()
        );

        const double translationWork =
            pressure_force_kernel_test_utils::global_dot(
                translationTrue, residualTrue, f.mesh->GetComm()
            );

        const double dotProductScale =
            std::fmax(residualNorm * translationNorm, 1.0);

        CAPTURE(component, translationWork, dotProductScale);

        CHECK(std::abs(translationWork) <= 5.0e-12 * dotProductScale);
    }
}

TEST_CASE(
    "Pressure Force Residual Respects byNODES Component Layout",
    tags::barotrope &tags::pressure &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());

    REQUIRE(f.displacementFes->GetOrdering() == mfem::Ordering::byNODES);

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 0.25);

    const mfem::Vector enthalpyTrue =
        pressure_force_kernel_test_utils::make_positive_asymmetric_enthalpy(f);

    const mfem::Vector displacementTrue =
        pressure_force_kernel_test_utils::make_zero_displacement(f);

    mfem::Vector residualTrue;

    mean_field::operators::kernels::apply_pressure_force_residual(
        f, *f.domainMapperStateless, barotrope, enthalpyTrue, displacementTrue,
        residualTrue
    );

    const int dimension = f.mesh->Dimension();

    REQUIRE(dimension == 3);

    mfem::DenseMatrix virtualWork(dimension, dimension);

    for (int component = 0; component < dimension; ++component) {
        for (int coordinate = 0; coordinate < dimension; ++coordinate) {
            const mfem::Vector affineTestTrue =
                pressure_force_kernel_test_utils::make_component_test_field(
                    f, component, coordinate
                );

            virtualWork(component, coordinate) =
                pressure_force_kernel_test_utils::global_dot(
                    affineTestTrue, residualTrue, f.mesh->GetComm()
                );
        }
    }

    double meanDiagonalWork = 0.0;

    for (int component = 0; component < dimension; ++component) {
        meanDiagonalWork += virtualWork(component, component);
    }

    meanDiagonalWork /= static_cast<double>(dimension);

    // INFO(
    //     "Affine pressure virtual-work tensor:\n"
    //     << virtualWork
    // );

    INFO("Mean diagonal virtual work = " << meanDiagonalWork);

    REQUIRE(
        std::abs(meanDiagonalWork) >
        100.0 * std::numeric_limits<double>::epsilon()
    );

    const double comparisonTolerance = 1.0e-8 * std::abs(meanDiagonalWork);

    for (int component = 0; component < dimension; ++component) {
        for (int coordinate = 0; coordinate < dimension; ++coordinate) {
            const double computedWork = virtualWork(component, coordinate);

            CAPTURE(
                component, coordinate, computedWork, meanDiagonalWork,
                comparisonTolerance
            );

            if (component == coordinate) {
                CHECK(
                    std::abs(computedWork - meanDiagonalWork) <=
                    comparisonTolerance
                );
            } else {
                CHECK(std::abs(computedWork) <= comparisonTolerance);
            }
        }
    }
}