module D = DataType
module T = Tools

module Data = struct

(** Data format *)

type 'a dataFormat = 'a ScTypes.dataFormat

let most_generic_format: type a. a dataFormat -> a dataFormat -> a dataFormat =
  begin fun a b -> match a, b with
  | ScTypes.Number, x -> x
  | x, ScTypes.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 typ_of_format: type a. a ScTypes.dataFormat -> a typ = function
  | ScTypes.Date -> Num
  | ScTypes.Number -> Num
  | ScTypes.String -> String
  | ScTypes.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 ScTypes.dataFormat = function
  | Num -> ScTypes.Date
  | String -> ScTypes.String
  | Bool -> ScTypes.Bool
  | List _ -> raise Errors.TypeError
  | Unit -> raise Errors.TypeError

type 'a returnType = 'a ScTypes.returnType

let specialize_result: type a. a ScTypes.returnType -> a ScTypes.dataFormat -> a ScTypes.returnType =
  begin fun a b -> match a, b with
  | ScTypes.Num (Some ScTypes.Date) as _1, _ -> _1
  | _, ScTypes.Date -> ScTypes.Num (Some ScTypes.Date)
  | x, y -> x
end

let typ_of_result: type a. a ScTypes.returnType -> a typ = function
  | ScTypes.Num _ -> Num
  | ScTypes.Bool -> Bool
  | ScTypes.Str -> String

(*** Values definitions *)

type 'a value =
  | Bool: D.Bool.t -> D.Bool.t value
  | Num: D.Num.t ScTypes.dataFormat * D.Num.t -> D.Num.t value
  | String: UTF8.t -> UTF8.t   value
  | List: 'a ScTypes.dataFormat * 'a list -> 'a list value
  | Matrix: 'a ScTypes.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
  | Matrix (t, l) -> l

(* 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)
  | Matrix (t, l) -> List (List (typ_of_format t))

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

let inject':
  type a. a ScTypes.returnType -> (unit -> a ScTypes.dataFormat) -> a -> a value =
  fun resultFormat f res -> begin match resultFormat, res with
    | ScTypes.Bool, x -> Bool x
    | ScTypes.Str, s -> String s
    | ScTypes.Num None, x -> Num (f (), x)
    | ScTypes.Num (Some v), x -> Num(v, x)
    end

