open Ast
open Ast_util
open Formula
open Typing

open Util
open Num_domain
open Enum_domain

type cmdline_opt = {
    widen_delay       : int;
    disjunct          : (id -> bool);
    verbose_ci        : bool;
}


module I (ED : ENUM_ENVIRONMENT_DOMAIN) (ND : NUMERICAL_ENVIRONMENT_DOMAIN) : sig

    type st

    val do_prog : cmdline_opt -> rooted_prog -> unit

end = struct

    (*  **************************
            Disjunction domain 
        **************************  *)

    type case_v = ED.t * ND.t
    type case = ED.item list

    type dd_v = {
        d_vars  : id list;
        data    : (case, case_v) Hashtbl.t;
    }

    (*
      print_dd : Formatter.t -> dd_v -> unit
    *)
    let print_aux fmt d_vars env =
      let print_it q (enum, num) =
        Format.fprintf fmt "@[<hov 4>";
        List.iter2
          (fun id  v -> Format.fprintf fmt "%a ≡ %s,@ "
              Formula_printer.print_id id v)
          d_vars q;
        Format.fprintf fmt "@ %a,@ %a@]@."
            ED.print enum ND.print num;
      in
      Hashtbl.iter print_it env
    let print_dd fmt dd = print_aux fmt dd.d_vars dd.data


    (*
      dd_bot : id list -> dd_v
    *)
    let dd_bot d_vars = { d_vars; data = Hashtbl.create 12 }

    (*
      dd_add_case : dd_v -> case_v -> unit
    *)
    let dd_add_case env (enum, num) =
      let rec aux enum vals = function
        | [] ->
          let case = List.rev vals in
          begin try let e0, n0 = Hashtbl.find env.data case in
            Hashtbl.replace env.data case (ED.join enum e0, ND.join num n0);
          with Not_found ->
            Hashtbl.add env.data case (enum, num);
          end
        | p::q ->
          List.iter
            (fun v -> match ED.apply_cons enum (E_EQ, p, EItem v) with
              | [en] -> aux en (v::vals) q
              | _ -> assert false)
            (ED.project enum p)
      in aux enum [] env.d_vars

    (*
      dd_singleton : id list -> case_v -> dd_v
    *)
    let dd_singleton d_vars d =
      let v = { d_vars; data = Hashtbl.create 1 } in
      dd_add_case v d;
      v

    (*
      dd_extract_case : dd_v -> case -> dd_v
    *)
    let dd_extract_case v case =
      try dd_singleton v.d_vars (Hashtbl.find v.data case)
      with Not_found -> dd_bot v.d_vars


    (*
      dd_apply_ncons : dd_v -> num_cons list -> dd_v
    *)
    let dd_apply_ncl v ncl =
        let vv = dd_bot v.d_vars in
        Hashtbl.iter (fun case (enum, num) ->
            let num = ND.apply_cl num ncl in
            if not (ND.is_bot num) then
              Hashtbl.add vv.data case (enum, num))
          v.data;
        vv

    (*
      dd_apply_econs : dd_v -> enum_cons -> dd_v
    *)
    let dd_apply_econs v cons =
        let vv = dd_bot v.d_vars in
        Hashtbl.iter (fun case (enum, num) ->
            List.iter
              (fun enum -> dd_add_case vv (enum, num))
              (ED.apply_cons enum cons))
          v.data;
        vv

    (*
      dd_join : dd_v -> dd_v -> dd_v
    *)
    let dd_join a b =
        let vv = { d_vars = a.d_vars; data = Hashtbl.copy a.data } in
        Hashtbl.iter (fun case (enum, num) ->
            try let e2, n2 = Hashtbl.find vv.data case in
              Hashtbl.replace vv.data case (ED.join e2 enum, ND.join n2 num)
            with Not_found -> Hashtbl.replace vv.data case (enum, num))
          b.data;
        vv

    (*
      dd_meet : dd_v -> dd_v -> dd_v
    *)
    let dd_meet a b =
        let vv = dd_bot a.d_vars in
        Hashtbl.iter (fun case (enum, num) ->
            try let e2, n2 = Hashtbl.find b.data case in
              let en, nn = ED.meet enum e2, ND.meet n2 num in
              if not (ND.is_bot nn) then Hashtbl.add vv.data case (en, nn)
            with _ -> ())
          a.data;
        vv

    (*
      dd_apply_cl : dd_v -> conslist -> dd_v
    *)
    let rec apply_cl env (ec, nc, r) =
        match r with
        | CLTrue ->
          let env_ec = List.fold_left dd_apply_econs env ec in
          let env_nc = dd_apply_ncl env_ec nc in
          (* Format.printf "[[ %a -> %a ]]@."
              print_dd env print_dd env_nc; *)
          env_nc
        | CLFalse -> dd_bot env.d_vars
        | CLAnd(a, b) ->
          let env = apply_cl env (ec, nc, a) in
          apply_cl env (ec, nc, b)
        | CLOr((eca, nca, ra), (ecb, ncb, rb)) ->
          dd_join (apply_cl env (ec@eca, nc@nca, ra))
                  (apply_cl env (ec@ecb, nc@ncb, rb))
      

    (*  ***************************
          Abstract interpretation 
        ***************************  *)

    type st = {
        rp          : rooted_prog;
        opt         : cmdline_opt;

        enum_vars   : (id * ED.item list) list;
        num_vars    : (id * bool) list;

        (* program expressions *)
        f           : bool_expr;
        cl          : conslist;
        f_g         : bool_expr;
        cl_g        : conslist;
        guarantees  : (id * bool_expr) list;

        (* abstract interpretation *)
        cycle       : (id * id * typ) list;     (* s'(x) = s(y) *)
        forget      : (id * typ) list;          (* s'(x) not specified *)
        d_vars      : id list;                  (* disjunction variables *)
        top         : case_v;
        env         : (case, case_v) Hashtbl.t;
        mutable delta : case list;
        widen       : (case, int) Hashtbl.t;
    }

    (*
      print_st : Formatter.t -> st -> unit
    *)
    let print_st fmt st = print_aux fmt st.d_vars st.env

    (*
      add_case : st -> case_v -> unit
    *)
    let add_case st (enum, num) =
      let rec aux enum vals = function
        | [] ->
          let case = List.rev vals in
          begin try let e0, n0 = Hashtbl.find st.env case in
            if (not (ED.subset enum e0)) || (not (ND.subset num n0)) then begin
              begin try
                let n = Hashtbl.find st.widen case in
                let jw = if n < st.opt.widen_delay then ND.join else ND.widen in
                Hashtbl.replace st.env case (ED.join e0 enum, jw n0 num);
                Hashtbl.replace st.widen case (n+1);
              with Not_found ->
                Hashtbl.replace st.env case (ED.join enum e0, ND.join num n0);
                Hashtbl.add st.widen case 0;
              end;
              if not (List.mem case st.delta) then st.delta <- case::st.delta
            end
          with Not_found ->
            Hashtbl.add st.env case (enum, num);
            if not (List.mem case st.delta) then st.delta <- case::st.delta
          end
        | p::q ->
          List.iter
            (fun v -> match ED.apply_cons enum (E_EQ, p, EItem v) with
              | [en] -> aux en (v::vals) q
              | _ -> assert false)
            (ED.project enum p)
      in aux enum [] st.d_vars



    (*
      init_state : int -> rooted_prog -> st
    *)
    let init_state opt rp =
      Format.printf "Vars: @[<hov>%a@]@.@."
          (print_list Ast_printer.print_typed_var ", ")
          rp.all_vars;

      let enum_vars = List.fold_left
          (fun v (_, id, t) -> match t with
              | TEnum ch -> (id, ch)::v | _ -> v)
          [] rp.all_vars in
      let num_vars = List.fold_left
          (fun v (_, id, t) -> match t with
              | TInt -> (id, false)::v | TReal -> (id, true)::v | _ -> v)
          [] rp.all_vars in

      let init_f = Transform.init_f_of_prog rp in
      Format.printf "Init formula: %a@.@." Formula_printer.print_expr init_f;
      let init_cl = conslist_of_f init_f in

      let guarantees = Transform.guarantees_of_prog rp in
      Format.printf "Guarantees:@.";
      List.iter (fun (id, f) ->
          Format.printf "  %s: %a@." id Formula_printer.print_expr f)
        guarantees;
      Format.printf "@.";

      let f = Formula.eliminate_not (Transform.f_of_prog rp false) in
      let f_g = Formula.eliminate_not (Transform.f_of_prog rp true) in
      Format.printf "Cycle formula:@.%a@.@." Formula_printer.print_expr f;

      let cl = Formula.conslist_of_f f in
      let cl_g = Formula.conslist_of_f f_g in
      Format.printf "Cycle conslist:@.%a@.@." Formula_printer.print_conslist cl;

      (* calculate cycle variables and forget variables *)
      let cycle = List.fold_left
        (fun q (_, id, ty) ->
            if id.[0] = 'N' then
              (String.sub id 1 (String.length id - 1), id, ty)::q
            else q)
        [] rp.all_vars
      in
      let forget = List.map (fun (_, id, ty) -> (id, ty))
          (List.filter
            (fun (_, id, _) ->
              not (List.exists (fun (_, id2, _) -> id2 = "N"^id) rp.all_vars))
            rp.all_vars)
      in

      (* calculate disjunction variables *)
      (* first approximation : use all enum variables appearing in init formula *)
      let rec calculate_dvars (ec, _, r) =
          let a = List.map (fun (_, id, _) -> id) ec in
          let rec aux = function
          | CLTrue | CLFalse -> []
          | CLAnd(u, v) ->
              aux u @ aux v
          | CLOr(u, v) ->
              calculate_dvars u @ calculate_dvars v
          in a @ aux r
      in
      let d_vars_0 = calculate_dvars init_cl in
      let d_vars_1 = List.filter
          (fun id -> opt.disjunct id && not (List.mem id d_vars_0) && id.[0] <> 'N')
          (List.map (fun (id, _) -> id) enum_vars) in
      let d_vars = d_vars_0 @ d_vars_1 in

      (* setup abstract data *)
      let top = ED.top enum_vars, ND.top num_vars in
      let delta = [] in
      let widen = Hashtbl.create 12 in

      (* calculate initial state and environment *)
      let init_s = dd_bot d_vars in

      let init_ec, _, _ = init_cl in
      let init_ed = List.fold_left
          (fun ed cons ->
            List.flatten (List.map (fun x -> ED.apply_cons x cons) ed))
          [ED.top enum_vars] (init_ec) in
      List.iter (fun ed -> dd_add_case init_s (ed, ND.top num_vars)) init_ed;
      let env_dd = apply_cl init_s init_cl in
      let env = env_dd.data in
      
      let st = 
      { rp; opt; enum_vars; num_vars;
        f; cl; f_g; cl_g; guarantees;
        cycle; forget; d_vars; top;
        env; delta; widen } in

      Hashtbl.iter (fun case _ -> st.delta <- case::st.delta) st.env;

      st


    (*
      pass_cycle : st -> case_v -> case_v
    *)
    let pass_cycle st (enum, num) =
        let assign_e, assign_n = List.fold_left
          (fun (ae, an) (a, b, t) -> match t with
            | TEnum _ -> (a, b)::ae, an
            | TInt | TReal -> ae, (a, NIdent b)::an)
          ([], []) st.cycle in

        let enum = ED.assign enum assign_e in
        let num = ND.assign num assign_n in

        List.fold_left
            (fun (enum, num) (var, t) -> match t with
              | TEnum _ -> ED.forgetvar enum var, num
              | TReal | TInt -> enum, ND.forgetvar num var)
            (enum, num) st.forget

    let dd_pass_cycle st (v : dd_v) =
      let vv = dd_bot v.d_vars in
      Hashtbl.iter (fun _ x -> dd_add_case vv (pass_cycle st x)) v.data;
      vv

    (*
      cycle : st -> cl -> dd_v -> dd_v
    *)
    let cycle st cl env =
      let env_f = apply_cl env cl in
      (*Format.printf "{{ %a ->@.%a }}@." print_dd env print_dd env_f;*)
      dd_pass_cycle st env_f


    (*
      set_target_case : st -> dd_v -> case -> dd_v
    *)
    let set_target_case st env case =
      List.fold_left2
        (fun env id v ->
          if List.exists (fun (id2, _) -> id2 = "N"^id) st.enum_vars then
            dd_apply_econs env (E_EQ, "N"^id, EItem v)
          else env)
        env st.d_vars case

    (*
      do_prog : rooted_prog -> unit
    *)
    let do_prog opt rp =
      let st = init_state opt rp in

      Format.printf "Init: %a@." print_st st;

      let n_ch_iter = ref 0 in

      (* chaotic iteration loop *)
      while st.delta <> [] do
        let case = List.hd st.delta in

        incr n_ch_iter;
        Format.printf "Chaotic iteration %d: @[<hov>[%a]@]@."
            !n_ch_iter
            (print_list Formula_printer.print_id ", ") case;
        
        let start = try Hashtbl.find st.env case with Not_found -> assert false in
        let start_dd = dd_singleton st.d_vars start in
        let f i =
          let i = dd_singleton st.d_vars i in
          let i' = set_target_case st i case in
          let q = dd_join start_dd (cycle st st.cl i') in
          try Hashtbl.find q.data case with Not_found -> assert false
        in

        let rec iter n (i_enum, i_num) =
          let fi_enum, fi_num = f (i_enum, i_num) in
          let j_enum, j_num =
            if n < st.opt.widen_delay then
              ED.join i_enum fi_enum, ND.join i_num fi_num
            else
              ED.join i_enum fi_enum, ND.widen i_num fi_num
          in
          if ED.eq j_enum i_enum && ND.eq j_num i_num
            then (i_enum, i_num)
            else iter (n+1) (j_enum, j_num)
        in
        let y = iter 0 start in
        let z = fix (fun (a, b) (c, d) -> ED.eq a c && ND.eq b d) f y in

        let i = dd_singleton st.d_vars z in
        let j = cycle st st.cl i in
        let cases = ref [] in
        Hashtbl.iter (fun case _ -> cases := case::(!cases)) j.data;
        List.iter
            (fun case ->
              let i' = set_target_case st i case in
              let j = cycle st st.cl i' in
              Hashtbl.iter (fun _ q -> add_case st q) j.data)
            !cases;

        st.delta <- List.filter ((<>) case) st.delta;

        if st.opt.verbose_ci then
          Format.printf "-> @[<hov>%a@]@." print_st st;

      done;

      let final = apply_cl { d_vars = st.d_vars; data = st.env } st.cl in
      Format.printf "Accessible: %a@." print_dd final;


      let check_guarantee (id, f) =
        let cl = Formula.conslist_of_f f in
        Format.printf "@[<hv 4>%s:@ %a ⇒ ⊥  @ "
          id Formula_printer.print_conslist cl;
        let z = apply_cl final cl in
        if Hashtbl.length z.data = 0 then
          Format.printf "OK@]@ "
        else
          Format.printf "FAIL@]@ "
      in
      if st.guarantees <> [] then begin
        Format.printf "Guarantee @[<v 0>";
        List.iter check_guarantee st.guarantees;
        Format.printf "@]@."
      end;

      if List.exists (fun (p, _, _) -> p) st.rp.all_vars then begin
        Format.printf "Probes: @[<v 0>";
        List.iter (fun (p, id, ty) ->
          if p then begin
            Format.printf "%a ∊ @[<v>" Formula_printer.print_id id;
            Hashtbl.iter (fun case (enum, num) ->
              begin match ty with
              | TInt | TReal ->
                Format.printf "%a" ND.print_itv (ND.project num id)
              | TEnum _ ->
                Format.printf "@[<hov>{ %a }@]"
                  (print_list Formula_printer.print_id ", ")
                  (ED.project enum id)
              end;
              Format.printf " when @[<hov>[%a]@]@ "
                (print_list Formula_printer.print_id ", ") case)
              final.data;
            Format.printf "@]@ "
          end)
          st.rp.all_vars;
        Format.printf "@]@."
      end

end