module GA = GrafiteAst;;
module PT = CicNotationPt;;
module A = Ast;;

let rec collect_arities_from_term = function
  | A.Constant name -> [name,0]
  | A.Variable name -> []
  | A.Function (name,l) -> 
      (name,List.length l) :: List.flatten (List.map collect_arities_from_term l)
;;

let rec collect_fv_from_term = function
  | A.Constant name -> []
  | A.Variable name -> [name]
  | A.Function (_,l) -> 
      List.flatten (List.map collect_fv_from_term l)
;;

let collect_arities_from_atom a = 
  let aux = function
    | A.Proposition name -> assert false
    | A.Predicate _ -> assert false
    | A.True -> []
    | A.False -> []
    | A.Eq (t1,t2) -> collect_arities_from_term t1 @ collect_arities_from_term t2
    | A.NotEq (t1,t2) -> collect_arities_from_term t1 @ collect_arities_from_term t2
  in
  HExtlib.list_uniq (List.sort (fun (a,_) (b,_) -> compare a b) (aux a))
;;
  
let collect_fv_from_atom a = 
  let aux = function
    | A.Proposition name -> assert false
    | A.Predicate _ -> assert false
    | A.True -> []
    | A.False -> []
    | A.Eq (t1,t2) -> collect_fv_from_term t1 @ collect_fv_from_term t2
    | A.NotEq (t1,t2) -> collect_fv_from_term t1 @ collect_fv_from_term t2
  in
  HExtlib.list_uniq (List.sort compare (aux a))
;;  

let rec convert_term = function
  | A.Variable x -> PT.Ident (x,None)
  | A.Constant x -> PT.Ident (x,None)
  | A.Function (name, args) -> 
      PT.Appl (PT.Ident (name,None) :: List.map convert_term args)
;;

let atom_of_formula = function
    | A.Disjunction _ -> assert false
    | A.NegAtom a -> a (* removes the negation *)
    | A.Atom a -> a
;;
  
let rec mk_arrow component = function
  | 0 -> PT.Ident (component,None)
  | n -> 
      PT.Binder 
        (`Forall,
          ((PT.Variable (PT.FreshVar "_")),Some (PT.Ident (component,None))),
          mk_arrow component (n-1))
;;

let build_ctx_for_arities arities t = 
  let rec aux = function
    | [] -> t
    | (name,nargs)::tl ->
        PT.Binder 
          (`Forall,
            (PT.Ident (name,None),Some (mk_arrow "A" nargs)),
            aux tl)
  in
  aux arities
;;

let convert_atom a = 
  let aux = function
  | A.Proposition _ -> assert false
  | A.Predicate (name,params) -> assert false
  | A.True -> PT.Ident ("True",None)
  | A.False -> PT.Ident ("False",None)
  | A.Eq (l,r) 
  | A.NotEq (l,r) -> (* removes the negation *)
      PT.Appl [
        PT.Ident ("eq",None); PT.Ident ("A",None); convert_term l; convert_term r]
  in
  build_ctx_for_arities (List.map (fun x -> (x,0)) (collect_fv_from_atom a)) (aux a)
;;

let rec convert_formula no_arities context f =
  let atom = atom_of_formula f in
  let t = convert_atom atom in
  let rec build_ctx n = function
    | [] -> t
    | hp::tl -> 
        PT.Binder 
          (`Forall,
            (PT.Ident ("H" ^ string_of_int n,None), 
                       Some (convert_formula true [] hp)), 
            build_ctx (n+1) tl)
  in
  let arities = if no_arities then [] else collect_arities_from_atom atom in
  build_ctx_for_arities arities (build_ctx 0 context) 
;;

let convert_ast statements context = function
  | A.Comment s -> 
      let s = String.sub s 1 (String.length s - 1) in
      if s.[String.length s - 1] = '\n' then
        s.[String.length s - 1] <- '\000';
      statements @ [GA.Comment (HExtlib.dummy_floc,GA.Note (HExtlib.dummy_floc,s))],
      context
  | A.Inclusion (s,_) ->
      statements @ [
        GA.Comment (
          HExtlib.dummy_floc, GA.Note (
            HExtlib.dummy_floc,"Inclusion of: " ^ s))], context
  | A.AnnotatedFormula (name,kind,f,_,_) -> 
      match kind with
      | A.Axiom 
      | A.Hypothesis ->
          statements, f::context
      | A.Negated_conjecture ->
          let f = 
            PT.Binder 
             (`Forall,
               (PT.Ident ("A",None),Some (PT.Sort `Set)), 
               convert_formula false context f)
          in
          let o = PT.Theorem (`Theorem,name,f,None) in
          statements @ [
            GA.Executable(
              HExtlib.dummy_floc,GA.Command(
                HExtlib.dummy_floc,GA.Obj(HExtlib.dummy_floc,o)))],
          context
      | A.Definition 
      | A.Lemma 
      | A.Theorem 
      | A.Conjecture
      | A.Lemma_conjecture 
      | A.Plain 
      | A.Unknown -> assert false
;;

(* OPTIONS *)
let tptppath = ref "./";;
let librarymode = ref false;;
let spec = [
  ("-tptppath", 
      Arg.String (fun x -> tptppath := x), 
      "Where to find the Axioms/ and Problems/ directory");
  ("-librarymode",
      Arg.Set librarymode,
      "... not supported yet")
]

(* HELPERS *)
let resolve s = 
  let resolved_name = 
    if Filename.check_suffix s ".p" then
      (assert (String.length s > 5);
      let prefix = String.sub s 0 3 in
      !tptppath ^ "/Problems/" ^ prefix ^ "/" ^ s)
    else
      !tptppath ^ "/" ^ s
  in
  if HExtlib.is_regular resolved_name then
    resolved_name
  else
    begin
      prerr_endline ("Unable to find " ^ s ^ " (" ^ resolved_name ^ ")");
      exit 1
    end
;;

(* MAIN *)
let _ =
  let usage = "Usage: tptp2grafite [options] file" in
  let inputfile = ref "" in
  Arg.parse spec (fun s -> inputfile := s) usage;
  if !inputfile = "" then 
    begin
      prerr_endline usage;
      exit 1
    end;
  let rec aux = function
    | [] -> []
    | ((A.Inclusion (file,_)) as hd) :: tl ->
        let file = resolve file in
        let lexbuf = Lexing.from_channel (open_in file) in
        let statements = Parser.main Lexer.yylex lexbuf in
        hd :: aux (statements @ tl)
    | hd::tl -> hd :: aux tl
  in
  let statements = aux [A.Inclusion (!inputfile ^ ".p",[])] in
  let grafite_ast_statements,_ = 
    List.fold_left 
      (fun (st, ctx) f -> 
        let newst, ctx = convert_ast st ctx f in
        newst, ctx)
      ([],[]) statements 
  in
  let pp t = 
    (* for a correct pp we should disambiguate the term... *)
    let term_pp = CicNotationPp.pp_term in 
    let lazy_term_pp = fun x -> assert false in
    let obj_pp = CicNotationPp.pp_obj in
    print_endline 
      (GrafiteAstPp.pp_statement ~term_pp ~lazy_term_pp ~obj_pp t)
  in
  List.iter pp grafite_ast_statements;
  exit 0