open Misc;;
open Iset;;
open Econd;;
open Cout;;
open Syntax;;

let rec get_cond_vars accu = function
  | Ein (v, _) -> v::accu
  | Eand cs -> List.fold_left get_cond_vars accu cs
  | Eor cs -> List.fold_left get_cond_vars accu cs
  | Efalse | Etrue -> []
;;

let rec get_vars = function
  | [] -> []
  | (cond,_,_)::cs -> (get_cond_vars [] cond) @@ (get_vars cs)
;;

let rec find_cut cur = function
  | [] -> cur
  | s::ss
     -> let rec refine accu = function
          | [] -> accu
          | p::ps -> let sp = conj s p in
                     if sp = Pos [] then refine (p::accu) ps
                     else
                     let ex = except p s in
                     if ex = Pos [] then accu @ (p::ps)
                     else refine ((except p s)::sp::accu) ps
        in
        find_cut (refine [] cur) ss
;;

let rec cut_cost_aux accu ss cut =
  match ss with
  | [] -> accu
  | s::ss1 -> let inter x y = (conj x y <> Pos []) in
              cut_cost_aux (accu + list_count (inter s) cut) ss1 cut
;;

let cut_cost ss cut =
   let x = List.length cut - 1 in
   if x <= 0 then max_int
   else (cut_cost_aux (- List.length ss) ss cut * 1000) / x
;;


let max_cut cs (vv, set) =
  let vals = list_rmap (fun (a,b,c) -> possible_values vv a) cs in
  let cut = find_cut [set] vals in
  (cut_cost vals cut, vv, cut)
;;


(* Find the set of variables that occur in more than half the sets *)
let majority vcs =
  let threshold = (List.length vcs) / 2 in
  let (neg,pos) = List.partition Iset.infinite vcs in
  let accu = ref (if List.length neg > threshold
                  then List.fold_left Iset.conj (Neg []) neg
                  else Pos [])
  in
  let getvars = function Pos x -> x | Neg x -> assert false in
  let allvars = List.concat (list_rmap getvars pos) in
  let addvar v = if list_count (fun x -> Iset.element v x) vcs > threshold
                 then accu := Iset.disj (Pos [v]) !accu
  in
  List.iter addvar allvars;
  !accu
;;

(* Get the vcs that intersect set *)
let relevant set vcs =
  let is_good s = Iset.conj s set <> Pos [] in
  List.filter is_good vcs
;;

let rec maximize vcs v =
  let cur = relevant v vcs in
  let newv = majority cur in
  if Iset.included newv v then v
  else maximize vcs (Iset.disj v newv)
;;

let rec scc_cut_aux vcs accu set =
  let subvcs = list_rmap (Iset.conj set) vcs in
  let subvcs = List.filter (fun x -> x <> Pos []) subvcs in
  let allvars = List.fold_left Iset.disj (Pos []) subvcs in
  match allvars with
  | Neg l
     -> let (neg,pos) = List.partition Iset.infinite vcs in
        let v = List.fold_left Iset.conj (Neg []) neg in
        let v = List.fold_left Iset.except v pos in
        let subset = Iset.conj (maximize subvcs v) set in
        if subset = Pos [] then set::accu
        else scc_cut_aux vcs (subset :: accu) (Iset.except set subset)
  | Pos (v::l)
     -> let subset = Iset.conj (maximize subvcs (Pos [v])) set in
        if subset = Pos [] then set::accu
        else scc_cut_aux vcs (subset :: accu) (Iset.except set subset)
  | Pos [] -> accu
;;

let scc_cut cs (vv, set) =
  let vcs = list_rmap (fun (a,b,c) -> possible_values vv a) cs in
  match vcs with
  | [] -> (cut_cost vcs [set], vv, [set])
  | x::t -> let cut = scc_cut_aux vcs [] set
            in (cut_cost vcs cut, vv, cut)
;;

let rec filter_clauses v set = function
  | [] -> []
  | (cond,st,stm)::cs
     -> let c = evaluate cond v set in
        if c = Efalse then filter_clauses v set cs
        else (c,st,stm) :: (filter_clauses v set cs)
;;

let rec test_cond = function
  | Ein (v, Pos []) -> Or []
  | Ein (v, Pos [x]) -> Equals (v, x)
  | Ein (v, Pos l) -> Or (list_rmap (fun x -> Equals (v, x)) l)
  | Ein (v, Neg []) -> And []
  | Ein (v, Neg l) -> assert false
  | Eand l -> And (list_rmap test_cond l)
  | Eor l -> Or (list_rmap test_cond l)
  | Etrue -> And []
  | Efalse -> Or []
;;

let test_cost (cond,st,stm) = cout_mem_cond (test_cond cond);;

let pos_test (cond,st,stm) = pos_cond cond;;

let make_char states x =
  match x with
  | Decision (None, "\"") -> []
  | Case ("state", arms, Decision (None, "\""))
     -> list_rmap (fun (Arm (x,y)) -> Rule (x,y)) arms
  | _ -> [Rule (states, x)]
;;

let xcost states x =
  if states <> [] then cout_mem_character (make_char states x)
  else cout_mem_st x
