// Throughput benchmark: items/second vs. graph topology and size.
//
// Topologies:
//   chain   — linear depth D: push → n[0..D-1] → pop
//   wide    — fanout<W>:      push → fanout → W parallel nodes → W pops
//   diamond — push → fanout<2> → 2×2 nodes → 2 pops
//
// Two scheduling modes for each topology:
//   private — each node owns a private ThreadPool(1)  [Node<>]
//   pool    — all nodes share one ThreadPool(T)        [PoolNode<> + shared pool]
//
// Usage: ./bench_pipeline | tee results.csv

#include <kpn/kpn.hpp>

#ifdef KPN_BENCH_TBB
#include <oneapi/tbb/flow_graph.h>
namespace tbb_flow = oneapi::tbb::flow;
#endif

#include <array>
#include <atomic>
#include <chrono>
#include <cstdio>
#include <memory>
#include <string>
#include <thread>
#include <vector>

using namespace kpn;
using namespace std::chrono_literals;
using sclock = std::chrono::steady_clock;

// ── configurable work ─────────────────────────────────────────────────────────

static std::atomic<int> g_work_us{0};

static int chain_fn(int x) {
    int us = g_work_us.load(std::memory_order_relaxed);
    if (us > 0) {
        auto end = sclock::now() + std::chrono::microseconds(us);
        while (sclock::now() < end);
    }
    return x;
}

using ChainNode     = Node<chain_fn, in<>, out<>>;
using PoolChainNode = PoolNode<chain_fn, in<>, out<>>;

// ── push helper: yield-spin on overflow (no artificial sleep latency) ─────────

static void push_retry(Channel<int>& ch, int val) {
    while (true) {
        try   { ch.push(val); return; }
        catch (const ChannelOverflowError&) { std::this_thread::yield(); }
        catch (const ChannelClosedError&)   { return; }
    }
}

// ── result ────────────────────────────────────────────────────────────────────

struct Result {
    const char* topology;
    int         size;
    int         work_us;
    int         threads;     // 0 = private (1 thread per node), N = shared pool size
    double      items_per_sec;
    double      overhead_us;
};

// ── chain ─────────────────────────────────────────────────────────────────────

static int items_for(int work_us, int depth = 1) {
    int effective = std::max(1, work_us) * std::max(1, depth);
    if (effective <= 1)     return 5000;
    if (effective <= 10)    return 3000;
    if (effective <= 100)   return 1000;
    if (effective <= 1000)  return 200;
    return 50;
}

static Result bench_chain(int depth, int work_us) {
    const int N   = items_for(work_us, depth);
    const int CAP = N;

    std::vector<std::shared_ptr<Channel<int>>> chs;
    for (int i = 0; i <= depth; ++i)
        chs.push_back(std::make_shared<Channel<int>>(CAP));

    std::vector<std::unique_ptr<ChainNode>> nodes;
    for (int i = 0; i < depth; ++i) {
        nodes.push_back(std::make_unique<ChainNode>(CAP));
        nodes.back()->set_input_channel<0>(chs[i]);
        nodes.back()->set_output_channel<0>(chs[i + 1].get());
    }

    for (auto& n : nodes) n->start();

    std::atomic<sclock::time_point> t1;
    std::thread reader([&] {
        for (int i = 0; i < N; ++i) chs.back()->pop();
        t1.store(sclock::now(), std::memory_order_release);
    });

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*chs[0], i);
    });

    pusher.join();
    reader.join();
    for (auto& n : nodes) n->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    // Subtract theoretical pipeline fill cost (depth-1)*W so that overhead
    // reflects only framework latency, not the expected pipeline startup time.
    double pipeline_us = static_cast<double>(work_us) * (N + depth - 1);
    double wus = (elapsed * 1e6 - pipeline_us) / N;
    return {"chain", depth, work_us, 0, N / elapsed, wus};
}

