open Printf

exception AssertFailure of string
exception MetaSubstFailure of string

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
;;

exception Occur;;

let rec force_does_not_occur subst to_be_restricted t =
 let module C = Cic in
 let more_to_be_restricted = ref [] in
 let rec aux k = function
      C.Rel r when List.mem (r+k) to_be_restricted -> raise Occur
    | C.Rel _
    | C.Sort _ as t -> t
    | C.Implicit -> assert false
    | C.Meta (n, l) ->
       (* we do not retrieve the term associated to ?n in subst since *)
       (* in this way we can restrict if something goes wrong         *)
       let l' =
         let i = ref 0 in
         List.map
           (function t ->
             incr i ;
             match t with
                None -> None
              | Some t ->
                 try
                   Some (aux k t)
                 with Occur ->
                   more_to_be_restricted := (n,!i) :: !more_to_be_restricted;
                   None)
           l
       in
        C.Meta (n, l')
    | C.Cast (te,ty) -> C.Cast (aux k te, aux k ty)
    | C.Prod (name,so,dest) -> C.Prod (name, aux k so, aux (k+1) dest)
    | C.Lambda (name,so,dest) -> C.Lambda (name, aux k so, aux (k+1) dest)
    | C.LetIn (name,so,dest) -> C.LetIn (name, aux k so, aux (k+1) dest)
    | C.Appl l -> C.Appl (List.map (aux k) l)
    | C.Var (uri,exp_named_subst) ->
        let exp_named_subst' =
          List.map (fun (uri,t) -> (uri, aux k t)) exp_named_subst
        in
        C.Var (uri, exp_named_subst')
    | C.Const (uri, exp_named_subst) ->
        let exp_named_subst' =
          List.map (fun (uri,t) -> (uri, aux k t)) exp_named_subst
        in
        C.Const (uri, exp_named_subst')
    | C.MutInd (uri,tyno,exp_named_subst) ->
        let exp_named_subst' =
          List.map (fun (uri,t) -> (uri, aux k t)) exp_named_subst
        in
        C.MutInd (uri, tyno, exp_named_subst')
    | C.MutConstruct (uri,tyno,consno,exp_named_subst) ->
        let exp_named_subst' =
          List.map (fun (uri,t) -> (uri, aux k t)) exp_named_subst
        in
        C.MutConstruct (uri, tyno, consno, exp_named_subst')
    | C.MutCase (uri,tyno,out,te,pl) ->
        C.MutCase (uri, tyno, aux k out, aux k te, List.map (aux k) pl)
    | C.Fix (i,fl) ->
       let len = List.length fl in
       let k_plus_len = k + len in
       let fl' =
         List.map
          (fun (name,j,ty,bo) -> (name, j, aux k ty, aux k_plus_len bo)) fl
       in
       C.Fix (i, fl')
    | C.CoFix (i,fl) ->
       let len = List.length fl in
       let k_plus_len = k + len in
       let fl' =
         List.map
          (fun (name,ty,bo) -> (name, aux k ty, aux k_plus_len bo)) fl
       in
       C.CoFix (i, fl')
 in
 let res = aux 0 t in
 (!more_to_be_restricted, res)
 
let rec restrict subst to_be_restricted metasenv =
  let names_of_context_indexes context indexes =
    String.concat ", "
      (List.map
        (fun i ->
          match List.nth context i with
          | None -> assert false
          | Some (n, _) -> CicPp.ppname n)
        indexes)
  in
  let force_does_not_occur_in_context to_be_restricted = function
    | None -> [], None
    | Some (name, Cic.Decl t) ->
        let (more_to_be_restricted, t') =
          force_does_not_occur subst to_be_restricted t
        in
        more_to_be_restricted, Some (name, Cic.Decl t)
    | Some (name, Cic.Def (bo, ty)) ->
        let (more_to_be_restricted, bo') =
          force_does_not_occur subst to_be_restricted bo
        in
        let more_to_be_restricted, ty' =
          match ty with
          | None ->  more_to_be_restricted, None
          | Some ty ->
              let more_to_be_restricted', ty' =
                force_does_not_occur subst to_be_restricted ty
              in
              more_to_be_restricted @ more_to_be_restricted',
              Some ty'
        in
        more_to_be_restricted, Some (name, Cic.Def (bo', ty'))
  in
  let rec erase i to_be_restricted n = function
    | [] -> [], to_be_restricted, []
    | hd::tl ->
        let restrict_me = List.mem i to_be_restricted in
        if restrict_me then
          let more_to_be_restricted, restricted, new_tl =
            erase (i+1) (i :: to_be_restricted) n tl
          in
          more_to_be_restricted, restricted, None :: new_tl
        else
          (try
            let more_to_be_restricted, hd' =
              force_does_not_occur_in_context to_be_restricted hd
            in
            let more_to_be_restricted', restricted, tl' =
              erase (i+1) to_be_restricted n tl
            in
            more_to_be_restricted @ more_to_be_restricted',
            restricted, hd' :: tl'
          with Occur ->
            let more_to_be_restricted, restricted, tl' =
              erase (i+1) (i :: to_be_restricted) n tl
            in
            more_to_be_restricted, restricted, None :: tl')
  in
  let (more_to_be_restricted, metasenv, subst) =
    List.fold_right
      (fun (n, context, t) (more, metasenv, subst) ->
        let to_be_restricted =
          List.map snd (List.filter (fun (m, _) -> m = n) to_be_restricted)
        in
        let (more_to_be_restricted, restricted, context') =
          erase 1 to_be_restricted n context
        in
        try
          let more_to_be_restricted', t' =
            force_does_not_occur subst restricted t
          in
          let metasenv' = (n, context', t') :: metasenv in
          (try
            let s = List.assoc n subst in
            try
              let more_to_be_restricted'', s' =
                force_does_not_occur subst restricted s
              in
              let subst' = (n, s') :: (List.remove_assoc n subst) in
              let more =
                more @ more_to_be_restricted @ more_to_be_restricted' @
                  more_to_be_restricted''
              in
              (more, metasenv', subst')
            with Occur ->
              raise (MetaSubstFailure (sprintf
                "Cannot restrict the context of the metavariable ?%d over the hypotheses %s since ?%d is already instantiated with %s and at least one of the hypotheses occurs in the substituted term"
                n (names_of_context_indexes context to_be_restricted) n
                (CicPp.ppterm s)))
           with Not_found -> (more @ more_to_be_restricted @ more_to_be_restricted', metasenv', subst))
        with Occur ->
          raise (MetaSubstFailure (sprintf
            "Cannot restrict the context of the metavariable ?%d over the hypotheses %s since metavariable's type depends on at least one of them"
            n (names_of_context_indexes context to_be_restricted))))
      metasenv ([], [], subst)
  in
  match more_to_be_restricted with
  | [] -> (metasenv, subst)
  | _ -> restrict subst more_to_be_restricted metasenv
;;

(*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 (MetaSubstFailure "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
         (* I do not consider the term associated to ?i in subst since *)
         (* in this way I can restrict if something goes wrong.        *)
          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
                   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
   let (metasenv, subst) = restrict subst !to_be_restricted metasenv in
    res, metasenv, subst
;;

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

let apply_subst_gen ~appl_fun subst term =
 let rec um_aux =
  let module C = Cic in
  let module S = CicSubstitution in 
   function
      C.Rel _ as t -> t
    | C.Var _  as t -> t
    | C.Meta (i, l) -> 
        (try
          let t = List.assoc i subst in
          um_aux (S.lift_meta l t)
        with Not_found -> (* not constrained variable, i.e. free in subst*)
          let l' =
            List.map (function None -> None | Some t -> Some (um_aux t)) l
          in
           C.Meta (i,l'))
    | C.Sort _ as t -> t
    | C.Implicit -> assert false
    | 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 (hd :: tl) -> appl_fun um_aux hd tl
    | C.Appl _ -> assert false
    | C.Const (uri,exp_named_subst) ->
       let exp_named_subst' =
         List.map (fun (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 (fun (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 (fun (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) ->
       let pl' = List.map um_aux pl in
       C.MutCase (sp, i, um_aux outty, um_aux t, pl')
    | C.Fix (i, fl) ->
       let fl' =
         List.map (fun (name, i, ty, bo) -> (name, i, um_aux ty, um_aux bo)) fl
       in
       C.Fix (i, fl')
    | C.CoFix (i, fl) ->
       let fl' =
         List.map (fun (name, ty, bo) -> (name, um_aux ty, um_aux bo)) fl
       in
       C.CoFix (i, fl')
 in
 um_aux term

let apply_subst =
  let appl_fun um_aux he tl =
    let tl' = List.map um_aux tl in
      begin
       match um_aux he with
          Cic.Appl l -> Cic.Appl (l@tl')
        | he' -> Cic.Appl (he'::tl')
      end
  in
  apply_subst_gen ~appl_fun

let ppterm subst term = CicPp.ppterm (apply_subst subst term)

(* 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 apply_subst_reducing meta_to_reduce =
  let appl_fun um_aux he tl =
    let tl' = List.map um_aux tl in
    let t' =
     match um_aux he with
        Cic.Appl l -> Cic.Appl (l@tl')
      | he' -> Cic.Appl (he'::tl')
    in
     begin
      match meta_to_reduce, he with
         Some (mtr,reductions_no), Cic.Meta (m,_) when m = mtr ->
          let rec beta_reduce =
           function
              (n,(Cic.Appl (Cic.Lambda (_,_,t)::he'::tl'))) when n > 0 ->
                let he'' = CicSubstitution.subst he' t in
                 if tl' = [] then
                  he''
                 else
                  beta_reduce (n-1,Cic.Appl(he''::tl'))
            | (_,t) -> t
          in
           beta_reduce (reductions_no,t')
       | _,_ -> t'
     end
  in
  apply_subst_gen ~appl_fun

let rec apply_subst_context subst context =
  List.fold_right
    (fun item context ->
      match item with
      | Some (n, Cic.Decl t) ->
          let t' = apply_subst subst t in
          Some (n, Cic.Decl t') :: context
      | Some (n, Cic.Def (t, ty)) ->
          let ty' =
            match ty with
            | None -> None
            | Some ty -> Some (apply_subst subst ty)
          in
          let t' = apply_subst subst t in
          Some (n, Cic.Def (t', ty')) :: context
      | None -> None :: context)
    context []

let apply_subst_metasenv subst metasenv =
  List.map
    (fun (n, context, ty) ->
      (n, apply_subst_context subst context, apply_subst subst ty))
    (List.filter
      (fun (i, _, _) -> not (List.exists (fun (j, _) -> (j = i)) subst))
      metasenv)

let ppterm subst term = CicPp.ppterm (apply_subst subst term)

let ppterm_in_context subst term name_context =
 CicPp.pp (apply_subst subst term) name_context

let ppcontext' ?(sep = "\n") subst context =
 let separate s = if s = "" then "" else s ^ sep in
  List.fold_right 
   (fun context_entry (i,name_context) ->
     match context_entry with
        Some (n,Cic.Decl t) ->
         sprintf "%s%s : %s" (separate i) (CicPp.ppname n)
          (ppterm_in_context subst t name_context), (Some n)::name_context
      | Some (n,Cic.Def (bo,ty)) ->
         sprintf "%s%s : %s := %s" (separate i) (CicPp.ppname n)
          (match ty with
              None -> "_"
            | Some ty -> ppterm_in_context subst ty name_context)
          (ppterm_in_context subst bo name_context), (Some n)::name_context
       | None ->
          sprintf "%s_ :? _" (separate i), None::name_context
    ) context ("",[])

let ppcontext ?sep subst context = fst (ppcontext' ?sep subst context)

let ppmetasenv ?(sep = "\n") metasenv subst =
  String.concat sep
    (List.map
      (fun (i, c, t) ->
        let context,name_context = ppcontext' ~sep:"; " subst c in
         sprintf "%s |- ?%d: %s" context i
          (ppterm_in_context subst t name_context))
      (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 lift subst n term =
  let term = apply_subst subst term in
  try
    CicSubstitution.lift n term
  with e ->
    raise (MetaSubstFailure ("Lift failure: " ^ Printexc.to_string e))

let subst subst t1 t2 =
  let t1 = apply_subst subst t1 in
  let t2 = apply_subst subst t2 in
  try
    CicSubstitution.subst t1 t2
  with e ->
    raise (MetaSubstFailure ("Subst failure: " ^ Printexc.to_string e))

let whd subst context term =
  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 subst context t1 t2 =
  let context = apply_subst_context subst context in
  let t1 = apply_subst subst t1 in
  let t2 = apply_subst subst t2 in
  CicReduction.are_convertible context t1 t2

let type_of_aux' metasenv subst context term =
  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))