module D = DataType
module T = Tools

let u = UTF8.from_utf8string

exception RegisteredFunction

(** Data format *)

type _ dataFormat =
  | Date: D.Num.t dataFormat    (* Date *)
  | Number: D.Num.t dataFormat  (* Number *)
  | String: UTF8.t dataFormat   (* String result, there is only one representation *)
  | Bool: D.Bool.t dataFormat   (* Boolean result *)

let most_generic_format: type a. a dataFormat -> a dataFormat -> a dataFormat =
  begin fun a b -> match a, b with
  | Number, x -> x
  | x, Number -> x
  | x, _ -> x
end

(*** Type definitions *)

type _ typ =
  | Unit: unit typ
  | Bool: D.Bool.t typ
  | Num: D.Num.t typ
  | String: UTF8.t typ
  | List: 'a typ -> 'a list typ

let t_bool= Bool
let t_int = Num
let t_string = String
let t_list t = List t

let typ_of_format: type a. a dataFormat -> a typ = function
  | Date -> Num
  | Number -> Num
  | String -> String
  | Bool -> Bool

let rec compare_typ: type a b. a typ -> b typ -> (a, b) T.cmp = begin fun a b ->
  match a, b with
  | Unit, Unit     -> T.Eq
  | Bool, Bool     -> T.Eq
  | Num, Num       -> T.Eq
  | String, String -> T.Eq
  | List l1, List l2 ->
      begin match compare_typ l1 l2 with
      | T.Lt -> T.Lt
      | T.Eq -> T.Eq
      | T.Gt -> T.Gt
      end
  | x, y -> if (T.Ex x) > (T.Ex y) then T.Gt else T.Lt
end

let rec print_typ: type a. Format.formatter -> a typ -> unit = fun printer typ -> match typ with
  | Unit ->   Format.fprintf printer "Unit"
  | Bool ->   Format.fprintf printer "Bool"
  | Num  ->   Format.fprintf printer "Num"
  | String -> Format.fprintf printer "String"
  | List t -> Format.fprintf printer "List[%a]"
                print_typ t

let default_format_for_type: type a. a typ -> a dataFormat = function
  | Num -> Date
  | String -> String
  | Bool -> Bool
  | List _ -> raise Errors.TypeError
  | Unit -> raise Errors.TypeError

(** Results format.
    Any value which can be encoded with different representation requires as
    many format than there are representations for this value.
*)

type _ result =
  | Numeric: D.Num.t result (* Any numeric format : the representation depends from the inputs  *)
  | Date: D.Num.t result    (* Date *)
  | Number: D.Num.t result  (* Number *)
  | String: UTF8.t result   (* String result, there is only one representation *)
  | Bool: D.Bool.t result   (* Boolean result *)

let f_num  = Numeric
let f_date = Date
let f_number = Number
let f_string = String
let f_bool = Bool

let specialize_result: type a. a result -> a dataFormat -> a result =
  begin fun a b -> match a, b with
  | Date, _ -> Date
  | _, Date -> Date
  | x, y -> x
end

let typ_of_result: type a. a result -> a typ = function
  | Numeric -> Num
  | Number -> Num
  | Date -> Num
  | Bool -> Bool
  | String -> String

let rec compare_result: type a b. a result -> b result -> (a, b) T.cmp = begin fun a b ->
  match a, b with
  | Bool, Bool     -> T.Eq
  | Numeric, Numeric-> T.Eq
  | String, String -> T.Eq
  | Number, Number -> T.Eq
  | Date, Date -> T.Eq
  | x, y -> if (T.Ex x) > (T.Ex y) then T.Gt else T.Lt

end
(*** Values definitions *)

type 'a value =
  | Bool: D.Bool.t -> D.Bool.t value
  | Num: D.Num.t dataFormat * D.Num.t -> D.Num.t  value
  | String: UTF8.t -> UTF8.t   value
  | List: 'a dataFormat * 'a list -> 'a list value
  | List2: 'a dataFormat * 'a list list -> 'a list list value

