(* Code for random-access-list from
 * http://www.ai.univie.ac.at/~markus/ocaml_sources/pure_fun-1.0-2/chp9.ml
 *)

exception Error

(*
exception Empty
exception Subscript
exception Impossible_pattern of string

let _ = Random.self_init()

let impossible_pat x = raise (Impossible_pattern x)

module type RANDOM_ACCESS_LIST = sig
  type 'a ra_list

  val empty : 'a ra_list
  val is_empty : 'a ra_list -> bool

  val cons : 'a -> 'a ra_list -> 'a ra_list
  val head : 'a ra_list -> 'a
  val tail : 'a ra_list -> 'a ra_list
    (* head and tail raise Empty if list is empty *)

  val lookup : int -> 'a ra_list -> 'a
  val update : int -> 'a -> 'a ra_list -> 'a ra_list
    (* lookup and update raise Subscript if index is out of bounds *)
end


module RaList : RANDOM_ACCESS_LIST = struct
  type 'a tree = Leaf of 'a | Node of int * 'a tree * 'a tree
  type 'a digit = Zero | One of 'a tree
  type 'a ra_list = 'a digit list

  let empty = []
  let is_empty ts = ts = []

  let size = function
    | Leaf x -> 1
    | Node (w, _, _) -> w

  let link t1 t2 = Node (size t1 + size t2, t1, t2)

  let rec cons_tree t = function
    | [] -> [One t]
    | Zero :: ts -> One t :: ts
    | One t' :: ts -> Zero :: cons_tree (link t t') ts

  let rec uncons_tree = function
    | [] -> raise Empty
    | [One t] -> t, []
    | One t :: ts -> t, Zero :: ts
    | Zero :: ts ->
        match uncons_tree ts with
        | Node (_, t1, t2), ts' -> t1, One t2 :: ts'
        | _ -> impossible_pat "uncons_tree"

  let cons x ts = cons_tree (Leaf x) ts

  let head ts =
    match uncons_tree ts with
    | Leaf x, _ -> x
    | _ -> impossible_pat "head"

  let tail ts = snd (uncons_tree ts)

  let rec lookup_tree i t = match i, t with
    | 0, Leaf x -> x
    | i, Leaf x -> raise Subscript
    | i, Node (w, t1, t2) ->
        if i < w/2 then lookup_tree i t1
        else lookup_tree (i - w/2) t2

  let rec update_tree i y t = match i, t with
    | 0, Leaf x -> Leaf y
    | _, Leaf x -> raise Subscript
    | _, Node (w, t1, t2) ->
        if i < w/2 then Node (w, update_tree i y t1, t2)
        else Node (w, t1, update_tree (i - w/2) y t2)

  let rec lookup i = function
    | [] -> raise Subscript
    | Zero :: ts -> lookup i ts
    | One t :: ts ->
        if i < size t then lookup_tree i t
        else lookup (i - size t) ts

  let rec update i y = function
    | [] -> raise Subscript
    | Zero :: ts -> Zero :: update i y ts
    | One t :: ts ->
        if i < size t then One (update_tree i y t) :: ts
        else One t :: update (i - size t) y ts
end
*)

type simple_type = string
type complex_type =
  | Map of complex_type list * complex_type
  | Pair of complex_type * complex_type
  | Simple of simple_type
type name = string
type probability = float
type variable = name * probability * simple_type
type func = name * complex_type

type exp =
  | Func of (name * exp list * complex_type) (* name, params, type *)
  | Var of (name * simple_type) (* name, type *)

let rec print_exp = function
  | Func(name, exp_list, typ) ->
    print_string name;
    print_string "(";
    if exp_list != [] then begin
      print_exp (List.hd exp_list);
      ignore (List.map (function e -> print_string "; "; print_exp e) exp_list);
    end;
    print_string ")"
  | Var(name,typ) -> print_string name


let rec random_item fromList =
  let rec aux depth = function
    | i::[] ->
      let choice = Random.int depth + 1 in
      if choice = depth then (i,-1)
      else (i,choice)
    | h::t ->
      let (i, choice) = aux (depth + 1) t in
      if choice = depth then (h, -1)
      else (i, choice)
    | _ -> raise Error
  in
  fst (aux 1 fromList)


(* Eliminates any non-targetType from the list, and then chooses *)
let rand_var var_list targetType =
  let rec aux_withrand targetType = function
    | [] -> []
    | ((name, prob, typ) as h)::t ->
      if typ = targetType then
        if Random.float 1.0 > prob then (* skip this one? *)
          h::(aux_withrand targetType t)
        else aux_withrand targetType t
      else aux_withrand targetType t
  in
  let rec aux targetType = function
    | [] -> []
    | ((name, prob, typ) as h)::t ->
      if typ = targetType then
        h::(aux targetType t)
      else aux targetType t
  in
  let narrowed_list = aux_withrand targetType var_list in
  if narrowed_list = [] then
    random_item (aux targetType var_list)
  else random_item narrowed_list


(* Grab a random function for the given targetType *)
let rand_func func_list targetType =
  let rec aux targetType = function
    | [] -> []
    | ((name, Map(from_type, to_type)) as h)::t ->
      if to_type = targetType then
        h::(aux targetType t)
      else aux targetType t
    | _ -> raise Error
  in
  random_item (aux targetType func_list)


(* take a number and a size and break it up into a list so that the sum of the
 * elements in the list equals the number. This should preferably be pretty
 * evenly distributed over the list, but every combination should also be
 * attainable. Right now it is not even enough. *)
let rec int_to_list num size =
  if (size < 2) or (num < 1) then [num]
  else
    (* let sub = (Random.int (num / size)) + 1 in *)
    let sub = (Random.int (num - size + 1)) + 1 in
    sub::int_to_list (num - sub) (size - 1)


(* Create a function provided with some restrictions *)
let solution
  vars terms nterms targetSize targetType pterm =
  if targetSize < 2 then (* We need to fill in variables *)
    let (name, prob, typ) = rand_var vars targetType in
    Var(name,typ)
  else
    let (name, typ) = rand_func
      (if Random.float 1.0 < pterm then nterms else terms)
      (Simple targetType) in
    Func(name,[],typ)


(* Do the demo *)
let demo () =
  let vars = [
    "nil",0.0,"list";
    "Xs1",0.0,"list";
    "Ys",0.0,"list";
    "X1",0.7,"int";
    "X2",0.0,"int";
  ] in
  let term_funcs = [
    "f1", Map([Simple "list"],Simple "list");
    "cons", Map([Simple "int"; Simple "list"], Simple "list");
    "add", Map([Simple "int"; Simple "int"], Simple "int");
    "mul", Map([Simple "int"; Simple "int"], Simple "int");
  ] in
  let nterm_funcs = [
    "g1", Map([Simple "int"; Simple "list"], Simple "list");
    "g2", Map([Simple "list"; Simple "list"], Simple "list");
  ] in
  let targetType = "list" in
  let targetSize = 8 in
  let pterm = 0.7 in
  print_exp (solution vars term_funcs nterm_funcs targetSize targetType pterm);
  print_newline()