static Result bench_chain_pool(int depth, int work_us, int pool_threads) {
    const int N   = items_for(work_us, depth);
    const int CAP = N;

    auto pool = std::make_shared<ThreadPool>(pool_threads);

    std::vector<std::shared_ptr<Channel<int>>> chs;
    for (int i = 0; i <= depth; ++i)
        chs.push_back(std::make_shared<Channel<int>>(CAP));

    std::vector<std::unique_ptr<PoolChainNode>> nodes;
    for (int i = 0; i < depth; ++i) {
        nodes.push_back(std::make_unique<PoolChainNode>(pool, CAP));
        nodes.back()->set_input_channel<0>(chs[i]);
        nodes.back()->set_output_channel<0>(chs[i + 1].get());
    }

    pool->start();
    for (auto& n : nodes) n->start();

    std::atomic<sclock::time_point> t1;
    std::thread reader([&] {
        for (int i = 0; i < N; ++i) chs.back()->pop();
        t1.store(sclock::now(), std::memory_order_release);
    });

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*chs[0], i);
    });

    pusher.join();
    reader.join();
    for (auto& n : nodes) n->stop();
    pool->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    double pipeline_us = static_cast<double>(work_us) * (N + depth - 1);
    double wus = (elapsed * 1e6 - pipeline_us) / N;
    return {"chain", depth, work_us, pool_threads, N / elapsed, wus};
}

// ── wide (fanout<W>) ──────────────────────────────────────────────────────────

template<std::size_t W>
static Result bench_wide(int work_us) {
    const int N   = items_for(work_us);
    const int CAP = N;

    auto src_ch = std::make_shared<Channel<int>>(CAP);
    auto fan    = std::make_unique<FanoutNode<int, W>>(CAP);
    fan->template set_input_channel<0>(src_ch);

    std::array<std::unique_ptr<ChainNode>,    W> nodes;
    std::array<std::shared_ptr<Channel<int>>, W> sink_chs;

    for (std::size_t i = 0; i < W; ++i) {
        nodes[i]    = std::make_unique<ChainNode>(CAP);
        sink_chs[i] = std::make_shared<Channel<int>>(CAP);
        nodes[i]->template set_output_channel<0>(sink_chs[i].get());
    }

    [&]<std::size_t... Is>(std::index_sequence<Is...>) {
        (fan->template set_output_channel<Is>(
             &nodes[Is]->template input_channel<0>()), ...);
    }(std::make_index_sequence<W>{});

    fan->start();
    for (auto& n : nodes) n->start();

    std::array<std::thread, W> readers;
    std::atomic<sclock::time_point> t1;
    std::atomic<int> readers_done{0};

    for (std::size_t w = 0; w < W; ++w) {
        readers[w] = std::thread([&, w] {
            for (int i = 0; i < N; ++i) sink_chs[w]->pop();
            if (readers_done.fetch_add(1, std::memory_order_acq_rel) + 1
                    == static_cast<int>(W))
                t1.store(sclock::now(), std::memory_order_release);
        });
    }

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*src_ch, i);
    });

    pusher.join();
    for (auto& r : readers) r.join();
    fan->stop();
    for (auto& n : nodes) n->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    double wus = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"wide", static_cast<int>(W), work_us, 0, N / elapsed, wus};
}

