open Printf

exception AssertFailure of string
exception MetaSubstFailure of string

exception RelToHiddenHypothesis (* TODO remove this exception *)

let debug_print = prerr_endline

type substitution = (int * Cic.term) list

let ppsubst subst =
  String.concat "\n"
    (List.map
      (fun (idx, term) -> Printf.sprintf "?%d := %s" idx (CicPp.ppterm term))
      subst)

(**** DELIFT ****)
(* the delift function takes in input a metavariable index, an ordered list of
 * optional terms [t1,...,tn] and a term t, and substitutes every tk = Some
 * (rel(nk)) with rel(k).  Typically, the list of optional terms is the explicit
 * substitution that is applied to a metavariable occurrence and the result of
 * the delift function is a term the implicit variable can be substituted with
 * to make the term [t] unifiable with the metavariable occurrence.  In general,
 * the problem is undecidable if we consider equivalence in place of alpha
 * convertibility. Our implementation, though, is even weaker than alpha
 * convertibility, since it replace the term [tk] if and only if [tk] is a Rel
 * (missing all the other cases). Does this matter in practice?
 * The metavariable index is the index of the metavariable that must not occur
 * in the term (for occur check).
 *)

exception NotInTheList;;

let position n =
  let rec aux k =
   function 
       [] -> raise NotInTheList
     | (Some (Cic.Rel m))::_ when m=n -> k
     | _::tl -> aux (k+1) tl in
  aux 1
;;
 
