(* Copyright (C) 2004, HELM Team.
 * 
 * This file is part of HELM, an Hypertextual, Electronic
 * Library of Mathematics, developed at the Computer Science
 * Department, University of Bologna, Italy.
 * 
 * HELM is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 * 
 * HELM is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with HELM; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston,
 * MA  02111-1307, USA.
 * 
 * For details, see the HELM World-Wide-Web page,
 * http://helm.cs.unibo.it/
 *)

open Printf

open DisambiguateTypes
open UriManager

exception No_choices of domain_item
exception NoWellTypedInterpretation

  (** raised when an environment is not enough informative to decide *)
exception Try_again

let debug = true
let debug_print = if debug then prerr_endline else ignore

let descr_of_domain_item = function
 | Id s -> s
 | Symbol (s, _) -> s
 | Num i -> string_of_int i

type test_result =
  | Ok of Cic.term * Cic.metasenv
  | Ko
  | Uncertain

let refine metasenv context term =
  let metasenv, term = CicMkImplicit.expand_implicits metasenv context term in
  debug_print (sprintf "TEST_INTERPRETATION: %s" (CicPp.ppterm term));
  try
    let term', _, metasenv' = CicRefine.type_of_aux' metasenv context term in
    Ok (term', metasenv')
  with
    | CicRefine.Uncertain _ ->
        debug_print ("%%% UNCERTAIN!!! " ^ CicPp.ppterm term) ;
        Uncertain
    | CicRefine.RefineFailure _ ->
        debug_print ("%%% PRUNED!!! " ^ CicPp.ppterm term) ;
        Ko

let resolve (env: environment) (item: domain_item) ?(num = "") ?(args = []) () =
  try
    snd (Environment.find item env) env num args
  with Not_found -> assert false

  (* TODO move it to Cic *)
let find_in_environment name context =
  let rec aux acc = function
    | [] -> raise Not_found
    | Cic.Name hd :: tl when hd = name -> acc
    | _ :: tl ->  aux (acc + 1) tl
  in
  aux 1 context