template<std::size_t W>
static Result bench_wide_pool(int work_us, int pool_threads) {
    const int N   = items_for(work_us);
    const int CAP = N;

    auto pool   = std::make_shared<ThreadPool>(pool_threads);
    auto src_ch = std::make_shared<Channel<int>>(CAP);
    auto fan    = std::make_unique<FanoutNode<int, W>>(CAP);
    fan->template set_input_channel<0>(src_ch);

    std::array<std::unique_ptr<PoolChainNode>, W> nodes;
    std::array<std::shared_ptr<Channel<int>>,  W> sink_chs;

    for (std::size_t i = 0; i < W; ++i) {
        nodes[i]    = std::make_unique<PoolChainNode>(pool, CAP);
        sink_chs[i] = std::make_shared<Channel<int>>(CAP);
        nodes[i]->template set_output_channel<0>(sink_chs[i].get());
    }

    [&]<std::size_t... Is>(std::index_sequence<Is...>) {
        (fan->template set_output_channel<Is>(
             &nodes[Is]->template input_channel<0>()), ...);
    }(std::make_index_sequence<W>{});

    fan->start();
    pool->start();
    for (auto& n : nodes) n->start();

    std::array<std::thread, W> readers;
    std::atomic<sclock::time_point> t1;
    std::atomic<int> readers_done{0};

    for (std::size_t w = 0; w < W; ++w) {
        readers[w] = std::thread([&, w] {
            for (int i = 0; i < N; ++i) sink_chs[w]->pop();
            if (readers_done.fetch_add(1, std::memory_order_acq_rel) + 1
                    == static_cast<int>(W))
                t1.store(sclock::now(), std::memory_order_release);
        });
    }

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*src_ch, i);
    });

    pusher.join();
    for (auto& r : readers) r.join();
    fan->stop();
    for (auto& n : nodes) n->stop();
    pool->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    double wus = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"wide", static_cast<int>(W), work_us, pool_threads, N / elapsed, wus};
}

// ── diamond ───────────────────────────────────────────────────────────────────

static Result bench_diamond(int work_us) {
    const int N   = items_for(work_us, 2);
    const int CAP = N;

    auto src_ch = std::make_shared<Channel<int>>(CAP);
    auto fan    = std::make_unique<FanoutNode<int, 2>>(CAP);
    fan->template set_input_channel<0>(src_ch);

    auto nL  = std::make_unique<ChainNode>(CAP);
    auto nR  = std::make_unique<ChainNode>(CAP);
    auto nL2 = std::make_unique<ChainNode>(CAP);
    auto nR2 = std::make_unique<ChainNode>(CAP);
    auto chL  = std::make_shared<Channel<int>>(CAP);
    auto chR  = std::make_shared<Channel<int>>(CAP);
    auto snkL = std::make_shared<Channel<int>>(CAP);
    auto snkR = std::make_shared<Channel<int>>(CAP);

    fan->template set_output_channel<0>(&nL->template input_channel<0>());
    fan->template set_output_channel<1>(&nR->template input_channel<0>());
    nL->set_output_channel<0>(chL.get());
    nR->set_output_channel<0>(chR.get());
    nL2->set_input_channel<0>(chL);
    nR2->set_input_channel<0>(chR);
    nL2->set_output_channel<0>(snkL.get());
    nR2->set_output_channel<0>(snkR.get());

    fan->start(); nL->start(); nR->start(); nL2->start(); nR2->start();

    std::atomic<sclock::time_point> t1;
    std::atomic<int> done{0};
    auto make_reader = [&](Channel<int>& ch) {
        return std::thread([&] {
            for (int i = 0; i < N; ++i) ch.pop();
            if (done.fetch_add(1, std::memory_order_acq_rel) + 1 == 2)
                t1.store(sclock::now(), std::memory_order_release);
        });
    };
    auto rL = make_reader(*snkL);
    auto rR = make_reader(*snkR);

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*src_ch, i);
    });

    pusher.join(); rL.join(); rR.join();
    fan->stop(); nL->stop(); nR->stop(); nL2->stop(); nR2->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    double wus = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"diamond", 4, work_us, 0, N / elapsed, wus};
}

