(* lambda.ml *)
(* Contains all the code needed for pure lambda calculus *)
(* 2001.10.29 - Merged Lambda.ml and Genome.ml *)

type exp =
    (* An integer number *)
  | Num of int
    (* Binding depth of a variable, de Bruijn style *)
  | Var of string * int
    (* Lambda expression *)
  | Lambda of string * exp
    (* Function application *)
  | Apply of exp * exp
    (* An Ocaml primitive, taking and returning an expression *)
    (* The string is a comment for printing and debugging *)
  | Prim of string * (exp -> exp)

exception NotLegal of string
exception Timeout

(* This is called when sigalrm goes off, raises a Timeout *)
let timeout n =
  if n = Sys.sigalrm then raise Timeout

let var_env : (string * exp) list ref = ref []
let func_env : (string * (exp -> exp)) list ref = ref []

let debug = false

let rec print = function
  | Num(num) ->
    print_int num
  | Var(name,num) ->
    if (name = "") then begin
      print_string "#";
      print_int num
    end else
      print_string name
  | Lambda(name,ex) ->
    print_string "\\";
    if not (name = "") then begin
      print_string name;
      print_string "."
    end;
    print ex
  | Apply(e1,e2) ->
    print_string "(";
    print e1;
    print_string " ";
    print e2;
    print_string ")"
  | Prim(name,f) ->
    print_string "!";
    print_string name

let rec print_caml = function
  | Num(num) ->
    print_int num
  | Var(name,num) ->
    if (name = "") then begin
      print_string "#";
      print_int num
    end else
      print_string name
  | Lambda(name,ex) ->
    print_string "function ";
    if not (name = "") then begin
      print_string name;
      print_string " -> "
    end;
    print_caml ex
  | Apply(e1,e2) ->
    print_string "(";
    print_caml e1;
    print_string " ";
    print_caml e2;
    print_string ")"
  | Prim(name,f) ->
    print_string name

(* Changes an expression to De Bruijn notation *)
(* Re-number variables *)
let rec bruijnize = function
  | Lambda(name,ex) ->
    Lambda("",bruijnize (replace name ex 0))
  | Apply(e1,e2) ->
    Apply(bruijnize e1, bruijnize e2)
  | _ as x -> x
  (*| Prim(name,f) -> Prim(name,f)
  | Num(num) -> Num(num)
  | Var(name,num) -> Var(name,num) *)

(* Does a recursive replacement in a bruijn re-write *)
and replace var body depth =
  match body with
  | Var(name,num) ->
    if name = var then Var(name,depth)
    else Var(name,num)
  | Lambda(name,ex) -> Lambda(name,replace var ex (depth + 1))
  | Apply(e1,e2) -> Apply(
      replace var e1 depth,
      replace var e2 depth)
  | _ as x -> x
  (*| Prim(name,f) -> Prim(name,f)
  | Num(num) -> Num(num) *)

let rec expand = function
  | Var(name,num) ->
    begin
      try
        let ex = List.assoc name !var_env in
        expand ex
      with Not_found -> Var(name,num)
    end
  | Lambda(name,ex) -> Lambda(name,expand ex)
  | Apply(e1,e2) -> begin
    match e1 with
    | Var(name,num) -> begin
      try
        let f = List.assoc name !func_env in
        Apply(Prim(name,f),expand e2)
      with Not_found -> (* raise (NotLegal ("Unbound variable " ^ v ^ ".")) *)
        Apply(expand e1, expand e2)
      end
    | _ as ee -> Apply(expand e1, expand e2)
    end
  | _ as x -> x
(*  | Num(n) -> Num(n)*)

let rec print_var_env = function
  | [] -> ()
  | (s,v)::e ->
    print_string "let ";
    print_string s;
    print_string " = ";
    print_caml v;
    print_newline();
    print_var_env e

let rand_term () =
  match Random.int 2 with
  | 0 -> Num(Random.int 10)
  | 1 -> Var("",Random.int 5)
  | _ -> raise (NotLegal "RANDOM MAKER ERROR")

let rec rand = function
  | 0 -> rand_term()
  | _ as n ->
    match Random.int 5 with
    | 0 -> rand_term()
    | 1 -> Var(fst (List.nth !var_env (Random.int (List.length !var_env))),0)
    | 2 -> Apply(Var(fst (List.nth !func_env (Random.int (List.length !func_env))),0),
                 rand (n - 1))
    | 3 -> Lambda("",rand (n - 1))
    | 4 -> Apply(rand (n - 1), rand (n - 1))
    | _ -> raise (NotLegal "RANDOM MAKER ERROR")

let add_func name f =
  func_env := (name,f)::!func_env

let add_var name v =
  var_env := (name,v)::!var_env


(* This just gets the depth of an expression. May be useless *)
let rec depth = function
  | Lambda(name,ex) -> 1 + depth ex
  | Apply(e1,e2) ->
    let a = depth e1 in
    let b = depth e2 in
    1 + if a > b then a else b
  | _ -> 1
    
let rec num_nodes = function
  | Lambda(name,ex) -> 1 + num_nodes ex
  | Apply(e1,e2) ->
    let a = num_nodes e1 in
    let b = num_nodes e2 in
    1 + a + b
  | _ -> 1

let debug_alpha_reduce body param depth =
  if debug then begin
    print_string "Alpha ";
    print_int depth;
    print_string ": ";
    print param;
    print_string " -> ";
    print body;
    print_newline()
  end


(* Alpha reduce body with param at the given depth *)
let rec alpha_reduce body param depth =
  debug_alpha_reduce body param depth;
  match body with
  | Var(name,num) ->
    if num = depth then param else body
  | Lambda(name,ex) -> Lambda(name,alpha_reduce ex param (depth + 1))
  | Apply(e1,e2) -> Apply(
      alpha_reduce e1 param depth,
      alpha_reduce e2 param depth)
  | _ -> body
  (*| Prim(f,s) -> body
  | Num(num) -> body *)


let debug_beta_reduce body param =
  if debug then begin
    print_string "Beta: ";
    print body;
    print_string " <- ";
    print param;
    print_newline()
  end


(* Do a beta reduction. May be extracted into the reduce function *)
let rec beta_reduce body param =
  debug_beta_reduce body param;
  match body with
  | Lambda(name,ex) -> alpha_reduce ex param 0
  | _ -> Apply(body,param)


(* Do a single reduction on an expression *)
let rec reduce = function
  | Apply(e1,e2) -> begin
    match e1 with
    | Prim(name,f) ->
      reduce (f (reduce e2))
    | _ ->
      let ex = beta_reduce (reduce e1) e2 in
      reduce ex
    end
  | _ as x -> x (* If it wasn't an application, do nothing *)


(* Reduce for a certain length of time before throwing an error *)
let timed_reduce ex length =
  Sys.set_signal Sys.sigalrm (Sys.Signal_handle timeout);
  try
    let _ = Unix.alarm(length) in (* Set the alarm for length seconds *)
    let new_ex = reduce ex in 
    let _ = Unix.alarm(0) in (* Turn off the alarm *)
    new_ex;
  with Timeout ->
    let _ = Unix.alarm(0) in (* Turn off the alarm *)
    raise Timeout

let make_pair first second =
  Lambda("",
    Apply(
      Apply(Var("",0),first),
      second))

let get_pair = function
  | Lambda(_,
      Apply(
        (Apply((Var(_,0)), (_ as first))),
        (_ as second))) -> (first,second)
  | _ -> (Var("",-1),Var("",-1))