(*
    ||M||  This file is part of HELM, an Hypertextual, Electronic        
    ||A||  Library of Mathematics, developed at the Computer Science     
    ||T||  Department, University of Bologna, Italy.                     
    ||I||                                                                
    ||T||  HELM is free software; you can redistribute it and/or         
    ||A||  modify it under the terms of the GNU General Public License   
    \   /  version 2 or (at your option) any later version.      
     \ /   This software is distributed as is, NO WARRANTY.     
      V_______________________________________________________________ *)

(* $Id$ *)

exception RefineFailure of (Stdpp.location * string) Lazy.t;;
exception Uncertain of (Stdpp.location * string) Lazy.t;;
exception AssertFailure of string Lazy.t;;

module C = NCic
module Ref = NReference

let indent = ref "";;
let inside c = indent := !indent ^ String.make 1 c;;
let outside () = indent := String.sub !indent 0 (String.length !indent -1);;

(*
let pp s = 
  prerr_endline (Printf.sprintf "%-20s" !indent ^ " " ^ Lazy.force s)
;;  
*)

let pp _ = ();;

let wrap_exc msg = function
  | NCicUnification.Uncertain _ -> Uncertain msg
  | NCicUnification.UnificationFailure _ -> RefineFailure msg
  | NCicTypeChecker.TypeCheckerFailure _ -> RefineFailure msg
  | e -> raise e
;;