static Result bench_diamond_pool(int work_us, int pool_threads) {
    const int N   = items_for(work_us, 2);
    const int CAP = N;

    auto pool   = std::make_shared<ThreadPool>(pool_threads);
    auto src_ch = std::make_shared<Channel<int>>(CAP);
    auto fan    = std::make_unique<FanoutNode<int, 2>>(CAP);
    fan->template set_input_channel<0>(src_ch);

    auto nL  = std::make_unique<PoolChainNode>(pool, CAP);
    auto nR  = std::make_unique<PoolChainNode>(pool, CAP);
    auto nL2 = std::make_unique<PoolChainNode>(pool, CAP);
    auto nR2 = std::make_unique<PoolChainNode>(pool, CAP);
    auto chL  = std::make_shared<Channel<int>>(CAP);
    auto chR  = std::make_shared<Channel<int>>(CAP);
    auto snkL = std::make_shared<Channel<int>>(CAP);
    auto snkR = std::make_shared<Channel<int>>(CAP);

    fan->template set_output_channel<0>(&nL->template input_channel<0>());
    fan->template set_output_channel<1>(&nR->template input_channel<0>());
    nL->set_output_channel<0>(chL.get());
    nR->set_output_channel<0>(chR.get());
    nL2->set_input_channel<0>(chL);
    nR2->set_input_channel<0>(chR);
    nL2->set_output_channel<0>(snkL.get());
    nR2->set_output_channel<0>(snkR.get());

    fan->start();
    pool->start();
    nL->start(); nR->start(); nL2->start(); nR2->start();

    std::atomic<sclock::time_point> t1;
    std::atomic<int> done{0};
    auto make_reader = [&](Channel<int>& ch) {
        return std::thread([&] {
            for (int i = 0; i < N; ++i) ch.pop();
            if (done.fetch_add(1, std::memory_order_acq_rel) + 1 == 2)
                t1.store(sclock::now(), std::memory_order_release);
        });
    };
    auto rL = make_reader(*snkL);
    auto rR = make_reader(*snkR);

    auto t0 = sclock::now();
    std::thread pusher([&] {
        for (int i = 0; i < N; ++i) push_retry(*src_ch, i);
    });

    pusher.join(); rL.join(); rR.join();
    fan->stop();
    nL->stop(); nR->stop(); nL2->stop(); nR2->stop();
    pool->stop();

    double elapsed = std::chrono::duration<double>(
        t1.load(std::memory_order_acquire) - t0).count();
    double wus = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"diamond", 4, work_us, pool_threads, N / elapsed, wus};
}

// ── TBB flow graph ────────────────────────────────────────────────────────────
#ifdef KPN_BENCH_TBB

static Result bench_chain_tbb(int depth, int work_us) {
    const int N = items_for(work_us, depth);

    tbb_flow::graph g;
    using FN = tbb_flow::function_node<int, int>;
    std::vector<std::unique_ptr<FN>> nodes;
    nodes.reserve(depth);
    for (int i = 0; i < depth; ++i)
        nodes.push_back(std::make_unique<FN>(g, tbb_flow::serial,
            [](int x) -> int { return chain_fn(x); }));
    for (int i = 0; i + 1 < depth; ++i)
        tbb_flow::make_edge(*nodes[i], *nodes[i + 1]);

    auto t0 = sclock::now();
    for (int i = 0; i < N; ++i) nodes[0]->try_put(i);
    g.wait_for_all();
    auto t1 = sclock::now();

    double elapsed    = std::chrono::duration<double>(t1 - t0).count();
    double pipeline_us = static_cast<double>(work_us) * (N + depth - 1);
    double wus        = (elapsed * 1e6 - pipeline_us) / N;
    return {"chain_tbb", depth, work_us, -1, N / elapsed, wus};
}

template<std::size_t W>
static Result bench_wide_tbb(int work_us) {
    const int N = items_for(work_us);

    tbb_flow::graph g;
    tbb_flow::broadcast_node<int> fan(g);
    using FN = tbb_flow::function_node<int, int>;
    std::array<std::unique_ptr<FN>, W> nodes;
    for (auto& n : nodes) {
        n = std::make_unique<FN>(g, tbb_flow::serial,
            [](int x) -> int { return chain_fn(x); });
        tbb_flow::make_edge(fan, *n);
    }

    auto t0 = sclock::now();
    for (int i = 0; i < N; ++i) fan.try_put(i);
    g.wait_for_all();
    auto t1 = sclock::now();

    double elapsed = std::chrono::duration<double>(t1 - t0).count();
    double wus     = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"wide_tbb", static_cast<int>(W), work_us, -1, N / elapsed, wus};
}

