open Ast
open Formula
open Num_domain
open Typing
open Util

type disjunct_policy =
  | DAlways
  | DNever
  | DSometimes

module type ENVIRONMENT_DOMAIN = sig
    type t
    type itv

    (* construction *)
    val init              : t
    val bottom            : t
    val vbottom           : t -> t      (* bottom value with same environment *)
    val is_bot            : t -> bool

    (* variable management *)
    val addvar            : t -> id -> Typing.typ -> t
    val rmvar             : t -> id -> t
    val vars              : t -> (id * Typing.typ) list

    val forgetvar         : t -> id -> t
    val var_itv           : t -> id -> itv
    val set_disjunct      : t -> id -> disjunct_policy -> t

    (* set_theoretic operations *)
    val join              : t -> t -> t   (* union *)
    val meet              : t -> t -> t   (* intersection *)
    val widen             : t -> t -> t

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

    (* abstract operations *)
    val apply_f           : t -> bool_expr -> t
    val apply_cl          : t -> conslist -> t
    val assign            : t -> (id * id) list -> t

    (* pretty-printing *)
    val print_vars  : Format.formatter -> t -> id list -> unit
    val print_all   : Format.formatter -> t -> unit
    val print_itv   : Format.formatter -> itv -> unit

end


module D_ignore_enum (ND : NUMERICAL_ENVIRONMENT_DOMAIN) : ENVIRONMENT_DOMAIN =
struct
    type t = ND.t
    type itv = ND.itv


    let init = ND.init
    let bottom = ND.bottom
    let vbottom = ND.vbottom
    let is_bot = ND.is_bot

    let addvar x id ty = ND.addvar x id (ty = TReal)
    let rmvar = ND.rmvar
    let vars x = List.map (fun (id, r) -> id, if r then TReal else TInt) (ND.vars x)

    let forgetvar = ND.forgetvar
    let var_itv = ND.var_itv
    let set_disjunct x _ _ = x

    let join = ND.join
    let meet = ND.meet
    let widen = ND.widen

    let subset = ND.subset
    let eq = ND.eq

    let rec apply_cl x (econs, cons, rest) = match rest with
        | CLTrue ->
          ND.apply_ncl x cons
        | CLFalse -> vbottom x
        | CLAnd(a, b) ->
          let y = apply_cl x (econs, cons, a) in
          apply_cl y (econs, cons, b)
        | CLOr((eca, ca, ra), (ecb, cb, rb)) ->
          let y = apply_cl x (econs@eca, cons@ca, ra) in
          let z = apply_cl x (econs@ecb, cons@cb, rb) in
          join y z

    let apply_f x f = apply_cl x (conslist_of_f f)

    let assign x v = ND.assign x (List.map (fun (id, id2) -> id, NIdent id2) v)

    let print_vars = ND.print_vars
    let print_all = ND.print_all
    let print_itv = ND.print_itv
end

