open Ast
open Formula

open Num_domain
open Enum_domain

open Varenv

module type ENVIRONMENT_DOMAIN = sig
    type t
    type itv

    (* constructors *)
    val top           : varenv -> t
    val bottom        : varenv -> t
    val is_top        : t -> bool
    val is_bot        : t -> bool

    (* variable management *)
    val varenv        : t -> varenv

    val forgetvar     : t -> id -> t
    val forgetvars    : t -> id list -> t
    val nproject      : t -> id -> itv    (* numerical project *)
    val eproject      : t -> id -> item list

    (* set-theoretic operations *)
    val join          : t -> t -> t
    val meet          : t -> t -> t
    val widen         : t -> t -> t

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

    (* abstract operations *)
    val apply_ec      : t -> enum_cons -> t
    val apply_ecl     : t -> enum_cons list -> t
    val apply_nc      : t -> num_cons -> t
    val apply_ncl     : t -> num_cons list -> t
    val eassign       : t -> (id * id) list -> t
    val nassign       : t -> (id * num_expr) list -> t

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

end

module NumDomainOnly(ND : NUMERICAL_ENVIRONMENT_DOMAIN) : ENVIRONMENT_DOMAIN =
struct

    type t = varenv * ND.t
    type itv = ND.itv

    let top ve = (ve, ND.top ve.nvars)
    let bottom ve = (ve, ND.top ve.nvars)
    let is_top (_, x) = ND.is_top x
    let is_bot (_, x) = ND.is_bot x

    let varenv (ve, x) = ve

    let forgetvar (ve, x) id =
      if List.exists (fun (i, _) -> i = id) ve.nvars
        then (ve, ND.forgetvar x id)
        else (ve, x)
    let forgetvars (ve, x) ids =
      (ve, List.fold_left ND.forgetvar x
        (List.filter (fun id -> List.exists (fun (i, _) -> i = id) ve.nvars) ids))
    let nproject (_, x) id = ND.project x id
    let eproject (ve, _) id = List.assoc id ve.evars

    let join (ve, a) (_, b) = (ve, ND.join a b)
    let meet (ve, a) (_, b) = (ve, ND.meet a b)
    let widen (ve, a) (_, b) = (ve, ND.widen a b)

    let subset (_, a) (_, b) = ND.subset a b
    let eq (_, a) (_, b) = ND.eq a b

    let apply_ec x _ = x
    let apply_ecl x _ = x
    let apply_nc (ve, x) c = (ve, ND.apply_cons x c)
    let apply_ncl (ve, x) cl = (ve, ND.apply_cl x cl)
    let eassign x _ = x
    let nassign (ve, x) ex = (ve, ND.assign x ex)

    let print fmt (_, x) = ND.print fmt x
    let print_vars fmt (ve, x) ids =
      ND.print_vars fmt x
        (List.filter
          (fun id -> List.exists (fun (i, _) -> i = id) ve.nvars) ids)
    let print_itv = ND.print_itv

end

module NumEnumDomain
  (ED : ENUM_ENVIRONMENT_DOMAIN)(ND : NUMERICAL_ENVIRONMENT_DOMAIN)
  : ENVIRONMENT_DOMAIN
