#include <fesa/constraints/essential_bc.hpp>

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <utility>
#include <vector>

#include <fesa/model/domain_builder.hpp>

#include <gtest/gtest.h>

namespace {

fesa::DofManager build_dofs(
    std::vector<fesa::PrescribedDof> prescribed) {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{0},
        fesa::EntityOrigin{"", "", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.set_step({"Load", std::move(prescribed), {}});
    auto result = std::move(builder).build();
    if (!result.domain.has_value()) {
        throw std::runtime_error{"Test Domain failed validation."};
    }
    return fesa::DofManager::build(*result.domain);
}

fesa::EquationSystem hand_system() {
    return {
        {
            6,
            {0, 3, 5, 6, 7, 8, 9},
            {0, 1, 2, 1, 2, 2, 3, 4, 5},
            {4.0, 1.0, 2.0, 3.0, 6.0, 5.0, 1.0, 1.0, 1.0},
        },
        {7.0, 8.0, 9.0, 0.0, 0.0, 0.0},
    };
}

TEST(EssentialBc, ZeroPrescribedValueDoesNotShiftReducedForce) {
    const std::vector<fesa::PrescribedDof> prescribed{
        {fesa::NodeId{0}, 2, 0.0},
    };
    const fesa::DofManager dofs = build_dofs(prescribed);

    const fesa::ConstraintResult result =
        fesa::eliminate_essential_bcs(hand_system(), dofs);

    ASSERT_TRUE(result.reduced_system.has_value());
    EXPECT_TRUE(result.diagnostics.empty());
    EXPECT_EQ(
        result.reduced_system->force,
        (std::vector<double>{7.0, 9.0, 0.0, 0.0, 0.0}));
}

TEST(
    ConstraintElimination,
    ShiftsNonzeroPrescribedValueAndPreservesOriginalSystem) {
    const std::vector<fesa::PrescribedDof> prescribed{
        {fesa::NodeId{0}, 2, 2.0},
    };
    const fesa::DofManager dofs = build_dofs(prescribed);
    const fesa::EquationSystem original = hand_system();
    const fesa::EquationSystem before = original;

    const fesa::ConstraintResult result =
        fesa::eliminate_essential_bcs(original, dofs);

    ASSERT_TRUE(result.reduced_system.has_value());
    EXPECT_TRUE(result.diagnostics.empty());
    const fesa::ReducedSystem& reduced = *result.reduced_system;
    EXPECT_EQ(reduced.stiffness.order, 5);
    EXPECT_EQ(
        reduced.stiffness.row_offsets,
        (std::vector<std::int32_t>{0, 2, 3, 4, 5, 6}));
    EXPECT_EQ(
        reduced.stiffness.column_indices,
        (std::vector<std::int32_t>{0, 1, 1, 2, 3, 4}));
    EXPECT_EQ(
        reduced.stiffness.values,
        (std::vector<double>{4.0, 2.0, 5.0, 1.0, 1.0, 1.0}));
    EXPECT_EQ(
        reduced.force,
        (std::vector<double>{5.0, -3.0, 0.0, 0.0, 0.0}));

    EXPECT_EQ(original.stiffness.order, before.stiffness.order);
    EXPECT_EQ(
        original.stiffness.row_offsets,
        before.stiffness.row_offsets);
    EXPECT_EQ(
        original.stiffness.column_indices,
        before.stiffness.column_indices);
    EXPECT_EQ(original.stiffness.values, before.stiffness.values);
    EXPECT_EQ(original.force, before.force);

    EXPECT_EQ(
        dofs.reconstruct_full(
            std::vector<double>{1.0, 3.0, 0.0, 0.0, 0.0}),
        (std::vector<double>{1.0, 2.0, 3.0, 0.0, 0.0, 0.0}));
}

TEST(ConstraintElimination, AcceptsAllDofsConstrained) {
    const std::vector<fesa::PrescribedDof> prescribed{
        {fesa::NodeId{0}, 1, 1.0},
        {fesa::NodeId{0}, 2, 2.0},
        {fesa::NodeId{0}, 3, 3.0},
        {fesa::NodeId{0}, 4, 4.0},
        {fesa::NodeId{0}, 5, 5.0},
        {fesa::NodeId{0}, 6, 6.0},
    };
    const fesa::DofManager dofs = build_dofs(prescribed);

    const fesa::ConstraintResult result =
        fesa::eliminate_essential_bcs(hand_system(), dofs);

    ASSERT_TRUE(result.reduced_system.has_value());
    const fesa::ReducedSystem& reduced = *result.reduced_system;
    EXPECT_EQ(reduced.stiffness.order, 0);
    EXPECT_EQ(
        reduced.stiffness.row_offsets,
        (std::vector<std::int32_t>{0}));
    EXPECT_TRUE(reduced.stiffness.column_indices.empty());
    EXPECT_TRUE(reduced.stiffness.values.empty());
    EXPECT_TRUE(reduced.force.empty());
    EXPECT_EQ(
        dofs.reconstruct_full({}),
        (std::vector<double>{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}));
}

TEST(ConstraintElimination, RejectsNonfiniteSystemData) {
    const fesa::DofManager dofs = build_dofs({});
    fesa::EquationSystem nonfinite_matrix = hand_system();
    nonfinite_matrix.stiffness.values[0] =
        std::numeric_limits<double>::quiet_NaN();
    fesa::EquationSystem nonfinite_force = hand_system();
    nonfinite_force.force[0] =
        std::numeric_limits<double>::infinity();

    const fesa::ConstraintResult matrix_result =
        fesa::eliminate_essential_bcs(nonfinite_matrix, dofs);
    const fesa::ConstraintResult force_result =
        fesa::eliminate_essential_bcs(nonfinite_force, dofs);

    ASSERT_FALSE(matrix_result.diagnostics.empty());
    EXPECT_EQ(
        matrix_result.diagnostics.front().code,
        "equation.invalid_system");
    ASSERT_FALSE(force_result.diagnostics.empty());
    EXPECT_EQ(
        force_result.diagnostics.front().code,
        "equation.invalid_system");
}

TEST(ConstraintElimination, RejectsNonfiniteReducedForce) {
    const fesa::DofManager dofs = build_dofs({
        {fesa::NodeId{0}, 2, 2.0},
    });
    fesa::EquationSystem overflowing = hand_system();
    overflowing.stiffness.values[1] =
        std::numeric_limits<double>::max();

    const fesa::ConstraintResult result =
        fesa::eliminate_essential_bcs(overflowing, dofs);

    EXPECT_FALSE(result.reduced_system.has_value());
    ASSERT_FALSE(result.diagnostics.empty());
    EXPECT_EQ(
        result.diagnostics.front().code,
        "equation.nonfinite_result");
}

TEST(Reaction, UsesOriginalFullEquilibriumEquation) {
    const std::vector<double> reaction = fesa::recover_reaction(
        hand_system(),
        std::vector<double>{1.0, 2.0, 3.0, 0.0, 0.0, 0.0});

    EXPECT_EQ(
        reaction,
        (std::vector<double>{5.0, 17.0, 20.0, 0.0, 0.0, 0.0}));
}

TEST(Reaction, RejectsNonfiniteFullDisplacement) {
    std::vector<double> displacement(6, 0.0);
    displacement[0] = std::numeric_limits<double>::quiet_NaN();

    EXPECT_THROW(
        static_cast<void>(
            fesa::recover_reaction(hand_system(), displacement)),
        std::invalid_argument);
}

TEST(Reaction, RejectsNonfiniteRecoveredValue) {
    fesa::EquationSystem overflowing = hand_system();
    overflowing.stiffness.values[0] =
        std::numeric_limits<double>::max();

    EXPECT_THROW(
        static_cast<void>(fesa::recover_reaction(
            overflowing,
            std::vector<double>{2.0, 0.0, 0.0, 0.0, 0.0, 0.0})),
        std::invalid_argument);
}

}  // namespace