#!/usr/bin/env python3
import json, random, os, sys
import sys; sys.setrecursionlimit(1000000)

OUT = os.path.dirname(os.path.abspath(__file__))

def var(n): return {"tag":"Var","name":n}
def lam(n,b): return {"tag":"Lam","param":n,"body":b}
def app(f,a): return {"tag":"App","fun":f,"arg":a}
def subst(n,s,t): return {"tag":"Subst","var":n,"repl":s,"body":t}

def gen(rng, scope, fv_pool, depth, binder_density, shadow_freq, fv_freq):
    if depth <= 0:
        if scope and rng.random() > fv_freq:
            return var(rng.choice(scope))
        return var(rng.choice(fv_pool) if rng.random() < 0.8 else "z")
    r = rng.random()
    if r < binder_density:
        if scope and rng.random() < shadow_freq:
            n = rng.choice(scope)
        else:
            n = "v" + str(rng.randrange(8))
        return lam(n, gen(rng, scope+[n], fv_pool, depth-1, binder_density, shadow_freq, fv_freq))
    if r < binder_density + 0.25:
        if rng.random() < 0.15:
            n = rng.choice(scope) if scope else "u"
            s = gen(rng, scope, fv_pool, depth-1, binder_density, shadow_freq, fv_freq)
            t = gen(rng, scope, fv_pool, depth-1, binder_density, shadow_freq, fv_freq)
            return subst(n, s, t)
        return app(gen(rng, scope, fv_pool, depth-1, binder_density, shadow_freq, fv_freq),
                   gen(rng, scope, fv_pool, depth-1, binder_density, shadow_freq, fv_freq))
    if scope and rng.random() > fv_freq:
        return var(rng.choice(scope))
    return var(rng.choice(fv_pool))

def rename_chain(rng, t, depth):
    names = ["a","b","c","d","e","f","g","h"]
    out = t
    for i in range(depth):
        out = lam(rng.choice(names), out)
    return out

corpus = []

def add(name, t1, t2, expected):
    corpus.append({"name":name, "t1":t1, "t2":t2, "expected":expected})

add("identical-var", var("x"), var("x"), True)
add("identical-lam", lam("x", var("x")), lam("x", var("x")), True)
add("renamed-binder", lam("x", var("x")), lam("y", var("y")), True)
add("nested-renamed", lam("x", lam("y", app(var("x"), var("y")))),
    lam("a", lam("b", app(var("a"), var("b")))), True)
add("shadowing-same", lam("x", lam("x", var("x"))), lam("y", lam("y", var("y"))), True)
add("shadowing-diff-binding", lam("x", lam("x", var("x"))), lam("x", lam("y", var("x"))), False)
add("free-match", app(var("f"), lam("x", var("x"))), app(var("f"), lam("y", var("y"))), True)
add("free-mismatch", app(var("f"), var("x")), app(var("g"), var("x")), False)
add("binder-collides-free", lam("x", app(var("x"), var("y"))), lam("y", app(var("y"), var("y"))), False)
add("binder-collides-free2", lam("y", app(var("y"), var("x"))), lam("z", app(var("z"), var("x"))), True)
add("deep-nest", lam("a", lam("b", lam("c", lam("d", app(var("a"), var("d")))))),
    lam("w", lam("x", lam("y", lam("z", app(var("w"), var("z")))))), True)
add("subst-basic", subst("x", var("y"), lam("z", var("x"))), lam("z", var("y")), True)
add("subst-capture-delayed", subst("x", var("z"), lam("z", var("x"))), lam("z", var("z")), False)
add("subst-capture-correct", subst("x", var("z"), lam("z", var("x"))), lam("w", var("z")), True)
add("subst-shadowed", subst("x", var("q"), lam("x", var("x"))), lam("x", var("x")), True)
add("subst-repeated",
    subst("x", var("a"), subst("y", var("b"), app(var("x"), var("y")))),
    app(var("a"), var("b")), True)
add("subst-under-binder-shift",
    subst("x", app(var("w"), var("v")), lam("w", lam("v", var("x")))),
    lam("a", lam("b", app(var("w"), var("v")))), True)
add("subst-nested-subst",
    subst("x", subst("y", var("k"), var("y")), lam("k", var("x"))),
    lam("q", var("k")), True)
