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
    type itv = Value_domain.itv

    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

    let var_itv x id = match x with
        | Bot -> Value_domain.Bot
        | Env env -> V.to_itv (get_var env id)

    (* Set-theoretic operations *)
    let f_in_merge f _ a b = match a, b with
      | Some a, None -> Some a
      | None, Some b -> Some b
      | Some a, Some b -> Some (f a b)
      | _ -> assert false

    let join a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.merge (f_in_merge V.join) m n)
    
    let meet a b = match a, b with
        | Bot, _ | _, Bot -> Bot
        | Env m, Env n ->
            strict (VarMap.merge (f_in_merge V.meet) m n)

    let widen a b = match a, b with
        | Bot, x | x, Bot -> x
        | Env m, Env n ->
            strict (VarMap.merge (f_in_merge V.widen) 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) = match rest with
        | CLTrue ->
          let apply_cl_l x = List.fold_left apply_cons x cons in
          fix eq apply_cl_l x
        | CLFalse -> vbottom x
        | CLAnd(a, b) ->
          let y = apply_cl x (cons, a) in
          apply_cl y (cons, b)
        | CLOr((ca, ra), (cb, rb)) ->
          let y = apply_cl x (cons@ca, ra) in
          let z = apply_cl x (cons@cb, rb) in
          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 "⊥"
        | Env env ->
            Format.fprintf fmt "@[<v 2>{ ";
            let l = List.map (fun id -> (get_var env id, id)) ids in
            let s = List.sort Pervasives.compare l in
            let rec bl = function
              | [] -> []
              | (v, id)::q when v <> V.top ->
                begin match bl q with
                | (vv, ids)::q when vv = v -> (v, id::ids)::q
                | r -> (v, [id])::r
                end
              | _::q -> bl q
            in
            let sbl = bl s in
            List.iteri
              (fun j (v, ids) ->
                if j > 0 then Format.fprintf fmt "@ ";
                Format.fprintf fmt "@[<hov 4>";
                List.iteri
                  (fun i id ->
                    if i > 0 then Format.fprintf fmt ",@ ";
                    Format.fprintf fmt "%a" Formula_printer.print_id id)
                  ids;
                match V.as_const v with
                | Some i -> Format.fprintf fmt " = %d@]" i
                | _ -> Format.fprintf fmt " ∊ %s@]" (V.to_string v))
              sbl;
            Format.fprintf fmt " }@]"

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

    let print_itv fmt x =
        Format.fprintf fmt "%s" (string_of_itv x)
    
end