open Formula
open Util
open Ast

open Value_domain
open Environment_domain

let debug = false

(* Restricted domain, only support integers *)

module D (V : VALUE_DOMAIN) : ENVIRONMENT_DOMAIN = struct
    type env = V.t VarMap.t

    type t = Env of env | Bot

    let init = Env VarMap.empty
    let bottom = Bot

    let get_var env id =
        try VarMap.find id env
        with Not_found -> V.top
    
    (* utilities *)

    let rec eval env = function
        | NIdent id -> get_var env id
        | NIntConst i -> V.const i
        | NRealConst f -> V.const (int_of_float f) (* TODO floats *)
        | NUnary (AST_UPLUS, e) -> eval env e
        | NUnary (AST_UMINUS, e) -> V.neg (eval env e)
        | NBinary (op, e1, e2) ->
            let f = match op with
            | AST_PLUS -> V.add
            | AST_MINUS -> V.sub
            | AST_MUL -> V.mul
            | AST_DIV -> V.div
            | AST_MOD -> V.rem
            in f (eval env e1) (eval env e2)

    let strict env =                    (*      env -> t        *)
        if VarMap.exists (fun _ x -> x = V.bottom) env
            then Bot
            else Env env
    let strict_apply f = function       (* (env -> t) -> t -> t *)
        | Bot -> Bot
        | Env e -> match f e with
            | Bot -> Bot
            | Env e -> strict e

    (* implementation *)

    let is_bot env = match env with
        | Bot -> true
        | Env env -> strict env = Bot

    let addvar x id _ = strict_apply (fun y -> Env (VarMap.add id V.top y)) x
    let forgetvar x id = strict_apply (fun y -> Env (VarMap.add id V.top y)) x
    let rmvar x id = strict_apply (fun y -> Env (VarMap.remove id y)) x

    let vars env = match env with
        | Bot -> []
        | Env env -> List.map (fun (k, _) -> (k, false)) (VarMap.bindings env)
    let vbottom _ = bottom

    (* Set-theoretic operations *)
    let join a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.join a b) m n)
    
    let meet a b = match a, b with
        | Bot, _ | _, Bot -> Bot
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.meet a b) m n)

    let widen a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.map2z (fun _ a b -> V.widen a b) m n)

    (* Inclusion and equality *)
    let subset a b = match a, b with
        | Bot, x -> true
        | Env _, Bot -> false
        | Env m, Env n ->
            VarMap.for_all2o
                (fun _ _ -> true)
                (fun _ v -> v = V.top)
                (fun _ a b -> V.subset a b)
                m n

    let eq a b = match a, b with
        | Bot, Bot -> true
        | Env m, Env n ->
            VarMap.for_all2o
                (fun _ _ -> false)
                (fun _ _ -> false)
                (fun _ a b -> a = b)
                m n
        | _ -> false

    (* Apply some formula to the environment *)
    let apply_cons env (expr, sign) =
        let inv_op = function
          | AST_LT -> AST_GT
          | AST_GT -> AST_LT
          | AST_LE -> AST_GE
          | AST_GE -> AST_LE
          | x -> x
        in
        let rec extract_var (lhs, op, rhs) =
          match lhs with
          | NIdent i -> [i, op, rhs]
          | NIntConst _ | NRealConst _ -> []
          | NUnary(AST_UPLUS, x) -> extract_var (x, op, rhs)
          | NUnary(AST_UMINUS, x) ->
            extract_var (x, inv_op op, NUnary(AST_UMINUS, x))
          | NBinary(AST_PLUS, a, b) ->
            extract_var (a, op, NBinary(AST_MINUS, rhs, b)) @
            extract_var (b, op, NBinary(AST_MINUS, rhs, a))
          | NBinary(AST_MINUS, a, b) ->
            extract_var (a, op, NBinary(AST_PLUS, rhs, b)) @
            extract_var (b, inv_op op, NBinary(AST_MINUS, a, rhs))
          | NBinary(AST_MUL, a, b) ->
            extract_var (a, op, NBinary(AST_DIV, rhs, b)) @
            extract_var (b, op, NBinary(AST_DIV, rhs, a))
          | NBinary(AST_DIV, a, b) ->
            extract_var (a, op, NBinary(AST_MUL, rhs, b)) @
            extract_var (b, inv_op op, NBinary(AST_DIV, a, rhs))
          | NBinary(AST_MOD, _, _) -> []
        in
        let zop = match sign with
          | CONS_EQ -> AST_EQ | CONS_NE -> AST_NE
          | CONS_GT -> AST_GT | CONS_GE -> AST_GE
        in
        let restrict_var env (i, op, expr) =
          strict_apply (fun env ->
            let v1, v2 = get_var env i, eval env expr in
            let v1' = match op with
            | AST_EQ -> V.meet v1 v2
            | AST_NE -> v1
            | AST_LE -> let u, _ = V.leq v1 v2 in u
            | AST_GE -> let _, v = V.leq v2 v1 in v
            | AST_LT -> let u, _ = V.leq v1 (V.sub v2 (V.const 1)) in u
            | AST_GT -> let _, v = V.leq (V.add v2 (V.const 1)) v1 in v
            in
            if debug then Format.printf
              "restrict %s %s @[<hv>%a@] : %s %s -> %s@." i 
              (Formula_printer.string_of_binary_rel op)
              Formula_printer.print_num_expr expr
              (V.to_string v1) (V.to_string v2) (V.to_string v1');
            Env (VarMap.add i v1' env))
          env
        in
        List.fold_left restrict_var env
          (extract_var (expr, zop, NIntConst 0))

    let rec apply_cl x (cons, rest) =
        let apply_cl_l x = List.fold_left apply_cons x cons in

        let y = fix eq apply_cl_l x in
        if debug then Format.printf "-[@.";
        let z = apply_cl_r y rest in
        if debug then Format.printf "-]@.";
        fix eq apply_cl_l z

    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) ->
          if debug then Format.printf "<<@.";
          let y = apply_cl x a in
          if debug then Format.printf "\\/@.";
          let z = apply_cl x b in
          if debug then Format.printf ">>@.";
          join y z

    let apply_f x f = apply_cl x (conslist_of_f f)

    (* Assignment *)
    let assign x exprs =
        let aux env =
          let vars = List.map (fun (id, v) -> (id, eval env v)) exprs in
          let env2 =
            List.fold_left (fun e (id, v) -> VarMap.add id v e)
              env vars
          in Env env2
        in
        strict_apply aux x

        
    (* pretty-printing *)
    let print_vars fmt env ids =
        match env with
        | Bot -> Format.fprintf fmt "bottom"
        | Env env ->
            Format.fprintf fmt "@[<hov 2>[";
            let _ = List.fold_left
                (fun nf id ->
                  let v = get_var env id in
                  if v <> V.top then begin
                    if nf then Format.fprintf fmt ",@ ";
                    Format.fprintf fmt "%s in %s" id (V.to_string (get_var env id));
                    true
                  end else nf)
                false
                ids
            in Format.fprintf fmt "]@]"

    let print_all fmt x =
        print_vars fmt x (List.map fst (vars x))
    
end