(** Get the value out of the box *)
let get_value_content: type a. a value -> a = function
  | Bool b -> b
  | Num (_, n) -> n
  | String s -> s
  | List (t, l) -> l
  | List2 (t, l) -> l

(** Create a value from a known type and an unboxed value *)
let build_value: type a. a dataFormat -> a -> a value = begin fun format content ->
  match (typ_of_format format), content with
    | Unit, _ -> raise Errors.TypeError
    | Bool, x -> Bool x
    | Num, x -> Num (format, x)
    | String, s -> String s
    | List t, l -> raise Errors.TypeError
end

(* Extract the type from a boxed value *)
let type_of_value: type a. a value -> a typ = function
  | Bool b -> Bool
  | Num (n, _) -> Num
  | String s -> String
  | List (t, l) -> List (typ_of_format t)
  | List2 (t, l) -> List (List (typ_of_format t))

let format_of_value: type a. a value -> a dataFormat = function
  | Bool b -> Bool
  | Num (f, _) -> f
  | String s -> String
  | List (t, l)  -> raise Errors.TypeError
  | List2 (t, l) -> raise Errors.TypeError

type existencialResult =
  | Result : 'a value -> existencialResult

(** Catalog for all functions *)
module C = struct

  (** This is the way the function is store in the map.
       We just the return type, and the function itself.

       For Fn1 and T1 constructors, we need to add extra information in the
       GADT signature in order to help the compiler: 'a could be any ('a * 'b),
       ('a * 'b * 'c) and so on…

       Instead of returning a signature with type 'a t_function, we have to
       force it as 'a typ t_function.
    *)
  type _ t_function =
    | Fn1: 'b result * ('a -> 'b) -> 'a typ t_function
    | Fn2: 'c result * ('a -> 'b -> 'c) -> ('a * 'b) t_function
    | Fn3: 'd result * ('a -> 'b -> 'c -> 'd) -> ('a * 'b * 'c) t_function

  (** This is the key for storing functions in the map.
   *)
  type _ sig_typ =
    | T1: 'a typ -> 'a typ t_function sig_typ
    | T2: 'a typ * 'b typ -> ('a * 'b) t_function sig_typ
    | T3: 'a typ * 'b typ * 'c typ -> ('a * 'b * 'c) t_function sig_typ

  let print_sig_typ: type a. Format.formatter -> a sig_typ -> unit = begin fun printer typ ->
    match typ with
    | T1 a -> Format.fprintf printer "(%a)"
        print_typ a
    | T2 (a, b) -> Format.fprintf printer "(%a, %a)"
        print_typ a
        print_typ b
    | T3 (a, b, c) -> Format.fprintf printer "(%a, %a, %a)"
        print_typ a
        print_typ b
        print_typ c
  end

  module ComparableSignature = struct

    type 'a t = 'a sig_typ

    (** Compare two signature *)
    let eq: type a b. a sig_typ -> b sig_typ -> (a, b) T.cmp = begin fun a b ->
      match a, b with
      | T1(a), T1(b) ->
        begin match compare_typ a b with
        | T.Lt -> T.Lt
        | T.Gt -> T.Gt
        | T.Eq -> T.Eq
        end
      | T2(a, b), T2(c, d) ->
        begin match (compare_typ a c) with
        | T.Lt -> T.Lt
        | T.Gt -> T.Gt
        | T.Eq ->
          begin match (compare_typ b d) with
          | T.Lt -> T.Lt
          | T.Gt -> T.Gt
          | T.Eq -> T.Eq
          end
        end
      | T3(a, b, c), T3(d, e, f) ->
        begin match (compare_typ a d) with
        | T.Lt -> T.Lt
        | T.Gt -> T.Gt
        | T.Eq ->
          begin match (compare_typ b e) with
          | T.Lt -> T.Lt
          | T.Gt -> T.Gt
          | T.Eq ->
            begin match (compare_typ c f) with
            | T.Lt -> T.Lt
            | T.Gt -> T.Gt
            | T.Eq -> T.Eq
            end
          end
        end
      | x, y -> if (T.Ex x) > (T.Ex y) then T.Gt else T.Lt
      end

  end

  module Catalog = Map.Make(String)
  module Functions = Tools.Map(ComparableSignature)


  (* This is the map which contains all the registered functions.
     Each name is binded with another map with contains the function for each
     signature.
   *)
  let (catalog:Functions.t Catalog.t ref) = ref Catalog.empty

  (**
     Register a function in the catalog. If the function is already defined,
     raise an exception.
   *)
  let register name signature f = begin

    let name' = String.uppercase_ascii name in
    let map = begin match Catalog.find name' !catalog with
    | exception Not_found ->
        Functions.singleton signature f
    | x ->
        (* We prevent any update to already registered function *)
        if (Functions.mem signature x) then
          raise RegisteredFunction
        else
          Functions.add signature f x
    end in

    catalog := Catalog.add name' map !catalog
  end

  let inject: type a. a result -> (unit -> a dataFormat) -> a -> existencialResult = fun resultFormat f res ->
    let (x:a value) = begin match resultFormat, res with
    | Bool, x -> Bool x
    | Numeric, x -> Num (f (), x)
    | Date, x -> Num(Date, x)
    | Number, x -> Num(Number, x)
    | String, s -> String s
    end in
    Result x

  (** Look in the catalog for a function with the given name and signature *)
  let find_function: type a. string -> a t_function sig_typ -> a t_function = begin fun name signature ->
    Catalog.find (String.uppercase_ascii name) !catalog
    |> Functions.find signature
  end

