open Ast
open Formula
open Util

exception Bot

module type ENUM_ENVIRONMENT_DOMAIN = sig
    type t

    type item = string

    (* construction *)
    val top           : (id * item list) list -> t

    (* variable management *)
    val vars          : t -> (id * item list) list

    val forgetvar     : t -> id -> t
    val project       : t -> id -> item list

    (* set-theoretic operations *)
    val join          : t -> t -> t   (* union *)
    val meet          : t -> t -> t   (* intersection, may raise Bot *)

    val subset        : t -> t -> bool
    val eq            : t -> t -> bool

    (* abstract operations *)
    val apply_cons    : t -> enum_cons -> t list (* may disjunct *)
    val assign        : t -> (id * id) list -> t

    (* pretty-printing *)
    val print         : Format.formatter -> t -> unit
end



module Valuation : ENUM_ENVIRONMENT_DOMAIN = struct
    type item = string

    type t = {
        vars      : (id * item list) list;
        value     : item VarMap.t;
    }

    let top vars = { vars; value = VarMap.empty }

    let vars x = x.vars

    let forgetvar x id =
        { x with value = VarMap.remove id x.value }
    let project x id =
        try [VarMap.find id x.value]
        with Not_found -> List.assoc id x.vars

    let join x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, Some b when a = b -> Some a
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }
    let meet x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, None | None, Some a -> Some a
            | Some a, Some b -> if a = b then Some a else raise Bot
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }

    let subset a b =
        VarMap.for_all
          (fun id v ->
              try v = VarMap.find id b.value
              with Not_found -> true)
          a.value
    let eq a b =
        VarMap.equal (=) a.value b.value

    let cc op = match op with
        | E_EQ -> (=)
        | E_NE -> (<>)

    let apply_cons x (op, id, e) =
        match e with
        | EItem s ->
          begin try let t = VarMap.find id x.value in
              if cc op s t then [x] else []
          with Not_found ->
            List.map (fun v -> { x with value = VarMap.add id v x.value })
              (List.filter (cc op s) (List.assoc id x.vars))
          end
        | EIdent id2 ->
          match
            (try Some (VarMap.find id x.value) with _  -> None),
            (try Some (VarMap.find id2 x.value) with _  -> None)
          with
          | None, None -> [x]
          | None, Some s ->
            List.map (fun v -> { x with value = VarMap.add id v x.value})
              (List.filter (cc op s) (List.assoc id x.vars))
          | Some s, None ->
            List.map (fun v -> { x with value = VarMap.add id2 v x.value})
              (List.filter (cc op s) (List.assoc id2 x.vars))
          | Some s, Some t ->
            if cc op s t then [x] else []
    
    let assign x idl =
        let v = List.fold_left
            (fun v (id, id2) ->
                try VarMap.add id (VarMap.find id2 x.value) v
                with Not_found -> VarMap.remove id v )
            x.value idl
        in { x with value = v }


    let print fmt x =
        let b = List.map (fun (x, y) -> y, x) (VarMap.bindings x.value) in
        let s = List.sort Pervasives.compare b in
        let rec bl = function
          | [] -> []
          | (v, id)::q ->
            begin match bl q with
            | (vv, ids)::q when vv = v -> (v, id::ids)::q
            | r -> (v, [id])::r
            end
        in
        let sbl = bl s in
        Format.fprintf fmt "@[<v 2>{ ";
        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;
            Format.fprintf fmt " ≡ %s@]" v)
          sbl;
        Format.fprintf fmt " }@]"

end