(* Copyright (C) 2000, 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://cs.unibo.it/helm/.
 *)

open Printf

exception UnificationFailure of string;;
exception Uncertain of string;;
exception AssertFailure of string;;

let debug_print = prerr_endline

let type_of_aux' metasenv subst context term =
  try
    CicMetaSubst.type_of_aux' metasenv subst context term
  with
  | CicMetaSubst.MetaSubstFailure msg ->
    raise (AssertFailure
      ((sprintf
        "Type checking error: %s in context\n%s\nand metasenv\n%s.\nException: %s.\nBroken invariant: unification must be invoked only on well typed terms"
        (CicMetaSubst.ppterm subst term)
        (CicMetaSubst.ppcontext subst context)
        (CicMetaSubst.ppmetasenv metasenv subst) msg)))

let rec beta_expand test_equality_only metasenv subst context t arg =
 let module S = CicSubstitution in
 let module C = Cic in
  let rec aux metasenv subst n context t' =
   try
    let subst,metasenv =
     fo_unif_subst test_equality_only subst context metasenv arg t'
    in
     subst,metasenv,C.Rel (1 + n)
   with
      Uncertain _
    | UnificationFailure _ ->
       match t' with
        | C.Rel m  -> subst,metasenv, if m <= n then C.Rel m else C.Rel (m+1)
        | C.Var (uri,exp_named_subst) ->
           let subst,metasenv,exp_named_subst' =
            aux_exp_named_subst metasenv subst n context exp_named_subst
           in
            subst,metasenv,C.Var (uri,exp_named_subst')
        | C.Meta (i,l) as t->
           (try
             let (_, t') = CicMetaSubst.lookup_subst i subst in
              aux metasenv subst n context (CicSubstitution.lift_meta l t')
            with
             CicMetaSubst.SubstNotFound _ -> subst,metasenv,t)
        | C.Sort _
        | C.Implicit _ as t -> subst,metasenv,t
        | C.Cast (te,ty) ->
           let subst,metasenv,te' = aux metasenv subst n context te in
           let subst,metasenv,ty' = aux metasenv subst n context ty in
           subst,metasenv,C.Cast (te', ty')
        | C.Prod (nn,s,t) ->
           let subst,metasenv,s' = aux metasenv subst n context s in
           let subst,metasenv,t' =
            aux metasenv subst (n+1) ((Some (nn, C.Decl s))::context) t
           in
            subst,metasenv,C.Prod (nn, s', t')
        | C.Lambda (nn,s,t) ->
           let subst,metasenv,s' = aux metasenv subst n context s in
           let subst,metasenv,t' =
            aux metasenv subst (n+1) ((Some (nn, C.Decl s))::context) t
           in
            subst,metasenv,C.Lambda (nn, s', t')
        | C.LetIn (nn,s,t) ->
           let subst,metasenv,s' = aux metasenv subst n context s in
           let subst,metasenv,t' =
            aux metasenv subst (n+1) ((Some (nn, C.Def (s,None)))::context) t
           in
            subst,metasenv,C.LetIn (nn, s', t')
        | C.Appl l ->
           let subst,metasenv,revl' =
            List.fold_left
             (fun (subst,metasenv,appl) t ->
               let subst,metasenv,t' = aux metasenv subst n context t in
                subst,metasenv,t'::appl
             ) (subst,metasenv,[]) l
           in
            subst,metasenv,C.Appl (List.rev revl')
        | C.Const (uri,exp_named_subst) ->
           let subst,metasenv,exp_named_subst' =
            aux_exp_named_subst metasenv subst n context exp_named_subst
           in
            subst,metasenv,C.Const (uri,exp_named_subst')
        | C.MutInd (uri,i,exp_named_subst) ->
           let subst,metasenv,exp_named_subst' =
            aux_exp_named_subst metasenv subst n context exp_named_subst
           in
            subst,metasenv,C.MutInd (uri,i,exp_named_subst')
        | C.MutConstruct (uri,i,j,exp_named_subst) ->
           let subst,metasenv,exp_named_subst' =
            aux_exp_named_subst metasenv subst n context exp_named_subst
           in
            subst,metasenv,C.MutConstruct (uri,i,j,exp_named_subst')
        | C.MutCase (sp,i,outt,t,pl) ->
           let subst,metasenv,outt' = aux metasenv subst n context outt in
           let subst,metasenv,t' = aux metasenv subst n context t in
           let subst,metasenv,revpl' =
            List.fold_left
             (fun (subst,metasenv,pl) t ->
               let subst,metasenv,t' = aux metasenv subst n context t in
                subst,metasenv,t'::pl
             ) (subst,metasenv,[]) pl
           in
           subst,metasenv,C.MutCase (sp,i,outt', t', List.rev revpl')
        | C.Fix (i,fl) ->
(*CSC: not implemented
           let tylen = List.length fl in
            let substitutedfl =
             List.map
              (fun (name,i,ty,bo) -> (name, i, aux n ty, aux (n+tylen) bo))
               fl
            in
             C.Fix (i, substitutedfl)
*) subst,metasenv,CicMetaSubst.lift subst 1 t'
        | C.CoFix (i,fl) ->
(*CSC: not implemented
           let tylen = List.length fl in
            let substitutedfl =
             List.map
              (fun (name,ty,bo) -> (name, aux n ty, aux (n+tylen) bo))
               fl
            in
             C.CoFix (i, substitutedfl)
*) subst,metasenv,CicMetaSubst.lift subst 1 t'

  and aux_exp_named_subst metasenv subst n context ens =
   List.fold_right
    (fun (uri,t) (subst,metasenv,l) ->
      let subst,metasenv,t' = aux metasenv subst n context t in
       subst,metasenv,(uri,t')::l) ens (subst,metasenv,[])
  in
  let argty = type_of_aux' metasenv subst context arg in
  let fresh_name =
   FreshNamesGenerator.mk_fresh_name
    metasenv context (Cic.Name "Heta") ~typ:argty
  in
   let subst,metasenv,t' = aux metasenv subst 0 context t in
    subst, metasenv, C.Appl [C.Lambda (fresh_name,argty,t') ; arg]

and beta_expand_many test_equality_only metasenv subst context t =
 List.fold_left
  (fun (subst,metasenv,t) arg ->
    beta_expand test_equality_only metasenv subst context t arg
  ) (subst,metasenv,t)

(* NUOVA UNIFICAZIONE *)
(* A substitution is a (int * Cic.term) list that associates a
   metavariable i with its body.
   A metaenv is a (int * Cic.term) list that associate a metavariable
   i with is type. 
   fo_unif_new takes a metasenv, a context, two terms t1 and t2 and gives back
   a new substitution which is _NOT_ unwinded. It must be unwinded before
   applying it. *)

and fo_unif_subst test_equality_only subst context metasenv t1 t2 =  
 let module C = Cic in
 let module R = CicMetaSubst in
 let module S = CicSubstitution in
  match (t1, t2) with
     (C.Meta (n,ln), C.Meta (m,lm)) when n=m ->
       let ok,subst,metasenv =
       try
        List.fold_left2
         (fun (b,subst,metasenv) t1 t2 ->
           if b then true,subst,metasenv else
            match t1,t2 with
               None,_
             | _,None -> true,subst,metasenv
             | Some t1', Some t2' ->
                (* First possibility:  restriction    *)
                (* Second possibility: unification    *)
                (* Third possibility:  convertibility *)
                if R.are_convertible subst context t1' t2' then
                 true,subst,metasenv
                else
                 (try
                   let subst,metasenv =
                    fo_unif_subst 
                     test_equality_only subst context metasenv t1' t2'
                   in
                    true,subst,metasenv
                 with
                  Not_found -> false,subst,metasenv)
         ) (true,subst,metasenv) ln lm
       with
        Invalid_argument _ ->
         raise (UnificationFailure (sprintf
           "Error trying to unify %s with %s: the lengths of the two local contexts do not match." (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
       in
        if ok then
          subst,metasenv
        else
          raise (UnificationFailure (sprintf
            "Error trying to unify %s with %s: the algorithm tried to check whether the two substitutions are convertible; if they are not, it tried to unify the two substitutions. No restriction was attempted."
            (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | (C.Meta (n,_), C.Meta (m,_)) when n>m ->
       fo_unif_subst test_equality_only subst context metasenv t2 t1
   | (C.Meta (n,l), t)   
   | (t, C.Meta (n,l)) ->
       let swap =
        match t1,t2 with
           C.Meta (n,_), C.Meta (m,_) when n < m -> false
         | _, C.Meta _ -> false
         | _,_ -> true
       in
       let lower = fun x y -> if swap then y else x in
       let upper = fun x y -> if swap then x else y in
       let fo_unif_subst_ordered 
        test_equality_only subst context metasenv m1 m2 =
          fo_unif_subst test_equality_only subst context metasenv 
           (lower m1 m2) (upper m1 m2)
       in
        begin
        try
          let (_, oldt) = CicMetaSubst.lookup_subst n subst in
          let lifted_oldt = S.lift_meta l oldt in
          let ty_lifted_oldt =
            type_of_aux' metasenv subst context lifted_oldt
          in
          let tyt = type_of_aux' metasenv subst context t in
          let (subst, metasenv) =
            fo_unif_subst_ordered test_equality_only subst context metasenv
              tyt ty_lifted_oldt
          in
          fo_unif_subst_ordered 
            test_equality_only subst context metasenv t lifted_oldt
        with CicMetaSubst.SubstNotFound _ ->
         (* First of all we unify the type of the meta with the type of the term *)
         let subst,metasenv =
          let (_,_,meta_type) =  CicUtil.lookup_meta n metasenv in
          (try
            let tyt = type_of_aux' metasenv subst context t in
             fo_unif_subst 
              test_equality_only 
               subst context metasenv tyt (S.lift_meta l meta_type)
          with AssertFailure _ ->
            (* TODO huge hack!!!!
             * we keep on unifying/refining in the hope that the problem will be
             * eventually solved. In the meantime we're breaking a big invariant:
             * the terms that we are unifying are no longer well typed in the
             * current context (in the worst case we could even diverge)
             *)
(*
prerr_endline "********* FROM NOW ON EVERY REASONABLE INVARIANT IS BROKEN.";
prerr_endline "********* PROCEED AT YOUR OWN RISK. AND GOOD LUCK." ;
*)
            (subst, metasenv))
         in
          let t',metasenv,subst =
           try
            CicMetaSubst.delift n subst context metasenv l t
           with
              (CicMetaSubst.MetaSubstFailure msg)-> raise(UnificationFailure msg)
            | (CicMetaSubst.Uncertain msg) -> raise (Uncertain msg)
          in
           let t'' =
            match t' with
               C.Sort (C.Type u) when not test_equality_only ->
                 let u' = CicUniv.fresh () in
                 let s = C.Sort (C.Type u') in
                  ignore (CicUniv.add_ge (upper u u') (lower u u')) ;
                 s
              | _ -> t'
           in
            (* Unifying the types may have already instantiated n. Let's check *)
            try
             let (_, oldt) = CicMetaSubst.lookup_subst n subst in
             let lifted_oldt = S.lift_meta l oldt in
              fo_unif_subst_ordered 
               test_equality_only subst context metasenv t lifted_oldt
            with
             CicMetaSubst.SubstNotFound _ -> 
               let (_, context, _) = CicUtil.lookup_meta n metasenv in
               let subst = (n, (context, t'')) :: subst in
               let metasenv =
(*                 CicMetaSubst.apply_subst_metasenv [n,(context, t'')] metasenv *)
                CicMetaSubst.apply_subst_metasenv subst metasenv
               in
               subst, metasenv
(*               (n,t'')::subst, metasenv *)
        end
   | (C.Var (uri1,exp_named_subst1),C.Var (uri2,exp_named_subst2))
   | (C.Const (uri1,exp_named_subst1),C.Const (uri2,exp_named_subst2)) ->
      if UriManager.eq uri1 uri2 then
       fo_unif_subst_exp_named_subst test_equality_only subst context metasenv
        exp_named_subst1 exp_named_subst2
      else
       raise (UnificationFailure (sprintf
        "Can't unify %s with %s due to different constants"
        (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | C.MutInd (uri1,i1,exp_named_subst1),C.MutInd (uri2,i2,exp_named_subst2) ->
      if UriManager.eq uri1 uri2 && i1 = i2 then
       fo_unif_subst_exp_named_subst test_equality_only subst context metasenv
        exp_named_subst1 exp_named_subst2
      else
       raise (UnificationFailure (sprintf
        "Can't unify %s with %s due to different inductive principles"
        (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | C.MutConstruct (uri1,i1,j1,exp_named_subst1),
     C.MutConstruct (uri2,i2,j2,exp_named_subst2) ->
      if UriManager.eq uri1 uri2 && i1 = i2 && j1 = j2 then
       fo_unif_subst_exp_named_subst test_equality_only subst context metasenv
        exp_named_subst1 exp_named_subst2
      else
       raise (UnificationFailure (sprintf
        "Can't unify %s with %s due to different inductive constructors"
        (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | (C.Implicit _, _) | (_, C.Implicit _) ->  assert false
   | (C.Cast (te,ty), t2) -> fo_unif_subst test_equality_only 
                              subst context metasenv te t2
   | (t1, C.Cast (te,ty)) -> fo_unif_subst test_equality_only 
                              subst context metasenv t1 te
   | (C.Prod (n1,s1,t1), C.Prod (_,s2,t2)) -> 
        (* TASSI: this is the only case in which we want == *)
       let subst',metasenv' = fo_unif_subst true 
                               subst context metasenv s1 s2 in
        fo_unif_subst test_equality_only 
	 subst' ((Some (n1,(C.Decl s1)))::context) metasenv' t1 t2
   | (C.Lambda (n1,s1,t1), C.Lambda (_,s2,t2)) -> 
        (* TASSI: ask someone a reason for not putting true here *)
       let subst',metasenv' = fo_unif_subst test_equality_only 
                               subst context metasenv s1 s2 in
        fo_unif_subst test_equality_only 
	 subst' ((Some (n1,(C.Decl s1)))::context) metasenv' t1 t2
   | (C.LetIn (_,s1,t1), t2)  
   | (t2, C.LetIn (_,s1,t1)) -> 
       fo_unif_subst 
        test_equality_only subst context metasenv t2 (S.subst s1 t1)
   | (C.Appl l1, C.Appl l2) -> 
       let subst,metasenv,t1',t2' =
        match l1,l2 with
           C.Meta (i,_)::_, C.Meta (j,_)::_ when i = j ->
             subst,metasenv,t1,t2
         (* In the first two cases when we reach the next begin ... end
            section useless work is done since, by construction, the list
            of arguments will be equal.
         *)
         | C.Meta (i,l)::args, _ ->
            let subst,metasenv,t2' =
             beta_expand_many test_equality_only metasenv subst context t2 args
            in
             subst,metasenv,t1,t2'
         | _, C.Meta (i,l)::args ->
            let subst,metasenv,t1' =
             beta_expand_many test_equality_only metasenv subst context t1 args
            in
             subst,metasenv,t1',t2
         | _,_ ->
             subst,metasenv,t1,t2
       in
        begin
         match t1',t2' with
            C.Appl l1, C.Appl l2 ->
             let lr1 = List.rev l1 in
             let lr2 = List.rev l2 in
             let rec fo_unif_l test_equality_only subst metasenv =
              function
                 [],_
               | _,[] -> assert false
               | ([h1],[h2]) ->
                   fo_unif_subst test_equality_only subst context metasenv h1 h2
               | ([h],l) 
               | (l,[h]) ->
                  fo_unif_subst 
	           test_equality_only subst context metasenv h (C.Appl (List.rev l))
               | ((h1::l1),(h2::l2)) -> 
                  let subst', metasenv' = 
                   fo_unif_subst test_equality_only subst context metasenv h1 h2
                  in 
                   fo_unif_l test_equality_only subst' metasenv' (l1,l2)
             in
              fo_unif_l test_equality_only subst metasenv (lr1, lr2) 
          | _ -> assert false
        end
   | (C.MutCase (_,_,outt1,t1',pl1), C.MutCase (_,_,outt2,t2',pl2))->
       let subst', metasenv' = 
        fo_unif_subst test_equality_only subst context metasenv outt1 outt2 in
       let subst'',metasenv'' = 
        fo_unif_subst test_equality_only subst' context metasenv' t1' t2' in
       (try
         List.fold_left2 
          (function (subst,metasenv) ->
            fo_unif_subst test_equality_only subst context metasenv
          ) (subst'',metasenv'') pl1 pl2 
        with
         Invalid_argument _ ->
          raise (UnificationFailure (sprintf
            "Error trying to unify %s with %s: the number of branches is not the same." (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2))))
   | (C.Rel _, _) | (_,  C.Rel _) ->
       if t1 = t2 then
         subst, metasenv
       else
        raise (UnificationFailure (sprintf
          "Can't unify %s with %s because they are not convertible"
          (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | (C.Sort _ ,_) | (_, C.Sort _)
   | (C.Const _, _) | (_, C.Const _)
   | (C.MutInd  _, _) | (_, C.MutInd _)
   | (C.MutConstruct _, _) | (_, C.MutConstruct _)
   | (C.Fix _, _) | (_, C.Fix _) 
   | (C.CoFix _, _) | (_, C.CoFix _) -> 
       if t1 = t2 || R.are_convertible subst context t1 t2 then
        subst, metasenv
       else
        raise (UnificationFailure (sprintf
          "Can't unify %s with %s because they are not convertible"
          (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))
   | (_,_) ->
       if R.are_convertible subst context t1 t2 then
        subst, metasenv
       else
        raise (UnificationFailure (sprintf
          "Can't unify %s with %s because they are not convertible"
          (CicMetaSubst.ppterm subst t1) (CicMetaSubst.ppterm subst t2)))

and fo_unif_subst_exp_named_subst test_equality_only subst context metasenv
 exp_named_subst1 exp_named_subst2
=
 try
  List.fold_left2
   (fun (subst,metasenv) (uri1,t1) (uri2,t2) ->
     assert (uri1=uri2) ;
     fo_unif_subst test_equality_only subst context metasenv t1 t2
   ) (subst,metasenv) exp_named_subst1 exp_named_subst2
 with
  Invalid_argument _ ->
   let print_ens ens =
    String.concat " ; "
     (List.map
       (fun (uri,t) ->
         UriManager.string_of_uri uri ^ " := " ^ (CicMetaSubst.ppterm subst t)
       ) ens) 
   in
    raise (UnificationFailure (sprintf
     "Error trying to unify the two explicit named substitutions (local contexts) %s and %s: their lengths is different." (print_ens exp_named_subst1) (print_ens exp_named_subst2)))

(* A substitution is a (int * Cic.term) list that associates a               *)
(* metavariable i with its body.                                             *)
(* metasenv is of type Cic.metasenv                                          *)
(* fo_unif takes a metasenv, a context, two terms t1 and t2 and gives back   *)
(* a new substitution which is already unwinded and ready to be applied and  *)
(* a new metasenv in which some hypothesis in the contexts of the            *)
(* metavariables may have been restricted.                                   *)
let fo_unif metasenv context t1 t2 = 
 fo_unif_subst false [] context metasenv t1 t2 ;;

let fo_unif_subst subst context metasenv t1 t2 =
  let enrich_msg msg =
    sprintf "Unification error unifying %s of type %s with %s of type %s in context\n%s\nand metasenv\n%s\nbecause %s"
      (CicMetaSubst.ppterm subst t1)
      (try
        CicPp.ppterm (type_of_aux' metasenv subst context t1)
      with _ -> "MALFORMED")
      (CicMetaSubst.ppterm subst t2)
      (try
        CicPp.ppterm (type_of_aux' metasenv subst context t2)
      with _ -> "MALFORMED")
      (CicMetaSubst.ppcontext subst context)
      (CicMetaSubst.ppmetasenv metasenv subst) msg
  in
  try
    fo_unif_subst false subst context metasenv t1 t2
  with
  | AssertFailure msg -> raise (AssertFailure (enrich_msg msg))
  | UnificationFailure msg -> raise (UnificationFailure (enrich_msg msg))
;;