open Shdl_types

(* Internal representation *)

module Ctx =
  struct
    type value = bool

    and index = string

    and element = User of (t option -> value list -> value list)
                | Base of (unit -> value list -> value list)
                | Variable of value list

    and t = {
      dict   : (index, element) Hashtbl.t ;
      parent : t option ;
    }

    exception Index_not_found of index
    exception Box_needed of index
    exception Variable_needed of index
    exception Overriding_error of index
    exception Undeclared_error of index
    exception Out_of_bound_error of index * int

    let create size parent =
      { dict = Hashtbl.create size ; parent = parent }

    let rec lookup t idx =
      try
        Hashtbl.find t.dict idx
      with Not_found ->
        begin
          match t.parent with
          | None -> raise (Index_not_found(idx))
          | Some p -> lookup p idx
        end

    let register t idx elem =
      if Hashtbl.mem t.dict idx then
        raise (Overriding_error idx)
      else
        Hashtbl.add t.dict idx elem

    let register_base t idx gen = register t idx (Base gen)

    let register_user t idx gen = register t idx (User gen)

    let declare = 
      let rec make_var size acc =
        if size = 0 then Variable(acc)
        else make_var (size-1) (false::acc)
      in
      fun t size idx ->
        let elem = (make_var size []) in
        register t idx elem


    let instanciate t idx =
      match lookup t idx with
      | User gen -> gen (Some t)
      | Base gen -> gen ()
      | other -> raise (Box_needed idx)

    let assign t idx content =
      if Hashtbl.mem t.dict idx then
        let elem = Variable content in
        let _ = Hashtbl.replace t.dict idx elem in
        content
      else
        raise (Undeclared_error idx)

    let assign_one =
      let rec replace l nth v acc =
        match nth, l with
        | 0, h::r -> (List.rev (v::acc))@r
        | n, h::r -> replace r (n-1) v (h::acc)
        | _       -> raise (Invalid_argument "")
      in
      fun t idx n value ->
        try
          match Hashtbl.find t.dict idx with
          | Variable vl ->
              Hashtbl.replace t.dict idx (Variable (replace vl n value []))
          | _ -> raise (Variable_needed idx)
        with
        | Not_found -> raise (Undeclared_error idx)
        | Invalid_argument _ -> raise (Out_of_bound_error (idx, n))


    let input =
      let rec sublist start size l acc =
        match start, size, l with
        | (0, n, h::t) when n > 0 -> sublist 0 (n-1) t (h::acc)
        | (n, _, h::t) when n > 0 -> sublist (n-1) size t acc
        | _                       -> List.rev acc
      in
      fun t inputs start size idx ->
        let content = sublist start size inputs [] in
        assign t idx content
      
    let print t =
      Hashtbl.iter
        ( fun idx elem -> match elem with
            | Variable vl ->
              begin
                print_string ("var" ^ " : " ^ idx ^ " : ") ;
                List.iter
                  ( fun b -> print_string ((string_of_bool b) ^ ";") )
                  vl ;
                print_newline ()
              end
            | _ -> print_endline ("box : " ^ idx) )
        t.dict
  end

(* Examples *)

let toplevel = Ctx.create 10 None

let not_box =
  let not_base = fun inputs -> List.map ( fun b -> not b ) inputs in
  fun () -> not_base

let _ = Ctx.register_base toplevel "not" not_box

let and_box =
  let and_base = fun inputs -> [List.for_all ( fun b -> b ) inputs] in
  fun () -> and_base

let _ = Ctx.register_base toplevel "and" and_box

let or_box =
  let or_base = fun inputs -> [List.exists ( fun b -> b ) inputs] in
  fun () -> or_base

