(* 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 Invalid_choice
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

let symbol_choices = Hashtbl.create 1023
let num_choices = ref []

let add_symbol_choice symbol codomain_item =
  let current_choices =
    try
      Hashtbl.find symbol_choices symbol
    with Not_found -> []
  in
  Hashtbl.replace symbol_choices symbol (codomain_item :: current_choices)
let add_num_choice choice = num_choices := choice :: !num_choices

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.MutCaseFixAndCofixRefineNotImplemented ->
        (* TODO remove this case as soon as refine is fully implemented *)
        (try
          debug_print (Printf.sprintf "TYPE CHECKER %s" (CicPp.ppterm term));
          let term' = CicTypeChecker.type_of_aux' metasenv context term in
          Ok (term',metasenv)
(*         with _ -> Ko) *)
        with _ -> Uncertain)
    | CicRefine.Uncertain _ ->
        debug_print ("%%% UNCERTAIN!!! " ^ CicPp.ppterm term) ;
        Uncertain
    | _ ->
        debug_print ("%%% PRUNED!!! " ^ CicPp.ppterm term) ;
        Ko

let resolve (env: environment) (item: domain_item) ?(num = "") ?(args = []) () =
  snd (Environment.find item env) env num args

  (* 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
    | CicTextualParser2Ast.LocatedTerm (loc, term) -> aux loc context term
    | CicTextualParser2Ast.Appl terms -> Cic.Appl (List.map (aux loc context) terms)
    | CicTextualParser2Ast.Appl_symbol (symb, i, args) ->
        let cic_args = List.map (aux loc context) args in
        resolve env (Symbol (symb, i)) ~args:cic_args ()
    | CicTextualParser2Ast.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) ] ())
    | CicTextualParser2Ast.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 (pat, term) =
          let rec do_branch' context = function
            | [] -> aux loc context term
            | hd :: tl ->
                let cic_body = do_branch' (Cic.Name hd :: context) tl in
                Cic.Lambda (Cic.Name hd, Cic.Implicit, cic_body)
          in
          match pat with
          | _ :: tl -> (* ignoring constructor *) do_branch' context tl
          | [] -> assert false
        in
        let (indtype_uri, indtype_no) =
          match resolve env (Id indty_ident) () with
          | Cic.MutInd (uri, tyno, _) -> uri, tyno
          | Cic.Implicit -> raise Try_again
          | _ -> raise Invalid_choice
        in
        Cic.MutCase (indtype_uri, indtype_no, cic_outtype, cic_term,
          (List.map do_branch branches))
    | CicTextualParser2Ast.LetIn (var, def, body) ->
        let cic_def = aux loc context def in
        let name = Cic.Name var in
        let cic_body = aux loc (name :: context) body in
        Cic.LetIn (name, cic_def, cic_body)
    | CicTextualParser2Ast.LetRec (kind, defs, body) ->
        let context' =
          List.fold_left (fun acc (var, _, _, _) -> Cic.Name var :: acc)
            context defs
        in
        let cic_body = aux loc context' body in
        let inductiveFuns =
          List.map
            (fun (var, body, typ, decr_idx) ->
              let cic_body = aux loc context' body in
              let cic_type = aux_option loc context typ in
              (var, decr_idx, cic_type, cic_body))
            defs
        in
        let counter = ref 0 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
           * indcutiveFun 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 ->
                Cic.LetIn (Cic.Name var, Cic.CoFix (!counter, funs), cic))
        in
        List.fold_right (build_term inductiveFuns) inductiveFuns cic_body
    | CicTextualParser2Ast.Ident (name, subst) ->
        (* TODO hanlde explicit substitutions *)
        (try
          let index = find_in_environment name context in
          if subst <> [] then
            CicTextualParser2.fail loc "Explicit substitutions not allowed here";
          Cic.Rel index
        with Not_found -> resolve env (Id name) ())
    | CicTextualParser2Ast.Num (num, i) -> resolve env (Num i) ~num ()
    | CicTextualParser2Ast.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)
    | CicTextualParser2Ast.Sort `Prop -> Cic.Sort Cic.Prop
    | CicTextualParser2Ast.Sort `Set -> Cic.Sort Cic.Set
    | CicTextualParser2Ast.Sort `Type -> Cic.Sort Cic.Type
    | CicTextualParser2Ast.Sort `CProp -> Cic.Sort Cic.CProp
  and aux_option loc context = function
    | None -> Cic.Implicit
    | Some term -> aux loc context term
  in
  match ast with
  | CicTextualParser2Ast.LocatedTerm (loc, term) -> aux loc context term
  | _ -> assert false

