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

#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::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_beam_domain(const bool branched) {
    fesa::DomainBuilder builder;
    const std::array<fesa::Vec3, 6> positions{{
        {0.0, 0.0, 0.0},
        {1.0, 0.0, 0.0},
        {2.0, 0.0, 0.0},
        {3.0, 0.0, 0.0},
        {1.0, 0.0, 1.0},
        {1.0, 1.0, 1.0},
    }};
    const std::size_t node_count = branched ? positions.size() : 4;
    for (std::size_t index = 0; index < node_count; ++index) {
        builder.add_node({
            fesa::NodeId{static_cast<std::int64_t>(index)},
            fesa::EntityOrigin{
                "BeamPart",
                "Beam-1",
                static_cast<std::int64_t>(index + 1),
            },
            positions[index],
        });
    }
    builder.add_material({
        fesa::MaterialId{0},
        "Steel",
        210.0e9,
        0.3,
    });
    builder.add_section({
        fesa::SectionId{0},
        "General",
        0.02,
        3.0e-5,
        4.0e-5,
        2.0e-5,
        0.015,
        0.016,
        fesa::ShearPropertySource::input,
        fesa::Vec3{0.0, 1.0, 0.0},
        {},
    });

    const std::array<fesa::BeamElement, 5> elements{{
        {
            fesa::ElementId{40},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 300},
            {fesa::NodeId{2}, fesa::NodeId{3}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
        {
            fesa::ElementId{10},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 100},
            {fesa::NodeId{0}, fesa::NodeId{1}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
        {
            fesa::ElementId{30},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 200},
            {fesa::NodeId{1}, fesa::NodeId{2}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
        {
            fesa::ElementId{0},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 500},
            {fesa::NodeId{1}, fesa::NodeId{5}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
        {
            fesa::ElementId{20},
            fesa::EntityOrigin{"BeamPart", "Beam-1", 400},
            {fesa::NodeId{1}, fesa::NodeId{4}},
            fesa::MaterialId{0},
            fesa::SectionId{0},
        },
    }};
    const std::size_t element_count = branched ? elements.size() : 3;
    for (std::size_t index = 0; index < element_count; ++index) {
        builder.add_beam_element(elements[index]);
    }

    return finish_domain(
        std::move(builder),
        {
            "Load",
            {},
            {
                {fesa::NodeId{3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}},
                {fesa::NodeId{1}, {1.0e16, 1.0, 0.0, 0.0, 0.0, 0.0}},
                {fesa::NodeId{1}, {-1.0e16, 2.0, 0.0, 0.0, 0.0, 0.0}},
            },
        });
}

fesa::Domain build_multiple_kernel_failure_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},
    });

    auto material = fesa::IsotropicElastic{
        fesa::MaterialId{0},
        "Invalid",
        std::numeric_limits<double>::max(),
        0.3,
    };
    builder.add_material(std::move(material));
    auto section = fesa::BeamSection{
        fesa::SectionId{0},
        "Invalid",
        std::numeric_limits<double>::max(),
        1.0,
        1.0,
        1.0,
        1.0,
        1.0,
        fesa::ShearPropertySource::input,
        fesa::Vec3{0.0, 1.0, 0.0},
        {},
    };
    builder.add_section(std::move(section));

    builder.add_beam_element({
        fesa::ElementId{0},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 20},
        {fesa::NodeId{0}, fesa::NodeId{1}},
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });
    builder.add_beam_element({
        fesa::ElementId{1},
        fesa::EntityOrigin{"BeamPart", "Beam-1", 10},
        {fesa::NodeId{0}, fesa::NodeId{1}},
        fesa::MaterialId{0},
        fesa::SectionId{0},
    });

    return finish_domain(std::move(builder), {"Load", {}, {}});
}

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

void expect_bitwise_equal(
    const fesa::EquationSystem& expected,
    const fesa::EquationSystem& actual) {
    EXPECT_EQ(actual.stiffness.order, expected.stiffness.order);
    EXPECT_EQ(
        actual.stiffness.row_offsets,
        expected.stiffness.row_offsets);
    EXPECT_EQ(
        actual.stiffness.column_indices,
        expected.stiffness.column_indices);
    EXPECT_EQ(bits(actual.stiffness.values), bits(expected.stiffness.values));
    EXPECT_EQ(bits(actual.force), bits(expected.force));
}

TEST(ParallelAssembly, MatchesSerialForFixedBeamChain) {
    const fesa::Domain domain = build_beam_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);
    const fesa::EquationSystem serial = fesa::assemble_serial(domain, dofs);

    expect_bitwise_equal(
        serial,
        fesa::assemble_parallel(domain, dofs, {1, 1}));
    expect_bitwise_equal(
        serial,
        fesa::assemble_parallel(domain, dofs, {4, 2}));
}

TEST(ParallelAssembly, MatchesSerialForFixedBranchedDomain) {
    const fesa::Domain domain = build_beam_domain(true);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);
    const fesa::EquationSystem serial = fesa::assemble_serial(domain, dofs);

    expect_bitwise_equal(
        serial,
        fesa::assemble_parallel(domain, dofs, {1, 2}));
    expect_bitwise_equal(
        serial,
        fesa::assemble_parallel(domain, dofs, {3, 1}));
}

TEST(TbbElementEvaluation, RejectsZeroExecutionLimits) {
    const fesa::Domain domain = build_beam_domain(false);
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    EXPECT_THROW(
        static_cast<void>(
            fesa::assemble_parallel(domain, dofs, {0, 1})),
        std::invalid_argument);
    EXPECT_THROW(
        static_cast<void>(
            fesa::assemble_parallel(domain, dofs, {1, 0})),
        std::invalid_argument);
}

TEST(TbbElementEvaluation, ReportsFirstFailureByCanonicalElementOrigin) {
    const fesa::Domain domain = build_multiple_kernel_failure_domain();
    const fesa::DofManager dofs = fesa::DofManager::build(domain);

    for (const std::size_t threads : {std::size_t{1}, std::size_t{4}}) {
        SCOPED_TRACE(::testing::Message{} << "threads=" << threads);
        try {
            static_cast<void>(
                fesa::assemble_parallel(domain, dofs, {threads, 1}));
            FAIL() << "Expected a Beam kernel failure.";
        } catch (const std::runtime_error& error) {
            const std::string message{error.what()};
            EXPECT_NE(
                message.find("Beam element 10 kernel failed"),
                std::string::npos);
            EXPECT_EQ(
                message.find("Beam element 20 kernel failed"),
                std::string::npos);
        }
    }
}

}  // namespace