(***************************************************************
 SIMPREX adds the simprec construct to Simprex
 ****************************************************************)

module Simprex = struct

  open Sexp
  open List

  exception SyntaxError of string
  exception Unbound of string list 

  (************************************************************
   Abstract Syntax
   ************************************************************)

  type var = string

  type pgm = Pgm of string list * exp (* param names, body *)

  and exp =
      Lit of int (* integer literal with value *)
    | Var of var (* variable reference *)
    | BinApp of binop * exp * exp (* primitive application with rator, rands *)
    | Bind of var * exp * exp (* bind name to value of defn in body *)
    | Sigma of var * exp * exp * exp (* name * lo * hi * body *)
    | Simprec of exp * var * var * exp * exp 
       (* zero-exp * n-var * ans-var * comb-exp * arg-exp *)

  and binop = 
    | Add | Sub | Mul | Div | Rem (* binary arithmetic ops *)


  (************************************************************
   Free Variables 
   ************************************************************)

  module S = Set.Make(String) (* String Sets *)

  (* val addElts : S.elt list -> S.t -> S.t *)
  let addElts strs set = (* add a list of elements to a set *)
    List.fold_right S.add strs set

  (* val listToSet : S.elt list -> S.t *)
  let listToSet strs = addElts strs S.empty

  (* val setToList : S.t -> S.elt list *)
  let setToList set = S.elements set 

  (* val varCheck : pgm -> unit *)
  let rec varCheck pgm = 
    let unbounds = freeVarsPgm pgm 
     in if S.is_empty unbounds then
          () (* OK *)
        else 
          raise (Unbound (S.elements unbounds))

  (* val freeVarsPgm : pgm -> S.t *)
  (* Returns the free variables of a program *)
  and freeVarsPgm (Pgm(fmls,body)) = 
    S.diff (freeVarsExp body) (listToSet fmls)

  (* val freeVarsExp : exp -> S.t *)
  (* Returns the free variables of an expression *)
  (* direct version *)
  and freeVarsExp e = 
    match e with 
      Lit i -> S.empty
    | Var s -> S.singleton s
    | BinApp(_,r1,r2) -> freeVarsExps [r1;r2]
    | Bind(name,defn,body) -> 
        S.union (freeVarsExp defn)
                     (S.diff (freeVarsExp body)
                                  (S.singleton name))
    | Sigma(name,lo,hi,body) -> 
        S.union (S.diff (freeVarsExp body)
			          (S.singleton name))
                     (S.union (freeVarsExp lo)
                                   (freeVarsExp hi))
    | Simprec(zero,n,a,comb,arg) -> S.empty (* replace this stub *)

  (* val freeVarsExps : exp list -> S.t *)
  (* Returns the free variables of a list of expressions *)
  (* direct version *)
  and freeVarsExps es = 
    List.fold_right S.union (List.map freeVarsExp es) S.empty

  (************************************************************
   Substitution & Renaming
   ************************************************************)

  (* val subst : exp -> exp Env.env -> exp *)
  let rec subst exp env = 
    match exp with 
      Lit i -> exp
    | Var v -> (match Env.lookup v env with 
	          Some e -> e
                | None -> exp)
    | BinApp(op,r1,r2) -> BinApp(op, subst r1 env, subst r2 env)
    | Bind(name,defn,body) -> 
        (* Take the simple approach of renaming every name. 
           With more work, could avoid renaming unless absolutely necessary. *)
	let name' = StringUtils.fresh name in 
          Bind(name', subst defn env, subst (rename1 name name' body) env)
	  (* note: could be cleverer and do a single substitution/renaming *)
    | Sigma(name,lo,hi,body) -> 
	let name' = StringUtils.fresh name in 
          Sigma(name', subst lo env, subst hi env, subst (rename1 name name' body) env)
    | Simprec(zero,n,a,comb,arg) -> exp (* replace this stub *)

  (* val subst1 : var -> exp -> exp -> exp *)
  and subst1 name newexp exp = subst exp (Env.make [name] [newexp])

  (* val subst_all : string list -> exp list -> exp -> exp *)
  and subst_all names newexps exp = subst exp (Env.make names newexps)

  (* val rename1 : var -> var -> exp -> exp *)
  and rename1 oldname newname exp = subst1 oldname (Var newname) exp 

  (* val rename_all : string list -> var list -> exp -> exp *)
  and rename_all oldnames newnames exp = 
    subst_all oldnames (List.map (fun s -> Var s) newnames) exp 

  (************************************************************
   Parsing from S-Expressions 
   ************************************************************)

  (* val sexpToPgm : Sexp.sexp -> pgm *)
  let rec sexpToPgm sexp = 
    match sexp with 
      Seq [Sym "simprex"; Seq formals; body] -> 
            Pgm(List.map symToString formals, sexpToExp body)
    | _ -> raise (SyntaxError ("invalid Simprex program: " ^ (sexpToString sexp)))

  (* val symToString : Sexp.sexp -> string *)
  and symToString sexp =
    match sexp with 
      Sym s -> s
    | _ -> raise (SyntaxError ("symToString: not a string -- " ^ (sexpToString sexp)))

  (* val sexpToExp : Sexp.sexp -> exp *)
  and sexpToExp sexp = 
    match sexp with 
      Int i -> Lit i
    | Sym s -> Var s
    | Seq [Sym p; rand1; rand2] -> 
        BinApp(stringToBinop p, sexpToExp rand1, sexpToExp rand2)
    | Seq [Sym "bind"; Sym name; defn; body] -> 
        Bind (name, sexpToExp defn, sexpToExp body)
    | Seq [Sym "sigma"; Sym name; lo; hi; body] -> 
        Sigma (name, sexpToExp lo, sexpToExp hi, sexpToExp body)
    (* add a clause here to parse simprec *)
    | _ ->  raise (SyntaxError ("invalid Simprex expression: "  
                                ^ (sexpToString sexp)))

  (* val stringToBinop : string -> binop *)
  and stringToBinop s = 
    match s with 
    | "+" -> Add
    | "-" -> Sub
    | "*" -> Mul
    | "/" -> Div
    | "%" -> Rem
    |  _ -> raise (SyntaxError ("invalid Simprex primop: " ^ s))

  (* val stringToExp : string -> exp *)
  and stringToExp s = sexpToExp (stringToSexp s) 

  (* val stringToPgm : string -> pgm *)
  and stringToPgm s = sexpToPgm (stringToSexp s)

  (************************************************************
   Unparsing to S-Expressions
   ************************************************************)

  (* val pgmToSexp : pgm -> Sexp.sexp *)
  let rec pgmToSexp p = 
    match p with 
      Pgm (fmls, e) -> 
        Seq [Sym "sigmex"; Seq(List.map (fun s -> Sym s) fmls); expToSexp e]

  (* val expToSexp : exp -> Sexp.sexp *)
  and expToSexp e = 
    match e with 
      Lit i -> Int i
    | Var s -> Sym s
    | BinApp (rator, rand1, rand2) -> 
        Seq [Sym (binopToString rator); expToSexp rand1; expToSexp rand2]
    | Bind(n,d,b) -> Seq [Sym "bind"; Sym n; expToSexp d; expToSexp b]
    | Sigma(name,lo,hi,body) -> 
        Seq [Sym "sigma"; Sym name; expToSexp lo; expToSexp hi; expToSexp body]
    | Simprec(zero,n,a,comb,arg) -> Seq[] (* replace this stub *)

  (* val binopToString : binop -> string *)
  and binopToString p = 
    match p with 
    | Add -> "+"
    | Sub -> "-" 
    | Mul -> "*" 
    | Div -> "/" 
    | Rem -> "%" 

  (* val expToString : exp -> string *)
  and expToString s = sexpToString (expToSexp s)

  (* val pgmToString : pgm -> string *)
  and pgmToString s = sexpToString (pgmToSexp s)

  (************************************************************
   Simple Testing
   ************************************************************)

  and freeVarsPgmString str = setToList (freeVarsPgm (stringToPgm str))
  and freeVarsExpString str = setToList (freeVarsExp (stringToExp str))
  and freeVarsPgmFile str = 
    setToList (freeVarsPgm (stringToPgm (File.fileToString str)))
  and freeVarsExpFile str = 
    setToList (freeVarsExp (stringToExp (File.fileToString str)))

  and substString expString envString = 
    StringUtils.println
      (expToString
        (subst (stringToExp expString)
               (stringToExpEnv envString)))

  and stringToExpEnv str =
    let bindings = 
      match stringToSexp str with 
        Seq binds -> List.map (fun bind -> 
                                 match bind with 
                                   Seq [Sym name; sexp] -> 
                                     (name, sexpToExp sexp) 
                                 | _ -> raise (Failure ("wrong format for binding"
							 ^ (sexpToString bind))))
                               binds
      |	_ -> raise (Failure ("wrong format for bindings" ^ str)) in 
    let (names,exps) = ListUtils.unzip bindings in 
    Env.make names exps 


end