import std/[os, random, sets]
import ego/named, ego/debruijn, ego/locally_nameless, ego/explicit_subst,
       ego/syntax, ego/jsonio

type Rep = enum rNamed, rDebruijn, rLn, rSubst

let reps = [rNamed, rDebruijn, rLn, rSubst]

proc check(cond: bool, msg: string) =
  if not cond:
    raise newException(AssertionDefect, msg)

proc expectRaises[T: CatchableError](msg: string, f: proc()) =
  try:
    f()
    raise newException(AssertionDefect, msg & ": expected exception")
  except T:
    discard

proc findCorpus(): string =
  for p in ["../shared/corpus.json", "shared/corpus.json",
            "/home/bulg/projects/ego/shared/corpus.json"]:
    if fileExists(p): return p
  raise newException(IOError, "corpus.json not found")

proc checkWith(rep: Rep, t1, t2: Named): bool =
  case rep
  of rNamed:
    named.alphaEquivalent(t1, t2)
  of rDebruijn:
    debruijn.alphaEquivalent(debruijn.convertFromNamed(t1),
                             debruijn.convertFromNamed(t2))
  of rLn:
    locally_nameless.alphaEquivalent(locally_nameless.convertFromNamed(t1),
                                     locally_nameless.convertFromNamed(t2))
  of rSubst:
    explicit_subst.alphaEquivalent(explicit_subst.convertFromNamed(t1),
                                   explicit_subst.convertFromNamed(t2))

proc pick(r: var Rand, xs: openArray[string]): string =
  xs[r.rand(xs.len - 1)]

proc genLeaf(r: var Rand, scope: seq[string]): Named =
  let fvs = ["p", "q", "g", "h"]
  if scope.len > 0 and r.rand(99) < 75:
    nVar(r.pick(scope))
  else:
    nVar(r.pick(fvs))

proc gen(r: var Rand, scope: seq[string], depth: int): Named =
  let names = ["a", "b", "c", "d", "e", "f"]
  if depth <= 0:
    return genLeaf(r, scope)
  let n = r.rand(99)
  if n < 45:
    let binder =
      if scope.len > 0 and r.rand(99) < 35: r.pick(scope) else: r.pick(names)
    return nLam(binder, gen(r, scope & @[binder], depth - 1))
  if n < 75:
    return nApp(genLeaf(r, scope), genLeaf(r, scope))
  if n < 85:
    let x = if scope.len == 0: "u" else: r.pick(scope)
    return nSubst(x, gen(r, scope, depth - 1), gen(r, scope, depth - 1))
  genLeaf(r, scope)

proc randomTerm(r: var Rand): Named =
  gen(r, @[], 2 + r.rand(3))

proc testCorpus() =
  for e in loadCorpus(findCorpus()):
    var first = false
    var haveFirst = false
    for rep in reps:
      let got = checkWith(rep, e.t1, e.t2)
      if e.expected != 2:
        check(got == (e.expected == 1), e.name & " expected mismatch")
      if haveFirst:
        check(got == first, e.name & " representation disagreement")
      else:
        first = got
        haveFirst = true