let exp_implicit metasenv context expty =
 let foo x = match expty with Some t -> `WithType t | None -> x in
 function      
  | `Closed -> NCicMetaSubst.mk_meta metasenv [] (foo `Term)
  | `Type -> NCicMetaSubst.mk_meta metasenv context (foo `Type)
  | `Term -> NCicMetaSubst.mk_meta metasenv context (foo `Term)
  | _ -> assert false
;;


let check_allowed_sort_elimination localise r orig =
   let mkapp he arg =
     match he with
     | C.Appl l -> C.Appl (l @ [arg])
     | t -> C.Appl [t;arg] in
   (* ctx, ind_type @ lefts, sort_of_ind_ty@lefts, outsort *)
   let rec aux metasenv subst context ind arity1 arity2 =
     (*D*)inside 'S'; try let rc = 
     let arity1 = NCicReduction.whd ~subst context arity1 in
     pp (lazy(NCicPp.ppterm ~subst ~metasenv ~context arity1 ^ "   elimsto   " ^
        NCicPp.ppterm ~subst ~metasenv ~context arity2 ^ "\nMENV:\n"^
        NCicPp.ppmetasenv ~subst metasenv));
     match arity1 with
     | C.Prod (name,so1,de1) (* , t ==?== C.Prod _ *) ->
        let metasenv, meta, _ = 
          NCicMetaSubst.mk_meta metasenv ((name,C.Decl so1)::context) `Type 
        in
        let metasenv, subst = 
          try NCicUnification.unify metasenv subst context 
                arity2 (C.Prod (name, so1, meta)) 
          with exc -> raise (wrap_exc (lazy (localise orig, Printf.sprintf
            "expected %s, found %s" (* XXX localizzare meglio *)
            (NCicPp.ppterm ~subst ~metasenv ~context (C.Prod (name, so1, meta)))
            (NCicPp.ppterm ~subst ~metasenv ~context arity2))) exc)
        in
        aux metasenv subst ((name, C.Decl so1)::context)
         (mkapp (NCicSubstitution.lift 1 ind) (C.Rel 1)) de1 meta
     | C.Sort _ (* , t ==?== C.Prod _ *) ->
        let metasenv, meta, _ = NCicMetaSubst.mk_meta metasenv [] `Type in
        let metasenv, subst = 
          try NCicUnification.unify metasenv subst context 
                arity2 (C.Prod ("_", ind, meta)) 
          with exc -> raise (wrap_exc (lazy (localise orig, Printf.sprintf
            "expected %s, found %s" (* XXX localizzare meglio *)
            (NCicPp.ppterm ~subst ~metasenv ~context (C.Prod ("_", ind, meta)))
            (NCicPp.ppterm ~subst ~metasenv ~context arity2))) exc)
        in
        (try NCicTypeChecker.check_allowed_sort_elimination
            ~metasenv ~subst r context ind arity1 arity2; 
            metasenv, subst
        with exc -> raise (wrap_exc (lazy (localise orig, 
          "Sort elimination not allowed ")) exc))
     | _ -> assert false
     (*D*)in outside(); rc with exc -> outside (); raise exc
   in
    aux
;;

let rec typeof 
  ?(localise=fun _ -> Stdpp.dummy_loc) 
  ~look_for_coercion metasenv subst context term expty 
=
  let force_ty metasenv subst context orig t infty expty =
    (*D*)inside 'F'; try let rc = 
    match expty with
    | Some expty ->
       (match t with
       | C.Implicit _ 
       | C.Lambda _ -> metasenv, subst, t, expty
       | _ ->
          pp (lazy (
          (NCicPp.ppterm ~metasenv ~subst ~context infty) ^  " === " ^
          (NCicPp.ppterm ~metasenv ~subst ~context expty)));
           try 
             let metasenv, subst =
               NCicUnification.unify metasenv subst context infty expty 
             in
             metasenv, subst, t, expty
           with exc -> 
             try_coercions ~look_for_coercion ~localise 
               metasenv subst context t orig infty expty true exc)
    | None -> metasenv, subst, t, infty
    (*D*)in outside(); rc with exc -> outside (); raise exc
  in
  let rec typeof_aux metasenv subst context expty = 
    fun t as orig -> 
    (*D*)inside 'R'; try let rc = 
    pp (lazy (NCicPp.ppterm ~metasenv ~subst ~context t));
    pp (lazy (NCicPp.ppmetasenv ~subst metasenv));
    let metasenv, subst, t, infty = 
    match t with
    | C.Rel n ->
        let infty = 
         (try
           match List.nth context (n - 1) with
           | (_,C.Decl ty) -> NCicSubstitution.lift n ty
           | (_,C.Def (_,ty)) -> NCicSubstitution.lift n ty
         with Failure _ -> 
           raise (RefineFailure (lazy (localise t,"unbound variable"))))
        in
        metasenv, subst, t, infty
    | C.Sort (C.Type [false,u]) -> metasenv,subst,t,(C.Sort (C.Type [true, u]))
    | C.Sort (C.Type _) -> 
        raise (AssertFailure (lazy ("Cannot type an inferred type: "^
          NCicPp.ppterm ~subst ~metasenv ~context t)))
    | C.Sort _ -> metasenv,subst,t,(C.Sort (C.Type NCicEnvironment.type0))
    | C.Implicit infos -> 
         let metasenv,t,ty = exp_implicit metasenv context expty infos in
         metasenv, subst, t, ty 
    | C.Meta (n,l) as t -> 
       let ty =
        try
         let _,_,_,ty = NCicUtils.lookup_subst n subst in ty 
        with NCicUtils.Subst_not_found _ -> try
         let _,_,ty = NCicUtils.lookup_meta n metasenv in 
         match ty with C.Implicit _ -> assert false | _ -> ty 
        with NCicUtils.Meta_not_found _ ->
         raise (AssertFailure (lazy (Printf.sprintf
          "%s not found" (NCicPp.ppterm ~subst ~metasenv ~context t))))
       in
       metasenv, subst, t, NCicSubstitution.subst_meta l ty
    | C.Const _ -> 
       metasenv, subst, t, NCicTypeChecker.typeof ~subst ~metasenv context t
    | C.Prod (name,(s as orig_s),(t as orig_t)) ->
       let metasenv, subst, s, s1 = typeof_aux metasenv subst context None s in
       let metasenv, subst, s, s1 = 
         force_to_sort ~look_for_coercion
           metasenv subst context s orig_s localise s1 in
       let context1 = (name,(C.Decl s))::context in
       let metasenv, subst, t, s2 = typeof_aux metasenv subst context1 None t in
       let metasenv, subst, t, s2 = 
         force_to_sort ~look_for_coercion
           metasenv subst context1 t orig_t localise s2 in
       let metasenv, subst, s, t, ty = 
         sort_of_prod ~look_for_coercion localise metasenv subst 
           context orig_s orig_t (name,s) t (s1,s2)
       in
       metasenv, subst, NCic.Prod(name,s,t), ty
    | C.Lambda (n,(s as orig_s),t) ->
       let exp_s, exp_ty_t =
         match expty with
         | None -> None, None
         | Some expty -> 
             match NCicReduction.whd ~subst context expty with
             | C.Prod (_,s,t) -> Some s, Some t
             | _ -> None, None (** XXX FUNCLASS |-> QUALCOSA *)
       in
       let metasenv, subst, s, ty_s = 
         typeof_aux metasenv subst context None s in
       let metasenv, subst, s, _ = 
         force_to_sort ~look_for_coercion
           metasenv subst context s orig_s localise ty_s in
       let (metasenv,subst), exp_ty_t = 
         match exp_s with 
         | Some exp_s -> 
             (try NCicUnification.unify metasenv subst context s exp_s,exp_ty_t
             with exc -> raise (wrap_exc (lazy (localise orig_s, Printf.sprintf
               "Source type %s was expected to be %s" (NCicPp.ppterm ~metasenv
               ~subst ~context s) (NCicPp.ppterm ~metasenv ~subst ~context
               exp_s))) exc))
         | None -> (metasenv, subst), None
       in
       let context = (n,C.Decl s) :: context in
       let metasenv, subst, t, ty_t = 
         typeof_aux metasenv subst context exp_ty_t t 
       in
       metasenv, subst, C.Lambda(n,s,t), C.Prod (n,s,ty_t)
    | C.LetIn (n,(ty as orig_ty),t,bo) ->
       let metasenv, subst, ty, ty_ty = 
         typeof_aux metasenv subst context None ty in
       let metasenv, subst, ty, _ = 
         force_to_sort ~look_for_coercion
           metasenv subst context ty orig_ty localise ty_ty in
       let metasenv, subst, t, _ = 
         typeof_aux metasenv subst context (Some ty) t in
       let context1 = (n, C.Def (t,ty)) :: context in
       let metasenv, subst, bo, bo_ty = 
         typeof_aux metasenv subst context1 None bo 
       in
       let bo_ty = NCicSubstitution.subst ~avoid_beta_redexes:true t bo_ty in
       metasenv, subst, C.LetIn (n, ty, t, bo), bo_ty
    | C.Appl ((he as orig_he)::(_::_ as args)) ->
       let metasenv, subst, he, ty_he = 
         typeof_aux metasenv subst context None he in
       eat_prods ~localise ~look_for_coercion 
         metasenv subst context orig_he he ty_he args
   | C.Appl _ -> raise (AssertFailure (lazy "Appl of length < 2"))
   | C.Match (Ref.Ref (_,Ref.Ind (_,tyno,_)) as r,
          outtype,(term as orig_term),pl) as orig ->
      let _, leftno, itl, _, _ = NCicEnvironment.get_checked_indtys r in
      let _, _, arity, cl = List.nth itl tyno in
      let constructorsno = List.length cl in
      let _, metasenv, args = 
        NCicMetaSubst.saturate metasenv subst context arity 0 in
      let ind = if args = [] then C.Const r else C.Appl (C.Const r::args) in
      let metasenv, subst, term, _ = 
        typeof_aux metasenv subst context (Some ind) term in
      let metasenv, subst, outtype, outsort = 
        typeof_aux metasenv subst context None outtype in
      let parameters, arguments = HExtlib.split_nth leftno args in
      (* let's control if the sort elimination is allowed: [(I q1 ... qr)|B] *)
      let ind =
        if parameters = [] then C.Const r
        else C.Appl ((C.Const r)::parameters) in
      let metasenv, subst, ind, ind_ty = 
        typeof_aux metasenv subst context None ind in
      let metasenv, subst = 
         check_allowed_sort_elimination localise r orig_term metasenv subst 
           context ind ind_ty outsort 
      in
      (* let's check if the type of branches are right *)
      if List.length pl <> constructorsno then
       raise (RefineFailure (lazy (localise orig, 
         "Wrong number of cases in a match")));
      let _, metasenv, subst, pl_rev =
        List.fold_left
          (fun (j, metasenv, subst, branches) p ->
              let cons = 
                let cons = Ref.mk_constructor j r in
                if parameters = [] then C.Const cons
                else C.Appl (C.Const cons::parameters)
              in
              let metasenv, subst, cons, ty_cons = 
                typeof_aux metasenv subst context None cons in
              let ty_branch = 
                NCicTypeChecker.type_of_branch 
                  ~subst context leftno outtype cons ty_cons in
              pp (lazy ("TYPEOFBRANCH: " ^
               NCicPp.ppterm ~metasenv ~subst ~context p ^ " ::: " ^
               NCicPp.ppterm ~metasenv ~subst ~context ty_branch ));
              let metasenv, subst, p, _ = 
                typeof_aux metasenv subst context (Some ty_branch) p in
              j+1, metasenv, subst, p :: branches)
          (1, metasenv, subst, []) pl
      in
      metasenv, subst, 
      C.Match (r, outtype, term, List.rev pl_rev),
      NCicReduction.head_beta_reduce (C.Appl (outtype::arguments@[term]))
    | C.Match _ as orig -> 
        prerr_endline (NCicPp.ppterm ~metasenv ~subst ~context orig);
        assert false
    in
    pp (lazy (NCicPp.ppterm ~metasenv ~subst ~context t ^ " :: "^
         NCicPp.ppterm ~metasenv ~subst ~context infty ));
      force_ty metasenv subst context orig t infty expty
    (*D*)in outside(); rc with exc -> outside (); raise exc
  in
    typeof_aux metasenv subst context expty term

and try_coercions 
  ~localise ~look_for_coercion 
  metasenv subst context t orig_t infty expty perform_unification exc 
=
  (* we try with a coercion *)
  let rec first exc = function
  | [] ->         
      raise (wrap_exc (lazy (localise orig_t, Printf.sprintf
        "The term %s has type %s but is here used with type %s"
        (NCicPp.ppterm ~metasenv ~subst ~context t)
        (NCicPp.ppterm ~metasenv ~subst ~context infty)
        (NCicPp.ppterm ~metasenv ~subst ~context expty))) exc)
  | (metasenv, newterm, newtype, meta)::tl ->
      try 
        let metasenv, subst = 
          NCicUnification.unify metasenv subst context meta t
        in
        let metasenv, subst = 
          if perform_unification then
            NCicUnification.unify metasenv subst context newtype expty
          else metasenv, subst
        in
        metasenv, subst, newterm, newtype
      with 
      | NCicUnification.UnificationFailure _ -> first exc tl
      | NCicUnification.Uncertain _ as exc -> first exc tl
  in 
    first exc
      (look_for_coercion metasenv subst context infty expty)

and force_to_sort 
  ~look_for_coercion metasenv subst context t orig_t localise ty 
=
  match NCicReduction.whd ~subst context ty with
  | C.Meta (_,(0,(C.Irl 0 | C.Ctx []))) as ty -> 
     metasenv, subst, t, ty
  | C.Meta (i,(_,(C.Irl 0 | C.Ctx []))) -> 
     metasenv, subst, t, C.Meta(i,(0,C.Irl 0))
  | C.Meta (i,(_,lc)) ->
     let len = match lc with C.Irl len->len | C.Ctx l->List.length l in
     let metasenv, subst, newmeta = 
       if len > 0 then
         NCicMetaSubst.restrict metasenv subst i (HExtlib.list_seq 1 (len+1)) 
       else metasenv, subst, i
     in
     metasenv, subst, t, C.Meta (newmeta,(0,C.Irl 0))
  | C.Sort _ as ty -> metasenv, subst, t, ty 
  | ty -> 
      try_coercions ~localise ~look_for_coercion metasenv subst context
        t orig_t ty (NCic.Sort (NCic.Type NCicEnvironment.type0)) false 
         (Uncertain (lazy (localise orig_t, 
         "The type of " ^ NCicPp.ppterm ~metasenv ~subst ~context t
         ^ " is not a sort: " ^ NCicPp.ppterm ~metasenv ~subst ~context ty)))

and sort_of_prod ~look_for_coercion
  localise metasenv subst context orig_s orig_t (name,s) t (t1, t2) 
=
   (* force to sort is done in the Prod case in typeof *)
   match t1, t2 with
   | C.Sort _, C.Sort C.Prop -> metasenv, subst, s, t, t2
   | C.Sort (C.Type u1), C.Sort (C.Type u2) ->
        metasenv, subst, s, t, C.Sort (C.Type (NCicEnvironment.max u1 u2)) 
   | C.Sort C.Prop,C.Sort (C.Type _) -> metasenv, subst, s, t, t2
   | C.Meta _, C.Sort _ 
   | C.Meta _, C.Meta (_,(_,_))
   | C.Sort _, C.Meta (_,(_,_)) -> metasenv, subst, s, t, t2 
   | x, (C.Sort _ | C.Meta _) | _, x -> 
      let y, context, orig = 
        if x == t1 then s, context, orig_s 
        else t, ((name,C.Decl s)::context), orig_t
      in
      raise (RefineFailure (lazy (localise orig,Printf.sprintf
        "%s is expected to be a type, but its type is %s that is not a sort" 
         (NCicPp.ppterm ~subst ~metasenv ~context y) 
         (NCicPp.ppterm ~subst ~metasenv ~context x))))

and eat_prods 
  ~localise ~look_for_coercion metasenv subst context orig_he he ty_he args 
=
  (*D*)inside 'E'; try let rc = 
  let rec aux metasenv subst args_so_far he ty_he = function 
  | []->metasenv, subst, NCicUntrusted.mk_appl he (List.rev args_so_far), ty_he
  | arg::tl ->
      match NCicReduction.whd ~subst context ty_he with 
      | C.Prod (_,s,t) ->
          let metasenv, subst, arg, _ = 
            typeof ~look_for_coercion ~localise 
              metasenv subst context arg (Some s) in
          let t = NCicSubstitution.subst ~avoid_beta_redexes:true arg t in
          aux metasenv subst (arg :: args_so_far) he t tl
      | C.Meta _
      | C.Appl (C.Meta _ :: _) as t ->
          let metasenv, subst, arg, ty_arg = 
            typeof ~look_for_coercion ~localise 
              metasenv subst context arg None in
          let metasenv, meta, _ = 
            NCicMetaSubst.mk_meta metasenv 
              (("_",C.Decl ty_arg) :: context) `Type
          in
          let flex_prod = C.Prod ("_", ty_arg, meta) in
          (* next line grants that ty_args is a type *)
          let metasenv,subst, flex_prod, _ =
           typeof ~look_for_coercion ~localise metasenv subst
            context flex_prod None in
          pp (lazy ( "UNIFICATION in CTX:\n"^ 
            NCicPp.ppcontext ~metasenv ~subst context
            ^ "\nOF: " ^
            NCicPp.ppterm ~metasenv ~subst ~context t ^  " === " ^
            NCicPp.ppterm ~metasenv ~subst ~context flex_prod ^ "\n"));
          let metasenv, subst = 
             try NCicUnification.unify metasenv subst context t flex_prod 
             with exc -> raise (wrap_exc (lazy (localise orig_he, Printf.sprintf
              ("The term %s has an inferred type %s, but is applied to the" ^^
               " argument %s of type %s")
              (NCicPp.ppterm ~metasenv ~subst ~context he)
              (NCicPp.ppterm ~metasenv ~subst ~context t)
              (NCicPp.ppterm ~metasenv ~subst ~context arg)
              (NCicPp.ppterm ~metasenv ~subst ~context ty_arg))) exc)
              (* XXX coerce to funclass *)
          in
          let meta = NCicSubstitution.subst ~avoid_beta_redexes:true arg meta in
          aux metasenv subst (arg :: args_so_far) he meta tl
      | C.Match (_,_,C.Meta _,_) 
      | C.Match (_,_,C.Appl (C.Meta _ :: _),_) 
      | C.Appl (C.Const (NReference.Ref (_, NReference.Fix _)) :: _) ->
          raise (Uncertain (lazy (localise orig_he, Printf.sprintf
            ("The term %s is here applied to %d arguments but expects " ^^
            "only %d arguments") (NCicPp.ppterm ~metasenv ~subst ~context he)
            (List.length args) (List.length args_so_far))))
      | ty ->
          let metasenv, subst, newhead, newheadty = 
            try_coercions ~localise ~look_for_coercion metasenv subst context
              (NCicUntrusted.mk_appl he (List.rev args_so_far)) orig_he ty
              (NCic.Prod ("_",NCic.Implicit `Term,NCic.Implicit `Term))
              false
              (RefineFailure (lazy (localise orig_he, Printf.sprintf
               ("The term %s is here applied to %d arguments but expects " ^^
               "only %d arguments") (NCicPp.ppterm ~metasenv ~subst ~context he)
               (List.length args) (List.length args_so_far))))
          in
          aux metasenv subst [] newhead newheadty (arg :: tl)
  in
   aux metasenv subst [] he ty_he args
  (*D*)in outside(); rc with exc -> outside (); raise exc
;;
(* vim:set foldmethod=marker: *)