#include <algorithm>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace rotational_displacement_force_context_test_utils {
    using Context =
        mean_field::operators::context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext;

    using Dependencies =
        mean_field::operators::context::rotational_displacement_force::RotationalDisplacementForceDependencies;

    [[nodiscard]] Dependencies make_dependencies() {
        return {
            .discretization = {.identity = 101, .revision = 3},
            .density        = {.identity = 103, .revision = 5},
            .displacement   = {.identity = 107, .revision = 7},
            .rotation       = {.identity = 109, .revision = 11}
        };
    }

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

        mfem::FunctionCoefficient coefficient([offset](const mfem::Vector &position) {
            return offset + 0.04 * position(0) - 0.03 * position(1) + 0.02 * position(2);
        });

        density.ProjectCoefficient(coefficient);

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

    [[nodiscard]] double relative_difference(
        const mfem::Vector &left,
        const mfem::Vector &right
    ) {
        mfem::Vector difference(left);
        difference -= right;

        return difference.Norml2() / std::max(right.Norml2(), 1.0e-30);
    }
} // namespace rotational_displacement_force_context_test_utils

TEST_CASE(
    "Rotational Displacement Force Context Applies Selective Invalidation",
    tags::rotation_context_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());

    auto dependencies = rotational_displacement_force_context_test_utils::make_dependencies();

    rotational_displacement_force_context_test_utils::Context context(f, *f.domainMapperStateless);

    mfem::Vector densityTrue      = rotational_displacement_force_context_test_utils::make_density(f, 0.83);
    mfem::Vector density          = context.GetDensityMap().gather(densityTrue);
    mfem::Vector displacementTrue = gravity_prepared_test_utils::make_displacement(f, 0.47);
    mfem::Vector displacement     = context.GetDisplacementMap().gather(displacementTrue);

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

    REQUIRE(context.IsPrepared());
    CHECK(context.MatchesDependencies(dependencies));
    CHECK(initialReport.preparedStaticDependencies);
    CHECK(initialReport.preparedGeometryState);
    CHECK(initialReport.preparedRotationDependencies);
    CHECK(initialReport.preparedBaseState);
    CHECK(initialReport.updatedDensity);
    CHECK(initialReport.updatedDisplacement);

    const mfem::Vector expectedDensityTrue      = context.GetDensityMap().scatter(density);
    const mfem::Vector expectedDisplacementTrue = context.GetDisplacementMap().scatter(displacement);

    CHECK(
        rotational_displacement_force_context_test_utils::relative_difference(
            context.GetBaseDensityTrue(), expectedDensityTrue
        ) < 1.0e-14
    );

    CHECK(
        rotational_displacement_force_context_test_utils::relative_difference(
            context.GetDisplacementTrue(), expectedDisplacementTrue
        ) < 1.0e-14
    );

    const auto unchangedReport = context.Prepare({.density = density, .displacement = displacement}, dependencies);

    CHECK_FALSE(unchangedReport.DidAnyWork());

    densityTrue = rotational_displacement_force_context_test_utils::make_density(f, 1.17);
    density     = context.GetDensityMap().gather(densityTrue);

    ++dependencies.density.revision;

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

    CHECK_FALSE(densityReport.preparedStaticDependencies);
    CHECK_FALSE(densityReport.preparedGeometryState);
    CHECK_FALSE(densityReport.preparedRotationDependencies);
    CHECK(densityReport.preparedBaseState);
    CHECK(densityReport.updatedDensity);
    CHECK_FALSE(densityReport.updatedDisplacement);

    const mfem::Vector densityAfterDensityRevision(context.GetBaseDensityTrue());

    displacementTrue = gravity_prepared_test_utils::make_displacement(f, 0.81);
    displacement     = context.GetDisplacementMap().gather(displacementTrue);

    ++dependencies.displacement.revision;

    const auto displacementReport = context.Prepare({.density = density, .displacement = displacement}, dependencies);

    CHECK_FALSE(displacementReport.preparedStaticDependencies);
    CHECK(displacementReport.preparedGeometryState);
    CHECK_FALSE(displacementReport.preparedRotationDependencies);
    CHECK(displacementReport.preparedBaseState);
    CHECK_FALSE(displacementReport.updatedDensity);
    CHECK(displacementReport.updatedDisplacement);

    CHECK(
        rotational_displacement_force_context_test_utils::relative_difference(
            context.GetBaseDensityTrue(), densityAfterDensityRevision
        ) < 1.0e-14
    );

    ++dependencies.rotation.revision;

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

    CHECK_FALSE(rotationReport.preparedStaticDependencies);
    CHECK_FALSE(rotationReport.preparedGeometryState);
    CHECK(rotationReport.preparedRotationDependencies);
    CHECK(rotationReport.preparedBaseState);
    CHECK_FALSE(rotationReport.updatedDensity);
    CHECK_FALSE(rotationReport.updatedDisplacement);

    const auto statistics = context.GetPreparationStatistics();

    CHECK(statistics.staticPreparations == 1);
    CHECK(statistics.geometryPreparations == 2);
    CHECK(statistics.rotationPreparations == 2);
    CHECK(statistics.baseStatePreparations == 4);
}

TEST_CASE(
    "Rotational Displacement Force Context Treats New Identities As New "
    "Dependency Streams",
    tags::rotation_context_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());

    auto dependencies = rotational_displacement_force_context_test_utils::make_dependencies();

    rotational_displacement_force_context_test_utils::Context context(f, *f.domainMapperStateless);

    const mfem::Vector density =
        context.GetDensityMap().gather(rotational_displacement_force_context_test_utils::make_density(f, 0.91));
    const mfem::Vector displacement =
        context.GetDisplacementMap().gather(gravity_prepared_test_utils::make_displacement(f, 0.39));

    context.Prepare({.density = density, .displacement = displacement}, dependencies);

    dependencies.density.identity += 1000;
    dependencies.density.revision = 0;

    const auto report             = context.Prepare({.density = density, .displacement = displacement}, dependencies);

    CHECK(report.preparedBaseState);
    CHECK(report.updatedDensity);
    CHECK_FALSE(report.preparedGeometryState);
    CHECK_FALSE(report.preparedRotationDependencies);
}