//! Benchmarking `takeaway` manually.
//!
//! # Benchmarking Strategy
//!
//! The objective is to measure the throughput of a particular task queue
//! implementation, in units of tasks per second.  Each implementation will be
//! sampled multiple times, each time for a different duration.  These samples
//! (mapping a duration to a number of executed tasks) are fitted to a linear
//! function, whose parameters are the estimated throughput and the overhead.

use std::{
    f32, fmt,
    io::Write,
    num::NonZeroUsize,
    ops::{ControlFlow, Div, Mul},
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
        mpsc,
    },
    task::{Poll, Waker},
    time::{Duration, Instant},
};

use takeaway::util::CachePadded;

//----------- Implementations --------------------------------------------------

/// An implementation of a task queue.
pub trait Implementation<T: Task>: Send + Sync + Sized + 'static {
    /// Thread-local state for this implementation.
    type Local: Send + 'static;

    /// The name of this implementation.
    const NAME: &'static str;

    /// Initialize the implementation.
    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>);

    /// Drive a worker thread.
    ///
    /// Return the number of locally executed tasks.
    fn drive(
        &self,
        local: Self::Local,
        context: &T::Context,
        limit: WorkerLimit,
    ) -> usize;

    /// Shut down a running instance.
    fn stop(&self);
}

/// An execution limit for a single worker thread.
#[derive(Copy, Clone, Debug)]
pub enum WorkerLimit {
    /// No limit.
    None,

    /// Execute this many tasks at most.
    TaskCount(usize),
}

pub struct Takeaway<T: takeaway::Task>(takeaway::Queue<T>);

impl<T: Task + takeaway::Task> Implementation<T> for Takeaway<T> {
    // The ID of each worker.
    type Local = usize;

    const NAME: &'static str = "takeaway";

    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>) {
        let queue = takeaway::Config::new(num_workers)
            .with_batch_size(batch_size)
            .build();
        let locals = (0..num_workers.get()).collect();
        (Self(queue), locals)
    }

    fn drive(
        &self,
        local: Self::Local,
        context: &T::Context,
        limit: WorkerLimit,
    ) -> usize {
        let mut executed = 0usize;
        let limit = match limit {
            WorkerLimit::None => usize::MAX,
            WorkerLimit::TaskCount(n) => n,
        };
        let mut worker = takeaway::Worker::new(&self.0, local);
        while executed < limit {
            match worker.poll(Waker::noop()) {
                Poll::Ready(Some(task)) => {
                    task.execute(context, |t| worker.enqueue_one(t));
                    executed += 1;
                }
                Poll::Ready(None) => break,
                Poll::Pending => {
                    // The worker ran out of tasks, add some more.
                    worker.enqueue((0..16).map(|_| T::random()));
                }
            }
        }
        executed
    }

    fn stop(&self) {
        self.0.shutdown();
    }
}

pub struct CrossbeamDeque<T> {
    /// The configured batch size.
    batch_size: NonZeroUsize,

    /// The stealer of every worker.
    stealers: Vec<crossbeam_deque::Stealer<T>>,

    /// Whether a shutdown has occurred.
    shutdown: Box<CachePadded<AtomicBool>>,
}

impl<T: Task> Implementation<T> for CrossbeamDeque<T> {
    type Local = crossbeam_deque::Worker<T>;

    const NAME: &'static str = "crossbeam-deque";

    fn init(
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
    ) -> (Self, Vec<Self::Local>) {
        let workers = (0..num_workers.get())
            .map(|_| crossbeam_deque::Worker::new_fifo())
            .collect::<Vec<_>>();
        let stealers = workers.iter().map(|w| w.stealer()).collect();
        let shutdown = Box::new(CachePadded::new(AtomicBool::new(false)));
        (
            Self {
                batch_size,
                stealers,
                shutdown,
            },
            workers,
        )
    }

