#include <fesa/assembly/serial_assembler.hpp>

#include <algorithm>
#include <array>
#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", {}, {}});
}

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))];
}

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(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, 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