let _ = Ctx.register_base toplevel "or" or_box
(*
let mux parent =
  let ctx = Ctx.create 5 parent in
  let _ = Ctx.declare ctx 1 "x"
  and _ = Ctx.declare ctx 1 "y"
  and _ = Ctx.declare ctx 1 "o"
  and _ = Ctx.declare ctx 2 "a"
  and _ = Ctx.declare ctx 1 "sel"
  and f_1 = Ctx.instanciate ctx "and"
  and f_2 = Ctx.instanciate ctx "not"
  and f_3 = Ctx.instanciate ctx "and"
  and f_4 = Ctx.instanciate ctx "or" in
  fun ins ->
    (* inputs *)
    let a   = Ctx.input ctx ins 0 2 "a"
    and sel = Ctx.input ctx ins 2 1 "sel" in
    (* definitions *)
    let x   = Ctx.assign ctx "x" ( f_1 ([] @ (f_2 sel) @ [List.nth a 0] ) ) in
    let y   = Ctx.assign ctx "y" ( f_3 ([] @ sel @ [List.nth a 1] ) ) in
    let o   = Ctx.assign ctx "o" ( f_4 ([] @ x @ y ) ) in
    (* result *)
    [] @ o

let _ = Ctx.register_user toplevel "mux" mux
*)
(** Generation part **)

type pos = int

type ident_repr = Ctx.index * pos option

type expr_repr = Repr_b of bool
               | Repr_v of ident_repr
               | Repr_c of (Ctx.value list -> Ctx.value list) * expr_repr list

type def_repr = Repr_d of ident_repr * expr_repr

(* Inputs helper *)
let inputs_fn ctx inputs (Ident(n,s)) ptr =
  let _ = Ctx.input ctx inputs ptr s n in
  ptr + s

(* Definition helpers *)
let convert_ident (Ident(n,s)) = (n, if s > 0 then Some s else None)

let convert_expr n_translator =
  let rec aux = function
    | Call(n,exprs) -> Repr_c(n_translator n, List.map aux exprs)
    | Boolean b -> Repr_b b
    | Value(i) -> Repr_v (convert_ident i)
  in
  aux

let convert_defs n_translator defs =
  List.map
    ( fun (Definition(i,e)) ->
      Repr_d(convert_ident i, convert_expr n_translator e) )
    defs

let expr_eval ctx =
  let rec aux = function
    | Repr_b b -> [b]
    | Repr_v(idx, pos) ->
      begin
        match Ctx.lookup ctx idx with
        | Ctx.Variable vl ->
          begin
            match pos with
            | None -> vl
            | Some n -> [List.nth vl n]
          end
        | _ -> raise (Ctx.Variable_needed idx)
      end
    | Repr_c(f, inputs) -> f (List.flatten (List.map aux inputs))
  in
  fun expr -> aux expr
  
let def_eval ctx (Repr_d((idx,pos), expr)) =
  match pos with
  | None -> ignore (Ctx.assign ctx idx (expr_eval ctx expr))
  | Some n -> Ctx.assign_one ctx idx n (List.hd (expr_eval ctx expr))

let generate_box box =
  let Box(_, ins, outs, vars, defs) = box in
  let all_ident = List.flatten [ins ; outs ; vars] in
  let ident_asc = List.map (fun (Ident(n,s)) -> (n,s)) all_ident in
  let ctx_size = List.fold_left (fun sum (_, a) -> sum + a) 0 ident_asc in
  
  fun parent ->
    let ctx = Ctx.create ctx_size parent in
    (* Declarations *)
    List.iter ( fun (n,s) -> Ctx.declare ctx s n ) ident_asc ;
    let defs_r = convert_defs ( fun call -> Ctx.instanciate ctx call ) defs in
    (* TODO: Get aux box *)
    fun inputs ->
      (* TODO: Inputs *)
      let _ = List.fold_right (inputs_fn ctx inputs) ins 0 in
      (* TODO: Definitions *)
      List.iter (def_eval ctx) defs_r ; 
      (* TODO: Ouputs *)
      List.fold_left
        ( fun acc (Ident(n,_)) ->
          acc@(expr_eval ctx (Repr_v(n, None))) )
        []
        outs