#include <fesa/assembly/contribution.hpp>
#include <fesa/assembly/serial_assembler.hpp>

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

#include <fesa/model/domain_builder.hpp>

#include <gtest/gtest.h>

namespace {

fesa::IsotropicElastic unit_material() {
    return {
        fesa::MaterialId{0},
        "Unit",
        1.0,
        0.25,
    };
}

fesa::BeamSection unit_section() {
    return {
        fesa::SectionId{0},
        "Unit",
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        fesa::ShearPropertySource::input,
        fesa::Vec3{0.0, 1.0, 0.0},
        {},
    };
}

fesa::Domain finish_domain(
    fesa::DomainBuilder builder,
    fesa::StepDefinition step) {
    builder.set_step(std::move(step));
    auto result = std::move(builder).build();
    if (!result.domain.has_value()) {
        throw std::runtime_error{"Test Domain failed validation."};
    }
    return std::move(*result.domain);
}

fesa::Domain build_chain_domain(const bool reverse_storage_order) {
    fesa::DomainBuilder builder;
    std::array<fesa::Node, 3> nodes{{
        {
            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},
        },
    }};
    if (reverse_storage_order) {
        std::ranges::reverse(nodes);
    }
    for (fesa::Node& node : nodes) {
        builder.add_node(std::move(node));
    }

    builder.add_material(unit_material());
    builder.add_section(unit_section());

    std::array<fesa::BeamElement, 2> elements{{
        {
            fesa::ElementId{1},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 200},
            {fesa::NodeId{10}, fesa::NodeId{20}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
        {
            fesa::ElementId{0},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 10},
            {fesa::NodeId{4}, fesa::NodeId{10}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
    }};
    if (reverse_storage_order) {
        std::ranges::reverse(elements);
    }
    for (fesa::BeamElement& element : elements) {
        builder.add_beam_element(std::move(element));
    }

    return finish_domain(
        std::move(builder),
        {
            "Load",
            {},
            {
                {
                    fesa::NodeId{20},
                    {3.0, 4.0, 5.0, 6.0, 7.0, 8.0},
                },
                {
                    fesa::NodeId{10},
                    {1.0, 2.0, 3.0, 4.0, 5.0, 6.0},
                },
                {
                    fesa::NodeId{10},
                    {-1.0, 10.0, -3.0, 0.0, 0.0, 0.0},
                },
            },
        });
}

fesa::Domain build_parallel_domain(
    const std::vector<std::int64_t>& element_labels) {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{4},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{10},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
        fesa::Vec3{1.0, 0.0, 0.0},
    });
    builder.add_material(unit_material());
    builder.add_section(unit_section());

    for (std::size_t index = 0; index < element_labels.size(); ++index) {
        builder.add_beam_element({
            fesa::ElementId{static_cast<std::int64_t>(index)},
            fesa::EntityOrigin{
                "BeamPart",
                "Beam-1",
                element_labels[index],
            },
            {fesa::NodeId{4}, fesa::NodeId{10}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        });
    }
    return finish_domain(std::move(builder), {"Load", {}, {}});
}

fesa::Domain build_rounding_domain(const bool large_element_first) {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{4},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{10},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
        fesa::Vec3{1.0, 0.0, 0.0},
    });

    for (std::size_t index = 0; index < 3; ++index) {
        fesa::IsotropicElastic material = unit_material();
        material.id =
            fesa::MaterialId{static_cast<std::int64_t>(index)};
        material.name = "Material-" + std::to_string(index);
        material.young = index == 2 ? 1.0e16 : 1.0;
        builder.add_material(std::move(material));
    }
    builder.add_section(unit_section());

    std::array<std::size_t, 3> storage_order{0, 1, 2};
    if (large_element_first) {
        storage_order = {2, 0, 1};
    }
    for (const std::size_t index : storage_order) {
        constexpr std::array<std::int64_t, 3> element_ids{1, 2, 0};
        builder.add_beam_element({
            fesa::ElementId{element_ids[index]},
            fesa::EntityOrigin{
                "BeamPart",
                "Beam-1",
                static_cast<std::int64_t>((index + 1) * 10),
            },
            {fesa::NodeId{4}, fesa::NodeId{10}},
            fesa::MaterialId{static_cast<std::int64_t>(index)},
            fesa::SectionId{0},
        });
    }
    return finish_domain(std::move(builder), {"Load", {}, {}});
}