static Result bench_diamond_tbb(int work_us) {
    const int N = items_for(work_us, 2);

    tbb_flow::graph g;
    tbb_flow::broadcast_node<int> fan(g);
    using FN = tbb_flow::function_node<int, int>;
    auto fn = [](int x) -> int { return chain_fn(x); };
    FN nL(g, tbb_flow::serial, fn), nR(g, tbb_flow::serial, fn);
    FN nL2(g, tbb_flow::serial, fn), nR2(g, tbb_flow::serial, fn);
    tbb_flow::make_edge(fan, nL);  tbb_flow::make_edge(fan, nR);
    tbb_flow::make_edge(nL, nL2);  tbb_flow::make_edge(nR, nR2);

    auto t0 = sclock::now();
    for (int i = 0; i < N; ++i) fan.try_put(i);
    g.wait_for_all();
    auto t1 = sclock::now();

    double elapsed = std::chrono::duration<double>(t1 - t0).count();
    double wus     = (elapsed * 1e6) / N - static_cast<double>(work_us);
    return {"diamond_tbb", 4, work_us, -1, N / elapsed, wus};
}
#endif // KPN_BENCH_TBB

// ── main ──────────────────────────────────────────────────────────────────────

int main() {
    const int work_amts[]  = {10, 100, 1000};
    const int pool_sizes[] = {1, 2, 4};

    std::fprintf(stderr, "%-12s %-8s %-10s %-8s %-18s %-20s\n",
                 "topology", "size", "work_us", "threads", "items/sec", "overhead_us/item");
    std::fprintf(stderr, "%s\n", std::string(78, '-').c_str());
    std::printf("topology,size,work_us,threads,items_per_sec,overhead_us_per_item\n");

    auto emit = [](const Result& r) {
        std::string sched = r.threads  < 0 ? "tbb"
                          : r.threads == 0 ? "priv"
                          : std::to_string(r.threads);
        std::fprintf(stderr, "%-12s %-8d %-10d %-8s %-18.0f %-20.1f\n",
                     r.topology, r.size, r.work_us, sched.c_str(),
                     r.items_per_sec, r.overhead_us);
        std::printf("%s,%d,%d,%s,%.0f,%.2f\n",
                    r.topology, r.size, r.work_us, sched.c_str(),
                    r.items_per_sec, r.overhead_us);
        std::fflush(stdout);
    };

    for (int w : work_amts) {
        g_work_us.store(w, std::memory_order_relaxed);
        std::fprintf(stderr, "\n── work_us=%-4d private pools ───────────────────────────────────────\n", w);

        for (int d : {1, 2, 4, 8, 16, 32}) emit(bench_chain(d, w));
        emit(bench_wide<1>(w));
        emit(bench_wide<2>(w));
        emit(bench_wide<3>(w));
        emit(bench_wide<4>(w));
        emit(bench_diamond(w));

        for (int pt : pool_sizes) {
            std::fprintf(stderr, "\n── work_us=%-4d shared pool (%d thread%s) ─────────────────────────────\n",
                         w, pt, pt == 1 ? "" : "s");
            for (int d : {1, 2, 4, 8, 16, 32}) emit(bench_chain_pool(d, w, pt));
            emit(bench_wide_pool<1>(w, pt));
            emit(bench_wide_pool<2>(w, pt));
            emit(bench_wide_pool<3>(w, pt));
            emit(bench_wide_pool<4>(w, pt));
            emit(bench_diamond_pool(w, pt));
        }

#ifdef KPN_BENCH_TBB
        std::fprintf(stderr, "\n── work_us=%-4d TBB flow graph ──────────────────────────────────────\n", w);
        for (int d : {1, 2, 4, 8, 16, 32}) emit(bench_chain_tbb(d, w));
        emit(bench_wide_tbb<1>(w));
        emit(bench_wide_tbb<2>(w));
        emit(bench_wide_tbb<3>(w));
        emit(bench_wide_tbb<4>(w));
        emit(bench_diamond_tbb(w));
#endif
    }
}