open Ego

let reps_of_string = function
  | "named" -> [ `Named ]
  | "debruijn" -> [ `Debruijn ]
  | "ln" -> [ `Ln ]
  | "subst" -> [ `Subst ]
  | "all" -> [ `Named; `Debruijn; `Ln; `Subst ]
  | r -> failwith ("unknown representation: " ^ r)

let rep_name = function
  | `Named -> "named"
  | `Debruijn -> "debruijn"
  | `Ln -> "ln"
  | `Subst -> "subst"

let check_with rep t1 t2 =
  match rep with
  | `Named -> Named.alpha_equivalent t1 t2
  | `Debruijn ->
      Debruijn.alpha_equivalent
        (Debruijn.convert_from_named t1)
        (Debruijn.convert_from_named t2)
  | `Ln ->
      Locally_nameless.alpha_equivalent
        (Locally_nameless.convert_from_named t1)
        (Locally_nameless.convert_from_named t2)
  | `Subst ->
      Explicit_subst.alpha_equivalent
        (Explicit_subst.convert_from_named t1)
        (Explicit_subst.convert_from_named t2)

let cmd_check t1s t2s rep =
  match (Syntax.parse_result t1s, Syntax.parse_result t2s) with
  | Error m, _ | _, Error m ->
      Printf.printf "PARSE-ERROR: %s\n" m;
      exit 2
  | Ok t1, Ok t2 ->
      List.iter
        (fun r ->
          let eq = check_with r t1 t2 in
          Printf.printf "%s: %s\n" (rep_name r) (if eq then "EQUIV" else "NOT-EQUIV"))
        (reps_of_string rep)

let cmd_corpus path =
  let entries = Json_io.load_corpus path in
  let failures = ref 0 in
  List.iter
    (fun (e : Json_io.corpus_entry) ->
      let results =
        List.map (fun r -> (r, check_with r e.t1 e.t2))
          [ `Named; `Debruijn; `Ln; `Subst ]
      in
      List.iter
        (fun (r, got) ->
          match e.expected with
          | Some exp when got <> exp ->
              incr failures;
              Printf.printf "FAIL %s [%s]: expected %s got %s\n" e.name
                (rep_name r)
                (if exp then "EQUIV" else "NOT-EQUIV")
                (if got then "EQUIV" else "NOT-EQUIV")
          | _ -> ())
        results;
      let agree =
        match results with
        | (_, a) :: tl -> List.for_all (fun (_, b) -> b = a) tl
        | [] -> true
      in
      if not agree then begin
        incr failures;
        let rs =
          String.concat ", "
            (List.map (fun (r, g) -> rep_name r ^ "=" ^ string_of_bool g) results)
        in
        Printf.printf "DISAGREE %s: %s\n" e.name rs
      end)
    entries;
  if !failures > 0 then begin
    Printf.printf "%d failure(s)\n" !failures;
    exit 1
  end
  else Printf.printf "corpus: %d entries, 0 failures\n" (List.length entries)

let now_us () = int_of_float (Unix.gettimeofday () *. 1e6)

let median arr =
  Array.sort compare arr;
  arr.(Array.length arr / 2)

type bench_result = {
  alpha_us : int;
  norm_us : int;
  traversals : int;
  max_depth : int;
}

let json_escape s =
  let b = Buffer.create (String.length s + 2) in
  String.iter
    (function
      | '"' -> Buffer.add_string b "\\\""
      | '\\' -> Buffer.add_string b "\\\\"
      | '\n' -> Buffer.add_string b "\\n"
      | '\r' -> Buffer.add_string b "\\r"
      | '\t' -> Buffer.add_string b "\\t"
      | c -> Buffer.add_char b c)
    s;
  Buffer.contents b

let write_results path impl reps =
  let oc = open_out_bin path in
  Fun.protect
    ~finally:(fun () -> close_out_noerr oc)
    (fun () ->
      Printf.fprintf oc "{\"%s\":{" (json_escape impl);
      reps
      |> List.iteri (fun ri (rep, benches) ->
             if ri > 0 then output_char oc ',';
             Printf.fprintf oc "\"%s\":{" (json_escape rep);
             benches
             |> List.iteri (fun bi (name, r) ->
                    if bi > 0 then output_char oc ',';
                    Printf.fprintf oc
                      "\"%s\":{\"alpha_us\":%d,\"norm_us\":%d,\"traversals\":%d,\"max_depth\":%d}"
                      (json_escape name) r.alpha_us r.norm_us r.traversals
                      r.max_depth);
             output_char oc '}');
      output_string oc "}}")