fesa::Domain build_orphan_node_domain() {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{1},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
        fesa::Vec3{1.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{2},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 3},
        fesa::Vec3{2.0, 0.0, 0.0},
    });
    builder.add_material(unit_material());
    builder.add_section(unit_section());
    builder.add_beam_element({
        fesa::ElementId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        {fesa::NodeId{0}, fesa::NodeId{1}},
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });
    return finish_domain(std::move(builder), {"Load", {}, {}});
}

double csr_value(
    const fesa::SymmetricCsr& matrix,
    std::size_t row,
    std::size_t column) {
    if (column < row) {
        std::swap(row, column);
    }
    const auto begin = matrix.column_indices.begin() + matrix.row_offsets[row];
    const auto end =
        matrix.column_indices.begin() + matrix.row_offsets[row + 1];
    const auto found = std::lower_bound(
        begin, end, static_cast<std::int32_t>(column));
    if (found == end || *found != static_cast<std::int32_t>(column)) {
        throw std::out_of_range{"CSR entry is not present."};
    }
    return matrix.values[static_cast<std::size_t>(
        std::distance(matrix.column_indices.begin(), found))];
}

std::vector<std::uint64_t> value_bits(
    const fesa::SymmetricCsr& matrix) {
    std::vector<std::uint64_t> bits;
    bits.reserve(matrix.values.size());
    for (const double value : matrix.values) {
        bits.push_back(std::bit_cast<std::uint64_t>(value));
    }
    return bits;
}

TEST(CanonicalContribution, OrdersByCoordinateElementAndLocalOrder) {
    const std::vector<fesa::MatrixContribution> contributions{
        {1, 2, fesa::ElementId{9}, 4, 9.0},
        {0, 2, fesa::ElementId{3}, 8, 3.8},
        {0, 2, fesa::ElementId{3}, 2, 3.2},
        {0, 2, fesa::ElementId{1}, 7, 1.7},
        {0, 0, fesa::ElementId{5}, 0, 5.0},
    };

    const std::vector<fesa::MatrixContribution> canonical =
        fesa::canonicalize_contributions(contributions);

    ASSERT_EQ(canonical.size(), contributions.size());
    EXPECT_EQ(canonical[0].row, 0);
    EXPECT_EQ(canonical[0].column, 0);
    EXPECT_EQ(canonical[0].element, fesa::ElementId{5});
    EXPECT_EQ(canonical[1].element, fesa::ElementId{1});
    EXPECT_EQ(canonical[2].element, fesa::ElementId{3});
    EXPECT_EQ(canonical[2].local_order, 2);
    EXPECT_EQ(canonical[3].element, fesa::ElementId{3});
    EXPECT_EQ(canonical[3].local_order, 8);
    EXPECT_EQ(canonical[4].row, 1);
    EXPECT_EQ(canonical[4].column, 2);
}

TEST(DeterministicMerge, IsBitwiseStableAcrossInputPermutations) {
    const std::vector<fesa::MatrixContribution> contributions{
        {0, 1, fesa::ElementId{2}, 0, 1.0},
        {0, 1, fesa::ElementId{0}, 0, 1.0e16},
        {0, 1, fesa::ElementId{1}, 0, -1.0e16},
        {1, 1, fesa::ElementId{1}, 1, -0.0},
        {1, 1, fesa::ElementId{0}, 1, +0.0},
    };
    const std::vector<fesa::MatrixContribution> canonical =
        fesa::canonicalize_contributions(contributions);
    const fesa::SymmetricCsr expected =
        fesa::merge_contributions(2, canonical);

    std::array<std::size_t, 5> permutation{0, 1, 2, 3, 4};
    do {
        std::vector<fesa::MatrixContribution> shuffled;
        shuffled.reserve(contributions.size());
        for (const std::size_t index : permutation) {
            shuffled.push_back(contributions[index]);
        }
        const std::vector<fesa::MatrixContribution> shuffled_canonical =
            fesa::canonicalize_contributions(shuffled);
        const fesa::SymmetricCsr actual =
            fesa::merge_contributions(2, shuffled_canonical);

        EXPECT_EQ(actual.order, expected.order);
        EXPECT_EQ(actual.row_offsets, expected.row_offsets);
        EXPECT_EQ(actual.column_indices, expected.column_indices);
        EXPECT_EQ(value_bits(actual), value_bits(expected));
    } while (std::ranges::next_permutation(permutation).found);

    EXPECT_EQ(
        std::bit_cast<std::uint64_t>(csr_value(expected, 0, 1)),
        std::bit_cast<std::uint64_t>(1.0));
    EXPECT_EQ(
        std::bit_cast<std::uint64_t>(csr_value(expected, 1, 1)),
        std::bit_cast<std::uint64_t>(+0.0));
}

