open Inference;;
open Utils;;

type result =
  | Failure
  | Success of Cic.term option * environment
;;


(*
let symbols_of_equality (_, (_, left, right), _, _) =
  TermSet.union (symbols_of_term left) (symbols_of_term right)
;;
*)

let symbols_of_equality ((_, (_, left, right, _), _, _) as equality) =
  let m1 = symbols_of_term left in
  let m = 
    TermMap.fold
      (fun k v res ->
         try
           let c = TermMap.find k res in
           TermMap.add k (c+v) res
         with Not_found ->
           TermMap.add k v res)
      (symbols_of_term right) m1
  in
(*   Printf.printf "symbols_of_equality %s:\n" *)
(*     (string_of_equality equality); *)
(*   TermMap.iter (fun k v -> Printf.printf "%s: %d\n" (CicPp.ppterm k) v) m; *)
(*   print_newline (); *)
  m
;;


module OrderedEquality =
struct
  type t = Inference.equality

  let compare eq1 eq2 =
    match meta_convertibility_eq eq1 eq2 with
    | true -> 0
    | false ->
        let _, (ty, left, right, _), _, _ = eq1
        and _, (ty', left', right', _), _, _ = eq2 in
        let weight_of t = fst (weight_of_term ~consider_metas:false t) in
        let w1 = (weight_of ty) + (weight_of left) + (weight_of right)
        and w2 = (weight_of ty') + (weight_of left') + (weight_of right') in
        match Pervasives.compare w1 w2 with
        | 0 -> Pervasives.compare eq1 eq2
        | res -> res
end

module EqualitySet = Set.Make(OrderedEquality);;


let weight_age_ratio = ref 0;; (* settable by the user from the command line *)
let weight_age_counter = ref !weight_age_ratio;;

let symbols_ratio = ref 0;;
let symbols_counter = ref 0;;


let select env passive active =
  let (neg_list, neg_set), (pos_list, pos_set) = passive in
  let remove eq l =
    List.filter (fun e -> not (e = eq)) l
  in
  if !weight_age_ratio > 0 then
    weight_age_counter := !weight_age_counter - 1;
  match !weight_age_counter with
  | 0 -> (
      weight_age_counter := !weight_age_ratio;
      match neg_list, pos_list with
      | hd::tl, pos ->
          (Negative, hd), ((tl, EqualitySet.remove hd neg_set), (pos, pos_set))
      | [], hd::tl ->
          (Positive, hd), (([], neg_set), (tl, EqualitySet.remove hd pos_set))
      | _, _ -> assert false
    )
  | _ when (!symbols_counter > 0) && (EqualitySet.is_empty neg_set) -> (
      symbols_counter := !symbols_counter - 1;
      let cardinality map =
        TermMap.fold (fun k v res -> res + v) map 0
      in
      match active with
      | (Negative, e)::_ ->
          let symbols = symbols_of_equality e in
          let card = cardinality symbols in
          let f equality (i, e) =
            let common, others =
              TermMap.fold
                (fun k v (r1, r2) ->
                   if TermMap.mem k symbols then
                     let c = TermMap.find k symbols in
                     let c1 = abs (c - v) in
                     let c2 = v - c1 in
                     r1 + c2, r2 + c1
                   else
                     r1, r2 + v)
                (symbols_of_equality equality) (0, 0)
            in
(*             Printf.printf "equality: %s, common: %d, others: %d\n" *)
(*               (string_of_equality ~env equality) common others; *)
            let c = others + (abs (common - card)) in
            if c < i then (c, equality)
            else (i, e)
          in
          let e1 = EqualitySet.min_elt pos_set in
          let initial =
            let common, others = 
              TermMap.fold
                (fun k v (r1, r2) ->
                   if TermMap.mem k symbols then
                     let c = TermMap.find k symbols in
                     let c1 = abs (c - v) in
                     let c2 = v - (abs (c - v)) in
                     r1 + c1, r2 + c2
                   else
                     r1, r2 + v)
                (symbols_of_equality e1) (0, 0)
            in
            (others + (abs (common - card))), e1
          in
          let _, current = EqualitySet.fold f pos_set initial in
(*           Printf.printf "\nsymbols-based selection: %s\n\n" *)
(*             (string_of_equality ~env current); *)
          (Positive, current),
          (([], neg_set),
           (remove current pos_list, EqualitySet.remove current pos_set))
      | _ ->
            let current = EqualitySet.min_elt pos_set in
          let passive =
            (neg_list, neg_set),
            (remove current pos_list, EqualitySet.remove current pos_set)
          in
          (Positive, current), passive
    )
  | _ ->
      symbols_counter := !symbols_ratio;
      let set_selection set = EqualitySet.min_elt set in
      if EqualitySet.is_empty neg_set then
        let current = set_selection pos_set in
        let passive =
          (neg_list, neg_set),
          (remove current pos_list, EqualitySet.remove current pos_set)
        in
        (Positive, current), passive
      else
        let current = set_selection neg_set in
        let passive =
          (remove current neg_list, EqualitySet.remove current neg_set),
          (pos_list, pos_set)
        in
        (Negative, current), passive
;;


let make_passive neg pos =
  let set_of equalities =
    List.fold_left (fun s e -> EqualitySet.add e s) EqualitySet.empty equalities
  in
  (neg, set_of neg), (pos, set_of pos)
;;


let add_to_passive passive (new_neg, new_pos) =
  let (neg_list, neg_set), (pos_list, pos_set) = passive in
  let ok set equality = not (EqualitySet.mem equality set) in
  let neg = List.filter (ok neg_set) new_neg
  and pos = List.filter (ok pos_set) new_pos in
  let add set equalities =
    List.fold_left (fun s e -> EqualitySet.add e s) set equalities
  in
  (neg @ neg_list, add neg_set neg), (pos_list @ pos, add pos_set pos)
;;


let passive_is_empty = function
  | ([], _), ([], _) -> true
  | _ -> false
;;


(* TODO: find a better way! *)
let maxmeta = ref 0;;

let infer env sign current active =
  let rec infer_negative current = function
    | [] -> [], []
    | (Negative, _)::tl -> infer_negative current tl
    | (Positive, equality)::tl ->
        let res = superposition_left env current equality in
        let neg, pos = infer_negative current tl in
        res @ neg, pos

  and infer_positive current = function
    | [] -> [], []
    | (Negative, equality)::tl ->
        let res = superposition_left env equality current in
        let neg, pos = infer_positive current tl in
        res @ neg, pos
    | (Positive, equality)::tl ->
        let maxm, res = superposition_right !maxmeta env current equality in
        let maxm, res' = superposition_right maxm env equality current in
        maxmeta := maxm;
        let neg, pos = infer_positive current tl in

(*         Printf.printf "risultato di superposition_right: %s %s\n%s\n\n" *)
(*           (string_of_equality ~env current) (string_of_equality ~env equality) *)
(*           (String.concat "\n" (List.map (string_of_equality ~env) res)); *)
(*         Printf.printf "risultato di superposition_right: %s %s\n%s\n\n" *)
(*           (string_of_equality ~env equality) (string_of_equality ~env current) *)
(*           (String.concat "\n" (List.map (string_of_equality ~env) res')); *)
        
        neg, res @ res' @ pos
  in
  match sign with
  | Negative -> infer_negative current active
  | Positive -> infer_positive current active
;;


let contains_empty env (negative, positive) =
  let metasenv, context, ugraph = env in
  try
    let (proof, _, _, _) =
      List.find
        (fun (proof, (ty, left, right, ordering), m, a) ->
           fst (CicReduction.are_convertible context left right ugraph))
        negative
    in
    true, Some proof
  with Not_found ->
    false, None
;;


let forward_simplify env (sign, current) ?passive active =
  let pn, pp =
    match passive with
    | None -> [], []
    | Some ((pn, _), (pp, _)) ->
        (List.map (fun e -> Negative, e) pn),
        (List.map (fun e -> Positive, e) pp)
  in
  let all = active @ pn @ pp in
  let rec find_duplicate sign current = function
    | [] -> false
    | (s, eq)::tl when s = sign ->
        if meta_convertibility_eq current eq then true
        else find_duplicate sign current tl
    | _::tl -> find_duplicate sign current tl
  in
(*   let duplicate = find_duplicate sign current all in *)
(*   if duplicate then *)
(*     None *)
(*   else *)
  let rec aux env (sign, current) = function
    | [] -> Some (sign, current)
    | (Negative, _)::tl -> aux env (sign, current) tl
    | (Positive, equality)::tl ->
        let newmeta, newcurrent =
          demodulation !maxmeta env current equality in
        maxmeta := newmeta;
        if is_identity env newcurrent then
          if sign = Negative then
            Some (sign, current)
          else
            None
        else if newcurrent <> current then
          aux env (sign, newcurrent) active
        else
          aux env (sign, newcurrent) tl
  in
  let res = aux env (sign, current) all in
  match res with
  | None -> None
  | Some (s, c) ->
      if find_duplicate s c all then
        None
      else
        let pred (sign, eq) =
          if sign <> s then false
          else subsumption env c eq
        in
        if List.exists pred all then None
        else res
;;


let forward_simplify_new env (new_neg, new_pos) ?passive active =
  let pn, pp =
    match passive with
    | None -> [], []
    | Some ((pn, _), (pp, _)) ->
        (List.map (fun e -> Negative, e) pn),
        (List.map (fun e -> Positive, e) pp)
  in
  let all = active @ pn @ pp in
  let remove_identities equalities =
    let ok eq = not (is_identity env eq) in
    List.filter ok equalities
  in
  let rec simpl all' target =
    match all' with
    | [] -> target
    | (Negative, _)::tl -> simpl tl target
    | (Positive, source)::tl ->
        let newmeta, newtarget = demodulation !maxmeta env target source in
        maxmeta := newmeta;
        if is_identity env newtarget then newtarget
        else if newtarget <> target then (
(*           Printf.printf "OK:\n%s\n%s\n" *)
(*             (string_of_equality ~env target) *)
(*             (string_of_equality ~env newtarget); *)
(*           print_newline (); *)
          simpl all newtarget
        )
        else simpl tl newtarget
  in
  let new_neg = List.map (simpl all) new_neg
  and new_pos = remove_identities (List.map (simpl all) new_pos) in
  let new_pos_set =
    List.fold_left (fun s e -> EqualitySet.add e s) EqualitySet.empty new_pos
  in
  let new_pos = EqualitySet.elements new_pos_set in
  let f sign' target (sign, eq) =
(*     Printf.printf "f %s <%s> (%s, <%s>)\n" *)
(*       (string_of_sign sign') (string_of_equality ~env target) *)
(*       (string_of_sign sign) (string_of_equality ~env eq); *)
    if sign <> sign' then false
    else subsumption env target eq
  in
(*   new_neg, new_pos *)
  (List.filter (fun e -> not (List.exists (f Negative e) all)) new_neg,
   List.filter (fun e -> not (List.exists (f Positive e) all)) new_pos)
;;


let backward_simplify_active env (new_neg, new_pos) active =
  let new_pos = List.map (fun e -> Positive, e) new_pos in
  let active, newa = 
    List.fold_right
      (fun (s, equality) (res, newn) ->
         match forward_simplify env (s, equality) new_pos with
         | None when s = Negative ->
             Printf.printf "\nECCO QUI: %s\n"
               (string_of_equality ~env equality);
             print_newline ();
             res, newn
         | None -> res, newn
         | Some (s, e) ->
             if equality = e then
               (s, e)::res, newn
             else
               res, (s, e)::newn)
      active ([], [])
  in
  let find eq1 where =
    List.exists (fun (s, e) -> meta_convertibility_eq eq1 e) where
  in
  let active, newa =
    let f (s, eq) res =
      if (is_identity env eq) || (find eq res) then res else (s, eq)::res
    in
    List.fold_right
      (fun (s, eq) res ->
         if (is_identity env eq) || (find eq res) then res else (s, eq)::res)
      active [],
    List.fold_right
      (fun (s, eq) (n, p) ->
         if (s <> Negative) && (is_identity env eq) then
           (n, p)
         else
           if s = Negative then eq::n, p
           else n, eq::p)
      newa ([], [])
  in
  match newa with
  | [], [] -> active, None
  | _ -> active, Some newa
;;


let backward_simplify_passive env (new_neg, new_pos) passive =
  let new_pos = List.map (fun e -> Positive, e) new_pos in
  let (nl, ns), (pl, ps) = passive in
  let f sign equality (resl, ress, newn) =
    match forward_simplify env (sign, equality) new_pos with
    | None -> resl, EqualitySet.remove equality ress, newn
    | Some (s, e) ->
        if equality = e then
          equality::resl, ress, newn
        else
          let ress = EqualitySet.remove equality ress in
          resl, ress, e::newn
  in
  let nl, ns, newn = List.fold_right (f Negative) nl ([], ns, [])
  and pl, ps, newp = List.fold_right (f Positive) pl ([], ps, []) in
  match newn, newp with
  | [], [] -> ((nl, ns), (pl, ps)), None
  | _, _ -> ((nl, ns), (pl, ps)), Some (newn, newp)
;;


let backward_simplify env new' ?passive active =
  let active, newa = backward_simplify_active env new' active in
  match passive with
  | None ->
      active, (([], EqualitySet.empty), ([], EqualitySet.empty)), newa, None
  | Some passive ->
      let passive, newp =
        backward_simplify_passive env new' passive in
      active, passive, newa, newp
;;


  
let rec given_clause env passive active =
  match passive_is_empty passive with
  | true -> Failure
  | false ->
(*       Printf.printf "before select\n"; *)
      let (sign, current), passive = select env passive active in
(*       Printf.printf "before simplification: sign: %s\ncurrent: %s\n\n" *)
(*         (string_of_sign sign) (string_of_equality ~env current); *)
      match forward_simplify env (sign, current) ~passive active with
(*       | None when sign = Negative -> *)
(*           Printf.printf "OK!!! %s %s" (string_of_sign sign) *)
(*             (string_of_equality ~env current); *)
(*           print_newline (); *)
(*           let proof, _, _, _ = current in *)
(*           Success (Some proof, env) *)
      | None ->
(*           Printf.printf "avanti... %s %s" (string_of_sign sign) *)
(*             (string_of_equality ~env current); *)
(*           print_newline (); *)
          given_clause env passive active
      | Some (sign, current) ->
          if (sign = Negative) && (is_identity env current) then (
            Printf.printf "OK!!! %s %s" (string_of_sign sign)
              (string_of_equality ~env current);
            print_newline ();
            let proof, _, _, _ = current in
            Success (Some proof, env)
          ) else (            
            print_endline "\n================================================";
            Printf.printf "selected: %s %s"
              (string_of_sign sign) (string_of_equality ~env current);
            print_newline ();

            let new' = infer env sign current active in
            let res, proof = contains_empty env new' in
            if res then
              Success (proof, env)
            else
              let new' = forward_simplify_new env new' active in
              let active =
                match sign with
                | Negative -> active
                | Positive ->
                    let active, _, newa, _ =
                      backward_simplify env ([], [current]) active
                    in
                    match newa with
                    | None -> active
                    | Some (n, p) ->
                        let nn = List.map (fun e -> Negative, e) n
                        and pp = List.map (fun e -> Positive, e) p in
                        nn @ active @ pp
              in
              let _ =
                Printf.printf "active:\n%s\n"
                  (String.concat "\n"
                     ((List.map
                         (fun (s, e) -> (string_of_sign s) ^ " " ^
                            (string_of_equality ~env e)) active)));
                print_newline ();
              in
              let _ =
                match new' with
                | neg, pos ->
                    Printf.printf "new':\n%s\n"
                      (String.concat "\n"
                         ((List.map
                             (fun e -> "Negative " ^
                                (string_of_equality ~env e)) neg) @
                            (List.map
                               (fun e -> "Positive " ^
                                  (string_of_equality ~env e)) pos)));
                    print_newline ();
              in
              match contains_empty env new' with
              | false, _ -> 
                  let active =
                    match sign with
                    | Negative -> (sign, current)::active
                    | Positive -> active @ [(sign, current)]
                  in
                  let passive = add_to_passive passive new' in
                  let (_, ns), (_, ps) = passive in
                  Printf.printf "passive:\n%s\n"
                    (String.concat "\n"
                       ((List.map (fun e -> "Negative " ^
                                     (string_of_equality ~env e))
                           (EqualitySet.elements ns)) @
                          (List.map (fun e -> "Positive " ^
                                       (string_of_equality ~env e))
                             (EqualitySet.elements ps))));
                  print_newline ();
                  given_clause env passive active
              | true, proof ->
                  Success (proof, env)
          )
;;


let rec given_clause_fullred env passive active =
  match passive_is_empty passive with
  | true -> Failure
  | false ->
(*       Printf.printf "before select\n"; *)
      let (sign, current), passive = select env passive active in
(*       Printf.printf "before simplification: sign: %s\ncurrent: %s\n\n" *)
(*         (string_of_sign sign) (string_of_equality ~env current); *)
      match forward_simplify env (sign, current) ~passive active with
      | None ->
          given_clause_fullred env passive active
      | Some (sign, current) ->
          if (sign = Negative) && (is_identity env current) then (
            Printf.printf "OK!!! %s %s" (string_of_sign sign)
              (string_of_equality ~env current);
            print_newline ();
            let proof, _, _, _ = current in
            Success (Some proof, env)
          ) else (
            print_endline "\n================================================";
            Printf.printf "selected: %s %s"
              (string_of_sign sign) (string_of_equality ~env current);
            print_newline ();

            let new' = infer env sign current active in

            let active =
              if is_identity env current then active
              else
                match sign with
                | Negative -> (sign, current)::active
                | Positive -> active @ [(sign, current)]
            in
(*             let _ = *)
(*               match new' with *)
(*               | neg, pos -> *)
(*                   Printf.printf "new' before simpl:\n%s\n" *)
(*                     (String.concat "\n" *)
(*                        ((List.map *)
(*                            (fun e -> "Negative " ^ *)
(*                               (string_of_equality ~env e)) neg) @ *)
(*                           (List.map *)
(*                              (fun e -> "Positive " ^ *)
(*                                 (string_of_equality ~env e)) pos))); *)
(*                   print_newline (); *)
(*             in *)
            let rec simplify new' active passive =
              let new' = forward_simplify_new env new' ~passive active in
              let active, passive, newa, retained =
                backward_simplify env new' ~passive active
              in
              match newa, retained with
              | None, None -> active, passive, new'
              | Some (n, p), None
              | None, Some (n, p) ->
                  let nn, np = new' in
                  simplify (nn @ n, np @ p) active passive
              | Some (n, p), Some (rn, rp) ->
                  let nn, np = new' in
                  simplify (nn @ n @ rn, np @ p @ rp) active passive
            in
            let active, passive, new' = simplify new' active passive in
            let _ =
              Printf.printf "active:\n%s\n"
                (String.concat "\n"
                   ((List.map
                       (fun (s, e) -> (string_of_sign s) ^ " " ^
                          (string_of_equality ~env e)) active)));
              print_newline ();
            in
            let _ =
              match new' with
              | neg, pos ->
                  Printf.printf "new':\n%s\n"
                    (String.concat "\n"
                       ((List.map
                           (fun e -> "Negative " ^
                              (string_of_equality ~env e)) neg) @
                          (List.map
                             (fun e -> "Positive " ^
                                (string_of_equality ~env e)) pos)));
                  print_newline ();
            in
            match contains_empty env new' with
            | false, _ -> 
                let passive = add_to_passive passive new' in
(*                 let (_, ns), (_, ps) = passive in *)
(*                 Printf.printf "passive:\n%s\n" *)
(*                   (String.concat "\n" *)
(*                      ((List.map (fun e -> "Negative " ^ *)
(*                                    (string_of_equality ~env e)) *)
(*                          (EqualitySet.elements ns)) @ *)
(*                         (List.map (fun e -> "Positive " ^ *)
(*                                      (string_of_equality ~env e)) *)
(*                            (EqualitySet.elements ps)))); *)
(*                 print_newline (); *)
                given_clause_fullred env passive active
            | true, proof ->
                Success (proof, env)
          )
;;


let get_from_user () =
  let dbd = Mysql.quick_connect
    ~host:"localhost" ~user:"helm" ~database:"mowgli" () in
  let rec get () =
    match read_line () with
    | "" -> []
    | t -> t::(get ())
  in
  let term_string = String.concat "\n" (get ()) in
  let env, metasenv, term, ugraph =
    List.nth (Disambiguate.Trivial.disambiguate_string dbd term_string) 0
  in
  term, metasenv, ugraph
;;


let given_clause_ref = ref given_clause;;


let main () =
  let module C = Cic in
  let module T = CicTypeChecker in
  let module PET = ProofEngineTypes in
  let module PP = CicPp in
  let term, metasenv, ugraph = get_from_user () in
  let proof = None, (1, [], term)::metasenv, C.Meta (1, []), term in
  let proof, goals =
    PET.apply_tactic (PrimitiveTactics.intros_tac ()) (proof, 1) in
  let goal = List.nth goals 0 in
  let _, metasenv, meta_proof, _ = proof in
  let _, context, goal = CicUtil.lookup_meta goal metasenv in
  let equalities, maxm = find_equalities context proof in
  maxmeta := maxm; (* TODO ugly!! *)
  let env = (metasenv, context, ugraph) in
  try
    let term_equality = equality_of_term meta_proof goal in
    let meta_proof, (eq_ty, left, right, ordering), _, _ = term_equality in
    let active = [] in
    let passive = make_passive [term_equality] equalities in
    Printf.printf "\ncurrent goal: %s\n"
      (string_of_equality ~env term_equality);
    Printf.printf "\ncontext:\n%s\n" (PP.ppcontext context);
    Printf.printf "\nmetasenv:\n%s\n" (print_metasenv metasenv);
    Printf.printf "\nequalities:\n%s\n"
      (String.concat "\n"
         (List.map
            (string_of_equality ~env)
            equalities));
    print_endline "--------------------------------------------------";
    let start = Unix.gettimeofday () in
    print_endline "GO!";
    let res = !given_clause_ref env passive active in
    let finish = Unix.gettimeofday () in
    match res with
    | Failure ->
        Printf.printf "NO proof found! :-(\n\n"
    | Success (Some proof, env) ->
        Printf.printf "OK, found a proof!:\n%s\n%.9f\n" (PP.ppterm proof)
          (finish -. start);
    | Success (None, env) ->
        Printf.printf "Success, but no proof?!?\n\n"
  with exc ->
    print_endline ("EXCEPTION: " ^ (Printexc.to_string exc));
;;


let configuration_file = ref "../../gTopLevel/gTopLevel.conf.xml";;

let _ =
  let set_ratio v = weight_age_ratio := (v+1); weight_age_counter := (v+1)
  and set_sel v = symbols_ratio := v; symbols_counter := v;
  and set_conf f = configuration_file := f
  and set_lpo () = Utils.compare_terms := lpo
  and set_kbo () = Utils.compare_terms := nonrec_kbo
  and set_fullred () = given_clause_ref := given_clause_fullred
  in
  Arg.parse [
    "-f", Arg.Unit set_fullred, "Use full-reduction strategy";
    
    "-r", Arg.Int set_ratio, "Weight-Age equality selection ratio (default: 0)";

    "-s", Arg.Int set_sel,
    "symbols-based selection ratio (relative to the weight ratio)";

    "-c", Arg.String set_conf, "Configuration file (for the db connection)";

    "-lpo", Arg.Unit set_lpo, "Use lpo term ordering";

    "-kbo", Arg.Unit set_kbo, "Use (non-recursive) kbo term ordering (default)";
  ] (fun a -> ()) "Usage:"
in
Helm_registry.load_from !configuration_file;
main ()