open Ego

let find_corpus () =
  let candidates =
    [ "../shared/corpus.json"; "shared/corpus.json";
      "/home/bulg/projects/ego/shared/corpus.json" ]
  in
  match List.find_opt Sys.file_exists candidates with
  | Some p -> p
  | None -> failwith "corpus.json not found"

let check_bool name exp got =
  if got <> exp then
    failwith
      (Printf.sprintf "%s: expected %b got %b" name exp got)

let check_raises name f =
  match f () with
  | _ -> failwith (name ^ ": expected exception")
  | exception _ -> ()

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 reps = [ `Named; `Debruijn; `Ln; `Subst ]

let test_corpus () =
  let entries = Json_io.load_corpus (find_corpus ()) in
  List.iter
    (fun (e : Json_io.corpus_entry) ->
      List.iter
        (fun r ->
          match e.Json_io.expected with
          | Some exp ->
              check_bool
                (Printf.sprintf "%s [%s]" e.name
                   (match r with
                    | `Named -> "named" | `Debruijn -> "debruijn"
                    | `Ln -> "ln" | `Subst -> "subst"))
                exp (check_with r e.t1 e.t2)
          | None -> ())
        reps)
    entries

let rng = Random.State.make [| 42 |]

let names = [| "a"; "b"; "c"; "d"; "e"; "f" |]
let fvs = [| "p"; "q"; "g"; "h" |]

let rec gen scope depth =
  let pick arr = arr.(Random.State.int rng (Array.length arr)) in
  if depth <= 0 then
    match scope with
    | [] -> Named.Var (pick fvs)
    | _ ->
        if Random.State.float rng 1.0 < 0.75 then Named.Var (pick (Array.of_list scope))
        else Named.Var (pick fvs)
  else
    let r = Random.State.float rng 1.0 in
    if r < 0.45 then
      let n =
        if scope <> [] && Random.State.float rng 1.0 < 0.35
        then pick (Array.of_list scope)
        else pick names
      in
      Named.Lam (n, gen (n :: scope) (depth - 1))
    else if r < 0.75 then
      Named.App (gen scope (depth - 1), gen scope (depth - 1))
    else if r < 0.85 then
      let n = match scope with [] -> "u" | _ -> pick (Array.of_list scope) in
      Named.Subst (n, gen scope (depth - 1), gen scope (depth - 1))
    else
      match scope with
      | [] -> Named.Var (pick fvs)
      | _ -> Named.Var (pick (Array.of_list scope))

let random_term () = gen [] (2 + Random.State.int rng 4)

let test_reflexivity () =
  for _ = 1 to 200 do
    let t = random_term () in
    List.iter
      (fun r ->
        check_bool "reflexive" true (check_with r t t))
      reps
  done

let test_symmetry () =
  for _ = 1 to 200 do
    let t1 = random_term () in
    let t2 = random_term () in
    List.iter
      (fun r ->
        check_bool "symmetric"
          (check_with r t1 t2) (check_with r t2 t1))
      reps
  done

