#!/usr/bin/env python3
"""
bench.py — benchmarks for lumbda.
Compares interpreter time against equivalent CPython.

Usage: python3 bench.py [-v]
"""
import sys, time, math
sys.setrecursionlimit(10000)
from lumbda import make_global_env, read_all, leval, PRELUDE, show, bc_compile_proc, Proc

VERBOSE = '-v' in sys.argv

def fresh():
    g = make_global_env()
    for e in read_all(PRELUDE): leval(e, g)
    return g

def run(src, env):
    result = None
    for e in read_all(src): result = leval(e, env)
    return result

def bench(name, lisp_src, python_fn, iters=3):
    """Time both lisp and python implementations; print results."""
    env = fresh()
    # Warm up
    run(lisp_src, env)
    python_fn()

    # Time Lisp
    lisp_times = []
    for _ in range(iters):
        env2 = fresh()
        t = time.perf_counter()
        result = run(lisp_src, env2)
        lisp_times.append(time.perf_counter() - t)
    lisp_best = min(lisp_times)

    # Time Python
    py_times = []
    for _ in range(iters):
        t = time.perf_counter()
        py_result = python_fn()
        py_times.append(time.perf_counter() - t)
    py_best = min(py_times)

    slowdown = lisp_best / py_best if py_best > 0 else float('inf')
    status = '✓' if result == py_result or str(result) == str(py_result) else '✗'

    print(f'{status} {name:<30}  lisp={lisp_best*1000:7.1f}ms  py={py_best*1000:7.1f}ms  '
          f'ratio={slowdown:5.1f}x')
    if VERBOSE:
        print(f'  lisp result: {show(result)!r}')
        print(f'  py   result: {py_result!r}')
    return lisp_best, py_best


print('lumbda benchmarks')
print('=' * 75)

# ── 1. Fibonacci (iterative, named let) ──────────────────────────────────────
bench(
    'fib(35) named-let',
    '''
    (define (fib n)
      (let loop ((a 0) (b 1) (i 0))
        (if (= i n) a (loop b (+ a b) (+ i 1)))))
    (fib 35)
    ''',
    lambda: (lambda a=0, b=1, i=0, n=35:
             [setattr(sys.modules[__name__], '_fib', None)] and
             (lambda: exec('''
def _fib(n):
    a, b = 0, 1
    for _ in range(n): a, b = b, a+b
    return a
''', globals()) or _fib(35))())()
)

# Simpler lambda-based benchmark
def py_fib(n=35):
    a, b = 0, 1
    for _ in range(n): a, b = b, a + b
    return a

bench(
    'fib(35) iterative',
    '''
    (define (fib n)
      (let loop ((a 0) (b 1) (i 0))
        (if (= i n) a (loop b (+ a b) (+ i 1)))))
    (fib 35)
    ''',
    lambda: py_fib(35)
)

# ── 2. fib(20) tree-recursive ─────────────────────────────────────────────────
def py_fib_rec(n):
    if n <= 1: return n
    return py_fib_rec(n-1) + py_fib_rec(n-2)

bench(
    'fib(20) tree-recursive',
    '''
    (define (fib n)
      (if (<= n 1) n (+ (fib (- n 1)) (fib (- n 2)))))
    (fib 20)
    ''',
    lambda: py_fib_rec(20)
)

# ── 3. Tak function ───────────────────────────────────────────────────────────
def py_tak(x, y, z):
    if y >= x: return z
    return py_tak(py_tak(x-1,y,z), py_tak(y-1,z,x), py_tak(z-1,x,y))

bench(
    'tak(15,10,6)',
    '''
    (define (tak x y z)
      (if (>= y x) z
          (tak (tak (- x 1) y z)
               (tak (- y 1) z x)
               (tak (- z 1) x y))))
    (tak 15 10 6)
    ''',
    lambda: py_tak(15, 10, 6)
)

# ── 4. Tail-recursive sum ─────────────────────────────────────────────────────
def py_sum(n):
    acc = 0
    while n > 0: acc += n; n -= 1
    return acc

bench(
    'sum-to(50000) tail-recursive',
    '''
    (define (sum-to n)
      (let loop ((i n) (acc 0))
        (if (= i 0) acc (loop (- i 1) (+ acc i)))))
    (sum-to 50000)
    ''',
    lambda: py_sum(50000)
)

# ── 5. List operations ────────────────────────────────────────────────────────
def py_list_ops():
    lst = list(range(5000))
    lst = list(reversed(lst))
    lst = [x for x in lst if x % 2 == 1]
    lst = [x * x for x in lst]
    return sum(lst)

bench(
    'list ops (5000 elements)',
    '''
    (define lst (iota 5000))
    (define lst (reverse lst))
    (define lst (filter odd? lst))
    (define lst (map (lambda (x) (* x x)) lst))
    (fold-left + 0 lst)
    ''',
    py_list_ops
)

# ── 6. Higher-order / closures ────────────────────────────────────────────────
def py_adder_factory():
    adders = [(lambda n: lambda x: x + n)(i) for i in range(100)]
    return sum(f(10) for f in adders)