let domain_of_term ~context ast =
  let rec aux loc context = function
    | CicTextualParser2Ast.LocatedTerm (_, term) -> aux loc context term
    | CicTextualParser2Ast.Appl terms ->
        List.fold_left (fun dom term -> Domain.union dom (aux loc context term))
          Domain.empty terms
    | CicTextualParser2Ast.Appl_symbol (symb, i, args) ->
        List.fold_left (fun dom term -> Domain.union dom (aux loc context term))
          (Domain.singleton (Symbol (symb, i))) args
    | CicTextualParser2Ast.Binder (_, var, typ, body) ->
        let type_dom = aux_option loc context typ in
        let body_dom = aux loc (var :: context) body in
        Domain.union type_dom body_dom
    | CicTextualParser2Ast.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 (pat, term) =
          match pat with
          | _ :: tl ->
              aux loc
                (List.fold_left (fun acc var -> (Cic.Name var) :: acc)
                  context tl)
                term
          | [] -> assert false
        in
        let branches_dom =
          List.fold_left (fun dom branch -> Domain.union dom (do_branch branch))
            Domain.empty branches
        in
        Domain.add (Id indty_ident)
          (Domain.union outtype_dom (Domain.union term_dom branches_dom))
    | CicTextualParser2Ast.LetIn (var, body, where) ->
        let body_dom = aux loc context body in
        let where_dom = aux loc (Cic.Name var :: context) where in
        Domain.union body_dom where_dom
    | CicTextualParser2Ast.LetRec (kind, defs, where) ->
        let context' =
          List.fold_left (fun acc (var, _, _, _) -> Cic.Name var :: acc)
            context defs
        in
        let where_dom = aux loc context' where in
        let defs_dom =
          List.fold_left
            (fun dom (_, body, typ, _) ->
              Domain.union (aux loc context' body) (aux_option loc context typ))
            Domain.empty defs
        in
        Domain.union where_dom defs_dom
    | CicTextualParser2Ast.Ident (name, subst) ->
        (* TODO hanlde explicit substitutions *)
        (try
          let index = find_in_environment name context in
          if subst <> [] then
            CicTextualParser2.fail loc "Explicit substitutions not allowed here";
          Domain.empty
        with Not_found -> Domain.singleton (Id name))
    | CicTextualParser2Ast.Num (num, i) -> Domain.singleton (Num i)
    | CicTextualParser2Ast.Meta (index, local_context) ->
        List.fold_left
          (fun dom term -> Domain.union dom (aux_option loc context term))
            Domain.empty local_context
    | CicTextualParser2Ast.Sort _ -> Domain.empty
  and aux_option loc context = function
    | None -> Domain.empty
    | Some t -> aux loc context t
  in
  match ast with
  | CicTextualParser2Ast.LocatedTerm (loc, term) -> aux loc context term
  | _ -> assert false

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
      C.output_html (`Msg (`T "Locate query:"));
      MQueryUtil.text_of_query
       (fun s -> C.output_html ~append_NL:false (`Msg (`T s)))
       "" query; 
      C.output_html (`Msg (`T "Result:"));
      MQueryUtil.text_of_result
        (fun s -> C.output_html (`Msg (`T s))) "" result;
      let uris' =
       match uris with
        | [] ->
           [UriManager.string_of_uri (C.input_or_locate_uri
            ~title:("URI matching \"" ^ id ^ "\" unknown."))]
        | [uri] -> [uri]
        | _ ->
            C.interactive_user_uri_choice ~selection_mode:`MULTIPLE
             ~ok:"Try every selection." ~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
    =
      let current_dom = (* TODO temporary, remove ASAP *)
        Environment.fold (fun item _ dom -> Domain.add item dom)
          current_env Domain.empty
      in
      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 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, _) ->
              (try Hashtbl.find symbol_choices symb with Not_found -> [])
          | Num instance -> !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 =
        let filled_env =
          Domain.fold
            (fun item env ->
              Environment.add item ("Implicit", fun _ _ _ -> Cic.Implicit) env)
            todo_dom current_env
        in
        try
          let cic_term =
            interpretate ~context:disambiguate_context ~env:filled_env term
          in
          refine metasenv context cic_term
        with
        | Try_again -> Uncertain
        | Invalid_choice -> Ko
      in
      (* (4) build all possible interpretations *)
      let rec aux current_env todo_dom =
        if Domain.is_empty todo_dom then
          match test_env current_env Domain.empty with
          | Ok (term, metasenv) -> [ current_env, term, metasenv ]
          | Ko | Uncertain -> []
        else
          let item = Domain.choose todo_dom in
          let remaining_dom = Domain.remove item todo_dom in
          debug_print (sprintf "CHOOSED ITEM: %s" (string_of_domain_item item));
          let choices = lookup_choices item in
          let rec filter = function
            | [] -> []
            | codomain_item :: tl ->
                let new_env = Environment.add item codomain_item current_env in
                (match test_env new_env remaining_dom with
                | Ok (term, metasenv) ->
                    (if Domain.is_empty remaining_dom then
                      [ new_env, term, metasenv ]
                    else
                      aux new_env remaining_dom)
                    @ filter tl
                | Uncertain ->
                    (if Domain.is_empty remaining_dom then
                      []
                    else
                      aux new_env remaining_dom)
                    @ filter tl
                | Ko -> filter tl)
          in
          filter choices
      in
      let (choosed_env, choosed_term, choosed_metasenv) =
        match aux current_env todo_dom with
        | [] -> raise NoWellTypedInterpretation
        | [ x ] ->
            debug_print "UNA SOLA SCELTA";
            x
        | 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))
                    (Domain.elements term_dom))
                l
            in
            let choosed = C.interactive_interpretation_choice choices in
            List.nth l choosed
      in
      (choosed_env, choosed_metasenv, choosed_term)

  end