type number = I of int | F of float
type op = MUL | ADD | DIV | SUB

type paren = OPN | CLS
type token = P of paren | N of number | O of op | S of string | D

type tree = BranchO of op * (tree list) | BranchD of string * (tree list) | LeafN of number | LeafS of string | LeafU

type stack = Num of number | Op of op * int

exception Parse_error
exception Compile_error

let is_digit ?(zero = true) c = match c with
  |'1' |'2' |'3' |'4' |'5' |'6' |'7' |'8' |'9' -> true
  |'0' when zero -> true
  |_ -> false

let is_char c = List.exists (fun x -> x = c) (BatString.to_list "abcdefghijklmnopqrstuvwxyz") 

let string_of_op op = match op with
  |MUL -> "MUL"
  |ADD -> "ADD"
  |SUB -> "SUB"
  |DIV -> "DIV"

let rec string_of_tokens toks = match toks with
  |[] -> ""
  |P OPN :: ts -> "OPEN " ^ (string_of_tokens ts)
  |P CLS :: ts -> "CLOSE " ^ (string_of_tokens ts)
  |N n :: ts -> "N " ^ (match n with |I i -> string_of_int i |F f -> string_of_float f) ^ " " ^ (string_of_tokens ts)
  |O op :: ts -> (string_of_op op) ^ " " ^ (string_of_tokens ts)
  |S str :: ts -> "SYM " ^ str ^ " " ^ (string_of_tokens ts)
  |D :: ts -> "DEF " ^ (string_of_tokens ts)

let rec string_of_tree tree = match tree with
  |LeafU -> "()"
  |LeafS s -> s
  |LeafN (I n) -> string_of_int n
  |LeafN (F n) -> string_of_float n
  |BranchD (s,subtree) -> "(DEF " ^ s ^ " (" ^ (String.concat " " (List.map string_of_tree subtree)) ^ "))"
  |BranchO (op, subtree) -> "(" ^ (string_of_op op) ^ " (" ^ (String.concat " " (List.map string_of_tree subtree)) ^ "))"

let string_of_number n = match n with
  |F f -> string_of_float f
  |I i -> string_of_int i

let lex s =
  let rec lex_r acc chars =
    match chars with
      |' '::cx -> lex_r acc cx
      |'('::cx -> lex_r (P OPN :: acc) cx
      |')'::cx -> lex_r (P CLS :: acc) cx
      |'+'::cx -> lex_r (O ADD :: acc) cx
      |'-'::cx -> lex_r (O SUB :: acc) cx
      |'/'::cx -> lex_r (O DIV :: acc) cx
      |'*'::cx -> lex_r (O MUL :: acc) cx
      |c::_ when (is_digit ~zero:false c) -> 
	 let beg = BatString.of_list (BatList.take_while is_digit chars) in
	 let rest = BatList.drop_while is_digit chars in begin 
	   match List.hd rest with
	     |'.' -> 
		lex_r 
		  (N (F (float_of_string (beg ^"."^ (BatString.of_list (BatList.take_while is_digit (List.tl rest)))))) :: acc) 
		  (BatList.drop_while is_digit (List.tl rest))
	     |_ -> lex_r (N (I (int_of_string beg)) :: acc) rest end
      |'0'::cx -> lex_r (N (I 0) :: acc) cx
      |'d'::'e'::'f'::cx ->
	 lex_r (D :: acc) cx
      |c::cx  -> lex_r (S (BatString.of_list (BatList.take_while is_char chars)) :: acc) (BatList.drop_while is_char chars)
      |[] -> acc in
    try Some (List.rev (lex_r [] (BatString.to_list s))) with Parse_error -> None

