(* *************** environments **************  *)

exception NotBound of string
val env0 = fn x => (print x; raise (NotBound x));

fun ext env x v = (fn y => if x=y then v else env y);


(******* Staged Version without monad **************** *)

datatype value = 
    VI of int 
  | VF of (value -> value)
  | VC of exp
and exp  =
    EI of int              (* integers *)
  | EA of exp * exp        (* applications *)
  | EL of string * exp     (* lambda-abstractions *)
  | EV of string           (* variables *)
  | EB of exp              (* brackets *)
  | ES of exp              (* escape *)
  | ER of exp              (* run *)
  | EC of string * value;  (* cross-stage constants *)


fun showe e =
case e of
  EI n => show n
| EA(f,x as EA(_,_)) => (showe f)^" ("^(showe x)^")"
| EA(f,x) => (showe f)^" "^(showe x)
| EL(x,e) => "(fn "^x^" => "^(showe e)^")"
| EV x => x
| EB e => "<"^(showe e)^">"
| ES(EV x) => "~"^x
| ES e => "~("^(showe e)^")"
| ER e => "run("^(showe e)^")"
| EC(s,v) => "%"^s;


val id = VF(fn x => x);
val compose = VF(fn (VF f) => VF (fn (VF g) => VF(fn x => f(g x))));
val square = VF (fn (VI x) => VI(x * x));
val bang =
  let fun fact n = if n=0 then 1 else n*(fact (n-1))
  in VF (fn (VI x) => VI(fact x)) end;

val ID = EL("x",EV "x");
val COMPOSE = EL("f",EL("g",EL("x",EA(EV "f",EA(EV "g",EV "x")))));

(* ************* Constants ********************* *)

fun g s f (VI x) (VI y) = f (x,y)
  | g s f a b = error ("Non int to oper in "^s);

val plus = VF(fn x => VF(fn y => VI(g "plus" (op +) x y)));
val minus = VF(fn x => VF(fn y => VI(g "minus" (op - ) x y)));
val times = VF(fn x => VF(fn y => VI(g "times" (op * ) x y)));

val env1 = ext env0 "*" (EC("*",times));

val SQUARE = EL("x",EA(EA(EV "*",EV "x"),EV "x"));

val Z = fn n => fn tf => fn ff => if n=0 then tf 0 else ff n;
val ifzero = 
  VF(fn (VI n) => 
  VF(fn (VF tf) =>
  VF(fn (VF ff) => if n=0 then tf(VI 0) else ff(VI n))));

val Y = let fun Y' f = f (fn v => (Y' f) v) in Y' end;
fun Y f = f (fn v => (Y f) v);

val recur = 
  let fun recur' (VF f) =
           f(VF (fn v=> (case (recur' (VF f)) of 
                            VF fp => fp v
                         | w => error "off end in recur")))
  in VF recur' end;


fun env0 x = EV x;

val env1 =
 ext(ext(ext(ext(ext env0 "+" (EC("+",plus))) 
             "-" (EC("-",minus))) 
         "*" (EC("*",times))) 
     "Z" (EC("if",ifzero))) 
 "Y" (EC("Y",recur));

(* ************************************ *)

val ctr = ref 0;
fun NextVar s = let val _ = ctr := (!ctr) + 1
                in s^(toString (! ctr)) end;



fun coverE env e =
case e of
  EI i => EI i
| EA(e1,e2) => EA(coverE env e1,coverE env e2)
| EL(x,e1) =>  EL(x,coverE env e1)
| EV y => env y
| EB e1 => EB(coverE env e1)
| ES e1 => ES(coverE env e1)
| ER e1 => ER(coverE env e1)
| EC(s,v) => EC(s,coverV env v)

and coverV env v =
case v of
  VI i => VI i
| VF f => VF((coverV env) o f)
| VC e => VC(coverE env e);



fun ev1 env e = 
case e of
  EI i => VI i
| EA(e1,e2) => (case (ev1 env e1,ev1 env e2) of (VF f,v) => f v)
| EL(x,e1) => VF(fn v => ev1 (ext env x (EC(x,v))) e1)
| EV x => (case env x of EC(_,v) => v | w => VC w)
| EB e1 => VC(eb1 1 env e1)
| ER e1 => (case ev1 env e1 of VC e2 => ev1 env0 e2)
| EC(s,v) => coverV env v
| ES e1 => error "escape at level 0"

and eb1 n env e = 
case e of
  EI i => EI i
| EA(e1,e2) => EA(eb1 n env e1,eb1 n env e2)
| EL(x,e1) => let val x' = NextVar x 
              in EL(x',eb1 n (ext env x (EV x')) e1) end
| EV y => env y
| EB e1 => EB(eb1 (n+1) env e1)
| ES e1 => if n=1
              then (case ev1 env e1 of VC e => e)
              else ES(eb1 (n-1) env e1)
| ER e1 => ER(eb1 n env e1)
| EC(s,v) => EC(s,coverV env v)