let interpretate ~context ~env ast =
  let rec aux loc context = function
    | CicAst.AttributedTerm (`Loc loc, term) ->
        aux loc context term
    | CicAst.AttributedTerm (_, term) -> aux loc context term
    | CicAst.Appl (CicAst.Symbol (symb, i) :: args) ->
        let cic_args = List.map (aux loc context) args in
        resolve env (Symbol (symb, i)) ~args:cic_args ()
    | CicAst.Appl terms -> Cic.Appl (List.map (aux loc context) terms)
    | CicAst.Binder (binder_kind, (var, typ), body) ->
        let cic_type = aux_option loc context typ in
        let cic_body = aux loc (var :: context) body in
        (match binder_kind with
        | `Lambda -> Cic.Lambda (var, cic_type, cic_body)
        | `Pi | `Forall -> Cic.Prod (var, cic_type, cic_body)
        | `Exists ->
            resolve env (Symbol ("exists", 0))
              ~args:[ cic_type; Cic.Lambda (var, cic_type, cic_body) ] ())
    | CicAst.Case (term, indty_ident, outtype, branches) ->
        let cic_term = aux loc context term in
        let cic_outtype = aux_option loc context outtype in
        let do_branch ((head, args), term) =
          let rec do_branch' context = function
            | [] -> aux loc context term
            | (name, typ) :: tl ->
                let cic_body = do_branch' (name :: context) tl in
                let typ =
                  match typ with
                  | None -> Cic.Implicit (Some `Type)
                  | Some typ -> aux loc context typ
                in
                Cic.Lambda (name, typ, cic_body)
          in
          do_branch' context args
        in
        let (indtype_uri, indtype_no) =
          match indty_ident with
          | Some indty_ident ->
              (match resolve env (Id indty_ident) () with
              | Cic.MutInd (uri, tyno, _) -> (uri, tyno)
              | Cic.Implicit _ -> raise Try_again
              | _ -> raise DisambiguateChoices.Invalid_choice)
          | None ->
              let fst_constructor =
                match branches with
                | ((head, _), _) :: _ -> head
                | [] -> raise DisambiguateChoices.Invalid_choice
              in
              (match resolve env (Id fst_constructor) () with
              | Cic.MutConstruct (indtype_uri, indtype_no, _, _) ->
                  (indtype_uri, indtype_no)
              | Cic.Implicit _ -> raise Try_again
              | _ -> raise DisambiguateChoices.Invalid_choice)
        in
        Cic.MutCase (indtype_uri, indtype_no, cic_outtype, cic_term,
          (List.map do_branch branches))
    | CicAst.LetIn ((name, typ), def, body) ->
        let cic_def = aux loc context def in
        let cic_def =
          match typ with
          | None -> cic_def
          | Some t -> Cic.Cast (cic_def, aux loc context t)
        in
        let cic_body = aux loc (name :: context) body in
        Cic.LetIn (name, cic_def, cic_body)
    | CicAst.LetRec (kind, defs, body) ->
        let context' =
          List.fold_left (fun acc ((name, _), _, _) -> name :: acc)
            context defs
        in
        let cic_body = aux loc context' body in
        let inductiveFuns =
          List.map
            (fun ((name, typ), body, decr_idx) ->
              let cic_body = aux loc context' body in
              let cic_type = aux_option loc context typ in
              let name =
                match name with
                | Cic.Anonymous ->
                    CicTextualParser2.fail loc
                      "Recursive functions cannot be anonymous"
                | Cic.Name name -> name
              in
              (name, decr_idx, cic_type, cic_body))
            defs
        in
        let counter = ref ~-1 in
        let build_term funs =
          (* this is the body of the fold_right function below. Rationale: Fix
           * and CoFix cases differs only in an additional index in the
           * inductiveFun list, see Cic.term *)
          match kind with
          | `Inductive ->
              (fun (var, _, _, _) cic ->
                incr counter;
                Cic.LetIn (Cic.Name var, Cic.Fix (!counter, funs), cic))
          | `CoInductive ->
              let funs =
                List.map (fun (name, _, typ, body) -> (name, typ, body)) funs
              in
              (fun (var, _, _, _) cic ->
                incr counter;
                Cic.LetIn (Cic.Name var, Cic.CoFix (!counter, funs), cic))
        in
        List.fold_right (build_term inductiveFuns) inductiveFuns cic_body
    | CicAst.Ident (name, subst) ->
        (try
          let index = find_in_environment name context in
          if subst <> None then
            CicTextualParser2.fail loc
              "Explicit substitutions not allowed here";
          Cic.Rel index
        with Not_found ->
          let cic = resolve env (Id name) () in
          let mk_subst uris =
            let ids_to_uris =
              List.map (fun uri -> UriManager.name_of_uri uri, uri) uris
            in
            (match subst with
            | Some subst ->
                List.map
                  (fun (s, term) ->
                    (try
                      List.assoc s ids_to_uris, aux loc context term
                     with Not_found ->
                       raise DisambiguateChoices.Invalid_choice))
                  subst
            | None -> List.map (fun uri -> uri, Cic.Implicit None) uris)
          in
          (* the try is for CicTypeChecker.typecheck *)
          (try 
            match cic with
            | Cic.Const (uri, []) ->
                let uris =
                  match CicEnvironment.get_obj uri with
                  (*match CicTypeChecker.typecheck uri with*)
                  | Cic.Constant (_, _, _, uris) -> uris
                  | _ -> assert false
                in
                Cic.Const (uri, mk_subst uris)
            | Cic.Var (uri, []) ->
                let uris =
                  match CicEnvironment.get_obj uri with
                  (*match CicTypeChecker.typecheck uri with*)
                  | Cic.Variable (_, _, _, uris) -> uris
                  | _ -> assert false
                in
                Cic.Var (uri, mk_subst uris)
            | Cic.MutInd (uri, i, []) ->
                let uris =
                  match CicEnvironment.get_obj uri with
                  (*match CicTypeChecker.typecheck uri with*)
                  | Cic.InductiveDefinition (_, uris, _) -> uris
                  | _ -> assert false
                in
                Cic.MutInd (uri, i, mk_subst uris)
            | Cic.MutConstruct (uri, i, j, []) ->
                let uris =
                  match CicEnvironment.get_obj uri with
                  (*match CicTypeChecker.typecheck uri with*)
                  | Cic.InductiveDefinition (_, uris, _) -> uris
                  | _ -> assert false
                in
                Cic.MutConstruct (uri, i, j, mk_subst uris)
            | Cic.Meta _ | Cic.Implicit _ as t ->
(*
                prerr_endline (sprintf
                  "Warning: %s must be instantiated with _[%s] but we do not enforce it"
                  (CicPp.ppterm t)
                  (String.concat "; "
                    (List.map
                      (fun (s, term) -> s ^ " := " ^ CicAstPp.pp_term term)
                      subst)));
*)
                t
            | _ ->
              raise DisambiguateChoices.Invalid_choice
           with 
             CicEnvironment.CircularDependency _ -> 
               raise DisambiguateChoices.Invalid_choice))
    | CicAst.Implicit -> Cic.Implicit None
    | CicAst.Num (num, i) -> resolve env (Num i) ~num ()
    | CicAst.Meta (index, subst) ->
        let cic_subst =
          List.map
            (function None -> None | Some term -> Some (aux loc context term))
            subst
        in
        Cic.Meta (index, cic_subst)
    | CicAst.Sort `Prop -> Cic.Sort Cic.Prop
    | CicAst.Sort `Set -> Cic.Sort Cic.Set
    | CicAst.Sort `Type -> Cic.Sort (Cic.Type (CicUniv.fresh())) (* TASSI *)
    | CicAst.Sort `CProp -> Cic.Sort Cic.CProp
    | CicAst.Symbol (symbol, instance) ->
        resolve env (Symbol (symbol, instance)) ()
  and aux_option loc context = function
    | None -> Cic.Implicit (Some `Type)
    | Some term -> aux loc context term
  in
  match ast with
  | CicAst.AttributedTerm (`Loc loc, term) -> aux loc context term
  | term -> aux (-1, -1) context term

