open Util
open Abstract_syntax_tree

module CD : Domain.S =
  struct

    exception DivideByZero
    exception TypeError
    exception Bot
    exception NotImplemented

    type tv =   (* type for a variable *)
      | Top
      | I of Z.t

    type ts =   (* type for an environment *)
      | Bot
      | Nb of tv VarMap.t

    let top_ts = Nb VarMap.empty

    let ts_union a b =
      begin match a, b with
      | Bot, Bot -> Bot
      | Nb a, Bot -> Nb a
      | Bot, Nb b -> Nb b
      | Nb a, Nb b ->
        Nb
          (VarMap.fold
            (fun var value a ->
              try match VarMap.find var a with
                | Top -> a
                | I x ->
                  match value with
                  | I y when y = x -> a
                  | _ -> VarMap.add var Top a
              with
              | Not_found -> VarMap.add var value a)
          b a)
      end

    (* must not raise exception *)
    let ts_add_bool_constraint expr (s : ts) =
      s (* TODO, not necessary... *)


    (* only evaluates numerical statements, raises
      TypeError when result is bool *)
    let rec eval_abs expr (s : tv VarMap.t) =
      match expr with
      | AST_unary(op, (i, _)) ->
        begin match eval_abs i s with
        | Top -> Top
        | I x ->
          match op with
          | AST_UNARY_PLUS -> I x
          | AST_UNARY_MINUS -> I (Z.neg x)
          | _ -> raise TypeError
        end
      | AST_binary(op, (a, _), (b, _)) ->
        begin match eval_abs a s, eval_abs b s with
        | I x, I  y ->
          begin match op with
          | AST_PLUS -> I (Z.add x y)
          | AST_MINUS -> I (Z.sub x y)
          | AST_MULTIPLY -> I (Z.mul x y)
          | AST_DIVIDE ->
            if y = Z.zero then raise DivideByZero;
            I (Z.div x y)
          | AST_MODULO ->
            if y = Z.zero then raise DivideByZero;
            I (Z.rem x y)
          | _ -> raise TypeError
          end
        | _ -> Top
        end
      | AST_identifier(id, _) ->
        begin
          try VarMap.find id s
          with _ -> Top
        end 
      | AST_int_const(s, _) ->
        I (Z.of_string s)
      | AST_bool_const(_) -> raise TypeError
      | AST_int_rand _ -> Top
      | _ -> raise NotImplemented (* extensions *)

    (* returns true if the expression can evaluate to true
      returns false if the expression always evaluates to false *)
    let rec eval_bool_abs expr (s : tv VarMap.t) =
      true (* TODO *)

    (* must not raise exception *)
    let rec interp_abs stat s =
      begin match s with
      | Bot -> Bot
      | Nb vars ->
        begin match stat with
        | AST_block b ->
          List.fold_left
            (fun ss (st, _) -> interp_abs st ss)
            s b
        | AST_assign ((id, _), (exp, _)) ->
          begin
            try
              let value = eval_abs exp vars in
              Nb (VarMap.add id value vars)
            with _ -> Bot
          end
        | AST_if ((cond, _), (tb, _), None) ->
          ts_union (Nb vars)
            (interp_abs tb
              (ts_add_bool_constraint cond (Nb vars)))
        | AST_if ((cond, l), (tb, _), Some (eb, _)) ->
          ts_union
            (interp_abs tb
              (ts_add_bool_constraint cond (Nb vars)))
            (interp_abs eb
              (ts_add_bool_constraint 
                  (AST_unary (AST_NOT, (cond, l)))
                (Nb vars)))
        | AST_while ((cond, _), (body, _)) ->
          let f s =
            ts_union
              s
              (ts_add_bool_constraint cond s)
          in
            fix f (Nb vars)
        | AST_HALT -> Bot
        | AST_assert (cond, _) ->
          if eval_bool_abs cond vars
            then Nb vars
            else begin
              print_string "Assert fail\n";
              Bot
            end
        | AST_print items ->
          List.iter
            (fun (var, _) ->
              print_string (var ^ " = ");
              try
                match VarMap.find var vars with
                | Top -> print_string "T\n"
                | I x -> print_string (Z.to_string x ^ "\n")
              with _ -> print_string "T\n")
            items;
          Nb vars
        | _ -> raise NotImplemented (* extensions *)
        end
      end

  end