end

(** Guess the format to use for the result function from the arguments given.
    The most specialized format take over the others.
*)
let guess_format_result: type a. a result -> existencialResult list -> unit -> a dataFormat =
  begin fun init_value values () ->

  let init_typ = typ_of_result init_value in

  (* fold over the arguments, and check if they have the same format *)
  let compare_format (currentResult: a result) (Result value): a result =

    (* If the argument as the same type as the result format, just the most specialized *)
    match compare_typ init_typ (type_of_value value) with
      | T.Eq -> begin match value with
        | Bool b -> Bool
        | String s -> String
        | Num (f, v) -> specialize_result currentResult f
        (* There is no possibility to get init_typ as List typ*)
        | List (f, v) -> raise Errors.TypeError
        | List2 (f, v) -> raise Errors.TypeError
        end
      (* The types differ, handle the special cases for Lists *)
      | _ ->
        begin match value with
        | List (f, v) ->
          begin match compare_typ init_typ (typ_of_format f) with
          | T.Eq -> specialize_result currentResult f
          | _ -> currentResult
          end
        | List2 (f, v) ->
          begin match compare_typ init_typ (typ_of_format f) with
          | T.Eq -> specialize_result currentResult f
          | _ -> currentResult
          end
        | _ -> currentResult
        end in

  begin match List.fold_left compare_format init_value values with
  | String -> String
  | Bool -> Bool
  | Number -> Number
  | Date -> Date
  | Numeric -> Number
  end

end

let register0 name returnType f =
  C.register name (C.T1(Unit)) (C.Fn1 (returnType, f))

let register1 name typ1 returnType f =
  C.register name (C.T1(typ1)) (C.Fn1 (returnType, f))

let register2 name (typ1, typ2) result f =
  C.register name (C.T2(typ1, typ2)) (C.Fn2 (result, f))

let register3 name (typ1, typ2, typ3) result f =
  C.register name (C.T3(typ1, typ2, typ3)) (C.Fn3 (result, f))

