#include <fesa/fem/dof_manager.hpp>

#include <array>
#include <cstdint>
#include <stdexcept>
#include <utility>
#include <vector>

#include <fesa/model/domain_builder.hpp>

#include <gtest/gtest.h>

namespace {

fesa::Domain build_domain(
    std::vector<fesa::Node> nodes,
    fesa::StepDefinition step,
    const std::array<fesa::NodeId, 2> element_nodes = {
        fesa::NodeId{4},
        fesa::NodeId{10},
    }) {
    fesa::DomainBuilder builder;
    for (fesa::Node& node : nodes) {
        builder.add_node(std::move(node));
    }
    builder.add_material({
        fesa::MaterialId{0},
        "Steel",
        210.0e9,
        0.3,
    });
    builder.add_section({
        fesa::SectionId{0},
        "General",
        0.04,
        1.2e-4,
        1.4e-4,
        2.0e-4,
        0.03,
        0.031,
        fesa::ShearPropertySource::input,
        fesa::Vec3{0.0, 1.0, 0.0},
        {},
    });
    builder.add_beam_element({
        fesa::ElementId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        element_nodes,
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });
    builder.set_step(std::move(step));

    auto result = std::move(builder).build();
    EXPECT_TRUE(result.diagnostics.empty());
    EXPECT_TRUE(result.domain.has_value());
    return std::move(*result.domain);
}

std::vector<fesa::Node> nodes_in_external_label_order() {
    return {
        {
            fesa::NodeId{20},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
            fesa::Vec3{2.0, 0.0, 0.0},
        },
        {
            fesa::NodeId{4},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 100},
            fesa::Vec3{0.0, 0.0, 0.0},
        },
        {
            fesa::NodeId{10},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 50},
            fesa::Vec3{1.0, 0.0, 0.0},
        },
    };
}

TEST(EquationNumbering, UsesInternalNodeIdOrderNotExternalLabelOrStorageOrder) {
    const fesa::Domain domain =
        build_domain(nodes_in_external_label_order(), {"Load", {}, {}});

    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    EXPECT_EQ(dofs.full_dof_count(), 18);
    EXPECT_EQ(dofs.free_equation_count(), 18);
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{4}, fesa::NodeDof::ux}),
        std::optional<std::size_t>{0});
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{4}, fesa::NodeDof::rz}),
        std::optional<std::size_t>{5});
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{10}, fesa::NodeDof::ux}),
        std::optional<std::size_t>{6});
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{20}, fesa::NodeDof::rz}),
        std::optional<std::size_t>{17});
}

TEST(DofManager, NumbersOnlyFreeDofsAndReconstructsPrescribedValues) {
    const fesa::Domain domain = build_domain(
        nodes_in_external_label_order(),
        {
            "Load",
            {
                {fesa::NodeId{4}, 1, 1.25},
                {fesa::NodeId{10}, 6, -0.5},
            },
            {},
        });

    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    EXPECT_EQ(dofs.full_dof_count(), 18);
    EXPECT_EQ(dofs.free_equation_count(), 16);
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{4}, fesa::NodeDof::ux}),
        std::nullopt);
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{4}, fesa::NodeDof::uy}),
        std::optional<std::size_t>{0});
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{10}, fesa::NodeDof::rx}),
        std::optional<std::size_t>{8});
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{10}, fesa::NodeDof::rz}),
        std::nullopt);
    EXPECT_EQ(
        dofs.equation({fesa::NodeId{20}, fesa::NodeDof::ux}),
        std::optional<std::size_t>{10});
    EXPECT_EQ(dofs.equation(std::size_t{0}), std::nullopt);
    EXPECT_EQ(
        dofs.prescribed_value(std::size_t{0}),
        std::optional<double>{1.25});
    EXPECT_EQ(
        dofs.equation(std::size_t{1}),
        std::optional<std::size_t>{0});
    EXPECT_EQ(
        dofs.prescribed_value(std::size_t{1}),
        std::nullopt);

    std::vector<double> reduced(dofs.free_equation_count());
    for (std::size_t equation = 0; equation < reduced.size(); ++equation) {
        reduced[equation] = 100.0 + static_cast<double>(equation);
    }

    EXPECT_EQ(
        dofs.reconstruct_full(reduced),
        (std::vector<double>{
            1.25,
            100.0,
            101.0,
            102.0,
            103.0,
            104.0,
            105.0,
            106.0,
            107.0,
            108.0,
            109.0,
            -0.5,
            110.0,
            111.0,
            112.0,
            113.0,
            114.0,
            115.0,
        }));
}

TEST(DofManager, MapsElementNodesToTwelveFullDofsInComponentOrder) {
    const fesa::Domain domain = build_domain(
        nodes_in_external_label_order(),
        {"Load", {}, {}},
        {fesa::NodeId{10}, fesa::NodeId{4}});
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    EXPECT_EQ(
        dofs.element_full_dofs(domain.beam_elements().front()),
        (std::array<std::size_t, 12>{
            6, 7, 8, 9, 10, 11,
            0, 1, 2, 3, 4, 5,
        }));
}

TEST(DofManager, RejectsInvalidAddressesAndReducedVectorSize) {
    const fesa::Domain domain =
        build_domain(nodes_in_external_label_order(), {"Load", {}, {}});
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    EXPECT_THROW(
        static_cast<void>(dofs.equation({
            fesa::NodeId{4},
            static_cast<fesa::NodeDof>(6),
        })),
        std::invalid_argument);
    EXPECT_THROW(
        static_cast<void>(
            dofs.equation({fesa::NodeId{99}, fesa::NodeDof::ux})),
        std::out_of_range);
    EXPECT_THROW(
        static_cast<void>(dofs.equation(dofs.full_dof_count())),
        std::out_of_range);
    EXPECT_THROW(
        static_cast<void>(
            dofs.prescribed_value(dofs.full_dof_count())),
        std::out_of_range);
    EXPECT_THROW(
        static_cast<void>(
            dofs.reconstruct_full(std::vector<double>(17, 0.0))),
        std::invalid_argument);
}

}  // namespace