open Formula
open Ast

open Util
open Environment_domain

open Apron

module D : ENVIRONMENT_DOMAIN = struct

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

    (* 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) ->
            texpr_of_nexpr e
        | NUnary(AST_UMINUS, e) ->
            (* 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),
                Texpr1.Int,
                Texpr1.Zero)
        | NBinary(op, e1, e2) ->
            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),
                  Texpr1.Int,
                  Texpr1.Zero)
    
    (* 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 init = Abstract1.top manager (Environment.make  [||] [||])
    let bottom = Abstract1.bottom manager (Environment.make [||] [||])
    let is_bot = Abstract1.is_bottom manager

    let addvar x id isfloat =
        let env = Abstract1.env x in
        let env2 = if isfloat then
            Environment.add env [||] [| Var.of_string id |]
          else
            Environment.add env [| Var.of_string id |] [||]
        in
        Abstract1.change_environment manager x env2 false
    let forgetvar x id =
        let v = [| Var.of_string id |] in
        Abstract1.forget_array manager x v false
    let rmvar x id =
        let v = [| Var.of_string id |] in
        let env = Abstract1.env x in
        let env2 = Environment.remove env v in
        let y = Abstract1.forget_array manager x v false in
        Abstract1.change_environment manager y env2 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)

    (* Apply some formula to the environment *)
    let rec apply_cl x (cons, rest) =
        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;
        let y = Abstract1.meet_tcons_array manager x ar in
        apply_cl_r y rest

    and apply_cl_r x = function
        | CLTrue -> x
        | CLFalse -> vbottom x
        | CLAnd(a, b) ->
          let y = apply_cl_r x a in
          apply_cl_r y b
        | CLOr(a, b) ->
          let y = apply_cl x a in
          let z = apply_cl x b in
          Abstract1.join manager y z

    let apply_f x f = apply_cl x (conslist_of_f f)
    
    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_all fmt 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_all fmt xx

end