let test_transitivity () =
  for _ = 1 to 100 do
    let t = random_term () in
    List.iter
      (fun r ->
        let norm_named n =
          match r with
          | `Named -> Named.normalise n
          | _ -> Named.normalise n
        in
        let a = check_with r t t in
        let b = check_with r (norm_named t) (norm_named t) in
        check_bool "transitive-self" a b)
      reps
  done

let test_fv_preserved () =
  for _ = 1 to 200 do
    let t1 = random_term () in
    let t2 = random_term () in
    if Named.alpha_equivalent t1 t2 then
      check_bool "fv equal" true
        (Named.SSet.equal (Named.free_variables (Named.normalise t1))
           (Named.free_variables (Named.normalise t2)))
  done

let test_roundtrip () =
  for _ = 1 to 200 do
    let t = random_term () in
    let t' = Debruijn.convert_to_named (Debruijn.convert_from_named t) in
    check_bool "roundtrip debruijn" true
      (Named.alpha_equivalent t' t);
    let t'' =
      Locally_nameless.convert_to_named
        (Locally_nameless.convert_from_named t)
    in
    check_bool "roundtrip ln" true
      (Named.alpha_equivalent t'' t);
    let t''' =
      Explicit_subst.convert_to_named
        (Explicit_subst.convert_from_named t)
    in
    check_bool "roundtrip subst" true
      (Named.alpha_equivalent t''' t)
  done

let test_cross_rep_normalise () =
  for _ = 1 to 100 do
    let t = random_term () in
    let n1 = Named.normalise t in
    let n2 =
      Debruijn.convert_to_named
        (Debruijn.normalise (Debruijn.convert_from_named t))
    in
    let n3 =
      Locally_nameless.convert_to_named
        (Locally_nameless.normalise (Locally_nameless.convert_from_named t))
    in
    let n4 =
      Explicit_subst.convert_to_named
        (Explicit_subst.normalise (Explicit_subst.convert_from_named t))
    in
    check_bool "named vs debruijn" true
      (Named.alpha_equivalent n1 n2);
    check_bool "named vs ln" true
      (Named.alpha_equivalent n1 n3);
    check_bool "named vs subst" true
      (Named.alpha_equivalent n1 n4)
  done

let test_parse_errors () =
  List.iter
    (fun s ->
      match Syntax.parse_result s with
      | Ok _ -> failwith ("should reject: " ^ s)
      | Error _ -> ())
    [ ""; "\\"; "\\x"; "\\x."; "(\\x. x"; "\\x. x)"; "[x := ] x"; "[ := y] x";
      "((x"; "[x := y]"; "."; "\\. x" ]

let test_malformed () =
  let open Locally_nameless in
  check_bool "dangling bvar" false (lc (BVar 0));
  check_bool "closed lam" true (lc (Lam (BVar 0)));
  check_bool "dangling deep" false
    (lc (Lam (Lam (BVar 2))));
  check_raises "convert raises"
    (fun () -> ignore (convert_to_named (BVar 3)));
  check_raises "alpha raises"
    (fun () -> ignore (alpha_equivalent (BVar 1) (BVar 1)))

let test_debruijn_scope () =
  let open Debruijn in
  check_bool "well scoped" true (well_scoped 0 (Lam (Var 0)));
  check_bool "ill scoped" false (well_scoped 0 (Var 0));
  check_raises "to_named raises"
    (fun () -> ignore (convert_to_named (Lam (Var 2))))

let test_capture_avoidance () =
  let parse s =
    match Syntax.parse_result s with
    | Ok t -> t
    | Error m -> failwith m
  in
  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) ]
  in
  List.iter
    (fun (s1, s2, exp) ->
      let t1 = parse s1 in
      let t2 = parse s2 in
      List.iter
        (fun r ->
          check_bool (s1 ^ " ~ " ^ s2) exp (check_with r t1 t2))
        reps)
    cases

let test_explicit_susp () =
  let open Explicit_subst in
  let t = of_named (Named.Subst ("x", Named.Var "z", Named.Lam ("z", Named.Var "x"))) in
  let n = normalise t in
  (match n with
  | Lam (y, Var z) ->
      check_bool "binder renamed" true (y <> "z");
      check_bool "body is free z" true (z = "z")
  | _ -> failwith "expected renamed lambda");
  check_bool "susp agrees with normal" true
    (alpha_equivalent_susp t (normalise t))

let () =
  let tests =
    [ ("corpus", test_corpus);
      ("reflexivity", test_reflexivity);
      ("symmetry", test_symmetry);
      ("transitivity", test_transitivity);
      ("fv preserved", test_fv_preserved);
      ("roundtrip", test_roundtrip);
      ("cross-rep normalise", test_cross_rep_normalise);
      ("parse errors", test_parse_errors);
      ("local closure", test_malformed);
      ("de bruijn scope", test_debruijn_scope);
      ("capture avoidance", test_capture_avoidance);
      ("suspended subst", test_explicit_susp) ]
  in
  List.iter
    (fun (name, f) ->
      f ();
      Printf.printf "ok %s\n" name)
    tests