= struct

  type t =
    | V of varenv * ED.t * ND.t
    | Bot of varenv
  type itv = ND.itv

  let strict = function
    | V(ve, enum, num) as x ->
      if ND.is_bot num || ED.is_bot enum
        then Bot ve
        else x
    | b -> b
  let strict_apply f x = match x with
    | V(ve, enum, num) ->
      begin
        try strict (f (ve, enum, num))
        with Enum_domain.Bot -> Bot ve
      end
    | Bot ve -> Bot ve

  let top ve = V(ve, ED.top ve.evars, ND.top ve.nvars)
  let bottom ve = Bot ve
  let is_top x = match strict x with
    | Bot _ -> false
    | V(_, enum, num) -> ED.is_top enum && ND.is_top num
  let is_bot x = match strict x with
    | Bot _ -> true
    | _ -> false

  let varenv = function
    | Bot x -> x
    | V(ve, _, _) -> ve

  let forgetvar x id =
    strict_apply (fun (ve, enum, num) ->
        if List.exists (fun (id2, _) -> id2 = id) ve.evars
          then V(ve, ED.forgetvar enum id, num)
          else V(ve, enum, ND.forgetvar num id))
      x

  let forgetvars x ids = 
    strict_apply (fun (ve, enum, num) ->
        let efv, nfv = List.fold_left
          (fun (efv, nfv) id ->
            if List.exists (fun (id2, _) -> id2 = id) ve.evars
              then id::efv, nfv
              else efv, id::nfv)
          ([], []) ids
        in
          V(ve, ED.forgetvars enum efv, List.fold_left ND.forgetvar num nfv))
      x

  let nproject x id = match x with
    | V(ve, enum, num) ->
        ND.project num id
    | Bot ve -> ND.project (ND.bottom ve.nvars) id
  let eproject x id = match x with
    | V(ve, enum, num) ->
        ED.project enum id
    | Bot ve -> []

  let join x y = match strict x, strict y with
    | Bot _, a | a, Bot _ -> a
    | V(ve, enum, num), V(_, enum', num') ->
      V(ve, ED.join enum enum', ND.join num num')

  let meet x y = match strict x, strict y with
    | Bot ve, a | a, Bot ve -> Bot ve
    | V(ve, enum, num), V(_, enum', num') ->
      try
        strict (V(ve, ED.meet enum enum', ND.meet num num'))
      with Enum_domain.Bot -> Bot ve

  let widen x y = match strict x, strict y with
    | Bot _, a | a, Bot _ -> a
    | V(ve, enum, num), V(_, enum', num') ->
      V(ve, ED.join enum enum', ND.widen num num')

  let subset x y = match strict x, strict y with
    | Bot _, a -> true
    | V(_, _, _), Bot _ -> false
    | V(_, enum, num), V(_, enum', num') ->
      ED.subset enum enum' && ND.subset num num'

  let eq x y = match strict x, strict y with
    | Bot _, Bot _ -> true
    | V(_, enum, num), V(_, enum', num') ->
      ED.eq enum enum' && ND.eq num num'
    | _ -> false

  let apply_ec x f =
    strict_apply
      (fun (ve, enum, num) -> V(ve, ED.apply_cl enum [f], num))
      x
  let apply_ecl x f =
    strict_apply (fun (ve, enum, num) -> V(ve, ED.apply_cl enum f, num)) x
  let apply_nc x f =
    strict_apply (fun (ve, enum, num) -> V(ve, enum, ND.apply_cons num f)) x
  let apply_ncl x f =
    strict_apply (fun (ve, enum, num) -> V(ve, enum, ND.apply_cl num f)) x
  let eassign x ass =
    strict_apply (fun  (ve, enum, num) -> V(ve, ED.assign enum ass, num)) x
  let nassign x ass =
    strict_apply (fun  (ve, enum, num) -> V(ve, enum, ND.assign num ass)) x

  let print_itv = ND.print_itv
  let print fmt x = match strict x with
    | Bot _ -> Format.fprintf fmt "⊥"
    | V(_, enum, num) ->
      Format.fprintf fmt "@[<hov 1>(%a,@ %a)@]" ED.print enum ND.print num
  let print_vars fmt x vars = match strict x with
    | Bot _ -> Format.fprintf fmt "⊥"
    | V(ve, enum, num) ->
      let nvars = List.filter
        (fun v -> List.exists (fun (x, _) -> x = v) ve.nvars)
        vars
      in
      Format.fprintf fmt "@[<hov 1>(%a,@ " ED.print enum;
      ND.print_vars fmt num nvars;
      Format.printf ")@]" 

end