module Ref = NReference

let cn_to_s = function
  | Cic.Anonymous -> "_"
  | Cic.Name s -> s
;;

type ctx = 
  | Ce of NCic.hypothesis 
  | Fix of Ref.reference * string * NCic.term

let splat mk_pi ctx t =
  List.fold_left
    (fun t c -> 
      match c with
      | Ce (name, NCic.Def (bo,ty)) -> NCic.LetIn (name, ty, bo, t)
      | Ce (name, NCic.Decl ty) when mk_pi -> NCic.Prod (name, ty, t)
      | Ce (name, NCic.Decl ty) -> NCic.Lambda (name, ty, t)
      | Fix (_,name,ty) when mk_pi -> NCic.Prod (name, ty, t)
      | Fix (_,name,ty) -> NCic.Lambda (name,ty,t))
    t ctx
;;

let context_tassonomy ctx = 
    let rec split inner acc acc1 = function 
      | Ce _ :: tl when inner -> split inner (acc+1) (acc1+1) tl
      | Fix _ ::tl -> split false acc (acc1+1) tl
      | _ as l -> acc, List.length l, acc1
    in
      split true 0 1 ctx
;;

let splat_args_for_rel ctx t = 
  let bound, free, primo_ce_dopo_fix = context_tassonomy ctx in
  if free = 0 then t 
  else
    let rec aux = function
      | 0 -> []
      | n -> 
         (match List.nth ctx (n+bound) with
         | Fix (refe, _, _) when (n+bound) < primo_ce_dopo_fix -> NCic.Const refe
         | Fix _ | Ce _ -> NCic.Rel (n+bound)) :: aux (n-1)
    in
    NCic.Appl (t:: aux free)
;;

let splat_args ctx t n_fix = 
  let bound, free, primo_ce_dopo_fix = context_tassonomy ctx in
  if ctx = [] then t
  else
   let rec aux = function
     | 0 -> []
     | n -> 
        (match List.nth ctx (n-1) with
         | Ce _ when n <= bound -> NCic.Rel n
         | Fix (refe, _, _) when n < primo_ce_dopo_fix ->
            splat_args_for_rel ctx (NCic.Const refe)
         | Fix _ | Ce _ -> NCic.Rel (n - n_fix)
        ) :: aux (n-1)
   in
   NCic.Appl (t:: aux (List.length ctx))
;;

exception Nothing_to_do;;

let fix_outty curi tyno t context outty =
 let leftno,rightno =
  match fst (CicEnvironment.get_obj CicUniv.oblivion_ugraph curi) with
     Cic.InductiveDefinition (tyl,_,leftno,_) ->
      let _,_,arity,_ = List.nth tyl tyno in
      let rec count_prods leftno context arity =
       match leftno, CicReduction.whd context arity with
          0, Cic.Sort _ -> 0
        | 0, Cic.Prod (name,so,ty) ->
           1 + count_prods 0 (Some (name, Cic.Decl so)::context) ty
        | n, Cic.Prod (name,so,ty) ->
           count_prods (leftno - 1) (Some (name, Cic.Decl so)::context) ty
        | _,_ -> assert false
      in