let domain_of_term ~context ast =
    (* "aux" keeps domain in reverse order and doesn't care about duplicates.
     * Domain item more in deep in the list will be processed first.
     *)
  let rec aux loc context = function
    | CicAst.AttributedTerm (`Loc loc, term) -> aux loc context term
    | CicAst.AttributedTerm (_, term) -> aux loc context term
    | CicAst.Appl terms ->
        List.fold_left (fun dom term -> aux loc context term @ dom) [] terms
    | CicAst.Binder (kind, (var, typ), body) ->
        let kind_dom =
          match kind with
          | `Exists -> [ Symbol ("exists", 0) ]
          | _ -> []
        in
        let type_dom = aux_option loc context typ in
        let body_dom = aux loc (var :: context) body in
        body_dom @ type_dom @ kind_dom
    | CicAst.Case (term, indty_ident, outtype, branches) ->
        let term_dom = aux loc context term in
        let outtype_dom = aux_option loc context outtype in
        let do_branch ((head, args), term) =
          let (term_context, args_domain) =
            List.fold_left
              (fun (cont, dom) (name, typ) ->
                (name :: cont,
                 (match typ with
                 | None -> dom
                 | Some typ -> aux loc cont typ @ dom)))
              (context, []) args
          in
          args_domain @ aux loc term_context term
        in
        let branches_dom =
          List.fold_left (fun dom branch -> do_branch branch @ dom) [] branches
        in
        branches_dom @ outtype_dom @ term_dom @
        (match indty_ident with None -> [] | Some ident -> [ Id ident ])
    | CicAst.LetIn ((var, typ), body, where) ->
        let body_dom = aux loc context body in
        let type_dom = aux_option loc context typ in
        let where_dom = aux loc (var :: context) where in
        where_dom @ type_dom @ body_dom
    | CicAst.LetRec (kind, defs, where) ->
        let context' =
          List.fold_left (fun acc ((var, typ), _, _) -> var :: acc)
            context defs
        in
        let where_dom = aux loc context' where in
        let defs_dom =
          List.fold_left
            (fun dom ((_, typ), body, _) ->
              aux loc context' body @ aux_option loc context typ)
            [] defs
        in
        where_dom @ defs_dom
    | CicAst.Ident (name, subst) ->
        (try
          let index = find_in_environment name context in
          if subst <> None then
            CicTextualParser2.fail loc
              "Explicit substitutions not allowed here"
          else
            []
        with Not_found ->
          (match subst with
          | None -> [Id name]
          | Some subst ->
              List.fold_left
                (fun dom (_, term) ->
                  let dom' = aux loc context term in
                  dom' @ dom)
                [Id name] subst))
    | CicAst.Implicit -> []
    | CicAst.Num (num, i) -> [ Num i ]
    | CicAst.Meta (index, local_context) ->
        List.fold_left (fun dom term -> aux_option loc context term @ dom) []
          local_context
    | CicAst.Sort _ -> []
    | CicAst.Symbol (symbol, instance) -> [ Symbol (symbol, instance) ]

  and aux_option loc context = function
    | None -> []
    | Some t -> aux loc context t
  in

    (* e.g. [5;1;1;1;2;3;4;1;2] -> [2;1;4;3;5] *)
  let rev_uniq =
    let module SortedItem =
      struct
        type t = DisambiguateTypes.domain_item
        let compare = Pervasives.compare
      end
    in
    let module Set = Set.Make (SortedItem) in
    fun l ->
      let rev_l = List.rev l in
      let (_, uniq_rev_l) =
        List.fold_left
          (fun (members, rev_l) elt ->
            if Set.mem elt members then
              (members, rev_l)
            else
              Set.add elt members, elt :: rev_l)
          (Set.empty, []) rev_l
      in
      List.rev uniq_rev_l
  in
            
  rev_uniq
    (match ast with
    | CicAst.AttributedTerm (`Loc loc, term) -> aux loc context term
    | term -> aux (-1, -1) context term)


  (* dom1 \ dom2 *)
let domain_diff dom1 dom2 =
(* let domain_diff = Domain.diff *)
  let is_in_dom2 =
    List.fold_left (fun pred elt -> (fun elt' -> elt' = elt || pred elt'))
      (fun _ -> false) dom2
  in
  List.filter (fun elt -> not (is_in_dom2 elt)) dom1

module Make (C: Callbacks) =
  struct
    let choices_of_id mqi_handle id =
     let query  =  MQueryGenerator.locate id in
     let result = MQueryInterpreter.execute mqi_handle query in
     let uris =
      List.map
       (function uri,_ ->
         MQueryMisc.wrong_xpointer_format_from_wrong_xpointer_format' uri
       ) result in
      let uris' =
       match uris with
        | [] ->
           [UriManager.string_of_uri (C.input_or_locate_uri
            ~title:("URI matching \"" ^ id ^ "\" unknown.") ~id ())]
        | [uri] -> [uri]
        | _ ->
            C.interactive_user_uri_choice ~selection_mode:`MULTIPLE
             ~ok:"Try selected." ~enable_button_for_non_vars:true
             ~title:"Ambiguous input." ~id
             ~msg: ("Ambiguous input \"" ^ id ^
                "\". Please, choose one or more interpretations:")
             uris
      in
      List.map
        (fun uri ->
          (uri,
           let term =
             try
               HelmLibraryObjects.term_of_uri (UriManager.uri_of_string uri)
             with _ -> assert false
            in
           fun _ _ _ -> term))
        uris'

    let disambiguate_term mqi_handle context metasenv term ~aliases:current_env
    =
      debug_print "NEW DISAMBIGUATE INPUT";
      let disambiguate_context =  (* cic context -> disambiguate context *)
        List.map
          (function None -> Cic.Anonymous | Some (name, _) -> name)
          context
      in
      let term_dom = domain_of_term ~context:disambiguate_context term in
      debug_print (sprintf "DISAMBIGUATION DOMAIN: %s"
        (string_of_domain term_dom));
      let current_dom =
        Environment.fold (fun item _ dom -> item :: dom) current_env []
      in
      let todo_dom = domain_diff term_dom current_dom in
      (* (2) lookup function for any item (Id/Symbol/Num) *)
      let lookup_choices =
        let id_choices = Hashtbl.create 1023 in
        fun item ->
        let choices =
          match item with
          | Id id ->
              (try
                Hashtbl.find id_choices id
              with Not_found ->
                let choices = choices_of_id mqi_handle id in
                Hashtbl.add id_choices id choices;
                choices)
          | Symbol (symb, _) -> DisambiguateChoices.lookup_symbol_choices symb
          | Num instance -> DisambiguateChoices.lookup_num_choices ()
        in
        if choices = [] then raise (No_choices item);
        choices
      in
      (* (3) test an interpretation filling with meta uninterpreted identifiers
       *)
      let test_env current_env todo_dom univ = 
        let filled_env =
          List.fold_left
            (fun env item ->
              Environment.add item
                ("Implicit",
                 (match item with
                 | Id _ | Num _ -> (fun _ _ _ -> Cic.Implicit (Some `Closed))
                 | Symbol _ -> (fun _ _ _ -> Cic.Implicit None))) env)
            current_env todo_dom 
        in
        try
          CicUniv.set_working univ; 
          let cic_term =
            interpretate ~context:disambiguate_context ~env:filled_env term
          in
           let k = refine metasenv context cic_term in
           let new_univ = CicUniv.get_working () in
            (k , new_univ )
        with
        | Try_again -> Uncertain,univ 
        | DisambiguateChoices.Invalid_choice -> Ko,univ 
      in
      (* (4) build all possible interpretations *)
      let rec aux current_env todo_dom base_univ =
        match todo_dom with
        | [] ->
            (match test_env current_env [] base_univ with
            | Ok (term, metasenv),new_univ -> 
                               [ current_env, metasenv, term, new_univ ]
            | Ko,_ | Uncertain,_ -> [])
        | item :: remaining_dom ->
            debug_print (sprintf "CHOOSED ITEM: %s"
             (string_of_domain_item item));
            let choices = lookup_choices item in
            let rec filter univ = function 
              | [] -> []
              | codomain_item :: tl ->
                  debug_print (sprintf "%s CHOSEN" (fst codomain_item)) ;
                  let new_env =
                    Environment.add item codomain_item current_env
                  in
                  (match test_env new_env remaining_dom univ with
                  | Ok (term, metasenv),new_univ ->
                      (match remaining_dom with
                      | [] -> [ new_env, metasenv, term, new_univ ]
                      | _ -> aux new_env remaining_dom new_univ )@ 
                        filter univ tl
                  | Uncertain,new_univ ->
                      (match remaining_dom with
                      | [] -> []
                      | _ -> aux new_env remaining_dom new_univ )@ 
                        filter univ tl
                  | Ko,_ -> filter univ tl)
            in
            filter base_univ choices 
      in
       let base_univ = CicUniv.get_working () in
      try
       match aux current_env todo_dom base_univ with
       | [] -> raise NoWellTypedInterpretation
       | [ e,me,t,u ] as l ->
           debug_print "UNA SOLA SCELTA";
           CicUniv.set_working u;
           [ e,me,t ]
       | l ->
           debug_print (sprintf "PIU' SCELTE (%d)" (List.length l));
           let choices =
             List.map
               (fun (env, _, _, _) ->
                 List.map
                   (fun domain_item ->
                     let description =
                       fst (Environment.find domain_item env)
                     in
                     (descr_of_domain_item domain_item, description))
                   term_dom)
               l
           in
           let choosed = C.interactive_interpretation_choice choices in
           let l' = List.map (List.nth l) choosed in
           match l' with
             [] -> assert false
           | [e,me,t,u] -> 
               CicUniv.set_working u;
               (*CicUniv.print_working_graph ();*)
               [e,me,t]
           | hd::tl -> (* ok, testlibrary... cosi' stampa MANY... bah *)
               List.map (fun (e,me,t,u) -> (e,me,t)) l'
     with
      CicEnvironment.CircularDependency s -> 
        raise (Failure "e chi la becca sta CircularDependency?");
  end