proc testProperties() =
  var r = initRand(42)
  for _ in 0 ..< 200:
    let t = randomTerm(r)
    for rep in reps:
      check(checkWith(rep, t, t), "reflexivity")
  for _ in 0 ..< 200:
    let t1 = randomTerm(r)
    let t2 = randomTerm(r)
    for rep in reps:
      check(checkWith(rep, t1, t2) == checkWith(rep, t2, t1), "symmetry")
  for _ in 0 ..< 100:
    let t = randomTerm(r)
    let n = named.normalise(t)
    for rep in reps:
      check(checkWith(rep, t, t) and checkWith(rep, n, n), "transitivity sample")
  for _ in 0 ..< 200:
    let t1 = randomTerm(r)
    let t2 = randomTerm(r)
    if named.alphaEquivalent(t1, t2):
      check(freeVars(named.normalise(t1)) == freeVars(named.normalise(t2)),
            "free variables preserved")
  for _ in 0 ..< 200:
    let t = randomTerm(r)
    check(named.alphaEquivalent(debruijn.convertToNamed(debruijn.convertFromNamed(t)), t),
          "debruijn roundtrip")
    check(named.alphaEquivalent(locally_nameless.convertToNamed(
          locally_nameless.convertFromNamed(t)), t), "ln roundtrip")
    check(named.alphaEquivalent(explicit_subst.convertToNamed(
          explicit_subst.convertFromNamed(t)), t), "subst roundtrip")
  for _ in 0 ..< 100:
    let t = randomTerm(r)
    let n1 = named.normalise(t)
    let n2 = debruijn.convertToNamed(debruijn.normalise(debruijn.convertFromNamed(t)))
    let n3 = locally_nameless.convertToNamed(
      locally_nameless.normalise(locally_nameless.convertFromNamed(t)))
    let n4 = explicit_subst.convertToNamed(
      explicit_subst.normalise(explicit_subst.convertFromNamed(t)))
    check(named.alphaEquivalent(n1, n2), "named vs debruijn normalise")
    check(named.alphaEquivalent(n1, n3), "named vs ln normalise")
    check(named.alphaEquivalent(n1, n4), "named vs subst normalise")

proc testParseErrors() =
  for s in ["", "\\", "\\x", "\\x.", "(\\x. x", "\\x. x)", "[x := ] x",
            "[ := y] x", "((x", "[x := y]", ".", "\\. x"]:
    let r = parseResult(s)
    check(not r.ok, "parser accepted " & s)

proc testMalformed() =
  check(not locally_nameless.lc(bVar(0)), "dangling bvar")
  check(locally_nameless.lc(lLam(bVar(0))), "closed lambda")
  check(not locally_nameless.lc(lLam(lLam(bVar(2)))), "dangling deep")
  expectRaises[NotLocallyClosed]("ln convert", proc () =
    discard locally_nameless.convertToNamed(bVar(3))
  )
  expectRaises[NotLocallyClosed]("ln alpha", proc () =
    discard locally_nameless.alphaEquivalent(bVar(1), bVar(1))
  )
  check(debruijn.wellScoped(dLam(dVar(0))), "well scoped db")
  check(not debruijn.wellScoped(dVar(0)), "ill scoped db")
  expectRaises[UnboundIndex]("db convert", proc () =
    discard debruijn.convertToNamed(dLam(dVar(2)))
  )

proc parsed(s: string): Named =
  let r = parseResult(s)
  if not r.ok:
    raise newException(ValueError, r.err)
  r.term

proc testCapture() =
  let cases = [
    ("[x := z] \\z. x", "\\w. z", true),
    ("[x := z] \\z. x", "\\z. z", false),
    ("[x := q] \\x. x", "\\x. x", true),
    ("[x := y] \\y. x y", "\\w. y w", true),
    ("[x := f y] \\y. \\z. x z", "\\a. \\b. f y b", true),
    ("[x := a] [y := b] x y", "a b", true)
  ]
  for (a, b, exp) in cases:
    let t1 = parsed(a)
    let t2 = parsed(b)
    for rep in reps:
      check(checkWith(rep, t1, t2) == exp, a & " ~ " & b)
  let t = explicit_subst.ofNamed(
    nSubst("x", nVar("z"), nLam("z", nVar("x"))))
  let n = explicit_subst.normalise(t)
  check(n.kind == ekLam, "expected lambda after explicit normalise")
  check(n.param != "z", "binder renamed")
  check(n.body.kind == ekVar and n.body.name == "z", "body is free z")
  check(explicit_subst.alphaEquivalentSusp(t, n), "susp agrees")

proc run(name: string, f: proc()) =
  f()
  echo "ok ", name

run("corpus", testCorpus)
run("properties", testProperties)
run("parse errors", testParseErrors)
run("malformed", testMalformed)
run("capture", testCapture)