let parse defs ts =
  let rec parse_args prefix toks = 
    let (branch, rest) = parse_r ~prefix:prefix toks in
      match List.hd rest with
	|P CLS -> ([branch], List.tl rest)
	|_ -> let (args, rest) = parse_args prefix rest in
	     (branch :: args, rest)
  and parse_r ?(prefix="") ?(allow_def=true) tokens =
    match tokens with
      |N n::ts -> (LeafN n, ts)
      |P OPN::D::S name::ts when allow_def ->
	 let arg_names = List.tl (BatList.take_while ((<>) (P CLS)) ts) in
	 let rest1 = List.tl (BatList.drop_while ((<>) (P CLS)) ts) in
	 let (body, rest2) = parse_r ~prefix:(name ^ "_") ~allow_def:false rest1 in
	 let fpre t = match t with |S s -> name ^ "_" ^ s |_ -> raise Parse_error in
	   Hashtbl.add defs name (List.map fpre arg_names, body);
	   (LeafU, rest2)
      |P OPN::S sym::ts ->
	 let (args, rest) = parse_args prefix ts in
	   (BranchD (sym, args), rest)
      |P OPN::O op::ts ->
	 let (args, rest) = parse_args prefix ts in
	   (BranchO (op, List.rev args), rest)
      |S s::ts -> (LeafS (prefix ^ s), ts)
      |_ -> raise Parse_error in
    try Some (fst (parse_r ts)) with _ -> None
	     
let compile defs tree = 
  let rec compile_r defs tree = match tree with 
    |LeafU -> []
    |LeafN n -> [Num n]
    |LeafS s -> begin match Hashtbl.mem defs s with
	|true -> compile_r defs (snd (Hashtbl.find defs s))
	|false -> print_endline ("can't find " ^ s); flush stdout; raise Compile_error end
    |BranchO (op,ls) when op = ADD || op = MUL -> 
       (List.concat (List.map (compile_r defs) ls))@[Op (op, List.length ls)]
    |BranchO (op,l1::l2::[]) when op = DIV || op = SUB -> (compile_r defs l1)@(compile_r defs l2)@[Op (op, 2)]
    |BranchD (sym,args) -> begin
       match Hashtbl.mem defs sym with
	 |false -> raise Compile_error
	 |true ->
	    let (symbs, fun_tree) = Hashtbl.find defs sym in
	    let loc = Hashtbl.copy defs in
	      List.iter2 (fun s a -> Hashtbl.add loc s ([], a)) symbs args;
	      compile_r loc fun_tree end
    |_ -> raise Compile_error in
    try Some (compile_r defs tree) with Compile_error -> None

let poly_op int_op float_op n1 n2 = match n1 with
  |I num1 -> begin match n2 with
       |I num2 -> I (int_op num1 num2)
       |F num2 -> F (float_op (float_of_int num1) num2) end
  |F num1 -> begin match n2 with
       |I num2 -> F (float_op num1 (float_of_int num2))
       |F num2 -> F (float_op num1 num2) end

let poly_div x1 x2 = match x1,x2 with
  |F _, I _ |F _, F _ |I _, F _ -> poly_op (/) (/.) x1 x2
  |I n1, I n2 when n1 mod n2 = 0 -> poly_op (/) (/.) x1 x2
  |I n1, I n2 -> poly_op (/) (/.) (F (float_of_int n1)) x2 (* need to make one a float to not cause it to be chopped *)

let run code =
  let rec run_r stack code =
    match code with
      |[] -> if List.length stack > 0 then string_of_number (List.hd stack) else ""
      |(Num n)::xs -> run_r (n::stack) xs
      |(Op (op, argc))::xs -> 
	 let args = BatList.take argc stack in
	 let rst = BatList.drop argc stack in
	   run_r ((BatList.reduce (
		     match op with
		       |ADD -> poly_op (+) (+.)
		       |SUB -> poly_op (-) (-.)
		       |MUL -> poly_op ( * ) ( *. )
		       |DIV -> poly_div) args)::rst) xs in
    run_r [] code

let eval defs s =
  match lex s with
    |None -> "Failed to lex."
    |Some toks -> 
       match parse defs toks with
	 |None -> "Failed to parse."
	 |Some tree ->
	    match compile defs tree with
	      |None -> "Failed to compile."
	      |Some code ->
		 run code

let rec top defs =
  let _ = print_string "# " in
  let r = read_line () in
  let v = eval defs r in
  let _ = print_endline v in
    top defs

let () =
  let defs = Hashtbl.create 20 in
    print_endline "OCScheme v0.01 by daniel patterson";
    print_endline "   for now, just an integer calculator with def's and scheme syntax, using a stack machine.";
    top defs