TEST(SparsePattern, BuildsExpectedTwoElementChainStructure) {
    const fesa::Domain domain = build_chain_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    const fesa::EquationSystem system =
        fesa::assemble_serial(domain, dofs);

    EXPECT_EQ(
        system.stiffness.row_offsets,
        (std::vector<std::int32_t>{
            0,
            12,
            23,
            33,
            42,
            50,
            57,
            69,
            80,
            90,
            99,
            107,
            114,
            120,
            125,
            129,
            132,
            134,
            135,
        }));

    std::vector<std::int32_t> expected_columns;
    for (std::int32_t row = 0; row < 18; ++row) {
        const std::int32_t last_column = row < 6 ? 11 : 17;
        for (std::int32_t column = row; column <= last_column; ++column) {
            expected_columns.push_back(column);
        }
    }
    EXPECT_EQ(system.stiffness.column_indices, expected_columns);
}

TEST(SymmetricCsr, StoresSortedUpperTriangleWithValidOffsets) {
    const fesa::Domain domain = build_chain_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);
    const fesa::SymmetricCsr matrix =
        fesa::assemble_serial(domain, dofs).stiffness;

    ASSERT_EQ(matrix.order, 18);
    ASSERT_EQ(matrix.row_offsets.size(), matrix.order + 1);
    ASSERT_EQ(matrix.row_offsets.front(), 0);
    ASSERT_EQ(
        static_cast<std::size_t>(matrix.row_offsets.back()),
        matrix.column_indices.size());
    ASSERT_EQ(matrix.column_indices.size(), matrix.values.size());

    for (std::size_t row = 0; row < matrix.order; ++row) {
        const auto begin =
            matrix.column_indices.begin() + matrix.row_offsets[row];
        const auto end =
            matrix.column_indices.begin() + matrix.row_offsets[row + 1];
        EXPECT_TRUE(std::ranges::is_sorted(begin, end));
        EXPECT_EQ(std::adjacent_find(begin, end), end);
        for (auto entry = begin; entry != end; ++entry) {
            EXPECT_GE(*entry, static_cast<std::int32_t>(row));
            EXPECT_LT(*entry, static_cast<std::int32_t>(matrix.order));
        }
    }
}

TEST(SparsePattern, IncludesZeroDiagonalForEveryOrphanNodeDof) {
    const fesa::Domain domain = build_orphan_node_domain();
    const fesa::SymmetricCsr matrix =
        fesa::assemble_serial(
            domain, fesa::DofManager::build(domain))
            .stiffness;

    ASSERT_EQ(matrix.order, 18);
    for (std::size_t row = 0; row < matrix.order; ++row) {
        const auto begin =
            matrix.column_indices.begin() + matrix.row_offsets[row];
        const auto end =
            matrix.column_indices.begin() + matrix.row_offsets[row + 1];
        const auto diagonal = std::lower_bound(
            begin, end, static_cast<std::int32_t>(row));
        ASSERT_NE(diagonal, end);
        EXPECT_EQ(*diagonal, static_cast<std::int32_t>(row));
    }
    for (std::size_t row = 12; row < 18; ++row) {
        EXPECT_DOUBLE_EQ(csr_value(matrix, row, row), 0.0);
    }
}

TEST(SerialAssembly, AssemblesHandCalculatedAxialChainAndFullLoad) {
    const fesa::Domain domain = build_chain_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    const fesa::EquationSystem system =
        fesa::assemble_serial(domain, dofs);

    EXPECT_DOUBLE_EQ(csr_value(system.stiffness, 0, 0), 1.0);
    EXPECT_DOUBLE_EQ(csr_value(system.stiffness, 0, 6), -1.0);
    EXPECT_DOUBLE_EQ(csr_value(system.stiffness, 6, 6), 2.0);
    EXPECT_DOUBLE_EQ(csr_value(system.stiffness, 6, 12), -1.0);
    EXPECT_DOUBLE_EQ(csr_value(system.stiffness, 12, 12), 1.0);

    EXPECT_EQ(
        system.force,
        (std::vector<double>{
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            0.0,
            12.0,
            0.0,
            4.0,
            5.0,
            6.0,
            3.0,
            4.0,
            5.0,
            6.0,
            7.0,
            8.0,
        }));
}