let call name args = begin
  let name' = UTF8.to_utf8string name in
  begin try match args with

  | [] ->
    let C.Fn1(ret, f) = C.find_function name' (C.T1 Unit) in
    C.inject ret (fun () -> raise Errors.TypeError) (f ())

  | (Result p1)::[] ->
    let C.Fn1(ret, f) =
      C.find_function name' (C.T1 (type_of_value p1)) in
    C.inject ret (guess_format_result ret args) (f (get_value_content p1))

  | (Result p1)::(Result p2)::[] ->
    let C.Fn2(ret, f) =
      C.find_function name' (C.T2 (type_of_value p1, type_of_value p2)) in
    C.inject ret (guess_format_result ret args) (f (get_value_content p1) (get_value_content p2))

  | (Result p1)::(Result p2)::(Result p3)::[] ->
    let C.Fn3(ret, f) =
      C.find_function name' (C.T3 (type_of_value p1, type_of_value p2, type_of_value p3)) in
    C.inject ret (guess_format_result ret args) (f (get_value_content p1) (get_value_content p2) (get_value_content p3))

  | _ -> raise Not_found
  with Not_found ->
      let signature = List.map (fun (Result x) ->
        let formatter = Format.str_formatter in
        print_typ formatter (type_of_value x);
        Format.flush_str_formatter ()) args in

      raise (Errors.Undefined (name, signature))
  end
end