bench(
    'closure factory (100 adders)',
    '''
    (define (make-adder n) (lambda (x) (+ x n)))
    (define adders (map make-adder (iota 100)))
    (fold-left + 0 (map (lambda (f) (f 10)) adders))
    ''',
    py_adder_factory
)

# ── 7. Hash table ─────────────────────────────────────────────────────────────
def py_hash_ops():
    h = {}
    for i in range(1000):
        h[i] = i * i
    return sum(h.get(i, 0) for i in range(1000))

bench(
    'hash-table (1000 set+ref)',
    '''
    (define h (make-hash-table))
    (do ((i 0 (+ i 1))) ((= i 1000))
      (hash-table-set! h i (* i i)))
    (do ((i 0 (+ i 1)) (s 0 (+ s (hash-table-ref/default h i 0))))
        ((= i 1000) s))
    ''',
    py_hash_ops
)

# ── 8. String operations ──────────────────────────────────────────────────────
def py_string_ops():
    parts = [str(i) for i in range(200)]
    joined = ', '.join(parts)
    return len(joined)

bench(
    'string-join (200 numbers)',
    '''
    (string-length
      (string-join (map number->string (iota 200)) ", "))
    ''',
    py_string_ops
)

# ── 9. Ackermann (small) ─────────────────────────────────────────────────────
def py_ack(m, n):
    if m == 0: return n + 1
    if n == 0: return py_ack(m-1, 1)
    return py_ack(m-1, py_ack(m, n-1))

bench(
    'ackermann(3,4)',
    '''
    (define (ack m n)
      (cond ((= m 0) (+ n 1))
            ((= n 0) (ack (- m 1) 1))
            (else    (ack (- m 1) (ack m (- n 1))))))
    (ack 3 4)
    ''',
    lambda: py_ack(3, 4)
)

# ── 10. Mergesort ─────────────────────────────────────────────────────────────
def py_msort(lst):
    if len(lst) <= 1: return lst
    mid = len(lst) // 2
    L = py_msort(lst[:mid]); R = py_msort(lst[mid:])
    result = []; i = j = 0
    while i < len(L) and j < len(R):
        if L[i] <= R[j]: result.append(L[i]); i += 1
        else: result.append(R[j]); j += 1
    return result + L[i:] + R[j:]

bench(
    'mergesort (200 elements)',
    '''
    (define (merge a b)
      (cond ((null? a) b) ((null? b) a)
            ((< (car a) (car b)) (cons (car a) (merge (cdr a) b)))
            (else (cons (car b) (merge a (cdr b))))))
    (define (split lst)
      (let loop ((l lst) (a (quote ())) (b (quote ())))
        (if (null? l) (list a b) (loop (cdr l) b (cons (car l) a)))))
    (define (msort lst)
      (if (or (null? lst) (null? (cdr lst))) lst
          (let ((h (split lst)))
            (merge (msort (car h)) (msort (cadr h))))))
    (length (msort (reverse (iota 200))))
    ''',
    lambda: len(py_msort(list(range(199, -1, -1))))
)

print('=' * 75)
print('ratio = lisp time / python time  (lower is better for lisp)')

# ═══════════════════════════════════════════════════════════════════════════
# Compiled (bytecode) benchmarks
# ═══════════════════════════════════════════════════════════════════════════

print()
print('COMPILED (bytecode) benchmarks')
print('=' * 75)

def compile_env(env):
    """Compile all user-defined Procs in env."""
    for k, v in list(env.b.items()):
        if isinstance(v, Proc):
            env.b[k] = bc_compile_proc(v, env)

def bench_compiled(name, lisp_src, python_fn, iters=3):
    """Time compiled lisp vs python. Separates definition from invocation."""
    # Parse all expressions; everything except the last is setup (defines)
    exprs = read_all(lisp_src)
    setup_exprs = exprs[:-1]
    call_expr = exprs[-1]

    lisp_times = []
    for _ in range(iters):
        env = fresh()
        for e in setup_exprs: leval(e, env)
        compile_env(env)
        t = time.perf_counter()
        result = leval(call_expr, env)
        lisp_times.append(time.perf_counter() - t)
    lisp_best = min(lisp_times)

    py_times = []
    for _ in range(iters):
        t = time.perf_counter()
        py_result = python_fn()
        py_times.append(time.perf_counter() - t)
    py_best = min(py_times)

    slowdown = lisp_best / py_best if py_best > 0 else float('inf')
    status = '✓' if result == py_result or str(result) == str(py_result) else '✗'
    print(f'{status} {name:<30}  bc={lisp_best*1000:7.1f}ms  py={py_best*1000:7.1f}ms  '
          f'ratio={slowdown:5.1f}x')
    return lisp_best, py_best

bench_compiled('fib(35) iterative', '''
    (define (fib n)
      (let loop ((a 0) (b 1) (i 0))
        (if (= i n) a (loop b (+ a b) (+ i 1)))))
    (fib 35)
''', lambda: py_fib(35))

