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

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

#include <fesa/constraints/essential_bc.hpp>
#include <fesa/fem/dof_manager.hpp>
#include <fesa/model/domain_builder.hpp>
#include <fesa/solvers/linear/pardiso_linear_solver.hpp>

#include <gtest/gtest.h>

namespace {

constexpr std::size_t branch_count = 24;
constexpr std::size_t repeat_count = 10;

fesa::Domain build_branched_cantilever() {
    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},
    });
    for (std::size_t index = 0; index < branch_count; ++index) {
        const double y = static_cast<double>(
            static_cast<int>(index % 7) - 3) * 0.35;
        const double z = static_cast<double>(
            static_cast<int>((index * 3) % 11) - 5) * 0.22 + 0.1;
        builder.add_node({
            fesa::NodeId{static_cast<std::int64_t>(index + 2)},
            fesa::EntityOrigin{
                "BeamPart",
                "Beam-1",
                static_cast<std::int64_t>(index + 3),
            },
            fesa::Vec3{
                2.0 + static_cast<double>(index) * 0.05,
                y,
                z,
            },
        });
    }
    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},
        {},
    });

    for (std::size_t offset = 0; offset < branch_count; ++offset) {
        const std::size_t index = branch_count - offset - 1;
        builder.add_beam_element({
            fesa::ElementId{static_cast<std::int64_t>(index + 1)},
            fesa::EntityOrigin{
                "BeamPart",
                "Beam-1",
                static_cast<std::int64_t>(index + 2),
            },
            {
                fesa::NodeId{1},
                fesa::NodeId{static_cast<std::int64_t>(index + 2)},
            },
            fesa::MaterialId{0},
            fesa::SectionId{0},
        });
    }
    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},
    });

    std::vector<fesa::PrescribedDof> fixed;
    fixed.reserve(6);
    for (std::uint8_t dof = 1; dof <= 6; ++dof) {
        fixed.push_back({fesa::NodeId{0}, dof, 0.0});
    }
    builder.set_step({
        "Load",
        std::move(fixed),
        {{
            fesa::NodeId{static_cast<std::int64_t>(branch_count + 1)},
            {1000.0, -2500.0, 1750.0, 20.0, -30.0, 40.0},
        }},
    });

    auto result = std::move(builder).build();
    if (!result.domain.has_value()) {
        throw std::runtime_error{"Test Domain failed validation."};
    }
    return std::move(*result.domain);
}

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

struct LinearState final {
    std::vector<double> displacement;
    std::vector<double> reaction;
};

LinearState solve_linear_system(
    const fesa::EquationSystem& system,
    const fesa::DofManager& dofs) {
    fesa::ConstraintResult constrained =
        fesa::eliminate_essential_bcs(system, dofs);
    if (!constrained.reduced_system.has_value()) {
        throw std::runtime_error{"Constraint elimination failed."};
    }

    fesa::PardisoLinearSolver solver;
    fesa::LinearSolveResult solved = solver.solve(
        constrained.reduced_system->stiffness,
        constrained.reduced_system->force);
    if (!solved.diagnostics.empty()) {
        throw std::runtime_error{"Linear solve failed."};
    }

    std::vector<double> displacement =
        dofs.reconstruct_full(solved.solution);
    return {
        displacement,
        fesa::recover_reaction(system, displacement),
    };
}

std::size_t available_concurrency() {
    return std::max(
        std::size_t{1},
        static_cast<std::size_t>(std::thread::hardware_concurrency()));
}

TEST(ThreadCountDeterminism, AssemblyAndLinearStateAreBitwiseStable) {
    const fesa::Domain domain = build_branched_cantilever();
    const fesa::DofManager dofs = fesa::DofManager::build(domain);
    const fesa::EquationSystem serial = fesa::assemble_serial(domain, dofs);
    const LinearState serial_state = solve_linear_system(serial, dofs);
    const std::array<std::size_t, 3> thread_counts{
        1,
        2,
        available_concurrency(),
    };

    for (std::size_t repeat = 0; repeat < repeat_count; ++repeat) {
        for (const std::size_t threads : thread_counts) {
            SCOPED_TRACE(::testing::Message{}
                         << "repeat=" << repeat << ", threads=" << threads);
            const fesa::EquationSystem parallel =
                fesa::assemble_parallel(domain, dofs, {threads, 1});
            expect_bitwise_equal(serial, parallel);

            const LinearState parallel_state =
                solve_linear_system(parallel, dofs);
            EXPECT_EQ(
                bits(parallel_state.displacement),
                bits(serial_state.displacement));
            EXPECT_EQ(
                bits(parallel_state.reaction),
                bits(serial_state.reaction));
        }
    }
}

}  // namespace