TEST(SerialAssembly, MergesDuplicateElementContributions) {
    const fesa::Domain single_domain = build_parallel_domain({10});
    const fesa::Domain duplicate_domain =
        build_parallel_domain({30, 10, 20});
    const fesa::EquationSystem single = fesa::assemble_serial(
        single_domain, fesa::DofManager::build(single_domain));
    const fesa::EquationSystem duplicate = fesa::assemble_serial(
        duplicate_domain, fesa::DofManager::build(duplicate_domain));

    EXPECT_EQ(
        duplicate.stiffness.row_offsets,
        single.stiffness.row_offsets);
    EXPECT_EQ(
        duplicate.stiffness.column_indices,
        single.stiffness.column_indices);
    ASSERT_EQ(duplicate.stiffness.values.size(), single.stiffness.values.size());
    for (std::size_t index = 0; index < single.stiffness.values.size();
         ++index) {
        const double value = single.stiffness.values[index];
        EXPECT_DOUBLE_EQ(
            duplicate.stiffness.values[index], (value + value) + value);
    }
}

TEST(SerialAssembly, ReducesByElementOriginNotDomainStorageOrder) {
    const fesa::Domain origin_order_domain =
        build_rounding_domain(false);
    const fesa::Domain storage_order_domain =
        build_rounding_domain(true);

    const double origin_order_value = csr_value(
        fesa::assemble_serial(
            origin_order_domain,
            fesa::DofManager::build(origin_order_domain))
            .stiffness,
        0,
        0);
    const double storage_order_value = csr_value(
        fesa::assemble_serial(
            storage_order_domain,
            fesa::DofManager::build(storage_order_domain))
            .stiffness,
        0,
        0);
    const double expected = (1.0 + 1.0) + 1.0e16;

    EXPECT_EQ(
        std::bit_cast<std::uint64_t>(origin_order_value),
        std::bit_cast<std::uint64_t>(expected));
    EXPECT_EQ(
        std::bit_cast<std::uint64_t>(storage_order_value),
        std::bit_cast<std::uint64_t>(expected));
}

TEST(SerialAssembly, IsIndependentOfDomainStorageAndExternalLabelOrder) {
    const fesa::Domain first_domain = build_chain_domain(false);
    const fesa::Domain second_domain = build_chain_domain(true);

    const fesa::EquationSystem first = fesa::assemble_serial(
        first_domain, fesa::DofManager::build(first_domain));
    const fesa::EquationSystem second = fesa::assemble_serial(
        second_domain, fesa::DofManager::build(second_domain));

    EXPECT_EQ(first.stiffness.order, second.stiffness.order);
    EXPECT_EQ(
        first.stiffness.row_offsets, second.stiffness.row_offsets);
    EXPECT_EQ(
        first.stiffness.column_indices,
        second.stiffness.column_indices);
    EXPECT_EQ(first.stiffness.values, second.stiffness.values);
    EXPECT_EQ(first.force, second.force);
}

TEST(SerialAssembly, PropagatesBeamKernelFailure) {
    fesa::DomainBuilder builder;
    builder.add_node({
        fesa::NodeId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        fesa::Vec3{0.0, 0.0, 0.0},
    });
    builder.add_node({
        fesa::NodeId{1},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 2},
        fesa::Vec3{1.0, 0.0, 0.0},
    });
    auto material = unit_material();
    material.young = std::numeric_limits<double>::max();
    builder.add_material(std::move(material));
    auto section = unit_section();
    section.area = std::numeric_limits<double>::max();
    builder.add_section(std::move(section));
    builder.add_beam_element({
        fesa::ElementId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 1},
        {fesa::NodeId{0}, fesa::NodeId{1}},
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });
    const fesa::Domain domain =
        finish_domain(std::move(builder), {"Load", {}, {}});

    try {
        static_cast<void>(fesa::assemble_serial(
            domain, fesa::DofManager::build(domain)));
        FAIL() << "Expected a Beam kernel failure.";
    } catch (const std::runtime_error& error) {
        EXPECT_NE(
            std::string{error.what()}.find("model.nonfinite_value"),
            std::string::npos);
    }
}

}  // namespace