open Inference;;
open Utils;;

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


type equality_sign = Negative | Positive;;

let string_of_sign = function
  | Negative -> "Negative"
  | Positive -> "Positive"
;;


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 set_selection = ref (fun set -> EqualitySet.min_elt set);;

let select env passive =
  let (neg_list, neg_set), (pos_list, pos_set) = passive 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
    )
  | _ ->
      let remove eq l =
        List.filter (fun e -> not (e = eq)) l
      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), 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) active =
  (* first step, remove already present equalities *)
  let duplicate =
    let rec aux = function
      | [] -> false
      | (s, eq)::tl when s = sign ->
          if meta_convertibility_eq current eq then true
          else aux tl
      | _::tl -> aux tl
    in
    aux active
  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, current =
            demodulation !maxmeta env current equality in
          maxmeta := newmeta;
          if is_identity env current then
            None
          else
            aux env (sign, current) tl
    in
    aux env (sign, current) active
;;


let forward_simplify_new env (new_neg, new_pos) active =
  let remove_identities equalities =
    let ok eq = not (is_identity env eq) in
    List.filter ok equalities
  in
  let rec simpl active target =
    match active with
    | [] -> target
    | (Negative, _)::tl -> simpl tl target
    | (Positive, source)::tl ->
        let newmeta, target = demodulation !maxmeta env target source in
        maxmeta := newmeta;
        if is_identity env target then target
        else simpl tl target
  in
  let new_neg = List.map (simpl active) new_neg
  and new_pos = List.map (simpl active) new_pos in
  new_neg, remove_identities new_pos
;;


let backward_simplify env (sign, current) active =
  match sign with
  | Negative -> active
  | Positive ->
      let active = 
        List.map
          (fun (s, equality) ->
             (*            match s with *)
             (*            | Negative -> s, equality *)
             (*            | Positive -> *)
             let newmeta, equality =
               demodulation !maxmeta env equality current in
             maxmeta := newmeta;
             s, equality)
          active
      in
      let active =
        List.filter (fun (s, eq) -> not (is_identity env eq)) active
      in
      let find eq1 where =
        List.exists (fun (s, e) -> meta_convertibility_eq eq1 e) where
      in
      List.fold_right
        (fun (s, eq) res -> if find eq res then res else (s, eq)::res)
        active []
;;


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 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) 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) ->
(*           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 new' = forward_simplify_new env new' active in
          
          let active =
            backward_simplify env (sign, current) active
(*             match new' with *)
(*             | [], [] -> backward_simplify env (sign, current) active *)
(*             | _ -> active *)
          in

          print_endline "\n================================================";
          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
              given_clause 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 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), _, _ = term_equality in
    let active = [] in
    let passive = make_passive [term_equality] equalities in
    Printf.printf "\ncurrent goal: %s ={%s} %s\n"
      (PP.ppterm left) (PP.ppterm eq_ty) (PP.ppterm right);
    Printf.printf "\ncontext:\n%s\n" (PP.ppcontext context);
    Printf.printf "\nmetasenv:\n%s\n" (print_metasenv metasenv);
    Printf.printf "\nequalities:\n";
    List.iter
      (function (_, (ty, t1, t2), _, _) ->
         let w1 = weight_of_term t1 in
         let w2 = weight_of_term t2 in
         let res = nonrec_kbo t1 t2 in
         Printf.printf "{%s}: %s<%s> %s %s<%s>\n" (PP.ppterm ty)
           (PP.ppterm t1) (string_of_weight w1)
           (string_of_comparison res)
           (PP.ppterm t2) (string_of_weight w2))
      equalities;
    print_endline "--------------------------------------------------";
    let start = Sys.time () in
    print_endline "GO!";
    let res = given_clause env passive active in
    let finish = Sys.time () 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; weight_age_counter := v
  and set_sel () = set_selection := (fun s -> EqualitySet.max_elt s)
  and set_conf f = configuration_file := f
  in
  Arg.parse [
    "-r", Arg.Int set_ratio, "Weight-Age equality selection ratio";

    "-i", Arg.Unit set_sel,
    "Inverse selection (select heaviest equalities before)";

    "-c", Arg.String set_conf, "Configuration file (for the db connection)";
  ] (fun a -> ()) "Usage:"
in
Helm_registry.load_from !configuration_file;
main ()