open Printf
exception Var_already_defined
exception Var_undefined of string

module type VARCONTEXT =
  functor (Vars : sig type var_t type array_t val cell2var : array_t -> int list -> var_t end) ->
sig
  type var_t = Vars.var_t
  type arr_t = Vars.array_t
  type t =
    {
      index2var : (int,var_t) Hashtbl.t;
      var2index : (var_t,int) Hashtbl.t;

      index2array : (int,arr_t) Hashtbl.t;
      array2index : (arr_t,int) Hashtbl.t;

      cells2vars : (int * int list, int) Hashtbl.t;

      mutable nextVarIndex : int;
      mutable nextArrayIndex : int;
    }

  val create : unit -> t

  val index_of_var : t -> var_t -> int
  val var_of_index : t -> int -> var_t
  val mem : t -> var_t -> bool

  val index_of_array : t -> arr_t -> int
  val array_of_index : t -> int -> arr_t
  val arraymem : t -> arr_t -> bool

  (* array index -> indices -> corresponding var index *)
  val index_of_cell : t -> int -> int list -> int

  val add : t -> var_t -> int
  val add_array : t -> arr_t -> unit
  val add_cell : t -> arr_t -> int list -> int

  val size : t -> int
  val iter : t -> (var_t -> int -> unit) -> unit
end

(**
 * a custom type to store array cells
 *)
module ArrayCell =
struct
  type arr = Leaf of int array | Node of arr array

  let rec make_array = function
    | [] -> failwith "an array cannot be empty"
    | [n] -> Leaf (Array.make n 0)
    | n::l -> Node (Array.init n (fun _ -> make_array l))
  
  let rec get_leaf = function
    | Leaf a -> begin function
      | [n] -> a.(n)
      | _ -> failwith "wrong array dimension"
    end
    | Node a -> begin function
      | n::l -> get_leaf a.(n) l
      | _ -> failwith "wrong array dimension"
    end

  let rec set_leaf value = function
    | Leaf a -> begin function
      | [n] -> a.(n) <- value
      | _ -> failwith "wrong array dimension"
    end
    | Node a -> begin function
      | n::l -> set_leaf value a.(n) l
      | _ -> failwith "wrong array dimension"
    end

  let rec _get_cells = function
    | Leaf a -> Array.to_list a
    | Node a -> Array.fold_left (fun acc x -> acc @ (_get_cells x)) [] a

  let get_cells x =
    List.sort Pervasives.compare (_get_cells x)

end

module VarContextFunctor =
  functor (Vars : sig type var_t type array_t val cell2var : array_t -> int list -> var_t end) ->
struct
  type var_t = Vars.var_t
  type arr_t = Vars.array_t
  
  type t =
    {
      index2var : (int,var_t) Hashtbl.t;
      var2index : (var_t,int) Hashtbl.t;

      index2array : (int,arr_t) Hashtbl.t;
      array2index : (arr_t,int) Hashtbl.t;

      mutable cells2vars : ArrayCell.arr array;

      mutable nextVarIndex : int;
      mutable nextArrayIndex : int;
    }

  let create () =
    {
      index2var = Hashtbl.create 16;
      var2index = Hashtbl.create 16;

      index2array = Hashtbl.create 16;
      array2index = Hashtbl.create 16;

      cells2vars = [||];

      nextVarIndex = 0;
      nextArrayIndex = 0;
    }

  let index_of_var vc var =
    Hashtbl.find vc.var2index var
  let var_of_index vc ind =
    Hashtbl.find vc.index2var ind
  let mem vc var =
    Hashtbl.mem vc.var2index var

  let index_of_array vc arr =
    Hashtbl.find vc.array2index arr
  let array_of_index vc ind =
    Hashtbl.find vc.index2array ind
  let arraymem vc arr =
    Hashtbl.mem vc.array2index arr

  let index_of_cell vc arrayId indices =
    ArrayCell.get_leaf vc.cells2vars.(arrayId) indices

  let add vc varName =
    if (Hashtbl.mem vc.var2index varName) then (
      raise Var_already_defined
    );
    let index = vc.nextVarIndex in
    vc.nextVarIndex <- vc.nextVarIndex + 1;
    Hashtbl.add vc.index2var index varName;
    Hashtbl.add vc.var2index varName index;
    index

  let add_array vc arrayName dims =
    if (Hashtbl.mem vc.array2index arrayName) then (
      raise Var_already_defined
    );
    let index = vc.nextArrayIndex in
    vc.nextArrayIndex <- vc.nextArrayIndex + 1;
    Hashtbl.add vc.index2array index arrayName;
    Hashtbl.add vc.array2index arrayName index;
    vc.cells2vars <- Array.append vc.cells2vars [|ArrayCell.make_array dims|]

  let add_cell vc arrayName indices =
    let arrayIndex = Hashtbl.find vc.array2index arrayName in
    let cellVar = Vars.cell2var arrayName indices in
    let cellIndex = add vc cellVar in
    ArrayCell.set_leaf cellIndex vc.cells2vars.(arrayIndex) indices;
    cellIndex

  let size vc = Hashtbl.length vc.var2index

  let iter vc f = Hashtbl.iter f vc.var2index
  
end

module StringArrayVars =
struct
  type var_t = string
  type array_t = string

  let cell2var arrayName indices =
    List.fold_left (fun s -> fun i -> s ^ "[" ^ (string_of_int i) ^ "]") arrayName indices
end

module VarContext = VarContextFunctor(StringArrayVars)

module ScopeStringArrayVars =
struct
  type var_t = int option * string
  type array_t = int option * string

  let cell2var (procRef,arrayName) indices =
    procRef, (List.fold_left (fun s -> fun i -> s ^ "[" ^ (string_of_int i) ^ "]") arrayName indices)
end

module ScopeVarContext =
struct
  include VarContextFunctor(ScopeStringArrayVars)

  (* read all elements, by prepending the process name (if any) to the
   * variable and array names
   *)
  let _prepend_procName procs s = function
    | None    -> s
    | Some i  -> procs.(i) ^ "." ^ s

  (* Convert the ScopeVarContext.t to VarContext.t, by prepending the process
   * names to variables and clocks. We just extract these elements from the
   * hash tables, along with their indices, so as to reinsert them in the
   * same order in the new VarContext. These contexts are only for pretty
   * printing and have no role in simulation.
   *)
  let to_vc : t -> string array -> VarContext.t = fun svc -> fun procs ->
    let result = VarContext.create () in
    let prep = _prepend_procName procs in
    (* Variable names *)
    let varList =
      List.sort compare
        (Hashtbl.fold (fun (p,v) index acc -> (index, prep v p)::acc) svc.var2index []) in
    List.iter (fun (_,name) -> let _ = VarContext.add result name in ()) varList;
    (* Array names *)
    let arrList =
      List.sort compare
        (Hashtbl.fold (fun (p,v) index acc -> (index, prep v p)::acc) svc.array2index []) in
    List.iter (fun (_,name) -> VarContext.add_array result name [0]) arrList;
    (* Do not forget the mapping from cells to vars *)
    { result with VarContext.cells2vars = svc.cells2vars }

end