let compare_format: type a b. a typ -> a ScTypes.returnType -> b value -> a ScTypes.returnType = begin
fun init_typ currentResult value ->

  (* If the argument as the same type as the result format, just select the most specialized *)
  match compare_typ init_typ (type_of_value value) with
    | T.Eq -> begin match value with
      | Bool b -> ScTypes.Bool
      | String s -> ScTypes.Str
      | Num (f, v) -> specialize_result currentResult f
      (* There is no possibility to get init_typ as List typ*)
      | List (f, v) -> raise Errors.TypeError
      | Matrix (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
      | Matrix (f, v) ->
        begin match compare_typ init_typ (typ_of_format f) with
        | T.Eq -> specialize_result currentResult f
        | _ -> currentResult
        end
      | _ -> currentResult
      end
  end

end

module C = Catalog.Make(Data)

let (catalog:C.t ref) = ref C.empty


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

(** 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 ScTypes.returnType -> existencialResult list -> unit -> a Data.dataFormat =
begin fun init_value values () ->

  let init_typ:a Data.typ = Data.typ_of_result init_value in

  (* fold over the arguments, and check if they have the same format *)
  let compare_format: a ScTypes.returnType -> existencialResult -> a ScTypes.returnType =
  fun currentResult (Result value) ->
    Data.compare_format init_typ currentResult value in

  begin match List.fold_left compare_format init_value values with
  | ScTypes.Str    -> ScTypes.String
  | ScTypes.Bool   -> ScTypes.Bool
  | ScTypes.Num None-> ScTypes.Number
  | ScTypes.Num (Some x)-> x
  end

end

let inject:
type a. a Data.returnType -> (unit -> a Data.dataFormat) -> a -> existencialResult =
fun resultFormat f res ->
  let (x:a Data.value) = Data.inject' resultFormat f res in
  Result x

let register0 name returnType f =
  catalog := C.register !catalog name (C.T1(Data.Unit)) (C.Fn1 (returnType, f))

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

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

let register3 name (typ1, typ2, typ3) result f =
  catalog := C.register !catalog 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 !catalog name' (C.T1 Data.Unit) in
    inject ret (fun () -> raise Errors.TypeError) (f ())

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

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

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

  | _ -> raise Not_found
  with Not_found ->
      let signature = List.map (fun (Result x) ->
        let formatter = Format.str_formatter in
        Data.print_typ formatter (Data.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 : ScTypes.result -> existencialResult = begin function
    | ScTypes.Result (ScTypes.Num (f, n)) -> Result (Data.Num (f, n))
    | ScTypes.Result (ScTypes.Bool b)    -> Result (Data.Bool b)
    | ScTypes.Result (ScTypes.Str s)     -> Result (Data.String s)
    | ScTypes.Error x -> raise x
  end in

  let add_elem: type a. a Data.typ -> a list * a Data.dataFormat -> ScTypes.result option -> a list * a Data.dataFormat =
  begin fun type_of (result, format_of) next ->
    let Result r = match next with
    | Some x -> extract_value x
    | None -> begin match type_of with
      | Data.Num  ->   Result (Data.Num (ScTypes.Number, D.Num.nan))
      | Data.Bool ->   Result (Data.Bool false)
      | Data.String -> Result (Data.String (UTF8.empty))
      | Data.List x  -> Result (Data.List ((Data.default_format_for_type x), []))
      | Data.Unit -> raise Errors.TypeError
    end in
    begin match Data.compare_typ type_of (Data.type_of_value r) with
    | T.Eq ->
      let l' = (Data.get_value_content r)::result in
      l' , (Data.most_generic_format (Data.format_of_value r) format_of)
    | _ -> raise Errors.TypeError
    end
  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 ->
            begin match v with
            | None -> raise Errors.TypeError
            | Some v -> extract_value v
            end
        | ScTypes.Refs.Array1 l ->
          (* Guess the list type from it's first defined element *)
          let Result r = extract_value (Tools.List.find_map (fun x -> x) l) in
          let format_of = Data.format_of_value r
          and type_of = Data.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 (Data.List (format, elems))
        | ScTypes.Refs.Array2 l ->
          (* Guess the list type from it's first defined element *)
          let Result r = extract_value (Tools.List.find_map2 (fun x -> x) l) in

          let format_of = Data.format_of_value r
          and type_of = Data.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, (Data.most_generic_format format_of format)
          ) ([], format_of) l in
          Result (Data.Matrix (format, elems))
        end

    (* Evaluate the expression *)
    | ScTypes.Expression e -> extract e
    | ScTypes.Value v -> extract_value (ScTypes.Result v)
    | ScTypes.Call (name, args) ->
        let args' = List.map extract args in
        call name args'
    end
  in
  let Result r = (extract value) in
  begin match r with
  | Data.Bool b ->  ScTypes.Result (ScTypes.boolean b)
  | Data.String s -> ScTypes.Result (ScTypes.string s)
  | Data.Num (format, n)  -> begin match ScTypes.get_numeric_type format with
      | ScTypes.Date -> ScTypes.Result (ScTypes.date n)
      | ScTypes.Number -> ScTypes.Result (ScTypes.number n)
    end
  | _ -> raise Errors.TypeError
  end
end

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


(* Register the standard functions *)
type 'a returnType = 'a ScTypes.returnType

let f_num = ScTypes.f_num
let f_date = ScTypes.f_date
let f_number = ScTypes.f_number
let f_string = ScTypes.f_string
let f_bool = ScTypes.f_bool

module Make_Compare(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

type 'a typ = 'a Data.typ
let t_bool: DataType.Bool.t typ = Data.Bool
let t_int:  DataType.Num.t  typ = Data.Num
let t_string: UTF8.t typ = Data.String
let t_list (t: 'a typ): 'a list typ = Data.List t

(* 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 if_: type a. bool -> a -> a -> a = fun a b c -> if a then b else c


let () = begin

  let module CompareNum = Make_Compare(D.Num) in
  Data.(
  CompareNum.register t_int;
  register0 "rand"              f_number  D.Num.rnd;

  register1 "+"  t_int          f_num     (fun x -> x);
  register1 "-"  t_int          f_num     D.Num.neg;    (* Unary negation *)
  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;

  register3 "if" (t_bool, t_int, t_int) f_number if_;
  register3 "if" (t_bool, t_bool, t_bool) f_bool if_;
  register3 "if" (t_bool, t_string, t_string) f_string if_;

  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_Compare(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_Compare(D.String) in
  CompareString.register t_string;

  (* Build a date *)
  register3 "date"  (t_int, t_int, t_int) f_date (
    fun year month day ->
      Date.get_julian_day
        (Num.int_of_num @@ Num.floor_num @@ D.Num.to_num year)
        (Num.int_of_num @@ Num.floor_num @@ D.Num.to_num month)
        (Num.int_of_num @@ Num.floor_num @@ D.Num.to_num day)
      |> D.Num.of_num
  )
  )

end