    fn drive(
        &self,
        local: Self::Local,
        context: &T::Context,
        limit: WorkerLimit,
    ) -> usize {
        let mut executed = 0usize;
        let limit = match limit {
            WorkerLimit::None => usize::MAX,
            WorkerLimit::TaskCount(n) => n,
        };
        'main: while executed < limit {
            let task = 'fetch: {
                // Stop if all tasks are complete.
                if self.shutdown.load(Ordering::Relaxed) {
                    break 'main;
                }

                // Pop a task from the local queue, if not empty.
                if let Some(task) = local.pop() {
                    break 'fetch task;
                }

                let split = rand::random_range(0..self.stealers.len());
                let mut stealers = self.stealers[split..]
                    .iter()
                    .chain(&self.stealers[..split]);

                match stealers.try_for_each(|stealer| {
                    match stealer.steal_batch_with_limit_and_pop(
                        &local,
                        self.batch_size.get(),
                    ) {
                        crossbeam_deque::Steal::Success(task) => {
                            ControlFlow::Break(task)
                        }
                        _ => ControlFlow::Continue(()),
                    }
                }) {
                    ControlFlow::Break(task) => break 'fetch task,
                    ControlFlow::Continue(()) => {}
                }

                // We ran out of tasks; add some more.
                for _ in 0..16 {
                    local.push(T::random());
                }
                continue 'main;
            };

            task.execute(context, |t| local.push(t));
            executed += 1;
        }
        executed
    }

    fn stop(&self) {
        self.shutdown.store(true, Ordering::Relaxed);
    }
}

//----------- Tasks ------------------------------------------------------------

/// A benchmarking task.
pub trait Task: Send + Sized + 'static {
    /// Context for executing tasks.
    type Context;

    /// The name of this task type.
    const NAME: &'static str;

    /// Generate a random task.
    fn random() -> Self;

    /// Execute the task.
    ///
    /// The returned tasks should be enqueued for future execution.
    fn execute(self, context: &Self::Context, enqueue: impl Fn(Self));
}

/// A trivial no-op task.
///
/// Each iteration performs a few arithmetic operations.
pub struct Noop {
    data: usize,
}

impl Task for Noop {
    type Context = ();
    const NAME: &'static str = "noop";

    fn random() -> Self {
        Self {
            data: rand::random::<u32>() as usize,
        }
    }

    fn execute(self, _context: &Self::Context, enqueue: impl Fn(Self)) {
        for i in 0..rand::random_range(0..4) {
            enqueue(Self {
                data: self.data.wrapping_mul(4).wrapping_add(i),
            })
        }
    }
}

impl takeaway::Task for Noop {
    type Priority = ();
    fn priority(&self) -> Self::Priority {}
}

/// A lightweight task.
///
/// Each iteration performs string formatting, on an ever-growing string.
pub struct Light {
    data: String,
}

impl Task for Light {
    type Context = ();
    const NAME: &'static str = "light";

    fn random() -> Self {
        Self {
            data: format!("{}", rand::random::<u32>() as usize),
        }
    }

    fn execute(self, _context: &Self::Context, enqueue: impl Fn(Self)) {
        for i in 0..rand::random_range(0..4) {
            let data = format!("{}{i}", self.data);
            let data = data.replace(['8', '9'], "");
            let data = usize::from_str_radix(&data, 8).unwrap();
            let data = format!("{data}");
            enqueue(Self { data })
        }
    }
}

impl takeaway::Task for Light {
    type Priority = ();
    fn priority(&self) -> Self::Priority {}
}

//----------- Running ----------------------------------------------------------

/// A benchmarker for a task queue implementation.
pub struct Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    /// Control of every benchmarking thread.
    ctrls: Vec<ChannelControl<I, I::Local>>,
}

struct ChannelControl<I, L> {
    start: mpsc::SyncSender<(Arc<I>, L, WorkerLimit)>,
    finish: mpsc::Receiver<(Instant, usize)>,
}

/// A limit for benchmarking.
#[derive(Copy, Clone, Debug)]
pub enum Limit {
    /// Run for a limited amount of time, in total.
    Time(Duration),

    /// Run a limited number of tasks, in total.
    TaskCount(usize),
}

impl Limit {
    /// Split off a fraction of the limit.
    pub fn split_off(&mut self, frac: f32) -> Self {
        match self {
            Self::Time(dur) => {
                let away = dur.mul_f32(frac);
                *dur -= away;
                Self::Time(away)
            }

            Self::TaskCount(num) => {
                let away = (*num as f32 * frac) as usize;
                *num -= away;
                Self::TaskCount(away)
            }
        }
    }
}

impl Mul<f32> for Limit {
    type Output = Self;

    fn mul(self, rhs: f32) -> Self::Output {
        match self {
            Self::Time(dur) => Self::Time(dur.mul_f32(rhs)),
            Self::TaskCount(num) => {
                Self::TaskCount(((num as f32) * rhs) as usize)
            }
        }
    }
}

impl Div<f32> for Limit {
    type Output = Self;

    fn div(self, rhs: f32) -> Self::Output {
        match self {
            Self::Time(dur) => Self::Time(dur.div_f32(rhs)),
            Self::TaskCount(num) => {
                Self::TaskCount(((num as f32) / rhs) as usize)
            }
        }
    }
}