let bench_rep rep t1 t2 iters =
  let alpha_ts = Array.make iters 0 in
  let norm_ts = Array.make iters 0 in
  let trav = ref 0 in
  let depth = ref 0 in
  for i = 0 to iters - 1 do
    let t0 = now_us () in
    let (r, c, d) =
      match rep with
      | `Named -> Named.alpha_equivalent_counted t1 t2
      | `Debruijn ->
          Debruijn.alpha_equivalent_counted
            (Debruijn.convert_from_named t1)
            (Debruijn.convert_from_named t2)
      | `Ln ->
          Locally_nameless.alpha_equivalent_counted
            (Locally_nameless.convert_from_named t1)
            (Locally_nameless.convert_from_named t2)
      | `Subst ->
          Explicit_subst.alpha_equivalent_counted
            (Explicit_subst.convert_from_named t1)
            (Explicit_subst.convert_from_named t2)
    in
    alpha_ts.(i) <- now_us () - t0;
    ignore r;
    if c > !trav then trav := c;
    if d > !depth then depth := d;
    let t0 = now_us () in
    let (_, c, d) =
      let drop (x, c, d) = (ignore x, c, d) in
      match rep with
      | `Named -> drop (Named.normalise_counted t1)
      | `Debruijn ->
          drop (Debruijn.normalise_counted (Debruijn.convert_from_named t1))
      | `Ln ->
          drop
            (Locally_nameless.normalise_counted
               (Locally_nameless.convert_from_named t1))
      | `Subst ->
          drop
            (Explicit_subst.normalise_counted
               (Explicit_subst.convert_from_named t1))
    in
    norm_ts.(i) <- now_us () - t0;
    if c > !trav then trav := c;
    if d > !depth then depth := d
  done;
  {
    alpha_us = median alpha_ts;
    norm_us = median norm_ts;
    traversals = !trav;
    max_depth = !depth;
  }

let cmd_bench path iters out_path =
  let benches = Json_io.load_bench path in
  let reps = [ `Named; `Debruijn; `Ln; `Subst ] in
  let pair_for name t =
    if Filename.check_suffix name "-a" then
      let base = String.sub name 0 (String.length name - 2) in
      match List.assoc_opt (base ^ "-b") benches with Some other -> other | None -> t
    else t
  in
  let rep_json =
    List.map
      (fun rep ->
        let bench_json =
          List.map
            (fun (name, t) ->
              let other = pair_for name t in
              (name, bench_rep rep t other iters))
            benches
        in
        (rep_name rep, bench_json))
      reps
  in
  write_results out_path "ocaml" rep_json;
  Printf.printf "bench: wrote %s\n" out_path

let () =
  let argv = Sys.argv in
  let argc = Array.length argv in
  if argc < 2 then begin
    prerr_endline "usage: ego check <t1> <t2> [-r rep] | ego corpus <file> | ego bench <file> [--iters N] [--out file]";
    exit 2
  end;
  match argv.(1) with
  | "check" ->
      let rep = ref "named" in
      let terms = ref [] in
      let i = ref 2 in
      while !i < argc do
        (match argv.(!i) with
        | "-r" ->
            incr i;
            rep := argv.(!i)
        | s -> terms := s :: !terms);
        incr i
      done;
      (match List.rev !terms with
      | [ t1; t2 ] -> cmd_check t1 t2 !rep
      | _ ->
          prerr_endline "check requires exactly two terms";
          exit 2)
  | "corpus" ->
      if argc < 3 then begin
        prerr_endline "corpus requires a file";
        exit 2
      end;
      cmd_corpus argv.(2)
  | "bench" ->
      let iters = ref 5 in
      let out = ref "results.json" in
      let file = ref None in
      let i = ref 2 in
      while !i < argc do
        (match argv.(!i) with
        | "--iters" ->
            incr i;
            iters := int_of_string argv.(!i)
        | "--out" ->
            incr i;
            out := argv.(!i)
        | s -> file := Some s);
        incr i
      done;
      (match !file with
      | Some f -> cmd_bench f !iters !out
      | None ->
          prerr_endline "bench requires a file";
          exit 2)
  | _ ->
      prerr_endline "unknown command";
      exit 2