let repr mapper value = begin

  (** Extract the value from a raw type.
      If the value is Undefined, raise an exception.
   *)
  let extract_value = begin function
    | ScTypes.Num (n,s) -> Result (Num (Number, (D.Num.of_num n)))
    | ScTypes.Bool b    -> Result (Bool b)
    | ScTypes.Date d    -> Result (Num (Date, (D.Num.of_num d)))
    | ScTypes.Str s     -> Result (String s)
    | ScTypes.Undefined -> raise Errors.TypeError
  end in

  (** Extract the value from a raw type.
      If the value is Undefined, provide a default result.
   *)
  let guess_value: type a. a typ -> ScTypes.types -> existencialResult = fun typ value -> begin
    try extract_value value with Errors.TypeError ->
      match typ with
      | Num  ->   Result (Num (Number, (D.Num.nan)))
      | Bool ->   Result (Bool false)
      | String -> Result (String (u""))
      | List x  -> Result (List ((default_format_for_type x), []))
      | Unit -> raise Errors.TypeError
  end in


  let add_elem: type a. a typ -> a list * a dataFormat -> ScTypes.types -> a list * a dataFormat =
  begin fun type_of (result, format_of) next ->
      let Result r = guess_value type_of next in
      begin match compare_typ type_of (type_of_value r) with
      | T.Eq ->
        let l' = (get_value_content r)::result in
        l' , (most_generic_format (format_of_value r) format_of)
      | _ -> raise Errors.TypeError
      end
  end in

  (* Return the result for any expression as an ScTypes.types result *)
  let rec get_repr: type a. a value -> ScTypes.types = begin function
  | Bool b ->  ScTypes.Bool b
  | Num (format, n)  ->  begin match format with
    | Number -> ScTypes.Num (D.Num.to_num n, None)
    | Date -> ScTypes.Date (D.Num.to_num n)
    | _ -> raise Errors.TypeError (* This pattern could be refuted *)
    end
  | String s -> ScTypes.Str s
  | List (t, l) ->
         List.hd l      (* Extract the first element *)
      |> build_value t  (* Convert it in boxed value *)
      |> get_repr       (* Return it's representation *)
  | List2 (t, l) ->
         List.hd l      (* Extract the first element *)
      |> List.hd
      |> build_value t  (* Convert it in boxed value *)
      |> get_repr       (* Return it's representation *)
  end in

  (** Extract the value from an expression.
      [extract typ expr] will evaluate the expression and return it. If the
      result cannot be evaluated (because of references pointing to missing
      values) a default value of type [typ] will be returned.
   *)
  let rec extract = begin function
    (* For a reference to an external we first extract the value pointed  *)
    | ScTypes.Ref r ->
        begin match mapper r with
        | ScTypes.Refs.Single v -> extract_value v
        | ScTypes.Refs.Array1 l ->

          (* Guess the list type from it's first defined element *)
          let Result r = extract_value (List.find ((!=) ScTypes.Undefined) l) in
          let format_of = format_of_value r in
          let type_of = type_of_value r in
          (* Build the list with all the elements *)
          let elems, format = List.fold_left (add_elem type_of) ([], format_of) l in
          Result (List (format, elems))
        | ScTypes.Refs.Array2 l ->
          (* Guess the list type from it's first defined element *)
          let Result r = extract_value (Tools.List.find2 ((!=) ScTypes.Undefined) l) in
          let format_of = format_of_value r in
          let type_of = type_of_value r in
          (* Build the list with all the elements *)
          let elems, format = List.fold_left (fun (result, format_of) elems ->
            let elems, format = List.fold_left (add_elem type_of) ([], format_of) elems in
            elems::result, (most_generic_format format_of format)
            ) ([], format_of) l in
          Result (List2 (format, elems))
        end

    (* Evaluate the expression *)
    | ScTypes.Expression e -> extract e
    | ScTypes.Value v -> extract_value v
    | ScTypes.Call (name, args) ->
        let args' = List.map extract args in
        call name args'
    end
  in
  let Result r = extract value in
  get_repr r
end

let wrap f =
  let old_catalog = !C.catalog in
  Tools.try_finally
    (fun () -> C.catalog := C.Catalog.empty; f ())
    (fun () -> C.catalog := old_catalog)

(* Register the standard functions *)

module MAKE(C: D.COMPARABLE) = struct

  let register t = begin
    register2 "="  (t, t) f_bool C.eq;
    register2 "<>" (t, t) f_bool C.neq;
    register2 ">"  (t, t) f_bool C.gt;
    register2 ">=" (t, t) f_bool C.ge;
    register2 "<"  (t, t) f_bool C.lt;
    register2 "<=" (t, t) f_bool C.le;
  end

end

(* Helper for list functions : reduce over a list of elements *)
let reduce name typ res f = begin
  register1 name (t_list typ) res (fun x ->
    List.fold_left f (List.hd x) x);
  register1 name (t_list (t_list typ)) res (fun x ->
    List.fold_left (List.fold_left f) (List.hd (List.hd x)) x);
end

(* Helper for list functions : fold over a list of elements *)
let fold name t_in t_out f init = begin
  register1 name (t_list t_in) t_out (fun x ->
    List.fold_left f init x);
  register1 name (t_list (t_list t_in)) t_out (fun x ->
    List.fold_left (List.fold_left f) init x);
end


let () = begin

  let module CompareNum = MAKE(D.Num) in
  CompareNum.register t_int;
  register0 "rand"              f_number D.Num.rnd;
  register2 "+"  (t_int, t_int) f_num D.Num.add;
  register2 "-"  (t_int, t_int) f_num D.Num.sub;
  register2 "*"  (t_int, t_int) f_number D.Num.mult;
  register2 "/"  (t_int, t_int) f_number D.Num.div;
  register2 "^"  (t_int, t_int) f_number D.Num.pow;

  register1 "abs" t_int         f_number D.Num.abs;

  fold "sum"     t_int f_number D.Num.add  (D.Num.of_num (Num.num_of_int 0));
  fold "product" t_int f_number D.Num.mult (D.Num.of_num (Num.num_of_int 1));

  reduce "min" t_int f_num D.Num.min; (* Minimum value from a list *)
  reduce "max" t_int f_num D.Num.max; (* Maximum value from a list *)

  let module CompareBool = MAKE(D.Bool) in
  CompareBool.register t_bool;
  register0 "true"                    f_bool (fun () -> D.Bool.true_);
  register0 "false"                   f_bool (fun () -> D.Bool.false_);
  register1 "not"   t_bool            f_bool D.Bool.not;
  register2 "and"   (t_bool, t_bool)  f_bool D.Bool.and_;
  register2 "or"    (t_bool, t_bool)  f_bool D.Bool.or_;
  register2 "xor"   (t_bool, t_bool)  f_bool D.Bool.neq;

  let module CompareString = MAKE(D.String) in
  CompareString.register t_string;

end