impl fmt::Display for Limit {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Self::Time(dur) => write!(f, "{:.1}s", dur.as_secs_f32()),
            Self::TaskCount(num) => write!(f, "{} M tasks", num as f32 / 1e6),
        }
    }
}

impl<I, T> Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    /// Initialize a [`Runner`].
    pub fn new() -> Self {
        let num = std::thread::available_parallelism().unwrap();
        let ctrls = (0..num.get())
            .map(|_| {
                let (start_tx, start_rx) = mpsc::sync_channel(0);
                let (finish_tx, finish_rx) = mpsc::sync_channel(0);
                std::thread::spawn(|| Self::drive(start_rx, finish_tx));
                ChannelControl {
                    start: start_tx,
                    finish: finish_rx,
                }
            })
            .collect();
        Self { ctrls }
    }

    /// Drive a thread for this runner.
    fn drive(
        start: mpsc::Receiver<(Arc<I>, I::Local, WorkerLimit)>,
        finish: mpsc::SyncSender<(Instant, usize)>,
    ) {
        while let Ok((global, local, limit)) = start.recv() {
            let num = global.drive(local, &(), limit);
            let _ = finish.send((Instant::now(), num));
        }
    }

    /// Benchmark this implementation.
    pub fn benchmark(
        &mut self,
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
        mut limit: Limit,
        warmup: bool,
    ) -> Report {
        eprintln!(
            "Benchmarking {}/{} {num_workers}t {batch_size}b...",
            I::NAME,
            T::NAME
        );

        let estimate = if warmup {
            let mut limit = limit.split_off(30.0 / 85.0);
            eprint!("  Warming up for {limit}...");
            std::io::stdout().flush().unwrap();

            let mut time = Duration::ZERO;
            for _ in 0..3 {
                time += self
                    .iterate(num_workers, batch_size, limit.split_off(0.333))
                    .latency;
            }
            Some(time.div_f32(30.0))
        } else {
            None
        };

        let mut total_steps = 0;
        let reports: Vec<_> = (1..=10)
            .map(|steps| {
                let limit =
                    limit.split_off(steps as f32 / (55 - total_steps) as f32);
                total_steps += steps;

                eprint!("\x1b[2K\r  Iterating for {limit}");
                if let Some(estimate) = estimate {
                    eprint!(" ({:.1}s)", (estimate * steps).as_secs_f32());
                }
                eprint!("...");
                std::io::stdout().flush().unwrap();

                self.iterate(num_workers, batch_size, limit)
            })
            .collect();

        let n = reports.len() as f64;
        let x = reports.iter().map(|r| r.latency.as_secs_f64());
        let y = reports.iter().map(|r| r.executed as f64);

        let xs = x.clone().sum::<f64>();
        let ys = y.clone().sum::<f64>();
        let xs2 = xs * xs;
        let x2s = x.clone().map(|x| x * x).sum::<f64>();
        let xys = x.clone().zip(y.clone()).map(|(x, y)| x * y).sum::<f64>();

        let throughput = (n * xys - xs * ys) / (n * x2s - xs2);

        eprintln!(
            "\x1b[2K\r  Throughput: {:.3} M tasks/s/thread",
            throughput / 1e6 / num_workers.get() as f64
        );

        Report {
            params: ReportParams {
                implementation: I::NAME,
                task: T::NAME,
                num_workers,
                batch_size,
            },
            throughput,
        }
    }

    /// Perform a single iteration of a benchmark.
    fn iterate(
        &mut self,
        num_workers: NonZeroUsize,
        batch_size: NonZeroUsize,
        limit: Limit,
    ) -> IterReport {
        let (global, locals) = I::init(num_workers, batch_size);
        let global = Arc::new(global);
        let ctrls = &self.ctrls[..num_workers.get()];

        let worker_limit = match limit {
            Limit::Time(_) => WorkerLimit::None,
            Limit::TaskCount(num) => WorkerLimit::TaskCount(num),
        };

        // Begin the benchmarking.
        let start = Instant::now();
        ctrls.iter().zip(locals).for_each(|(ctrl, local)| {
            ctrl.start
                .send((global.clone(), local, worker_limit))
                .unwrap();
        });

        if let Limit::Time(dur) = limit {
            // Wait for the requested duration.
            std::thread::sleep(dur);

            // End the benchmarking.
            global.stop();
        }

        let mut executed = 0usize;
        let ends: Vec<_> = ctrls
            .iter()
            .map(|ctrl| {
                let (end, local_executed) = ctrl.finish.recv().unwrap();
                executed += local_executed;
                end
            })
            .collect();
        let end = Instant::now();

        // Prepare the report.
        IterReport {
            params: ReportParams {
                implementation: I::NAME,
                task: T::NAME,
                num_workers,
                batch_size,
            },
            latency: end - start,
            _total_work_time: ends.iter().map(|&e| e - start).sum(),
            _termination_latency_cost: ends.iter().map(|&e| end - e).sum(),
            executed,
        }
    }
}