module D (ND : NUMERICAL_ENVIRONMENT_DOMAIN) : ENVIRONMENT_DOMAIN = struct

    module Valuation = struct
      type t = string VarMap.t
      let compare = VarMap.compare Pervasives.compare

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

      let print fmt x =
          let b = List.map (fun (x, y) -> y, x) (VarMap.bindings x) 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
    module VMap = Mapext.Make(Valuation)

    type t = {
        nvars     : (id * bool) list;
        evars     : (id * string list) list;
        nd_init   : ND.t;
        nd_bottom : ND.t;
        value     : ND.t VMap.t;
        disj_v    : id list;
        nodisj_v  : id list;
    }
    type itv =
      | IN of ND.itv VMap.t
      | IE of SSet.t


    let init = {
        nvars = []; evars = []; disj_v = []; nodisj_v = [];
        nd_init = ND.init;
        nd_bottom = ND.bottom;
        value = VMap.singleton VarMap.empty ND.init;
      }
    let bottom = {
        nvars = []; evars = []; disj_v = []; nodisj_v = [];
        nd_init = ND.init;
        nd_bottom = ND.bottom;
        value = VMap.empty;
      }
    let vbottom x = { x with value = VMap.empty }
    
    let is_bot x = VMap.is_empty x.value || VMap.for_all (fun _ v -> ND.is_bot v) x.value

    let strict value =
        VMap.filter (fun _ v -> not (ND.is_bot v)) value

    let add_case x (enum, num) =
        let value = x.value in
        let enum = VarMap.filter (fun k v -> not (List.mem k x.nodisj_v)) enum in
        let v = 
          if VMap.exists (fun e0 n0 -> Valuation.subset enum e0 && ND.subset num n0) value || ND.is_bot num then
            value
          else
            try VMap.add enum (ND.join num (VMap.find enum value)) value
            with Not_found -> VMap.add enum num value
        in { x with value = v }

    let addvar x id ty = match ty with
        | TInt | TReal ->
          let is_real = (ty = TReal) in
          { x with
            nvars     = (id, is_real)::(x.nvars);
            nd_init   = ND.addvar x.nd_init id is_real;
            nd_bottom = ND.addvar x.nd_bottom id is_real;
            value     = VMap.map (fun v -> ND.addvar v id is_real) x.value
          }
        | TEnum options ->
          { x with
            evars     = (id, options)::x.evars
          }

    let vars x =
        List.map (fun (id, r) -> id, if r then TReal else TInt) x.nvars @
        List.map (fun (id, o) -> id, TEnum o) x.evars

    let forgetvar x id =
        if List.mem_assoc id x.evars then
          VMap.fold
              (fun enum num value ->
                let enum' = VarMap.remove id enum in
                add_case value (enum', num))
              x.value
              { x with value = VMap.empty }
        else
          { x with value = strict @@ VMap.map (fun v -> ND.forgetvar v id) x.value }

    let rmvar x id =
        if List.mem_assoc id x.evars then
          let y = forgetvar x id in
          { y with
            evars = List.remove_assoc id x.evars }
        else
          { x with
            nvars = List.remove_assoc id x.nvars;
            value = strict @@ VMap.map (fun v -> ND.rmvar v id) x.value }

    let var_itv x id =
        if List.mem_assoc id x.nvars
          then IN (VMap.map (fun v -> ND.var_itv v id) x.value)
          else not_implemented "var_itv for enum variable"

    let set_disjunct x id p =
        { x with
          disj_v = (if p = DAlways then id::x.disj_v
                      else List.filter ((<>) id) x.disj_v);
          nodisj_v = (if p = DNever then id::x.nodisj_v
                      else List.filter ((<>) id) x.nodisj_v);
        }

    let join x y =
        VMap.fold
          (fun enum num value -> add_case value (enum, num))
          x.value
          y
    let meet x y = not_implemented "meet for enum+num domain"
    let widen x y =
        { x with
          value = VMap.merge
            (fun _ a b -> match a, b with
              | None, Some a -> Some a
              | Some a, None -> Some a
              | Some a, Some b -> Some (ND.widen a b)
              | _ -> assert false)
            x.value y.value }

    (* Some cases of subset/equality not detected *)
    let subset a b =
        VMap.for_all
          (fun enum num ->
              VMap.exists
                (fun enum_b num_b ->
                  Valuation.subset enum enum_b && ND.subset num num_b)
                b.value)
          a.value
    let eq a b = subset a b && subset b a

    (* Constraints ! *)
    let apply_ec x (op, a, b) =
        VMap.fold
            (fun enum num value ->
              match a, b with
              | EItem x, EItem y ->
                  if (x = y && op = E_EQ) || (x <> y && op = E_NE)
                    then add_case value (enum, num)
                    else value
              | EItem u, EIdent i | EIdent i, EItem u ->
                  begin try let y = VarMap.find i enum in
                    if (u = y && op = E_EQ) || (u <> y && op = E_NE)
                      then add_case value (enum, num)
                      else value
                  with Not_found ->
                    if op = E_EQ
                      then add_case value (VarMap.add i u enum, num)
                      else (* MAKE A DISJUNCTION *)
                        List.fold_left
                          (fun value item ->
                              add_case value (VarMap.add i item enum, num))
                          value
                          (List.filter ((<>) u) (List.assoc i x.evars))
                  end
              | EIdent i, EIdent j ->
                  begin
                    try let x = VarMap.find i enum in
                      try let y = VarMap.find j enum in
                        if (x = y && op = E_EQ) || (x <> y && op = E_NE)
                          then add_case value (enum, num)
                          else value
                      with Not_found (* J not found *) ->
                        add_case value (VarMap.add j x enum, num)
                    with Not_found (* I not found *) ->
                      try let y = VarMap.find j enum in
                        add_case value (VarMap.add i y enum, num)
                      with Not_found (* J not found *) ->
                        if op = E_EQ
                        then (* MAKE A DISJUNCTION ! *)
                            List.fold_left
                              (fun value item ->
                                let enum = VarMap.add i item enum in
                                let enum = VarMap.add j item enum in
                                add_case value (enum, num))
                              value
                              (List.assoc i x.evars)
                        else
                          add_case value (enum, num)
                  end)
            x.value
            { x with value = VMap.empty }

    let apply_ecl x cons =
        let f x = List.fold_left apply_ec x cons in
        fix eq f x
            

    let apply_ncl x econs =
        { x with value = strict @@ VMap.map (fun x -> ND.apply_ncl x econs) x.value }

    let rec apply_cl x (econs, ncons, rest) = match rest with
        | CLTrue ->
          let y = apply_ecl x econs in
          let z = apply_ncl y ncons in
          z
        | CLFalse -> vbottom x
        | CLAnd(a, b) ->
          let y = apply_cl x (econs, ncons, a) in
          apply_cl y (econs, ncons, b)
        | CLOr((eca, ca, ra), (ecb, cb, rb)) ->
          let y = apply_cl x (econs@eca, ncons@ca, ra) in
          let z = apply_cl x (econs@ecb, ncons@cb, rb) in
          join y z

    let apply_f x f = apply_cl x (conslist_of_f f)

    let assign x e =
      let ee, ne = List.partition (fun (id, id2) -> List.mem_assoc id x.evars) e in
      let ne_e = List.map (fun (id, id2) -> id, NIdent id2) ne in

      VMap.fold
          (fun enum0 num value ->
              let enum = List.fold_left
                  (fun enum (id, id2) ->
                      try let y = VarMap.find id2 enum0 in
                        VarMap.add id y enum
                      with Not_found ->
                        VarMap.remove id enum)
                  enum0 ee
              in
              add_case value (enum, ND.assign num ne_e))
          x.value
          { x with value = VMap.empty }

    let print_all fmt x =
        Format.fprintf fmt "@[<v 0>";
        VMap.iter
          (fun enum num ->
            Format.fprintf fmt "@[<hv 2>{ %a ∧@ %a }@]@ "
              Valuation.print enum ND.print_all num)
          x.value;
        Format.fprintf fmt "@]"

    let print_vars fmt x idl = print_all fmt x (* TODO *)

    let print_itv fmt = function
        | IN x ->
          Format.fprintf fmt "@[<v 0>";
          VMap.iter
            (fun enum i ->
              Format.fprintf fmt "%a when %a@ "
                ND.print_itv i Valuation.print enum)
            x;
          Format.fprintf fmt "@]"
        | IE x ->
          Format.fprintf fmt "@[<hov 2>{ ";
          SSet.iter (fun x -> Format.fprintf fmt "%s@ " x) x;
          Format.fprintf fmt "}@]"

end