(*prerr_endline (UriManager.string_of_uri curi);
prerr_endline ("LEFTNO: " ^ string_of_int leftno ^ "  " ^ CicPp.ppterm arity);*)
       leftno, count_prods leftno [] arity
   | _ -> assert false in
 let ens,args =
  let tty,_= CicTypeChecker.type_of_aux' [] context t CicUniv.oblivion_ugraph in
  match CicReduction.whd context tty with
     Cic.MutInd (_,_,ens) -> ens,[]
   | Cic.Appl (Cic.MutInd (_,_,ens)::args) ->
      ens,fst (HExtlib.split_nth leftno args)
   | _ -> assert false
 in
  let rec aux n irl context outsort =
   match n, CicReduction.whd context outsort with
      0, Cic.Prod _ -> raise Nothing_to_do
    | 0, _ ->
       let irl = List.rev irl in
       let ty = CicSubstitution.lift rightno (Cic.MutInd (curi,tyno,ens)) in
       let ty =
        if args = [] && irl = [] then ty
        else
         Cic.Appl (ty::(List.map (CicSubstitution.lift rightno) args)@irl) in
       let he = CicSubstitution.lift (rightno + 1) outty in
       let t =
        if irl = [] then he
        else Cic.Appl (he::List.map (CicSubstitution.lift 1) irl)
       in
        Cic.Lambda (Cic.Anonymous, ty, t)
    | n, Cic.Prod (name,so,ty) ->
       let ty' =
        aux (n - 1) (Cic.Rel n::irl) (Some (name, Cic.Decl so)::context) ty
       in
        Cic.Lambda (name,so,ty')
    | _,_ -> assert false
  in
(*prerr_endline ("RIGHTNO = " ^ string_of_int rightno ^ " OUTTY = " ^ CicPp.ppterm outty);*)
   let outsort =
    fst (CicTypeChecker.type_of_aux' [] context outty CicUniv.oblivion_ugraph)
   in
    try aux rightno [] context outsort
    with Nothing_to_do -> outty
(*prerr_endline (CicPp.ppterm outty ^ " <==> " ^ CicPp.ppterm outty');*)
;;

let fix_outtype t =
 let module C = Cic in
 let rec aux context =
  function
     C.Rel _ as t -> t
   | C.Var (uri,exp_named_subst) ->
      let exp_named_subst' =
       List.map (function i,t -> i, (aux context t)) exp_named_subst in
        C.Var (uri,exp_named_subst')
   | C.Implicit _
   | C.Meta _ -> assert false
   | C.Sort _ as t -> t
   | C.Cast (v,t) -> C.Cast (aux context v, aux context t)
   | C.Prod (n,s,t) ->
        C.Prod (n, aux context s, aux ((Some (n, C.Decl s))::context) t)
   | C.Lambda (n,s,t) ->
       C.Lambda (n, aux context s, aux ((Some (n, C.Decl s))::context) t)
   | C.LetIn (n,s,ty,t) ->
      C.LetIn
       (n, aux context s, aux context ty,
        aux ((Some (n, C.Def(s,ty)))::context) t)
   | C.Appl l -> C.Appl (List.map (aux context) l)
   | C.Const (uri,exp_named_subst) ->
      let exp_named_subst' =
       List.map (function i,t -> i, (aux context t)) exp_named_subst
      in
       C.Const (uri,exp_named_subst')
   | C.MutInd (uri,tyno,exp_named_subst) ->
      let exp_named_subst' =
       List.map (function i,t -> i, (aux context 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 (function i,t -> i, (aux context t)) exp_named_subst
      in
       C.MutConstruct (uri, tyno, consno, exp_named_subst')
   | C.MutCase (uri, tyno, outty, term, patterns) ->
      let outty = fix_outty uri tyno term context outty in
       C.MutCase (uri, tyno, aux context outty,
        aux context term, List.map (aux context) patterns)
   | C.Fix (funno, funs) ->
      let tys,_ =
        List.fold_left
          (fun (types,len) (n,_,ty,_) ->
            ((Some (C.Name n,(C.Decl (CicSubstitution.lift len ty)))))::types,
              len+1
	  ) ([],0) funs
      in
       C.Fix (funno,
        List.map
         (fun (name, indidx, ty, bo) ->
           (name, indidx, aux context ty, aux (tys@context) bo)
         ) funs
      )
   | C.CoFix (funno, funs) ->
      let tys,_ =
        List.fold_left
          (fun (types,len) (n,ty,_) ->
            ((Some (C.Name n,(C.Decl (CicSubstitution.lift len ty)))))::types,
              len+1
	  ) ([],0) funs
      in
       C.CoFix (funno,
        List.map
         (fun (name, ty, bo) ->
           (name, aux context ty, aux (tys@context) bo)
         ) funs
       )
 in
  aux [] t
;;

(* we are lambda-lifting also variables that do not occur *)
(* ctx does not distinguish successive blocks of cofix, since there may be no
 *   lambda separating them *)
let convert_term uri t = 
  let rec aux octx (ctx : ctx list) n_fix uri = function
    | Cic.CoFix (k, fl) ->
        let buri = 
          UriManager.uri_of_string 
           (UriManager.buri_of_uri uri^"/"^
            UriManager.name_of_uri uri ^ "___" ^
             string_of_int (List.length ctx)^".con")
        in
        let bctx, fixpoints_tys, tys, _ = 
          List.fold_right 
            (fun (name,ty,_) (ctx, fixpoints, tys, idx) -> 
              let ty, fixpoints_ty = aux octx ctx n_fix uri ty in
              let r = Ref.reference_of_ouri buri(Ref.CoFix idx) in
              Fix (r,name,ty) :: ctx, fixpoints_ty @ fixpoints,ty::tys,idx+1)
            fl ([], [], [], 0)
        in
        let bctx = bctx @ ctx in
        let n_fl = List.length fl in
        let boctx,_ =
         List.fold_left
          (fun (types,len) (n,ty,_) ->
             (Some (Cic.Name n,(Cic.Decl (CicSubstitution.lift len ty)))::types,
              len+1)) (octx,0) fl
        in
        let fl, fixpoints =
          List.fold_right2 
            (fun (name,_,bo) ty  (l,fixpoints) -> 
               let bo, fixpoints_bo = aux boctx bctx n_fl buri bo in
               (([],name,~-1,splat true ctx ty, splat false ctx bo)::l),
               fixpoints_bo @ fixpoints)
            fl tys ([],fixpoints_tys)
        in
        let obj = 
          NUri.nuri_of_ouri buri,0,[],[],
            NCic.Fixpoint (false, fl, (`Generated, `Definition)) 
        in
        splat_args ctx 
         (NCic.Const (Ref.reference_of_ouri buri (Ref.CoFix k)))
         n_fix,
        fixpoints @ [obj]
    | Cic.Fix (k, fl) ->
        let buri = 
          UriManager.uri_of_string 
           (UriManager.buri_of_uri uri^"/"^
            UriManager.name_of_uri uri ^ "___" ^
             string_of_int (List.length ctx)^".con")
        in
        let bad_bctx, fixpoints_tys, tys, _ = 
          List.fold_right 
            (fun (name,recno,ty,_) (bctx, fixpoints, tys, idx) -> 
              let ty, fixpoints_ty = aux octx ctx n_fix uri ty in
              let r =  (* recno is dummy here, must be lifted by the ctx len *)
                Ref.reference_of_ouri buri (Ref.Fix (idx,recno)) 
              in
              Fix (r,name,ty) :: bctx, fixpoints_ty@fixpoints,ty::tys,idx+1)
            fl ([], [], [], 0)
        in
        let _, free, _ = context_tassonomy (bad_bctx @ ctx) in
        let bctx = 
          List.map (function 
            | Fix (Ref.Ref (_,_,Ref.Fix (idx, recno)),name, ty) ->
              Fix (Ref.reference_of_ouri buri(Ref.Fix (idx,recno+free)),name,ty)
            | _ -> assert false) bad_bctx @ ctx
        in
        let n_fl = List.length fl in
        let boctx,_ =
         List.fold_left
          (fun (types,len) (n,_,ty,_) ->
             (Some (Cic.Name n,(Cic.Decl (CicSubstitution.lift len ty)))::types,
              len+1)) (octx,0) fl
        in
        let rno_k = ref 0 in
        let fl, fixpoints,_ =
          List.fold_right2 
            (fun (name,rno,_,bo) ty (l,fixpoints,idx) -> 
               let bo, fixpoints_bo = aux boctx bctx n_fl buri bo in
               let rno = rno + free in
               if idx = k then rno_k := rno;
               (([],name,rno,splat true ctx ty, splat false ctx bo)::l),
               fixpoints_bo @ fixpoints,idx+1)
            fl tys ([],fixpoints_tys,0)
        in
        let obj = 
          NUri.nuri_of_ouri buri,max_int,[],[],
            NCic.Fixpoint (true, fl, (`Generated, `Definition)) 
        in
        splat_args ctx
          (NCic.Const 
            (Ref.reference_of_ouri buri (Ref.Fix (k,!rno_k))))
          n_fix,
        fixpoints @ [obj]
    | Cic.Rel n ->
        let bound, _, primo_ce_dopo_fix = context_tassonomy ctx in
        (match List.nth ctx (n-1) with
        | Fix (r,_,_) when n < primo_ce_dopo_fix -> 
            splat_args_for_rel ctx (NCic.Const r), []
        | Ce _ when n <= bound -> NCic.Rel n, []
        | Fix _ (* BUG 3 fix nested *) 
        | Ce _ -> NCic.Rel (n-n_fix), [])
    | Cic.Lambda (name, (s as old_s), t) ->
        let s, fixpoints_s = aux octx ctx n_fix uri s in
        let ctx = Ce (cn_to_s name, NCic.Decl s) :: ctx in
        let octx = Some (name, Cic.Decl old_s) :: octx in
        let t, fixpoints_t = aux octx ctx n_fix uri t in
        NCic.Lambda (cn_to_s name, s, t), fixpoints_s @ fixpoints_t
    | Cic.Prod (name, (s as old_s), t) ->
        let s, fixpoints_s = aux octx ctx n_fix uri s in
        let ctx = Ce (cn_to_s name, NCic.Decl s) :: ctx in
        let octx = Some (name, Cic.Decl old_s) :: octx in
        let t, fixpoints_t = aux octx ctx n_fix uri t in
        NCic.Prod (cn_to_s name, s, t), fixpoints_s @ fixpoints_t
    | Cic.LetIn (name, (te as old_te), (ty as old_ty), t) ->
        let te, fixpoints_s = aux octx ctx n_fix uri te in
        let ty, fixpoints_ty = aux octx ctx n_fix uri ty in
        let ctx = Ce (cn_to_s name, NCic.Def (te, ty)) :: ctx in
        let octx = Some (name, Cic.Def (old_te, old_ty)) :: octx in
        let t, fixpoints_t = aux octx ctx n_fix uri t in
        NCic.LetIn (cn_to_s name, ty, te, t), 
        fixpoints_s @ fixpoints_t @ fixpoints_ty
    | Cic.Cast (t,ty) ->
        let t, fixpoints_t = aux octx ctx n_fix uri t in
        let ty, fixpoints_ty = aux octx ctx n_fix uri ty in
        NCic.LetIn ("cast", ty, t, NCic.Rel 1), fixpoints_t @ fixpoints_ty
    | Cic.Sort Cic.Prop -> NCic.Sort NCic.Prop,[]
    | Cic.Sort Cic.CProp -> NCic.Sort NCic.CProp,[]
    | Cic.Sort (Cic.Type _) -> NCic.Sort (NCic.Type 0),[] 
    | Cic.Sort Cic.Set -> NCic.Sort (NCic.Type 0),[] 
       (* calculate depth in the univ_graph*)
    | Cic.Appl l -> 
        let l, fixpoints =
          List.fold_right 
             (fun t (l,acc) -> 
               let t, fixpoints = aux octx ctx n_fix uri t in 
               (t::l,fixpoints@acc))
             l ([],[])
        in
        (match l with
        | (NCic.Appl l1)::l2 -> NCic.Appl (l1@l2), fixpoints
        | _ -> NCic.Appl l, fixpoints)
    | Cic.Const (curi, ens) -> 
       aux_ens curi octx ctx n_fix uri ens
        (match fst(CicEnvironment.get_obj CicUniv.oblivion_ugraph curi) with
        | Cic.Constant (_,Some _,_,_,_) ->
               NCic.Const (Ref.reference_of_ouri curi Ref.Def)
        | Cic.Constant (_,None,_,_,_) ->
               NCic.Const (Ref.reference_of_ouri curi Ref.Decl)
        | _ -> assert false)
    | Cic.MutInd (curi, tyno, ens) -> 
       aux_ens curi octx ctx n_fix uri ens
        (NCic.Const (Ref.reference_of_ouri curi (Ref.Ind tyno)))
    | Cic.MutConstruct (curi, tyno, consno, ens) -> 
       aux_ens curi octx ctx n_fix uri ens
        (NCic.Const (Ref.reference_of_ouri curi (Ref.Con (tyno,consno))))
    | Cic.MutCase (curi, tyno, outty, t, branches) ->
        let r = Ref.reference_of_ouri curi (Ref.Ind tyno) in
        let outty, fixpoints_outty = aux octx ctx n_fix uri outty in
        let t, fixpoints_t = aux octx ctx n_fix uri t in
        let branches, fixpoints =
          List.fold_right 
             (fun t (l,acc) -> 
               let t, fixpoints = aux octx ctx n_fix uri t in 
               (t::l,fixpoints@acc))
             branches ([],[])
        in
        NCic.Match (r,outty,t,branches), fixpoints_outty@fixpoints_t@fixpoints
    | Cic.Implicit _ | Cic.Meta _ | Cic.Var _ -> assert false
  and aux_ens curi octx ctx n_fix uri ens he =
   match ens with
      [] -> he,[]
    | _::_ ->
      let params =
       match fst (CicEnvironment.get_obj CicUniv.oblivion_ugraph curi) with
          Cic.Constant (_,_,_,params,_)
        | Cic.InductiveDefinition (_,params,_,_) -> params
        | Cic.Variable _
        | Cic.CurrentProof _ -> assert false
      in
      let ens,objs =
       List.fold_right
        (fun uri (l,objs) ->
          let t = List.assoc uri ens in
          let t,o = aux octx ctx n_fix uri t in
           t::l, o@objs
        ) params ([],[])
      in
       NCic.Appl (he::ens),objs
  in
   aux [] [] 0 uri t
;;

let cook mode vars t =
 let t = fix_outtype t in
 let varsno = List.length vars in
 let t = CicSubstitution.lift varsno t in
 let rec aux n acc l =
  let subst =
   snd(List.fold_left (fun (i,res) uri -> i+1,(uri,Cic.Rel i)::res) (1,[]) acc)
  in
  match l with
     [] -> CicSubstitution.subst_vars subst t
   | uri::uris ->
    let bo,ty =
     match fst (CicEnvironment.get_obj CicUniv.oblivion_ugraph uri) with
        Cic.Variable (_,bo,ty,_,_) ->
         HExtlib.map_option fix_outtype bo, fix_outtype ty
      | _ -> assert false in
    let ty = CicSubstitution.subst_vars subst ty in
    let bo = HExtlib.map_option (CicSubstitution.subst_vars subst) bo in
    let id = Cic.Name (UriManager.name_of_uri uri) in
    let t = aux (n-1) (uri::acc) uris in
     match bo,ty,mode with
        None,ty,`Lambda -> Cic.Lambda (id,ty,t)
      | None,ty,`Pi -> Cic.Prod (id,ty,t)
      | Some bo,ty,_ -> Cic.LetIn (id,bo,ty,t)
 in
  aux varsno [] vars
;;

let convert_obj_aux uri = function
 | Cic.Constant (name, None, ty, vars, _) ->
     let ty = cook `Pi vars ty in
     let nty, fixpoints = convert_term uri ty in
     assert(fixpoints = []);
     NCic.Constant ([], name, None, nty, (`Provided,`Theorem,`Regular)),
     fixpoints
 | Cic.Constant (name, Some bo, ty, vars, _) ->
     let bo = cook `Lambda vars bo in
     let ty = cook `Pi vars ty in
     let nbo, fixpoints_bo = convert_term uri bo in
     let nty, fixpoints_ty = convert_term uri ty in
     assert(fixpoints_ty = []);
     NCic.Constant ([], name, Some nbo, nty, (`Provided,`Theorem,`Regular)),
     fixpoints_bo @ fixpoints_ty
 | Cic.InductiveDefinition (itl,vars,leftno,_) -> 
     let ind = let _,x,_,_ = List.hd itl in x in
     let itl, fix_itl = 
       List.fold_right
         (fun (name, _, ty, cl) (itl,acc) ->
            let ty = cook `Pi vars ty in
            let ty, fix_ty = convert_term uri ty in
            let cl, fix_cl = 
              List.fold_right
               (fun (name, ty) (cl,acc) -> 
                 let ty = cook `Pi vars ty in
                 let ty, fix_ty = convert_term uri ty in
                 ([], name, ty)::cl, acc @ fix_ty)
               cl ([],[])
            in
            ([], name, ty, cl)::itl, fix_ty @ fix_cl @ acc)
         itl ([],[])
     in
     NCic.Inductive(ind, leftno + List.length vars, itl, (`Provided, `Regular)),
     fix_itl
 | Cic.Variable _ 
 | Cic.CurrentProof _ -> assert false
;;

let convert_obj uri obj = 
  let o, fixpoints = convert_obj_aux uri obj in
  let obj = NUri.nuri_of_ouri uri,max_int, [], [], o in
  fixpoints @ [obj]
;;