(*CSC: this restriction function is utterly wrong, since it does not check  *)
(*CSC: that the variable that is going to be restricted does not occur free *)
(*CSC: in a part of the sequent that is not going to be restricted.         *)
(*CSC: In particular, the whole approach is wrong; if restriction can fail  *)
(*CSC: (as indeed it is the case), we can not collect all the restrictions  *)
(*CSC: and restrict everything at the end ;-(                               *)
let restrict to_be_restricted =
  let rec erase i n = 
    function
        [] -> []
      |        _::tl when List.mem (n,i) to_be_restricted ->
          None::(erase (i+1) n tl) 
      | he::tl -> he::(erase (i+1) n tl) in
  let rec aux =
    function 
        [] -> []
      |        (n,context,t)::tl -> (n,erase 1 n context,t)::(aux tl) in
  aux
;;

(*CSC: maybe we should rename delift in abstract, as I did in my dissertation *)
let delift n subst context metasenv l t =
 let module S = CicSubstitution in
  let to_be_restricted = ref [] in
  let rec deliftaux k =
   let module C = Cic in
    function
       C.Rel m -> 
         if m <=k then
          C.Rel m   (*CSC: che succede se c'e' un Def? Dovrebbe averlo gia' *)
                    (*CSC: deliftato la regola per il LetIn                 *)
                    (*CSC: FALSO! La regola per il LetIn non lo fa          *)
         else
          (match List.nth context (m-k-1) with
            Some (_,C.Def (t,_)) ->
             (*CSC: Hmmm. This bit of reduction is not in the spirit of    *)
             (*CSC: first order unification. Does it help or does it harm? *)
             deliftaux k (S.lift m t)
          | Some (_,C.Decl t) ->
             (*CSC: The following check seems to be wrong!             *)
             (*CSC: B:Set |- ?2 : Set                                  *)
             (*CSC: A:Set ; x:?2[A/B] |- ?1[x/A] =?= x                 *)
             (*CSC: Why should I restrict ?2 over B? The instantiation *)
             (*CSC: ?1 := A is perfectly reasonable and well-typed.    *)
             (*CSC: Thus I comment out the following two lines that    *)
             (*CSC: are the incriminated ones.                         *)
             (*(* It may augment to_be_restricted *)
               ignore (deliftaux k (S.lift m t)) ;*)
             (*CSC: end of bug commented out                           *)
             C.Rel ((position (m-k) l) + k)
          | None -> raise RelToHiddenHypothesis)
     | C.Var (uri,exp_named_subst) ->
        let exp_named_subst' =
         List.map (function (uri,t) -> uri,deliftaux k t) exp_named_subst
        in
         C.Var (uri,exp_named_subst')
     | C.Meta (i, l1) as t -> 
        if i = n then
          raise (MetaSubstFailure (sprintf
            "Cannot unify the metavariable ?%d with a term that has as subterm %s in which the same metavariable occurs (occur check)"
            i (CicPp.ppterm t)))
        else
          (try
            deliftaux k (S.lift_meta l (List.assoc i subst))
          with Not_found ->
            let rec deliftl j =
             function
                [] -> []
              | None::tl -> None::(deliftl (j+1) tl)
              | (Some t)::tl ->
                 let l1' = (deliftl (j+1) tl) in
                  try
                   Some (deliftaux k t)::l1'
                  with
                     RelToHiddenHypothesis
                   | NotInTheList
                   | MetaSubstFailure _ ->
                      to_be_restricted := (i,j)::!to_be_restricted ; None::l1'
            in
             let l' = deliftl 1 l1 in
              C.Meta(i,l'))
     | C.Sort _ as t -> t
     | C.Implicit as t -> t
     | C.Cast (te,ty) -> C.Cast (deliftaux k te, deliftaux k ty)
     | C.Prod (n,s,t) -> C.Prod (n, deliftaux k s, deliftaux (k+1) t)
     | C.Lambda (n,s,t) -> C.Lambda (n, deliftaux k s, deliftaux (k+1) t)
     | C.LetIn (n,s,t) -> C.LetIn (n, deliftaux k s, deliftaux (k+1) t)
     | C.Appl l -> C.Appl (List.map (deliftaux k) l)
     | C.Const (uri,exp_named_subst) ->
        let exp_named_subst' =
         List.map (function (uri,t) -> uri,deliftaux k t) exp_named_subst
        in
         C.Const (uri,exp_named_subst')
     | C.MutInd (uri,typeno,exp_named_subst) ->
        let exp_named_subst' =
         List.map (function (uri,t) -> uri,deliftaux k t) exp_named_subst
        in
         C.MutInd (uri,typeno,exp_named_subst')
     | C.MutConstruct (uri,typeno,consno,exp_named_subst) ->
        let exp_named_subst' =
         List.map (function (uri,t) -> uri,deliftaux k t) exp_named_subst
        in
         C.MutConstruct (uri,typeno,consno,exp_named_subst')
     | C.MutCase (sp,i,outty,t,pl) ->
        C.MutCase (sp, i, deliftaux k outty, deliftaux k t,
         List.map (deliftaux k) pl)
     | C.Fix (i, fl) ->
        let len = List.length fl in
        let liftedfl =
         List.map
          (fun (name, i, ty, bo) ->
           (name, i, deliftaux k ty, deliftaux (k+len) bo))
           fl
        in
         C.Fix (i, liftedfl)
     | C.CoFix (i, fl) ->
        let len = List.length fl in
        let liftedfl =
         List.map
          (fun (name, ty, bo) -> (name, deliftaux k ty, deliftaux (k+len) bo))
           fl
        in
         C.CoFix (i, liftedfl)
  in
   let res =
    try
     deliftaux 0 t
    with
     NotInTheList ->
      (* This is the case where we fail even first order unification. *)
      (* The reason is that our delift function is weaker than first  *)
      (* order (in the sense of alpha-conversion). See comment above  *)
      (* related to the delift function.                              *)
debug_print "!!!!!!!!!!! First Order UnificationFailure, but maybe it could have been successful even in a first order setting (no conversion, only alpha convertibility)! Please, implement a better delift function !!!!!!!!!!!!!!!!" ;
      raise (MetaSubstFailure (sprintf
        "Error trying to abstract %s over [%s]: the algorithm only tried to abstract over bound variables"
        (CicPp.ppterm t)
        (String.concat "; "
          (List.map
            (function Some t -> CicPp.ppterm t | None -> "_")
            l))))
   in
    res, restrict !to_be_restricted metasenv
;;

(**** END OF DELIFT ****)

let rec unwind metasenv subst unwinded t =
 let unwinded = ref unwinded in
 let frozen = ref [] in
 let rec um_aux metasenv =
  let module C = Cic in
  let module S = CicSubstitution in 
   function
      C.Rel _ as t -> t,metasenv
    | C.Var _  as t -> t,metasenv
    | C.Meta (i,l) -> 
        (try
          S.lift_meta l (List.assoc i !unwinded), metasenv
         with Not_found ->
           if List.mem i !frozen then
             raise (MetaSubstFailure
              "Failed to unify due to cyclic constraints (occur check)")
           else
            let saved_frozen = !frozen in 
            frozen := i::!frozen ;
            let res =
             try
              let t = List.assoc i subst in
              let t',metasenv' = um_aux metasenv t in
              let _,metasenv'' =
               let (_,canonical_context,_) = CicUtil.lookup_meta i metasenv in
                delift i subst canonical_context metasenv' l t'
              in
               unwinded := (i,t')::!unwinded ;
               S.lift_meta l t', metasenv'
             with Not_found ->
               (* not constrained variable, i.e. free in subst*)
               let l',metasenv' =
                List.fold_right
                 (fun t (tl,metasenv) ->
                   match t with
                      None -> None::tl,metasenv
                    | Some t -> 
                       let t',metasenv' = um_aux metasenv t in
                        (Some t')::tl, metasenv'
                 ) l ([],metasenv)
               in
                C.Meta (i,l'), metasenv'
            in
            frozen := saved_frozen ;
            res
        ) 
    | C.Sort _
    | C.Implicit as t -> t,metasenv
    | C.Cast (te,ty) ->
       let te',metasenv' = um_aux metasenv te in
       let ty',metasenv'' = um_aux metasenv' ty in
       C.Cast (te',ty'),metasenv''
    | C.Prod (n,s,t) ->
       let s',metasenv' = um_aux metasenv s in
       let t',metasenv'' = um_aux metasenv' t in
       C.Prod (n, s', t'), metasenv''
    | C.Lambda (n,s,t) ->
       let s',metasenv' = um_aux metasenv s in
       let t',metasenv'' = um_aux metasenv' t in
       C.Lambda (n, s', t'), metasenv''
    | C.LetIn (n,s,t) ->
       let s',metasenv' = um_aux metasenv s in
       let t',metasenv'' = um_aux metasenv' t in
       C.LetIn (n, s', t'), metasenv''
    | C.Appl (he::tl) ->
       let tl',metasenv' =
        List.fold_right
         (fun t (tl,metasenv) ->
           let t',metasenv' = um_aux metasenv t in
            t'::tl, metasenv'
         ) tl ([],metasenv)
       in
        begin
         match um_aux metasenv' he with
            (C.Appl l, metasenv'') -> C.Appl (l@tl'),metasenv''
          | (he', metasenv'') -> C.Appl (he'::tl'),metasenv''
        end
    | C.Appl _ -> assert false
    | C.Const (uri,exp_named_subst) ->
       let exp_named_subst', metasenv' =
        List.fold_right
         (fun (uri,t) (tl,metasenv) ->
           let t',metasenv' = um_aux metasenv t in
            (uri,t')::tl, metasenv'
         ) exp_named_subst ([],metasenv)
       in
        C.Const (uri,exp_named_subst'),metasenv'
    | C.MutInd (uri,typeno,exp_named_subst) ->
       let exp_named_subst', metasenv' =
        List.fold_right
         (fun (uri,t) (tl,metasenv) ->
           let t',metasenv' = um_aux metasenv t in
            (uri,t')::tl, metasenv'
         ) exp_named_subst ([],metasenv)
       in
        C.MutInd (uri,typeno,exp_named_subst'),metasenv'
    | C.MutConstruct (uri,typeno,consno,exp_named_subst) ->
       let exp_named_subst', metasenv' =
        List.fold_right
         (fun (uri,t) (tl,metasenv) ->
           let t',metasenv' = um_aux metasenv t in
            (uri,t')::tl, metasenv'
         ) exp_named_subst ([],metasenv)
       in
        C.MutConstruct (uri,typeno,consno,exp_named_subst'),metasenv'
    | C.MutCase (sp,i,outty,t,pl) ->
       let outty',metasenv' = um_aux metasenv outty in
       let t',metasenv'' = um_aux metasenv' t in
       let pl',metasenv''' =
        List.fold_right
         (fun p (pl,metasenv) ->
           let p',metasenv' = um_aux metasenv p in
            p'::pl, metasenv'
         ) pl ([],metasenv'')
       in
        C.MutCase (sp, i, outty', t', pl'),metasenv'''
    | C.Fix (i, fl) ->
       let len = List.length fl in
       let liftedfl,metasenv' =
        List.fold_right
         (fun (name, i, ty, bo) (fl,metasenv) ->
           let ty',metasenv' = um_aux metasenv ty in
           let bo',metasenv'' = um_aux metasenv' bo in
            (name, i, ty', bo')::fl,metasenv''
         ) fl ([],metasenv)
       in
        C.Fix (i, liftedfl),metasenv'
    | C.CoFix (i, fl) ->
       let len = List.length fl in
       let liftedfl,metasenv' =
        List.fold_right
         (fun (name, ty, bo) (fl,metasenv) ->
           let ty',metasenv' = um_aux metasenv ty in
           let bo',metasenv'' = um_aux metasenv' bo in
            (name, ty', bo')::fl,metasenv''
         ) fl ([],metasenv)
       in
        C.CoFix (i, liftedfl),metasenv'
 in
  let t',metasenv' = um_aux metasenv t in
   t',metasenv',!unwinded 

let apply_subst subst t = 
 (* metasenv will not be used nor modified. So, let's use a dummy empty one *)
 let metasenv = [] in
  let (t',_,_) = unwind metasenv [] subst t in
   t'

(* apply_subst_reducing subst (Some (mtr,reductions_no)) t              *)
(* performs as (apply_subst subst t) until it finds an application of   *)
(* (META [meta_to_reduce]) that, once unwinding is performed, creates   *)
(* a new beta-redex; in this case up to [reductions_no] consecutive     *)
(* beta-reductions are performed.                                       *)
(* Hint: this function is usually called when [reductions_no]           *)
(*  eta-expansions have been performed and the head of the new          *)
(*  application has been unified with (META [meta_to_reduce]):          *)
(*  during the unwinding the eta-expansions are undone.                 *)

let rec apply_subst_context subst =
    List.map (function
      | Some (n, Cic.Decl t) -> Some (n, Cic.Decl (apply_subst subst t))
      | Some (n, Cic.Def (t, ty)) ->
          let ty' =
            match ty with
            | None -> None
            | Some ty -> Some (apply_subst subst ty)
          in
          Some (n, Cic.Def (apply_subst subst t, ty'))
      | None -> None)

let rec apply_subst_reducing subst meta_to_reduce t =
 let rec um_aux =
  let module C = Cic in
  let module S = CicSubstitution in 
   function
      C.Rel _
    | C.Var _  as t -> t
    | C.Meta (i,l) as t ->
       (try
         S.lift_meta l (List.assoc i subst)
        with Not_found ->
          C.Meta (i,l))
    | C.Sort _ as t -> t
    | C.Implicit as t -> t
    | C.Cast (te,ty) -> C.Cast (um_aux te, um_aux ty)
    | C.Prod (n,s,t) -> C.Prod (n, um_aux s, um_aux t)
    | C.Lambda (n,s,t) -> C.Lambda (n, um_aux s, um_aux t)
    | C.LetIn (n,s,t) -> C.LetIn (n, um_aux s, um_aux t)
    | C.Appl (he::tl) ->
       let tl' = List.map um_aux tl in
        let t' =
         match um_aux he with
            C.Appl l -> C.Appl (l@tl')
          | _ as he' -> C.Appl (he'::tl')
        in
         begin
          match meta_to_reduce,he with
             Some (mtr,reductions_no), C.Meta (m,_) when m = mtr ->
              let rec beta_reduce =
               function
                  (n,(C.Appl (C.Lambda (_,_,t)::he'::tl'))) when n > 0 ->
                    let he'' = CicSubstitution.subst he' t in
                     if tl' = [] then
                      he''
                     else
                      beta_reduce (n-1,C.Appl(he''::tl'))
                | (_,t) -> t
              in
               beta_reduce (reductions_no,t')
           | _,_ -> t'
         end
    | C.Appl _ -> assert false
    | C.Const (uri,exp_named_subst) ->
       let exp_named_subst' =
        List.map (function (uri,t) -> (uri,um_aux t)) exp_named_subst
       in
        C.Const (uri,exp_named_subst')
    | C.MutInd (uri,typeno,exp_named_subst) ->
       let exp_named_subst' =
        List.map (function (uri,t) -> (uri,um_aux t)) exp_named_subst
       in
        C.MutInd (uri,typeno,exp_named_subst')
    | C.MutConstruct (uri,typeno,consno,exp_named_subst) ->
       let exp_named_subst' =
        List.map (function (uri,t) -> (uri,um_aux t)) exp_named_subst
       in
        C.MutConstruct (uri,typeno,consno,exp_named_subst')
    | C.MutCase (sp,i,outty,t,pl) ->
       C.MutCase (sp, i, um_aux outty, um_aux t,
        List.map um_aux pl)
    | C.Fix (i, fl) ->
       let len = List.length fl in
       let liftedfl =
        List.map
         (fun (name, i, ty, bo) -> (name, i, um_aux ty, um_aux bo))
          fl
       in
        C.Fix (i, liftedfl)
    | C.CoFix (i, fl) ->
       let len = List.length fl in
       let liftedfl =
        List.map
         (fun (name, ty, bo) -> (name, um_aux ty, um_aux bo))
          fl
       in
        C.CoFix (i, liftedfl)
 in
   um_aux t

let ppcontext ?(sep = "\n") subst context =
  String.concat sep
    (List.rev_map (function
      | Some (n, Cic.Decl t) ->
          sprintf "%s : %s"
            (CicPp.ppname n) (CicPp.ppterm (apply_subst subst t))
      | Some (n, Cic.Def (t, ty)) ->
          sprintf "%s : %s := %s"
            (CicPp.ppname n)
            (match ty with
            | None -> "_"
            | Some ty -> CicPp.ppterm (apply_subst subst ty))
            (CicPp.ppterm (apply_subst subst t))
      | None -> "_")
      context)

let ppmetasenv ?(sep = "\n") subst metasenv =
  String.concat sep
    (List.map
      (fun (i, c, t) ->
        sprintf "%s |- ?%d: %s" (ppcontext ~sep:"; " subst c) i
          (CicPp.ppterm (apply_subst subst t)))
      (List.filter
        (fun (i, _, _) -> not (List.exists (fun (j, _) -> (j = i)) subst))
        metasenv))

(* UNWIND THE MGU INSIDE THE MGU *)
let unwind_subst metasenv subst =
  List.fold_left
   (fun (unwinded,metasenv) (i,_) ->
     let (_,canonical_context,_) = CicUtil.lookup_meta i metasenv in
     let identity_relocation_list =
      CicMkImplicit.identity_relocation_list_for_metavariable canonical_context
     in
      let (_,metasenv',subst') =
       unwind metasenv subst unwinded (Cic.Meta (i,identity_relocation_list))
      in
       subst',metasenv'
   ) ([],metasenv) subst

(* From now on we recreate a kernel abstraction where substitutions are part of
 * the calculus *)

let whd metasenv subst context term =
  (* TODO unwind's result is thrown away :-( *)
  let (subst, _) = unwind_subst metasenv subst in
  let term = apply_subst subst term in
  let context = apply_subst_context subst context in
  try
    CicReduction.whd context term
  with e ->
    raise (MetaSubstFailure ("Weak head reduction failure: " ^
      Printexc.to_string e))

let are_convertible metasenv subst context t1 t2 =
  (* TODO unwind's result is thrown away :-( *)
  let (subst, _) = unwind_subst metasenv subst in
  let context = apply_subst_context subst context in
  let (t1, t2) = (apply_subst subst t1, apply_subst subst t2) in
  CicReduction.are_convertible context t1 t2

let type_of_aux' metasenv subst context term =
  (* TODO unwind's result is thrown away :-( *)
  let (subst, _) = unwind_subst metasenv subst in
  let term = apply_subst subst term in
  let context = apply_subst_context subst context in
  let metasenv =
    List.map
      (fun (i, c, t) -> (i, apply_subst_context subst c, apply_subst subst t))
      (List.filter
        (fun (i, _, _) -> not (List.exists (fun (j, _) -> (j = i)) subst))
        metasenv)
  in
  try
    CicTypeChecker.type_of_aux' metasenv context term
  with CicTypeChecker.TypeCheckerFailure msg ->
    raise (MetaSubstFailure ("Type checker failure: " ^ msg))