;;

let restrict context v set =
  let old = List.assoc v context in
  let others = List.remove_assoc v context in
  (v, (Iset.conj old set)) :: others
;;

let rec make_cascade cond iftrue iffalse =
  match cond with
  | Etrue -> iftrue
  | Efalse -> iffalse
  | Ein (v, Pos l) -> If (test_cond cond, iftrue, [], iffalse)
  | Ein (v, Neg l) -> If (test_cond (Ein (v, Pos l)), iffalse, [], iftrue)
  | Eor [] -> iffalse
  | Eor (h::t) -> make_cascade h iftrue (make_cascade (Eor t) iftrue iffalse)
  | Eand [] -> iftrue
  | Eand (h::t) -> make_cascade h (make_cascade (Eand t) iftrue iffalse) iffalse
;;

let rec estimate_test_cost = function
  | Etrue | Efalse -> 0
  | Ein (_, Pos l) -> List.length l
  | Ein (_, Neg l) -> List.length l
  | Eor l -> List.fold_left (fun x y -> x + (estimate_test_cost y)) 0 l
  | Eand l -> List.fold_left (fun x y -> x + (estimate_test_cost y)) 0 l
;;

let rec normalize_cut (cost,v,part) =
  let newpart = list_rmap Iset.normalize part in
  (cost, v, list_sortu newpart)
;;

(* context = set of possible values for each variable
   states = the list of states if this is the toplevel statement, [] otherwise
*)
let rec synthesis context states = function
  | [] -> Decision (None, "\"")             (* empty character *)
  | [(cond,st,stm)] -> Decision (st, stm)
  | cs
     -> let cuts = (list_rmap (max_cut cs) context)
                   @ (list_rmap (scc_cut cs) context) in
        let cuts = List.filter (fun (a,b,c) -> List.length c > 1) cuts in
        let cuts = list_sortu (list_rmap normalize_cut cuts) in
        let cuts = Sort.list (fun (a,b,c) (d,e,f) -> a < d) cuts in
        let goodcuts =
          if !yourelate then begin
            (* prerr_endline "Hurrying up..."; *)
            list_prefix 1 cuts
          end else list_prefix 3 cuts
        in
        let goodcuts =
          let is_statecut (a,b,c) = b = "state" in
          if states <> [] && not (List.exists is_statecut goodcuts)
          then try (List.find is_statecut cuts) :: goodcuts
               with Not_found -> goodcuts
          else goodcuts
        in
        let cut_cand = list_rmap (do_cut states context cs) goodcuts in
        let candidates =
          if cut_cand = [] || List.length cs <= 3
          then (do_test context cs) :: cut_cand
          else cut_cand
        in
        let cost_cand = list_rmap (fun x -> (xcost states x, x)) candidates in
        begin match Sort.list (fun (a,b) (c,d) -> a < c) cost_cand with
        | [] -> assert false
        | (_,c)::_ -> c
        end

and do_cut states context cs (cost, v, cut) =
  let (default, others) =
    let (neg, pos) = List.partition Iset.infinite cut in
    match neg with
    | [defset]
       -> let d = synthesis (restrict context v defset) []
                            (filter_clauses v defset cs)
          in (d, pos)
    | [] when states <> [] -> (Decision (None, "\""), cut)
    | [] -> let set = List.hd cut in
            let d = synthesis (restrict context v set) []
                              (filter_clauses v set cs)
            in (d, List.tl cut)
    | _ -> assert false
  in
  let make_arm set =
    let stm = synthesis (restrict context v set) [] (filter_clauses v set cs) in
    match set with Pos s -> Arm (list_sortu s, stm)
                 | Neg s -> assert false
  in
  let arms = list_rmap make_arm others in
  if states <> [] then Case (v, arms, default)
  else Xlsimplif.optimize_case v arms default
      (* Case (v, arms, default) *)

and do_test context cs =
  let (pos, neg) = List.partition pos_test cs in
  if pos = [] then begin
    let costs = list_rmap (fun (a,b,c) -> (estimate_test_cost a,(a,b,c))) neg in
    let sorted = Sort.list (fun (a,b) (c,d) -> a < c) costs in
    match sorted with
    | [] -> assert false
    | (_,(cond,st,stm))::t
       -> let others = list_rmap snd t in
          let iftrue = Decision (st, stm) in
          let iffalse = synthesis context [] others in
          make_cascade cond iftrue iffalse
  end else begin
    let pos = Sort.list (fun c1 c2 -> test_cost c1 < test_cost c2) pos in
    match pos with
    | [] -> assert false
    | (c,st,stm)::cs
       -> let else_branch = synthesis context [] (neg @ cs) in
          If (test_cond c, Decision (st, stm), [], else_branch)
  end
;;

let find_best states cs =
  let vars = list_sortu ("state" :: (get_vars cs)) in
  let context = list_rmap (fun x -> (x, Neg [])) vars in
  let ctx = restrict context "state" (Pos states) in
  let cs = filter_clauses "state" (Pos states) cs in
  make_char states (synthesis ctx states cs)
;;