impl<I, T> Default for Runner<I, T>
where
    I: Implementation<T>,
    T: Task<Context = ()>,
{
    fn default() -> Self {
        Self::new()
    }
}

/// A benchmark report.
#[derive(Clone, Debug)]
pub struct Report {
    /// The benchmark parameters.
    pub params: ReportParams,

    /// The estimated throughput, in units of tasks per second per thread.
    pub throughput: f64,
}

impl Report {
    pub fn write(&self, w: &mut impl fmt::Write) -> fmt::Result {
        writeln!(
            w,
            "{}\t{}\t{}\t{}\t{}",
            self.params.implementation,
            self.params.task,
            self.params.num_workers,
            self.params.batch_size,
            self.throughput
        )
    }
}

/// Parameters for a report.
#[derive(Clone, Debug)]
pub struct ReportParams {
    /// The name of the task queue implementation.
    pub implementation: &'static str,

    /// The name of the task type.
    pub task: &'static str,

    /// The number of workers.
    pub num_workers: NonZeroUsize,

    /// The batch size.
    pub batch_size: NonZeroUsize,
}

impl fmt::Display for ReportParams {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}/{} {}t {}b",
            self.implementation, self.task, self.num_workers, self.batch_size
        )
    }
}

/// A report for a single iteration.
#[derive(Clone, Debug)]
struct IterReport {
    /// The iteration parameters.
    params: ReportParams,

    /// The overall latency.
    latency: Duration,

    /// The total duration across the worker threads.
    ///
    /// This is the sum of the per-thread work time, including any overhead.
    _total_work_time: Duration,

    /// The termination inefficiency.
    ///
    /// After a shutdown is initiated, all workers continue operating for a
    /// short duration of time.  In the best case, that work is distributed
    /// across all workers evenly, so they all end simultaneously.  This field
    /// measures the difference in global latency between the optimal case and
    /// reality; i.e. how much latency could be optimized in termination.
    _termination_latency_cost: Duration,

    /// The number of tasks executed.
    executed: usize,
}

impl fmt::Display for IterReport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let throughput = self.executed as f64
            / self.latency.as_secs_f64()
            / self.params.num_workers.get() as f64
            / 1e6;

        write!(f, "{throughput:.3} M tasks/s latency")
    }
}

//----------- main() -----------------------------------------------------------

pub fn evaluate<T>(report: &mut impl fmt::Write, limit: Limit, warmup: bool)
where
    T: Task<Context = ()> + takeaway::Task,
{
    let mut takeaway = Runner::<Takeaway<_>, T>::new();
    let mut crossbeam_deque = Runner::<CrossbeamDeque<_>, T>::new();

    let parallelism = std::thread::available_parallelism().unwrap();
    for num_workers in 1..=parallelism.get() {
        let num_workers = NonZeroUsize::new(num_workers).unwrap();

        for batch_size in [4, 16, 64, 256, 4096] {
            let batch_size = NonZeroUsize::new(batch_size).unwrap();

            takeaway
                .benchmark(num_workers, batch_size, limit, warmup)
                .write(report)
                .unwrap();
            crossbeam_deque
                .benchmark(num_workers, batch_size, limit, warmup)
                .write(report)
                .unwrap();
        }
    }
}

fn main() {
    let limit = Limit::Time(Duration::from_secs_f32(8.5));
    let mut report = String::new();
    // Runner::<Takeaway<_>, Light>::new()
    //     .benchmark(1.try_into().unwrap(), 256.try_into().unwrap(), limit, true)
    //     .write(&mut report)
    //     .unwrap();
    // Runner::<CrossbeamDeque<_>, Light>::new()
    //     .benchmark(1.try_into().unwrap(), 256.try_into().unwrap(), limit, true)
    //     .write(&mut report)
    //     .unwrap();
    evaluate::<Noop>(&mut report, limit, true);
    evaluate::<Light>(&mut report, limit, true);
    print!("{report}");
}
