(** CEK benchmark — measures throughput of the CEK evaluator on tight loops.

    Usage:
      dune exec bin/bench_cek.exe
      dune exec bin/bench_cek.exe -- 5      (5 runs each)
*)

open Sx_types
open Sx_parser

let parse_one s =
  let exprs = parse_all s in
  match exprs with
  | e :: _ -> e
  | [] -> failwith "empty parse"

let parse_many s = parse_all s

let bench_run name setup expr iters =
  let env = Sx_types.make_env () in
  (* Run setup forms in env *)
  List.iter (fun e -> ignore (Sx_ref.eval_expr e (Env env))) setup;
  let times = ref [] in
  for _ = 1 to iters do
    Gc.full_major ();
    let t0 = Unix.gettimeofday () in
    let _r = Sx_ref.eval_expr expr (Env env) in
    let t1 = Unix.gettimeofday () in
    times := (t1 -. t0) :: !times
  done;
  let sorted = List.sort compare !times in
  let median = List.nth sorted (iters / 2) in
  let min_t = List.nth sorted 0 in
  let max_t = List.nth sorted (iters - 1) in
  Printf.printf "  %-22s  min=%8.2fms  median=%8.2fms  max=%8.2fms\n%!"
    name (min_t *. 1000.0) (median *. 1000.0) (max_t *. 1000.0);
  median

let () =
  let iters =
    if Array.length Sys.argv > 1
    then int_of_string Sys.argv.(1)
    else 5
  in
  Printf.printf "CEK benchmark (%d runs each, taking median)\n%!" iters;
  Printf.printf "==========================================\n%!";

  (* fib 18 — recursive function call benchmark, smallish *)
  let fib_setup = parse_many "(define (fib n) (if (< n 2) n (+ (fib (- n 1)) (fib (- n 2)))))" in
  let fib_expr = parse_one "(fib 18)" in
  let _ = bench_run "fib(18)" fib_setup fib_expr iters in

  (* loop 5000 — tight let loop *)
  let loop_setup = parse_many "(define (loop n acc) (if (= n 0) acc (loop (- n 1) (+ acc 1))))" in
  let loop_expr = parse_one "(loop 5000 0)" in
  let _ = bench_run "loop(5000)" loop_setup loop_expr iters in

  (* map+square over 1000 elem list *)
  let map_setup = parse_many "(define (range-list n) (let loop ((i 0) (acc (list))) (if (= i n) acc (loop (+ i 1) (cons i acc))))) (define xs (range-list 1000))" in
  let map_expr = parse_one "(map (fn (x) (* x x)) xs)" in
  let _ = bench_run "map sq xs(1000)" map_setup map_expr iters in

  (* reduce + over 2000 elem list *)
  let red_setup = parse_many "(define (range-list n) (let loop ((i 0) (acc (list))) (if (= i n) acc (loop (+ i 1) (cons i acc))))) (define ys (range-list 2000))" in
  let red_expr = parse_one "(reduce + 0 ys)" in
  let _ = bench_run "reduce + ys(2000)" red_setup red_expr iters in

  (* let-heavy: many bindings + if *)
  let lh_setup = parse_many "(define (lh n) (let ((a 1) (b 2) (c 3) (d 4)) (if (= n 0) (+ a b c d) (lh (- n 1)))))" in
  let lh_expr = parse_one "(lh 2000)" in
  let _ = bench_run "let-heavy(2000)" lh_setup lh_expr iters in

  Printf.printf "\nDone.\n%!"