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

#include <bit>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <iostream>
#include <stdexcept>
#include <thread>
#include <utility>
#include <vector>

#include <fesa/fem/dof_manager.hpp>
#include <fesa/model/domain_builder.hpp>

namespace {

constexpr std::size_t benchmark_element_count = 1000;

fesa::Domain build_beam_chain() {
    fesa::DomainBuilder builder;
    for (std::size_t index = 0; index <= benchmark_element_count; ++index) {
        builder.add_node({
            fesa::NodeId{static_cast<std::int64_t>(index)},
            fesa::EntityOrigin{
                "BenchmarkPart",
                "Benchmark-1",
                static_cast<std::int64_t>(index + 1),
            },
            fesa::Vec3{static_cast<double>(index), 0.0, 0.0},
        });
    }
    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 index = 0; index < benchmark_element_count; ++index) {
        builder.add_beam_element({
            fesa::ElementId{static_cast<std::int64_t>(index)},
            fesa::EntityOrigin{
                "BenchmarkPart",
                "Benchmark-1",
                static_cast<std::int64_t>(index + 1),
            },
            {
                fesa::NodeId{static_cast<std::int64_t>(index)},
                fesa::NodeId{static_cast<std::int64_t>(index + 1)},
            },
            fesa::MaterialId{0},
            fesa::SectionId{0},
        });
    }
    builder.set_step({"Benchmark", {}, {}});

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

template <typename Operation>
std::pair<fesa::EquationSystem, double> measure(Operation&& operation) {
    const auto start = std::chrono::steady_clock::now();
    fesa::EquationSystem result = std::forward<Operation>(operation)();
    const auto finish = std::chrono::steady_clock::now();
    const double milliseconds =
        std::chrono::duration<double, std::milli>(finish - start).count();
    return {std::move(result), milliseconds};
}

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

bool same_system(
    const fesa::EquationSystem& left,
    const fesa::EquationSystem& right) {
    return left.stiffness.order == right.stiffness.order &&
           left.stiffness.row_offsets == right.stiffness.row_offsets &&
           left.stiffness.column_indices == right.stiffness.column_indices &&
           bits(left.stiffness.values) == bits(right.stiffness.values) &&
           bits(left.force) == bits(right.force);
}

std::size_t available_concurrency() {
    const std::size_t available =
        static_cast<std::size_t>(std::thread::hardware_concurrency());
    return available == 0 ? 1 : available;
}

}  // namespace

int main() {
    try {
        const fesa::Domain domain = build_beam_chain();
        const fesa::DofManager dofs = fesa::DofManager::build(domain);
        const std::size_t threads = available_concurrency();

        auto [serial, serial_ms] = measure(
            [&] { return fesa::assemble_serial(domain, dofs); });
        auto [parallel, parallel_ms] = measure([&] {
            return fesa::assemble_parallel(domain, dofs, {threads, 1});
        });

        if (!same_system(serial, parallel)) {
            std::cerr << "Serial and parallel assembly results differ.\n";
            return 1;
        }

        std::cout << "elements=" << domain.beam_elements().size()
                  << " serial_ms=" << serial_ms
                  << " parallel_ms=" << parallel_ms
                  << " parallel_threads=" << threads << '\n';
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "Assembly benchmark failed: " << error.what() << '\n';
        return 1;
    }
}