(* *************** Now a monadic version ******************** *)

datatype 'a M = M of (int -> ('a * int));

fun return x = M(fn n => (x,n));
fun bind (M f) g =
  M(fn n => let val (a,n1) = f n
                val M h = g a
            in h n1 end);
val m = Mon(return,bind);

fun newvar s = M(fn n => (s^(toString n),n+1));

fun fmap f e = Do m { x <- e; Return m (f x) };
  
(* ************************************************* *)

datatype V = 
    VI of int 
  | VF of (V -> V M)
  | VC of exp
and exp  =
    EI of int              (* integers *)
  | EA of exp * exp        (* applications *)
  | EL of string * exp     (* lambda-abstractions *)
  | EV of string           (* variables *)
  | EB of exp              (* brackets *)
  | ES of exp              (* escape *)
  | ER of exp              (* run *)
  | EC of string * V;      (* cross-stage constants *)


fun showe e =
case e of
  EI n => show n
| EA(f,x as EA(_,_)) => (showe f)^" ("^(showe x)^")"
| EA(f,x) => (showe f)^" "^(showe x)
| EL(x,e) => "(fn "^x^" => "^(showe e)^")"
| EV x => x
| EB e => "<"^(showe e)^">"
| ES(EV x) => "~"^x
| ES e => "~("^(showe e)^")"
| ER e => "run("^(showe e)^")"
| EC(s,v) => "%"^s;



(* *************** environments **************  *)

type env =string -> V;

exception NotBound of string
val env0 = fn x => (print x; raise (NotBound x));

fun ext env x v = (fn y => if x=y then v else env y);

(* ************** Covers ******************** *)


fun coverE env e =
case e of
  EI i => EI i
| EA(e1,e2) => EA(coverE env e1,coverE env e2)
| EL(x,e1) =>  EL(x,coverE env e1)
| EV y => env y
| EB e1 => EB(coverE env e1)
| ES e1 => ES(coverE env e1)
| ER e1 => ER(coverE env e1)
| EC(s,v) => EC(s,coverV env v)

and coverV env v =
case v of
  VI i => VI i
| VF f => VF(fmap (coverV env) o f)
| VC e => VC(coverE env e);




(* ***************** interpreters ******************* *)

fun apply f x = 
  (case (f,x) of (VF f,v) => f v);


fun ev2 e env = 
case e of
  EI i => Return m (VI i)
| EA(e1,e2) => 
   Do m { f <- ev2 e1 env
        ; x <- ev2 e2 env
        ; apply f x
        }
| EL(x,e1) => Return m (VF(fn v => ev2 e1 (ext env x (EC(x,v)))))
| EV x =>
   (case env x of EC(_,v) => 
      Return m v
    | w => Return m (VC w))
| EB e1 => 
   Do m { c <- eb2 1 e1 env
        ; Return m (VC c)
        } 
| ER e1 => 
   Do m { VC c <- ev2 e1 env 
        ; ev2 c env0
        }
| EC(s,v) => Return m (coverV env v)
| ES e1 => error "escape at level 0"

and eb2 n e env = 
case e of
  EI i => Return m (EI i)
| EA(e1,e2) => 
   Do m { f <- eb2 n e1 env
        ; x <- eb2 n e2 env
        ; Return m (EA(f,x)) }
