// Normalization by evaluation
//
// I was playing around with mogensen-scott encoding of the STLC and got carried away.
//
// Inspired Andras Kovacs' by elaboration-zoo  typecheck-HOAS-names
//
// By inspired, I mean that I wrote this after studing Kovacs' stuff and then
// cheated and looked up some of the answers when I got to the infer/check bit.
// Any bugs are mine, though.

// make function types easier to write
type F0<A> = () => A;
type F1<A, B> = (_: A) => B;
type F2<A, B, C> = (a: A, b: B) => C;
type F3<A, B, C, D> = (a: A, b: B, c: C) => D;
type F4<A, B, C, D, E> = (a: A, b: B, c: C, d: D) => E;

const fail = (msg: string) => { throw new Error(msg) }

// Term is both Raw and Term, but let never gets quoted back.
type Term = <M>(
  v: F1<string, M>,
  lam: F2<string, Term, M>,
  app: F2<Term, Term, M>,
  tlet: F4<string, Term, Term, Term, M>,
  tpi: F3<string, Term, Term, M>,
  tu: F0<M>
) => M;
let tvar: F1<string, Term> = (n) => (fv, fl, fa) => fv(n);
let tlam: F2<string, Term, Term> = (n, t) => (fv, fl, fa) => fl(n, t);
let tapp: F2<Term, Term, Term> = (t, u) => (fv, fl, fa) => fa(t, u);
let tlet: F4<string, Term, Term, Term, Term> = (n, a, t, u) => (fv, fl, fa, flet) => flet(n, a, t, u);
let tpi: F3<string, Term, Term, Term> = (n, ty, sc) => (fv, fl, fa, flet, fpi) => fpi(n, ty, sc);
let tu: Term = (fv, fl, fa, flet, fpi, fu) => fu();

type Env = [string, Val][];
let lookup = (env: Env, name: string) => env.find((x) => x[0] == name)?.[1];

type Val = <C>(
  nvar: F1<string, C>,
  napp: F2<Val, Val, C>,
  vlam: F1<F1<Val, Val>, C>,
  vpi: F2<Val, F1<Val, Val>, C>,
  vu: F0<C>
) => C;
let nvar: F1<string, Val> = (n) => (nv, na) => nv(n);
let napp: F2<Val, Val, Val> = (t, u) => (nv, na) => na(t, u);
let vlam: F1<F1<Val, Val>, Val> = (sc) => (nv, na, nl) => nl(sc);
let vpi: F2<Val, F1<Val, Val>, Val> = (ty, sc) => (nv, na, nl, npi) => npi(ty, sc);
let vu: Val = (nv, na, nl, npi, vu) => vu();

let vapp = (t:Val, u: Val): Val =>
  t(
    () => napp(t, u),
    () => napp(t, u),
    (sc) => sc(u),
    (_, sc) => sc(u),
    () => napp(t, u)
  );

let ev = (env: Env, tm: Term): Val =>
  tm(
    (name) => lookup(env, name) || nvar(name),
    (n, sc) => vlam((v) => ev([[n, v], ...env], sc)),
    (t, u) => vapp(ev(env, t), ev(env, u)),
    (n, ty, t, sc) => ev([[n, ev(env, t)], ...env], sc),
    // Hold onto ty so we can quote it back
    (n, ty, sc) => vpi(ev(env, ty), (v) => ev([[n, v], ...env], sc)),
    () => vu
  );

let quote = (l: number, v: Val): Term =>
  v(
    (n) => tvar(n),
    (t, u) => tapp(quote(l, t), quote(l, u)),
    (sc) => tlam("_" + l, quote(l + 1, sc(nvar("_" + l)))),
    (ty, sc) => tpi("_" + l, quote(l, ty), quote(l + 1, sc(nvar("_" + l)))),
    () => tu
  );

let nf = (t: Term): Term => quote(0, ev([], t));

// same shape, but the Val is the type
type Ctx = Env;

let conv = (env: Env, ty1: Val, ty2: Val): boolean => {
  let push = (x: string): Env => [[x, nvar(x)], ...env];
  let no = () => false;
  // the u, VLam case
  let eta = (sc: F1<Val, Val>) => fresh(env, (x) => conv(push(x), vapp(ty1, nvar(x)), sc(nvar(x))));
  return ty1(
    (n) => ty2((n2) => n == n2, no, eta, no, no),
    (t, u) => ty2(no, (t2, u2) => conv(env, t, t2) && conv(env, u, u2), eta, no, no),
    (sc) =>
      fresh(env, (x) => {
        let eta2 = () => conv(push(x), sc(nvar(x)), vapp(ty2, nvar(x)));
        return ty2(eta2, eta2, (sc2) => conv(push(x), sc(nvar(x)), sc2(nvar(x))), eta2, eta2);
      }),
    (a, b) => ty2(no, no, eta, (a2, b2) => conv(env, a, a2) && fresh(env, (x) => conv(push(x), b(nvar(x)), b2(nvar(x)))), no),
    () => ty2(no, no, eta, no, () => true)
  );
};

let notpi = (ty: Val) => () => fail(`expected pi type, got ${showVal(ty)}`);

let fresh = <A>(e: Env, f: F1<string, A>): A => f(`__${e.length}`);

