import std/[algorithm, os, json, monotimes, strutils]
import ego/named, ego/debruijn, ego/locally_nameless, ego/explicit_subst,
       ego/syntax, ego/jsonio

type Rep = enum rNamed, rDebruijn, rLn, rSubst

const repNames = ["named", "debruijn", "ln", "subst"]

proc repsOf(s: string): seq[Rep] =
  case s
  of "named": @[rNamed]
  of "debruijn": @[rDebruijn]
  of "ln": @[rLn]
  of "subst": @[rSubst]
  of "all": @[rNamed, rDebruijn, rLn, rSubst]
  else: raise newException(ValueError, "unknown representation: " & s)

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 cmdCheck(args: seq[string], repSel: string): int =
  let r1 = parseResult(args[0])
  let r2 = parseResult(args[1])
  if not r1.ok:
    echo "PARSE-ERROR: ", r1.err
    return 2
  if not r2.ok:
    echo "PARSE-ERROR: ", r2.err
    return 2
  for rep in repsOf(repSel):
    let eq = checkWith(rep, r1.term, r2.term)
    echo repNames[rep.ord], ": ", (if eq: "EQUIV" else: "NOT-EQUIV")
  0

proc cmdCorpus(path: string): int =
  let entries = loadCorpus(path)
  var failures = 0
  let reps = [rNamed, rDebruijn, rLn, rSubst]
  for e in entries:
    var results: seq[(Rep, bool)]
    for rep in reps:
      results.add((rep, checkWith(rep, e.t1, e.t2)))
    for (rep, got) in results:
      if e.expected != 2 and got != (e.expected == 1):
        inc failures
        echo "FAIL ", e.name, " [", repNames[rep.ord], "]: expected ",
             (if e.expected == 1: "EQUIV" else: "NOT-EQUIV"), " got ",
             (if got: "EQUIV" else: "NOT-EQUIV")
    let first = results[0][1]
    var agree = true
    for (_, g) in results:
      if g != first: agree = false
    if not agree:
      inc failures
      var parts: seq[string]
      for (rep, g) in results:
        parts.add(repNames[rep.ord] & "=" & $g)
      echo "DISAGREE ", e.name, ": ", parts.join(", ")
  if failures > 0:
    echo failures, " failure(s)"
    1
  else:
    echo "corpus: ", entries.len, " entries, 0 failures"
    0

proc nowUs(): int64 =
  let m = getMonoTime()
  m.ticks div 1000

proc median(xs: var seq[int64]): int64 =
  xs.sort()
  xs[xs.len div 2]

proc benchRep(rep: Rep, t1, t2: Named, iters: int): JsonNode =
  var alphaTs: seq[int64]
  var normTs: seq[int64]
  var trav = 0
  var depth = 0
  var memBefore = getOccupiedMem()
  for _ in 0 ..< iters:
    GC_fullCollect()
    memBefore = getOccupiedMem()
    let t0 = nowUs()
    case rep
    of rNamed:
      let (r, c) = named.alphaEquivalentCounted(t1, t2)
      alphaTs.add(nowUs() - t0)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
      if r: discard
    of rDebruijn:
      let (r, c) = debruijn.alphaEquivalentCounted(
        debruijn.convertFromNamed(t1), debruijn.convertFromNamed(t2))
      alphaTs.add(nowUs() - t0)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
      if r: discard
    of rLn:
      let (r, c) = locally_nameless.alphaEquivalentCounted(
        locally_nameless.convertFromNamed(t1),
        locally_nameless.convertFromNamed(t2))
      alphaTs.add(nowUs() - t0)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
      if r: discard
    of rSubst:
      let (r, c) = explicit_subst.alphaEquivalentCounted(
        explicit_subst.convertFromNamed(t1),
        explicit_subst.convertFromNamed(t2))
      alphaTs.add(nowUs() - t0)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
      if r: discard
    let t1n = nowUs()
    case rep
    of rNamed:
      let (_, c) = named.normaliseCounted(t1)
      normTs.add(nowUs() - t1n)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
    of rDebruijn:
      let (_, c) = debruijn.normaliseCounted(debruijn.convertFromNamed(t1))
      normTs.add(nowUs() - t1n)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
    of rLn:
      let (_, c) = locally_nameless.normaliseCounted(
        locally_nameless.convertFromNamed(t1))
      normTs.add(nowUs() - t1n)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
    of rSubst:
      let (_, c) = explicit_subst.normaliseCounted(
        explicit_subst.convertFromNamed(t1))
      normTs.add(nowUs() - t1n)
      trav = max(trav, c.traversals)
      depth = max(depth, c.maxDepth)
  let memAfter = getOccupiedMem()
  %* {"alpha_us": median(alphaTs), "norm_us": median(normTs),
      "traversals": trav, "max_depth": depth,
      "mem_delta_bytes": memAfter - memBefore}

proc cmdBench(path: string, iters: int, outPath: string): int =
  let benches = loadBench(path)
  var repJson = newJObject()
  let reps = [rNamed, rDebruijn, rLn, rSubst]
  for rep in reps:
    var benchJson = newJObject()
    for (name, t) in benches:
      var other = t
      if name.endsWith("-a"):
        let base = name[0 ..< name.len - 2]
        for (n2, t2) in benches:
          if n2 == base & "-b": other = t2
      benchJson[name] = benchRep(rep, t, other, iters)
    repJson[repNames[rep.ord]] = benchJson
  let outDoc = %* {"nim": repJson}
  writeFile(outPath, outDoc.pretty)
  echo "bench: wrote ", outPath
  0

proc main(): int =
  let args = commandLineParams()
  if args.len < 1:
    stderr.writeLine("usage: ego check <t1> <t2> [-r rep] | ego corpus <file> | ego bench <file> [--iters N] [--out file]")
    return 2
  case args[0]
  of "check":
    var rep = "named"
    var terms: seq[string]
    var i = 1
    while i < args.len:
      if args[i] == "-r":
        inc i
        rep = args[i]
      else:
        terms.add(args[i])
      inc i
    if terms.len != 2:
      stderr.writeLine("check requires exactly two terms")
      return 2
    cmdCheck(terms, rep)
  of "corpus":
    if args.len < 2:
      stderr.writeLine("corpus requires a file")
      return 2
    cmdCorpus(args[1])
  of "bench":
    var iters = 5
    var outPath = "results.json"
    var file = ""
    var i = 1
    while i < args.len:
      if args[i] == "--iters":
        inc i
        iters = parseInt(args[i])
      elif args[i] == "--out":
        inc i
        outPath = args[i]
      else:
        file = args[i]
      inc i
    if file == "":
      stderr.writeLine("bench requires a file")
      return 2
    cmdBench(file, iters, outPath)
  else:
    stderr.writeLine("unknown command")
    2


quit(main())