open Formula
open Ast

open Util
open Num_domain

open Apron

module ND : NUMERICAL_ENVIRONMENT_DOMAIN = struct

    (* manager *)
    type man = Polka.loose Polka.t
    let manager = Polka.manager_alloc_loose ()

    type itv = Interval.t

    (* abstract elements *)
    type t = man Abstract1.t

    (* direct translation of numerical expressions into Apron tree expressions *)
    (* TODO : some variables have type real and not int... *)
    let rec texpr_of_nexpr = function
        | NIdent id -> Texpr1.Var (Var.of_string id)
        | NIntConst i -> Texpr1.Cst (Coeff.s_of_mpqf (Mpqf.of_int i))
        | NRealConst r -> Texpr1.Cst (Coeff.s_of_mpqf (Mpqf.of_float r))
        | NUnary(AST_UPLUS, e, is_real) ->
            Texpr1.Unop(
                Texpr1.Cast,
                texpr_of_nexpr e,
                (if is_real then Texpr1.Real else Texpr1.Int),
                Texpr1.Rnd)
        | NUnary(AST_UMINUS, e, is_real) ->
            (* APRON bug ? Unary negate seems to not work correctly... *)
            Texpr1.Binop(
                Texpr1.Sub,
                Texpr1.Cst(Coeff.s_of_mpqf (Mpqf.of_string "0")),
                (texpr_of_nexpr e),
                (if is_real then Texpr1.Real else Texpr1.Int),
                Texpr1.Rnd)
        | NBinary(op, e1, e2, is_real) ->
            let f = match op with
            | AST_PLUS -> Texpr1.Add
            | AST_MINUS -> Texpr1.Sub
            | AST_MUL -> Texpr1.Mul
            | AST_DIV -> Texpr1.Div
            | AST_MOD -> Texpr1.Mod
            in
            Texpr1.Binop(
                  f,
                  (texpr_of_nexpr e1),
                  (texpr_of_nexpr e2),
                  (if is_real then Texpr1.Real else Texpr1.Int),
                  Texpr1.Rnd)
    
    (* direct translation of constraints into Apron constraints *)
    let tcons_of_cons env (eq, op) =
        let op = match op with
          | CONS_EQ -> Tcons1.EQ
          | CONS_NE -> Tcons1.DISEQ
          | CONS_GT -> Tcons1.SUP
          | CONS_GE -> Tcons1.SUPEQ
        in Tcons1.make (Texpr1.of_expr env (texpr_of_nexpr eq)) op

    (* implementation *)
    let v_env vars = 
        let intv = List.map
              (fun (s, _) -> Var.of_string s)
              (List.filter (fun (_, n) -> not n) vars) in
        let realv = List.map
              (fun (s, _) -> Var.of_string s)
              (List.filter (fun (_, n) -> n) vars) in
        (Environment.make (Array.of_list intv) (Array.of_list realv))

    let top vars =
      Abstract1.top manager (v_env vars)
    let bottom vars = Abstract1.bottom manager (v_env vars)
    let is_bot = Abstract1.is_bottom manager
    let is_top = Abstract1.is_top manager

    let forgetvar x id =
        let v = [| Var.of_string id |] in
        Abstract1.forget_array manager x v false

    let vars x =
        let (ivt, fvt) = Environment.vars (Abstract1.env x) in
        let ivl = List.map Var.to_string (Array.to_list ivt) in
        let fvl = List.map Var.to_string (Array.to_list fvt) in
        (List.map (fun x -> x, false) ivl) @ (List.map (fun x -> x, true) fvl)

    let vbottom x =
        Abstract1.bottom manager (Abstract1.env x)

    let project x id =
        Abstract1.bound_variable manager x (Var.of_string id)

    (* Apply some formula to the environment *)
    let apply_cl x cons =
        let env = Abstract1.env x in
        let tca = Array.of_list (List.map (tcons_of_cons env) cons) in
        let ar = Tcons1.array_make env (Array.length tca) in
        Array.iteri (Tcons1.array_set ar) tca;
        Abstract1.meet_tcons_array manager x ar

    let apply_cons x cons = apply_cl x [cons]


    let assign x eqs =
      let env = Abstract1.env x in
      let vars = Array.of_list
        (List.map (fun (id, _) -> Var.of_string id) eqs) in
      let vals = Array.of_list
        (List.map (fun (_, v) -> Texpr1.of_expr env (texpr_of_nexpr v)) eqs) in
      Abstract1.assign_texpr_array manager x vars vals None

    
    (* Ensemblistic operations *)
    let join = Abstract1.join manager
    let meet = Abstract1.meet manager
    let widen = Abstract1.widening manager

    let subset = Abstract1.is_leq manager
    let eq = Abstract1.is_eq manager

    (* Pretty-printing *)
    let print fmt x =
        Abstract1.minimize manager x;
        Abstract1.print fmt x;
        Format.fprintf fmt "@?"

    let print_vars fmt x idl = 
        let prevars = vars x in
        let rm_vars_l = List.filter (fun (id, _) -> not (List.mem id idl)) prevars in
        let rm_vars = Array.of_list
          (List.map (fun (id, _) -> Var.of_string id) rm_vars_l) in
        let xx = Abstract1.forget_array manager x rm_vars false in
        print fmt xx

    let print_itv = Interval.print

end