// v1 - we'll just check/infer.  Later, let's return an elaborated term
let check = (env: Env, ctx: Ctx, t: Term, ty: Val): unknown =>
  t(
    () => conv(env, infer(env, ctx, t), ty) || fail(`conv ${showVal(infer(env, ctx, t))} ${showVal(ty)}`),
    (n, sc) =>
      ty(
        notpi(ty), // ezoo does check/infer to throw
        notpi(ty),
        notpi(ty),
        (a, b) => fresh(env, (x) => check([[n, nvar(x)], ...env], [[n, a], ...ctx], sc, b(nvar(x)))),
        notpi(ty)
      ), // lam
    () => conv(env, infer(env, ctx, t), ty), // app
    (n, a, t, u) => {
      check(env, ctx, a, vu);
      let va = ev(env, a);
      check(env, ctx, t, va);
      check([[n, ev(env, t)], ...env], [[n, va], ...ctx], u, ty);
    }, // let
    () => conv(env, infer(env, ctx, t), ty), // pi
    () => conv(env, infer(env, ctx, t), ty)
  );

let infer = (env: Env, ctx: Ctx, t: Term): Val =>
  t(
    (n) => lookup(ctx, n) || fail(`undefined ${n}`),
    (n, t) => fail(`can't infer lambda`),
    (t, u) => {
      let tty = infer(env, ctx, t);
      return tty(notpi(tty), notpi(tty), notpi(tty), (a, b) => (check(env, ctx, u, a), b(ev(env, u))), notpi(tty));
    },
    (n, a, t, u) => {
      check(env, ctx, a, vu);
      let va = ev(env, a);
      check(env, ctx, t, va);
      return infer([[n, ev(env, t)], ...env], [[n, va], ...ctx], u);
    },
    (n, a, b) => {
      check(env, ctx, a, vu);
      check([[n, nvar(n)], ...env], [[n, ev(env, a)], ...ctx], b, vu);
      return vu;
    }, // pi
    () => vu
  );

let showTerm = (t: Term): string =>
  t(
    (n) => n,
    (n, sc) => `(λ ${n}. ${showTerm(sc)})`,
    (t, u) => `(${showTerm(t)} ${showTerm(u)})`,
    (n, a, t, u) => `(let ${n} : ${showTerm(a)} = ${showTerm(t)}; ${showTerm(u)})`,
    (n, a, b) => `(∏(${n} : ${showTerm(a)}). ${showTerm(b)})`,
    () => "U"
  );
let showVal = (v: Val): string => showTerm(quote(0, v));
// using Π a : A . B for telescopes to keep the parser simple
let parse = (x: string): Term => {
  let toks = x.match(/\w+|[&|=]+|\S/g)!;
  let next = () => toks.shift()!;
  let must = (s: string) => (toks[0] == s ? next() : fail(`Expected ${s} got ${next()}`));
  let stop = ") ; λ . =".split(" ");
  return (function expr(prec: number): Term {
    let left: Term;
    let t = next();
    // C doesn't guarantee those arguments are evaluated in order, I hope JS does.
    if (t == "let") left = ((n, _1, t, _2, a, _3, sc) => tlet(n, t, a, sc))(next(), must(":"), expr(0), must("="), expr(0), must(";"), expr(0));
    else if (t == "λ") left = ((n, _, sc) => tlam(n, sc))(next(), must("."), expr(0));
    else if (t == "U") left = tu;
    // the greek keymap and Opt-Sh-P give me two different characters
    else if (t == "Π" || t == "∏") left = ((n, _1, a, _2, b) => tpi(n, a, b))(next(), must(":"), expr(0), must("."), expr(0));
    else if (t == "(") left = ((t, _) => t)(expr(0), must(")"));
    else left = tvar(t); // should check t is ident
    while (prec == 0 && toks[0] && !stop.includes(toks[0])) left = tapp(left, expr(1));
    return left;
  })(0);
};

function testnf(s: string) {
  let tm = parse(s);
  console.log(showTerm(tm));
  console.log(showTerm(nf(tm)));
}
testnf("λx.x y z");
testnf("(λx.y x) z");
testnf("(λz.λx.x z) x");

function testInfer(s: string) {
  console.log("parsing", s);
  let tm = parse(s);
  console.log("parsed", showTerm(tm));
  let ty = infer([], [], tm);
  console.log(showTerm(nf(tm)), ":", showVal(ty));
}

parse("Π Α : U . A");

// I'll probably want some sugar on the ∏ and maybe the λ
testnf(`
let id : ∏  A : U . ∏ x : A . A = λ A . λ x . x ;
let const : ∏ A : U . ∏ B : U . ∏ x : A . ∏ y : B. A = λ A . λ B . λ x . λ y . x;
let Nat : U = ∏ N : U . ∏ f : (∏ _ : N . N) . ∏ _ : N . N;
let five : Nat = λ N . λ s . λ z . s (s (s (s (s z))));
let add : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N s (b N s z);
let mul : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N (b N s) z;
let ten : Nat = add five five;
let hundred : Nat = mul ten ten;
let thousand : Nat = mul ten hundred;
thousand
`);

testInfer(`
let id : ∏  A : U . ∏ x : A . A = λ A . λ x . x ;
let const : ∏ A : U . ∏ B : U . ∏ x : A . ∏ y : B. A = λ A . λ B . λ x . λ y . x;
let Nat : U = ∏ N : U . ∏ f : (∏ _ : N . N) . ∏ _ : N . N;
let five : Nat = λ N . λ s . λ z . s (s (s (s (s z))));
let add : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N s (b N s z);
let mul : ∏ _ : Nat . ∏ _ : Nat . Nat = λa.λb.λN.λs.λz. a N (b N s) z;
let ten : Nat = add five five;
let hundred : Nat = mul ten ten;
let thousand : Nat = mul ten hundred;
U
`);