bench_compiled('fib(20) tree-recursive', '''
    (define (fib n)
      (if (<= n 1) n (+ (fib (- n 1)) (fib (- n 2)))))
    (fib 20)
''', lambda: py_fib_rec(20))

bench_compiled('tak(15,10,6)', '''
    (define (tak x y z)
      (if (>= y x) z
          (tak (tak (- x 1) y z)
               (tak (- y 1) z x)
               (tak (- z 1) x y))))
    (tak 15 10 6)
''', lambda: py_tak(15, 10, 6))

bench_compiled('sum-to(50000) tail-recursive', '''
    (define (sum-to n)
      (let loop ((i n) (acc 0))
        (if (= i 0) acc (loop (- i 1) (+ acc i)))))
    (sum-to 50000)
''', lambda: py_sum(50000))

bench_compiled('list ops (5000 elements)', '''
    (define lst (iota 5000))
    (define lst (reverse lst))
    (define lst (filter odd? lst))
    (define lst (map (lambda (x) (* x x)) lst))
    (fold-left + 0 lst)
''', py_list_ops)

bench_compiled('hash-table (1000 set+ref)', '''
    (define h (make-hash-table))
    (do ((i 0 (+ i 1))) ((= i 1000))
      (hash-table-set! h i (* i i)))
    (do ((i 0 (+ i 1)) (s 0 (+ s (hash-table-ref/default h i 0))))
        ((= i 1000) s))
''', py_hash_ops)

bench_compiled('ackermann(3,4)', '''
    (define (ack m n)
      (cond ((= m 0) (+ n 1))
            ((= n 0) (ack (- m 1) 1))
            (else    (ack (- m 1) (ack m (- n 1))))))
    (ack 3 4)
''', lambda: py_ack(3, 4))

bench_compiled('mergesort (200 elements)', '''
    (define (merge a b)
      (cond ((null? a) b) ((null? b) a)
            ((< (car a) (car b)) (cons (car a) (merge (cdr a) b)))
            (else (cons (car b) (merge a (cdr b))))))
    (define (split lst)
      (let loop ((l lst) (a (quote ())) (b (quote ())))
        (if (null? l) (list a b) (loop (cdr l) b (cons (car l) a)))))
    (define (msort lst)
      (if (or (null? lst) (null? (cdr lst))) lst
          (let ((h (split lst)))
            (merge (msort (car h)) (msort (cadr h))))))
    (length (msort (reverse (iota 200))))
''', lambda: len(py_msort(list(range(199, -1, -1)))))

print('=' * 75)
print('ratio = compiled time / python time  (lower is better)')

# ═══════════════════════════════════════════════════════════════════════════
# Auto-compile mode (zero-effort: just define and go)
# ═══════════════════════════════════════════════════════════════════════════

print()
print('AUTO-COMPILE mode (auto-compile! #t)')
print('=' * 75)

from lumbda import _auto_compile

def bench_auto(name, lisp_src, python_fn, iters=3):
    """Time auto-compiled lisp vs python. Defines + runs in one pass."""
    exprs = read_all(lisp_src)

    lisp_times = []
    for _ in range(iters):
        env = fresh()
        _auto_compile[0] = True
        for e in exprs[:-1]: leval(e, env)
        t = time.perf_counter()
        result = leval(exprs[-1], env)
        lisp_times.append(time.perf_counter() - t)
        _auto_compile[0] = False
    lisp_best = min(lisp_times)

    py_times = []
    for _ in range(iters):
        t = time.perf_counter()
        py_result = python_fn()
        py_times.append(time.perf_counter() - t)
    py_best = min(py_times)

    slowdown = lisp_best / py_best if py_best > 0 else float('inf')
    status = '✓' if result == py_result or str(result) == str(py_result) else '✗'
    print(f'{status} {name:<30}  auto={lisp_best*1000:6.1f}ms  py={py_best*1000:6.1f}ms  '
          f'ratio={slowdown:5.1f}x')
    return lisp_best, py_best

bench_auto('fib(35) iterative', '''
    (define (fib n)
      (let loop ((a 0) (b 1) (i 0))
        (if (= i n) a (loop b (+ a b) (+ i 1)))))
    (fib 35)
''', lambda: py_fib(35))

bench_auto('fib(20) tree-recursive', '''
    (define (fib n)
      (if (<= n 1) n (+ (fib (- n 1)) (fib (- n 2)))))
    (fib 20)
''', lambda: py_fib_rec(20))

bench_auto('sum-to(50000) tail-recursive', '''
    (define (sum-to n)
      (let loop ((i n) (acc 0))
        (if (= i 0) acc (loop (- i 1) (+ acc i)))))
    (sum-to 50000)
''', lambda: py_sum(50000))

bench_auto('ackermann(3,4)', '''
    (define (ack m n)
      (cond ((= m 0) (+ n 1))
            ((= n 0) (ack (- m 1) 1))
            (else    (ack (- m 1) (ack m (- n 1))))))
    (ack 3 4)
''', lambda: py_ack(3, 4))

print('=' * 75)