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


(*
    Simple domain : one valuation for each variable, or T
*)

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


(*
    Less simple domain : a set of possible values for each variable, or T
*)

module MultiValuation : ENUM_ENVIRONMENT_DOMAIN = struct
    type item = string

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

    let sort = List.sort Pervasives.compare
    let uniq x = uniq_sorted (sort x)

    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 -> Some (uniq (a@b))
            | _ -> 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 ->
              begin match List.filter (fun x -> List.mem x b) a with
              | [] -> raise Bot | l -> Some l end
            | _ -> None)
          x1.value x2.value in
        { x1 with value = v }

    let subset a b =
        VarMap.for_all
          (fun id v ->
              try let v2 = VarMap.find id b.value in
                  List.for_all (fun x -> List.mem x v2) v
              with Not_found -> true)
          a.value
    let eq a b = subset a b && subset b a

    let apply_cons x (op, id, e) =
        let op = match op with | E_EQ -> (=) | E_NE -> (<>) in

        match e with
        | EItem s ->
          let pv = project x id in
          begin match List.filter (op s) pv with
          | [] -> []
          | vals -> [{ x with value = VarMap.add id vals x.value }]
          end
        | EIdent id2 ->
          let v1 = project x id in
          let v2 = project x id2 in
          List.fold_left
            (fun l v1 ->
                let v2 = List.filter (op v1) v2 in
                let x = { x with value = VarMap.add id [v1] x.value } in
                match v2 with
                | [] -> l
                | _ -> { x with value = VarMap.add id2 v2 x.value }::l)
            [] v1
    
    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 ->
            let v = sort v in
            if v <> sort (List.assoc id x.vars) then
              match bl q with
              | (vv, ids)::q when vv = v -> (v, id::ids)::q
              | r -> (v, [id])::r
            else
              bl q
        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;
            match v with
            | [v0] -> Format.fprintf fmt " ≡ %s@]" v0
            | l -> Format.fprintf fmt " ∊ @[<hov>{ %a }@]@]" (print_list Format.pp_print_string ", ") l)
          sbl;
        Format.fprintf fmt " }@]"

end


