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

open Printf

let fresh_binder =
  let counter = ref ~-1 in
  function
    | true ->
        incr counter;
        Cic.Name ("elim" ^ string_of_int !counter)
    | _ -> Cic.Anonymous

  (** verifies if a given inductive type occurs in a term in target position *)
let rec recursive uri typeno subst = function
  | Cic.Prod (_, _, target) -> recursive uri typeno subst target
  | Cic.MutInd (uri', typeno', subst')
  | Cic.Appl (Cic.MutInd  (uri', typeno', subst') :: _) ->
      UriManager.eq uri uri' && typeno = typeno' && subst = subst'
(*   | Cic.Appl args -> List.exists (recursive uri typeno subst) args *)
  | _ -> false

let unfold_appl = function
  | Cic.Appl ((Cic.Appl args) :: tl) -> Cic.Appl (args @ tl)
  | t -> t

let rec split l n =
 match (l,n) with
    (l,0) -> ([], l)
  | (he::tl, n) -> let (l1,l2) = split tl (n-1) in (he::l1,l2)
  | (_,_) -> assert false

  (** build elimination principle part related to a single constructor
  * @param paramsno number of Prod to ignore in this constructor (i.e. number of
  * inductive parameters)
  * @param dependent true if we are in the dependent case (i.e. sort <> Prop) *)
let rec delta (uri, typeno, subst) dependent paramsno consno t p args =
  assert (subst = []);
  match t with
  | Cic.MutInd (uri', typeno', subst') when
    UriManager.eq uri uri' && typeno = typeno' && subst = subst' ->
      if dependent then
        (match args with
        | [] -> assert false
        | [arg] -> unfold_appl (Cic.Appl [p; arg])
        | _ -> unfold_appl (Cic.Appl [p; unfold_appl (Cic.Appl args)]))
      else
        p
  | Cic.Appl (Cic.MutInd (uri', typeno', subst') :: tl) when
    UriManager.eq uri uri' && typeno = typeno' && subst = subst' ->
      let (lparams, rparams) = split tl paramsno in
      if dependent then
        (match args with
        | [] -> assert false
        | [arg] -> unfold_appl (Cic.Appl (p :: rparams @ [arg]))
        | _ ->
            unfold_appl (Cic.Appl (p ::
              rparams @ [unfold_appl (Cic.Appl args)])))
      else  (* non dependent *)
        (match rparams with
        | [] -> p
        | _ -> Cic.Appl (p :: rparams))
  | Cic.Prod (binder, src, tgt) ->
      if recursive uri typeno subst src then
        let args = List.map (CicSubstitution.lift 2) args in
        let phi =
          let src = CicSubstitution.lift 1 src in
          delta (uri, typeno, subst) dependent paramsno consno src
            (CicSubstitution.lift 1 p) [Cic.Rel 1]
        in
        let tgt = CicSubstitution.lift 1 tgt in
        Cic.Prod (fresh_binder dependent, src,
          Cic.Prod (Cic.Anonymous, phi,
            delta (uri, typeno, subst) dependent paramsno consno tgt
              (CicSubstitution.lift 2 p) (args @ [Cic.Rel 2])))
      else  (* non recursive *)
        let args = List.map (CicSubstitution.lift 1) args in
        Cic.Prod (fresh_binder dependent, src,
          delta (uri, typeno, subst) dependent paramsno consno tgt
            (CicSubstitution.lift 1 p) (args @ [Cic.Rel 1]))
  | _ -> assert false

let rec strip_left_params consno leftno = function
  | t when leftno = 0 -> t (* no need to lift, the term is (hopefully) closed *)
  | Cic.Prod (_, _, tgt) (* when leftno > 0 *) ->
      (* after stripping the parameters we lift of consno. consno is 1 based so,
      * the first constructor will be lifted by 1 (for P), the second by 2 (1
      * for P and 1 for the 1st constructor), and so on *)
      if leftno = 1 then
        CicSubstitution.lift consno tgt
      else
        strip_left_params consno (leftno - 1) tgt
  | _ -> assert false

let delta (ury, typeno, subst) dependent paramsno consno t p args =
  let t = strip_left_params consno paramsno t in
  delta (ury, typeno, subst) dependent paramsno consno t p args

let rec add_params indno ty eliminator =
  if indno = 0 then
    eliminator
  else
    match ty with
    | Cic.Prod (binder, src, tgt) ->
        Cic.Prod (binder, src, add_params (indno - 1) tgt eliminator)
    | _ -> assert false

let rec mk_rels consno = function
  | 0 -> []
  | n -> Cic.Rel (n+consno) :: mk_rels consno (n-1)

let rec strip_pi = function
  | Cic.Prod (_, _, tgt) -> strip_pi tgt
  | t -> t

let rec count_pi = function
  | Cic.Prod (_, _, tgt) -> count_pi tgt + 1
  | t -> 0

let rec type_of_p sort dependent leftno indty = function
  | Cic.Prod (n, src, tgt) when leftno = 0 ->
      Cic.Prod (n, src, type_of_p sort dependent leftno indty tgt)
  | Cic.Prod (_, _, tgt) -> type_of_p sort dependent (leftno - 1) indty tgt
  | t ->
      if dependent then
        Cic.Prod (Cic.Anonymous, indty, Cic.Sort sort)
      else
        Cic.Sort sort

let rec add_right_pi dependent strip liftno liftfrom rightno indty = function
  | Cic.Prod (_, src, tgt) when strip = 0 ->
      Cic.Prod (fresh_binder true,
        CicSubstitution.lift_from (liftfrom + 1) liftno src,
        add_right_pi dependent strip liftno (liftfrom + 1) rightno indty tgt)
  | Cic.Prod (_, _, tgt) ->
      add_right_pi dependent (strip - 1) liftno liftfrom rightno indty tgt
  | t ->
      if dependent then
        Cic.Prod (fresh_binder dependent,
          CicSubstitution.lift_from (rightno + 1) liftno indty,
          Cic.Appl (Cic.Rel (1 + liftno + rightno) :: mk_rels 0 (rightno + 1)))
      else
        Cic.Prod (Cic.Anonymous,
          CicSubstitution.lift_from (rightno + 1) liftno indty,
          if rightno = 0 then
            Cic.Rel (1 + liftno + rightno)
          else
            Cic.Appl (Cic.Rel (1 + liftno + rightno) :: mk_rels 1 rightno))

let rec add_right_lambda dependent strip liftno liftfrom rightno indty case =
function
  | Cic.Prod (_, src, tgt) when strip = 0 ->
      Cic.Lambda (fresh_binder true,
        CicSubstitution.lift_from (liftfrom + 1) liftno src,
        add_right_lambda dependent strip liftno (liftfrom + 1) rightno indty
          case tgt)
  | Cic.Prod (_, _, tgt) ->
      add_right_lambda dependent (strip - 1) liftno liftfrom rightno indty
        case tgt
  | t ->
      Cic.Lambda (fresh_binder true,
        CicSubstitution.lift_from (rightno + 1) liftno indty,
        case)

exception Failure of string

let string_of_sort = function
  | Cic.Prop -> "Prop"
  | Cic.CProp -> "CProp"
  | Cic.Set -> "Set"
  | Cic.Type _ -> "Type"

let elim_of ?(sort = Cic.Type (CicUniv.fresh ())) uri typeno =
  let (obj, univ) = (CicEnvironment.get_obj uri CicUniv.empty_ugraph) in
  let subst = [] in
  match obj with
  | Cic.InductiveDefinition (indTypes, params, leftno) ->
      let (name, inductive, ty, constructors) =
        try
          List.nth indTypes typeno
        with Failure _ -> assert false
      in
      let paramsno = count_pi ty in (* number of (left or right) parameters *)
      let rightno = paramsno - leftno in
      let dependent = (strip_pi ty <> Cic.Sort Cic.Prop) in
      let conslen = List.length constructors in
      let consno = ref (conslen + 1) in
      if (not dependent) && (sort <> Cic.Prop) && (conslen > 1) then
        raise (Failure (sprintf "can't eliminate from Prop to %s"
          (string_of_sort sort)));
      let indty =
        let indty = Cic.MutInd (uri, typeno, subst) in
        if paramsno = 0 then
          indty
        else
          Cic.Appl (indty :: mk_rels 0 paramsno)
      in
      let mk_constructor consno =
        let constructor = Cic.MutConstruct (uri, typeno, consno, subst) in
        if leftno = 0 then
          constructor
        else
          Cic.Appl (constructor :: mk_rels consno leftno)
      in
      let eliminator =
        let p_ty = type_of_p sort dependent leftno indty ty in
        let final_ty =
          add_right_pi dependent leftno (conslen + 1) 1 rightno indty ty
        in
        Cic.Prod (Cic.Name "P", p_ty,
          (List.fold_right
            (fun (_, constructor) acc ->
              decr consno;
              let p = Cic.Rel !consno in
              Cic.Prod (Cic.Anonymous,
                (delta (uri, typeno, subst) dependent leftno !consno
                  constructor p [mk_constructor !consno]),
                acc))
            constructors final_ty))
      in
      add_params leftno ty eliminator
  | _ -> assert false

let rec branch (uri, typeno, subst) insource paramsno t fix head args =
  assert (subst = []);
  match t with
  | Cic.MutInd (uri', typeno', subst') when
    UriManager.eq uri uri' && typeno = typeno' && subst = subst' ->
      let head = if insource then fix else head in
      (match args with
      | [] -> head
      | _ -> Cic.Appl (head :: args))
  | Cic.Appl (Cic.MutInd (uri', typeno', subst') :: tl) when
    UriManager.eq uri uri' && typeno = typeno' && subst = subst' ->
      let (lparams, rparams) = split tl paramsno in
      (match args with
      | [] when insource && rparams = [] -> fix
      | [] when insource -> Cic.Appl (fix :: rparams)
      | _ when insource -> Cic.Appl (fix :: rparams @ args)
      | _ -> Cic.Appl (head :: rparams @ args))
  | Cic.Prod (binder, src, tgt) ->
      if recursive uri typeno subst src then
        let args = List.map (CicSubstitution.lift 1) args in
        let phi =
          let fix = CicSubstitution.lift 1 fix in
          branch (uri, typeno, subst) true paramsno src fix head
            (args @ [Cic.Rel 1])
        in
        let tgt = CicSubstitution.lift 1 tgt in
        Cic.Lambda (fresh_binder true, src,
          branch (uri, typeno, subst) insource paramsno tgt
            (CicSubstitution.lift 1 fix) (CicSubstitution.lift 1 head)
            (args @ [Cic.Rel 1; phi]))
      else  (* non recursive *)
        let args = List.map (CicSubstitution.lift 1) args in
        Cic.Lambda (fresh_binder true, src,
          branch (uri, typeno, subst) insource paramsno tgt
          (CicSubstitution.lift 1 fix) (CicSubstitution.lift 1 head)
            (args @ [Cic.Rel 1]))
  | _ -> assert false

let branch (uri, typeno, subst) insource liftno paramsno t fix head args =
  let t = strip_left_params liftno paramsno t in
  branch (uri, typeno, subst) insource paramsno t fix head args

let body_of ?(sort = Cic.Type (CicUniv.fresh ())) uri typeno =
  let (obj, univ) = (CicEnvironment.get_obj uri CicUniv.empty_ugraph) in
  let subst = [] in
  match obj with
  | Cic.InductiveDefinition (indTypes, params, leftno) ->
      let (name, inductive, ty, constructors) =
        try
          List.nth indTypes typeno
        with Failure _ -> assert false
      in
      let paramsno = count_pi ty in (* number of (left or right) parameters *)
      let rightno = paramsno - leftno in
      let dependent = (strip_pi ty <> Cic.Sort Cic.Prop) in
      let conslen = List.length constructors in
      let consno = ref (conslen + 1) in
      if (not dependent) && (sort <> Cic.Prop) && (conslen > 1) then
        raise (Failure (sprintf "can't eliminate from Prop to %s"
          (string_of_sort sort)));
      let indty =
        let indty = Cic.MutInd (uri, typeno, subst) in
        if paramsno = 0 then
          indty
        else
          Cic.Appl (indty :: mk_rels 0 paramsno)
      in
      let mk_constructor consno =
        let constructor = Cic.MutConstruct (uri, typeno, consno, subst) in
        if leftno = 0 then
          constructor
        else
          Cic.Appl (constructor :: mk_rels consno leftno)
      in
      let eliminator =
        let p_ty = type_of_p sort dependent leftno indty ty in
        let final_ty =
          add_right_pi dependent leftno (conslen + 1) 1 rightno indty ty
        in
        let fix = Cic.Rel (rightno + 2) in
        let (_, branches) =
          List.fold_right
            (fun (_, ty) (shift, branches) ->
              let head = Cic.Rel (rightno + shift + 2) in
              let b =
                branch (uri, typeno, subst) false (rightno + conslen + 3) leftno
                  ty fix head []
              in
              (shift + 1,  b :: branches))
            constructors (1, [])
        in
        let case =
          Cic.MutCase (uri, typeno, Cic.Rel (conslen + rightno + 3), Cic.Rel 1,
            branches)
        in
        let fix_body =
          add_right_lambda dependent leftno (conslen + 1) 1 rightno
            indty case ty
        in
        let fix = Cic.Fix (0, ["f", 0, final_ty, fix_body]) in
        Cic.Lambda (Cic.Name "P", p_ty,
          (List.fold_right
            (fun (_, constructor) acc ->
              decr consno;
              let p = Cic.Rel !consno in
              Cic.Lambda (fresh_binder true,
                (delta (uri, typeno, subst) dependent leftno !consno
                  constructor p [mk_constructor !consno]),
                acc))
            constructors fix))
      in
      add_params leftno ty eliminator
  | _ -> assert false