| EL(x,e1) => 
   Do m { x' <- newvar x
        ; body <- eb2 n e1 (ext env x (EV x'))
        ; Return m (EL(x',body)) }
| EV y => Return m (env y)
| EB e1 => 
   Do m { c <- eb2 (n+1) e1 env
        ; Return m (EB c) }
| ES e1 => 
   if n=1
      then Do m { VC c <- ev2 e1 env; Return m c }
      else Do m { c <- eb2 (n-1) e1 env; Return m (ES c) }
| ER e1 => 
   Do m { c <- eb2 n e1 env; Return m (ER c) }
| EC(s,v) => Return m (EC(s,coverV env v));


(* *************** Now a staged version *************** *)


fun F f = Return m (VF f);

fun ev3 e env = 
case e of
  EI i => <Return m (VI ~(lift i))>
| EA(e1,e2) => 
  <Do m { f <- ~(ev3 e1 env)
        ; x <- ~(ev3 e2 env)
        ; apply f x
        }>
| EL(x,e1) => <F (fn v => ~(ev3 e1 (ext env x <EC(~(lift x),v)>)))>
| EV x =>
   (case env x of <EC(~s,~v)> => 
      <Return m ~v>
    | w => <Return m (VC ~w)>)
| EB e1 => 
  <Do m { c <- ~(eb3 1 e1 env)
        ; Return m (VC c)
        }> 
| ER e1 => 
  <Do m { VC c <- ~(ev3 e1 env) 
        ; run (ev3 c env0)
        }>
| EC(s,v) => <Return m (coverV (fn x => run (env x)) v)>
| ES e1 => error "escape at level 0"

and eb3 n e env = 
case e of
  EI i => <Return m (EI ~(lift i))>
| EA(e1,e2) => 
   <Do m { f <- ~(eb3 n e1 env)
         ; x <- ~(eb3 n e2 env)
         ; Return m (EA(f,x)) }>
| EL(x,e1) => 
   <Do m { x' <- newvar ~(lift x)
        ; body <- ~(eb3 n e1 (ext env x <EV x'>))
        ; Return m (EL(x',body)) }>
| EV y => <Return m ~(env y)>
| EB e1 => 
   <Do m { c <- ~(eb3 (n+1) e1 env)
         ; Return m (EB c) }>
| ES e1 => 
   if n=1
      then <Do m { VC c <- ~(ev3 e1 env); Return m c }>
      else <Do m { c <- ~(eb3 (n-1) e1 env); Return m (ES c) }>
| ER e1 => 
   <Do m { c <- ~(eb3 n e1 env); Return m (ER c) }>
| EC(s,v) => <Return m (EC(s,coverV (fn x => run(env x)) v))>


(* ************** Tests *************************** *)


val H2 = let fun minus x y = EA(EA(EV "-",x),y)
             fun plus x y  = EA(EA(EV "+",x),y)
in  EA(EV "Y",EL("h1",EL("n",EL("z",
      EA(EA(EA(EV "Z",EV "n")
           ,EL("w",EV "z"))
        ,EL("n",EB(EA(EL("x",ES(EA(EA(EV "h1",minus (EV "n") (EI 1))
                                  ,EB(plus (EV "x") (ES (EV "z"))))))
                    ,EV "n"))))))))
end;

val IT = EA(EA(H2,EI 3),EB(EI 4));
 

val puzzle = run ((run <fn a => ~( (fn x => <x>) (fn w => <a>) ) 5>) 3);

(* (fn x => <x>) *)
val t1 = EL("x",EB(EV "x")); 

(* (fn w => <a>) *)
val t2 = EL("w",EB(EV "a")); 

(*  <fn a => ~( (fn x => <x>) (fn w => <a>) ) 5> *)
val t3 = EB(EL("a",EA(ES(EA(t1,t2)),EI 5)));

val puzzle = ER(EA(ER t3,EI 3));

(* *********************************************** *)


fun Y f x = f(fn y => Y f y) x
fun env0 "Y" = <(EC("Y",VI 6))>
  | env0 "Z" = <(EC("N",VI 7))>
  | env0 "+" = <(EC("+",VI 7))>
  | env0 "-" = <(EC("-",VI 7))>;

fun go exp = (ev3 exp env0);


val x1 = go puzzle;
val M x2 = run x1;
val x3 = x2 1;

(* 

-| val x1 = 
<Do %m
    { a <- %newvar "a"
    ; d <- %F (fn b => Do %m
                          { c <- Return %m EC("x",b)
                          ; Return %m VC c
                          })
    ; g <- %F (fn e => Do %m
                          { f <- Return %m EV a
                          ; Return %m VC f
                          })
    ; VC h <- %apply d g
    ; i <- Return %m EI 5
    ; j <- Return %m EA(h,i)
    ; k <- Return %m EL(a,j)
    ; VC l <- Return %m VC k
    ; m <- run %ev3 l (%env0)
    ; n <- Return %m VI 3
    ; VC o <- %apply m n
    ; run %ev3 o (%env0)
    }>
 : <V M>

-| go puzzle;
val it = 
<Do %m
    { a <- %newvar "a"
    ; c <- %F (fn b => Return %m (VC (EC("x",b))))
    ; e <- %F (fn d => Return %m (VC (EV a)))
    ; VC f <- %apply c e
    ; VC g <- Return %m (VC (EL(a,EA(f,EI 5))))
    ; h <- run %ev3 g %env0
    ; VC i <- %apply h (VI 3)
    ; run %ev3 i %env0
    }>
 : <V M>


<Do %m
    { p <- %F
           (fn a =>
           %F
           (fn b =>
           %F
           (fn c =>
               Do %m
                  { d <- %apply (VI 7) b
                  ; f <- %F (fn e => Return %m c)
                  ; g <- %apply d f
                  ; o <- %F
                          (fn h =>
                              Do %m
                                 { i <- %newvar "x"
                                 ; j <- %apply (VI 7) h
                                 ; k <- %apply j (VI 1)
                                 ; l <- %apply a k
                                 ; VC m <- Return %m c
                                 ; VC n <- %apply
                                            l
                                            (VC
                                              (EA(EA(EC("+",VI 7),EV i),m)))
                                 ; Return %m (VC (EA(EL(i,n),EC("n",h))))
                                 })
                  ; %apply g o
                  })))
    ; q <- %apply (VI 6) p
    ; r <- %apply q (VI 3)
    ; %apply r (VC (EI 4))
    }>
 : <V M>



*)