(*
    Complicated domain : a set of values for each variables, plus some
    constraints on couples of variables
      (eg. (x, y) in [ tt, tt ; ff, ff ]

    WARNING : This domain is not proved to be safe. In particular, it may represent
    a set of contraints that imply Bot (ie that are impossible) without raising the
    Bot exception. (there is potentially exponential cost in the checking that all
    the constraints are coherent, and I have no idea of an algorithm for doing that
    check.)

    Therefore, do not use this domain unless you know what you are doing (which is
    probably not the case anyway).
*)

module MultiValuationCCons : ENUM_ENVIRONMENT_DOMAIN = struct

    module VarC = struct
        type t = id * id
        let compare = Pervasives.compare
    end
    module VarCMap = Mapext.Make(VarC)

    type item = string

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

        (* in ccond (x, y) -> l, must have x < y (textual order on identifiers) -> unicity *)
        ccons     : (item * item) list VarCMap.t;
    }

    let sort = List.sort Pervasives.compare
    let uniq x = uniq_sorted (sort x)
    let list_inter x y = List.filter (fun k -> List.mem k y) x

    let all_couples l1 l2 =
        List.flatten
          (List.map (fun x -> List.map (fun y -> x, y) l2) l1)

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

    let vars x = x.vars

    let forgetvar x id =
        (* TODO : try to find a substitution variable so that some contraints can be propagated.
                  this is important so that cycle passing can be done correctly ! *)
        { x with
          value = VarMap.remove id x.value;
          ccons = VarCMap.filter (fun (a, b) _ -> a <> id && b <> id) x.ccons }

    let project x id =
        try VarMap.find id x.value
        with Not_found -> List.assoc id x.vars

    let project2 x id1 id2 =
        try
          let id1', id2' = ord_couple (id1, id2) in
          if id1' = id1 then
            VarCMap.find (id1, id2) x.ccons
          else
            List.map (fun (a, b) -> b, a) (VarCMap.find (id1', id2') x.ccons)
        with _ ->
          let v1, v2 = project x id1, project x id2 in
          all_couples v1 v2

    let strict x =
        let rec f x =
          (*
              - if (x, y) in [ a, b1 ; a, b2 ; ... ; a, bn ],
                  replace this by x = a and y in { b1, ..., bn }
              - filter (x, y) in [ ... ] with x in specified itv, y in specified itv
          *)
          let usefull, vv, cc =
            VarCMap.fold
              (fun (x, y) l (usefull, vv, cc) ->
                  let p1, p2 = uniq (List.map fst l), uniq (List.map snd l) in
                  let p1 = try list_inter p1 (VarMap.find x vv) with _ -> p1 in
                  let p2 = try list_inter p2 (VarMap.find y vv) with _ -> p2 in

                  if p1 = [] || p2 = [] then raise Bot;
                  let vv = VarMap.add x p1 (VarMap.add y p2 vv) in

                  if List.length p1 = 1 || List.length p2 = 1
                  then
                    true, vv, cc
                  else
                    match List.filter (fun (u, v) -> List.mem u p1 && List.mem v p2) l with
                    | [] -> raise Bot
                    | l2 -> List.length l2 < List.length l, vv, VarCMap.add (x, y) l2 cc)
              x.ccons (false, x.value, VarCMap.empty) in

          let x = { x with value = vv; ccons = cc } in
          if usefull then f x else x
        in
        f x

    let join x1 x2 =
        let v = VarMap.merge
          (fun _ a b -> match a, b with
            | Some a, Some b -> Some (uniq (a@b))
            | _ -> None)
          x1.value x2.value in
        let x = { x1 with value = v } in

        let cc = VarCMap.merge
          (fun (id1, id2) l1 l2 ->
              let v1, v2 = project x1 id1, project x id2 in
              let ac = all_couples v1 v2 in
              let c = List.filter
                  (fun q ->
                    (match l1 with Some l -> List.mem q l | _ -> true) ||
                    (match l2 with Some l -> List.mem q l | _ -> true))
                  ac in 
              if List.length c < List.length ac then Some c else None)
          x1.ccons x2.ccons in

        strict { x with ccons = cc }

    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 ->
              begin match list_inter a b with
              | [] -> raise Bot | l -> Some l end
            | _ -> None)
          x1.value x2.value in
        let x = { x1 with value = v } in 

        let cc = VarCMap.merge
          (fun (id1, id2) l1 l2 ->
              let v1, v2 = project x id1, project x id2 in
              let ac = all_couples v1 v2 in
              let c1 = match l1 with Some l -> list_inter l ac | None -> ac in
              let c2 = match l2 with Some l -> list_inter l ac | None -> ac in
              match list_inter c1 c2 with
              | [] -> raise Bot
              | l -> if List.length l < List.length ac then Some l else None)
          x1.ccons x2.ccons in

        strict { x with ccons = cc }

    let subset a b =
        VarMap.for_all
          (fun id v ->
              let vs = project a id in
              List.for_all (fun c -> List.mem c vs) v)
          b.value
        &&
        VarCMap.for_all
          (fun (id1, id2) l ->
              let l2 = project2 a id1 id2 in
              List.for_all (fun c -> List.mem c l) l2)
          b.ccons
    let eq a b =
      subset a b && subset b a

    let apply_cons x (op, id, e) =
        let op = match op with | E_EQ -> (=) | E_NE -> (<>) in
        
        match e with
        | EItem s ->
          let pv = project x id in
          begin match List.filter (op s) pv with
          | [] -> []
          | vals -> try [strict { x with value = VarMap.add id vals x.value }] with Bot -> []
          end
        | EIdent id2 ->
          let id, id2 = ord_couple (id, id2) in
          let c = project2 x id id2 in
          let ok_c = List.filter (fun (a, b) -> op a b) c in
          try
            [match uniq (List.map fst ok_c), uniq (List.map snd ok_c) with
              | [], _ | _, [] -> raise Bot
              | [a], q -> strict { x with value = VarMap.add id [a] (VarMap.add id2 q x.value) }
              | q, [b] -> strict { x with value = VarMap.add id q (VarMap.add id2 [b] x.value) }
              | _ -> strict { x with ccons = VarCMap.add (id, id2) ok_c x.ccons}
            ]
          with Bot -> []
    
    let assign x idl =
        let x2 = List.fold_left (fun x (v, _) -> forgetvar x v) x idl in
        let v = List.fold_left
            (fun v (id, id2) ->
                try VarMap.add id (VarMap.find id2 x.value) v
                with Not_found -> v)
            x2.value idl in
        let c = VarCMap.fold
            (fun (v1, v2) l c ->
                let v1' = try List.assoc v1 idl with _ -> v1 in
                let v2' = try List.assoc v2 idl with _ -> v2 in
                VarCMap.add (v1', v2') l c)
            x.ccons x2.ccons in
        strict { x with value = v; ccons = c }


    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 ->
            let v = sort v in
            if v <> sort (List.assoc id x.vars) then
              match bl q with
              | (vv, ids)::q when vv = v -> (v, id::ids)::q
              | r -> (v, [id])::r
            else
              bl q
        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;
            match v with
            | [v0] -> Format.fprintf fmt " ≡ %s@]" v0
            | l -> Format.fprintf fmt " ∊ @[<hov>{ %a }@]@]" (print_list Format.pp_print_string ", ") l)
          sbl;
        Format.fprintf fmt " }@]"

end