add("subst-not-equiv", subst("x", var("y"), var("x")), var("z"), False)
add("eta-ish-not-equiv", lam("x", app(var("f"), var("x"))), var("f"), False)
add("app-order", app(var("a"), var("b")), app(var("b"), var("a")), False)
add("normalises-equiv",
    app(lam("x", var("x")), lam("y", var("y"))),
    lam("z", var("z")), True)
add("normalises-equiv2",
    app(lam("x", lam("y", var("x"))), var("f")),
    lam("y", var("f")), True)

rng = random.Random(42)
fv_pool = ["f","g","h","p","q"]
for i in range(40):
    t = gen(rng, [], fv_pool, rng.randrange(3, 7), 0.45, 0.4, 0.3)
    add(f"rand-self-{i}", t, t, True)
for i in range(30):
    t = gen(rng, [], fv_pool, rng.randrange(3, 7), 0.45, 0.4, 0.3)
    c=[0]
    def alpha_rename(t, env):
        if t["tag"]=="Var":
            for old,new in reversed(env):
                if old==t["name"]: return var(new)
            return t
        if t["tag"]=="App": return app(alpha_rename(t["fun"],env), alpha_rename(t["arg"],env))
        if t["tag"]=="Subst":
            c[0]+=1; n=f"r{c[0]}"
            return subst(n, alpha_rename(t["repl"],env), alpha_rename(t["body"], env+[(t["var"],n)]))
        c[0]+=1; n = f"r{c[0]}"
        return lam(n, alpha_rename(t["body"], env+[(t["param"], n)]))
    t2 = alpha_rename(t, [])
    add(f"rand-renamed-{i}", t, t2, True)
for i in range(20):
    t1 = gen(rng, [], fv_pool, 5, 0.45, 0.4, 0.3)
    t2 = gen(rng, [], fv_pool, 5, 0.45, 0.4, 0.3)
    add(f"rand-any-{i}", t1, t2, None)

with open(os.path.join(OUT, "corpus.json"), "w") as f:
    json.dump(corpus, f, indent=1)

rng = random.Random(7)
def bench_chain(n):
    t = var("x0")
    for i in range(n): t = lam(f"b{i%5}", t)
    return t
def bench_apps(n):
    t = var("f")
    for i in range(n): t = app(t, var("x"))
    return t
def bench_shadow(n):
    t = var("x")
    for i in range(n): t = lam("x", t)
    return t
def bench_fv(n):
    t = var("s")
    for i in range(n): t = app(t, var(f"fv{i%50}"))
    return t
def bench_shift(n):
    t = var("x")
    for i in range(n):
        t = subst("x", var(f"w{i%3}"), lam(f"w{i%3}", t))
    return t
def bench_bigsubst(n):
    inner = var("x")
    for i in range(n): inner = app(inner, var(f"g{i%20}"))
    t = var("y")
    for i in range(20): t = lam(f"p{i}", t)
    return subst("y", inner, t)
def bench_diff_at_leaves(n):
    t = var("q")
    for i in range(n): t = lam(f"a{i}", app(t, var(f"x{i}")))
    t2 = var("DIFFERENT")
    for i in range(n): t2 = lam(f"a{i}", app(t2, var(f"x{i}")))
    return t, t2

def deep_rename(t, c=[0]):
    if t["tag"]=="Var": return t
    if t["tag"]=="App": return app(deep_rename(t["fun"],c), deep_rename(t["arg"],c))
    if t["tag"]=="Subst": return subst(t["var"], deep_rename(t["repl"],c), deep_rename(t["body"],c))
    c[0]+=1
    return lam(f"N{c[0]}", deep_rename(t["body"],c))

benches = {}
sizes = {"small":60, "medium":250, "large":1000}
for label,n in sizes.items():
    benches[f"chain-{label}"] = bench_chain(n)
    benches[f"apps-{label}"] = bench_apps(n)
    benches[f"shadow-{label}"] = bench_shadow(n)
    benches[f"fv-{label}"] = bench_fv(n)
    benches[f"shift-{label}"] = bench_shift(n // 5 if label!="large" else 120)
    benches[f"bigsubst-{label}"] = bench_bigsubst(n)
    a,b = bench_diff_at_leaves(n)
    benches[f"leaves-{label}-a"] = a
    benches[f"leaves-{label}-b"] = b
    benches[f"allbinders-{label}-a"] = bench_chain(n)
    benches[f"allbinders-{label}-b"] = deep_rename(bench_chain(n))

with open(os.path.join(OUT, "bench.json"), "w") as f:
    json.dump(benches, f)

print(f"corpus: {len(corpus)} pairs, benches: {len(benches)} terms")