open Nock_lib.Noun
open Nock_lib.Nock

(** Benchmark utilities *)

let time_ms () =
  Unix.gettimeofday () *. 1000.0

let bench_nock name subject formula iterations =
  (* Warmup *)
  for _i = 1 to 100 do
    let _ = nock subject formula in ()
  done;

  (* Actual benchmark *)
  Gc.compact ();
  let start = time_ms () in

  for _i = 1 to iterations do
    let _result = nock subject formula in ()
  done;

  let finish = time_ms () in
  let total = finish -. start in
  let per_iter = total /. (float_of_int iterations) in
  let ops_per_sec = 1000.0 /. per_iter in

  Printf.printf "%-30s %8d iterations in %10.2f ms (%10.6f ms/iter, %10.0f ops/sec)\n"
    name iterations total per_iter ops_per_sec

(** Benchmarks *)

let () =
  Printf.printf "Nock Benchmark - OCaml Implementation\n";
  Printf.printf "======================================\n\n";

  let iterations = 1_000_000 in  (* 1M iterations for fast ops *)
  let slow_iters = 100_000 in    (* 100K for slower ops *)

  (* Benchmark 0: slot lookup *)
  begin
    let subject = cell (atom 42) (atom 99) in
    let formula = cell (atom 0) (atom 2) in  (* [0 2] - get head *)
    bench_nock "Opcode 0: slot/fragment" subject formula iterations
  end;

  (* Benchmark 1: constant *)
  begin
    let subject = atom 0 in
    let formula = cell (atom 1) (atom 42) in  (* [1 42] *)
    bench_nock "Opcode 1: constant" subject formula iterations
  end;

  (* Benchmark 3: is-cell *)
  begin
    let subject = atom 0 in
    let formula = cell (atom 3) (cell (atom 1) (atom 42)) in  (* [3 [1 42]] *)
    bench_nock "Opcode 3: is-cell (atom)" subject formula iterations
  end;

  (* Benchmark 4: increment *)
  begin
    let subject = atom 0 in
    let formula = cell (atom 4) (cell (atom 1) (atom 1000)) in  (* [4 [1 1000]] *)
    bench_nock "Opcode 4: increment" subject formula iterations
  end;

  (* Benchmark 5: equality *)
  begin
    let subject = atom 0 in
    (* [5 [1 42] [1 42]] *)
    let formula = cell (atom 5) (cell (cell (atom 1) (atom 42)) (cell (atom 1) (atom 42))) in
    bench_nock "Opcode 5: equality (equal)" subject formula iterations
  end;

  (* Benchmark 6: if-then-else *)
  begin
    let subject = atom 0 in
    (* [6 [1 0] [1 11] [1 22]] *)
    let formula = cell (atom 6)
      (cell (cell (atom 1) (atom 0))
        (cell (cell (atom 1) (atom 11))
          (cell (atom 1) (atom 22)))) in
    bench_nock "Opcode 6: if-then-else" subject formula iterations
  end;

  (* Benchmark 7: composition *)
  begin
    let subject = atom 42 in
    (* [7 [1 99] [0 1]] *)
    let formula = cell (atom 7) (cell (cell (atom 1) (atom 99)) (cell (atom 0) (atom 1))) in
    bench_nock "Opcode 7: composition" subject formula iterations
  end;

  (* Benchmark 8: push *)
  begin
    let subject = atom 42 in
    (* [8 [1 99] [0 1]] *)
    let formula = cell (atom 8) (cell (cell (atom 1) (atom 99)) (cell (atom 0) (atom 1))) in
    bench_nock "Opcode 8: push" subject formula iterations
  end;

  (* Benchmark: Decrement-like operation (slower) *)
  begin
    (* [6 [5 [0 1] [1 0]] [1 0] [8 [1 0] [4 [0 3]]]] *)
    (* This is: if(subject == 0) 0 else subject+1 (simplified) *)
    let dec_fol = cell (atom 6)
      (cell (cell (atom 5) (cell (cell (atom 0) (atom 1)) (cell (atom 1) (atom 0))))
        (cell (cell (atom 1) (atom 0))
          (cell (atom 8) (cell (cell (atom 1) (atom 0)) (cell (atom 4) (cell (atom 0) (atom 3))))))) in

    let subject = atom 10 in
    bench_nock "Complex: decrement loop" subject dec_fol slow_iters
  end;

  (* Benchmark: Tree construction *)
  begin
    let subject = atom 0 in
    (* [[1 1] [1 2]] - constructs a cell *)
    let formula = cell (cell (atom 1) (atom 1)) (cell (atom 1) (atom 2)) in
    bench_nock "Cell construction" subject formula iterations
  end;

  (* Benchmark: Deep slot lookup *)
  begin
    (* Build a deep tree: [[[[1 2] 3] 4] 5] *)
    let subject = cell (cell (cell (cell (atom 1) (atom 2)) (atom 3)) (atom 4)) (atom 5) in
    let formula = cell (atom 0) (atom 16) in  (* slot 16 = deepest left (1) *)
    bench_nock "Deep slot lookup (depth 4)" subject formula iterations
  end;

  Printf.printf "\n"